From 1b6c59cc98c014d1b45859fcb29116a78de77d76 Mon Sep 17 00:00:00 2001 From: Tai An Date: Mon, 13 Apr 2026 23:44:12 -0700 Subject: [PATCH 1/3] 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"): From 186d04e1c8f27ba8ddeff266fbc512d64362e474 Mon Sep 17 00:00:00 2001 From: Tai An Date: Tue, 21 Apr 2026 18:25:14 -0700 Subject: [PATCH 2/3] fix(profiling): make annotate_pipeline transient by returning restore() callable Class-level patches are saved before application and the function now returns a restore() callable that undoes them, making the patches explicitly transient. This addresses reviewer feedback that class-level patches without cleanup are non-transient. --- examples/profiling/profiling_utils.py | 40 +++++++++++++++++++++------ 1 file changed, 31 insertions(+), 9 deletions(-) diff --git a/examples/profiling/profiling_utils.py b/examples/profiling/profiling_utils.py index 04d7b38a5e56..07c7d6f2810e 100644 --- a/examples/profiling/profiling_utils.py +++ b/examples/profiling/profiling_utils.py @@ -29,14 +29,18 @@ def annotate_pipeline(pipe): """Apply profiler annotations to key pipeline methods. Monkey-patches bound methods so they appear as named spans in the trace. - Non-invasive — no source modifications required. + 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 + `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. + + Returns: + A restore callable that undoes all patches and restores the + original method definitions, making the annotation transient. """ annotations = [ ("transformer", "forward", "transformer_forward"), @@ -45,6 +49,8 @@ def annotate_pipeline(pipe): ("scheduler", "step", "scheduler_step"), ] + saved = [] # (target, method_name, original_value, is_class_patch) + # Annotate sub-component methods for component_name, method_name, label in annotations: component = getattr(pipe, component_name, None) @@ -54,20 +60,36 @@ def annotate_pipeline(pipe): if method is None: continue 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)) + # Patch at the class level so the descriptor protocol correctly + # re-binds the wrapper to whichever instance accesses it. This + # prevents instance-isolation bugs when a component is + # shallow-copied. The original class attribute is saved so the + # patch can be reversed when restore() is called. + cls = type(component) + original = cls.__dict__.get(method_name) + setattr(cls, method_name, annotate(method.__func__, label)) + saved.append((cls, method_name, original, True)) else: + original = component.__dict__.get(method_name) setattr(component, method_name, annotate(method, label)) + saved.append((component, method_name, original, False)) # Annotate pipeline-level methods if hasattr(pipe, "encode_prompt"): + original = pipe.__dict__.get("encode_prompt") pipe.encode_prompt = annotate(pipe.encode_prompt, "encode_prompt") + saved.append((pipe, "encode_prompt", original, False)) + + def restore(): + """Undo all patches applied by annotate_pipeline, restoring originals.""" + for target, name, original, is_class in saved: + if original is None: + if name in vars(target): + delattr(target, name) + else: + setattr(target, name, original) + return restore def flush(): gc.collect() From 09c1e1e3a17abfbe373a9e06650e5254bc73cabc Mon Sep 17 00:00:00 2001 From: Tai An Date: Tue, 21 Apr 2026 18:26:03 -0700 Subject: [PATCH 3/3] fix(profiling): call restore() in PipelineProfiler.run() after profiling Store the restore callable from annotate_pipeline() and call it at the end of run() so class-level method patches are cleaned up. --- examples/profiling/profiling_utils.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/examples/profiling/profiling_utils.py b/examples/profiling/profiling_utils.py index 07c7d6f2810e..5a2ffd6a7c8b 100644 --- a/examples/profiling/profiling_utils.py +++ b/examples/profiling/profiling_utils.py @@ -169,7 +169,9 @@ def setup_pipeline(self, annotate=True): pipe.set_progress_bar_config(disable=True) if annotate: - annotate_pipeline(pipe) + self._restore_annotations = annotate_pipeline(pipe) + else: + self._restore_annotations = None return pipe def run(self): @@ -215,7 +217,9 @@ def run(self): ) ) - # Cleanup + # Cleanup -- restore patched methods so class-level patches don't persist + if self._restore_annotations is not None: + self._restore_annotations() pipe.to("cpu") del pipe flush()