scripts/_probe_encode.py

"""Decide how the encoder is loaded and released.

The Qwen3 text encoder is ~8 GB in bfloat16 - the single biggest allocation of a
CPU render, and it is only needed once, before the denoising loop. If
`Flux2KleinPipeline.encode_prompt` can produce the embeds the loop consumes, the
encoder can be unloaded before the transformer runs and the peak resident set is
max(8 GB, 3 GB) instead of the sum.

Run:  python scripts/_probe_encode.py
"""
import inspect
import re

from diffusers import Flux2KleinPipeline

cls = Flux2KleinPipeline
print("has encode_prompt:", hasattr(cls, "encode_prompt"))
if hasattr(cls, "encode_prompt"):
    print("signature:", inspect.signature(cls.encode_prompt))

src = inspect.getsource(cls.__call__)
for marker in ("prompt_embeds", "self.text_encoder", "self.tokenizer", "negative_prompt_embeds"):
    hits = len(re.findall(re.escape(marker), src))
    print("  __call__ references %-22s %d time(s)" % (marker, hits))

lines = inspect.getsource(cls).splitlines()
for i, line in enumerate(lines):
    if "def encode_prompt" in line:
        print("encode_prompt source starts at relative line", i)
        for offset in range(i, min(i + 40, len(lines))):
            print("    " + lines[offset])
        break