From 1b6c59cc98c014d1b45859fcb29116a78de77d76 Mon Sep 17 00:00:00 2001 From: Tai An Date: Mon, 13 Apr 2026 23:44:12 -0700 Subject: [PATCH] fix(profiling): patch sub-component methods at class level to prevent instance-isolation bugs --- examples/profiling/profiling_utils.py | 19 ++++++++++++++++++- 1 file changed, 18 insertions(+), 1 deletion(-) diff --git a/examples/profiling/profiling_utils.py b/examples/profiling/profiling_utils.py index 1c7d59d42fde..04d7b38a5e56 100644 --- a/examples/profiling/profiling_utils.py +++ b/examples/profiling/profiling_utils.py @@ -1,5 +1,6 @@ import functools import gc +import inspect import logging import os from dataclasses import dataclass, field @@ -29,6 +30,13 @@ def annotate_pipeline(pipe): Monkey-patches bound methods so they appear as named spans in the trace. Non-invasive — no source modifications required. + + Sub-component methods are patched at the **class level** (via + setattr(type(component), ...)) rather than on the instance. This + ensures Python's descriptor protocol re-binds the wrapper to whichever + instance accesses it, so shallow-copied components (e.g. the duplicated + audio_scheduler inside LTX2) call their own logic rather than the + original instance's. """ annotations = [ ("transformer", "forward", "transformer_forward"), @@ -45,7 +53,16 @@ def annotate_pipeline(pipe): method = getattr(component, method_name, None) if method is None: continue - setattr(component, method_name, annotate(method, label)) + if inspect.ismethod(method): + # Wrap the underlying function and patch at the class level so + # that the descriptor protocol correctly rebinds the wrapper to + # whichever instance accesses it. This prevents instance- + # isolation bugs when a component is shallow-copied after + # annotation (e.g. audio_scheduler = copy.copy(self.scheduler) + # in the LTX2 pipeline). + setattr(type(component), method_name, annotate(method.__func__, label)) + else: + setattr(component, method_name, annotate(method, label)) # Annotate pipeline-level methods if hasattr(pipe, "encode_prompt"):