diffusers-workflow 0.4.0__py3-none-any.whl

dw/teacache.py

@@ -0,0 +1,381 @@
+"""
+TeaCache - Training-free inference acceleration for diffusion transformers.
+
+Caches intermediate transformer computations and skips redundant steps
+when the input hasn't changed significantly between timesteps.
+
+Based on: https://github.com/ali-vilab/TeaCache
+Adapted from: https://github.com/Teriks/dgenerate (Apache 2.0)
+
+Implemented: Flux (FluxTransformer2DModel)
+Registry includes: Mochi, LTX-Video, CogVideoX, Lumina2, HunyuanVideo, Wan2.1
+(these require custom forward functions to be added)
+
+Each model requires a custom forward function because transformer architectures
+differ. The core caching algorithm is the same: extract a signal from the first
+block's normalization, compare via polynomial rescaling, skip if below threshold.
+"""
+
+import json
+import typing
+import logging
+from pathlib import Path
+from contextlib import contextmanager
+
+import torch
+import numpy as np
+from diffusers.models.modeling_outputs import Transformer2DModelOutput
+from diffusers.utils import (
+    USE_PEFT_BACKEND,
+    scale_lora_layers,
+    unscale_lora_layers,
+)
+
+logger = logging.getLogger("dw")
+
+
+# ---------------------------------------------------------------------------
+# Model registry loaded from JSON
+# ---------------------------------------------------------------------------
+
+_REGISTRY_PATH = Path(__file__).parent / "teacache_models.json"
+
+
+def _load_registry():
+    """Load the model registry from the JSON file."""
+    with open(_REGISTRY_PATH) as f:
+        return json.load(f)
+
+
+def _get_model_info(transformer, variant=None):
+    """Look up model info from registry.
+
+    Args:
+        transformer: The transformer model instance.
+        variant: Optional explicit variant name (e.g., "wan2.1_t2v_1.3b").
+            If None, uses class_defaults mapping.
+
+    Returns:
+        dict with coefficients, default_threshold, threshold_guide.
+    """
+    registry = _load_registry()
+    class_name = transformer.__class__.__name__
+
+    if variant is not None:
+        info = registry["models"].get(variant)
+        if info is None:
+            available = ", ".join(registry["models"].keys())
+            raise ValueError(
+                f"TeaCache variant '{variant}' not found. Available: {available}"
+            )
+        return info
+
+    # Look up default variant for this class
+    default_variant = registry["class_defaults"].get(class_name)
+    if default_variant is None:
+        supported_classes = ", ".join(registry["class_defaults"].keys())
+        raise ValueError(
+            f"TeaCache does not support {class_name}. Supported: {supported_classes}"
+        )
+
+    return registry["models"][default_variant]
+
+
+# ---------------------------------------------------------------------------
+# Forward function factories, one per supported transformer architecture.
+# ---------------------------------------------------------------------------
+
+
+def _create_flux_teacache_forward(num_inference_steps, rel_l1_thresh, coefficients):
+    """Create TeaCache forward for FluxTransformer2DModel."""
+    cnt = 0
+    accumulated_rel_l1_distance = 0
+    previous_modulated_input = None
+    previous_residual = None
+    previous_timestep = None
+    rescale_func = np.poly1d(coefficients)
+
+    def teacache_forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor = None,
+        pooled_projections: torch.Tensor = None,
+        timestep: torch.LongTensor = None,
+        img_ids: torch.Tensor = None,
+        txt_ids: torch.Tensor = None,
+        guidance: torch.Tensor = None,
+        joint_attention_kwargs: typing.Optional[typing.Dict[str, typing.Any]] = None,
+        controlnet_block_samples=None,
+        controlnet_single_block_samples=None,
+        return_dict: bool = True,
+        controlnet_blocks_repeat: bool = False,
+    ) -> typing.Union[torch.FloatTensor, Transformer2DModelOutput]:
+        nonlocal \
+            cnt, \
+            accumulated_rel_l1_distance, \
+            previous_modulated_input, \
+            previous_residual, \
+            previous_timestep
+
+        # TeaCache assumes exactly one transformer forward call per denoising
+        # step. Pipelines running true classifier-free guidance (e.g. Flux with
+        # negative_prompt + true_cfg_scale > 1) call the transformer twice per
+        # step -- once for the conditional pass and once for the unconditional
+        # pass -- using the identical timestep both times. That second call
+        # would silently share/corrupt previous_modulated_input and
+        # previous_residual across the two passes, so detect it and fail loudly
+        # instead of producing a corrupted image.
+        if (
+            timestep is not None
+            and previous_timestep is not None
+            and timestep.shape == previous_timestep.shape
+            and torch.equal(timestep, previous_timestep)
+        ):
+            raise RuntimeError(
+                "TeaCache does not support true classifier-free guidance "
+                "(negative_prompt with true_cfg_scale > 1); disable one of them. "
+                "Detected two transformer forward calls with an identical "
+                "timestep within a single denoising step, which would corrupt "
+                "TeaCache's cached state."
+            )
+        if timestep is not None:
+            previous_timestep = timestep.detach().clone()
+
+        if joint_attention_kwargs is not None:
+            joint_attention_kwargs = joint_attention_kwargs.copy()
+            lora_scale = joint_attention_kwargs.pop("scale", 1.0)
+        else:
+            lora_scale = 1.0
+
+        if USE_PEFT_BACKEND:
+            scale_lora_layers(self, lora_scale)
+
+        hidden_states = self.x_embedder(hidden_states)
+
+        timestep = timestep.to(hidden_states.dtype) * 1000
+        if guidance is not None:
+            guidance = guidance.to(hidden_states.dtype) * 1000
+
+        temb = (
+            self.time_text_embed(timestep, pooled_projections)
+            if guidance is None
+            else self.time_text_embed(timestep, guidance, pooled_projections)
+        )
+        encoder_hidden_states = self.context_embedder(encoder_hidden_states)
+
+        if txt_ids.ndim == 3:
+            txt_ids = txt_ids[0]
+        if img_ids.ndim == 3:
+            img_ids = img_ids[0]
+
+        ids = torch.cat((txt_ids, img_ids), dim=0)
+        image_rotary_emb = self.pos_embed(ids)
+
+        if (
+            joint_attention_kwargs is not None
+            and "ip_adapter_image_embeds" in joint_attention_kwargs
+        ):
+            ip_adapter_image_embeds = joint_attention_kwargs.pop(
+                "ip_adapter_image_embeds"
+            )
+            ip_hidden_states = self.encoder_hid_proj(ip_adapter_image_embeds)
+            joint_attention_kwargs.update({"ip_hidden_states": ip_hidden_states})
+
+        # TeaCache: extract cache signal from first block's normalization
+        inp = hidden_states.clone()
+        temb_ = temb.clone()
+        modulated_inp, gate_msa, shift_mlp, scale_mlp, gate_mlp = (
+            self.transformer_blocks[0].norm1(inp, emb=temb_)
+        )
+
+        # Decide whether to compute or reuse cached result
+        if cnt == 0 or cnt == num_inference_steps - 1:
+            should_calc = True
+            accumulated_rel_l1_distance = 0
+        else:
+            relative_diff = (
+                (
+                    (modulated_inp - previous_modulated_input).abs().mean()
+                    / previous_modulated_input.abs().mean()
+                )
+                .cpu()
+                .item()
+            )
+            accumulated_rel_l1_distance += rescale_func(relative_diff)
+
+            if accumulated_rel_l1_distance < rel_l1_thresh:
+                should_calc = False
+            else:
+                should_calc = True
+                accumulated_rel_l1_distance = 0
+
+        previous_modulated_input = modulated_inp
+        cnt += 1
+        if cnt == num_inference_steps:
+            cnt = 0
+
+        if not should_calc:
+            hidden_states += previous_residual
+        else:
+            ori_hidden_states = hidden_states.clone()
+
+            # No gradient-checkpointing branch: this forward only runs under
+            # Pipeline.run's @torch.inference_mode(), so grads are never enabled
+            for index_block, block in enumerate(self.transformer_blocks):
+                encoder_hidden_states, hidden_states = block(
+                    hidden_states=hidden_states,
+                    encoder_hidden_states=encoder_hidden_states,
+                    temb=temb,
+                    image_rotary_emb=image_rotary_emb,
+                    joint_attention_kwargs=joint_attention_kwargs,
+                )
+
+                if controlnet_block_samples is not None:
+                    interval_control = len(self.transformer_blocks) / len(
+                        controlnet_block_samples
+                    )
+                    interval_control = int(np.ceil(interval_control))
+                    if controlnet_blocks_repeat:
+                        hidden_states = (
+                            hidden_states
+                            + controlnet_block_samples[
+                                index_block % len(controlnet_block_samples)
+                            ]
+                        )
+                    else:
+                        hidden_states = (
+                            hidden_states
+                            + controlnet_block_samples[index_block // interval_control]
+                        )
+
+            for index_block, block in enumerate(self.single_transformer_blocks):
+                encoder_hidden_states, hidden_states = block(
+                    hidden_states=hidden_states,
+                    encoder_hidden_states=encoder_hidden_states,
+                    temb=temb,
+                    image_rotary_emb=image_rotary_emb,
+                    joint_attention_kwargs=joint_attention_kwargs,
+                )
+
+                if controlnet_single_block_samples is not None:
+                    interval_control = len(self.single_transformer_blocks) / len(
+                        controlnet_single_block_samples
+                    )
+                    interval_control = int(np.ceil(interval_control))
+                    hidden_states[:, encoder_hidden_states.shape[1] :, ...] = (
+                        hidden_states[:, encoder_hidden_states.shape[1] :, ...]
+                        + controlnet_single_block_samples[
+                            index_block // interval_control
+                        ]
+                    )
+
+            previous_residual = hidden_states - ori_hidden_states
+
+        hidden_states = self.norm_out(hidden_states, temb)
+        output = self.proj_out(hidden_states)
+
+        if USE_PEFT_BACKEND:
+            unscale_lora_layers(self, lora_scale)
+
+        if not return_dict:
+            return (output,)
+
+        return Transformer2DModelOutput(sample=output)
+
+    return teacache_forward
+
+
+# Map transformer class names to their forward factory functions.
+# Models in the JSON registry without a factory here will get an informative error.
+_FORWARD_FACTORIES = {
+    "FluxTransformer2DModel": _create_flux_teacache_forward,
+}
+
+
+# ---------------------------------------------------------------------------
+# Public API
+# ---------------------------------------------------------------------------
+
+
+@contextmanager
+def teacache_context(
+    pipeline, num_inference_steps, rel_l1_thresh=None, coefficients=None, variant=None
+):
+    """Context manager that enables TeaCache on a pipeline's transformer.
+
+    Auto-detects the transformer type and applies the appropriate
+    TeaCache forward function. Restores original forward on exit.
+
+    Args:
+        pipeline: A DiffusionPipeline with a .transformer attribute
+        num_inference_steps: Number of inference steps (must match pipeline call)
+        rel_l1_thresh: Cache threshold override. If None, uses model default.
+            Higher = more speedup, more quality loss.
+        coefficients: Polynomial coefficients override. If None, uses model default.
+            List of 5 floats for np.poly1d rescaling function.
+        variant: Explicit model variant name (e.g., "wan2.1_t2v_1.3b").
+            Required when a transformer class has multiple variants (CogVideoX, Wan).
+            If None, uses class_defaults from the registry.
+    """
+    transformer = pipeline.transformer
+    class_name = transformer.__class__.__name__
+
+    # Look up model info from registry
+    model_info = _get_model_info(transformer, variant)
+
+    # Check we have a forward implementation for this class
+    factory = _FORWARD_FACTORIES.get(class_name)
+    if factory is None:
+        supported = ", ".join(_FORWARD_FACTORIES.keys())
+        raise ValueError(
+            f"No TeaCache forward implementation for {class_name}. "
+            f"Implemented: {supported}. "
+            f"The model is in the registry but needs a custom forward function."
+        )
+
+    # Use overrides or defaults
+    if rel_l1_thresh is None:
+        rel_l1_thresh = model_info["default_threshold"]
+    if coefficients is None:
+        coefficients = model_info["coefficients"]
+
+    teacache_forward_fn = factory(num_inference_steps, rel_l1_thresh, coefficients)
+
+    # accelerate's enable_model_cpu_offload/enable_sequential_cpu_offload installs
+    # an AlignDevicesHook via add_hook_to_module (accelerate/hooks.py), which
+    # replaces transformer.forward with a wrapper closing over module and the
+    # true original forward (stashed as transformer._old_forward, still bound to
+    # the instance). That wrapper is what moves the module's weights to the
+    # execution device in its pre_forward before calling _old_forward. If we
+    # clobber transformer.forward like the no-hook path below, we remove that
+    # wrapper entirely: the CPU-resident module then receives CUDA inputs and
+    # raises a device-mismatch RuntimeError. Instead, when a hook is present we
+    # wrap what the hook considers "the real forward" -- _old_forward -- so the
+    # call chain stays hook.forward -> pre_forward (places weights) ->
+    # teacache_forward -> post_forward.
+    has_hook = hasattr(transformer, "_hf_hook") and hasattr(transformer, "_old_forward")
+
+    if has_hook:
+        original_forward = transformer._old_forward
+        transformer._old_forward = teacache_forward_fn.__get__(
+            transformer, transformer.__class__
+        )
+    else:
+        original_forward = transformer.forward
+        transformer.forward = teacache_forward_fn.__get__(
+            transformer, transformer.__class__
+        )
+
+    logger.info(
+        f"TeaCache enabled for {class_name}: "
+        f"steps={num_inference_steps}, threshold={rel_l1_thresh}"
+    )
+
+    try:
+        yield pipeline
+    finally:
+        if has_hook:
+            transformer._old_forward = original_forward
+        else:
+            transformer.forward = original_forward
+        logger.debug("TeaCache disabled, original forward restored")

dw/teacache_models.json

@@ -0,0 +1,99 @@
+{
+    "$comment": "TeaCache model registry. Coefficients are 4th-degree polynomial coefficients for np.poly1d rescaling. Source: https://github.com/ali-vilab/TeaCache",
+    "models": {
+        "flux": {
+            "transformer_class": "FluxTransformer2DModel",
+            "coefficients": [4.98651651e+02, -2.83781631e+02, 5.58554382e+01, -3.82021401e+00, 2.64230861e-01],
+            "default_threshold": 0.6,
+            "threshold_guide": "0.25=~1.5x, 0.4=~1.8x, 0.6=~2.0x, 0.8=~2.25x"
+        },
+        "hunyuan_video": {
+            "transformer_class": "HunyuanVideoTransformer3DModel",
+            "coefficients": [7.33226126e+02, -4.01131952e+02, 6.75869174e+01, -3.14987800e+00, 9.61237896e-02],
+            "default_threshold": 0.15,
+            "threshold_guide": "0.1=~1.6x, 0.15=~2.1x"
+        },
+        "mochi": {
+            "transformer_class": "MochiTransformer3DModel",
+            "coefficients": [-3.51241319e+03, 8.11675948e+02, -6.09400215e+01, 2.42429681e+00, 3.05291719e-03],
+            "default_threshold": 0.09,
+            "threshold_guide": "0.06=~1.5x, 0.09=~2.1x"
+        },
+        "ltx_video": {
+            "transformer_class": "LTXVideoTransformer3DModel",
+            "coefficients": [2.14700694e+01, -1.28016453e+01, 2.31279151e+00, 7.92487521e-01, 9.69274326e-03],
+            "default_threshold": 0.05,
+            "threshold_guide": "0.03=~1.6x, 0.05=~2.1x"
+        },
+        "cogvideox_2b": {
+            "transformer_class": "CogVideoXTransformer3DModel",
+            "coefficients": [-3.10658903e+01, 2.54732368e+01, -5.92380459e+00, 1.75769064e+00, -3.61568434e-03],
+            "default_threshold": 0.1,
+            "threshold_guide": "0.1=~1.3x, 0.2=~1.8x"
+        },
+        "cogvideox_5b": {
+            "transformer_class": "CogVideoXTransformer3DModel",
+            "coefficients": [-1.53880483e+03, 8.43202495e+02, -1.34363087e+02, 7.97131516e+00, -5.23162339e-02],
+            "default_threshold": 0.1,
+            "threshold_guide": "0.1=~1.3x, 0.2=~1.8x"
+        },
+        "cogvideox1.5_5b": {
+            "transformer_class": "CogVideoXTransformer3DModel",
+            "coefficients": [2.50210439e+02, -1.65061612e+02, 3.57804877e+01, -7.81551492e-01, 3.58559703e-02],
+            "default_threshold": 0.2,
+            "threshold_guide": "0.1=~1.3x, 0.2=~1.8x, 0.3=~2.1x"
+        },
+        "cogvideox1.5_5b_i2v": {
+            "transformer_class": "CogVideoXTransformer3DModel",
+            "coefficients": [1.22842302e+02, -1.04088754e+02, 2.62981677e+01, -3.06009921e-01, 3.71213220e-02],
+            "default_threshold": 0.1,
+            "threshold_guide": "0.1=~1.5x, 0.2=~2.2x, 0.3=~2.7x"
+        },
+        "lumina2": {
+            "transformer_class": "Lumina2Transformer2DModel",
+            "coefficients": [393.76566581, -603.50993606, 209.10239044, -23.00726601, 0.86377344],
+            "default_threshold": 0.3,
+            "threshold_guide": "0.2=~1.25x, 0.3=~1.56x, 0.4=~2.08x, 0.5=~2.5x"
+        },
+        "lumina2_v2": {
+            "transformer_class": "Lumina2Transformer2DModel",
+            "coefficients": [225.7042019806413, -608.8453716535591, 304.1869942338369, 124.21267720116742, -1.4089066892956552],
+            "default_threshold": 0.3,
+            "threshold_guide": "0.2=~1.5x, 0.3=~1.6x, 0.5=~1.8x, 1.1=~2.1x"
+        },
+        "wan2.1_t2v_1.3b": {
+            "transformer_class": "WanTransformer3DModel",
+            "coefficients": [2.39676752e+03, -1.31110545e+03, 2.01331979e+02, -8.29855975e+00, 1.37887774e-01],
+            "default_threshold": 0.08,
+            "threshold_guide": "0.05=~1.5x, 0.07=~1.6x, 0.08=~2.0x"
+        },
+        "wan2.1_t2v_14b": {
+            "transformer_class": "WanTransformer3DModel",
+            "coefficients": [-5784.54975374, 5449.50911966, -1811.16591783, 256.27178429, -13.02252404],
+            "default_threshold": 0.2,
+            "threshold_guide": "0.14=~1.4x, 0.15=~1.8x, 0.2=~2.0x"
+        },
+        "wan2.1_i2v_480p": {
+            "transformer_class": "WanTransformer3DModel",
+            "coefficients": [-3.02331670e+02, 2.23948934e+02, -5.25463970e+01, 5.87348440e+00, -2.01973289e-01],
+            "default_threshold": 0.26,
+            "threshold_guide": "0.13=~1.6x, 0.19=~2.0x, 0.26=~2.5x"
+        },
+        "wan2.1_i2v_720p": {
+            "transformer_class": "WanTransformer3DModel",
+            "coefficients": [-114.36346466, 65.26524496, -18.82220707, 4.91518089, -0.23412683],
+            "default_threshold": 0.3,
+            "threshold_guide": "0.18=~1.7x, 0.2=~1.9x, 0.3=~2.4x"
+        }
+    },
+    "class_defaults": {
+        "$comment": "Default variant used when looking up by transformer class name (no explicit variant specified)",
+        "FluxTransformer2DModel": "flux",
+        "HunyuanVideoTransformer3DModel": "hunyuan_video",
+        "MochiTransformer3DModel": "mochi",
+        "LTXVideoTransformer3DModel": "ltx_video",
+        "CogVideoXTransformer3DModel": "cogvideox1.5_5b",
+        "Lumina2Transformer2DModel": "lumina2",
+        "WanTransformer3DModel": "wan2.1_t2v_14b"
+    }
+}

dw/test.py

@@ -0,0 +1,29 @@
+import os
+from .workflow import workflow_from_file
+from . import startup
+
+
+def main():
+    workflow = workflow_from_file(
+        os.path.join(
+            os.path.dirname(os.path.abspath(__file__)), "workflows", "test.json"
+        ),
+        "./outputs",
+    )
+
+    try:
+        startup("DEBUG")
+        workflow.validate()
+    except Exception as e:
+        print(f"Error validating workflow: {e}")
+        exit(1)
+
+    try:
+        workflow.run({})
+    except Exception as e:
+        print(f"Error running workflow: {e}")
+        exit(1)
+
+
+if __name__ == "__main__":
+    main()