Skip to content
Open
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
19 changes: 18 additions & 1 deletion examples/profiling/profiling_utils.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
import functools
import gc
import inspect
import logging
import os
from dataclasses import dataclass, field
Expand Down Expand Up @@ -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"),
Expand All @@ -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
Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Class-level patches should be avoided as they're not transient.

# 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"):
Expand Down
Loading