Migrate to canonical diffusers PR 14371 + public MiniMaxAI/MiniMax-H3 weights, add prompt upsampling toggle

README.md

@@ -17,20 +17,19 @@ Joint video **and** soundtrack out of a single denoising pass, at **bfloat16 wit
 
 This Space is the denoising half: the 61.73 GiB transformer and the two autoencoders. The 62.14 GiB Qwen3-VL
 conditioner runs in
-[`minimax-h3-conditioner`](https://huggingface.co/spaces/diffusers-internal-dev/minimax-h3-conditioner), which this
-Space calls over the gradio API for every request.
+[`qwen3vl-conditioner`](https://huggingface.co/spaces/multimodalart/qwen3vl-conditioner), which this
+Space calls over the gradio API for every request. The weights are the public
+[`MiniMaxAI/MiniMax-H3`](https://huggingface.co/MiniMaxAI/MiniMax-H3) diffusers checkpoint.
 
 ## Why split
 
 MiniMax-H3 is 195.9 GiB in bfloat16 and a ZeroGPU Space is evicted at **150 GB of storage**. An unquantized single
-Space is therefore impossible — the existing demos
-([`minimax-h3`](https://huggingface.co/spaces/diffusers-internal-dev/minimax-h3),
-[`-fp8`](https://huggingface.co/spaces/diffusers-internal-dev/minimax-h3-fp8)) run NVFP4 and float8 weights for that
-reason alone. Cut the `MiniMaxH3Blocks` sequence at its `text_encoder` step and both halves fit unquantized:
+Space is therefore impossible, which is why quantized demos of it run NVFP4 or float8 weights. Cut the
+`MiniMaxH3Blocks` sequence at its `text_encoder` step and both halves fit unquantized:
 
 | Space | Subfolders | Download | Resident |
 |---|---|---|---|
-| [`minimax-h3-conditioner`](https://huggingface.co/spaces/diffusers-internal-dev/minimax-h3-conditioner) | `text_encoder/` + `tokenizer/` + `processor/` | 66.7 GB | 62.15 GiB bf16 |
+| [`qwen3vl-conditioner`](https://huggingface.co/spaces/multimodalart/qwen3vl-conditioner) | `text_encoder/` + `tokenizer/` + `processor/` | 66.7 GB | 62.15 GiB bf16 |
 | this one | `transformer/` + `vae/` + `audio_vae/` | 77.3 GB | 61.73 GiB bf16 + 10.43 GiB float32 |
 
 Besides the quality argument, unquantized weights are the ones AoTI can export; an NVFP4 checkpoint cannot be
@@ -59,19 +58,22 @@ Space pays, where the whole cost is the traffic auto-offload has to move.
 
 ## How the split is expressed
 
-`MiniMaxH3Blocks` is a `SequentialPipelineBlocks` of eight steps:
+`MiniMaxH3Blocks` is a `SequentialPipelineBlocks` whose branches are picked per request — and per `workflow=` — from
+the inputs:
 
 ```
-setup -> text_encoder -> vae_encoder -> prepare_layout -> prepare_latents -> set_timesteps -> denoise -> decode
+before_encode -> text_encoder -> vae_encoder -> denoise -> after_denoise -> decode
 ```
 
+where `denoise` is itself `prepare_layout -> prepare_latents -> set_timesteps -> denoise`.
+
 `h3_split_blocks.py` subclasses it with the `text_encoder` step removed. Dropping the step drops the three
 components it declares, so `load_components` resolves `transformer` / `vae` / `audio_vae` / the two schedulers out of
 the shared `modular_model_index.json` and never fetches the conditioner — and `prompt_embeds` and `text_token_tags`
 become ordinary required inputs of the pipeline call:
 
 ```py
-pipe = MiniMaxH3GeneratorBlocks().init_pipeline("diffusers-internal-dev/MiniMax-H3")
+pipe = MiniMaxH3GeneratorBlocks().init_pipeline("MiniMaxAI/MiniMax-H3")
 pipe.load_components(dtype=torch.bfloat16)
 state = pipe(prompt_embeds=..., text_token_tags=..., height=768, width=1344, num_frames=124, num_inference_steps=30)
 ```
@@ -80,10 +82,12 @@ The wire format is exactly those two tensors — `(1, num_text_tokens, 5120)` bf
 carried as one safetensors file with the resolved `height` / `width` / `num_frames` in its metadata header. A
 text-only request is 246 KB of it; one 768x1344 keyframe adds 1016 vision rows and takes it to 10.7 MB.
 
-The `setup` step runs on **both** halves. It owns no component (PIL and arithmetic) and it resolves the canvas, the
-`17 * n + 5` frame count and the keyframes placed onto that canvas — which the conditioner needs to build its vision
-blocks and this Space needs to encode with the video VAE. It is deterministic, and the conditioner returns the plan
-it resolved so this Space pins the same canvas rather than re-deriving it.
+The keyframe `resize` step runs on **both** halves. It owns no pretrained component (PIL and arithmetic) and it puts
+the keyframes onto the target canvas — which the conditioner needs to build its vision blocks and this Space needs to
+encode with the video VAE. It is deterministic, and the conditioner returns the plan it resolved so this Space pins
+the same canvas rather than re-deriving it. Two things that step no longer does, and that both halves therefore do
+themselves: EXIF-transposing a keyframe into upright RGB, and snapping `num_frames` to `17 * n + 5` — the frame count
+is resolved by the layout step, which lives on this side of the cut.
 
 ## Nothing is paid for with GPU time
 
@@ -147,18 +151,25 @@ one-time `PIPE.to("cuda")` is inside the first row's 339 s and does not reappear
 
 | Variable | Default | Meaning |
 |---|---|---|
-| `H3_CONDITIONER` | `diffusers-internal-dev/minimax-h3-conditioner` | The Space this one asks for embeddings. |
+| `H3_MODEL_REPO` | `MiniMaxAI/MiniMax-H3` | The diffusers-layout checkpoint. Public. |
+| `H3_CONDITIONER` | `multimodalart/qwen3vl-conditioner` | The Space this one asks for embeddings. |
 | `H3_PLACEMENT` | `lazy` | `lazy` moves all 72.16 GiB onto the card on the first GPU call and leaves it there; `offload` hands placement to `ComponentsManager.enable_auto_cpu_offload` instead. |
 | `H3_ATTENTION` | `_native_cudnn` | cuDNN's fused kernel, 10–20% faster than the SDPA default and needs nothing installed. flash-attention 3 is sm90-only and this pool is sm120. |
 | `H3_GPU_DURATION` | `900` | Seconds per request; the pool applies a 1.5 duration factor. |
 | `H3_GPU_SIZE` | `xlarge` | ZeroGPU allocation size. `large` does not fit. |
 
-## Required secret
+## Secrets
 
-`HF_TOKEN` — `diffusers-internal-dev/MiniMax-H3` is private, and so is the conditioner Space this one calls.
+The weights and the conditioner Space are both public, so neither needs a token. `HF_TOKEN` is still read by
+`h3_aoti`, whose compiled-package repository is private — without it, set `H3_AOTI=0`.
 
 ## Where diffusers comes from
 
-MiniMax-H3 is modular-only and not in a released `diffusers`, so the integration branch's `src/diffusers` tree is
-vendored here as a top-level `diffusers/` package; the working directory comes first on `sys.path`, so there is no
-install step. `requirements.txt` only carries what that tree imports.
+MiniMax-H3 is modular-only and not in a released `diffusers`, so `requirements.txt` installs it from the canonical
+pull request, [huggingface/diffusers#14371](https://github.com/huggingface/diffusers/pull/14371), pinned to the
+**commit** `665f5782` (`refs/pull/14371/head` at deploy time) rather than to the moving `minimax-h3-refactor` branch.
+
+That PR is a WIP: it needs **re-pinning whenever it updates**, and `h3_split_blocks.py` — which subclasses its block
+classes to cut the pipeline in two — has to be re-checked against the new head at the same time. The PR refactored
+the blocks into one workflow-selected pipeline, so block names and the shape of the split are exactly what a new head
+is liable to move.

app.py

@@ -13,7 +13,7 @@ import traceback
 import spaces
 import gradio as gr
 
-MODEL_REPO = os.environ.get("H3_MODEL_REPO", "diffusers-internal-dev/MiniMax-H3")
+MODEL_REPO = os.environ.get("H3_MODEL_REPO", "MiniMaxAI/MiniMax-H3")
 CONDITIONER_SPACE = os.environ.get("H3_CONDITIONER", "multimodalart/qwen3vl-conditioner")
 # `lazy` moves all 72.16 GiB onto the card on the first GPU call and leaves it there; `offload` hands placement to
 # `ComponentsManager.enable_auto_cpu_offload` instead. Neither puts anything on the card at *startup*, which is
@@ -85,8 +85,9 @@ def load_models() -> str | None:
     """Load the denoising half. At **startup**, but *not* onto the card.
 
     `MiniMaxH3GeneratorBlocks` declares `transformer`, `vae`, `audio_vae`, `scheduler`, `audio_scheduler` and
-    `video_processor`, so `load_components` fetches exactly those subfolders out of the shared
-    `modular_model_index.json` — `text_encoder/` and `transformer_ref/` are never touched.
+    `video_processor` as its pretrained components (plus an `image_processor` built from config), so
+    `load_components` fetches exactly those subfolders out of the shared `modular_model_index.json` —
+    `text_encoder/` and `transformer_ref/` are never touched.
 
     Both autoencoders carry `_keep_in_fp32_modules` over every module, so the `dtype` below is refused for them and
     they stay float32: a bfloat16 audio VAE decodes the soundtrack roughly 20 dB too quiet.
@@ -105,11 +106,6 @@ def load_models() -> str | None:
     if PIPE is not None or LOAD_ERROR is not None:
         return LOAD_ERROR
 
-    token = os.environ.get("HF_TOKEN")
-    if not token:
-        LOAD_ERROR = f"**`HF_TOKEN` secret is missing** and `{MODEL_REPO}` is private. Add it and restart."
-        return LOAD_ERROR
-
     started = time.time()
     try:
         import torch
@@ -121,7 +117,9 @@ def load_models() -> str | None:
         blocks = MiniMaxH3GeneratorBlocks()
         print(f"[gen] loading {[c.name for c in blocks.expected_components]} from {MODEL_REPO} ...", flush=True)
         pipe = blocks.init_pipeline(MODEL_REPO, components_manager=manager, collection="h3")
-        pipe.load_components(dtype=torch.bfloat16, token=token)
+        # `MiniMaxAI/MiniMax-H3` is public, so the weights are fetched without a token. `HF_TOKEN` is still read by
+        # `h3_aoti`, whose compiled-package repository is not.
+        pipe.load_components(dtype=torch.bfloat16)
         pipe.transformer.set_attention_backend(ATTENTION)
 
         # Still startup, still free: an AoTI package carries no weights and opens its compiled archive lazily inside
@@ -182,8 +180,14 @@ def conditioner():
     return CLIENT
 
 
-def encode_remote(prompt, image_path, last_image_path, canvas, num_frames):
-    """Ask the conditioner Space for `prompt_embeds` + `text_token_tags`. Off this Space's GPU time entirely."""
+def encode_remote(prompt, image_path, last_image_path, canvas, num_frames, rewrite_prompt=False):
+    """Ask the conditioner Space for `prompt_embeds` + `text_token_tags`. Off this Space's GPU time entirely.
+
+    `rewrite_prompt` is the conditioner's prompt upsampling: it rewrites the request into MiniMax-H3's trained format
+    with its own Qwen3-VL and encodes *that*, handing the rewrite back under the plan's `refined_prompt`. It runs on
+    the conditioner's GPU booking, and this whole call happens before `_generate` books a card here, so it costs this
+    Space's `get_duration` nothing.
+    """
     from gradio_client import handle_file
     from safetensors import safe_open
 
@@ -193,6 +197,7 @@ def encode_remote(prompt, image_path, last_image_path, canvas, num_frames):
         last_image_path=handle_file(last_image_path) if last_image_path else None,
         canvas=canvas,
         num_frames=num_frames,
+        rewrite_prompt=bool(rewrite_prompt),
         api_name="/encode",
     )
     with safe_open(path, framework="pt") as handle:
@@ -250,7 +255,8 @@ def _generate(prompt_embeds, text_token_tags, image, last_image, height, width,
     return state.get("videos")[0], state.get("audio")[0].cpu(), state.get("sampling_rate")
 
 
-def generate(prompt, image_path=None, last_image_path=None, canvas=DEFAULT_CANVAS, duration=5, steps=28, seed=42, progress=gr.Progress(track_tqdm=True)):
+def generate(prompt, image_path=None, last_image_path=None, canvas=DEFAULT_CANVAS, duration=5, steps=28, seed=42, upsample=False, progress=gr.Progress(track_tqdm=True)):
+    """One request. `upsample` is appended last and defaults off, so an existing API client is untouched by it."""
     if LOAD_ERROR:
         raise gr.Error(LOAD_ERROR)
     if PIPE is None:
@@ -258,27 +264,34 @@ def generate(prompt, image_path=None, last_image_path=None, canvas=DEFAULT_CANVA
     if not prompt or not prompt.strip():
         raise gr.Error("MiniMax-H3 always takes a prompt, keyframes or not.")
 
-    from PIL import Image
+    from PIL import Image, ImageOps
 
     from diffusers.utils import encode_video
 
     num_frames = snap_frames(duration)
 
-    progress(0.0, desc=f"Conditioning on {CONDITIONER_SPACE} ...")
+    progress(0.0, desc=f"Upsampling the prompt on {CONDITIONER_SPACE} ..." if upsample else f"Conditioning on {CONDITIONER_SPACE} ...")
     conditioned = time.time()
     prompt_embeds, text_token_tags, metadata, plan = encode_remote(
-        prompt, image_path, last_image_path, canvas, num_frames
+        prompt, image_path, last_image_path, canvas, num_frames, rewrite_prompt=upsample
     )
     condition_seconds = time.time() - conditioned
     height, width, num_frames = (int(metadata[key]) for key in ("height", "width", "num_frames"))
+    refined = plan.get("refined_prompt") or ""
+
+    # EXIF-transposed and in RGB before the blocks see them, which the resize step that replaced the old setup step no
+    # longer does itself. Both halves have to prepare a keyframe the same way or the conditioning latents encoded here
+    # would not be of the image the conditioner looked at.
+    def keyframe(path):
+        return ImageOps.exif_transpose(Image.open(path)).convert("RGB") if path else None
 
     progress(0.1, desc=f"Denoising {steps} steps at {width}x{height}, {num_frames} frames ...")
     started = time.time()
     frames, audio, sampling_rate = _generate(
         prompt_embeds,
         text_token_tags,
-        Image.open(image_path) if image_path else None,
-        Image.open(last_image_path) if last_image_path else None,
+        keyframe(image_path),
+        keyframe(last_image_path),
         height,
         width,
         num_frames,
@@ -294,11 +307,12 @@ def generate(prompt, image_path=None, last_image_path=None, canvas=DEFAULT_CANVA
 
     report = (
         f"`{width}x{height}`, {num_frames} frames ({num_frames / FPS:.3f} s), {int(steps)} steps · "
-        f"conditioner {condition_seconds:.0f}s ({plan['num_text_tokens']} tokens) · "
+        f"conditioner {condition_seconds:.0f}s ({plan['num_text_tokens']} tokens"
+        f"{', upsampled' if refined else ''}) · "
         f"denoise + decode {generate_seconds:.0f}s ({generate_seconds / int(steps):.1f} s/step) · seed {int(seed)}"
     )
     print(f"[gen] {report}", flush=True)
-    return path, report
+    return path, report, refined
 
 
 
@@ -378,10 +392,22 @@ with gr.Blocks(title="MiniMax-H3") as demo:
                 duration = gr.Slider(label="Duration (s)", minimum=2, maximum=MAX_UI_DURATION, step=1, value=5)
                 steps = gr.Slider(label="Steps", minimum=10, maximum=40, step=1, value=28)
                 seed = gr.Number(label="Seed", value=42, precision=0)
-            
+                upsample = gr.Checkbox(
+                    label="Upsample prompt",
+                    value=False,
+                    info="Rewrites the prompt into the model's trained format with the conditioner's Qwen3-VL before encoding.",
+                )
+
         with gr.Column():
             video = gr.Video(label="Video + soundtrack")
             report = gr.Markdown(visible=False)
+            with gr.Accordion("Upsampled prompt", open=False):
+                upsampled = gr.Textbox(
+                    show_label=False,
+                    lines=8,
+                    interactive=False,
+                    placeholder="Turn on “Upsample prompt” to see the rewrite that was encoded.",
+                )
 
     image.upload(_fit_keyframe, [image, canvas], [image, canvas])
 
@@ -394,16 +420,19 @@ with gr.Blocks(title="MiniMax-H3") as demo:
             ["A slow seamless camera move from the first view to the last", "examples/first.png", "examples/last.png", "1344x768 · 16:9 full"],
         ],
         inputs=[prompt, image, last_image, canvas],
-        outputs=[video, report],
+        outputs=[video, report, upsampled],
         fn=generate,
         cache_examples=True,
         cache_mode="lazy",
     )
 
+    # `upsample` is appended *after* every input that was already here and every existing input keeps its position, so
+    # a positional API client that predates it keeps working and simply takes the default. Same on the way out: the
+    # video and the report stay first and the upsampled prompt is appended last.
     run.click(
         generate,
-        [prompt, image, last_image, canvas, duration, steps, seed],
-        [video, report],
+        [prompt, image, last_image, canvas, duration, steps, seed, upsample],
+        [video, report, upsampled],
         api_name="generate",
     )
 

diffusers/__init__.py

@@ -1,1750 +0,0 @@
-__version__ = "0.40.0.dev0"
-
-from typing import TYPE_CHECKING
-
-from .utils import (
-    DIFFUSERS_SLOW_IMPORT,
-    OptionalDependencyNotAvailable,
-    _LazyModule,
-    is_accelerate_available,
-    is_auto_round_available,
-    is_bitsandbytes_available,
-    is_gguf_available,
-    is_librosa_available,
-    is_note_seq_available,
-    is_nvidia_modelopt_available,
-    is_onnx_available,
-    is_opencv_available,
-    is_optimum_quanto_available,
-    is_scipy_available,
-    is_sdnq_available,
-    is_sentencepiece_available,
-    is_torch_available,
-    is_torchao_available,
-    is_torchsde_available,
-    is_transformers_available,
-    is_transformers_version,
-)
-
-
-# Lazy Import based on
-# https://github.com/huggingface/transformers/blob/main/src/transformers/__init__.py
-
-# When adding a new object to this init, please add it to `_import_structure`. The `_import_structure` is a dictionary submodule to list of object names,
-# and is used to defer the actual importing for when the objects are requested.
-# This way `import diffusers` provides the names in the namespace without actually importing anything (and especially none of the backends).
-
-_import_structure = {
-    "configuration_utils": ["ConfigMixin"],
-    "guiders": [],
-    "hooks": [],
-    "loaders": ["FromOriginalModelMixin"],
-    "models": [],
-    "modular_pipelines": [],
-    "pipelines": [],
-    "quantizers.pipe_quant_config": ["PipelineQuantizationConfig"],
-    "quantizers.quantization_config": [],
-    "schedulers": [],
-    "utils": [
-        "OptionalDependencyNotAvailable",
-        "is_inflect_available",
-        "is_invisible_watermark_available",
-        "is_librosa_available",
-        "is_note_seq_available",
-        "is_onnx_available",
-        "is_scipy_available",
-        "is_torch_available",
-        "is_torchsde_available",
-        "is_transformers_available",
-        "is_transformers_version",
-        "is_unidecode_available",
-        "logging",
-    ],
-}
-
-try:
-    if not is_torch_available() and not is_accelerate_available() and not is_bitsandbytes_available():
-        raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
-    from .utils import dummy_bitsandbytes_objects
-
-    _import_structure["utils.dummy_bitsandbytes_objects"] = [
-        name for name in dir(dummy_bitsandbytes_objects) if not name.startswith("_")
-    ]
-else:
-    _import_structure["quantizers.quantization_config"].append("BitsAndBytesConfig")
-
-try:
-    if not is_torch_available() and not is_accelerate_available() and not is_gguf_available():
-        raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
-    from .utils import dummy_gguf_objects
-
-    _import_structure["utils.dummy_gguf_objects"] = [
-        name for name in dir(dummy_gguf_objects) if not name.startswith("_")
-    ]
-else:
-    _import_structure["quantizers.quantization_config"].append("GGUFQuantizationConfig")
-
-try:
-    if not is_torch_available() and not is_accelerate_available() and not is_torchao_available():
-        raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
-    from .utils import dummy_torchao_objects
-
-    _import_structure["utils.dummy_torchao_objects"] = [
-        name for name in dir(dummy_torchao_objects) if not name.startswith("_")
-    ]
-else:
-    _import_structure["quantizers.quantization_config"].append("TorchAoConfig")
-
-try:
-    if not is_torch_available() and not is_accelerate_available() and not is_optimum_quanto_available():
-        raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
-    from .utils import dummy_optimum_quanto_objects
-
-    _import_structure["utils.dummy_optimum_quanto_objects"] = [
-        name for name in dir(dummy_optimum_quanto_objects) if not name.startswith("_")
-    ]
-else:
-    _import_structure["quantizers.quantization_config"].append("QuantoConfig")
-
-try:
-    if not is_torch_available() and not is_accelerate_available() and not is_nvidia_modelopt_available():
-        raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
-    from .utils import dummy_nvidia_modelopt_objects
-
-    _import_structure["utils.dummy_nvidia_modelopt_objects"] = [
-        name for name in dir(dummy_nvidia_modelopt_objects) if not name.startswith("_")
-    ]
-else:
-    _import_structure["quantizers.quantization_config"].append("NVIDIAModelOptConfig")
-
-try:
-    if not is_torch_available():
-        raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
-    from .utils import dummy_nunchaku_lite_objects
-
-    _import_structure["utils.dummy_nunchaku_lite_objects"] = [
-        name for name in dir(dummy_nunchaku_lite_objects) if not name.startswith("_")
-    ]
-else:
-    _import_structure["quantizers.quantization_config"].append("NunchakuLiteQuantizationConfig")
-
-try:
-    if not is_auto_round_available():
-        raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
-    from .utils import dummy_auto_round_objects
-
-    _import_structure["utils.dummy_auto_round_objects"] = [
-        name for name in dir(dummy_auto_round_objects) if not name.startswith("_")
-    ]
-else:
-    _import_structure["quantizers.quantization_config"].append("AutoRoundConfig")
-
-try:
-    if not is_torch_available() and not is_accelerate_available() and not is_sdnq_available():
-        raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
-    from .utils import dummy_sdnq_objects
-
-    _import_structure["utils.dummy_sdnq_objects"] = [
-        name for name in dir(dummy_sdnq_objects) if not name.startswith("_")
-    ]
-else:
-    _import_structure["quantizers.quantization_config"].append("SDNQConfig")
-
-try:
-    if not is_onnx_available():
-        raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
-    from .utils import dummy_onnx_objects  # noqa F403
-
-    _import_structure["utils.dummy_onnx_objects"] = [
-        name for name in dir(dummy_onnx_objects) if not name.startswith("_")
-    ]
-
-else:
-    _import_structure["pipelines"].extend(["OnnxRuntimeModel"])
-
-try:
-    if not is_torch_available():
-        raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
-    from .utils import dummy_pt_objects  # noqa F403
-
-    _import_structure["utils.dummy_pt_objects"] = [name for name in dir(dummy_pt_objects) if not name.startswith("_")]
-
-else:
-    _import_structure["guiders"].extend(
-        [
-            "AdaptiveProjectedGuidance",
-            "AdaptiveProjectedMixGuidance",
-            "AutoGuidance",
-            "BaseGuidance",
-            "ClassifierFreeGuidance",
-            "ClassifierFreeZeroStarGuidance",
-            "FrequencyDecoupledGuidance",
-            "PerturbedAttentionGuidance",
-            "SkipLayerGuidance",
-            "SmoothedEnergyGuidance",
-            "TangentialClassifierFreeGuidance",
-        ]
-    )
-    _import_structure["hooks"].extend(
-        [
-            "FasterCacheConfig",
-            "FirstBlockCacheConfig",
-            "HookRegistry",
-            "LayerSkipConfig",
-            "MagCacheConfig",
-            "PyramidAttentionBroadcastConfig",
-            "SmoothedEnergyGuidanceConfig",
-            "TaylorSeerCacheConfig",
-            "TextKVCacheConfig",
-            "apply_faster_cache",
-            "apply_first_block_cache",
-            "apply_layer_skip",
-            "apply_mag_cache",
-            "apply_pyramid_attention_broadcast",
-            "apply_taylorseer_cache",
-            "apply_text_kv_cache",
-        ]
-    )
-    _import_structure["image_processor"] = [
-        "InpaintProcessor",
-        "IPAdapterMaskProcessor",
-        "PixArtImageProcessor",
-        "VaeImageProcessor",
-        "VaeImageProcessorLDM3D",
-    ]
-    _import_structure["models"].extend(
-        [
-            "AceStepTransformer1DModel",
-            "AllegroTransformer3DModel",
-            "AnimaTextConditioner",
-            "AnyFlowFARTransformer3DModel",
-            "AnyFlowTransformer3DModel",
-            "AsymmetricAutoencoderKL",
-            "AttentionBackendName",
-            "AuraFlowTransformer2DModel",
-            "AutoencoderDC",
-            "AutoencoderKL",
-            "AutoencoderKLAllegro",
-            "AutoencoderKLCogVideoX",
-            "AutoencoderKLCosmos",
-            "AutoencoderKLFlux2",
-            "AutoencoderKLHunyuanImage",
-            "AutoencoderKLHunyuanImageRefiner",
-            "AutoencoderKLHunyuanVideo",
-            "AutoencoderKLHunyuanVideo15",
-            "AutoencoderKLKVAE",
-            "AutoencoderKLKVAEVideo",
-            "AutoencoderKLLTX2Audio",
-            "AutoencoderKLLTX2Video",
-            "AutoencoderKLLTXVideo",
-            "AutoencoderKLMagvit",
-            "AutoencoderKLMiniMaxH3",
-            "AutoencoderKLMiniMaxH3Audio",
-            "AutoencoderKLMochi",
-            "AutoencoderKLQwenImage",
-            "AutoencoderKLTemporalDecoder",
-            "AutoencoderKLWan",
-            "AutoencoderOobleck",
-            "AutoencoderRAE",
-            "AutoencoderTiny",
-            "AutoencoderVidTok",
-            "AutoModel",
-            "BriaFiboTransformer2DModel",
-            "BriaTransformer2DModel",
-            "CacheMixin",
-            "ChromaTransformer2DModel",
-            "ChronoEditTransformer3DModel",
-            "CogVideoXTransformer3DModel",
-            "CogView3PlusTransformer2DModel",
-            "CogView4Transformer2DModel",
-            "ConsisIDTransformer3DModel",
-            "ConsistencyDecoderVAE",
-            "ContextParallelConfig",
-            "ControlNetModel",
-            "ControlNetUnionModel",
-            "ControlNetXSAdapter",
-            "Cosmos3AVAEAudioTokenizer",
-            "Cosmos3OmniTransformer",
-            "CosmosControlNetModel",
-            "CosmosTransformer3DModel",
-            "DiTTransformer2DModel",
-            "DreamLiteTransformer2DModel",
-            "DreamLiteUNetModel",
-            "EasyAnimateTransformer3DModel",
-            "ErnieImageTransformer2DModel",
-            "Flux2Transformer2DModel",
-            "FluxControlNetModel",
-            "FluxMultiControlNetModel",
-            "FluxTransformer2DModel",
-            "GlmImageTransformer2DModel",
-            "HeliosTransformer3DModel",
-            "HiDreamImageTransformer2DModel",
-            "HunyuanDiT2DControlNetModel",
-            "HunyuanDiT2DModel",
-            "HunyuanDiT2DMultiControlNetModel",
-            "HunyuanImageTransformer2DModel",
-            "HunyuanVideo15Transformer3DModel",
-            "HunyuanVideoFramepackTransformer3DModel",
-            "HunyuanVideoTransformer3DModel",
-            "I2VGenXLUNet",
-            "Ideogram4Transformer2DModel",
-            "JoyImageEditPlusTransformer3DModel",
-            "JoyImageEditTransformer3DModel",
-            "Kandinsky3UNet",
-            "Kandinsky5Transformer3DModel",
-            "Krea2Transformer2DModel",
-            "LatteTransformer3DModel",
-            "LongCatAudioDiTTransformer",
-            "LongCatAudioDiTVae",
-            "LongCatImageTransformer2DModel",
-            "LTX2VideoTransformer3DModel",
-            "LTXVideoTransformer3DModel",
-            "Lumina2Transformer2DModel",
-            "LuminaNextDiT2DModel",
-            "MiniMaxH3Transformer3DModel",
-            "MochiTransformer3DModel",
-            "ModelMixin",
-            "MotifVideoTransformer3DModel",
-            "MotionAdapter",
-            "MultiAdapter",
-            "MultiControlNetModel",
-            "NucleusMoEImageTransformer2DModel",
-            "OmniGenTransformer2DModel",
-            "OvisImageTransformer2DModel",
-            "ParallelConfig",
-            "PixArtTransformer2DModel",
-            "PriorTransformer",
-            "PRXTransformer2DModel",
-            "QwenImageControlNetModel",
-            "QwenImageMultiControlNetModel",
-            "QwenImageTransformer2DModel",
-            "SanaControlNetModel",
-            "SanaTransformer2DModel",
-            "SanaVideoTransformer3DModel",
-            "SD3ControlNetModel",
-            "SD3MultiControlNetModel",
-            "SD3Transformer2DModel",
-            "SkyReelsV2Transformer3DModel",
-            "SparseControlNetModel",
-            "StableAudioDiTModel",
-            "StableCascadeUNet",
-            "T2IAdapter",
-            "T5FilmDecoder",
-            "Transformer2DModel",
-            "TransformerTemporalModel",
-            "UNet1DModel",
-            "UNet2DConditionModel",
-            "UNet2DModel",
-            "UNet3DConditionModel",
-            "UNetControlNetXSModel",
-            "UNetMotionModel",
-            "UNetSpatioTemporalConditionModel",
-            "UVit2DModel",
-            "VQModel",
-            "WanAnimateTransformer3DModel",
-            "WanTransformer3DModel",
-            "WanVACETransformer3DModel",
-            "ZImageControlNetModel",
-            "ZImageTransformer2DModel",
-            "attention_backend",
-        ]
-    )
-    _import_structure["modular_pipelines"].extend(
-        [
-            "AutoPipelineBlocks",
-            "ComponentsManager",
-            "ComponentSpec",
-            "ConditionalPipelineBlocks",
-            "ConfigSpec",
-            "InputParam",
-            "LoopSequentialPipelineBlocks",
-            "ModularPipeline",
-            "ModularPipelineBlocks",
-            "OutputParam",
-            "SequentialPipelineBlocks",
-        ]
-    )
-    _import_structure["optimization"] = [
-        "get_constant_schedule",
-        "get_constant_schedule_with_warmup",
-        "get_cosine_schedule_with_warmup",
-        "get_cosine_with_hard_restarts_schedule_with_warmup",
-        "get_linear_schedule_with_warmup",
-        "get_polynomial_decay_schedule_with_warmup",
-        "get_scheduler",
-    ]
-    _import_structure["pipelines"].extend(
-        [
-            "AudioPipelineOutput",
-            "AutoPipelineForImage2Image",
-            "AutoPipelineForInpainting",
-            "AutoPipelineForText2Audio",
-            "AutoPipelineForText2Image",
-            "ConsistencyModelPipeline",
-            "DanceDiffusionPipeline",
-            "DDIMPipeline",
-            "DDPMPipeline",
-            "DiffusionPipeline",
-            "DiTPipeline",
-            "ImagePipelineOutput",
-            "KarrasVePipeline",
-            "LDMPipeline",
-            "LDMSuperResolutionPipeline",
-            "PNDMPipeline",
-            "RePaintPipeline",
-            "ScoreSdeVePipeline",
-            "StableDiffusionMixin",
-        ]
-    )
-    _import_structure["quantizers"] = ["DiffusersQuantizer"]
-    _import_structure["schedulers"].extend(
-        [
-            "AmusedScheduler",
-            "BlockRefinementScheduler",
-            "BlockRefinementSchedulerOutput",
-            "CMStochasticIterativeScheduler",
-            "CogVideoXDDIMScheduler",
-            "CogVideoXDPMScheduler",
-            "DDIMInverseScheduler",
-            "DDIMParallelScheduler",
-            "DDIMScheduler",
-            "DDPMParallelScheduler",
-            "DDPMScheduler",
-            "DDPMWuerstchenScheduler",
-            "DEISMultistepScheduler",
-            "DiscreteDDIMScheduler",
-            "DiscreteDDIMSchedulerOutput",
-            "DPMSolverMultistepInverseScheduler",
-            "DPMSolverMultistepScheduler",
-            "DPMSolverSinglestepScheduler",
-            "EDMDPMSolverMultistepScheduler",
-            "EDMEulerScheduler",
-            "EntropyBoundScheduler",
-            "EntropyBoundSchedulerOutput",
-            "EulerAncestralDiscreteScheduler",
-            "EulerDiscreteScheduler",
-            "FlowMapEulerDiscreteScheduler",
-            "FlowMatchEulerDiscreteScheduler",
-            "FlowMatchHeunDiscreteScheduler",
-            "FlowMatchLCMScheduler",
-            "HeliosDMDScheduler",
-            "HeliosScheduler",
-            "HeunDiscreteScheduler",
-            "IPNDMScheduler",
-            "KarrasVeScheduler",
-            "KDPM2AncestralDiscreteScheduler",
-            "KDPM2DiscreteScheduler",
-            "LCMScheduler",
-            "LTXEulerAncestralRFScheduler",
-            "MiniMaxH3Scheduler",
-            "PNDMScheduler",
-            "RePaintScheduler",
-            "SASolverScheduler",
-            "SchedulerMixin",
-            "SCMScheduler",
-            "ScoreSdeVeScheduler",
-            "TCDScheduler",
-            "UnCLIPScheduler",
-            "UniPCMultistepScheduler",
-            "VQDiffusionScheduler",
-        ]
-    )
-    _import_structure["training_utils"] = ["EMAModel"]
-    _import_structure["video_processor"] = ["VideoProcessor"]
-
-try:
-    if not (is_torch_available() and is_scipy_available()):
-        raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
-    from .utils import dummy_torch_and_scipy_objects  # noqa F403
-
-    _import_structure["utils.dummy_torch_and_scipy_objects"] = [
-        name for name in dir(dummy_torch_and_scipy_objects) if not name.startswith("_")
-    ]
-
-else:
-    _import_structure["schedulers"].extend(["LMSDiscreteScheduler"])
-
-try:
-    if not (is_torch_available() and is_torchsde_available()):
-        raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
-    from .utils import dummy_torch_and_torchsde_objects  # noqa F403
-
-    _import_structure["utils.dummy_torch_and_torchsde_objects"] = [
-        name for name in dir(dummy_torch_and_torchsde_objects) if not name.startswith("_")
-    ]
-
-else:
-    _import_structure["schedulers"].extend(["CosineDPMSolverMultistepScheduler", "DPMSolverSDEScheduler"])
-
-try:
-    if not (is_torch_available() and is_transformers_available()):
-        raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
-    from .utils import dummy_torch_and_transformers_objects  # noqa F403
-
-    _import_structure["utils.dummy_torch_and_transformers_objects"] = [
-        name for name in dir(dummy_torch_and_transformers_objects) if not name.startswith("_")
-    ]
-
-else:
-    _import_structure["modular_pipelines"].extend(
-        [
-            "AnimaAutoBlocks",
-            "AnimaModularPipeline",
-            "Cosmos3DistilledBlocks",
-            "Cosmos3DistilledModularPipeline",
-            "Cosmos3OmniBlocks",
-            "Cosmos3OmniModularPipeline",
-            "ErnieImageAutoBlocks",
-            "ErnieImageModularPipeline",
-            "Flux2AutoBlocks",
-            "Flux2KleinAutoBlocks",
-            "Flux2KleinBaseAutoBlocks",
-            "Flux2KleinBaseModularPipeline",
-            "Flux2KleinModularPipeline",
-            "Flux2ModularPipeline",
-            "FluxAutoBlocks",
-            "FluxKontextAutoBlocks",
-            "FluxKontextModularPipeline",
-            "FluxModularPipeline",
-            "HeliosAutoBlocks",
-            "HeliosModularPipeline",
-            "HeliosPyramidAutoBlocks",
-            "HeliosPyramidDistilledAutoBlocks",
-            "HeliosPyramidDistilledModularPipeline",
-            "HeliosPyramidModularPipeline",
-            "HunyuanVideo15AutoBlocks",
-            "HunyuanVideo15ModularPipeline",
-            "Ideogram4AutoBlocks",
-            "Ideogram4ModularPipeline",
-            "Krea2AutoBlocks",
-            "Krea2ModularPipeline",
-            "Krea2TurboAutoBlocks",
-            "Krea2TurboModularPipeline",
-            "LTXAutoBlocks",
-            "LTXModularPipeline",
-            "MiniMaxH3Blocks",
-            "MiniMaxH3ModularPipeline",
-            "MiniMaxH3Ref2VABlocks",
-            "MiniMaxH3Ref2VAModularPipeline",
-            "QwenImageAutoBlocks",
-            "QwenImageEditAutoBlocks",
-            "QwenImageEditModularPipeline",
-            "QwenImageEditPlusAutoBlocks",
-            "QwenImageEditPlusModularPipeline",
-            "QwenImageLayeredAutoBlocks",
-            "QwenImageLayeredModularPipeline",
-            "QwenImageModularPipeline",
-            "StableDiffusion3AutoBlocks",
-            "StableDiffusion3ModularPipeline",
-            "StableDiffusionXLAutoBlocks",
-            "StableDiffusionXLModularPipeline",
-            "Wan22Blocks",
-            "Wan22Image2VideoBlocks",
-            "Wan22Image2VideoModularPipeline",
-            "Wan22ModularPipeline",
-            "WanBlocks",
-            "WanImage2VideoAutoBlocks",
-            "WanImage2VideoModularPipeline",
-            "WanModularPipeline",
-            "ZImageAutoBlocks",
-            "ZImageModularPipeline",
-        ]
-    )
-    _import_structure["pipelines"].extend(
-        [
-            "AceStepAudioTokenDetokenizer",
-            "AceStepAudioTokenizer",
-            "AceStepConditionEncoder",
-            "AceStepPipeline",
-            "AllegroPipeline",
-            "AltDiffusionImg2ImgPipeline",
-            "AltDiffusionPipeline",
-            "AmusedImg2ImgPipeline",
-            "AmusedInpaintPipeline",
-            "AmusedPipeline",
-            "AnimateDiffControlNetPipeline",
-            "AnimateDiffPAGPipeline",
-            "AnimateDiffPipeline",
-            "AnimateDiffSDXLPipeline",
-            "AnimateDiffSparseControlNetPipeline",
-            "AnimateDiffVideoToVideoControlNetPipeline",
-            "AnimateDiffVideoToVideoPipeline",
-            "AnyFlowFARPipeline",
-            "AnyFlowPipeline",
-            "AudioLDM2Pipeline",
-            "AudioLDM2ProjectionModel",
-            "AudioLDM2UNet2DConditionModel",
-            "AudioLDMPipeline",
-            "AuraFlowPipeline",
-            "BlipDiffusionControlNetPipeline",
-            "BlipDiffusionPipeline",
-            "BriaFiboEditPipeline",
-            "BriaFiboPipeline",
-            "BriaPipeline",
-            "ChromaImg2ImgPipeline",
-            "ChromaInpaintPipeline",
-            "ChromaPipeline",
-            "ChronoEditPipeline",
-            "CLIPImageProjection",
-            "CogVideoXFunControlPipeline",
-            "CogVideoXImageToVideoPipeline",
-            "CogVideoXPipeline",
-            "CogVideoXVideoToVideoPipeline",
-            "CogView3PlusPipeline",
-            "CogView4ControlPipeline",
-            "CogView4Pipeline",
-            "ConsisIDPipeline",
-            "Cosmos2_5_PredictBasePipeline",
-            "Cosmos2_5_TransferPipeline",
-            "Cosmos2TextToImagePipeline",
-            "Cosmos2VideoToWorldPipeline",
-            "Cosmos3OmniPipeline",
-            "CosmosActionCondition",
-            "CosmosTextToWorldPipeline",
-            "CosmosVideoToWorldPipeline",
-            "CycleDiffusionPipeline",
-            "DiffusionGemmaPipeline",
-            "DiffusionGemmaPipelineOutput",
-            "DreamLiteMobilePipeline",
-            "DreamLitePipeline",
-            "DreamLitePipelineOutput",
-            "EasyAnimateControlPipeline",
-            "EasyAnimateInpaintPipeline",
-            "EasyAnimatePipeline",
-            "ErnieImagePipeline",
-            "Flux2KleinInpaintPipeline",
-            "Flux2KleinKVPipeline",
-            "Flux2KleinPipeline",
-            "Flux2Pipeline",
-            "FluxControlImg2ImgPipeline",
-            "FluxControlInpaintPipeline",
-            "FluxControlNetImg2ImgPipeline",
-            "FluxControlNetInpaintPipeline",
-            "FluxControlNetPipeline",
-            "FluxControlPipeline",
-            "FluxFillPipeline",
-            "FluxImg2ImgPipeline",
-            "FluxInpaintPipeline",
-            "FluxKontextInpaintPipeline",
-            "FluxKontextPipeline",
-            "FluxPipeline",
-            "FluxPriorReduxPipeline",
-            "GlmImagePipeline",
-            "HeliosPipeline",
-            "HeliosPyramidPipeline",
-            "HiDreamImagePipeline",
-            "HunyuanDiTControlNetPipeline",
-            "HunyuanDiTPAGPipeline",
-            "HunyuanDiTPipeline",
-            "HunyuanImagePipeline",
-            "HunyuanImageRefinerPipeline",
-            "HunyuanSkyreelsImageToVideoPipeline",
-            "HunyuanVideo15ImageToVideoPipeline",
-            "HunyuanVideo15Pipeline",
-            "HunyuanVideoFramepackPipeline",
-            "HunyuanVideoImageToVideoPipeline",
-            "HunyuanVideoPipeline",
-            "I2VGenXLPipeline",
-            "Ideogram4Pipeline",
-            "Ideogram4PromptEnhancerHead",
-            "IFImg2ImgPipeline",
-            "IFImg2ImgSuperResolutionPipeline",
-            "IFInpaintingPipeline",
-            "IFInpaintingSuperResolutionPipeline",
-            "IFPipeline",
-            "IFSuperResolutionPipeline",
-            "ImageTextPipelineOutput",
-            "JoyImageEditPipeline",
-            "JoyImageEditPipelineOutput",
-            "JoyImageEditPlusPipeline",
-            "JoyImageEditPlusPipelineOutput",
-            "Kandinsky3Img2ImgPipeline",
-            "Kandinsky3Pipeline",
-            "Kandinsky5I2IPipeline",
-            "Kandinsky5I2VPipeline",
-            "Kandinsky5T2IPipeline",
-            "Kandinsky5T2VPipeline",
-            "KandinskyCombinedPipeline",
-            "KandinskyImg2ImgCombinedPipeline",
-            "KandinskyImg2ImgPipeline",
-            "KandinskyInpaintCombinedPipeline",
-            "KandinskyInpaintPipeline",
-            "KandinskyPipeline",
-            "KandinskyPriorPipeline",
-            "KandinskyV22CombinedPipeline",
-            "KandinskyV22ControlnetImg2ImgPipeline",
-            "KandinskyV22ControlnetPipeline",
-            "KandinskyV22Img2ImgCombinedPipeline",
-            "KandinskyV22Img2ImgPipeline",
-            "KandinskyV22InpaintCombinedPipeline",
-            "KandinskyV22InpaintPipeline",
-            "KandinskyV22Pipeline",
-            "KandinskyV22PriorEmb2EmbPipeline",
-            "KandinskyV22PriorPipeline",
-            "Krea2Pipeline",
-            "LatentConsistencyModelImg2ImgPipeline",
-            "LatentConsistencyModelPipeline",
-            "LattePipeline",
-            "LDMTextToImagePipeline",
-            "LEditsPPPipelineStableDiffusion",
-            "LEditsPPPipelineStableDiffusionXL",
-            "LLaDA2Pipeline",
-            "LLaDA2PipelineOutput",
-            "LongCatAudioDiTPipeline",
-            "LongCatImageEditPipeline",
-            "LongCatImagePipeline",
-            "LTX2ConditionPipeline",
-            "LTX2HDRPipeline",
-            "LTX2ImageToVideoPipeline",
-            "LTX2InContextPipeline",
-            "LTX2LatentUpsamplePipeline",
-            "LTX2Pipeline",
-            "LTXConditionPipeline",
-            "LTXI2VLongMultiPromptPipeline",
-            "LTXImageToVideoPipeline",
-            "LTXLatentUpsamplePipeline",
-            "LTXPipeline",
-            "LucyEditPipeline",
-            "Lumina2Pipeline",
-            "Lumina2Text2ImgPipeline",
-            "LuminaPipeline",
-            "LuminaText2ImgPipeline",
-            "MarigoldDepthPipeline",
-            "MarigoldIntrinsicsPipeline",
-            "MarigoldNormalsPipeline",
-            "MochiPipeline",
-            "MotifVideoImage2VideoPipeline",
-            "MotifVideoPipeline",
-            "MotifVideoPipelineOutput",
-            "MusicLDMPipeline",
-            "NucleusMoEImagePipeline",
-            "OmniGenPipeline",
-            "OvisImagePipeline",
-            "PaintByExamplePipeline",
-            "PIAPipeline",
-            "PixArtAlphaPipeline",
-            "PixArtSigmaPAGPipeline",
-            "PixArtSigmaPipeline",
-            "PRXPipeline",
-            "PRXPixelPipeline",
-            "QwenImageControlNetInpaintPipeline",
-            "QwenImageControlNetPipeline",
-            "QwenImageEditInpaintPipeline",
-            "QwenImageEditPipeline",
-            "QwenImageEditPlusPipeline",
-            "QwenImageImg2ImgPipeline",
-            "QwenImageInpaintPipeline",
-            "QwenImageLayeredPipeline",
-            "QwenImagePipeline",
-            "ReduxImageEncoder",
-            "SanaControlNetPipeline",
-            "SanaImageToVideoPipeline",
-            "SanaPAGPipeline",
-            "SanaPipeline",
-            "SanaSprintImg2ImgPipeline",
-            "SanaSprintPipeline",
-            "SanaVideoPipeline",
-            "SanaVideoPipeline",
-            "SemanticStableDiffusionPipeline",
-            "ShapEImg2ImgPipeline",
-            "ShapEPipeline",
-            "SkyReelsV2DiffusionForcingImageToVideoPipeline",
-            "SkyReelsV2DiffusionForcingPipeline",
-            "SkyReelsV2DiffusionForcingVideoToVideoPipeline",
-            "SkyReelsV2ImageToVideoPipeline",
-            "SkyReelsV2Pipeline",
-            "StableAudioPipeline",
-            "StableAudioProjectionModel",
-            "StableCascadeCombinedPipeline",
-            "StableCascadeDecoderPipeline",
-            "StableCascadePriorPipeline",
-            "StableDiffusion3ControlNetInpaintingPipeline",
-            "StableDiffusion3ControlNetPipeline",
-            "StableDiffusion3Img2ImgPipeline",
-            "StableDiffusion3InpaintPipeline",
-            "StableDiffusion3PAGImg2ImgPipeline",
-            "StableDiffusion3PAGImg2ImgPipeline",
-            "StableDiffusion3PAGPipeline",
-            "StableDiffusion3Pipeline",
-            "StableDiffusionAdapterPipeline",
-            "StableDiffusionAttendAndExcitePipeline",
-            "StableDiffusionControlNetImg2ImgPipeline",
-            "StableDiffusionControlNetInpaintPipeline",
-            "StableDiffusionControlNetPAGInpaintPipeline",
-            "StableDiffusionControlNetPAGPipeline",
-            "StableDiffusionControlNetPipeline",
-            "StableDiffusionControlNetXSPipeline",
-            "StableDiffusionDepth2ImgPipeline",
-            "StableDiffusionDiffEditPipeline",
-            "StableDiffusionGLIGENPipeline",
-            "StableDiffusionGLIGENTextImagePipeline",
-            "StableDiffusionImageVariationPipeline",
-            "StableDiffusionImg2ImgPipeline",
-            "StableDiffusionInpaintPipeline",
-            "StableDiffusionInpaintPipelineLegacy",
-            "StableDiffusionInstructPix2PixPipeline",
-            "StableDiffusionLatentUpscalePipeline",
-            "StableDiffusionLDM3DPipeline",
-            "StableDiffusionModelEditingPipeline",
-            "StableDiffusionPAGImg2ImgPipeline",
-            "StableDiffusionPAGInpaintPipeline",
-            "StableDiffusionPAGPipeline",
-            "StableDiffusionPanoramaPipeline",
-            "StableDiffusionParadigmsPipeline",
-            "StableDiffusionPipeline",
-            "StableDiffusionPipelineSafe",
-            "StableDiffusionPix2PixZeroPipeline",
-            "StableDiffusionSAGPipeline",
-            "StableDiffusionUpscalePipeline",
-            "StableDiffusionXLAdapterPipeline",
-            "StableDiffusionXLControlNetImg2ImgPipeline",
-            "StableDiffusionXLControlNetInpaintPipeline",
-            "StableDiffusionXLControlNetPAGImg2ImgPipeline",
-            "StableDiffusionXLControlNetPAGPipeline",
-            "StableDiffusionXLControlNetPipeline",
-            "StableDiffusionXLControlNetUnionImg2ImgPipeline",
-            "StableDiffusionXLControlNetUnionInpaintPipeline",
-            "StableDiffusionXLControlNetUnionPipeline",
-            "StableDiffusionXLControlNetXSPipeline",
-            "StableDiffusionXLImg2ImgPipeline",
-            "StableDiffusionXLInpaintPipeline",
-            "StableDiffusionXLInstructPix2PixPipeline",
-            "StableDiffusionXLPAGImg2ImgPipeline",
-            "StableDiffusionXLPAGInpaintPipeline",
-            "StableDiffusionXLPAGPipeline",
-            "StableDiffusionXLPipeline",
-            "StableUnCLIPImg2ImgPipeline",
-            "StableUnCLIPPipeline",
-            "StableVideoDiffusionPipeline",
-            "TextToVideoSDPipeline",
-            "TextToVideoZeroPipeline",
-            "TextToVideoZeroSDXLPipeline",
-            "UnCLIPImageVariationPipeline",
-            "UnCLIPPipeline",
-            "UniDiffuserModel",
-            "UniDiffuserPipeline",
-            "UniDiffuserTextDecoder",
-            "VersatileDiffusionDualGuidedPipeline",
-            "VersatileDiffusionImageVariationPipeline",
-            "VersatileDiffusionPipeline",
-            "VersatileDiffusionTextToImagePipeline",
-            "VideoToVideoSDPipeline",
-            "VisualClozeGenerationPipeline",
-            "VisualClozePipeline",
-            "VQDiffusionPipeline",
-            "WanAnimatePipeline",
-            "WanImageToVideoPipeline",
-            "WanPipeline",
-            "WanVACEPipeline",
-            "WanVideoToVideoPipeline",
-            "WuerstchenCombinedPipeline",
-            "WuerstchenDecoderPipeline",
-            "WuerstchenPriorPipeline",
-            "ZImageControlNetInpaintPipeline",
-            "ZImageControlNetPipeline",
-            "ZImageImg2ImgPipeline",
-            "ZImageInpaintPipeline",
-            "ZImageOmniPipeline",
-            "ZImagePipeline",
-        ]
-    )
-
-
-try:
-    if not (is_torch_available() and is_transformers_available() and is_opencv_available()):
-        raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
-    from .utils import dummy_torch_and_transformers_and_opencv_objects  # noqa F403
-
-    _import_structure["utils.dummy_torch_and_transformers_and_opencv_objects"] = [
-        name for name in dir(dummy_torch_and_transformers_and_opencv_objects) if not name.startswith("_")
-    ]
-
-else:
-    _import_structure["pipelines"].extend(["ConsisIDPipeline"])
-
-try:
-    if not (is_torch_available() and is_transformers_available() and is_sentencepiece_available()):
-        raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
-    from .utils import dummy_torch_and_transformers_and_sentencepiece_objects  # noqa F403
-
-    _import_structure["utils.dummy_torch_and_transformers_and_sentencepiece_objects"] = [
-        name for name in dir(dummy_torch_and_transformers_and_sentencepiece_objects) if not name.startswith("_")
-    ]
-
-else:
-    _import_structure["pipelines"].extend(["KolorsImg2ImgPipeline", "KolorsPAGPipeline", "KolorsPipeline"])
-
-try:
-    if not (is_torch_available() and is_transformers_available() and is_onnx_available()):
-        raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
-    from .utils import dummy_torch_and_transformers_and_onnx_objects  # noqa F403
-
-    _import_structure["utils.dummy_torch_and_transformers_and_onnx_objects"] = [
-        name for name in dir(dummy_torch_and_transformers_and_onnx_objects) if not name.startswith("_")
-    ]
-
-else:
-    _import_structure["pipelines"].extend(
-        [
-            "OnnxStableDiffusionImg2ImgPipeline",
-            "OnnxStableDiffusionInpaintPipeline",
-            "OnnxStableDiffusionInpaintPipelineLegacy",
-            "OnnxStableDiffusionPipeline",
-            "OnnxStableDiffusionUpscalePipeline",
-            "StableDiffusionOnnxPipeline",
-        ]
-    )
-
-try:
-    if not (is_torch_available() and is_librosa_available()):
-        raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
-    from .utils import dummy_torch_and_librosa_objects  # noqa F403
-
-    _import_structure["utils.dummy_torch_and_librosa_objects"] = [
-        name for name in dir(dummy_torch_and_librosa_objects) if not name.startswith("_")
-    ]
-
-else:
-    _import_structure["pipelines"].extend(["AudioDiffusionPipeline", "Mel"])
-
-try:
-    if not (is_transformers_available() and is_torch_available() and is_note_seq_available()):
-        raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
-    from .utils import dummy_transformers_and_torch_and_note_seq_objects  # noqa F403
-
-    _import_structure["utils.dummy_transformers_and_torch_and_note_seq_objects"] = [
-        name for name in dir(dummy_transformers_and_torch_and_note_seq_objects) if not name.startswith("_")
-    ]
-
-
-else:
-    _import_structure["pipelines"].extend(["SpectrogramDiffusionPipeline"])
-
-try:
-    if not (is_note_seq_available()):
-        raise OptionalDependencyNotAvailable()
-except OptionalDependencyNotAvailable:
-    from .utils import dummy_note_seq_objects  # noqa F403
-
-    _import_structure["utils.dummy_note_seq_objects"] = [
-        name for name in dir(dummy_note_seq_objects) if not name.startswith("_")
-    ]
-
-
-else:
-    _import_structure["pipelines"].extend(["MidiProcessor"])
-
-if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT:
-    from .configuration_utils import ConfigMixin
-    from .quantizers import PipelineQuantizationConfig
-
-    try:
-        if not is_bitsandbytes_available():
-            raise OptionalDependencyNotAvailable()
-    except OptionalDependencyNotAvailable:
-        from .utils.dummy_bitsandbytes_objects import *
-    else:
-        from .quantizers.quantization_config import BitsAndBytesConfig
-
-    try:
-        if not is_gguf_available():
-            raise OptionalDependencyNotAvailable()
-    except OptionalDependencyNotAvailable:
-        from .utils.dummy_gguf_objects import *
-    else:
-        from .quantizers.quantization_config import GGUFQuantizationConfig
-
-    try:
-        if not is_torchao_available():
-            raise OptionalDependencyNotAvailable()
-    except OptionalDependencyNotAvailable:
-        from .utils.dummy_torchao_objects import *
-    else:
-        from .quantizers.quantization_config import TorchAoConfig
-
-    try:
-        if not is_optimum_quanto_available():
-            raise OptionalDependencyNotAvailable()
-    except OptionalDependencyNotAvailable:
-        from .utils.dummy_optimum_quanto_objects import *
-    else:
-        from .quantizers.quantization_config import QuantoConfig
-
-    try:
-        if not is_nvidia_modelopt_available():
-            raise OptionalDependencyNotAvailable()
-    except OptionalDependencyNotAvailable:
-        from .utils.dummy_nvidia_modelopt_objects import *
-    else:
-        from .quantizers.quantization_config import NVIDIAModelOptConfig
-
-    try:
-        if not is_torch_available():
-            raise OptionalDependencyNotAvailable()
-    except OptionalDependencyNotAvailable:
-        from .utils.dummy_nunchaku_lite_objects import *
-    else:
-        from .quantizers.quantization_config import NunchakuLiteQuantizationConfig
-
-    try:
-        if not is_auto_round_available():
-            raise OptionalDependencyNotAvailable()
-    except OptionalDependencyNotAvailable:
-        from .utils.dummy_auto_round_objects import *
-    else:
-        from .quantizers.quantization_config import AutoRoundConfig
-
-    try:
-        if not is_sdnq_available():
-            raise OptionalDependencyNotAvailable()
-    except OptionalDependencyNotAvailable:
-        from .utils.dummy_sdnq_objects import *
-    else:
-        from .quantizers.quantization_config import SDNQConfig
-
-    try:
-        if not is_onnx_available():
-            raise OptionalDependencyNotAvailable()
-    except OptionalDependencyNotAvailable:
-        from .utils.dummy_onnx_objects import *  # noqa F403
-    else:
-        from .pipelines import OnnxRuntimeModel
-
-    try:
-        if not is_torch_available():
-            raise OptionalDependencyNotAvailable()
-    except OptionalDependencyNotAvailable:
-        from .utils.dummy_pt_objects import *  # noqa F403
-    else:
-        from .guiders import (
-            AdaptiveProjectedGuidance,
-            AdaptiveProjectedMixGuidance,
-            AutoGuidance,
-            BaseGuidance,
-            ClassifierFreeGuidance,
-            ClassifierFreeZeroStarGuidance,
-            FrequencyDecoupledGuidance,
-            PerturbedAttentionGuidance,
-            SkipLayerGuidance,
-            SmoothedEnergyGuidance,
-            TangentialClassifierFreeGuidance,
-        )
-        from .hooks import (
-            FasterCacheConfig,
-            FirstBlockCacheConfig,
-            HookRegistry,
-            LayerSkipConfig,
-            MagCacheConfig,
-            PyramidAttentionBroadcastConfig,
-            SmoothedEnergyGuidanceConfig,
-            TaylorSeerCacheConfig,
-            TextKVCacheConfig,
-            apply_faster_cache,
-            apply_first_block_cache,
-            apply_layer_skip,
-            apply_mag_cache,
-            apply_pyramid_attention_broadcast,
-            apply_taylorseer_cache,
-            apply_text_kv_cache,
-        )
-        from .image_processor import (
-            InpaintProcessor,
-            IPAdapterMaskProcessor,
-            PixArtImageProcessor,
-            VaeImageProcessor,
-            VaeImageProcessorLDM3D,
-        )
-        from .models import (
-            AceStepTransformer1DModel,
-            AllegroTransformer3DModel,
-            AnimaTextConditioner,
-            AnyFlowFARTransformer3DModel,
-            AnyFlowTransformer3DModel,
-            AsymmetricAutoencoderKL,
-            AttentionBackendName,
-            AuraFlowTransformer2DModel,
-            AutoencoderDC,
-            AutoencoderKL,
-            AutoencoderKLAllegro,
-            AutoencoderKLCogVideoX,
-            AutoencoderKLCosmos,
-            AutoencoderKLFlux2,
-            AutoencoderKLHunyuanImage,
-            AutoencoderKLHunyuanImageRefiner,
-            AutoencoderKLHunyuanVideo,
-            AutoencoderKLHunyuanVideo15,
-            AutoencoderKLKVAE,
-            AutoencoderKLKVAEVideo,
-            AutoencoderKLLTX2Audio,
-            AutoencoderKLLTX2Video,
-            AutoencoderKLLTXVideo,
-            AutoencoderKLMagvit,
-            AutoencoderKLMiniMaxH3,
-            AutoencoderKLMiniMaxH3Audio,
-            AutoencoderKLMochi,
-            AutoencoderKLQwenImage,
-            AutoencoderKLTemporalDecoder,
-            AutoencoderKLWan,
-            AutoencoderOobleck,
-            AutoencoderRAE,
-            AutoencoderTiny,
-            AutoencoderVidTok,
-            AutoModel,
-            BriaFiboTransformer2DModel,
-            BriaTransformer2DModel,
-            CacheMixin,
-            ChromaTransformer2DModel,
-            ChronoEditTransformer3DModel,
-            CogVideoXTransformer3DModel,
-            CogView3PlusTransformer2DModel,
-            CogView4Transformer2DModel,
-            ConsisIDTransformer3DModel,
-            ConsistencyDecoderVAE,
-            ContextParallelConfig,
-            ControlNetModel,
-            ControlNetUnionModel,
-            ControlNetXSAdapter,
-            Cosmos3AVAEAudioTokenizer,
-            Cosmos3OmniTransformer,
-            CosmosControlNetModel,
-            CosmosTransformer3DModel,
-            DiTTransformer2DModel,
-            DreamLiteTransformer2DModel,
-            DreamLiteUNetModel,
-            EasyAnimateTransformer3DModel,
-            ErnieImageTransformer2DModel,
-            Flux2Transformer2DModel,
-            FluxControlNetModel,
-            FluxMultiControlNetModel,
-            FluxTransformer2DModel,
-            GlmImageTransformer2DModel,
-            HeliosTransformer3DModel,
-            HiDreamImageTransformer2DModel,
-            HunyuanDiT2DControlNetModel,
-            HunyuanDiT2DModel,
-            HunyuanDiT2DMultiControlNetModel,
-            HunyuanImageTransformer2DModel,
-            HunyuanVideo15Transformer3DModel,
-            HunyuanVideoFramepackTransformer3DModel,
-            HunyuanVideoTransformer3DModel,
-            I2VGenXLUNet,
-            Ideogram4Transformer2DModel,
-            JoyImageEditPlusTransformer3DModel,
-            JoyImageEditTransformer3DModel,
-            Kandinsky3UNet,
-            Kandinsky5Transformer3DModel,
-            Krea2Transformer2DModel,
-            LatteTransformer3DModel,
-            LongCatAudioDiTTransformer,
-            LongCatAudioDiTVae,
-            LongCatImageTransformer2DModel,
-            LTX2VideoTransformer3DModel,
-            LTXVideoTransformer3DModel,
-            Lumina2Transformer2DModel,
-            LuminaNextDiT2DModel,
-            MiniMaxH3Transformer3DModel,
-            MochiTransformer3DModel,
-            ModelMixin,
-            MotifVideoTransformer3DModel,
-            MotionAdapter,
-            MultiAdapter,
-            MultiControlNetModel,
-            NucleusMoEImageTransformer2DModel,
-            OmniGenTransformer2DModel,
-            OvisImageTransformer2DModel,
-            ParallelConfig,
-            PixArtTransformer2DModel,
-            PriorTransformer,
-            PRXTransformer2DModel,
-            QwenImageControlNetModel,
-            QwenImageMultiControlNetModel,
-            QwenImageTransformer2DModel,
-            SanaControlNetModel,
-            SanaTransformer2DModel,
-            SanaVideoTransformer3DModel,
-            SD3ControlNetModel,
-            SD3MultiControlNetModel,
-            SD3Transformer2DModel,
-            SkyReelsV2Transformer3DModel,
-            SparseControlNetModel,
-            StableAudioDiTModel,
-            T2IAdapter,
-            T5FilmDecoder,
-            Transformer2DModel,
-            TransformerTemporalModel,
-            UNet1DModel,
-            UNet2DConditionModel,
-            UNet2DModel,
-            UNet3DConditionModel,
-            UNetControlNetXSModel,
-            UNetMotionModel,
-            UNetSpatioTemporalConditionModel,
-            UVit2DModel,
-            VQModel,
-            WanAnimateTransformer3DModel,
-            WanTransformer3DModel,
-            WanVACETransformer3DModel,
-            ZImageControlNetModel,
-            ZImageTransformer2DModel,
-            attention_backend,
-        )
-        from .modular_pipelines import (
-            AutoPipelineBlocks,
-            ComponentsManager,
-            ComponentSpec,
-            ConditionalPipelineBlocks,
-            ConfigSpec,
-            InputParam,
-            LoopSequentialPipelineBlocks,
-            ModularPipeline,
-            ModularPipelineBlocks,
-            OutputParam,
-            SequentialPipelineBlocks,
-        )
-        from .optimization import (
-            get_constant_schedule,
-            get_constant_schedule_with_warmup,
-            get_cosine_schedule_with_warmup,
-            get_cosine_with_hard_restarts_schedule_with_warmup,
-            get_linear_schedule_with_warmup,
-            get_polynomial_decay_schedule_with_warmup,
-            get_scheduler,
-        )
-        from .pipelines import (
-            AudioPipelineOutput,
-            AutoPipelineForImage2Image,
-            AutoPipelineForInpainting,
-            AutoPipelineForText2Audio,
-            AutoPipelineForText2Image,
-            BlipDiffusionControlNetPipeline,
-            BlipDiffusionPipeline,
-            CLIPImageProjection,
-            ConsistencyModelPipeline,
-            DanceDiffusionPipeline,
-            DDIMPipeline,
-            DDPMPipeline,
-            DiffusionPipeline,
-            DiTPipeline,
-            ImagePipelineOutput,
-            KarrasVePipeline,
-            LDMPipeline,
-            LDMSuperResolutionPipeline,
-            PNDMPipeline,
-            RePaintPipeline,
-            ScoreSdeVePipeline,
-            StableDiffusionMixin,
-        )
-        from .quantizers import DiffusersQuantizer
-        from .schedulers import (
-            AmusedScheduler,
-            BlockRefinementScheduler,
-            BlockRefinementSchedulerOutput,
-            CMStochasticIterativeScheduler,
-            CogVideoXDDIMScheduler,
-            CogVideoXDPMScheduler,
-            DDIMInverseScheduler,
-            DDIMParallelScheduler,
-            DDIMScheduler,
-            DDPMParallelScheduler,
-            DDPMScheduler,
-            DDPMWuerstchenScheduler,
-            DEISMultistepScheduler,
-            DiscreteDDIMScheduler,
-            DiscreteDDIMSchedulerOutput,
-            DPMSolverMultistepInverseScheduler,
-            DPMSolverMultistepScheduler,
-            DPMSolverSinglestepScheduler,
-            EDMDPMSolverMultistepScheduler,
-            EDMEulerScheduler,
-            EntropyBoundScheduler,
-            EntropyBoundSchedulerOutput,
-            EulerAncestralDiscreteScheduler,
-            EulerDiscreteScheduler,
-            FlowMapEulerDiscreteScheduler,
-            FlowMatchEulerDiscreteScheduler,
-            FlowMatchHeunDiscreteScheduler,
-            FlowMatchLCMScheduler,
-            HeliosDMDScheduler,
-            HeliosScheduler,
-            HeunDiscreteScheduler,
-            IPNDMScheduler,
-            KarrasVeScheduler,
-            KDPM2AncestralDiscreteScheduler,
-            KDPM2DiscreteScheduler,
-            LCMScheduler,
-            LTXEulerAncestralRFScheduler,
-            MiniMaxH3Scheduler,
-            PNDMScheduler,
-            RePaintScheduler,
-            SASolverScheduler,
-            SchedulerMixin,
-            SCMScheduler,
-            ScoreSdeVeScheduler,
-            TCDScheduler,
-            UnCLIPScheduler,
-            UniPCMultistepScheduler,
-            VQDiffusionScheduler,
-        )
-        from .training_utils import EMAModel
-        from .video_processor import VideoProcessor
-
-    try:
-        if not (is_torch_available() and is_scipy_available()):
-            raise OptionalDependencyNotAvailable()
-    except OptionalDependencyNotAvailable:
-        from .utils.dummy_torch_and_scipy_objects import *  # noqa F403
-    else:
-        from .schedulers import LMSDiscreteScheduler
-
-    try:
-        if not (is_torch_available() and is_torchsde_available()):
-            raise OptionalDependencyNotAvailable()
-    except OptionalDependencyNotAvailable:
-        from .utils.dummy_torch_and_torchsde_objects import *  # noqa F403
-    else:
-        from .schedulers import CosineDPMSolverMultistepScheduler, DPMSolverSDEScheduler
-
-    try:
-        if not (is_torch_available() and is_transformers_available()):
-            raise OptionalDependencyNotAvailable()
-    except OptionalDependencyNotAvailable:
-        from .utils.dummy_torch_and_transformers_objects import *  # noqa F403
-    else:
-        from .modular_pipelines import (
-            AnimaAutoBlocks,
-            AnimaModularPipeline,
-            Cosmos3DistilledBlocks,
-            Cosmos3DistilledModularPipeline,
-            Cosmos3OmniBlocks,
-            Cosmos3OmniModularPipeline,
-            ErnieImageAutoBlocks,
-            ErnieImageModularPipeline,
-            Flux2AutoBlocks,
-            Flux2KleinAutoBlocks,
-            Flux2KleinBaseAutoBlocks,
-            Flux2KleinBaseModularPipeline,
-            Flux2KleinModularPipeline,
-            Flux2ModularPipeline,
-            FluxAutoBlocks,
-            FluxKontextAutoBlocks,
-            FluxKontextModularPipeline,
-            FluxModularPipeline,
-            HeliosAutoBlocks,
-            HeliosModularPipeline,
-            HeliosPyramidAutoBlocks,
-            HeliosPyramidDistilledAutoBlocks,
-            HeliosPyramidDistilledModularPipeline,
-            HeliosPyramidModularPipeline,
-            HunyuanVideo15AutoBlocks,
-            HunyuanVideo15ModularPipeline,
-            Ideogram4AutoBlocks,
-            Ideogram4ModularPipeline,
-            Krea2AutoBlocks,
-            Krea2ModularPipeline,
-            Krea2TurboAutoBlocks,
-            Krea2TurboModularPipeline,
-            LTXAutoBlocks,
-            LTXModularPipeline,
-            MiniMaxH3Blocks,
-            MiniMaxH3ModularPipeline,
-            MiniMaxH3Ref2VABlocks,
-            MiniMaxH3Ref2VAModularPipeline,
-            QwenImageAutoBlocks,
-            QwenImageEditAutoBlocks,
-            QwenImageEditModularPipeline,
-            QwenImageEditPlusAutoBlocks,
-            QwenImageEditPlusModularPipeline,
-            QwenImageLayeredAutoBlocks,
-            QwenImageLayeredModularPipeline,
-            QwenImageModularPipeline,
-            StableDiffusion3AutoBlocks,
-            StableDiffusion3ModularPipeline,
-            StableDiffusionXLAutoBlocks,
-            StableDiffusionXLModularPipeline,
-            Wan22Blocks,
-            Wan22Image2VideoBlocks,
-            Wan22Image2VideoModularPipeline,
-            Wan22ModularPipeline,
-            WanBlocks,
-            WanImage2VideoAutoBlocks,
-            WanImage2VideoModularPipeline,
-            WanModularPipeline,
-            ZImageAutoBlocks,
-            ZImageModularPipeline,
-        )
-        from .pipelines import (
-            AceStepAudioTokenDetokenizer,
-            AceStepAudioTokenizer,
-            AceStepConditionEncoder,
-            AceStepPipeline,
-            AllegroPipeline,
-            AltDiffusionImg2ImgPipeline,
-            AltDiffusionPipeline,
-            AmusedImg2ImgPipeline,
-            AmusedInpaintPipeline,
-            AmusedPipeline,
-            AnimateDiffControlNetPipeline,
-            AnimateDiffPAGPipeline,
-            AnimateDiffPipeline,
-            AnimateDiffSDXLPipeline,
-            AnimateDiffSparseControlNetPipeline,
-            AnimateDiffVideoToVideoControlNetPipeline,
-            AnimateDiffVideoToVideoPipeline,
-            AnyFlowFARPipeline,
-            AnyFlowPipeline,
-            AudioLDM2Pipeline,
-            AudioLDM2ProjectionModel,
-            AudioLDM2UNet2DConditionModel,
-            AudioLDMPipeline,
-            AuraFlowPipeline,
-            BriaFiboEditPipeline,
-            BriaFiboPipeline,
-            BriaPipeline,
-            ChromaImg2ImgPipeline,
-            ChromaInpaintPipeline,
-            ChromaPipeline,
-            ChronoEditPipeline,
-            CLIPImageProjection,
-            CogVideoXFunControlPipeline,
-            CogVideoXImageToVideoPipeline,
-            CogVideoXPipeline,
-            CogVideoXVideoToVideoPipeline,
-            CogView3PlusPipeline,
-            CogView4ControlPipeline,
-            CogView4Pipeline,
-            ConsisIDPipeline,
-            Cosmos2_5_PredictBasePipeline,
-            Cosmos2_5_TransferPipeline,
-            Cosmos2TextToImagePipeline,
-            Cosmos2VideoToWorldPipeline,
-            Cosmos3OmniPipeline,
-            CosmosActionCondition,
-            CosmosTextToWorldPipeline,
-            CosmosVideoToWorldPipeline,
-            CycleDiffusionPipeline,
-            DiffusionGemmaPipeline,
-            DiffusionGemmaPipelineOutput,
-            DreamLiteMobilePipeline,
-            DreamLitePipeline,
-            DreamLitePipelineOutput,
-            EasyAnimateControlPipeline,
-            EasyAnimateInpaintPipeline,
-            EasyAnimatePipeline,
-            ErnieImagePipeline,
-            Flux2KleinInpaintPipeline,
-            Flux2KleinKVPipeline,
-            Flux2KleinPipeline,
-            Flux2Pipeline,
-            FluxControlImg2ImgPipeline,
-            FluxControlInpaintPipeline,
-            FluxControlNetImg2ImgPipeline,
-            FluxControlNetInpaintPipeline,
-            FluxControlNetPipeline,
-            FluxControlPipeline,
-            FluxFillPipeline,
-            FluxImg2ImgPipeline,
-            FluxInpaintPipeline,
-            FluxKontextInpaintPipeline,
-            FluxKontextPipeline,
-            FluxPipeline,
-            FluxPriorReduxPipeline,
-            GlmImagePipeline,
-            HeliosPipeline,
-            HeliosPyramidPipeline,
-            HiDreamImagePipeline,
-            HunyuanDiTControlNetPipeline,
-            HunyuanDiTPAGPipeline,
-            HunyuanDiTPipeline,
-            HunyuanImagePipeline,
-            HunyuanImageRefinerPipeline,
-            HunyuanSkyreelsImageToVideoPipeline,
-            HunyuanVideo15ImageToVideoPipeline,
-            HunyuanVideo15Pipeline,
-            HunyuanVideoFramepackPipeline,
-            HunyuanVideoImageToVideoPipeline,
-            HunyuanVideoPipeline,
-            I2VGenXLPipeline,
-            Ideogram4Pipeline,
-            Ideogram4PromptEnhancerHead,
-            IFImg2ImgPipeline,
-            IFImg2ImgSuperResolutionPipeline,
-            IFInpaintingPipeline,
-            IFInpaintingSuperResolutionPipeline,
-            IFPipeline,
-            IFSuperResolutionPipeline,
-            ImageTextPipelineOutput,
-            JoyImageEditPipeline,
-            JoyImageEditPipelineOutput,
-            JoyImageEditPlusPipeline,
-            JoyImageEditPlusPipelineOutput,
-            Kandinsky3Img2ImgPipeline,
-            Kandinsky3Pipeline,
-            Kandinsky5I2IPipeline,
-            Kandinsky5I2VPipeline,
-            Kandinsky5T2IPipeline,
-            Kandinsky5T2VPipeline,
-            KandinskyCombinedPipeline,
-            KandinskyImg2ImgCombinedPipeline,
-            KandinskyImg2ImgPipeline,
-            KandinskyInpaintCombinedPipeline,
-            KandinskyInpaintPipeline,
-            KandinskyPipeline,
-            KandinskyPriorPipeline,
-            KandinskyV22CombinedPipeline,
-            KandinskyV22ControlnetImg2ImgPipeline,
-            KandinskyV22ControlnetPipeline,
-            KandinskyV22Img2ImgCombinedPipeline,
-            KandinskyV22Img2ImgPipeline,
-            KandinskyV22InpaintCombinedPipeline,
-            KandinskyV22InpaintPipeline,
-            KandinskyV22Pipeline,
-            KandinskyV22PriorEmb2EmbPipeline,
-            KandinskyV22PriorPipeline,
-            Krea2Pipeline,
-            LatentConsistencyModelImg2ImgPipeline,
-            LatentConsistencyModelPipeline,
-            LattePipeline,
-            LDMTextToImagePipeline,
-            LEditsPPPipelineStableDiffusion,
-            LEditsPPPipelineStableDiffusionXL,
-            LLaDA2Pipeline,
-            LLaDA2PipelineOutput,
-            LongCatAudioDiTPipeline,
-            LongCatImageEditPipeline,
-            LongCatImagePipeline,
-            LTX2ConditionPipeline,
-            LTX2HDRPipeline,
-            LTX2ImageToVideoPipeline,
-            LTX2InContextPipeline,
-            LTX2LatentUpsamplePipeline,
-            LTX2Pipeline,
-            LTXConditionPipeline,
-            LTXI2VLongMultiPromptPipeline,
-            LTXImageToVideoPipeline,
-            LTXLatentUpsamplePipeline,
-            LTXPipeline,
-            LucyEditPipeline,
-            Lumina2Pipeline,
-            Lumina2Text2ImgPipeline,
-            LuminaPipeline,
-            LuminaText2ImgPipeline,
-            MarigoldDepthPipeline,
-            MarigoldIntrinsicsPipeline,
-            MarigoldNormalsPipeline,
-            MochiPipeline,
-            MotifVideoImage2VideoPipeline,
-            MotifVideoPipeline,
-            MotifVideoPipelineOutput,
-            MusicLDMPipeline,
-            NucleusMoEImagePipeline,
-            OmniGenPipeline,
-            OvisImagePipeline,
-            PaintByExamplePipeline,
-            PIAPipeline,
-            PixArtAlphaPipeline,
-            PixArtSigmaPAGPipeline,
-            PixArtSigmaPipeline,
-            PRXPipeline,
-            PRXPixelPipeline,
-            QwenImageControlNetInpaintPipeline,
-            QwenImageControlNetPipeline,
-            QwenImageEditInpaintPipeline,
-            QwenImageEditPipeline,
-            QwenImageEditPlusPipeline,
-            QwenImageImg2ImgPipeline,
-            QwenImageInpaintPipeline,
-            QwenImageLayeredPipeline,
-            QwenImagePipeline,
-            ReduxImageEncoder,
-            SanaControlNetPipeline,
-            SanaImageToVideoPipeline,
-            SanaPAGPipeline,
-            SanaPipeline,
-            SanaSprintImg2ImgPipeline,
-            SanaSprintPipeline,
-            SanaVideoPipeline,
-            SemanticStableDiffusionPipeline,
-            ShapEImg2ImgPipeline,
-            ShapEPipeline,
-            SkyReelsV2DiffusionForcingImageToVideoPipeline,
-            SkyReelsV2DiffusionForcingPipeline,
-            SkyReelsV2DiffusionForcingVideoToVideoPipeline,
-            SkyReelsV2ImageToVideoPipeline,
-            SkyReelsV2Pipeline,
-            StableAudioPipeline,
-            StableAudioProjectionModel,
-            StableCascadeCombinedPipeline,
-            StableCascadeDecoderPipeline,
-            StableCascadePriorPipeline,
-            StableDiffusion3ControlNetInpaintingPipeline,
-            StableDiffusion3ControlNetPipeline,
-            StableDiffusion3Img2ImgPipeline,
-            StableDiffusion3InpaintPipeline,
-            StableDiffusion3PAGImg2ImgPipeline,
-            StableDiffusion3PAGPipeline,
-            StableDiffusion3Pipeline,
-            StableDiffusionAdapterPipeline,
-            StableDiffusionAttendAndExcitePipeline,
-            StableDiffusionControlNetImg2ImgPipeline,
-            StableDiffusionControlNetInpaintPipeline,
-            StableDiffusionControlNetPAGInpaintPipeline,
-            StableDiffusionControlNetPAGPipeline,
-            StableDiffusionControlNetPipeline,
-            StableDiffusionControlNetXSPipeline,
-            StableDiffusionDepth2ImgPipeline,
-            StableDiffusionDiffEditPipeline,
-            StableDiffusionGLIGENPipeline,
-            StableDiffusionGLIGENTextImagePipeline,
-            StableDiffusionImageVariationPipeline,
-            StableDiffusionImg2ImgPipeline,
-            StableDiffusionInpaintPipeline,
-            StableDiffusionInpaintPipelineLegacy,
-            StableDiffusionInstructPix2PixPipeline,
-            StableDiffusionLatentUpscalePipeline,
-            StableDiffusionLDM3DPipeline,
-            StableDiffusionModelEditingPipeline,
-            StableDiffusionPAGImg2ImgPipeline,
-            StableDiffusionPAGInpaintPipeline,
-            StableDiffusionPAGPipeline,
-            StableDiffusionPanoramaPipeline,
-            StableDiffusionParadigmsPipeline,
-            StableDiffusionPipeline,
-            StableDiffusionPipelineSafe,
-            StableDiffusionPix2PixZeroPipeline,
-            StableDiffusionSAGPipeline,
-            StableDiffusionUpscalePipeline,
-            StableDiffusionXLAdapterPipeline,
-            StableDiffusionXLControlNetImg2ImgPipeline,
-            StableDiffusionXLControlNetInpaintPipeline,
-            StableDiffusionXLControlNetPAGImg2ImgPipeline,
-            StableDiffusionXLControlNetPAGPipeline,
-            StableDiffusionXLControlNetPipeline,
-            StableDiffusionXLControlNetUnionImg2ImgPipeline,
-            StableDiffusionXLControlNetUnionInpaintPipeline,
-            StableDiffusionXLControlNetUnionPipeline,
-            StableDiffusionXLControlNetXSPipeline,
-            StableDiffusionXLImg2ImgPipeline,
-            StableDiffusionXLInpaintPipeline,
-            StableDiffusionXLInstructPix2PixPipeline,
-            StableDiffusionXLPAGImg2ImgPipeline,
-            StableDiffusionXLPAGInpaintPipeline,
-            StableDiffusionXLPAGPipeline,
-            StableDiffusionXLPipeline,
-            StableUnCLIPImg2ImgPipeline,
-            StableUnCLIPPipeline,
-            StableVideoDiffusionPipeline,
-            TextToVideoSDPipeline,
-            TextToVideoZeroPipeline,
-            TextToVideoZeroSDXLPipeline,
-            UnCLIPImageVariationPipeline,
-            UnCLIPPipeline,
-            UniDiffuserModel,
-            UniDiffuserPipeline,
-            UniDiffuserTextDecoder,
-            VersatileDiffusionDualGuidedPipeline,
-            VersatileDiffusionImageVariationPipeline,
-            VersatileDiffusionPipeline,
-            VersatileDiffusionTextToImagePipeline,
-            VideoToVideoSDPipeline,
-            VisualClozeGenerationPipeline,
-            VisualClozePipeline,
-            VQDiffusionPipeline,
-            WanAnimatePipeline,
-            WanImageToVideoPipeline,
-            WanPipeline,
-            WanVACEPipeline,
-            WanVideoToVideoPipeline,
-            WuerstchenCombinedPipeline,
-            WuerstchenDecoderPipeline,
-            WuerstchenPriorPipeline,
-            ZImageControlNetInpaintPipeline,
-            ZImageControlNetPipeline,
-            ZImageImg2ImgPipeline,
-            ZImageInpaintPipeline,
-            ZImageOmniPipeline,
-            ZImagePipeline,
-        )
-
-    try:
-        if not (is_torch_available() and is_transformers_available() and is_sentencepiece_available()):
-            raise OptionalDependencyNotAvailable()
-    except OptionalDependencyNotAvailable:
-        from .utils.dummy_torch_and_transformers_and_sentencepiece_objects import *  # noqa F403
-    else:
-        from .pipelines import KolorsImg2ImgPipeline, KolorsPAGPipeline, KolorsPipeline
-
-    try:
-        if not (is_torch_available() and is_transformers_available() and is_opencv_available()):
-            raise OptionalDependencyNotAvailable()
-    except OptionalDependencyNotAvailable:
-        from .utils.dummy_torch_and_transformers_and_opencv_objects import *  # noqa F403
-    else:
-        from .pipelines import ConsisIDPipeline
-
-    try:
-        if not (is_torch_available() and is_transformers_available() and is_onnx_available()):
-            raise OptionalDependencyNotAvailable()
-    except OptionalDependencyNotAvailable:
-        from .utils.dummy_torch_and_transformers_and_onnx_objects import *  # noqa F403
-    else:
-        from .pipelines import (
-            OnnxStableDiffusionImg2ImgPipeline,
-            OnnxStableDiffusionInpaintPipeline,
-            OnnxStableDiffusionInpaintPipelineLegacy,
-            OnnxStableDiffusionPipeline,
-            OnnxStableDiffusionUpscalePipeline,
-            StableDiffusionOnnxPipeline,
-        )
-
-    try:
-        if not (is_torch_available() and is_librosa_available()):
-            raise OptionalDependencyNotAvailable()
-    except OptionalDependencyNotAvailable:
-        from .utils.dummy_torch_and_librosa_objects import *  # noqa F403
-    else:
-        from .pipelines import AudioDiffusionPipeline, Mel
-
-    try:
-        if not (is_transformers_available() and is_torch_available() and is_note_seq_available()):
-            raise OptionalDependencyNotAvailable()
-    except OptionalDependencyNotAvailable:
-        from .utils.dummy_transformers_and_torch_and_note_seq_objects import *  # noqa F403
-    else:
-        from .pipelines import SpectrogramDiffusionPipeline
-
-    try:
-        if not (is_note_seq_available()):
-            raise OptionalDependencyNotAvailable()
-    except OptionalDependencyNotAvailable:
-        from .utils.dummy_note_seq_objects import *  # noqa F403
-    else:
-        from .pipelines import MidiProcessor
-
-else:
-    import sys
-
-    sys.modules[__name__] = _LazyModule(
-        __name__,
-        globals()["__file__"],
-        _import_structure,
-        module_spec=__spec__,
-        extra_objects={"__version__": __version__},
-    )

diffusers/callbacks.py

@@ -1,244 +0,0 @@
-from typing import Any
-
-from .configuration_utils import ConfigMixin, register_to_config
-from .utils import CONFIG_NAME
-
-
-class PipelineCallback(ConfigMixin):
-    """
-    Base class for all the official callbacks used in a pipeline. This class provides a structure for implementing
-    custom callbacks and ensures that all callbacks have a consistent interface.
-
-    Please implement the following:
-        `tensor_inputs`: This should return a list of tensor inputs specific to your callback. You will only be able to
-        include
-            variables listed in the `._callback_tensor_inputs` attribute of your pipeline class.
-        `callback_fn`: This method defines the core functionality of your callback.
-    """
-
-    config_name = CONFIG_NAME
-
-    @register_to_config
-    def __init__(self, cutoff_step_ratio=1.0, cutoff_step_index=None):
-        super().__init__()
-
-        if (cutoff_step_ratio is None and cutoff_step_index is None) or (
-            cutoff_step_ratio is not None and cutoff_step_index is not None
-        ):
-            raise ValueError("Either cutoff_step_ratio or cutoff_step_index should be provided, not both or none.")
-
-        if cutoff_step_ratio is not None and (
-            not isinstance(cutoff_step_ratio, float) or not (0.0 <= cutoff_step_ratio <= 1.0)
-        ):
-            raise ValueError("cutoff_step_ratio must be a float between 0.0 and 1.0.")
-
-    @property
-    def tensor_inputs(self) -> list[str]:
-        raise NotImplementedError(f"You need to set the attribute `tensor_inputs` for {self.__class__}")
-
-    def callback_fn(self, pipeline, step_index, timesteps, callback_kwargs) -> dict[str, Any]:
-        raise NotImplementedError(f"You need to implement the method `callback_fn` for {self.__class__}")
-
-    def __call__(self, pipeline, step_index, timestep, callback_kwargs) -> dict[str, Any]:
-        return self.callback_fn(pipeline, step_index, timestep, callback_kwargs)
-
-
-class MultiPipelineCallbacks:
-    """
-    This class is designed to handle multiple pipeline callbacks. It accepts a list of PipelineCallback objects and
-    provides a unified interface for calling all of them.
-    """
-
-    def __init__(self, callbacks: list[PipelineCallback]):
-        self.callbacks = callbacks
-
-    @property
-    def tensor_inputs(self) -> list[str]:
-        return [input for callback in self.callbacks for input in callback.tensor_inputs]
-
-    def __call__(self, pipeline, step_index, timestep, callback_kwargs) -> dict[str, Any]:
-        """
-        Calls all the callbacks in order with the given arguments and returns the final callback_kwargs.
-        """
-        for callback in self.callbacks:
-            callback_kwargs = callback(pipeline, step_index, timestep, callback_kwargs)
-
-        return callback_kwargs
-
-
-class SDCFGCutoffCallback(PipelineCallback):
-    """
-    Callback function for Stable Diffusion Pipelines. After certain number of steps (set by `cutoff_step_ratio` or
-    `cutoff_step_index`), this callback will disable the CFG.
-
-    Note: This callback mutates the pipeline by changing the `_guidance_scale` attribute to 0.0 after the cutoff step.
-    """
-
-    tensor_inputs = ["prompt_embeds"]
-
-    def callback_fn(self, pipeline, step_index, timestep, callback_kwargs) -> dict[str, Any]:
-        cutoff_step_ratio = self.config.cutoff_step_ratio
-        cutoff_step_index = self.config.cutoff_step_index
-
-        # Use cutoff_step_index if it's not None, otherwise use cutoff_step_ratio
-        cutoff_step = (
-            cutoff_step_index if cutoff_step_index is not None else int(pipeline.num_timesteps * cutoff_step_ratio)
-        )
-
-        if step_index == cutoff_step:
-            prompt_embeds = callback_kwargs[self.tensor_inputs[0]]
-            prompt_embeds = prompt_embeds[-1:]  # "-1" denotes the embeddings for conditional text tokens.
-
-            pipeline._guidance_scale = 0.0
-
-            callback_kwargs[self.tensor_inputs[0]] = prompt_embeds
-        return callback_kwargs
-
-
-class SDXLCFGCutoffCallback(PipelineCallback):
-    """
-    Callback function for the base Stable Diffusion XL Pipelines. After certain number of steps (set by
-    `cutoff_step_ratio` or `cutoff_step_index`), this callback will disable the CFG.
-
-    Note: This callback mutates the pipeline by changing the `_guidance_scale` attribute to 0.0 after the cutoff step.
-    """
-
-    tensor_inputs = [
-        "prompt_embeds",
-        "add_text_embeds",
-        "add_time_ids",
-    ]
-
-    def callback_fn(self, pipeline, step_index, timestep, callback_kwargs) -> dict[str, Any]:
-        cutoff_step_ratio = self.config.cutoff_step_ratio
-        cutoff_step_index = self.config.cutoff_step_index
-
-        # Use cutoff_step_index if it's not None, otherwise use cutoff_step_ratio
-        cutoff_step = (
-            cutoff_step_index if cutoff_step_index is not None else int(pipeline.num_timesteps * cutoff_step_ratio)
-        )
-
-        if step_index == cutoff_step:
-            prompt_embeds = callback_kwargs[self.tensor_inputs[0]]
-            prompt_embeds = prompt_embeds[-1:]  # "-1" denotes the embeddings for conditional text tokens.
-
-            add_text_embeds = callback_kwargs[self.tensor_inputs[1]]
-            add_text_embeds = add_text_embeds[-1:]  # "-1" denotes the embeddings for conditional pooled text tokens
-
-            add_time_ids = callback_kwargs[self.tensor_inputs[2]]
-            add_time_ids = add_time_ids[-1:]  # "-1" denotes the embeddings for conditional added time vector
-
-            pipeline._guidance_scale = 0.0
-
-            callback_kwargs[self.tensor_inputs[0]] = prompt_embeds
-            callback_kwargs[self.tensor_inputs[1]] = add_text_embeds
-            callback_kwargs[self.tensor_inputs[2]] = add_time_ids
-
-        return callback_kwargs
-
-
-class SDXLControlnetCFGCutoffCallback(PipelineCallback):
-    """
-    Callback function for the Controlnet Stable Diffusion XL Pipelines. After certain number of steps (set by
-    `cutoff_step_ratio` or `cutoff_step_index`), this callback will disable the CFG.
-
-    Note: This callback mutates the pipeline by changing the `_guidance_scale` attribute to 0.0 after the cutoff step.
-    """
-
-    tensor_inputs = [
-        "prompt_embeds",
-        "add_text_embeds",
-        "add_time_ids",
-        "image",
-    ]
-
-    def callback_fn(self, pipeline, step_index, timestep, callback_kwargs) -> dict[str, Any]:
-        cutoff_step_ratio = self.config.cutoff_step_ratio
-        cutoff_step_index = self.config.cutoff_step_index
-
-        # Use cutoff_step_index if it's not None, otherwise use cutoff_step_ratio
-        cutoff_step = (
-            cutoff_step_index if cutoff_step_index is not None else int(pipeline.num_timesteps * cutoff_step_ratio)
-        )
-
-        if step_index == cutoff_step:
-            prompt_embeds = callback_kwargs[self.tensor_inputs[0]]
-            prompt_embeds = prompt_embeds[-1:]  # "-1" denotes the embeddings for conditional text tokens.
-
-            add_text_embeds = callback_kwargs[self.tensor_inputs[1]]
-            add_text_embeds = add_text_embeds[-1:]  # "-1" denotes the embeddings for conditional pooled text tokens
-
-            add_time_ids = callback_kwargs[self.tensor_inputs[2]]
-            add_time_ids = add_time_ids[-1:]  # "-1" denotes the embeddings for conditional added time vector
-
-            # For Controlnet
-            image = callback_kwargs[self.tensor_inputs[3]]
-            image = image[-1:]
-
-            pipeline._guidance_scale = 0.0
-
-            callback_kwargs[self.tensor_inputs[0]] = prompt_embeds
-            callback_kwargs[self.tensor_inputs[1]] = add_text_embeds
-            callback_kwargs[self.tensor_inputs[2]] = add_time_ids
-            callback_kwargs[self.tensor_inputs[3]] = image
-
-        return callback_kwargs
-
-
-class IPAdapterScaleCutoffCallback(PipelineCallback):
-    """
-    Callback function for any pipeline that inherits `IPAdapterMixin`. After certain number of steps (set by
-    `cutoff_step_ratio` or `cutoff_step_index`), this callback will set the IP Adapter scale to `0.0`.
-
-    Note: This callback mutates the IP Adapter attention processors by setting the scale to 0.0 after the cutoff step.
-    """
-
-    tensor_inputs = []
-
-    def callback_fn(self, pipeline, step_index, timestep, callback_kwargs) -> dict[str, Any]:
-        cutoff_step_ratio = self.config.cutoff_step_ratio
-        cutoff_step_index = self.config.cutoff_step_index
-
-        # Use cutoff_step_index if it's not None, otherwise use cutoff_step_ratio
-        cutoff_step = (
-            cutoff_step_index if cutoff_step_index is not None else int(pipeline.num_timesteps * cutoff_step_ratio)
-        )
-
-        if step_index == cutoff_step:
-            pipeline.set_ip_adapter_scale(0.0)
-        return callback_kwargs
-
-
-class SD3CFGCutoffCallback(PipelineCallback):
-    """
-    Callback function for Stable Diffusion 3 Pipelines. After certain number of steps (set by `cutoff_step_ratio` or
-    `cutoff_step_index`), this callback will disable the CFG.
-
-    Note: This callback mutates the pipeline by changing the `_guidance_scale` attribute to 0.0 after the cutoff step.
-    """
-
-    tensor_inputs = ["prompt_embeds", "pooled_prompt_embeds"]
-
-    def callback_fn(self, pipeline, step_index, timestep, callback_kwargs) -> dict[str, Any]:
-        cutoff_step_ratio = self.config.cutoff_step_ratio
-        cutoff_step_index = self.config.cutoff_step_index
-
-        # Use cutoff_step_index if it's not None, otherwise use cutoff_step_ratio
-        cutoff_step = (
-            cutoff_step_index if cutoff_step_index is not None else int(pipeline.num_timesteps * cutoff_step_ratio)
-        )
-
-        if step_index == cutoff_step:
-            prompt_embeds = callback_kwargs[self.tensor_inputs[0]]
-            prompt_embeds = prompt_embeds[-1:]  # "-1" denotes the embeddings for conditional text tokens.
-
-            pooled_prompt_embeds = callback_kwargs[self.tensor_inputs[1]]
-            pooled_prompt_embeds = pooled_prompt_embeds[
-                -1:
-            ]  # "-1" denotes the embeddings for conditional pooled text tokens.
-
-            pipeline._guidance_scale = 0.0
-
-            callback_kwargs[self.tensor_inputs[0]] = prompt_embeds
-            callback_kwargs[self.tensor_inputs[1]] = pooled_prompt_embeds
-        return callback_kwargs

diffusers/commands/__init__.py

@@ -1,27 +0,0 @@
-# Copyright 2026 The HuggingFace Team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-from abc import ABC, abstractmethod
-from argparse import ArgumentParser
-
-
-class BaseDiffusersCLICommand(ABC):
-    @staticmethod
-    @abstractmethod
-    def register_subcommand(parser: ArgumentParser):
-        raise NotImplementedError()
-
-    @abstractmethod
-    def run(self):
-        raise NotImplementedError()

diffusers/commands/custom_blocks.py

@@ -1,140 +0,0 @@
-# Copyright 2026 The HuggingFace Team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-"""`diffusers-cli custom_blocks` — package a local `ModularPipelineBlocks` subclass for the Hub.
-
-Parses `block.py` (or `--block_module_name`), instantiates the chosen block, and calls `save_pretrained` in the current
-working directory.
-"""
-
-import ast
-import importlib.util
-import os
-from argparse import ArgumentParser, Namespace
-from pathlib import Path
-
-from ..utils import logging
-from . import BaseDiffusersCLICommand
-
-
-EXPECTED_PARENT_CLASSES = ["ModularPipelineBlocks"]
-
-
-def conversion_command_factory(args: Namespace):
-    return CustomBlocksCommand(args.block_module_name, args.block_class_name)
-
-
-class CustomBlocksCommand(BaseDiffusersCLICommand):
-    @staticmethod
-    def register_subcommand(parser: ArgumentParser):
-        from argparse import RawDescriptionHelpFormatter
-
-        epilog = (
-            "Examples\n"
-            "  $ diffusers-cli custom_blocks\n"
-            "  $ diffusers-cli custom_blocks --block_module_name my_block.py\n"
-            "  $ diffusers-cli custom_blocks --block_module_name my_block.py --block_class_name MyDenoiseBlock\n"
-            "\n"
-            "Learn more\n"
-            "  Use `diffusers-cli <command> --help` for more information about a command.\n"
-            "  Read the documentation at https://huggingface.co/docs/diffusers\n"
-        )
-
-        conversion_parser = parser.add_parser(
-            "custom_blocks",
-            help="Package a local ModularPipelineBlocks subclass for the Hub.",
-            usage="\n  diffusers-cli custom_blocks [options]",
-            epilog=epilog,
-            formatter_class=RawDescriptionHelpFormatter,
-        )
-        conversion_parser._optionals.title = "Options"
-        conversion_parser.add_argument(
-            "--block_module_name",
-            type=str,
-            default="block.py",
-            help="Module filename in which the custom block will be implemented.",
-        )
-        conversion_parser.add_argument(
-            "--block_class_name",
-            type=str,
-            default=None,
-            help="Name of the custom block. If provided None, we will try to infer it.",
-        )
-        conversion_parser.set_defaults(func=conversion_command_factory)
-
-    def __init__(self, block_module_name: str = "block.py", block_class_name: str = None):
-        self.logger = logging.get_logger("diffusers-cli/custom_blocks")
-        self.block_module_name = Path(block_module_name)
-        self.block_class_name = block_class_name
-
-    def run(self):
-        # determine the block to be saved.
-        out = self._get_class_names(self.block_module_name)
-        classes_found = list({cls for cls, _ in out})
-
-        if self.block_class_name is not None:
-            child_class, parent_class = self._choose_block(out, self.block_class_name)
-            if child_class is None and parent_class is None:
-                raise ValueError(
-                    "`block_class_name` could not be retrieved. Available classes from "
-                    f"{self.block_module_name}:\n{classes_found}"
-                )
-        else:
-            self.logger.info(
-                f"Found classes: {classes_found} will be using {classes_found[0]}. "
-                "If this needs to be changed, re-run the command specifying `block_class_name`."
-            )
-            child_class, parent_class = out[0][0], out[0][1]
-
-        # dynamically get the custom block and initialize it to call `save_pretrained` in the current directory.
-        # the user is responsible for running it, so I guess that is safe?
-        module_name = f"__dynamic__{self.block_module_name.stem}"
-        spec = importlib.util.spec_from_file_location(module_name, str(self.block_module_name))
-        module = importlib.util.module_from_spec(spec)
-        spec.loader.exec_module(module)
-        getattr(module, child_class)().save_pretrained(os.getcwd())
-
-    def _choose_block(self, candidates, chosen=None):
-        for cls, base in candidates:
-            if cls == chosen:
-                return cls, base
-        return None, None
-
-    def _get_class_names(self, file_path):
-        source = file_path.read_text(encoding="utf-8")
-        try:
-            tree = ast.parse(source, filename=file_path)
-        except SyntaxError as e:
-            raise ValueError(f"Could not parse {file_path!r}: {e}") from e
-
-        results: list[tuple[str, str]] = []
-        for node in tree.body:
-            if not isinstance(node, ast.ClassDef):
-                continue
-
-            base_names = [bname for b in node.bases if (bname := self._get_base_name(b)) is not None]
-
-            for allowed in EXPECTED_PARENT_CLASSES:
-                if allowed in base_names:
-                    results.append((node.name, allowed))
-
-        return results
-
-    def _get_base_name(self, node: ast.expr):
-        if isinstance(node, ast.Name):
-            return node.id
-        elif isinstance(node, ast.Attribute):
-            val = self._get_base_name(node.value)
-            return f"{val}.{node.attr}" if val else node.attr
-        return None

diffusers/commands/diffusers_cli.py

@@ -1,69 +0,0 @@
-#!/usr/bin/env python
-# Copyright 2026 The HuggingFace Team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-from argparse import ArgumentParser
-
-from huggingface_hub.cli._output import OutputFormat, out
-
-from .custom_blocks import CustomBlocksCommand
-from .env import EnvironmentCommand
-from .fp16_safetensors import FP16SafetensorsCommand
-from .run import RunCommand
-from .schema import SchemaCommand
-from .skills import SkillsCommand
-
-
-def main():
-    parser = ArgumentParser(
-        prog="diffusers-cli",
-        usage="\n  diffusers-cli [--format <fmt>] <command> [options]",
-    )
-    parser._optionals.title = "Options"
-    parser.add_argument(
-        "--format",
-        choices=[m.value for m in OutputFormat],
-        default=OutputFormat.auto.value,
-        help=(
-            "Output format. 'auto' (default) picks 'agent' when an AI coding agent is detected "
-            "(via CLAUDECODE/CURSOR_AI/AIDER_AI_CONTEXT/... env vars) and 'human' otherwise. "
-            "Must appear before the subcommand."
-        ),
-    )
-    commands_parser = parser.add_subparsers(title="Commands", metavar="<command>")
-
-    # Register commands
-    EnvironmentCommand.register_subcommand(commands_parser)
-    FP16SafetensorsCommand.register_subcommand(commands_parser)
-    CustomBlocksCommand.register_subcommand(commands_parser)
-    RunCommand.register_subcommand(commands_parser)
-    SchemaCommand.register_subcommand(commands_parser)
-    SkillsCommand.register_subcommand(commands_parser)
-
-    # Let's go
-    args = parser.parse_args()
-
-    out.set_mode(OutputFormat(args.format))
-
-    if not hasattr(args, "func"):
-        parser.print_help()
-        exit(1)
-
-    # Run
-    service = args.func(args)
-    service.run()
-
-
-if __name__ == "__main__":
-    main()

diffusers/commands/env.py

@@ -1,185 +0,0 @@
-# Copyright 2026 The HuggingFace Team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-import importlib.metadata
-import platform
-import subprocess
-from argparse import ArgumentParser
-
-import huggingface_hub
-
-from .. import __version__ as version
-from ..utils import (
-    is_accelerate_available,
-    is_bitsandbytes_available,
-    is_gguf_available,
-    is_google_colab,
-    is_nvidia_modelopt_available,
-    is_optimum_quanto_available,
-    is_peft_available,
-    is_safetensors_available,
-    is_torch_available,
-    is_torchao_available,
-    is_transformers_available,
-    is_xformers_available,
-)
-from . import BaseDiffusersCLICommand
-
-
-# (display name, availability_fn, pypi distribution name for importlib.metadata.version)
-_QUANTIZATION_BACKENDS = (
-    ("bitsandbytes", is_bitsandbytes_available, "bitsandbytes"),
-    ("gguf", is_gguf_available, "gguf"),
-    ("optimum-quanto", is_optimum_quanto_available, "optimum-quanto"),
-    ("torchao", is_torchao_available, "torchao"),
-    ("nvidia-modelopt", is_nvidia_modelopt_available, "nvidia-modelopt"),
-)
-
-
-def info_command_factory(_):
-    return EnvironmentCommand()
-
-
-class EnvironmentCommand(BaseDiffusersCLICommand):
-    @staticmethod
-    def register_subcommand(parser: ArgumentParser) -> None:
-        download_parser = parser.add_parser(
-            "env",
-            help="Print versions of diffusers and its dependencies (for bug reports).",
-            usage="\n  diffusers-cli env",
-        )
-        download_parser._optionals.title = "Options"
-        download_parser.set_defaults(func=info_command_factory)
-
-    def run(self) -> dict:
-        hub_version = huggingface_hub.__version__
-
-        safetensors_version = "not installed"
-        if is_safetensors_available():
-            import safetensors
-
-            safetensors_version = safetensors.__version__
-
-        pt_version = "not installed"
-        pt_cuda_available = "NA"
-        if is_torch_available():
-            import torch
-
-            pt_version = torch.__version__
-            pt_cuda_available = torch.cuda.is_available()
-
-        transformers_version = "not installed"
-        if is_transformers_available():
-            import transformers
-
-            transformers_version = transformers.__version__
-
-        accelerate_version = "not installed"
-        if is_accelerate_available():
-            import accelerate
-
-            accelerate_version = accelerate.__version__
-
-        peft_version = "not installed"
-        if is_peft_available():
-            import peft
-
-            peft_version = peft.__version__
-
-        quantization_versions = {}
-        for backend_name, is_available_fn, dist_name in _QUANTIZATION_BACKENDS:
-            if not is_available_fn():
-                continue
-            try:
-                quantization_versions[backend_name] = importlib.metadata.version(dist_name)
-            except importlib.metadata.PackageNotFoundError:
-                quantization_versions[backend_name] = "N/A"
-
-        xformers_version = "not installed"
-        if is_xformers_available():
-            import xformers
-
-            xformers_version = xformers.__version__
-
-        platform_info = platform.platform()
-
-        is_google_colab_str = "Yes" if is_google_colab() else "No"
-
-        accelerator = "NA"
-        if platform.system() in {"Linux", "Windows"}:
-            try:
-                sp = subprocess.Popen(
-                    ["nvidia-smi", "--query-gpu=gpu_name,memory.total", "--format=csv,noheader"],
-                    stdout=subprocess.PIPE,
-                    stderr=subprocess.PIPE,
-                )
-                out_str, _ = sp.communicate()
-                out_str = out_str.decode("utf-8")
-
-                if len(out_str) > 0:
-                    accelerator = out_str.strip()
-            except FileNotFoundError:
-                pass
-        elif platform.system() == "Darwin":  # Mac OS
-            try:
-                sp = subprocess.Popen(
-                    ["system_profiler", "SPDisplaysDataType"],
-                    stdout=subprocess.PIPE,
-                    stderr=subprocess.PIPE,
-                )
-                out_str, _ = sp.communicate()
-                out_str = out_str.decode("utf-8")
-
-                start = out_str.find("Chipset Model:")
-                if start != -1:
-                    start += len("Chipset Model:")
-                    end = out_str.find("\n", start)
-                    accelerator = out_str[start:end].strip()
-
-                    start = out_str.find("VRAM (Total):")
-                    if start != -1:
-                        start += len("VRAM (Total):")
-                        end = out_str.find("\n", start)
-                        accelerator += " VRAM: " + out_str[start:end].strip()
-            except FileNotFoundError:
-                pass
-        else:
-            print("It seems you are running an unusual OS. Could you fill in the accelerator manually?")
-
-        info = {
-            "🤗 Diffusers version": version,
-            "Platform": platform_info,
-            "Running on Google Colab?": is_google_colab_str,
-            "Python version": platform.python_version(),
-            "PyTorch version (GPU?)": f"{pt_version} ({pt_cuda_available})",
-            "Huggingface_hub version": hub_version,
-            "Transformers version": transformers_version,
-            "Accelerate version": accelerate_version,
-            "PEFT version": peft_version,
-            **{f"{name} version": ver for name, ver in quantization_versions.items()},
-            "Safetensors version": safetensors_version,
-            "xFormers version": xformers_version,
-            "Accelerator": accelerator,
-            "Using GPU in script?": "<fill in>",
-            "Using distributed or parallel set-up in script?": "<fill in>",
-        }
-
-        print("\nCopy-and-paste the text below in your GitHub issue and FILL OUT the two last points.\n")
-        print(self.format_dict(info))
-
-        return info
-
-    @staticmethod
-    def format_dict(d: dict) -> str:
-        return "\n".join([f"- {prop}: {val}" for prop, val in d.items()]) + "\n"

diffusers/commands/fp16_safetensors.py

@@ -1,144 +0,0 @@
-# Copyright 2026 The HuggingFace Team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-"""
-Usage example:
-    diffusers-cli fp16_safetensors --ckpt_id=openai/shap-e --fp16 --use_safetensors
-"""
-
-import glob
-import json
-import warnings
-from argparse import ArgumentParser, Namespace
-from importlib import import_module
-
-import huggingface_hub
-import torch
-from huggingface_hub import hf_hub_download
-from packaging import version
-
-from ..utils import logging
-from . import BaseDiffusersCLICommand
-
-
-def conversion_command_factory(args: Namespace):
-    warnings.warn(
-        "`diffusers-cli fp16_safetensors` is deprecated and will be removed in a future version. "
-        "Convert weights to fp16 safetensors directly with `safetensors.torch.save_file` or via "
-        "`pipeline.save_pretrained(..., safe_serialization=True, variant='fp16')`.",
-        FutureWarning,
-        stacklevel=2,
-    )
-    if args.use_auth_token:
-        warnings.warn(
-            "The `--use_auth_token` flag is deprecated and will be removed in a future version."
-            "Authentication is now handled automatically if the user is logged in."
-        )
-    return FP16SafetensorsCommand(args.ckpt_id, args.fp16, args.use_safetensors)
-
-
-class FP16SafetensorsCommand(BaseDiffusersCLICommand):
-    @staticmethod
-    def register_subcommand(parser: ArgumentParser):
-        conversion_parser = parser.add_parser(
-            "fp16_safetensors",
-            help="[DEPRECATED] Convert a Hub checkpoint's weights to fp16 safetensors and push back as a PR.",
-            usage="\n  diffusers-cli fp16_safetensors [options]",
-        )
-        conversion_parser._optionals.title = "Options"
-        conversion_parser.add_argument(
-            "--ckpt_id",
-            type=str,
-            help="Repo id of the checkpoints on which to run the conversion. Example: 'openai/shap-e'.",
-        )
-        conversion_parser.add_argument(
-            "--fp16", action="store_true", help="If serializing the variables in FP16 precision."
-        )
-        conversion_parser.add_argument(
-            "--use_safetensors", action="store_true", help="If serializing in the safetensors format."
-        )
-        conversion_parser.add_argument(
-            "--use_auth_token",
-            action="store_true",
-            help="When working with checkpoints having private visibility. When used `hf auth login` needs to be run beforehand.",
-        )
-        conversion_parser.set_defaults(func=conversion_command_factory)
-
-    def __init__(self, ckpt_id: str, fp16: bool, use_safetensors: bool):
-        self.logger = logging.get_logger("diffusers-cli/fp16_safetensors")
-        self.ckpt_id = ckpt_id
-        self.local_ckpt_dir = f"/tmp/{ckpt_id}"
-        self.fp16 = fp16
-
-        self.use_safetensors = use_safetensors
-
-        if not self.use_safetensors and not self.fp16:
-            raise NotImplementedError(
-                "When `use_safetensors` and `fp16` both are False, then this command is of no use."
-            )
-
-    def run(self):
-        if version.parse(huggingface_hub.__version__) < version.parse("0.9.0"):
-            raise ImportError(
-                "The huggingface_hub version must be >= 0.9.0 to use this command. Please update your huggingface_hub"
-                " installation."
-            )
-        else:
-            from huggingface_hub import create_commit
-            from huggingface_hub._commit_api import CommitOperationAdd
-
-        model_index = hf_hub_download(repo_id=self.ckpt_id, filename="model_index.json")
-        with open(model_index, "r") as f:
-            pipeline_class_name = json.load(f)["_class_name"]
-        pipeline_class = getattr(import_module("diffusers"), pipeline_class_name)
-        self.logger.info(f"Pipeline class imported: {pipeline_class_name}.")
-
-        # Load the appropriate pipeline. We could have used `DiffusionPipeline`
-        # here, but just to avoid potential edge cases.
-        pipeline = pipeline_class.from_pretrained(
-            self.ckpt_id, torch_dtype=torch.float16 if self.fp16 else torch.float32
-        )
-        pipeline.save_pretrained(
-            self.local_ckpt_dir,
-            safe_serialization=True if self.use_safetensors else False,
-            variant="fp16" if self.fp16 else None,
-        )
-        self.logger.info(f"Pipeline locally saved to {self.local_ckpt_dir}.")
-
-        # Fetch all the paths.
-        if self.fp16:
-            modified_paths = glob.glob(f"{self.local_ckpt_dir}/*/*.fp16.*")
-        elif self.use_safetensors:
-            modified_paths = glob.glob(f"{self.local_ckpt_dir}/*/*.safetensors")
-
-        # Prepare for the PR.
-        commit_message = f"Serialize variables with FP16: {self.fp16} and safetensors: {self.use_safetensors}."
-        operations = []
-        for path in modified_paths:
-            operations.append(CommitOperationAdd(path_in_repo="/".join(path.split("/")[4:]), path_or_fileobj=path))
-
-        # Open the PR.
-        commit_description = (
-            "Variables converted by the [`diffusers`' `fp16_safetensors`"
-            " CLI](https://github.com/huggingface/diffusers/blob/main/src/diffusers/commands/fp16_safetensors.py)."
-        )
-        hub_pr_url = create_commit(
-            repo_id=self.ckpt_id,
-            operations=operations,
-            commit_message=commit_message,
-            commit_description=commit_description,
-            repo_type="model",
-            create_pr=True,
-        ).pr_url
-        self.logger.info(f"PR created here: {hub_pr_url}.")

diffusers/commands/run.py

@@ -1,1227 +0,0 @@
-# Copyright 2026 The HuggingFace Team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-"""`diffusers-cli run` — single agentic entry point.
-
-Runs any diffusers pipeline (standard or modular) by forwarding `--pipeline-kwargs` verbatim, saves the output by
-detecting its runtime type, and can submit the same call to an HF Sandbox via `--remote`.
-"""
-
-from __future__ import annotations
-
-import json
-import os
-import sys
-from argparse import ArgumentParser, Namespace, _SubParsersAction
-from pathlib import Path
-from typing import Any
-
-from huggingface_hub.cli._output import out
-
-from diffusers.models.attention_dispatch import _HUB_KERNELS_REGISTRY
-from diffusers.utils import load_image, load_video, logging
-
-from . import BaseDiffusersCLICommand
-
-
-logger = logging.get_logger("diffusers-cli/run")
-
-
-# ---------------------------------------------------------------------------
-# Constants
-# ---------------------------------------------------------------------------
-
-DEFAULT_OUTPUT_DIR = str(Path.home() / ".diffusers" / "cli" / "run" / "outputs")
-DTYPE_CHOICES = ("auto", "float16", "fp16", "bfloat16", "bf16", "float32", "fp32")
-CPU_OFFLOAD_CHOICES = ("model", "group")
-
-
-ATTENTION_BACKEND_CHOICES = ("default", *sorted(b.value for b in _HUB_KERNELS_REGISTRY))
-
-# Kwarg keys whose string value gets auto-loaded before being passed to the pipeline call.
-# Images resolve via `diffusers.utils.load_image` → PIL.Image.Image; videos resolve via
-# `diffusers.utils.load_video` → list[PIL.Image.Image].
-_IMAGE_INPUT_KEYS = (
-    "image",
-    "mask_image",
-    "control_image",
-    "ip_adapter_image",
-    "image_2",
-)
-_VIDEO_INPUT_KEYS = (
-    "video",
-    "control_video",
-)
-_AUDIO_INPUT_KEYS = (
-    "initial_audio_waveforms",
-    "reference_audio",
-    "src_audio",
-)
-
-# Pipeline attribute prefixes that identify a denoiser submodule. Matches base names
-# (`transformer`, `unet`) and their numbered variants (`transformer_2`, etc.).
-_DENOISER_COMPONENT_KEYS = ("transformer", "unet")
-
-_DEFAULT_REMOTE_DEPS = (
-    "diffusers",
-    "accelerate",
-    "transformers",
-    "safetensors",
-    "sentencepiece",  # required by several text-encoder tokenizers (T5, LLaMA, …)
-    "ftfy",  # required by older CLIP text-encoder paths
-)
-
-# Base sandbox image — provides torch + CUDA so `uv pip install --system`
-# only has to add the small Python deps. cuda12.8 is the highest cuda12.x tag
-# below the HF Jobs host driver's CUDA 12.9 max.
-_DEFAULT_REMOTE_IMAGE = "pytorch/pytorch:2.10.0-cuda12.8-cudnn9-runtime"
-
-# Installed console-script name invoked inside the sandbox after the deps land.
-_CONTAINER_CLI_BINARY = "diffusers-cli"
-
-# Working directories inside the sandbox: local media from `--pipeline-kwargs` is uploaded
-# under _SANDBOX_INPUTS_DIR, and the sandbox CLI is told to write its outputs under
-# _SANDBOX_OUTPUTS_DIR so we can download them back afterwards.
-_SANDBOX_INPUTS_DIR = "/tmp/diffusers-cli/inputs"
-_SANDBOX_OUTPUTS_DIR = "/tmp/diffusers-cli/outputs"
-
-RUN_ID_ENV = "DIFFUSERS_CLI_RUN_ID"
-
-# Namespace keys that control *how* a remote run is dispatched, not what the sandbox CLI
-# runs. They are stripped when forwarding argv to the sandbox.
-REMOTE_KEYS = frozenset(
-    {
-        "remote",
-        "flavor",
-        "timeout",
-        "dependencies",
-        "namespace",
-        "image",
-        "keep_alive",
-        "sandbox_id",
-        "idle_timeout",
-        "volume",
-        "func",
-        "format",  # top-level --format is a local rendering flag; never forward to the sandbox
-    }
-)
-
-
-# ---------------------------------------------------------------------------
-# Argparse helpers
-# ---------------------------------------------------------------------------
-
-
-def _add_loading_arguments(parser: ArgumentParser) -> None:
-    parser.add_argument("--model", "-m", required=True, help="Model id on the Hugging Face Hub or local path.")
-    parser.add_argument(
-        "--device-map",
-        default=None,
-        help=(
-            "Component placement. Accepts a torch device string (`cuda`, `cuda:0`, `cpu`, `mps`), "
-            "`balanced` for pipeline-level auto-split across visible GPUs, or a JSON dict of "
-            '`{"<component>": <device>}` for explicit per-component placement. Auto-detected if omitted.'
-        ),
-    )
-    parser.add_argument("--dtype", default="auto", choices=DTYPE_CHOICES, help="Torch dtype for pipeline weights.")
-    parser.add_argument("--variant", default=None, help='Optional weight variant (e.g. "fp16").')
-    parser.add_argument("--revision", default=None, help="Model revision (branch, tag, or commit SHA).")
-    parser.add_argument("--token", default=None, help="Hugging Face token for gated/private models.")
-    parser.add_argument("--trust-remote-code", action="store_true", help="Allow custom code from the Hub.")
-    parser.add_argument(
-        "--lora",
-        action="append",
-        default=None,
-        metavar="JSON",
-        help=(
-            "JSON dict describing a LoRA adapter to attach after the pipeline loads. Repeat to stack "
-            'multiple adapters. Format: \'{"lora_id": "<id>", "lora_scale": <float>}\'. `lora_scale` '
-            "defaults to 1.0; `adapter_name` is optional (auto-generated as `lora_<i>` when stacking)."
-        ),
-    )
-
-
-def _add_optimization_arguments(parser: ArgumentParser) -> None:
-    parser.add_argument(
-        "--cpu-offload",
-        choices=CPU_OFFLOAD_CHOICES,
-        default=None,
-        help=(
-            "Offload pipeline components to CPU during inference. "
-            "'model' uses enable_model_cpu_offload, "
-            "'group' uses pipeline.enable_group_offload(leaf_level, use_stream=True)."
-        ),
-    )
-    parser.add_argument(
-        "--attention-backend",
-        choices=ATTENTION_BACKEND_CHOICES,
-        default="default",
-        help=(
-            "Override the attention backend on the transformer/UNet. "
-            "Only Hub-hosted kernels are exposed — they auto-download on first use."
-        ),
-    )
-    parser.add_argument("--vae-tiling", action="store_true", help="Enable VAE tiling (lower peak VRAM).")
-    parser.add_argument("--vae-slicing", action="store_true", help="Enable VAE slicing (lower peak VRAM).")
-    parser.add_argument(
-        "--context-parallel",
-        action="store_true",
-        help=(
-            "Enable Ulysses-style context parallelism (ulysses_anything mode). "
-            "Requires a DiT-based pipeline and launching the CLI under torchrun with ≥2 GPUs."
-        ),
-    )
-    parser.add_argument(
-        "--compile",
-        nargs="?",
-        const='{"fullgraph": true}',
-        default=None,
-        metavar="JSON",
-        help=(
-            "torch.compile every denoiser submodule on the pipeline. Accepts an optional JSON "
-            'object of kwargs forwarded to `torch.compile`, e.g. \'{"mode": "max-autotune", '
-            '"fullgraph": true}\'. Bare `--compile` uses `fullgraph=true`. Adds a one-time '
-            "compilation cost on the first step but speeds up every subsequent step — worth it "
-            "for multi-step generation (50+ steps)."
-        ),
-    )
-
-
-def _add_output_arguments(parser: ArgumentParser) -> None:
-    parser.add_argument(
-        "--output",
-        "-o",
-        default=None,
-        help=(
-            "Output file or directory. Defaults to "
-            "~/.diffusers/cli/run/outputs/diffusers-run-<YYYYMMDDTHHMMSS>-<short-uuid>/<NNNN>.<ext>."
-        ),
-    )
-    parser.add_argument(
-        "--push-to",
-        default=None,
-        help=(
-            "Upload the generated files to this HF bucket after saving (created if missing). Accepts "
-            "an HF bucket id (`<namespace>/<name>`), an `hf://buckets/<namespace>/<name>[/<subpath>]` "
-            "URI, or a browser URL for the same — a subpath is used as a folder prefix. Under --remote "
-            "the upload runs inside the sandbox; without an explicit --output the bucket becomes the "
-            "sole destination and nothing is downloaded back."
-        ),
-    )
-
-
-def _add_remote_arguments(parser: ArgumentParser) -> None:
-    parser.add_argument(
-        "--remote",
-        action="store_true",
-        help="Run this command in a Hugging Face Sandbox instead of on the local machine.",
-    )
-    parser.add_argument(
-        "--flavor",
-        default="a10g-small",
-        help="HF Sandbox hardware flavor for --remote (e.g. a10g-small, a100-large, cpu-basic).",
-    )
-    parser.add_argument(
-        "--timeout",
-        default="10m",
-        help="Max wallclock for the run command inside the sandbox (e.g. 30m, 2h). Defaults to 10m.",
-    )
-    parser.add_argument(
-        "--dependencies",
-        action="append",
-        default=None,
-        help="Extra pip dependencies to install in the sandbox. Repeat to add multiple.",
-    )
-    parser.add_argument(
-        "--namespace",
-        default=None,
-        help="HF namespace to create the sandbox under (defaults to the current user).",
-    )
-    parser.add_argument(
-        "--image",
-        default=None,
-        help=(
-            "Sandbox image for --remote (defaults to "
-            f"{_DEFAULT_REMOTE_IMAGE!r}). Must provide torch + CUDA; the CLI installs the "
-            "small Python deps on top via `uv pip install --system`."
-        ),
-    )
-    parser.add_argument(
-        "--keep-alive",
-        action="store_true",
-        help=(
-            "Don't terminate the sandbox after the run. Its id is printed so a later --remote run "
-            "can reconnect with --sandbox-id and reuse the warm deps/weights/compile cache."
-        ),
-    )
-    parser.add_argument(
-        "--sandbox-id",
-        default=None,
-        help=(
-            "Reconnect to an existing sandbox (from a prior --keep-alive run) instead of creating a new "
-            "one, reusing its warm deps/weights/compile cache. Implies --keep-alive; stop it with "
-            "`hf sandbox kill <id>`."
-        ),
-    )
-    parser.add_argument(
-        "--idle-timeout",
-        default="10m",
-        help=(
-            "Auto-shutdown the sandbox after this much inactivity (e.g. 30m, 1h). Defaults to 10m. "
-            "Only applied on new sandbox creation — ignored when reconnecting via --sandbox-id."
-        ),
-    )
-    parser.add_argument(
-        "--volume",
-        action="append",
-        default=None,
-        metavar="BUCKET_ID[:MOUNT_PATH]",
-        help=(
-            "Mount an HF bucket into the sandbox as a read-write directory. Repeatable. Format: "
-            "`<namespace>/<name>` (mounts at `/mnt/buckets/<namespace>/<name>`) or "
-            "`<namespace>/<name>:/some/path` for a custom path. Reference mounted files from "
-            "--pipeline-kwargs like any other local path. Applied only on new sandbox creation — "
-            "ignored when reconnecting via --sandbox-id."
-        ),
-    )
-
-
-# ---------------------------------------------------------------------------
-# Pipeline loading + optimization
-# ---------------------------------------------------------------------------
-
-
-def _resolve_dtype(name: str | None):
-    if name in (None, "auto"):
-        return "auto"
-    import torch
-
-    mapping = {
-        "fp32": torch.float32,
-        "float32": torch.float32,
-        "fp16": torch.float16,
-        "float16": torch.float16,
-        "bf16": torch.bfloat16,
-        "bfloat16": torch.bfloat16,
-    }
-    if name not in mapping:
-        raise ValueError(f"Unknown dtype: {name}")
-    return mapping[name]
-
-
-def _resolve_device_map(raw: str | None) -> str | dict:
-    """Parse `--device-map` into a value acceptable by `from_pretrained(device_map=...)`.
-
-    Returns a JSON dict if the value looks like one, `"balanced"` verbatim, or a single-device string (e.g. `"cuda"`,
-    `"cuda:1"`, `"cpu"`, `"mps"`). Auto-detects when `raw is None`, pinning to `cuda:$LOCAL_RANK` under torchrun.
-    """
-    if raw is None:
-        from diffusers.utils.torch_utils import torch_device
-
-        if torch_device == "cuda":
-            local_rank = os.environ.get("LOCAL_RANK")
-            if local_rank is not None:
-                import torch
-
-                torch.cuda.set_device(int(local_rank))
-                return f"cuda:{local_rank}"
-        return torch_device
-
-    if raw.strip().startswith("{"):
-        try:
-            parsed = json.loads(raw)
-        except json.JSONDecodeError as e:
-            raise SystemExit(f"--device-map must be a device string or a JSON dict: {e}") from e
-        if not isinstance(parsed, dict):
-            raise SystemExit("--device-map JSON must decode to an object.")
-        return parsed
-
-    return raw
-
-
-def _apply_cpu_offload(pipeline: Any, mode: str, device_map: str | dict) -> None:
-    """Apply model or group CPU offload. Requires a single-device target (not balanced or dict)."""
-    if not isinstance(device_map, str) or device_map == "balanced":
-        raise SystemExit(
-            "--cpu-offload requires --device-map to be a single device string (e.g. 'cuda'); "
-            f"got {device_map!r}. balanced/dict placement is incompatible with CPU offload."
-        )
-
-    if mode == "model":
-        pipeline.enable_model_cpu_offload(device=device_map)
-    elif mode == "group":
-        import torch
-
-        pipeline.enable_group_offload(
-            onload_device=torch.device(device_map),
-            offload_type="leaf_level",
-            use_stream=True,
-        )
-
-
-def _set_attention_backend(pipeline: Any, backend: str) -> None:
-    transformer = getattr(pipeline, "transformer", None)
-    if transformer is None or not hasattr(transformer, "set_attention_backend"):
-        logger.warning(
-            f"--attention-backend is only supported on transformer-based pipelines; "
-            f"{type(pipeline).__name__} uses the legacy UNet attention path."
-        )
-        return
-    try:
-        transformer.set_attention_backend(backend)
-    except (ValueError, ImportError, RuntimeError) as e:
-        logger.warning(
-            f"Attention backend {backend!r} could not be set on {type(transformer).__name__}: "
-            f"{type(e).__name__}: {e}. Falling back to the model's default backend."
-        )
-
-
-def _enable_context_parallel(pipeline: Any) -> None:
-    import torch
-
-    if not torch.distributed.is_available():
-        raise SystemExit("--context-parallel requires a torch build with distributed support.")
-
-    if not torch.distributed.is_initialized():
-        # Hybrid backend: ulysses_anything's per-rank size coordination wants Gloo on CPU
-        # (avoids H2D/D2H for a tiny int tensor); the main attention all-to-all stays on NCCL.
-        torch.distributed.init_process_group(backend="cpu:gloo,cuda:nccl")
-
-    transformer = getattr(pipeline, "transformer", None)
-    if transformer is None or not hasattr(transformer, "enable_parallelism"):
-        raise SystemExit(
-            "--context-parallel requires a DiT-based pipeline. "
-            f"{type(pipeline).__name__} does not expose a `transformer` with `enable_parallelism`."
-        )
-
-    from diffusers import ContextParallelConfig
-
-    transformer.enable_parallelism(
-        config=ContextParallelConfig(
-            ulysses_degree=torch.distributed.get_world_size(),
-            ring_degree=1,
-            ulysses_anything=True,
-        )
-    )
-
-
-def _apply_optimizations(pipeline: Any, args: Namespace) -> None:
-    """Apply VAE tiling/slicing, attention backend, context-parallel, and torch.compile toggles."""
-    vae = getattr(pipeline, "vae", None)
-    if args.vae_tiling and vae is not None and hasattr(vae, "enable_tiling"):
-        vae.enable_tiling()
-    if args.vae_slicing and vae is not None and hasattr(vae, "enable_slicing"):
-        vae.enable_slicing()
-    if args.attention_backend != "default":
-        _set_attention_backend(pipeline, args.attention_backend)
-    if args.context_parallel:
-        _enable_context_parallel(pipeline)
-    if args.compile is not None:
-        if args.context_parallel:
-            logger.warning("--compile is currently not supported with --context-parallel; skipping compile.")
-        else:
-            _compile_denoiser(pipeline, args.compile)
-
-
-def _compile_denoiser(pipeline: Any, compile_spec: str) -> None:
-    """Compile every `transformer*` and `unet*` submodule on the pipeline.
-
-    `compile_spec` is the raw JSON string from `--compile` (`"{}"` for bare flag). Decoded into kwargs and forwarded
-    verbatim to the compile call.
-
-    Prefers regional compilation via `module.compile_repeated_blocks(**kwargs)` — only compiles the repeated inner
-    blocks (the bulk of the compute), much faster first-step latency than compiling the whole module. Falls back to
-    full `torch.compile` if the model doesn't expose `_repeated_blocks`.
-    """
-    import torch
-
-    try:
-        compile_kwargs = json.loads(compile_spec)
-    except json.JSONDecodeError as e:
-        raise SystemExit(f"--compile must be valid JSON: {e}") from e
-    if not isinstance(compile_kwargs, dict):
-        raise SystemExit("--compile must decode to a JSON object.")
-
-    for attr in dir(pipeline):
-        if not any(attr.startswith(key) for key in _DENOISER_COMPONENT_KEYS):
-            continue
-        module = getattr(pipeline, attr, None)
-        if not isinstance(module, torch.nn.Module):
-            continue
-
-        if getattr(module, "_repeated_blocks", None):
-            # Regional compile — only the repeated blocks. Mutates `module` in place.
-            module.compile_repeated_blocks(**compile_kwargs)
-        else:
-            # No regional metadata declared; fall back to compiling the whole module.
-            setattr(pipeline, attr, torch.compile(module, **compile_kwargs))
-
-
-def _load_lora(pipeline: Any, args: Namespace) -> None:
-    """Attach one or more LoRA adapters. Each `--lora` value is a JSON dict.
-
-    Per-entry fields: `lora_id` (required), `lora_scale` (optional float, default 1.0), `adapter_name` (optional;
-    auto-generated as `lora_<i>` when stacking). Multiple `--lora` flags stack via a single `set_adapters(...)` call at
-    the end.
-    """
-    if not args.lora:
-        return
-    specs = []
-    for raw in args.lora:
-        try:
-            parsed = json.loads(raw)
-        except json.JSONDecodeError as e:
-            raise SystemExit(f"--lora must be valid JSON: {e}") from e
-        if not isinstance(parsed, dict):
-            raise SystemExit(f"--lora must decode to a JSON object; got {type(parsed).__name__}.")
-        specs.append(parsed)
-    if not hasattr(pipeline, "load_lora_weights"):
-        raise SystemExit(f"{type(pipeline).__name__} does not support LoRA loading.")
-
-    names: list[str] = []
-    scales: list[float] = []
-    for i, spec in enumerate(specs):
-        lora_id = spec.get("lora_id")
-        if not lora_id:
-            raise SystemExit(f"--lora entry {i} is missing 'lora_id'.")
-        adapter_name = spec.get("adapter_name") or (f"lora_{i}" if len(specs) > 1 else "default")
-        pipeline.load_lora_weights(lora_id, adapter_name=adapter_name)
-        names.append(adapter_name)
-        scales.append(float(spec.get("lora_scale", 1.0)))
-
-    if hasattr(pipeline, "set_adapters"):
-        pipeline.set_adapters(names, adapter_weights=scales)
-
-
-def _load_pipeline(args: Namespace) -> Any:
-    import diffusers
-
-    # Detect modular repos by trying the standard config; `ModularPipeline` repos ship
-    # `modular_model_index.json` instead of `model_index.json`, so `load_config` OSErrors.
-    try:
-        diffusers.DiffusionPipeline.load_config(args.model, token=args.token, revision=args.revision)
-        modular = False
-    except OSError:
-        modular = True
-
-    dtype = _resolve_dtype(args.dtype)
-    device_map = _resolve_device_map(args.device_map)
-    common_kwargs: dict[str, Any] = {
-        "trust_remote_code": args.trust_remote_code,
-    }
-    if dtype != "auto":
-        common_kwargs["torch_dtype"] = dtype
-    if args.variant:
-        common_kwargs["variant"] = args.variant
-    if args.token:
-        common_kwargs["token"] = args.token
-    # CPU offload sets up its own placement hooks, so leave weights on CPU at load time.
-    if not args.cpu_offload:
-        common_kwargs["device_map"] = device_map
-
-    if modular:
-        # ModularPipeline.from_pretrained fetches only the pipeline config; component
-        # weights come in via load_components(). `revision` scopes the config fetch,
-        # so it stays on from_pretrained — each ComponentSpec pins its own revision,
-        # and forwarding a global `revision` to load_components() would override those.
-        pipeline = diffusers.ModularPipeline.from_pretrained(
-            args.model,
-            trust_remote_code=args.trust_remote_code,
-            token=args.token,
-            revision=args.revision,
-        )
-        pipeline.load_components(**common_kwargs)
-    else:
-        pipeline = diffusers.DiffusionPipeline.from_pretrained(args.model, revision=args.revision, **common_kwargs)
-
-    _load_lora(pipeline, args)
-    if args.cpu_offload:
-        _apply_cpu_offload(pipeline, args.cpu_offload, device_map)
-    _apply_optimizations(pipeline, args)
-
-    return pipeline
-
-
-# ---------------------------------------------------------------------------
-# Pipeline call helpers
-# ---------------------------------------------------------------------------
-
-
-def _parse_pipeline_kwargs(raw: str | None) -> dict[str, Any]:
-    if not raw:
-        return {}
-    try:
-        parsed = json.loads(raw)
-    except json.JSONDecodeError as e:
-        raise SystemExit(f"--pipeline-kwargs must be valid JSON: {e}") from e
-    if not isinstance(parsed, dict):
-        raise SystemExit("--pipeline-kwargs must decode to a JSON object.")
-    return parsed
-
-
-def _load_audio(url_or_path: str) -> tuple[Any, int]:
-    """Load audio from a URL or local path via torchaudio. Returns `(waveform, sampling_rate)`."""
-    import torchaudio
-
-    if url_or_path.startswith(("http://", "https://")):
-        import io
-
-        import httpx
-
-        from ..utils.constants import DIFFUSERS_REQUEST_TIMEOUT
-
-        resp = httpx.get(url_or_path, follow_redirects=True, timeout=DIFFUSERS_REQUEST_TIMEOUT)
-        resp.raise_for_status()
-        return torchaudio.load(io.BytesIO(resp.content))
-    return torchaudio.load(url_or_path)
-
-
-def _resolve_media_inputs(call_kwargs: dict[str, Any]) -> None:
-    """Replace string paths/URLs at known media-input keys with loaded tensors.
-
-    Images resolve to `PIL.Image.Image` via `load_image`; videos to `list[PIL.Image.Image]` via `load_video`; audio to
-    a `torch.Tensor` via `_load_audio` (also auto-sets the paired sampling-rate kwarg for `initial_audio_waveforms` if
-    the user didn't supply it). A `list[str]` at any key is treated as a batch: each entry is loaded and the value
-    becomes a list of loaded objects. Non-string, non-list values pass through untouched.
-    """
-
-    def _is_string_list(v: Any) -> bool:
-        return isinstance(v, list) and bool(v) and all(isinstance(x, str) for x in v)
-
-    for key in _IMAGE_INPUT_KEYS:
-        value = call_kwargs.get(key)
-        if isinstance(value, str):
-            call_kwargs[key] = load_image(value)
-        elif _is_string_list(value):
-            call_kwargs[key] = [load_image(v) for v in value]
-    for key in _VIDEO_INPUT_KEYS:
-        value = call_kwargs.get(key)
-        if isinstance(value, str):
-            call_kwargs[key] = load_video(value)
-        elif _is_string_list(value):
-            call_kwargs[key] = [load_video(v) for v in value]
-    for key in _AUDIO_INPUT_KEYS:
-        value = call_kwargs.get(key)
-        if isinstance(value, str):
-            waveform, sr = _load_audio(value)
-            call_kwargs[key] = waveform
-            if key == "initial_audio_waveforms" and "initial_audio_sampling_rate" not in call_kwargs:
-                call_kwargs["initial_audio_sampling_rate"] = sr
-        elif _is_string_list(value):
-            pairs = [_load_audio(v) for v in value]
-            call_kwargs[key] = [w for w, _ in pairs]
-            if key == "initial_audio_waveforms" and "initial_audio_sampling_rate" not in call_kwargs:
-                # All batched waveforms must share a sampling rate; use the first entry's.
-                call_kwargs["initial_audio_sampling_rate"] = pairs[0][1]
-
-
-def _get_generator(seed: int | None, device: str):
-    if seed is None:
-        return None
-    import torch
-
-    generator_device = "cpu" if device == "mps" else device
-    return torch.Generator(device=generator_device).manual_seed(seed)
-
-
-def _unwrap_pipeline_output(result: Any) -> Any:
-    """Unwrap a pipeline-output object into the raw payload the saver can dispatch on."""
-    if hasattr(result, "images"):
-        return result.images
-    if hasattr(result, "frames"):
-        return result.frames[0]
-    if hasattr(result, "audios"):
-        return result.audios
-    return result
-
-
-# ---------------------------------------------------------------------------
-# Output saving (dispatch by type)
-# ---------------------------------------------------------------------------
-
-
-def _get_or_create_run_id() -> str:
-    """Return the current run's id, creating one if not yet set.
-
-    Format: `diffusers-run-<YYYYMMDDTHHMMSS>-<6-char-uuid>`. Same id is reused as the local output subdirectory, the
-    remote bucket prefix, and the container-side `RUN_ID_ENV` so a run's artifacts are traceable end-to-end.
-    """
-    import uuid
-    from datetime import datetime
-
-    existing = os.environ.get(RUN_ID_ENV)
-    if existing:
-        return existing
-    run_id = f"diffusers-run-{datetime.now().strftime('%Y%m%dT%H%M%S')}-{uuid.uuid4().hex[:6]}"
-    os.environ[RUN_ID_ENV] = run_id
-    return run_id
-
-
-def _resolve_output_paths(task: str, num: int, explicit: str | None, ext: str) -> list[Path]:
-    if explicit is None:
-        base = Path(DEFAULT_OUTPUT_DIR) / _get_or_create_run_id()
-        base.mkdir(parents=True, exist_ok=True)
-        return [base / f"{i:04d}.{ext}" for i in range(num)]
-
-    p = Path(explicit)
-    if explicit.endswith(os.sep) or p.is_dir():
-        p.mkdir(parents=True, exist_ok=True)
-        return [p / f"{i:04d}.{ext}" for i in range(num)]
-
-    p.parent.mkdir(parents=True, exist_ok=True)
-    if num == 1:
-        return [p]
-    stem, suffix = p.stem, p.suffix or f".{ext}"
-    return [p.with_name(f"{stem}-{i:04d}{suffix}") for i in range(num)]
-
-
-def _as_pil_list(value: Any):
-    try:
-        from PIL.Image import Image as PILImage
-    except ImportError:
-        return None
-    if isinstance(value, PILImage):
-        return [value]
-    if isinstance(value, (list, tuple)) and value and all(isinstance(v, PILImage) for v in value):
-        return list(value)
-    return None
-
-
-def _as_frame_sequence(value: Any):
-    try:
-        from PIL.Image import Image as PILImage
-    except ImportError:
-        PILImage = None  # type: ignore[assignment]
-
-    if isinstance(value, (list, tuple)) and len(value) >= 2:
-        first = value[0]
-        if PILImage is not None and isinstance(first, PILImage):
-            return list(value)
-        try:
-            import numpy as np
-
-            if isinstance(first, np.ndarray):
-                return list(value)
-        except ImportError:
-            pass
-    return None
-
-
-def _as_audio_arrays(value: Any):
-    try:
-        import numpy as np
-    except ImportError:
-        return None
-    if isinstance(value, np.ndarray) and value.ndim <= 2:
-        return [value]
-    if isinstance(value, (list, tuple)) and value and all(isinstance(v, np.ndarray) for v in value):
-        return list(value)
-    return None
-
-
-def _save_audio_arrays(audios, sampling_rate: int, args: Namespace, task: str) -> list[str]:
-    """Write each numpy audio array to a 16-bit PCM WAV at `sampling_rate` Hz.
-
-    Uses the stdlib `wave` module so no scipy dependency is required.
-    """
-    import wave
-
-    import numpy as np
-
-    paths = _resolve_output_paths(task, len(audios), args.output, ext="wav")
-    saved: list[str] = []
-    for audio, path in zip(audios, paths):
-        data = np.asarray(audio)
-        if data.dtype.kind == "f":
-            data = (np.clip(data, -1.0, 1.0) * 32767).astype(np.int16)
-        else:
-            data = data.astype(np.int16)
-        if data.ndim == 1:
-            n_channels = 1
-        else:
-            # Heuristic: shorter axis is channels (interleaved layout for `wave` is
-            # samples × channels, so transpose if needed).
-            if data.shape[0] < data.shape[-1]:
-                data = data.T
-            n_channels = data.shape[1]
-        with wave.open(str(path), "wb") as w:
-            w.setnchannels(n_channels)
-            w.setsampwidth(2)  # 16-bit PCM
-            w.setframerate(sampling_rate)
-            w.writeframes(data.tobytes())
-        saved.append(str(path))
-    return saved
-
-
-def _save_output(value: Any, args: Namespace, task: str) -> list[str]:
-    """Save `value` by dispatching on its runtime type."""
-    pil_images = _as_pil_list(value)
-    if pil_images is not None:
-        paths = _resolve_output_paths(task, len(pil_images), args.output, ext="png")
-        for img, path in zip(pil_images, paths):
-            img.save(path)
-        return [str(p) for p in paths]
-
-    frames = _as_frame_sequence(value)
-    if frames is not None:
-        from diffusers.utils import export_to_video
-
-        path = _resolve_output_paths(task, 1, args.output, ext="mp4")[0]
-        export_to_video(frames, str(path), fps=args.fps)
-        return [str(path)]
-
-    audios = _as_audio_arrays(value)
-    if audios is not None:
-        return _save_audio_arrays(audios, args.sampling_rate or 16000, args, task)
-
-    path = _resolve_output_paths(task, 1, args.output, ext="json")[0]
-    Path(path).write_text(json.dumps(value, default=str, indent=2))
-    return [str(path)]
-
-
-# ---------------------------------------------------------------------------
-# Hub bucket upload (--push-to)
-# ---------------------------------------------------------------------------
-
-
-def _parse_push_to(spec: str) -> tuple[str, str]:
-    """Split `--push-to` into a bucket id and an optional subpath prefix.
-
-    Accepts an HF bucket id (`<namespace>/<name>[/<subpath>]`), a canonical
-    `hf://buckets/<namespace>/<name>[/<subpath>]` URI, or a Hub web URL for the same. Non-bucket URIs (models,
-    datasets, spaces) are rejected — `--push-to` targets storage buckets only.
-    """
-    from huggingface_hub import parse_hf_uri
-
-    # Bare shorthand → canonical URI so a single parser handles every accepted form.
-    if not spec.startswith(("hf://", "http://", "https://")):
-        spec = f"hf://buckets/{spec.strip('/')}"
-    uri = parse_hf_uri(spec)
-    if not uri.is_bucket:
-        raise SystemExit(f"--push-to must point at a bucket; got {uri.type!r} URI {spec!r}.")
-    return uri.id, uri.path_in_repo
-
-
-def _push_outputs(args: Namespace, saved_paths: list[str], task: str) -> dict[str, Any] | None:
-    """Upload `saved_paths` to the `--push-to` bucket. Returns a summary or None."""
-    if not args.push_to:
-        return None
-
-    from huggingface_hub import HfApi
-
-    bucket_id, subpath = _parse_push_to(args.push_to)
-    api = HfApi(token=args.token)
-    api.create_bucket(bucket_id, exist_ok=True)
-
-    run_id = _get_or_create_run_id()
-    prefix = f"{subpath}/{run_id}" if subpath else run_id
-    add = [(local, f"{prefix}/{Path(local).name}") for local in saved_paths]
-    api.batch_bucket_files(bucket_id, add=add)
-
-    uploaded = [f"hf://buckets/{bucket_id}/{dest}" for _, dest in add]
-    return {"bucket_id": bucket_id, "uploaded": uploaded}
-
-
-# ---------------------------------------------------------------------------
-# Remote execution (HF Sandbox)
-# ---------------------------------------------------------------------------
-
-
-def _build_task_kwargs(args: Namespace) -> dict[str, Any]:
-    """Pick out the kwargs the sandbox CLI should invoke the task with."""
-    out: dict[str, Any] = {}
-    for key, value in vars(args).items():
-        if key in REMOTE_KEYS or value is None or value is False:
-            continue
-        out[key] = value
-    return out
-
-
-def _kwargs_to_argv(task: str, task_kwargs: dict[str, Any]) -> list[str]:
-    """Render `task_kwargs` as the argv list the sandbox CLI's argparse will see."""
-    argv: list[str] = [task]
-    for key, value in task_kwargs.items():
-        flag = "--" + key.replace("_", "-")
-        if value is True:
-            argv.append(flag)
-        elif isinstance(value, list):
-            for item in value:
-                argv.extend([flag, str(item)])
-        else:
-            argv.extend([flag, str(value)])
-    return argv
-
-
-def _duration_to_seconds(value: str) -> float:
-    """Parse a duration like `30s`, `10m`, `2h` (or a bare number of seconds) into seconds."""
-    value = value.strip()
-    units = {"s": 1, "m": 60, "h": 3600}
-    if value and value[-1] in units:
-        return float(value[:-1]) * units[value[-1]]
-    return float(value)
-
-
-def _upload_inputs_to_sandbox(args: Namespace, sbx: Any, run_id: str) -> None:
-    """Upload local media paths in `--pipeline-kwargs` into the sandbox and rewrite the JSON in place.
-
-    Walks known image/video/audio-input keys; any string value that resolves to a local file is uploaded to
-    `<_SANDBOX_INPUTS_DIR>/<run_id>/<key>_<basename>` and the JSON path is rewritten to that in-sandbox path. URLs,
-    `hf://` URIs, and non-existent paths pass through untouched.
-    """
-    if not args.pipeline_kwargs:
-        return
-    try:
-        parsed = json.loads(args.pipeline_kwargs)
-    except json.JSONDecodeError:
-        return  # the sandbox CLI will fail loudly with a parse error later
-    if not isinstance(parsed, dict):
-        return
-
-    def _upload_one(key: str, index: int | None, local_str: str) -> str:
-        # `index` is None for scalar entries, an int for list entries (used to disambiguate names).
-        local = Path(local_str)
-        suffix = f"_{index}" if index is not None else ""
-        remote_path = f"{_SANDBOX_INPUTS_DIR}/{run_id}/{key}{suffix}_{local.name}"
-        sbx.files.upload(str(local), remote_path)
-        return remote_path
-
-    uploaded = 0
-    for key in (*_IMAGE_INPUT_KEYS, *_VIDEO_INPUT_KEYS, *_AUDIO_INPUT_KEYS):
-        value = parsed.get(key)
-        if isinstance(value, str) and Path(value).is_file():
-            parsed[key] = _upload_one(key, None, value)
-            uploaded += 1
-        elif isinstance(value, list):
-            # Batched inputs: upload each local path, leave URLs/hf:// URIs alone.
-            new_list = list(value)
-            for i, entry in enumerate(value):
-                if isinstance(entry, str) and Path(entry).is_file():
-                    new_list[i] = _upload_one(key, i, entry)
-                    uploaded += 1
-            parsed[key] = new_list
-
-    if uploaded:
-        logger.info(f"uploaded {uploaded} local input file(s) to the sandbox")
-        args.pipeline_kwargs = json.dumps(parsed)
-
-
-def _download_outputs_from_sandbox(sbx: Any, sandbox_dir: str, local_dir: Path) -> list[str]:
-    """Download every file the sandbox CLI wrote under `sandbox_dir` into `local_dir`."""
-    local_dir.mkdir(parents=True, exist_ok=True)
-    saved: list[str] = []
-    for entry in sbx.files.list(sandbox_dir):
-        if entry.type != "file":
-            continue
-        target = local_dir / Path(entry.path).name
-        sbx.files.download(entry.path, str(target))
-        saved.append(str(target))
-    return saved
-
-
-def _maybe_submit_remote(args: Namespace, task: str) -> bool:
-    """If `--remote` was set, run this invocation inside an HF Sandbox and return True."""
-    if not args.remote:
-        return False
-
-    import shlex
-    import time
-
-    from huggingface_hub import get_token
-    from huggingface_hub.utils import send_telemetry
-
-    import diffusers
-
-    try:
-        from huggingface_hub import Sandbox
-    except ImportError:
-        raise SystemExit(
-            "--remote requires huggingface_hub>=1.23 for HF Sandbox support. "
-            "Upgrade with `pip install -U huggingface_hub`."
-        )
-
-    if Path(args.model).exists():
-        raise SystemExit(
-            f"--model {args.model!r} is a local path; the sandbox can't see it. "
-            "Pass a Hub repo id so the sandbox can download it."
-        )
-
-    hf_token = args.token or get_token()
-    run_id = _get_or_create_run_id()
-
-    # An explicit --push-to means the bucket is the user's destination, so skip the local
-    # download unless they also asked for a local path via --output.
-    user_bucket = bool(args.push_to)
-    download_locally = (not user_bucket) or (args.output is not None)
-    local_dir = Path(args.output) if args.output else Path(DEFAULT_OUTPUT_DIR) / run_id
-
-    use_existing_sandbox = bool(args.sandbox_id)
-    keep_alive = args.keep_alive or use_existing_sandbox
-    if use_existing_sandbox and args.volume:
-        logger.warning(
-            "--volume is ignored when reconnecting to an existing sandbox (mounts are set at creation time)."
-        )
-    if use_existing_sandbox:
-        logger.info(f"reconnecting to sandbox {args.sandbox_id!r}...")
-        sbx = Sandbox.connect(args.sandbox_id, token=hf_token)
-    else:
-        logger.info(f"creating sandbox on flavor={args.flavor!r}...")
-        create_kwargs: dict[str, Any] = {
-            "image": args.image or _DEFAULT_REMOTE_IMAGE,
-            "flavor": args.flavor,
-            "forward_hf_token": True,
-            "token": hf_token,
-            "env": {
-                "HF_ENABLE_PARALLEL_LOADING": "1",
-                "DIFFUSERS_VERBOSITY": os.environ.get("DIFFUSERS_VERBOSITY", "info"),
-            },
-            "idle_timeout": args.idle_timeout,
-        }
-        if args.volume:
-            from huggingface_hub import Volume
-
-            volumes = []
-            for spec in args.volume:
-                bucket_id, sep, mount_path = spec.partition(":")
-                if not sep:
-                    mount_path = f"/mnt/buckets/{bucket_id}"
-                if bucket_id.count("/") != 1:
-                    raise SystemExit(f"--volume: bucket id must be <namespace>/<name>, got {bucket_id!r}")
-                if not mount_path.startswith("/"):
-                    raise SystemExit(f"--volume: mount path must be absolute, got {mount_path!r}")
-                volumes.append(Volume(type="bucket", source=bucket_id, mount_path=mount_path))
-            create_kwargs["volumes"] = volumes
-        if args.namespace is not None:
-            create_kwargs["namespace"] = args.namespace
-        sbx = Sandbox.create(**create_kwargs)
-
-    def _stream(chunk: str) -> None:
-        sys.stderr.write(chunk)
-        sys.stderr.flush()
-
-    exit_code = 0
-    saved: list[str] = []
-    run_seconds = 0.0
-    try:
-        _upload_inputs_to_sandbox(args, sbx, run_id)
-
-        dependencies = list(_DEFAULT_REMOTE_DEPS)
-        if args.dependencies:
-            dependencies.extend(args.dependencies)
-        # --break-system-packages bypasses PEP 668; harmless in a throwaway sandbox. uv is a
-        # near no-op when the deps are already satisfied, so this stays cheap on a reused sandbox.
-        install_cmd = shlex.join(["uv", "pip", "install", "--system", "--break-system-packages", *dependencies])
-        logger.info("installing dependencies in the sandbox...")
-        sbx.run(install_cmd, on_stdout=_stream, on_stderr=_stream)
-
-        # Per-run outputs subdirectory so a reused sandbox doesn't leak files from prior runs
-        # into this run's download set.
-        sandbox_output_dir = f"{_SANDBOX_OUTPUTS_DIR}/{run_id}"
-        task_kwargs = _build_task_kwargs(args)
-        task_kwargs["output"] = sandbox_output_dir + "/"
-        cli_argv = _kwargs_to_argv(task, task_kwargs)
-        # Suppress the container CLI's own `out.result(...)` payload — the outer wrapper owns the
-        # final structured output for --remote runs.
-        format_argv = ["--format", "quiet"]
-        # torchrun wraps the CLI for --context-parallel so torch.distributed initializes across
-        # every visible GPU before the run command starts.
-        if args.context_parallel:
-            cli_argv = [
-                "torchrun",
-                "--nproc-per-node=gpu",
-                "-m",
-                "diffusers.commands.diffusers_cli",
-                *format_argv,
-                *cli_argv,
-            ]
-        else:
-            cli_argv = [_CONTAINER_CLI_BINARY, *format_argv, *cli_argv]
-
-        started = time.perf_counter()
-        # Per-invocation env: RUN_ID_ENV must be fresh each run. Sandbox.create-time env is
-        # baked in and would go stale on reused sandboxes, silently reusing the initial run's
-        # bucket prefix in `_push_outputs`.
-        result = sbx.run(
-            cli_argv,
-            env={RUN_ID_ENV: run_id},
-            on_stdout=_stream,
-            on_stderr=_stream,
-            timeout=_duration_to_seconds(args.timeout),
-            check=False,
-        )
-        run_seconds = time.perf_counter() - started
-        exit_code = result.exit_code
-
-        if exit_code == 0 and download_locally:
-            saved = _download_outputs_from_sandbox(sbx, sandbox_output_dir, local_dir)
-    finally:
-        if keep_alive:
-            logger.info(
-                f"sandbox {sbx.id} kept alive — reconnect with "
-                f"`--remote --sandbox-id {sbx.id}`, stop with `hf sandbox kill {sbx.id}`."
-            )
-        else:
-            sbx.kill()
-
-    send_telemetry(
-        topic="diffusers/cli/run/remote",
-        library_name="diffusers",
-        library_version=diffusers.__version__,
-    )
-
-    payload: dict[str, Any] = {
-        "exit_code": exit_code,
-        "run_seconds": round(run_seconds, 1),
-    }
-    if keep_alive:
-        payload["sandbox_id"] = sbx.id
-    if download_locally:
-        payload["outputs"] = saved
-    if args.push_to:
-        bucket_id, subpath = _parse_push_to(args.push_to)
-        prefix = f"{subpath}/{run_id}" if subpath else run_id
-        payload["pushed-to"] = f"hf://buckets/{bucket_id}/{prefix}/"
-    out.result("remote-run", **payload)
-
-    if exit_code != 0:
-        raise SystemExit(f"remote run failed with exit code {exit_code}")
-    return True
-
-
-# ---------------------------------------------------------------------------
-# Subcommand
-# ---------------------------------------------------------------------------
-
-
-class RunCommand(BaseDiffusersCLICommand):
-    task = "run"
-
-    @staticmethod
-    def register_subcommand(subparsers: _SubParsersAction) -> None:
-        from argparse import RawDescriptionHelpFormatter
-
-        epilog = (
-            "Examples\n"
-            "  $ diffusers-cli run -m black-forest-labs/FLUX.1-dev --dtype bf16 \\\n"
-            '      --pipeline-kwargs \'{"prompt": "a cat on the moon"}\'\n'
-            "  $ diffusers-cli run -m black-forest-labs/FLUX.1-dev --dtype bf16 \\\n"
-            '      --pipeline-kwargs \'{"prompt": "make the fur grey", "image": "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/cat.png"}\'\n'
-            "  $ diffusers-cli run -m black-forest-labs/FLUX.1-dev --dtype bf16 \\\n"
-            '      --pipeline-kwargs \'{"prompt": "a tiny cat"}\' \\\n'
-            '      --lora \'{"lora_id": "alvdansen/littletinies", "lora_scale": 0.8}\'\n'
-            "  $ diffusers-cli run -m black-forest-labs/FLUX.1-dev --dtype bf16 \\\n"
-            '      --pipeline-kwargs \'{"prompt": "a cat"}\' --remote --flavor a100-large\n'
-            "  $ diffusers-cli run -m black-forest-labs/FLUX.1-dev --dtype bf16 --context-parallel \\\n"
-            '      --pipeline-kwargs \'{"prompt": "a cat"}\' --remote --flavor 4xa100-large\n'
-            "\n"
-            "Learn more\n"
-            "  Use `diffusers-cli <command> --help` for more information about a command.\n"
-            "  Read the documentation at https://huggingface.co/docs/diffusers\n"
-        )
-
-        parser: ArgumentParser = subparsers.add_parser(
-            "run",
-            help="Run any diffusers pipeline locally or remotely in an HF Sandbox.",
-            usage="\n  diffusers-cli run [options]",
-            epilog=epilog,
-            formatter_class=RawDescriptionHelpFormatter,
-        )
-        parser._optionals.title = "Options"
-        _add_loading_arguments(parser)
-        _add_optimization_arguments(parser)
-        parser.add_argument(
-            "--pipeline-kwargs",
-            default=None,
-            help=(
-                "JSON object of kwargs passed to the pipeline call. String values at known "
-                f"image-input keys ({', '.join(_IMAGE_INPUT_KEYS)}) are auto-loaded as PIL images; "
-                f"video-input keys ({', '.join(_VIDEO_INPUT_KEYS)}) are auto-loaded as frame lists; "
-                f"audio-input keys ({', '.join(_AUDIO_INPUT_KEYS)}) are auto-loaded via torchaudio."
-            ),
-        )
-        parser.add_argument(
-            "--output-key",
-            default=None,
-            help="For modular pipelines: name of the intermediate to extract (passed as `output=` to the call).",
-        )
-        parser.add_argument("--seed", type=int, default=None, help="Random seed for reproducibility.")
-        parser.add_argument(
-            "--fps",
-            type=int,
-            default=8,
-            help="FPS used when the output happens to be a frame sequence.",
-        )
-        parser.add_argument(
-            "--sampling-rate",
-            type=int,
-            default=None,
-            help="Sample rate used when the output happens to be an audio array.",
-        )
-        _add_remote_arguments(parser)
-        _add_output_arguments(parser)
-        parser.set_defaults(func=RunCommand)
-
-    def __init__(self, args: Namespace):
-        self.args = args
-
-    def run(self) -> None:
-        import diffusers
-
-        _get_or_create_run_id()  # populate RUN_ID_ENV so local output dir + remote bucket prefix agree
-
-        call_kwargs = _parse_pipeline_kwargs(self.args.pipeline_kwargs)
-
-        if _maybe_submit_remote(self.args, self.task):
-            return
-
-        # Resolve media before loading pipeline weights so dead URLs / missing files fail
-        # fast — cheap to fetch, expensive to load a 20GB model just to hit a 404.
-        _resolve_media_inputs(call_kwargs)
-        pipeline = _load_pipeline(self.args)
-        is_modular = isinstance(pipeline, diffusers.ModularPipeline)
-
-        if self.args.output_key is not None:
-            call_kwargs["output"] = self.args.output_key
-
-        device = pipeline.device.type if hasattr(pipeline, "device") else "cpu"
-        generator = _get_generator(self.args.seed, device)
-        if generator is not None:
-            call_kwargs["generator"] = generator
-
-        try:
-            result = pipeline(**call_kwargs)
-
-            # Under torchrun, ranks > 0 produce identical output to rank 0 (CP shards the
-            # transformer compute but ranks reduce to the same final tensors). Save/push/print
-            # from rank 0 only to avoid clobbering bucket files 4x and printing 4x.
-            if os.environ.get("RANK", "0") == "0":
-                savable = result if is_modular else _unwrap_pipeline_output(result)
-                saved = _save_output(savable, self.args, self.task)
-                pushed = _push_outputs(self.args, saved, self.task)
-
-                out.result(
-                    self.task,
-                    model=self.args.model,
-                    device=device,
-                    pipeline_class=type(pipeline).__name__,
-                    modular=is_modular,
-                    outputs=saved,
-                    pushed=pushed,
-                    seed=self.args.seed,
-                    output_key=self.args.output_key,
-                )
-        finally:
-            import torch
-
-            if torch.distributed.is_available() and torch.distributed.is_initialized():
-                torch.distributed.destroy_process_group()

diffusers/commands/schema.py

@@ -1,287 +0,0 @@
-# Copyright 2026 The HuggingFace Team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-"""`diffusers-cli schema` — print the input schema for any pipeline repo.
-
-Tries `DiffusionPipeline.config_name` first (so standard repos get their `__call__` signature introspected); falls back
-to `ModularPipelineBlocks.from_pretrained` for modular repos. No weights are downloaded — only the small index file
-(and any custom block code if `--trust-remote-code` is set).
-"""
-
-from __future__ import annotations
-
-import inspect
-import re
-from argparse import ArgumentParser, Namespace, _SubParsersAction
-from typing import Any
-
-from huggingface_hub.cli._output import OutputFormat, out
-
-from ..utils import logging
-from . import BaseDiffusersCLICommand
-
-
-logger = logging.get_logger("diffusers-cli/schema")
-
-
-def _schema(args: Namespace) -> None:
-    """Print the pipeline's input schema.
-
-    Tries `DiffusionPipeline.config_name` (= `model_index.json`) first; if present, introspects the declared pipeline
-    class's `__call__` signature. Otherwise falls back to `ModularPipelineBlocks.from_pretrained` and reads the
-    block-declared `inputs`. No weights downloaded either way.
-    """
-    import diffusers
-
-    try:
-        index = diffusers.DiffusionPipeline.load_config(args.model, token=args.token, revision=args.revision)
-    except OSError:
-        index = None
-
-    if index is not None:
-        class_name = index.get("_class_name")
-        if class_name is None:
-            raise SystemExit(
-                f"{diffusers.DiffusionPipeline.config_name} for {args.model!r} has no `_class_name` field."
-            )
-        pipeline_cls = getattr(diffusers, class_name, None)
-        if pipeline_cls is None:
-            raise SystemExit(
-                f"Pipeline class {class_name!r} declared in {diffusers.DiffusionPipeline.config_name} "
-                "is not exported by the installed diffusers."
-            )
-
-        sig = inspect.signature(pipeline_cls.__call__)
-        descriptions = _parse_docstring_args(pipeline_cls.__call__.__doc__) if args.verbose else {}
-        schema: list[dict[str, Any]] = []
-        for name, param in sig.parameters.items():
-            if name == "self":
-                continue
-            if param.kind in (inspect.Parameter.VAR_POSITIONAL, inspect.Parameter.VAR_KEYWORD):
-                continue
-            has_default = param.default is not inspect.Parameter.empty
-            schema.append(
-                {
-                    "name": name,
-                    "type_hint": str(param.annotation) if param.annotation is not inspect.Parameter.empty else None,
-                    "default": param.default if has_default else None,
-                    "required": not has_default,
-                    "description": descriptions.get(name, ""),
-                }
-            )
-    else:
-        kwargs: dict[str, Any] = {"trust_remote_code": args.trust_remote_code}
-        if args.revision:
-            kwargs["revision"] = args.revision
-        if args.token:
-            kwargs["token"] = args.token
-
-        # If the repo declares custom code + external dependencies, surface them upfront so
-        # the user knows what to install before we hit an ImportError inside from_pretrained.
-        _warn_custom_block_requirements(args)
-
-        try:
-            blocks = diffusers.ModularPipelineBlocks.from_pretrained(args.model, **kwargs)
-        except Exception as e:
-            hint = "\nPass --trust-remote-code if it ships custom block code." if not args.trust_remote_code else ""
-            raise SystemExit(
-                f"Could not read schema for {args.model!r}: no {diffusers.DiffusionPipeline.config_name} and "
-                f"loading as a modular pipeline failed with:\n  {type(e).__name__}: {e}{hint}"
-            ) from e
-
-        class_name = type(blocks).__name__
-        schema = [
-            {
-                "name": p.name,
-                "type_hint": str(p.type_hint) if p.type_hint is not None else None,
-                "default": p.default,
-                "required": p.required,
-                "description": p.description,
-            }
-            for p in blocks.inputs
-        ]
-
-    if out.mode == OutputFormat.json:
-        out.dict({"task": "schema", "model": args.model, "pipeline_class": class_name, "inputs": schema})
-    elif out.mode == OutputFormat.agent:
-        out.table(schema, headers=["name", "required", "type_hint", "default", "description"])
-    else:
-        out.text(f"{class_name} ({args.model}) inputs:")
-        for entry in schema:
-            tag = "required" if entry["required"] else f"optional, default={entry['default']!r}"
-            out.text(f"  {entry['name']}  ({tag})")
-            if entry["type_hint"]:
-                out.text(f"    type: {entry['type_hint']}")
-            if entry["description"]:
-                out.text(f"    desc: {entry['description']}")
-
-
-def _warn_custom_block_requirements(args: Namespace) -> None:
-    """Warn upfront when a modular block ships custom code with declared external dependencies.
-
-    Reads `modular_config.json` if present; if it has an `auto_map` (custom code) and a non-empty `requirements`
-    list/dict, prints a heads-up. `from_pretrained` will otherwise fail with an `ImportError` deep in the loader stack
-    when a listed dep is missing.
-    """
-    import diffusers
-
-    try:
-        config = diffusers.ModularPipelineBlocks.load_config(args.model, token=args.token, revision=args.revision)
-    except Exception:
-        return  # no modular_config.json or unreachable — nothing to warn about
-    if not isinstance(config, dict):
-        return
-    if not config.get("auto_map"):
-        return
-    requirements = config.get("requirements")
-    if not requirements:
-        return
-
-    # `requirements` may be a dict {name: version} or (older repos) a list of [name, version] pairs.
-    if isinstance(requirements, dict):
-        pairs = list(requirements.items())
-    elif isinstance(requirements, list):
-        pairs = [(item[0], item[1]) for item in requirements if isinstance(item, (list, tuple)) and len(item) >= 2]
-    else:
-        pairs = []
-    if not pairs:
-        return
-
-    formatted = ", ".join(f"{name}=={version}" for name, version in pairs)
-    logger.warning(
-        f"{args.model!r} ships custom block code with external dependencies: {formatted}. "
-        "You will need to install these in order to determine the pipeline schema."
-    )
-
-
-def _parse_docstring_args(docstring: str | None) -> dict[str, str]:
-    """Extract per-argument descriptions from a Google-style `Args:` block.
-
-    Returns a `{name: description}` mapping. Best-effort — unrecognised formats just yield an empty dict rather than
-    raising.
-    """
-    if not docstring:
-        return {}
-
-    lines = docstring.expandtabs().splitlines()
-    start = None
-    section_indent = 0
-    for i, line in enumerate(lines):
-        if line.strip() in ("Args:", "Arguments:", "Parameters:"):
-            start = i + 1
-            section_indent = len(line) - len(line.lstrip())
-            break
-    if start is None:
-        return {}
-
-    descriptions: dict[str, str] = {}
-    current_name: str | None = None
-    current_lines: list[str] = []
-    arg_indent: int | None = None
-    name_pattern = re.compile(r"^(\w+)\s*(?:\([^)]*\))?\s*:?\s*(.*)$")
-
-    def _flush() -> None:
-        if current_name and current_lines:
-            descriptions[current_name] = " ".join(s.strip() for s in current_lines).strip()
-
-    for line in lines[start:]:
-        if not line.strip():
-            continue
-        indent = len(line) - len(line.lstrip())
-        # A new top-level section ends the Args block.
-        if indent <= section_indent and line.strip().endswith(":"):
-            break
-        if arg_indent is None:
-            arg_indent = indent
-        if indent == arg_indent:
-            _flush()
-            current_lines = []
-            match = name_pattern.match(line.strip())
-            if match:
-                current_name = match.group(1)
-                tail = match.group(2).strip()
-                if tail:
-                    current_lines.append(tail)
-            else:
-                current_name = None
-        elif current_name is not None and indent > arg_indent:
-            current_lines.append(line.strip())
-    _flush()
-    return descriptions
-
-
-class SchemaCommand(BaseDiffusersCLICommand):
-    task = "schema"
-
-    @staticmethod
-    def register_subcommand(subparsers: _SubParsersAction) -> None:
-        from argparse import RawDescriptionHelpFormatter
-
-        epilog = (
-            "Examples\n"
-            "  $ diffusers-cli schema -m stabilityai/stable-diffusion-xl-base-1.0\n"
-            "  $ diffusers-cli schema -m black-forest-labs/FLUX.1-dev --verbose\n"
-            "  $ diffusers-cli --format json schema -m stabilityai/stable-diffusion-xl-base-1.0\n"
-            "\n"
-            "Learn more\n"
-            "  Use `diffusers-cli <command> --help` for more information about a command.\n"
-            "  Read the documentation at https://huggingface.co/docs/diffusers\n"
-        )
-
-        parser: ArgumentParser = subparsers.add_parser(
-            "schema",
-            help="Print the input schema for a diffusers pipeline repo. No weights downloaded.",
-            usage="\n  diffusers-cli schema [options]",
-            epilog=epilog,
-            formatter_class=RawDescriptionHelpFormatter,
-        )
-        parser._optionals.title = "Options"
-        parser.add_argument(
-            "--model",
-            "-m",
-            required=True,
-            help="Model id on the Hugging Face Hub or local path.",
-        )
-        parser.add_argument(
-            "--revision",
-            default=None,
-            help="Model revision (branch, tag, or commit SHA).",
-        )
-        parser.add_argument(
-            "--token",
-            default=None,
-            help="Hugging Face token for gated/private models.",
-        )
-        parser.add_argument(
-            "--trust-remote-code",
-            action="store_true",
-            help="Allow custom code from the Hub (required for modular pipelines that ship block code).",
-        )
-        parser.add_argument(
-            "--verbose",
-            "-v",
-            action="store_true",
-            help=(
-                "Also include per-argument descriptions from the pipeline's __call__ docstring. "
-                "Modular pipelines always include block-declared descriptions; --verbose populates "
-                "the equivalent field for standard pipelines by parsing the Google-style Args: block."
-            ),
-        )
-        parser.set_defaults(func=SchemaCommand)
-
-    def __init__(self, args: Namespace):
-        self.args = args
-
-    def run(self) -> None:
-        _schema(self.args)

diffusers/commands/skills.py

@@ -1,344 +0,0 @@
-# Copyright 2026 The HuggingFace Team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""`diffusers-cli skills` — install Agent Skills bundles.
-
-Skill bundles live under `.ai/skills/<name>/` in the diffusers repo and follow the Agent Skills standard: a directory
-containing `SKILL.md` (plus optional resources). Installs to `.agents/skills/<name>/` which Claude, Codex, and Cursor
-all discover.
-"""
-
-from __future__ import annotations
-
-import os
-import shutil
-from argparse import ArgumentParser, Namespace, _SubParsersAction
-from pathlib import Path
-
-import httpx
-from huggingface_hub.cli._output import out
-
-from ..utils import logging
-from ..utils.constants import DIFFUSERS_REQUEST_TIMEOUT
-from . import BaseDiffusersCLICommand
-
-
-logger = logging.get_logger("diffusers-cli/skills")
-
-
-_REGISTRY_BASE = "https://api.github.com/repos/huggingface/diffusers/contents/.ai/skills"
-_REGISTRY_REF = "main"
-
-# Native skill-discovery paths per agent. Claude Code reads only `.claude/skills/`; Codex and
-# Cursor read `.agents/skills/` (Cursor also honors `.claude/skills/` via compat, but installing
-# to `.agents/skills/` is the portable choice for both).
-_CLAUDE_SKILLS_DIR = Path(".claude") / "skills"
-_AGENTS_SKILLS_DIR = Path(".agents") / "skills"
-
-# Env vars set by each agent when it launches the CLI. Values are the install path to use.
-_AGENT_ENV_TO_DIR: dict[str, Path] = {
-    "CLAUDECODE": _CLAUDE_SKILLS_DIR,
-    "CLAUDE_CODE": _CLAUDE_SKILLS_DIR,
-    "CODEX_SANDBOX": _AGENTS_SKILLS_DIR,
-    "CURSOR_AI": _AGENTS_SKILLS_DIR,
-}
-# When no agent env var is set, install to every native path so whichever agent the user
-# later switches to picks the skill up.
-_ALL_INSTALL_DIRS: tuple[Path, ...] = (_CLAUDE_SKILLS_DIR, _AGENTS_SKILLS_DIR)
-
-# Empty marker dropped inside each installed skill dir so `update` can distinguish our
-# installs from user-placed skills at the same paths.
-_MANAGED_MARKER_FILE = ".diffusers-skill-managed"
-
-
-# ---------------------------------------------------------------------------
-# Registry fetch
-# ---------------------------------------------------------------------------
-
-
-def _registry_url(name: str = "") -> str:
-    """API URL for the registry root, or for a single skill bundle when `name` is given."""
-    path = f"/{name}" if name else ""
-    return f"{_REGISTRY_BASE}{path}?ref={_REGISTRY_REF}"
-
-
-def _fetch_json(url: str) -> list[dict]:
-    try:
-        resp = httpx.get(url, timeout=DIFFUSERS_REQUEST_TIMEOUT)
-        resp.raise_for_status()
-        return resp.json()
-    except httpx.HTTPStatusError as e:
-        if e.response.status_code == 404:
-            raise SystemExit(f"Not found in registry: {url}") from e
-        raise SystemExit(f"Registry fetch failed: HTTP {e.response.status_code} {e.response.reason_phrase}") from e
-    except httpx.HTTPError as e:
-        raise SystemExit(f"Could not reach registry: {e}") from e
-
-
-def _walk_skill_files(name: str) -> list[tuple[str, str]]:
-    files: list[tuple[str, str]] = []
-
-    def _walk(api_url: str, prefix: str) -> None:
-        for entry in _fetch_json(api_url):
-            if entry["type"] == "file":
-                files.append((f"{prefix}{entry['name']}", entry["download_url"]))
-            elif entry["type"] == "dir":
-                _walk(entry["url"], f"{prefix}{entry['name']}/")
-
-    _walk(_registry_url(name), "")
-    return files
-
-
-def _download_skill_bundle(name: str) -> dict[str, bytes]:
-    files = _walk_skill_files(name)
-    if not files:
-        raise SystemExit(f"Skill '{name}' has no files in the registry.")
-    bundle: dict[str, bytes] = {}
-    for rel_path, url in files:
-        resp = httpx.get(url, timeout=DIFFUSERS_REQUEST_TIMEOUT)
-        resp.raise_for_status()
-        bundle[rel_path] = resp.content
-    return bundle
-
-
-# ---------------------------------------------------------------------------
-# Install / discovery
-# ---------------------------------------------------------------------------
-
-
-def _detect_install_dirs() -> tuple[Path, ...]:
-    """Pick where to install based on the launching agent.
-
-    If we detect a specific agent from its env var, install only there. If nothing is detected, install to every native
-    path so any agent picks the skill up later.
-    """
-    for env_var, skills_dir in _AGENT_ENV_TO_DIR.items():
-        if os.environ.get(env_var):
-            return (skills_dir,)
-    return _ALL_INSTALL_DIRS
-
-
-def _install_skill(name: str, bundle: dict[str, bytes], root: Path, skills_dir: Path, force: bool) -> Path:
-    skill_dir = root / skills_dir / name
-    if skill_dir.exists():
-        if not force:
-            raise SystemExit(f"Skill already installed at {skill_dir}. Use --force to reinstall.")
-        shutil.rmtree(skill_dir)
-    skill_dir.mkdir(parents=True, exist_ok=True)
-    for rel_path, data in bundle.items():
-        target = skill_dir / rel_path
-        target.parent.mkdir(parents=True, exist_ok=True)
-        target.write_bytes(data)
-    (skill_dir / _MANAGED_MARKER_FILE).touch()
-    return skill_dir
-
-
-def _has_local_changes(skill_dir: Path, bundle: dict[str, bytes]) -> bool:
-    """True if the installed skill has any file that differs from `bundle` or has extra files.
-
-    The marker file is ignored. Compares raw bytes so a whitespace-only edit still counts as dirty.
-    """
-    on_disk: dict[str, bytes] = {}
-    for path in skill_dir.rglob("*"):
-        if not path.is_file():
-            continue
-        rel = str(path.relative_to(skill_dir))
-        if rel == _MANAGED_MARKER_FILE:
-            continue
-        on_disk[rel] = path.read_bytes()
-    return on_disk != bundle
-
-
-def _discover_installed(root: Path) -> list[tuple[Path, str]]:
-    """Return `(skills_dir, name)` pairs for every managed install under `root`."""
-    found: list[tuple[Path, str]] = []
-    for skills_dir in _ALL_INSTALL_DIRS:
-        skills_root = root / skills_dir
-        if not skills_root.exists():
-            continue
-        for d in sorted(skills_root.iterdir()):
-            if d.is_dir() and (d / _MANAGED_MARKER_FILE).exists():
-                found.append((skills_dir, d.name))
-    return found
-
-
-class SkillsCommand(BaseDiffusersCLICommand):
-    @staticmethod
-    def register_subcommand(subparsers: _SubParsersAction) -> None:
-        parser: ArgumentParser = subparsers.add_parser(
-            "skills",
-            help="Manage Agent Skills for AI assistants.",
-            usage="\n  diffusers-cli skills <add|list|update|preview> [options]",
-        )
-        parser._optionals.title = "Options"
-        actions = parser.add_subparsers(dest="skills_action", required=True, metavar="<action>")
-
-        add = actions.add_parser("add", help="Download and install a skill.")
-        add.add_argument(
-            "name",
-            nargs="?",
-            default=None,
-            help="Skill name (e.g. diffusers-cli, custom-blocks). Omit and pass --all to install every skill.",
-        )
-        add.add_argument(
-            "--all",
-            dest="install_all",
-            action="store_true",
-            help="Install every skill in the registry. Mutually exclusive with a positional name.",
-        )
-        add.add_argument(
-            "--global",
-            "-g",
-            dest="install_global",
-            action="store_true",
-            help="Install globally (user-level) instead of in the current project directory.",
-        )
-        add.add_argument("--force", action="store_true", help="Overwrite existing skills in the destination.")
-        add.set_defaults(func=SkillsCommand)
-
-        list_action = actions.add_parser("list", help="List available skills in the registry.")
-        list_action.set_defaults(func=SkillsCommand)
-
-        update = actions.add_parser("update", help="Re-download and reinstall managed skills.")
-        update.add_argument(
-            "name",
-            nargs="?",
-            default=None,
-            help="Optional installed skill name to update. Omit to update every managed skill.",
-        )
-        update.add_argument(
-            "--global",
-            "-g",
-            dest="install_global",
-            action="store_true",
-            help="Update skills installed globally (user-level) instead of the current project.",
-        )
-        update.add_argument(
-            "--force",
-            action="store_true",
-            help="Overwrite skills even if they have local modifications since install.",
-        )
-        update.set_defaults(func=SkillsCommand)
-
-        preview = actions.add_parser("preview", help="Print a skill's SKILL.md from the registry.")
-        preview.add_argument("name", help="Skill name to preview.")
-        preview.set_defaults(func=SkillsCommand)
-
-    def __init__(self, args: Namespace):
-        self.args = args
-
-    def run(self) -> None:
-        if self.args.skills_action == "add":
-            self._add()
-        elif self.args.skills_action == "list":
-            self._list()
-        elif self.args.skills_action == "update":
-            self._update()
-        elif self.args.skills_action == "preview":
-            self._preview()
-
-    def _add(self) -> None:
-        if self.args.install_all and self.args.name:
-            raise SystemExit("--all and a positional skill name are mutually exclusive.")
-        if not self.args.install_all and not self.args.name:
-            raise SystemExit("Pass a skill name (e.g. diffusers-cli) or --all to install every skill.")
-
-        root = Path.home() if self.args.install_global else Path.cwd()
-        install_dirs = _detect_install_dirs()
-        names = self._resolve_names()
-
-        installed: list[str] = []
-        failed: list[str] = []
-        for name in names:
-            try:
-                bundle = _download_skill_bundle(name)
-                for skills_dir in install_dirs:
-                    _install_skill(name, bundle, root, skills_dir, self.args.force)
-                installed.append(name)
-            except (SystemExit, httpx.HTTPError) as e:
-                # Downgrade to a warning so one broken skill doesn't abort the batch.
-                logger.warning(f"Skipping skill {name!r}: {e}")
-                failed.append(name)
-
-        if not installed:
-            raise SystemExit(f"No skills installed. Failed: {failed}")
-        out.result(
-            f"Installed {len(installed)} skill(s)",
-            installed=", ".join(installed),
-            failed=", ".join(failed) if failed else None,
-            paths=", ".join(str(root / d) for d in install_dirs),
-        )
-
-    def _update(self) -> None:
-        root = Path.home() if self.args.install_global else Path.cwd()
-        installed = _discover_installed(root)
-        if self.args.name is not None:
-            installed = [entry for entry in installed if entry[1] == self.args.name]
-            if not installed:
-                raise SystemExit(f"No installed skill named {self.args.name!r} found under {root}.")
-        if not installed:
-            raise SystemExit(f"No managed skills found under {root}.")
-
-        # Group by skill name so we redownload each bundle once even if it's installed to
-        # multiple locations (e.g. both .claude/skills/ and .agents/skills/).
-        by_name: dict[str, list[Path]] = {}
-        for skills_dir, name in installed:
-            by_name.setdefault(name, []).append(skills_dir)
-
-        updated: list[str] = []
-        failed: list[str] = []
-        skipped: list[str] = []
-        for name, dirs in sorted(by_name.items()):
-            try:
-                bundle = _download_skill_bundle(name)
-                for skills_dir in dirs:
-                    skill_dir = root / skills_dir / name
-                    if not self.args.force and _has_local_changes(skill_dir, bundle):
-                        logger.warning(
-                            f"Skill {name!r} at {skill_dir} has local modifications; "
-                            "skipping. Pass --force to overwrite them."
-                        )
-                        skipped.append(name)
-                        continue
-                    _install_skill(name, bundle, root, skills_dir, force=True)
-                    updated.append(name)
-            except (SystemExit, httpx.HTTPError) as e:
-                logger.warning(f"Skipping skill {name!r}: {e}")
-                failed.append(name)
-
-        out.result(
-            f"Updated {len(updated)} skill(s)",
-            updated=", ".join(updated),
-            skipped=", ".join(skipped) if skipped else None,
-            failed=", ".join(failed) if failed else None,
-        )
-
-    def _preview(self) -> None:
-        bundle = _download_skill_bundle(self.args.name)
-        skill_md = bundle.get("SKILL.md")
-        if skill_md is None:
-            raise SystemExit(f"Skill {self.args.name!r} has no SKILL.md in the registry.")
-        print(skill_md.decode())
-
-    def _list(self) -> None:
-        entries = _fetch_json(_registry_url())
-        skills = [{"name": e["name"]} for e in entries if e["type"] == "dir" and not e["name"].startswith(".")]
-        if not skills:
-            raise SystemExit("No skills found in registry.")
-        out.table(skills, headers=["name"])
-
-    def _resolve_names(self) -> list[str]:
-        if self.args.install_all:
-            entries = _fetch_json(_registry_url())
-            return sorted(e["name"] for e in entries if e["type"] == "dir" and not e["name"].startswith("."))
-        return [self.args.name]

diffusers/configuration_utils.py

@@ -1,752 +0,0 @@
-# coding=utf-8
-# Copyright 2025 The HuggingFace Inc. team.
-# Copyright (c) 2022, NVIDIA CORPORATION.  All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""ConfigMixin base class and utilities."""
-
-import functools
-import importlib
-import inspect
-import json
-import os
-import re
-from collections import OrderedDict
-from pathlib import Path
-from typing import Any
-
-import numpy as np
-from huggingface_hub import DDUFEntry, create_repo, hf_hub_download
-from huggingface_hub.utils import (
-    EntryNotFoundError,
-    HfHubHTTPError,
-    RepositoryNotFoundError,
-    RevisionNotFoundError,
-    validate_hf_hub_args,
-)
-from typing_extensions import Self
-
-from . import __version__
-from .utils import (
-    HUGGINGFACE_CO_RESOLVE_ENDPOINT,
-    DummyObject,
-    deprecate,
-    extract_commit_hash,
-    http_user_agent,
-    logging,
-)
-
-
-logger = logging.get_logger(__name__)
-
-_re_configuration_file = re.compile(r"config\.(.*)\.json")
-
-
-class FrozenDict(OrderedDict):
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-
-        for key, value in self.items():
-            setattr(self, key, value)
-
-        self.__frozen = True
-
-    def __delitem__(self, *args, **kwargs):
-        raise Exception(f"You cannot use ``__delitem__`` on a {self.__class__.__name__} instance.")
-
-    def setdefault(self, *args, **kwargs):
-        raise Exception(f"You cannot use ``setdefault`` on a {self.__class__.__name__} instance.")
-
-    def pop(self, *args, **kwargs):
-        raise Exception(f"You cannot use ``pop`` on a {self.__class__.__name__} instance.")
-
-    def update(self, *args, **kwargs):
-        raise Exception(f"You cannot use ``update`` on a {self.__class__.__name__} instance.")
-
-    def __setattr__(self, name, value):
-        if hasattr(self, "__frozen") and self.__frozen:
-            raise Exception(f"You cannot use ``__setattr__`` on a {self.__class__.__name__} instance.")
-        super().__setattr__(name, value)
-
-    def __setitem__(self, name, value):
-        if hasattr(self, "__frozen") and self.__frozen:
-            raise Exception(f"You cannot use ``__setattr__`` on a {self.__class__.__name__} instance.")
-        super().__setitem__(name, value)
-
-
-class ConfigMixin:
-    r"""
-    Base class for all configuration classes. All configuration parameters are stored under `self.config`. Also
-    provides the [`~ConfigMixin.from_config`] and [`~ConfigMixin.save_config`] methods for loading, downloading, and
-    saving classes that inherit from [`ConfigMixin`].
-
-    Class attributes:
-        - **config_name** (`str`) -- A filename under which the config should stored when calling
-          [`~ConfigMixin.save_config`] (should be overridden by parent class).
-        - **ignore_for_config** (`list[str]`) -- A list of attributes that should not be saved in the config (should be
-          overridden by subclass).
-        - **has_compatibles** (`bool`) -- Whether the class has compatible classes (should be overridden by subclass).
-        - **_deprecated_kwargs** (`list[str]`) -- Keyword arguments that are deprecated. Note that the `init` function
-          should only have a `kwargs` argument if at least one argument is deprecated (should be overridden by
-          subclass).
-    """
-
-    config_name = None
-    ignore_for_config = []
-    has_compatibles = False
-
-    _deprecated_kwargs = []
-    _auto_class = None
-
-    @classmethod
-    def register_for_auto_class(cls, auto_class="AutoModel"):
-        """
-        Register this class with the given auto class so that it can be loaded with `AutoModel.from_pretrained(...,
-        trust_remote_code=True)`.
-
-        When the config is saved, the resulting `config.json` will include an `auto_map` entry mapping the auto class
-        to this class's module and class name.
-
-        Args:
-            auto_class (`str` or type, *optional*, defaults to `"AutoModel"`):
-                The auto class to register this class with. Can be a string (e.g. `"AutoModel"`) or the class itself.
-                Currently only `"AutoModel"` is supported.
-
-        Example:
-
-        ```python
-        from diffusers import ModelMixin, ConfigMixin
-
-
-        class MyCustomModel(ModelMixin, ConfigMixin): ...
-
-
-        MyCustomModel.register_for_auto_class("AutoModel")
-        ```
-        """
-        if auto_class != "AutoModel":
-            raise ValueError(f"Only 'AutoModel' is supported, got '{auto_class}'.")
-
-        cls._auto_class = auto_class
-
-    def register_to_config(self, **kwargs):
-        if self.config_name is None:
-            raise NotImplementedError(f"Make sure that {self.__class__} has defined a class name `config_name`")
-        # Special case for `kwargs` used in deprecation warning added to schedulers
-        # TODO: remove this when we remove the deprecation warning, and the `kwargs` argument,
-        # or solve in a more general way.
-        kwargs.pop("kwargs", None)
-
-        if not hasattr(self, "_internal_dict"):
-            internal_dict = kwargs
-        else:
-            previous_dict = dict(self._internal_dict)
-            internal_dict = {**self._internal_dict, **kwargs}
-            logger.debug(f"Updating config from {previous_dict} to {internal_dict}")
-
-        self._internal_dict = FrozenDict(internal_dict)
-
-    def __getattr__(self, name: str) -> Any:
-        """The only reason we overwrite `getattr` here is to gracefully deprecate accessing
-        config attributes directly. See https://github.com/huggingface/diffusers/pull/3129
-
-        This function is mostly copied from PyTorch's __getattr__ overwrite:
-        https://pytorch.org/docs/stable/_modules/torch/nn/modules/module.html#Module
-        """
-
-        is_in_config = "_internal_dict" in self.__dict__ and hasattr(self.__dict__["_internal_dict"], name)
-        is_attribute = name in self.__dict__
-
-        if is_in_config and not is_attribute:
-            deprecation_message = f"Accessing config attribute `{name}` directly via '{type(self).__name__}' object attribute is deprecated. Please access '{name}' over '{type(self).__name__}'s config object instead, e.g. 'scheduler.config.{name}'."
-            deprecate("direct config name access", "1.0.0", deprecation_message, standard_warn=False)
-            return self._internal_dict[name]
-
-        raise AttributeError(f"'{type(self).__name__}' object has no attribute '{name}'")
-
-    def save_config(self, save_directory: str | os.PathLike, push_to_hub: bool = False, **kwargs):
-        """
-        Save a configuration object to the directory specified in `save_directory` so that it can be reloaded using the
-        [`~ConfigMixin.from_config`] class method.
-
-        Args:
-            save_directory (`str` or `os.PathLike`):
-                Directory where the configuration JSON file is saved (will be created if it does not exist).
-            push_to_hub (`bool`, *optional*, defaults to `False`):
-                Whether or not to push your model to the Hugging Face Hub after saving it. You can specify the
-                repository you want to push to with `repo_id` (will default to the name of `save_directory` in your
-                namespace).
-            kwargs (`dict[str, Any]`, *optional*):
-                Additional keyword arguments passed along to the [`~utils.PushToHubMixin.push_to_hub`] method.
-        """
-        if os.path.isfile(save_directory):
-            raise AssertionError(f"Provided path ({save_directory}) should be a directory, not a file")
-
-        os.makedirs(save_directory, exist_ok=True)
-
-        # If we save using the predefined names, we can load using `from_config`
-        output_config_file = os.path.join(save_directory, self.config_name)
-
-        self.to_json_file(output_config_file)
-        logger.info(f"Configuration saved in {output_config_file}")
-
-        if push_to_hub:
-            commit_message = kwargs.pop("commit_message", None)
-            private = kwargs.pop("private", None)
-            create_pr = kwargs.pop("create_pr", False)
-            token = kwargs.pop("token", None)
-            repo_id = kwargs.pop("repo_id", save_directory.split(os.path.sep)[-1])
-            repo_id = create_repo(repo_id, exist_ok=True, private=private, token=token).repo_id
-            subfolder = kwargs.pop("subfolder", None)
-
-            self._upload_folder(
-                save_directory,
-                repo_id,
-                token=token,
-                commit_message=commit_message,
-                create_pr=create_pr,
-                subfolder=subfolder,
-            )
-
-    @classmethod
-    def from_config(
-        cls, config: FrozenDict | dict[str, Any] = None, return_unused_kwargs=False, **kwargs
-    ) -> Self | tuple[Self, dict[str, Any]]:
-        r"""
-        Instantiate a Python class from a config dictionary.
-
-        Parameters:
-            config (`dict[str, Any]`):
-                A config dictionary from which the Python class is instantiated. Make sure to only load configuration
-                files of compatible classes.
-            return_unused_kwargs (`bool`, *optional*, defaults to `False`):
-                Whether kwargs that are not consumed by the Python class should be returned or not.
-            kwargs (remaining dictionary of keyword arguments, *optional*):
-                Can be used to update the configuration object (after it is loaded) and initiate the Python class.
-                `**kwargs` are passed directly to the underlying scheduler/model's `__init__` method and eventually
-                overwrite the same named arguments in `config`.
-
-        Returns:
-            [`ModelMixin`] or [`SchedulerMixin`]:
-                A model or scheduler object instantiated from a config dictionary.
-
-        Examples:
-
-        ```python
-        >>> from diffusers import DDPMScheduler, DDIMScheduler, PNDMScheduler
-
-        >>> # Download scheduler from huggingface.co and cache.
-        >>> scheduler = DDPMScheduler.from_pretrained("google/ddpm-cifar10-32")
-
-        >>> # Instantiate DDIM scheduler class with same config as DDPM
-        >>> scheduler = DDIMScheduler.from_config(scheduler.config)
-
-        >>> # Instantiate PNDM scheduler class with same config as DDPM
-        >>> scheduler = PNDMScheduler.from_config(scheduler.config)
-        ```
-        """
-        # <===== TO BE REMOVED WITH DEPRECATION
-        # TODO(Patrick) - make sure to remove the following lines when config=="model_path" is deprecated
-        if "pretrained_model_name_or_path" in kwargs:
-            config = kwargs.pop("pretrained_model_name_or_path")
-
-        if config is None:
-            raise ValueError("Please make sure to provide a config as the first positional argument.")
-        # ======>
-
-        if not isinstance(config, dict):
-            deprecation_message = "It is deprecated to pass a pretrained model name or path to `from_config`."
-            if "Scheduler" in cls.__name__:
-                deprecation_message += (
-                    f"If you were trying to load a scheduler, please use {cls}.from_pretrained(...) instead."
-                    " Otherwise, please make sure to pass a configuration dictionary instead. This functionality will"
-                    " be removed in v1.0.0."
-                )
-            elif "Model" in cls.__name__:
-                deprecation_message += (
-                    f"If you were trying to load a model, please use {cls}.load_config(...) followed by"
-                    f" {cls}.from_config(...) instead. Otherwise, please make sure to pass a configuration dictionary"
-                    " instead. This functionality will be removed in v1.0.0."
-                )
-            deprecate("config-passed-as-path", "1.0.0", deprecation_message, standard_warn=False)
-            config, kwargs = cls.load_config(pretrained_model_name_or_path=config, return_unused_kwargs=True, **kwargs)
-
-        init_dict, unused_kwargs, hidden_dict = cls.extract_init_dict(config, **kwargs)
-
-        # Allow dtype to be specified on initialization
-        if "dtype" in unused_kwargs:
-            init_dict["dtype"] = unused_kwargs.pop("dtype")
-
-        # add possible deprecated kwargs
-        for deprecated_kwarg in cls._deprecated_kwargs:
-            if deprecated_kwarg in unused_kwargs:
-                init_dict[deprecated_kwarg] = unused_kwargs.pop(deprecated_kwarg)
-
-        # Return model and optionally state and/or unused_kwargs
-        model = cls(**init_dict)
-
-        # make sure to also save config parameters that might be used for compatible classes
-        # update _class_name
-        if "_class_name" in hidden_dict:
-            hidden_dict["_class_name"] = cls.__name__
-
-        model.register_to_config(**hidden_dict)
-
-        # add hidden kwargs of compatible classes to unused_kwargs
-        unused_kwargs = {**unused_kwargs, **hidden_dict}
-
-        if return_unused_kwargs:
-            return (model, unused_kwargs)
-        else:
-            return model
-
-    @classmethod
-    def get_config_dict(cls, *args, **kwargs):
-        deprecation_message = (
-            f" The function get_config_dict is deprecated. Please use {cls}.load_config instead. This function will be"
-            " removed in version v1.0.0"
-        )
-        deprecate("get_config_dict", "1.0.0", deprecation_message, standard_warn=False)
-        return cls.load_config(*args, **kwargs)
-
-    @classmethod
-    @validate_hf_hub_args
-    def load_config(
-        cls,
-        pretrained_model_name_or_path: str | os.PathLike,
-        return_unused_kwargs=False,
-        return_commit_hash=False,
-        **kwargs,
-    ) -> tuple[dict[str, Any], dict[str, Any]]:
-        r"""
-        Load a model or scheduler configuration.
-
-        Parameters:
-            pretrained_model_name_or_path (`str` or `os.PathLike`, *optional*):
-                Can be either:
-
-                    - A string, the *model id* (for example `google/ddpm-celebahq-256`) of a pretrained model hosted on
-                      the Hub.
-                    - A path to a *directory* (for example `./my_model_directory`) containing model weights saved with
-                      [`~ConfigMixin.save_config`].
-
-            cache_dir (`str | os.PathLike`, *optional*):
-                Path to a directory where a downloaded pretrained model configuration is cached if the standard cache
-                is not used.
-            force_download (`bool`, *optional*, defaults to `False`):
-                Whether or not to force the (re-)download of the model weights and configuration files, overriding the
-                cached versions if they exist.
-            proxies (`dict[str, str]`, *optional*):
-                A dictionary of proxy servers to use by protocol or endpoint, for example, `{'http': 'foo.bar:3128',
-                'http://hostname': 'foo.bar:4012'}`. The proxies are used on each request.
-            output_loading_info(`bool`, *optional*, defaults to `False`):
-                Whether or not to also return a dictionary containing missing keys, unexpected keys and error messages.
-            local_files_only (`bool`, *optional*, defaults to `False`):
-                Whether to only load local model weights and configuration files or not. If set to `True`, the model
-                won't be downloaded from the Hub.
-            token (`str` or *bool*, *optional*):
-                The token to use as HTTP bearer authorization for remote files. If `True`, the token generated from
-                `diffusers-cli login` (stored in `~/.huggingface`) is used.
-            revision (`str`, *optional*, defaults to `"main"`):
-                The specific model version to use. It can be a branch name, a tag name, a commit id, or any identifier
-                allowed by Git.
-            subfolder (`str`, *optional*, defaults to `""`):
-                The subfolder location of a model file within a larger model repository on the Hub or locally.
-            return_unused_kwargs (`bool`, *optional*, defaults to `False):
-                Whether unused keyword arguments of the config are returned.
-            return_commit_hash (`bool`, *optional*, defaults to `False):
-                Whether the `commit_hash` of the loaded configuration are returned.
-
-        Returns:
-            `dict`:
-                A dictionary of all the parameters stored in a JSON configuration file.
-
-        """
-        cache_dir = kwargs.pop("cache_dir", None)
-        local_dir = kwargs.pop("local_dir", None)
-        local_dir_use_symlinks = kwargs.pop("local_dir_use_symlinks", "auto")
-        force_download = kwargs.pop("force_download", False)
-        proxies = kwargs.pop("proxies", None)
-        token = kwargs.pop("token", None)
-        local_files_only = kwargs.pop("local_files_only", False)
-        revision = kwargs.pop("revision", None)
-        _ = kwargs.pop("mirror", None)
-        subfolder = kwargs.pop("subfolder", None)
-        user_agent = kwargs.pop("user_agent", {})
-        dduf_entries: dict[str, DDUFEntry] | None = kwargs.pop("dduf_entries", None)
-
-        user_agent = {**user_agent, "file_type": "config"}
-        user_agent = http_user_agent(user_agent)
-
-        pretrained_model_name_or_path = str(pretrained_model_name_or_path)
-
-        if cls.config_name is None:
-            raise ValueError(
-                "`self.config_name` is not defined. Note that one should not load a config from "
-                "`ConfigMixin`. Please make sure to define `config_name` in a class inheriting from `ConfigMixin`"
-            )
-        # Custom path for now
-        if dduf_entries:
-            if subfolder is not None:
-                raise ValueError(
-                    "DDUF file only allow for 1 level of directory (e.g transformer/model1/model.safetentors is not allowed). "
-                    "Please check the DDUF structure"
-                )
-            config_file = cls._get_config_file_from_dduf(pretrained_model_name_or_path, dduf_entries)
-        elif os.path.isfile(pretrained_model_name_or_path):
-            config_file = pretrained_model_name_or_path
-        elif os.path.isdir(pretrained_model_name_or_path):
-            if subfolder is not None and os.path.isfile(
-                os.path.join(pretrained_model_name_or_path, subfolder, cls.config_name)
-            ):
-                config_file = os.path.join(pretrained_model_name_or_path, subfolder, cls.config_name)
-            elif os.path.isfile(os.path.join(pretrained_model_name_or_path, cls.config_name)):
-                # Load from a PyTorch checkpoint
-                config_file = os.path.join(pretrained_model_name_or_path, cls.config_name)
-            else:
-                raise EnvironmentError(
-                    f"Error no file named {cls.config_name} found in directory {pretrained_model_name_or_path}."
-                )
-        else:
-            try:
-                # Load from URL or cache if already cached
-                config_file = hf_hub_download(
-                    pretrained_model_name_or_path,
-                    filename=cls.config_name,
-                    cache_dir=cache_dir,
-                    force_download=force_download,
-                    proxies=proxies,
-                    local_files_only=local_files_only,
-                    token=token,
-                    user_agent=user_agent,
-                    subfolder=subfolder,
-                    revision=revision,
-                    local_dir=local_dir,
-                    local_dir_use_symlinks=local_dir_use_symlinks,
-                )
-            except RepositoryNotFoundError:
-                raise EnvironmentError(
-                    f"{pretrained_model_name_or_path} is not a local folder and is not a valid model identifier"
-                    " listed on 'https://huggingface.co/models'\nIf this is a private repository, make sure to pass a"
-                    " token having permission to this repo with `token` or log in with `hf auth login`."
-                )
-            except RevisionNotFoundError:
-                raise EnvironmentError(
-                    f"{revision} is not a valid git identifier (branch name, tag name or commit id) that exists for"
-                    " this model name. Check the model page at"
-                    f" 'https://huggingface.co/{pretrained_model_name_or_path}' for available revisions."
-                )
-            except EntryNotFoundError:
-                raise EnvironmentError(
-                    f"{pretrained_model_name_or_path} does not appear to have a file named {cls.config_name}."
-                )
-            except HfHubHTTPError as err:
-                raise EnvironmentError(
-                    "There was a specific connection error when trying to load"
-                    f" {pretrained_model_name_or_path}:\n{err}"
-                )
-            except ValueError:
-                raise EnvironmentError(
-                    f"We couldn't connect to '{HUGGINGFACE_CO_RESOLVE_ENDPOINT}' to load this model, couldn't find it"
-                    f" in the cached files and it looks like {pretrained_model_name_or_path} is not the path to a"
-                    f" directory containing a {cls.config_name} file.\nCheckout your internet connection or see how to"
-                    " run the library in offline mode at"
-                    " 'https://huggingface.co/docs/diffusers/installation#offline-mode'."
-                )
-            except EnvironmentError:
-                raise EnvironmentError(
-                    f"Can't load config for '{pretrained_model_name_or_path}'. If you were trying to load it from "
-                    "'https://huggingface.co/models', make sure you don't have a local directory with the same name. "
-                    f"Otherwise, make sure '{pretrained_model_name_or_path}' is the correct path to a directory "
-                    f"containing a {cls.config_name} file"
-                )
-        try:
-            config_dict = cls._dict_from_json_file(config_file, dduf_entries=dduf_entries)
-
-            commit_hash = extract_commit_hash(config_file)
-        except (json.JSONDecodeError, UnicodeDecodeError):
-            raise EnvironmentError(f"It looks like the config file at '{config_file}' is not a valid JSON file.")
-
-        if not (return_unused_kwargs or return_commit_hash):
-            return config_dict
-
-        outputs = (config_dict,)
-
-        if return_unused_kwargs:
-            outputs += (kwargs,)
-
-        if return_commit_hash:
-            outputs += (commit_hash,)
-
-        return outputs
-
-    @staticmethod
-    def _get_init_keys(input_class):
-        return set(dict(inspect.signature(input_class.__init__).parameters).keys())
-
-    @classmethod
-    def extract_init_dict(cls, config_dict, **kwargs):
-        # Skip keys that were not present in the original config, so default __init__ values were used
-        used_defaults = config_dict.get("_use_default_values", [])
-        config_dict = {k: v for k, v in config_dict.items() if k not in used_defaults and k != "_use_default_values"}
-
-        # 0. Copy origin config dict
-        original_dict = dict(config_dict.items())
-
-        # 1. Retrieve expected config attributes from __init__ signature
-        expected_keys = cls._get_init_keys(cls)
-        expected_keys.remove("self")
-        # remove general kwargs if present in dict
-        if "kwargs" in expected_keys:
-            expected_keys.remove("kwargs")
-
-        # 2. Remove attributes that cannot be expected from expected config attributes
-        # remove keys to be ignored
-        if len(cls.ignore_for_config) > 0:
-            expected_keys = expected_keys - set(cls.ignore_for_config)
-
-        # load diffusers library to import compatible and original scheduler
-        diffusers_library = importlib.import_module(__name__.split(".")[0])
-
-        if cls.has_compatibles:
-            compatible_classes = [c for c in cls._get_compatibles() if not isinstance(c, DummyObject)]
-        else:
-            compatible_classes = []
-
-        expected_keys_comp_cls = set()
-        for c in compatible_classes:
-            expected_keys_c = cls._get_init_keys(c)
-            expected_keys_comp_cls = expected_keys_comp_cls.union(expected_keys_c)
-        expected_keys_comp_cls = expected_keys_comp_cls - cls._get_init_keys(cls)
-        config_dict = {k: v for k, v in config_dict.items() if k not in expected_keys_comp_cls}
-
-        # remove attributes from orig class that cannot be expected
-        orig_cls_name = config_dict.pop("_class_name", cls.__name__)
-        if (
-            isinstance(orig_cls_name, str)
-            and orig_cls_name != cls.__name__
-            and hasattr(diffusers_library, orig_cls_name)
-        ):
-            orig_cls = getattr(diffusers_library, orig_cls_name)
-            unexpected_keys_from_orig = cls._get_init_keys(orig_cls) - expected_keys
-            config_dict = {k: v for k, v in config_dict.items() if k not in unexpected_keys_from_orig}
-        elif not isinstance(orig_cls_name, str) and not isinstance(orig_cls_name, (list, tuple)):
-            raise ValueError(
-                "Make sure that the `_class_name` is of type string or list of string (for custom pipelines)."
-            )
-
-        # remove private attributes
-        config_dict = {k: v for k, v in config_dict.items() if not k.startswith("_")}
-
-        # remove quantization_config
-        config_dict = {k: v for k, v in config_dict.items() if k != "quantization_config"}
-
-        # 3. Create keyword arguments that will be passed to __init__ from expected keyword arguments
-        init_dict = {}
-        for key in expected_keys:
-            # if config param is passed to kwarg and is present in config dict
-            # it should overwrite existing config dict key
-            if key in kwargs and key in config_dict:
-                config_dict[key] = kwargs.pop(key)
-
-            if key in kwargs:
-                # overwrite key
-                init_dict[key] = kwargs.pop(key)
-            elif key in config_dict:
-                # use value from config dict
-                init_dict[key] = config_dict.pop(key)
-
-        # 4. Give nice warning if unexpected values have been passed
-        if len(config_dict) > 0:
-            logger.warning(
-                f"The config attributes {config_dict} were passed to {cls.__name__}, "
-                "but are not expected and will be ignored. Please verify your "
-                f"{cls.config_name} configuration file."
-            )
-
-        # 5. Give nice info if config attributes are initialized to default because they have not been passed
-        passed_keys = set(init_dict.keys())
-        if len(expected_keys - passed_keys) > 0:
-            logger.info(
-                f"{expected_keys - passed_keys} was not found in config. Values will be initialized to default values."
-            )
-
-        # 6. Define unused keyword arguments
-        unused_kwargs = {**config_dict, **kwargs}
-
-        # 7. Define "hidden" config parameters that were saved for compatible classes
-        hidden_config_dict = {k: v for k, v in original_dict.items() if k not in init_dict}
-
-        return init_dict, unused_kwargs, hidden_config_dict
-
-    @classmethod
-    def _dict_from_json_file(cls, json_file: str | os.PathLike, dduf_entries: dict[str, DDUFEntry] | None = None):
-        if dduf_entries:
-            text = dduf_entries[json_file].read_text()
-        else:
-            with open(json_file, "r", encoding="utf-8") as reader:
-                text = reader.read()
-        return json.loads(text)
-
-    def __repr__(self):
-        return f"{self.__class__.__name__} {self.to_json_string()}"
-
-    @property
-    def config(self) -> dict[str, Any]:
-        """
-        Returns the config of the class as a frozen dictionary
-
-        Returns:
-            `dict[str, Any]`: Config of the class.
-        """
-        return self._internal_dict
-
-    def to_json_string(self) -> str:
-        """
-        Serializes the configuration instance to a JSON string.
-
-        Returns:
-            `str`:
-                String containing all the attributes that make up the configuration instance in JSON format.
-        """
-        config_dict = self._internal_dict if hasattr(self, "_internal_dict") else {}
-        config_dict["_class_name"] = self.__class__.__name__
-        config_dict["_diffusers_version"] = __version__
-
-        def to_json_saveable(value):
-            if isinstance(value, np.ndarray):
-                value = value.tolist()
-            elif isinstance(value, Path):
-                value = value.as_posix()
-            elif hasattr(value, "to_dict") and callable(value.to_dict):
-                value = value.to_dict()
-            elif isinstance(value, list):
-                value = [to_json_saveable(v) for v in value]
-            return value
-
-        if "quantization_config" in config_dict:
-            config_dict["quantization_config"] = (
-                config_dict.quantization_config.to_dict()
-                if not isinstance(config_dict.quantization_config, dict)
-                else config_dict.quantization_config
-            )
-
-        config_dict = {k: to_json_saveable(v) for k, v in config_dict.items()}
-        # Don't save "_ignore_files" or "_use_default_values"
-        config_dict.pop("_ignore_files", None)
-        config_dict.pop("_use_default_values", None)
-        # pop the `_pre_quantization_dtype` as torch.dtypes are not serializable.
-        _ = config_dict.pop("_pre_quantization_dtype", None)
-
-        if getattr(self, "_auto_class", None) is not None:
-            module = self.__class__.__module__.split(".")[-1]
-            auto_map = config_dict.get("auto_map", {})
-            auto_map[self._auto_class] = f"{module}.{self.__class__.__name__}"
-            config_dict["auto_map"] = auto_map
-
-        return json.dumps(config_dict, indent=2, sort_keys=True) + "\n"
-
-    def to_json_file(self, json_file_path: str | os.PathLike):
-        """
-        Save the configuration instance's parameters to a JSON file.
-
-        Args:
-            json_file_path (`str` or `os.PathLike`):
-                Path to the JSON file to save a configuration instance's parameters.
-        """
-        with open(json_file_path, "w", encoding="utf-8") as writer:
-            writer.write(self.to_json_string())
-
-    @classmethod
-    def _get_config_file_from_dduf(cls, pretrained_model_name_or_path: str, dduf_entries: dict[str, DDUFEntry]):
-        # paths inside a DDUF file must always be "/"
-        config_file = (
-            cls.config_name
-            if pretrained_model_name_or_path == ""
-            else "/".join([pretrained_model_name_or_path, cls.config_name])
-        )
-        if config_file not in dduf_entries:
-            raise ValueError(
-                f"We did not manage to find the file {config_file} in the dduf file. We only have the following files {dduf_entries.keys()}"
-            )
-        return config_file
-
-
-def register_to_config(init):
-    r"""
-    Decorator to apply on the init of classes inheriting from [`ConfigMixin`] so that all the arguments are
-    automatically sent to `self.register_for_config`. To ignore a specific argument accepted by the init but that
-    shouldn't be registered in the config, use the `ignore_for_config` class variable
-
-    Warning: Once decorated, all private arguments (beginning with an underscore) are trashed and not sent to the init!
-    """
-
-    @functools.wraps(init)
-    def inner_init(self, *args, **kwargs):
-        # Ignore private kwargs in the init.
-        init_kwargs = {k: v for k, v in kwargs.items() if not k.startswith("_")}
-        config_init_kwargs = {k: v for k, v in kwargs.items() if k.startswith("_")}
-        if not isinstance(self, ConfigMixin):
-            raise RuntimeError(
-                f"`@register_for_config` was applied to {self.__class__.__name__} init method, but this class does "
-                "not inherit from `ConfigMixin`."
-            )
-
-        ignore = getattr(self, "ignore_for_config", [])
-        # Get positional arguments aligned with kwargs
-        new_kwargs = {}
-        signature = inspect.signature(init)
-        parameters = {
-            name: p.default for i, (name, p) in enumerate(signature.parameters.items()) if i > 0 and name not in ignore
-        }
-        for arg, name in zip(args, parameters.keys()):
-            new_kwargs[name] = arg
-
-        # Then add all kwargs
-        new_kwargs.update(
-            {
-                k: init_kwargs.get(k, default)
-                for k, default in parameters.items()
-                if k not in ignore and k not in new_kwargs
-            }
-        )
-
-        # Take note of the parameters that were not present in the loaded config
-        if len(set(new_kwargs.keys()) - set(init_kwargs)) > 0:
-            new_kwargs["_use_default_values"] = list(set(new_kwargs.keys()) - set(init_kwargs))
-
-        new_kwargs = {**config_init_kwargs, **new_kwargs}
-        getattr(self, "register_to_config")(**new_kwargs)
-        init(self, *args, **init_kwargs)
-
-    return inner_init
-
-
-class LegacyConfigMixin(ConfigMixin):
-    r"""
-    A subclass of `ConfigMixin` to resolve class mapping from legacy classes (like `Transformer2DModel`) to more
-    pipeline-specific classes (like `DiTTransformer2DModel`).
-    """
-
-    @classmethod
-    def from_config(cls, config: FrozenDict | dict[str, Any] = None, return_unused_kwargs=False, **kwargs):
-        # To prevent dependency import problem.
-        from .models.model_loading_utils import _fetch_remapped_cls_from_config
-
-        # resolve remapping
-        remapped_class = _fetch_remapped_cls_from_config(config, cls)
-
-        if remapped_class is cls:
-            return super(LegacyConfigMixin, remapped_class).from_config(config, return_unused_kwargs, **kwargs)
-        else:
-            return remapped_class.from_config(config, return_unused_kwargs, **kwargs)

diffusers/dependency_versions_check.py

@@ -1,34 +0,0 @@
-# Copyright 2026 The HuggingFace Team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-from .dependency_versions_table import deps
-from .utils.versions import require_version, require_version_core
-
-
-# define which module versions we always want to check at run time
-# (usually the ones defined in `install_requires` in setup.py)
-#
-# order specific notes:
-# - tqdm must be checked before tokenizers
-
-pkgs_to_check_at_runtime = "python requests filelock numpy".split()
-for pkg in pkgs_to_check_at_runtime:
-    if pkg in deps:
-        require_version_core(deps[pkg])
-    else:
-        raise ValueError(f"can't find {pkg} in {deps.keys()}, check dependency_versions_table.py")
-
-
-def dep_version_check(pkg, hint=None):
-    require_version(deps[pkg], hint)

diffusers/dependency_versions_table.py

@@ -1,57 +0,0 @@
-# THIS FILE HAS BEEN AUTOGENERATED. To update:
-# 1. modify the `_deps` dict in setup.py
-# 2. run `make deps_table_update`
-deps = {
-    "Pillow": "Pillow",
-    "accelerate": "accelerate>=0.31.0",
-    "datasets": "datasets",
-    "filelock": "filelock",
-    "ftfy": "ftfy",
-    "hf-doc-builder": "hf-doc-builder>=0.3.0",
-    "httpx": "httpx<1.0.0",
-    "huggingface-hub": "huggingface-hub>=1.23.0,<2.0",
-    "requests-mock": "requests-mock==1.10.0",
-    "importlib_metadata": "importlib_metadata",
-    "invisible-watermark": "invisible-watermark>=0.2.0",
-    "isort": "isort>=5.5.4",
-    "Jinja2": "Jinja2",
-    "torchsde": "torchsde",
-    "note_seq": "note_seq",
-    "librosa": "librosa",
-    "llvmlite": "llvmlite>=0.40.0",
-    "numba": "numba>=0.57.0",
-    "numpy": "numpy",
-    "parameterized": "parameterized",
-    "peft": "peft>=0.17.0",
-    "protobuf": "protobuf>=3.20.3,<4",
-    "pytest": "pytest",
-    "pytest-timeout": "pytest-timeout",
-    "pytest-xdist": "pytest-xdist",
-    "python": "python>=3.10.0",
-    "ruff": "ruff==0.9.10",
-    "safetensors": "safetensors>=0.8.0",
-    "sentencepiece": "sentencepiece>=0.1.91,!=0.1.92",
-    "GitPython": "GitPython<3.1.19",
-    "scipy": "scipy",
-    "onnx": "onnx",
-    "optimum_quanto": "optimum_quanto>=0.2.6",
-    "gguf": "gguf>=0.10.0",
-    "auto-round": "auto-round>=0.13.0",
-    "torchao": "torchao>=0.7.0",
-    "bitsandbytes": "bitsandbytes>=0.43.3",
-    "nvidia_modelopt[hf]": "nvidia_modelopt[hf]>=0.33.1",
-    "sdnq": "sdnq>=0.2.2",
-    "regex": "regex!=2019.12.17",
-    "requests": "requests",
-    "tensorboard": "tensorboard",
-    "tiktoken": "tiktoken>=0.7.0",
-    "torch": "torch>=2.6",
-    "torchvision": "torchvision",
-    "transformers": "transformers>=4.41.2",
-    "urllib3": "urllib3<=2.0.0",
-    "black": "black",
-    "phonemizer": "phonemizer",
-    "opencv-python": "opencv-python",
-    "timm": "timm",
-    "flashpack": "flashpack",
-}

diffusers/experimental/README.md

@@ -1,5 +0,0 @@
-# 🧨 Diffusers Experimental
-
-We are adding experimental code to support novel applications and usages of the Diffusers library.
-Currently, the following experiments are supported:
-* Reinforcement learning via an implementation of the [Diffuser](https://huggingface.co/papers/2205.09991) model.

diffusers/experimental/__init__.py

@@ -1 +0,0 @@
-from .rl import ValueGuidedRLPipeline

diffusers/experimental/rl/__init__.py

@@ -1 +0,0 @@
-from .value_guided_sampling import ValueGuidedRLPipeline

diffusers/experimental/rl/value_guided_sampling.py

@@ -1,153 +0,0 @@
-# Copyright 2026 The HuggingFace Team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-import numpy as np
-import torch
-import tqdm
-
-from ...models.unets.unet_1d import UNet1DModel
-from ...pipelines import DiffusionPipeline
-from ...utils.dummy_pt_objects import DDPMScheduler
-from ...utils.torch_utils import randn_tensor
-
-
-class ValueGuidedRLPipeline(DiffusionPipeline):
-    r"""
-    Pipeline for value-guided sampling from a diffusion model trained to predict sequences of states.
-
-    This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods
-    implemented for all pipelines (downloading, saving, running on a particular device, etc.).
-
-    Parameters:
-        value_function ([`UNet1DModel`]):
-            A specialized UNet for fine-tuning trajectories base on reward.
-        unet ([`UNet1DModel`]):
-            UNet architecture to denoise the encoded trajectories.
-        scheduler ([`SchedulerMixin`]):
-            A scheduler to be used in combination with `unet` to denoise the encoded trajectories. Default for this
-            application is [`DDPMScheduler`].
-        env ():
-            An environment following the OpenAI gym API to act in. For now only Hopper has pretrained models.
-    """
-
-    def __init__(
-        self,
-        value_function: UNet1DModel,
-        unet: UNet1DModel,
-        scheduler: DDPMScheduler,
-        env,
-    ):
-        super().__init__()
-
-        self.register_modules(value_function=value_function, unet=unet, scheduler=scheduler, env=env)
-
-        self.data = env.get_dataset()
-        self.means = {}
-        for key in self.data.keys():
-            try:
-                self.means[key] = self.data[key].mean()
-            except:  # noqa: E722
-                pass
-        self.stds = {}
-        for key in self.data.keys():
-            try:
-                self.stds[key] = self.data[key].std()
-            except:  # noqa: E722
-                pass
-        self.state_dim = env.observation_space.shape[0]
-        self.action_dim = env.action_space.shape[0]
-
-    def normalize(self, x_in, key):
-        return (x_in - self.means[key]) / self.stds[key]
-
-    def de_normalize(self, x_in, key):
-        return x_in * self.stds[key] + self.means[key]
-
-    def to_torch(self, x_in):
-        if isinstance(x_in, dict):
-            return {k: self.to_torch(v) for k, v in x_in.items()}
-        elif torch.is_tensor(x_in):
-            return x_in.to(self.unet.device)
-        return torch.tensor(x_in, device=self.unet.device)
-
-    def reset_x0(self, x_in, cond, act_dim):
-        for key, val in cond.items():
-            x_in[:, key, act_dim:] = val.clone()
-        return x_in
-
-    def run_diffusion(self, x, conditions, n_guide_steps, scale):
-        batch_size = x.shape[0]
-        y = None
-        for i in tqdm.tqdm(self.scheduler.timesteps):
-            # create batch of timesteps to pass into model
-            timesteps = torch.full((batch_size,), i, device=self.unet.device, dtype=torch.long)
-            for _ in range(n_guide_steps):
-                with torch.enable_grad():
-                    x.requires_grad_()
-
-                    # permute to match dimension for pre-trained models
-                    y = self.value_function(x.permute(0, 2, 1), timesteps).sample
-                    grad = torch.autograd.grad([y.sum()], [x])[0]
-
-                    posterior_variance = self.scheduler._get_variance(i)
-                    model_std = torch.exp(0.5 * posterior_variance)
-                    grad = model_std * grad
-
-                grad[timesteps < 2] = 0
-                x = x.detach()
-                x = x + scale * grad
-                x = self.reset_x0(x, conditions, self.action_dim)
-
-            prev_x = self.unet(x.permute(0, 2, 1), timesteps).sample.permute(0, 2, 1)
-
-            # TODO: verify deprecation of this kwarg
-            x = self.scheduler.step(prev_x, i, x)["prev_sample"]
-
-            # apply conditions to the trajectory (set the initial state)
-            x = self.reset_x0(x, conditions, self.action_dim)
-            x = self.to_torch(x)
-        return x, y
-
-    def __call__(self, obs, batch_size=64, planning_horizon=32, n_guide_steps=2, scale=0.1):
-        # normalize the observations and create  batch dimension
-        obs = self.normalize(obs, "observations")
-        obs = obs[None].repeat(batch_size, axis=0)
-
-        conditions = {0: self.to_torch(obs)}
-        shape = (batch_size, planning_horizon, self.state_dim + self.action_dim)
-
-        # generate initial noise and apply our conditions (to make the trajectories start at current state)
-        x1 = randn_tensor(shape, device=self.unet.device)
-        x = self.reset_x0(x1, conditions, self.action_dim)
-        x = self.to_torch(x)
-
-        # run the diffusion process
-        x, y = self.run_diffusion(x, conditions, n_guide_steps, scale)
-
-        # sort output trajectories by value
-        sorted_idx = y.argsort(0, descending=True).squeeze()
-        sorted_values = x[sorted_idx]
-        actions = sorted_values[:, :, : self.action_dim]
-        actions = actions.detach().cpu().numpy()
-        denorm_actions = self.de_normalize(actions, key="actions")
-
-        # select the action with the highest value
-        if y is not None:
-            selected_index = 0
-        else:
-            # if we didn't run value guiding, select a random action
-            selected_index = np.random.randint(0, batch_size)
-
-        denorm_actions = denorm_actions[selected_index, 0]
-        return denorm_actions

diffusers/guiders/__init__.py

@@ -1,31 +0,0 @@
-# Copyright 2026 The HuggingFace Team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-
-from ..utils import is_torch_available, logging
-
-
-if is_torch_available():
-    from .adaptive_projected_guidance import AdaptiveProjectedGuidance
-    from .adaptive_projected_guidance_mix import AdaptiveProjectedMixGuidance
-    from .auto_guidance import AutoGuidance
-    from .classifier_free_guidance import ClassifierFreeGuidance
-    from .classifier_free_zero_star_guidance import ClassifierFreeZeroStarGuidance
-    from .frequency_decoupled_guidance import FrequencyDecoupledGuidance
-    from .guider_utils import BaseGuidance
-    from .magnitude_aware_guidance import MagnitudeAwareGuidance
-    from .perturbed_attention_guidance import PerturbedAttentionGuidance
-    from .skip_layer_guidance import SkipLayerGuidance
-    from .smoothed_energy_guidance import SmoothedEnergyGuidance
-    from .tangential_classifier_free_guidance import TangentialClassifierFreeGuidance

diffusers/guiders/adaptive_projected_guidance.py

@@ -1,253 +0,0 @@
-# Copyright 2026 The HuggingFace Team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-from __future__ import annotations
-
-import math
-from typing import TYPE_CHECKING
-
-import torch
-
-from ..configuration_utils import register_to_config
-from .guider_utils import BaseGuidance, GuiderOutput, rescale_noise_cfg
-
-
-if TYPE_CHECKING:
-    from ..modular_pipelines.modular_pipeline import BlockState
-
-
-class AdaptiveProjectedGuidance(BaseGuidance):
-    """
-    Adaptive Projected Guidance (APG): https://huggingface.co/papers/2410.02416
-
-    Args:
-        guidance_scale (`float`, defaults to `7.5`):
-            The scale parameter for classifier-free guidance. Higher values result in stronger conditioning on the text
-            prompt, while lower values allow for more freedom in generation. Higher values may lead to saturation and
-            deterioration of image quality.
-        adaptive_projected_guidance_momentum (`float`, defaults to `None`):
-            The momentum parameter for the adaptive projected guidance. Disabled if set to `None`.
-        adaptive_projected_guidance_rescale (`float`, defaults to `15.0`):
-            The rescale factor applied to the noise predictions. This is used to improve image quality and fix
-        adaptive_projected_guidance_norm_dim (`int` or `tuple[int]`, *optional*):
-            Dimension(s) over which to compute the APG norm and projection. If omitted, all non-batch dimensions are
-            used, preserving the original behavior.
-        guidance_rescale (`float`, defaults to `0.0`):
-            The rescale factor applied to the noise predictions. This is used to improve image quality and fix
-            overexposure. Based on Section 3.4 from [Common Diffusion Noise Schedules and Sample Steps are
-            Flawed](https://huggingface.co/papers/2305.08891).
-        use_original_formulation (`bool`, defaults to `False`):
-            Whether to use the original formulation of classifier-free guidance as proposed in the paper. By default,
-            we use the diffusers-native implementation that has been in the codebase for a long time. See
-            [~guiders.classifier_free_guidance.ClassifierFreeGuidance] for more details.
-        start (`float`, defaults to `0.0`):
-            The fraction of the total number of denoising steps after which guidance starts.
-        stop (`float`, defaults to `1.0`):
-            The fraction of the total number of denoising steps after which guidance stops.
-    """
-
-    _input_predictions = ["pred_cond", "pred_uncond"]
-
-    @register_to_config
-    def __init__(
-        self,
-        guidance_scale: float = 7.5,
-        adaptive_projected_guidance_momentum: float | None = None,
-        adaptive_projected_guidance_rescale: float = 15.0,
-        adaptive_projected_guidance_norm_dim: int | tuple[int, ...] | None = None,
-        eta: float = 1.0,
-        guidance_rescale: float = 0.0,
-        use_original_formulation: bool = False,
-        start: float = 0.0,
-        stop: float = 1.0,
-        enabled: bool = True,
-    ):
-        super().__init__(start, stop, enabled)
-
-        self.guidance_scale = guidance_scale
-        self.adaptive_projected_guidance_momentum = adaptive_projected_guidance_momentum
-        self.adaptive_projected_guidance_rescale = adaptive_projected_guidance_rescale
-        self.adaptive_projected_guidance_norm_dim = adaptive_projected_guidance_norm_dim
-        self.eta = eta
-        self.guidance_rescale = guidance_rescale
-        self.use_original_formulation = use_original_formulation
-        self.momentum_buffer = None
-
-    def prepare_inputs(self, data: dict[str, tuple[torch.Tensor, torch.Tensor]]) -> list["BlockState"]:
-        if self._step == 0:
-            if self.adaptive_projected_guidance_momentum is not None:
-                self.momentum_buffer = MomentumBuffer(self.adaptive_projected_guidance_momentum)
-        tuple_indices = [0] if self.num_conditions == 1 else [0, 1]
-        data_batches = []
-        for tuple_idx, input_prediction in zip(tuple_indices, self._input_predictions):
-            data_batch = self._prepare_batch(data, tuple_idx, input_prediction)
-            data_batches.append(data_batch)
-        return data_batches
-
-    def prepare_inputs_from_block_state(
-        self, data: "BlockState", input_fields: dict[str, str | tuple[str, str]]
-    ) -> list["BlockState"]:
-        if self._step == 0:
-            if self.adaptive_projected_guidance_momentum is not None:
-                self.momentum_buffer = MomentumBuffer(self.adaptive_projected_guidance_momentum)
-        tuple_indices = [0] if self.num_conditions == 1 else [0, 1]
-        data_batches = []
-        for tuple_idx, input_prediction in zip(tuple_indices, self._input_predictions):
-            data_batch = self._prepare_batch_from_block_state(input_fields, data, tuple_idx, input_prediction)
-            data_batches.append(data_batch)
-        return data_batches
-
-    def forward(self, pred_cond: torch.Tensor, pred_uncond: torch.Tensor | None = None) -> GuiderOutput:
-        pred = None
-
-        if not self._is_apg_enabled():
-            pred = pred_cond
-        else:
-            pred = normalized_guidance(
-                pred_cond,
-                pred_uncond,
-                self.guidance_scale,
-                self.momentum_buffer,
-                self.eta,
-                self.adaptive_projected_guidance_rescale,
-                self.use_original_formulation,
-                self.adaptive_projected_guidance_norm_dim,
-            )
-
-        if self.guidance_rescale > 0.0:
-            pred = rescale_noise_cfg(pred, pred_cond, self.guidance_rescale)
-
-        return GuiderOutput(pred=pred, pred_cond=pred_cond, pred_uncond=pred_uncond)
-
-    @property
-    def is_conditional(self) -> bool:
-        return self._count_prepared == 1
-
-    @property
-    def num_conditions(self) -> int:
-        num_conditions = 1
-        if self._is_apg_enabled():
-            num_conditions += 1
-        return num_conditions
-
-    def _is_apg_enabled(self) -> bool:
-        if not self._enabled:
-            return False
-
-        is_within_range = True
-        if self._num_inference_steps is not None:
-            skip_start_step = int(self._start * self._num_inference_steps)
-            skip_stop_step = int(self._stop * self._num_inference_steps)
-            is_within_range = skip_start_step <= self._step < skip_stop_step
-
-        is_close = False
-        if self.use_original_formulation:
-            is_close = math.isclose(self.guidance_scale, 0.0)
-        else:
-            is_close = math.isclose(self.guidance_scale, 1.0)
-
-        return is_within_range and not is_close
-
-
-class MomentumBuffer:
-    def __init__(self, momentum: float):
-        self.momentum = momentum
-        self.running_average = 0
-
-    def update(self, update_value: torch.Tensor):
-        new_average = self.momentum * self.running_average
-        self.running_average = update_value + new_average
-
-    def __repr__(self) -> str:
-        """
-        Returns a string representation showing momentum, shape, statistics, and a slice of the running_average.
-        """
-        if isinstance(self.running_average, torch.Tensor):
-            shape = tuple(self.running_average.shape)
-
-            # Calculate statistics
-            with torch.no_grad():
-                stats = {
-                    "mean": self.running_average.mean().item(),
-                    "std": self.running_average.std().item(),
-                    "min": self.running_average.min().item(),
-                    "max": self.running_average.max().item(),
-                }
-
-            # Get a slice (max 3 elements per dimension)
-            slice_indices = tuple(slice(None, min(3, dim)) for dim in shape)
-            sliced_data = self.running_average[slice_indices]
-
-            # Format the slice for display (convert to float32 for numpy compatibility with bfloat16)
-            slice_str = str(sliced_data.detach().float().cpu().numpy())
-            if len(slice_str) > 200:  # Truncate if too long
-                slice_str = slice_str[:200] + "..."
-
-            stats_str = ", ".join([f"{k}={v:.4f}" for k, v in stats.items()])
-
-            return (
-                f"MomentumBuffer(\n"
-                f"  momentum={self.momentum},\n"
-                f"  shape={shape},\n"
-                f"  stats=[{stats_str}],\n"
-                f"  slice={slice_str}\n"
-                f")"
-            )
-        else:
-            return f"MomentumBuffer(momentum={self.momentum}, running_average={self.running_average})"
-
-
-def normalized_guidance(
-    pred_cond: torch.Tensor,
-    pred_uncond: torch.Tensor,
-    guidance_scale: float,
-    momentum_buffer: MomentumBuffer | None = None,
-    eta: float = 1.0,
-    norm_threshold: float = 0.0,
-    use_original_formulation: bool = False,
-    norm_dim: int | tuple[int, ...] | None = None,
-):
-    diff = pred_cond - pred_uncond
-    if norm_dim is None:
-        dim = [-i for i in range(1, len(diff.shape))]
-    elif isinstance(norm_dim, int):
-        dim = [norm_dim]
-    else:
-        dim = list(norm_dim)
-
-    if momentum_buffer is not None:
-        momentum_buffer.update(diff)
-        diff = momentum_buffer.running_average
-
-    if norm_threshold > 0:
-        ones = torch.ones_like(diff)
-        diff_norm = diff.norm(p=2, dim=dim, keepdim=True)
-        scale_factor = torch.minimum(ones, norm_threshold / diff_norm)
-        diff = diff * scale_factor
-
-    if diff.device.type in {"mps", "npu"}:
-        v0, v1 = diff.cpu().double(), pred_cond.cpu().double()
-    else:
-        v0, v1 = diff.double(), pred_cond.double()
-    v1 = torch.nn.functional.normalize(v1, dim=dim)
-    v0_parallel = (v0 * v1).sum(dim=dim, keepdim=True) * v1
-    v0_orthogonal = v0 - v0_parallel
-    diff_parallel = v0_parallel.to(device=diff.device, dtype=diff.dtype)
-    diff_orthogonal = v0_orthogonal.to(device=diff.device, dtype=diff.dtype)
-    normalized_update = diff_orthogonal + eta * diff_parallel
-
-    pred = pred_cond if use_original_formulation else pred_uncond
-    pred = pred + guidance_scale * normalized_update
-
-    return pred

diffusers/guiders/adaptive_projected_guidance_mix.py

@@ -1,297 +0,0 @@
-# Copyright 2026 The HuggingFace Team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-import math
-from typing import TYPE_CHECKING
-
-import torch
-
-from ..configuration_utils import register_to_config
-from .guider_utils import BaseGuidance, GuiderOutput, rescale_noise_cfg
-
-
-if TYPE_CHECKING:
-    from ..modular_pipelines.modular_pipeline import BlockState
-
-
-class AdaptiveProjectedMixGuidance(BaseGuidance):
-    """
-    Adaptive Projected Guidance (APG) https://huggingface.co/papers/2410.02416 combined with Classifier-Free Guidance
-    (CFG). This guider is used in HunyuanImage2.1 https://github.com/Tencent-Hunyuan/HunyuanImage-2.1
-
-    Args:
-        guidance_scale (`float`, defaults to `7.5`):
-            The scale parameter for classifier-free guidance. Higher values result in stronger conditioning on the text
-            prompt, while lower values allow for more freedom in generation. Higher values may lead to saturation and
-            deterioration of image quality.
-        adaptive_projected_guidance_momentum (`float`, defaults to `None`):
-            The momentum parameter for the adaptive projected guidance. Disabled if set to `None`.
-        adaptive_projected_guidance_rescale (`float`, defaults to `15.0`):
-            The rescale factor applied to the noise predictions for adaptive projected guidance. This is used to
-            improve image quality and fix
-        guidance_rescale (`float`, defaults to `0.0`):
-            The rescale factor applied to the noise predictions for classifier-free guidance. This is used to improve
-            image quality and fix overexposure. Based on Section 3.4 from [Common Diffusion Noise Schedules and Sample
-            Steps are Flawed](https://huggingface.co/papers/2305.08891).
-        use_original_formulation (`bool`, defaults to `False`):
-            Whether to use the original formulation of classifier-free guidance as proposed in the paper. By default,
-            we use the diffusers-native implementation that has been in the codebase for a long time. See
-            [~guiders.classifier_free_guidance.ClassifierFreeGuidance] for more details.
-        start (`float`, defaults to `0.0`):
-            The fraction of the total number of denoising steps after which the classifier-free guidance starts.
-        stop (`float`, defaults to `1.0`):
-            The fraction of the total number of denoising steps after which the classifier-free guidance stops.
-        adaptive_projected_guidance_start_step (`int`, defaults to `5`):
-            The step at which the adaptive projected guidance starts (before this step, classifier-free guidance is
-            used, and momentum buffer is updated).
-        enabled (`bool`, defaults to `True`):
-            Whether this guidance is enabled.
-    """
-
-    _input_predictions = ["pred_cond", "pred_uncond"]
-
-    @register_to_config
-    def __init__(
-        self,
-        guidance_scale: float = 3.5,
-        guidance_rescale: float = 0.0,
-        adaptive_projected_guidance_scale: float = 10.0,
-        adaptive_projected_guidance_momentum: float = -0.5,
-        adaptive_projected_guidance_rescale: float = 10.0,
-        eta: float = 0.0,
-        use_original_formulation: bool = False,
-        start: float = 0.0,
-        stop: float = 1.0,
-        adaptive_projected_guidance_start_step: int = 5,
-        enabled: bool = True,
-    ):
-        super().__init__(start, stop, enabled)
-
-        self.guidance_scale = guidance_scale
-        self.guidance_rescale = guidance_rescale
-        self.adaptive_projected_guidance_scale = adaptive_projected_guidance_scale
-        self.adaptive_projected_guidance_momentum = adaptive_projected_guidance_momentum
-        self.adaptive_projected_guidance_rescale = adaptive_projected_guidance_rescale
-        self.eta = eta
-        self.adaptive_projected_guidance_start_step = adaptive_projected_guidance_start_step
-        self.use_original_formulation = use_original_formulation
-        self.momentum_buffer = None
-
-    def prepare_inputs(self, data: dict[str, tuple[torch.Tensor, torch.Tensor]]) -> list["BlockState"]:
-        if self._step == 0:
-            if self.adaptive_projected_guidance_momentum is not None:
-                self.momentum_buffer = MomentumBuffer(self.adaptive_projected_guidance_momentum)
-        tuple_indices = [0] if self.num_conditions == 1 else [0, 1]
-        data_batches = []
-        for tuple_idx, input_prediction in zip(tuple_indices, self._input_predictions):
-            data_batch = self._prepare_batch(data, tuple_idx, input_prediction)
-            data_batches.append(data_batch)
-        return data_batches
-
-    def prepare_inputs_from_block_state(
-        self, data: "BlockState", input_fields: dict[str, str | tuple[str, str]]
-    ) -> list["BlockState"]:
-        if self._step == 0:
-            if self.adaptive_projected_guidance_momentum is not None:
-                self.momentum_buffer = MomentumBuffer(self.adaptive_projected_guidance_momentum)
-        tuple_indices = [0] if self.num_conditions == 1 else [0, 1]
-        data_batches = []
-        for tuple_idx, input_prediction in zip(tuple_indices, self._input_predictions):
-            data_batch = self._prepare_batch_from_block_state(input_fields, data, tuple_idx, input_prediction)
-            data_batches.append(data_batch)
-        return data_batches
-
-    def forward(self, pred_cond: torch.Tensor, pred_uncond: torch.Tensor | None = None) -> GuiderOutput:
-        pred = None
-
-        # no guidance
-        if not self._is_cfg_enabled():
-            pred = pred_cond
-
-        # CFG + update momentum buffer
-        elif not self._is_apg_enabled():
-            if self.momentum_buffer is not None:
-                update_momentum_buffer(pred_cond, pred_uncond, self.momentum_buffer)
-            # CFG + update momentum buffer
-            shift = pred_cond - pred_uncond
-            pred = pred_cond if self.use_original_formulation else pred_uncond
-            pred = pred + self.guidance_scale * shift
-
-        # APG
-        elif self._is_apg_enabled():
-            pred = normalized_guidance(
-                pred_cond,
-                pred_uncond,
-                self.adaptive_projected_guidance_scale,
-                self.momentum_buffer,
-                self.eta,
-                self.adaptive_projected_guidance_rescale,
-                self.use_original_formulation,
-            )
-
-        if self.guidance_rescale > 0.0:
-            pred = rescale_noise_cfg(pred, pred_cond, self.guidance_rescale)
-
-        return GuiderOutput(pred=pred, pred_cond=pred_cond, pred_uncond=pred_uncond)
-
-    @property
-    def is_conditional(self) -> bool:
-        return self._count_prepared == 1
-
-    @property
-    def num_conditions(self) -> int:
-        num_conditions = 1
-        if self._is_apg_enabled() or self._is_cfg_enabled():
-            num_conditions += 1
-        return num_conditions
-
-    # Copied from diffusers.guiders.classifier_free_guidance.ClassifierFreeGuidance._is_cfg_enabled
-    def _is_cfg_enabled(self) -> bool:
-        if not self._enabled:
-            return False
-
-        is_within_range = True
-        if self._num_inference_steps is not None:
-            skip_start_step = int(self._start * self._num_inference_steps)
-            skip_stop_step = int(self._stop * self._num_inference_steps)
-            is_within_range = skip_start_step <= self._step < skip_stop_step
-
-        is_close = False
-        if self.use_original_formulation:
-            is_close = math.isclose(self.guidance_scale, 0.0)
-        else:
-            is_close = math.isclose(self.guidance_scale, 1.0)
-
-        return is_within_range and not is_close
-
-    def _is_apg_enabled(self) -> bool:
-        if not self._enabled:
-            return False
-
-        if not self._is_cfg_enabled():
-            return False
-
-        is_within_range = False
-        if self._step is not None:
-            is_within_range = self._step > self.adaptive_projected_guidance_start_step
-
-        is_close = False
-        if self.use_original_formulation:
-            is_close = math.isclose(self.adaptive_projected_guidance_scale, 0.0)
-        else:
-            is_close = math.isclose(self.adaptive_projected_guidance_scale, 1.0)
-
-        return is_within_range and not is_close
-
-    def get_state(self):
-        state = super().get_state()
-        state["momentum_buffer"] = self.momentum_buffer
-        state["is_apg_enabled"] = self._is_apg_enabled()
-        state["is_cfg_enabled"] = self._is_cfg_enabled()
-        return state
-
-
-# Copied from diffusers.guiders.adaptive_projected_guidance.MomentumBuffer
-class MomentumBuffer:
-    def __init__(self, momentum: float):
-        self.momentum = momentum
-        self.running_average = 0
-
-    def update(self, update_value: torch.Tensor):
-        new_average = self.momentum * self.running_average
-        self.running_average = update_value + new_average
-
-    def __repr__(self) -> str:
-        """
-        Returns a string representation showing momentum, shape, statistics, and a slice of the running_average.
-        """
-        if isinstance(self.running_average, torch.Tensor):
-            shape = tuple(self.running_average.shape)
-
-            # Calculate statistics
-            with torch.no_grad():
-                stats = {
-                    "mean": self.running_average.mean().item(),
-                    "std": self.running_average.std().item(),
-                    "min": self.running_average.min().item(),
-                    "max": self.running_average.max().item(),
-                }
-
-            # Get a slice (max 3 elements per dimension)
-            slice_indices = tuple(slice(None, min(3, dim)) for dim in shape)
-            sliced_data = self.running_average[slice_indices]
-
-            # Format the slice for display (convert to float32 for numpy compatibility with bfloat16)
-            slice_str = str(sliced_data.detach().float().cpu().numpy())
-            if len(slice_str) > 200:  # Truncate if too long
-                slice_str = slice_str[:200] + "..."
-
-            stats_str = ", ".join([f"{k}={v:.4f}" for k, v in stats.items()])
-
-            return (
-                f"MomentumBuffer(\n"
-                f"  momentum={self.momentum},\n"
-                f"  shape={shape},\n"
-                f"  stats=[{stats_str}],\n"
-                f"  slice={slice_str}\n"
-                f")"
-            )
-        else:
-            return f"MomentumBuffer(momentum={self.momentum}, running_average={self.running_average})"
-
-
-def update_momentum_buffer(
-    pred_cond: torch.Tensor,
-    pred_uncond: torch.Tensor,
-    momentum_buffer: MomentumBuffer | None = None,
-):
-    diff = pred_cond - pred_uncond
-    if momentum_buffer is not None:
-        momentum_buffer.update(diff)
-
-
-def normalized_guidance(
-    pred_cond: torch.Tensor,
-    pred_uncond: torch.Tensor,
-    guidance_scale: float,
-    momentum_buffer: MomentumBuffer | None = None,
-    eta: float = 1.0,
-    norm_threshold: float = 0.0,
-    use_original_formulation: bool = False,
-):
-    if momentum_buffer is not None:
-        update_momentum_buffer(pred_cond, pred_uncond, momentum_buffer)
-        diff = momentum_buffer.running_average
-    else:
-        diff = pred_cond - pred_uncond
-
-    dim = [-i for i in range(1, len(diff.shape))]
-
-    if norm_threshold > 0:
-        ones = torch.ones_like(diff)
-        diff_norm = diff.norm(p=2, dim=dim, keepdim=True)
-        scale_factor = torch.minimum(ones, norm_threshold / diff_norm)
-        diff = diff * scale_factor
-
-    v0, v1 = diff.double(), pred_cond.double()
-    v1 = torch.nn.functional.normalize(v1, dim=dim)
-    v0_parallel = (v0 * v1).sum(dim=dim, keepdim=True) * v1
-    v0_orthogonal = v0 - v0_parallel
-    diff_parallel, diff_orthogonal = v0_parallel.type_as(diff), v0_orthogonal.type_as(diff)
-    normalized_update = diff_orthogonal + eta * diff_parallel
-
-    pred = pred_cond if use_original_formulation else pred_uncond
-    pred = pred + guidance_scale * normalized_update
-
-    return pred

diffusers/guiders/auto_guidance.py

@@ -1,198 +0,0 @@
-# Copyright 2026 The HuggingFace Team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-from __future__ import annotations
-
-import math
-from typing import TYPE_CHECKING, Any
-
-import torch
-
-from ..configuration_utils import register_to_config
-from ..hooks import HookRegistry, LayerSkipConfig
-from ..hooks.layer_skip import _apply_layer_skip_hook
-from .guider_utils import BaseGuidance, GuiderOutput, rescale_noise_cfg
-
-
-if TYPE_CHECKING:
-    from ..modular_pipelines.modular_pipeline import BlockState
-
-
-class AutoGuidance(BaseGuidance):
-    """
-    AutoGuidance: https://huggingface.co/papers/2406.02507
-
-    Args:
-        guidance_scale (`float`, defaults to `7.5`):
-            The scale parameter for classifier-free guidance. Higher values result in stronger conditioning on the text
-            prompt, while lower values allow for more freedom in generation. Higher values may lead to saturation and
-            deterioration of image quality.
-        auto_guidance_layers (`int` or `list[int]`, *optional*):
-            The layer indices to apply skip layer guidance to. Can be a single integer or a list of integers. If not
-            provided, `skip_layer_config` must be provided.
-        auto_guidance_config (`LayerSkipConfig` or `list[LayerSkipConfig]`, *optional*):
-            The configuration for the skip layer guidance. Can be a single `LayerSkipConfig` or a list of
-            `LayerSkipConfig`. If not provided, `skip_layer_guidance_layers` must be provided.
-        dropout (`float`, *optional*):
-            The dropout probability for autoguidance on the enabled skip layers (either with `auto_guidance_layers` or
-            `auto_guidance_config`). If not provided, the dropout probability will be set to 1.0.
-        guidance_rescale (`float`, defaults to `0.0`):
-            The rescale factor applied to the noise predictions. This is used to improve image quality and fix
-            overexposure. Based on Section 3.4 from [Common Diffusion Noise Schedules and Sample Steps are
-            Flawed](https://huggingface.co/papers/2305.08891).
-        use_original_formulation (`bool`, defaults to `False`):
-            Whether to use the original formulation of classifier-free guidance as proposed in the paper. By default,
-            we use the diffusers-native implementation that has been in the codebase for a long time. See
-            [~guiders.classifier_free_guidance.ClassifierFreeGuidance] for more details.
-        start (`float`, defaults to `0.0`):
-            The fraction of the total number of denoising steps after which guidance starts.
-        stop (`float`, defaults to `1.0`):
-            The fraction of the total number of denoising steps after which guidance stops.
-    """
-
-    _input_predictions = ["pred_cond", "pred_uncond"]
-
-    @register_to_config
-    def __init__(
-        self,
-        guidance_scale: float = 7.5,
-        auto_guidance_layers: int | list[int] | None = None,
-        auto_guidance_config: LayerSkipConfig | list[LayerSkipConfig] | dict[str, Any] = None,
-        dropout: float | None = None,
-        guidance_rescale: float = 0.0,
-        use_original_formulation: bool = False,
-        start: float = 0.0,
-        stop: float = 1.0,
-        enabled: bool = True,
-    ):
-        super().__init__(start, stop, enabled)
-
-        self.guidance_scale = guidance_scale
-        self.auto_guidance_layers = auto_guidance_layers
-        self.auto_guidance_config = auto_guidance_config
-        self.dropout = dropout
-        self.guidance_rescale = guidance_rescale
-        self.use_original_formulation = use_original_formulation
-
-        is_layer_or_config_provided = auto_guidance_layers is not None or auto_guidance_config is not None
-        is_layer_and_config_provided = auto_guidance_layers is not None and auto_guidance_config is not None
-        if not is_layer_or_config_provided:
-            raise ValueError(
-                "Either `auto_guidance_layers` or `auto_guidance_config` must be provided to enable AutoGuidance."
-            )
-        if is_layer_and_config_provided:
-            raise ValueError("Only one of `auto_guidance_layers` or `auto_guidance_config` can be provided.")
-        if auto_guidance_config is None and dropout is None:
-            raise ValueError("`dropout` must be provided if `auto_guidance_layers` is provided.")
-
-        if auto_guidance_layers is not None:
-            if isinstance(auto_guidance_layers, int):
-                auto_guidance_layers = [auto_guidance_layers]
-            if not isinstance(auto_guidance_layers, list):
-                raise ValueError(
-                    f"Expected `auto_guidance_layers` to be an int or a list of ints, but got {type(auto_guidance_layers)}."
-                )
-            auto_guidance_config = [
-                LayerSkipConfig(layer, fqn="auto", dropout=dropout) for layer in auto_guidance_layers
-            ]
-
-        if isinstance(auto_guidance_config, dict):
-            auto_guidance_config = LayerSkipConfig.from_dict(auto_guidance_config)
-
-        if isinstance(auto_guidance_config, LayerSkipConfig):
-            auto_guidance_config = [auto_guidance_config]
-
-        if not isinstance(auto_guidance_config, list):
-            raise ValueError(
-                f"Expected `auto_guidance_config` to be a LayerSkipConfig or a list of LayerSkipConfig, but got {type(auto_guidance_config)}."
-            )
-        elif isinstance(next(iter(auto_guidance_config), None), dict):
-            auto_guidance_config = [LayerSkipConfig.from_dict(config) for config in auto_guidance_config]
-
-        self.auto_guidance_config = auto_guidance_config
-        self._auto_guidance_hook_names = [f"AutoGuidance_{i}" for i in range(len(self.auto_guidance_config))]
-
-    def prepare_models(self, denoiser: torch.nn.Module) -> None:
-        self._count_prepared += 1
-        if self._is_ag_enabled() and self.is_unconditional:
-            for name, config in zip(self._auto_guidance_hook_names, self.auto_guidance_config):
-                _apply_layer_skip_hook(denoiser, config, name=name)
-
-    def cleanup_models(self, denoiser: torch.nn.Module) -> None:
-        if self._is_ag_enabled() and self.is_unconditional:
-            for name in self._auto_guidance_hook_names:
-                registry = HookRegistry.check_if_exists_or_initialize(denoiser)
-                registry.remove_hook(name, recurse=True)
-
-    def prepare_inputs(self, data: dict[str, tuple[torch.Tensor, torch.Tensor]]) -> list["BlockState"]:
-        tuple_indices = [0] if self.num_conditions == 1 else [0, 1]
-        data_batches = []
-        for tuple_idx, input_prediction in zip(tuple_indices, self._input_predictions):
-            data_batch = self._prepare_batch(data, tuple_idx, input_prediction)
-            data_batches.append(data_batch)
-        return data_batches
-
-    def prepare_inputs_from_block_state(
-        self, data: "BlockState", input_fields: dict[str, str | tuple[str, str]]
-    ) -> list["BlockState"]:
-        tuple_indices = [0] if self.num_conditions == 1 else [0, 1]
-        data_batches = []
-        for tuple_idx, input_prediction in zip(tuple_indices, self._input_predictions):
-            data_batch = self._prepare_batch_from_block_state(input_fields, data, tuple_idx, input_prediction)
-            data_batches.append(data_batch)
-        return data_batches
-
-    def forward(self, pred_cond: torch.Tensor, pred_uncond: torch.Tensor | None = None) -> GuiderOutput:
-        pred = None
-
-        if not self._is_ag_enabled():
-            pred = pred_cond
-        else:
-            shift = pred_cond - pred_uncond
-            pred = pred_cond if self.use_original_formulation else pred_uncond
-            pred = pred + self.guidance_scale * shift
-
-        if self.guidance_rescale > 0.0:
-            pred = rescale_noise_cfg(pred, pred_cond, self.guidance_rescale)
-
-        return GuiderOutput(pred=pred, pred_cond=pred_cond, pred_uncond=pred_uncond)
-
-    @property
-    def is_conditional(self) -> bool:
-        return self._count_prepared == 1
-
-    @property
-    def num_conditions(self) -> int:
-        num_conditions = 1
-        if self._is_ag_enabled():
-            num_conditions += 1
-        return num_conditions
-
-    def _is_ag_enabled(self) -> bool:
-        if not self._enabled:
-            return False
-
-        is_within_range = True
-        if self._num_inference_steps is not None:
-            skip_start_step = int(self._start * self._num_inference_steps)
-            skip_stop_step = int(self._stop * self._num_inference_steps)
-            is_within_range = skip_start_step <= self._step < skip_stop_step
-
-        is_close = False
-        if self.use_original_formulation:
-            is_close = math.isclose(self.guidance_scale, 0.0)
-        else:
-            is_close = math.isclose(self.guidance_scale, 1.0)
-
-        return is_within_range and not is_close

diffusers/guiders/classifier_free_guidance.py

@@ -1,156 +0,0 @@
-# Copyright 2026 The HuggingFace Team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-from __future__ import annotations
-
-import math
-from typing import TYPE_CHECKING
-
-import torch
-
-from ..configuration_utils import register_to_config
-from .guider_utils import BaseGuidance, GuiderOutput, rescale_noise_cfg
-
-
-if TYPE_CHECKING:
-    from ..modular_pipelines.modular_pipeline import BlockState
-
-
-class ClassifierFreeGuidance(BaseGuidance):
-    """
-    Implements Classifier-Free Guidance (CFG) for diffusion models.
-
-    Reference: https://huggingface.co/papers/2207.12598
-
-    CFG improves generation quality and prompt adherence by jointly training models on both conditional and
-    unconditional data, then combining predictions during inference. This allows trading off between quality (high
-    guidance) and diversity (low guidance).
-
-    **Two CFG Formulations:**
-
-    1. **Original formulation** (from paper):
-       ```
-       x_pred = x_cond + guidance_scale * (x_cond - x_uncond)
-       ```
-       Moves conditional predictions further from unconditional ones.
-
-    2. **Diffusers-native formulation** (default, from Imagen paper):
-       ```
-       x_pred = x_uncond + guidance_scale * (x_cond - x_uncond)
-       ```
-       Moves unconditional predictions toward conditional ones, effectively suppressing negative features (e.g., "bad
-       quality", "watermarks"). Equivalent in theory but more intuitive.
-
-    Use `use_original_formulation=True` to switch to the original formulation.
-
-    Args:
-        guidance_scale (`float`, defaults to `7.5`):
-            CFG scale applied by this guider during post-processing. Higher values = stronger prompt conditioning but
-            may reduce quality. Typical range: 1.0-20.0.
-        guidance_rescale (`float`, defaults to `0.0`):
-            Rescaling factor to prevent overexposure from high guidance scales. Based on [Common Diffusion Noise
-            Schedules and Sample Steps are Flawed](https://huggingface.co/papers/2305.08891). Range: 0.0 (no rescaling)
-            to 1.0 (full rescaling).
-        use_original_formulation (`bool`, defaults to `False`):
-            If `True`, uses the original CFG formulation from the paper. If `False` (default), uses the
-            diffusers-native formulation from the Imagen paper.
-        start (`float`, defaults to `0.0`):
-            Fraction of denoising steps (0.0-1.0) after which CFG starts. Use > 0.0 to disable CFG in early denoising
-            steps.
-        stop (`float`, defaults to `1.0`):
-            Fraction of denoising steps (0.0-1.0) after which CFG stops. Use < 1.0 to disable CFG in late denoising
-            steps.
-        enabled (`bool`, defaults to `True`):
-            Whether CFG is enabled. Set to `False` to disable CFG entirely (uses only conditional predictions).
-    """
-
-    _input_predictions = ["pred_cond", "pred_uncond"]
-
-    @register_to_config
-    def __init__(
-        self,
-        guidance_scale: float = 7.5,
-        guidance_rescale: float = 0.0,
-        use_original_formulation: bool = False,
-        start: float = 0.0,
-        stop: float = 1.0,
-        enabled: bool = True,
-    ):
-        super().__init__(start, stop, enabled)
-
-        self.guidance_scale = guidance_scale
-        self.guidance_rescale = guidance_rescale
-        self.use_original_formulation = use_original_formulation
-
-    def prepare_inputs(self, data: dict[str, tuple[torch.Tensor, torch.Tensor]]) -> list["BlockState"]:
-        tuple_indices = [0] if self.num_conditions == 1 else [0, 1]
-        data_batches = []
-        for tuple_idx, input_prediction in zip(tuple_indices, self._input_predictions):
-            data_batch = self._prepare_batch(data, tuple_idx, input_prediction)
-            data_batches.append(data_batch)
-        return data_batches
-
-    def prepare_inputs_from_block_state(
-        self, data: "BlockState", input_fields: dict[str, str | tuple[str, str]]
-    ) -> list["BlockState"]:
-        tuple_indices = [0] if self.num_conditions == 1 else [0, 1]
-        data_batches = []
-        for tuple_idx, input_prediction in zip(tuple_indices, self._input_predictions):
-            data_batch = self._prepare_batch_from_block_state(input_fields, data, tuple_idx, input_prediction)
-            data_batches.append(data_batch)
-        return data_batches
-
-    def forward(self, pred_cond: torch.Tensor, pred_uncond: torch.Tensor | None = None) -> GuiderOutput:
-        pred = None
-
-        if not self._is_cfg_enabled():
-            pred = pred_cond
-        else:
-            shift = pred_cond - pred_uncond
-            pred = pred_cond if self.use_original_formulation else pred_uncond
-            pred = pred + self.guidance_scale * shift
-
-        if self.guidance_rescale > 0.0:
-            pred = rescale_noise_cfg(pred, pred_cond, self.guidance_rescale)
-
-        return GuiderOutput(pred=pred, pred_cond=pred_cond, pred_uncond=pred_uncond)
-
-    @property
-    def is_conditional(self) -> bool:
-        return self._count_prepared == 1
-
-    @property
-    def num_conditions(self) -> int:
-        num_conditions = 1
-        if self._is_cfg_enabled():
-            num_conditions += 1
-        return num_conditions
-
-    def _is_cfg_enabled(self) -> bool:
-        if not self._enabled:
-            return False
-
-        is_within_range = True
-        if self._num_inference_steps is not None:
-            skip_start_step = int(self._start * self._num_inference_steps)
-            skip_stop_step = int(self._stop * self._num_inference_steps)
-            is_within_range = skip_start_step <= self._step < skip_stop_step
-
-        is_close = False
-        if self.use_original_formulation:
-            is_close = math.isclose(self.guidance_scale, 0.0)
-        else:
-            is_close = math.isclose(self.guidance_scale, 1.0)
-
-        return is_within_range and not is_close

diffusers/guiders/classifier_free_zero_star_guidance.py

@@ -1,164 +0,0 @@
-# Copyright 2026 The HuggingFace Team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-from __future__ import annotations
-
-import math
-from typing import TYPE_CHECKING
-
-import torch
-
-from ..configuration_utils import register_to_config
-from .guider_utils import BaseGuidance, GuiderOutput, rescale_noise_cfg
-
-
-if TYPE_CHECKING:
-    from ..modular_pipelines.modular_pipeline import BlockState
-
-
-class ClassifierFreeZeroStarGuidance(BaseGuidance):
-    """
-    Classifier-free Zero* (CFG-Zero*): https://huggingface.co/papers/2503.18886
-
-    This is an implementation of the Classifier-Free Zero* guidance technique, which is a variant of classifier-free
-    guidance. It proposes zero initialization of the noise predictions for the first few steps of the diffusion
-    process, and also introduces an optimal rescaling factor for the noise predictions, which can help in improving the
-    quality of generated images.
-
-    The authors of the paper suggest setting zero initialization in the first 4% of the inference steps.
-
-    Args:
-        guidance_scale (`float`, defaults to `7.5`):
-            The scale parameter for classifier-free guidance. Higher values result in stronger conditioning on the text
-            prompt, while lower values allow for more freedom in generation. Higher values may lead to saturation and
-            deterioration of image quality.
-        zero_init_steps (`int`, defaults to `1`):
-            The number of inference steps for which the noise predictions are zeroed out (see Section 4.2).
-        guidance_rescale (`float`, defaults to `0.0`):
-            The rescale factor applied to the noise predictions. This is used to improve image quality and fix
-            overexposure. Based on Section 3.4 from [Common Diffusion Noise Schedules and Sample Steps are
-            Flawed](https://huggingface.co/papers/2305.08891).
-        use_original_formulation (`bool`, defaults to `False`):
-            Whether to use the original formulation of classifier-free guidance as proposed in the paper. By default,
-            we use the diffusers-native implementation that has been in the codebase for a long time. See
-            [~guiders.classifier_free_guidance.ClassifierFreeGuidance] for more details.
-        start (`float`, defaults to `0.01`):
-            The fraction of the total number of denoising steps after which guidance starts.
-        stop (`float`, defaults to `0.2`):
-            The fraction of the total number of denoising steps after which guidance stops.
-    """
-
-    _input_predictions = ["pred_cond", "pred_uncond"]
-
-    @register_to_config
-    def __init__(
-        self,
-        guidance_scale: float = 7.5,
-        zero_init_steps: int = 1,
-        guidance_rescale: float = 0.0,
-        use_original_formulation: bool = False,
-        start: float = 0.0,
-        stop: float = 1.0,
-        enabled: bool = True,
-    ):
-        super().__init__(start, stop, enabled)
-
-        self.guidance_scale = guidance_scale
-        self.zero_init_steps = zero_init_steps
-        self.guidance_rescale = guidance_rescale
-        self.use_original_formulation = use_original_formulation
-
-    def prepare_inputs(self, data: dict[str, tuple[torch.Tensor, torch.Tensor]]) -> list["BlockState"]:
-        tuple_indices = [0] if self.num_conditions == 1 else [0, 1]
-        data_batches = []
-        for tuple_idx, input_prediction in zip(tuple_indices, self._input_predictions):
-            data_batch = self._prepare_batch(data, tuple_idx, input_prediction)
-            data_batches.append(data_batch)
-        return data_batches
-
-    def prepare_inputs_from_block_state(
-        self, data: "BlockState", input_fields: dict[str, str | tuple[str, str]]
-    ) -> list["BlockState"]:
-        tuple_indices = [0] if self.num_conditions == 1 else [0, 1]
-        data_batches = []
-        for tuple_idx, input_prediction in zip(tuple_indices, self._input_predictions):
-            data_batch = self._prepare_batch_from_block_state(input_fields, data, tuple_idx, input_prediction)
-            data_batches.append(data_batch)
-        return data_batches
-
-    def forward(self, pred_cond: torch.Tensor, pred_uncond: torch.Tensor | None = None) -> GuiderOutput:
-        pred = None
-
-        # YiYi Notes: add default behavior for self._enabled == False
-        if not self._enabled:
-            pred = pred_cond
-
-        elif self._step < self.zero_init_steps:
-            pred = torch.zeros_like(pred_cond)
-        elif not self._is_cfg_enabled():
-            pred = pred_cond
-        else:
-            pred_cond_flat = pred_cond.flatten(1)
-            pred_uncond_flat = pred_uncond.flatten(1)
-            alpha = cfg_zero_star_scale(pred_cond_flat, pred_uncond_flat)
-            alpha = alpha.view(-1, *(1,) * (len(pred_cond.shape) - 1))
-            pred_uncond = pred_uncond * alpha
-            shift = pred_cond - pred_uncond
-            pred = pred_cond if self.use_original_formulation else pred_uncond
-            pred = pred + self.guidance_scale * shift
-
-        if self.guidance_rescale > 0.0:
-            pred = rescale_noise_cfg(pred, pred_cond, self.guidance_rescale)
-
-        return GuiderOutput(pred=pred, pred_cond=pred_cond, pred_uncond=pred_uncond)
-
-    @property
-    def is_conditional(self) -> bool:
-        return self._count_prepared == 1
-
-    @property
-    def num_conditions(self) -> int:
-        num_conditions = 1
-        if self._is_cfg_enabled():
-            num_conditions += 1
-        return num_conditions
-
-    def _is_cfg_enabled(self) -> bool:
-        if not self._enabled:
-            return False
-
-        is_within_range = True
-        if self._num_inference_steps is not None:
-            skip_start_step = int(self._start * self._num_inference_steps)
-            skip_stop_step = int(self._stop * self._num_inference_steps)
-            is_within_range = skip_start_step <= self._step < skip_stop_step
-
-        is_close = False
-        if self.use_original_formulation:
-            is_close = math.isclose(self.guidance_scale, 0.0)
-        else:
-            is_close = math.isclose(self.guidance_scale, 1.0)
-
-        return is_within_range and not is_close
-
-
-def cfg_zero_star_scale(cond: torch.Tensor, uncond: torch.Tensor, eps: float = 1e-8) -> torch.Tensor:
-    cond_dtype = cond.dtype
-    cond = cond.float()
-    uncond = uncond.float()
-    dot_product = torch.sum(cond * uncond, dim=1, keepdim=True)
-    squared_norm = torch.sum(uncond**2, dim=1, keepdim=True) + eps
-    # st_star = v_cond^T * v_uncond / ||v_uncond||^2
-    scale = dot_product / squared_norm
-    return scale.to(dtype=cond_dtype)

diffusers/guiders/frequency_decoupled_guidance.py

@@ -1,335 +0,0 @@
-# Copyright 2026 The HuggingFace Team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-from __future__ import annotations
-
-import math
-from typing import TYPE_CHECKING
-
-import torch
-
-from ..configuration_utils import register_to_config
-from ..utils import is_kornia_available
-from .guider_utils import BaseGuidance, GuiderOutput, rescale_noise_cfg
-
-
-if TYPE_CHECKING:
-    from ..modular_pipelines.modular_pipeline import BlockState
-
-
-_CAN_USE_KORNIA = is_kornia_available()
-
-
-if _CAN_USE_KORNIA:
-    from kornia.geometry import pyrup as upsample_and_blur_func
-    from kornia.geometry.transform import build_laplacian_pyramid as build_laplacian_pyramid_func
-else:
-    upsample_and_blur_func = None
-    build_laplacian_pyramid_func = None
-
-
-def project(v0: torch.Tensor, v1: torch.Tensor, upcast_to_double: bool = True) -> tuple[torch.Tensor, torch.Tensor]:
-    """
-    Project vector v0 onto vector v1, returning the parallel and orthogonal components of v0. Implementation from paper
-    (Algorithm 2).
-    """
-    # v0 shape: [B, ...]
-    # v1 shape: [B, ...]
-    # Assume first dim is a batch dim and all other dims are channel or "spatial" dims
-    all_dims_but_first = list(range(1, len(v0.shape)))
-    if upcast_to_double:
-        dtype = v0.dtype
-        v0, v1 = v0.double(), v1.double()
-    v1 = torch.nn.functional.normalize(v1, dim=all_dims_but_first)
-    v0_parallel = (v0 * v1).sum(dim=all_dims_but_first, keepdim=True) * v1
-    v0_orthogonal = v0 - v0_parallel
-    if upcast_to_double:
-        v0_parallel = v0_parallel.to(dtype)
-        v0_orthogonal = v0_orthogonal.to(dtype)
-    return v0_parallel, v0_orthogonal
-
-
-def build_image_from_pyramid(pyramid: list[torch.Tensor]) -> torch.Tensor:
-    """
-    Recovers the data space latents from the Laplacian pyramid frequency space. Implementation from the paper
-    (Algorithm 2).
-    """
-    # pyramid shapes: [[B, C, H, W], [B, C, H/2, W/2], ...]
-    img = pyramid[-1]
-    for i in range(len(pyramid) - 2, -1, -1):
-        img = upsample_and_blur_func(img) + pyramid[i]
-    return img
-
-
-class FrequencyDecoupledGuidance(BaseGuidance):
-    """
-    Frequency-Decoupled Guidance (FDG): https://huggingface.co/papers/2506.19713
-
-    FDG is a technique similar to (and based on) classifier-free guidance (CFG) which is used to improve generation
-    quality and condition-following in diffusion models. Like CFG, during training we jointly train the model on both
-    conditional and unconditional data, and use a combination of the two during inference. (If you want more details on
-    how CFG works, you can check out the CFG guider.)
-
-    FDG differs from CFG in that the normal CFG prediction is instead decoupled into low- and high-frequency components
-    using a frequency transform (such as a Laplacian pyramid). The CFG update is then performed in frequency space
-    separately for the low- and high-frequency components with different guidance scales. Finally, the inverse
-    frequency transform is used to map the CFG frequency predictions back to data space (e.g. pixel space for images)
-    to form the final FDG prediction.
-
-    For images, the FDG authors found that using low guidance scales for the low-frequency components retains sample
-    diversity and realistic color composition, while using high guidance scales for high-frequency components enhances
-    sample quality (such as better visual details). Therefore, they recommend using low guidance scales (low w_low) for
-    the low-frequency components and high guidance scales (high w_high) for the high-frequency components. As an
-    example, they suggest w_low = 5.0 and w_high = 10.0 for Stable Diffusion XL (see Table 8 in the paper).
-
-    As with CFG, Diffusers implements the scaling and shifting on the unconditional prediction based on the [Imagen
-    paper](https://huggingface.co/papers/2205.11487), which is equivalent to what the original CFG paper proposed in
-    theory. [x_pred = x_uncond + scale * (x_cond - x_uncond)]
-
-    The `use_original_formulation` argument can be set to `True` to use the original CFG formulation mentioned in the
-    paper. By default, we use the diffusers-native implementation that has been in the codebase for a long time.
-
-    Args:
-        guidance_scales (`list[float]`, defaults to `[10.0, 5.0]`):
-            The scale parameter for frequency-decoupled guidance for each frequency component, listed from highest
-            frequency level to lowest. Higher values result in stronger conditioning on the text prompt, while lower
-            values allow for more freedom in generation. Higher values may lead to saturation and deterioration of
-            image quality. The FDG authors recommend using higher guidance scales for higher frequency components and
-            lower guidance scales for lower frequency components (so `guidance_scales` should typically be sorted in
-            descending order).
-        guidance_rescale (`float` or `list[float]`, defaults to `0.0`):
-            The rescale factor applied to the noise predictions. This is used to improve image quality and fix
-            overexposure. Based on Section 3.4 from [Common Diffusion Noise Schedules and Sample Steps are
-            Flawed](https://huggingface.co/papers/2305.08891). If a list is supplied, it should be the same length as
-            `guidance_scales`.
-        parallel_weights (`float` or `list[float]`, *optional*):
-            Optional weights for the parallel component of each frequency component of the projected CFG shift. If not
-            set, the weights will default to `1.0` for all components, which corresponds to using the normal CFG shift
-            (that is, equal weights for the parallel and orthogonal components). If set, a value in `[0, 1]` is
-            recommended. If a list is supplied, it should be the same length as `guidance_scales`.
-        use_original_formulation (`bool`, defaults to `False`):
-            Whether to use the original formulation of classifier-free guidance as proposed in the paper. By default,
-            we use the diffusers-native implementation that has been in the codebase for a long time. See
-            [~guiders.classifier_free_guidance.ClassifierFreeGuidance] for more details.
-        start (`float` or `list[float]`, defaults to `0.0`):
-            The fraction of the total number of denoising steps after which guidance starts. If a list is supplied, it
-            should be the same length as `guidance_scales`.
-        stop (`float` or `list[float]`, defaults to `1.0`):
-            The fraction of the total number of denoising steps after which guidance stops. If a list is supplied, it
-            should be the same length as `guidance_scales`.
-        guidance_rescale_space (`str`, defaults to `"data"`):
-            Whether to performance guidance rescaling in `"data"` space (after the full FDG update in data space) or in
-            `"freq"` space (right after the CFG update, for each freq level). Note that frequency space rescaling is
-            speculative and may not produce expected results. If `"data"` is set, the first `guidance_rescale` value
-            will be used; otherwise, per-frequency-level guidance rescale values will be used if available.
-        upcast_to_double (`bool`, defaults to `True`):
-            Whether to upcast certain operations, such as the projection operation when using `parallel_weights`, to
-            float64 when performing guidance. This may result in better performance at the cost of increased runtime.
-    """
-
-    _input_predictions = ["pred_cond", "pred_uncond"]
-
-    @register_to_config
-    def __init__(
-        self,
-        guidance_scales: list[float] | tuple[float] = [10.0, 5.0],
-        guidance_rescale: float | list[float] | tuple[float] = 0.0,
-        parallel_weights: float | list[float] | tuple[float] | None = None,
-        use_original_formulation: bool = False,
-        start: float | list[float] | tuple[float] = 0.0,
-        stop: float | list[float] | tuple[float] = 1.0,
-        guidance_rescale_space: str = "data",
-        upcast_to_double: bool = True,
-        enabled: bool = True,
-    ):
-        if not _CAN_USE_KORNIA:
-            raise ImportError(
-                "The `FrequencyDecoupledGuidance` guider cannot be instantiated because the `kornia` library on which "
-                "it depends is not available in the current environment. You can install `kornia` with `pip install "
-                "kornia`."
-            )
-
-        # Set start to earliest start for any freq component and stop to latest stop for any freq component
-        min_start = start if isinstance(start, float) else min(start)
-        max_stop = stop if isinstance(stop, float) else max(stop)
-        super().__init__(min_start, max_stop, enabled)
-
-        self.guidance_scales = guidance_scales
-        self.levels = len(guidance_scales)
-
-        if isinstance(guidance_rescale, float):
-            self.guidance_rescale = [guidance_rescale] * self.levels
-        elif len(guidance_rescale) == self.levels:
-            self.guidance_rescale = guidance_rescale
-        else:
-            raise ValueError(
-                f"`guidance_rescale` has length {len(guidance_rescale)} but should have the same length as "
-                f"`guidance_scales` ({len(self.guidance_scales)})"
-            )
-        # Whether to perform guidance rescaling in frequency space (right after the CFG update) or data space (after
-        # transforming from frequency space back to data space)
-        if guidance_rescale_space not in ["data", "freq"]:
-            raise ValueError(
-                f"Guidance rescale space is {guidance_rescale_space} but must be one of `data` or `freq`."
-            )
-        self.guidance_rescale_space = guidance_rescale_space
-
-        if parallel_weights is None:
-            # Use normal CFG shift (equal weights for parallel and orthogonal components)
-            self.parallel_weights = [1.0] * self.levels
-        elif isinstance(parallel_weights, float):
-            self.parallel_weights = [parallel_weights] * self.levels
-        elif len(parallel_weights) == self.levels:
-            self.parallel_weights = parallel_weights
-        else:
-            raise ValueError(
-                f"`parallel_weights` has length {len(parallel_weights)} but should have the same length as "
-                f"`guidance_scales` ({len(self.guidance_scales)})"
-            )
-
-        self.use_original_formulation = use_original_formulation
-        self.upcast_to_double = upcast_to_double
-
-        if isinstance(start, float):
-            self.guidance_start = [start] * self.levels
-        elif len(start) == self.levels:
-            self.guidance_start = start
-        else:
-            raise ValueError(
-                f"`start` has length {len(start)} but should have the same length as `guidance_scales` "
-                f"({len(self.guidance_scales)})"
-            )
-        if isinstance(stop, float):
-            self.guidance_stop = [stop] * self.levels
-        elif len(stop) == self.levels:
-            self.guidance_stop = stop
-        else:
-            raise ValueError(
-                f"`stop` has length {len(stop)} but should have the same length as `guidance_scales` "
-                f"({len(self.guidance_scales)})"
-            )
-
-    def prepare_inputs(self, data: dict[str, tuple[torch.Tensor, torch.Tensor]]) -> list["BlockState"]:
-        tuple_indices = [0] if self.num_conditions == 1 else [0, 1]
-        data_batches = []
-        for tuple_idx, input_prediction in zip(tuple_indices, self._input_predictions):
-            data_batch = self._prepare_batch(data, tuple_idx, input_prediction)
-            data_batches.append(data_batch)
-        return data_batches
-
-    def prepare_inputs_from_block_state(
-        self, data: "BlockState", input_fields: dict[str, str | tuple[str, str]]
-    ) -> list["BlockState"]:
-        tuple_indices = [0] if self.num_conditions == 1 else [0, 1]
-        data_batches = []
-        for tuple_idx, input_prediction in zip(tuple_indices, self._input_predictions):
-            data_batch = self._prepare_batch_from_block_state(input_fields, data, tuple_idx, input_prediction)
-            data_batches.append(data_batch)
-        return data_batches
-
-    def forward(self, pred_cond: torch.Tensor, pred_uncond: torch.Tensor | None = None) -> GuiderOutput:
-        pred = None
-
-        if not self._is_fdg_enabled():
-            pred = pred_cond
-        else:
-            # Apply the frequency transform (e.g. Laplacian pyramid) to the conditional and unconditional predictions.
-            pred_cond_pyramid = build_laplacian_pyramid_func(pred_cond, self.levels)
-            pred_uncond_pyramid = build_laplacian_pyramid_func(pred_uncond, self.levels)
-
-            # From high frequencies to low frequencies, following the paper implementation
-            pred_guided_pyramid = []
-            parameters = zip(self.guidance_scales, self.parallel_weights, self.guidance_rescale)
-            for level, (guidance_scale, parallel_weight, guidance_rescale) in enumerate(parameters):
-                if self._is_fdg_enabled_for_level(level):
-                    # Get the cond/uncond preds (in freq space) at the current frequency level
-                    pred_cond_freq = pred_cond_pyramid[level]
-                    pred_uncond_freq = pred_uncond_pyramid[level]
-
-                    shift = pred_cond_freq - pred_uncond_freq
-
-                    # Apply parallel weights, if used (1.0 corresponds to using the normal CFG shift)
-                    if not math.isclose(parallel_weight, 1.0):
-                        shift_parallel, shift_orthogonal = project(shift, pred_cond_freq, self.upcast_to_double)
-                        shift = parallel_weight * shift_parallel + shift_orthogonal
-
-                    # Apply CFG update for the current frequency level
-                    pred = pred_cond_freq if self.use_original_formulation else pred_uncond_freq
-                    pred = pred + guidance_scale * shift
-
-                    if self.guidance_rescale_space == "freq" and guidance_rescale > 0.0:
-                        pred = rescale_noise_cfg(pred, pred_cond_freq, guidance_rescale)
-
-                    # Add the current FDG guided level to the FDG prediction pyramid
-                    pred_guided_pyramid.append(pred)
-                else:
-                    # Add the current pred_cond_pyramid level as the "non-FDG" prediction
-                    pred_guided_pyramid.append(pred_cond_freq)
-
-            # Convert from frequency space back to data (e.g. pixel) space by applying inverse freq transform
-            pred = build_image_from_pyramid(pred_guided_pyramid)
-
-            # If rescaling in data space, use the first elem of self.guidance_rescale as the "global" rescale value
-            # across all freq levels
-            if self.guidance_rescale_space == "data" and self.guidance_rescale[0] > 0.0:
-                pred = rescale_noise_cfg(pred, pred_cond, self.guidance_rescale[0])
-
-        return GuiderOutput(pred=pred, pred_cond=pred_cond, pred_uncond=pred_uncond)
-
-    @property
-    def is_conditional(self) -> bool:
-        return self._count_prepared == 1
-
-    @property
-    def num_conditions(self) -> int:
-        num_conditions = 1
-        if self._is_fdg_enabled():
-            num_conditions += 1
-        return num_conditions
-
-    def _is_fdg_enabled(self) -> bool:
-        if not self._enabled:
-            return False
-
-        is_within_range = True
-        if self._num_inference_steps is not None:
-            skip_start_step = int(self._start * self._num_inference_steps)
-            skip_stop_step = int(self._stop * self._num_inference_steps)
-            is_within_range = skip_start_step <= self._step < skip_stop_step
-
-        is_close = False
-        if self.use_original_formulation:
-            is_close = all(math.isclose(guidance_scale, 0.0) for guidance_scale in self.guidance_scales)
-        else:
-            is_close = all(math.isclose(guidance_scale, 1.0) for guidance_scale in self.guidance_scales)
-
-        return is_within_range and not is_close
-
-    def _is_fdg_enabled_for_level(self, level: int) -> bool:
-        if not self._enabled:
-            return False
-
-        is_within_range = True
-        if self._num_inference_steps is not None:
-            skip_start_step = int(self.guidance_start[level] * self._num_inference_steps)
-            skip_stop_step = int(self.guidance_stop[level] * self._num_inference_steps)
-            is_within_range = skip_start_step <= self._step < skip_stop_step
-
-        is_close = False
-        if self.use_original_formulation:
-            is_close = math.isclose(self.guidance_scales[level], 0.0)
-        else:
-            is_close = math.isclose(self.guidance_scales[level], 1.0)
-
-        return is_within_range and not is_close

diffusers/guiders/guider_utils.py

@@ -1,396 +0,0 @@
-# Copyright 2026 The HuggingFace Team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-from __future__ import annotations
-
-import os
-from typing import TYPE_CHECKING, Any
-
-import torch
-from huggingface_hub.utils import validate_hf_hub_args
-from typing_extensions import Self
-
-from ..configuration_utils import ConfigMixin
-from ..utils import BaseOutput, PushToHubMixin, get_logger
-
-
-if TYPE_CHECKING:
-    from ..modular_pipelines.modular_pipeline import BlockState
-
-
-GUIDER_CONFIG_NAME = "guider_config.json"
-
-
-logger = get_logger(__name__)  # pylint: disable=invalid-name
-
-
-class BaseGuidance(ConfigMixin, PushToHubMixin):
-    r"""Base class providing the skeleton for implementing guidance techniques."""
-
-    config_name = GUIDER_CONFIG_NAME
-    _input_predictions = None
-    _identifier_key = "__guidance_identifier__"
-
-    def __init__(self, start: float = 0.0, stop: float = 1.0, enabled: bool = True):
-        logger.warning(
-            "Guiders are currently an experimental feature under active development. The API is subject to breaking changes in future releases."
-        )
-
-        self._start = start
-        self._stop = stop
-        self._step: int = None
-        self._num_inference_steps: int = None
-        self._timestep: torch.LongTensor = None
-        self._count_prepared = 0
-        self._input_fields: dict[str, str | tuple[str, str]] = None
-        self._enabled = enabled
-
-        if not (0.0 <= start < 1.0):
-            raise ValueError(f"Expected `start` to be between 0.0 and 1.0, but got {start}.")
-        if not (start <= stop <= 1.0):
-            raise ValueError(f"Expected `stop` to be between {start} and 1.0, but got {stop}.")
-
-        if self._input_predictions is None or not isinstance(self._input_predictions, list):
-            raise ValueError(
-                "`_input_predictions` must be a list of required prediction names for the guidance technique."
-            )
-
-    def new(self, **kwargs):
-        """
-        Creates a copy of this guider instance, optionally with modified configuration parameters.
-
-        Args:
-            **kwargs: Configuration parameters to override in the new instance. If no kwargs are provided,
-                returns an exact copy with the same configuration.
-
-        Returns:
-            A new guider instance with the same (or updated) configuration.
-
-        Example:
-            ```python
-            # Create a CFG guider
-            guider = ClassifierFreeGuidance(guidance_scale=3.5)
-
-            # Create an exact copy
-            same_guider = guider.new()
-
-            # Create a copy with different start step, keeping other config the same
-            new_guider = guider.new(guidance_scale=5)
-            ```
-        """
-        return self.__class__.from_config(self.config, **kwargs)
-
-    def disable(self):
-        self._enabled = False
-
-    def enable(self):
-        self._enabled = True
-
-    def set_state(self, step: int, num_inference_steps: int, timestep: torch.LongTensor) -> None:
-        self._step = step
-        self._num_inference_steps = num_inference_steps
-        self._timestep = timestep
-        self._count_prepared = 0
-
-    def get_state(self) -> dict[str, Any]:
-        """
-        Returns the current state of the guidance technique as a dictionary. The state variables will be included in
-        the __repr__ method. Returns:
-            `dict[str, Any]`: A dictionary containing the current state variables including:
-                - step: Current inference step
-                - num_inference_steps: Total number of inference steps
-                - timestep: Current timestep tensor
-                - count_prepared: Number of times prepare_models has been called
-                - enabled: Whether the guidance is enabled
-                - num_conditions: Number of conditions
-        """
-        state = {
-            "step": self._step,
-            "num_inference_steps": self._num_inference_steps,
-            "timestep": self._timestep,
-            "count_prepared": self._count_prepared,
-            "enabled": self._enabled,
-            "num_conditions": self.num_conditions,
-        }
-        return state
-
-    def __repr__(self) -> str:
-        """
-        Returns a string representation of the guidance object including both config and current state.
-        """
-        # Get ConfigMixin's __repr__
-        str_repr = super().__repr__()
-
-        # Get current state
-        state = self.get_state()
-
-        # Format each state variable on its own line with indentation
-        state_lines = []
-        for k, v in state.items():
-            # Convert value to string and handle multi-line values
-            v_str = str(v)
-            if "\n" in v_str:
-                # For multi-line values (like MomentumBuffer), indent subsequent lines
-                v_lines = v_str.split("\n")
-                v_str = v_lines[0] + "\n" + "\n".join(["    " + line for line in v_lines[1:]])
-            state_lines.append(f"  {k}: {v_str}")
-
-        state_str = "\n".join(state_lines)
-
-        return f"{str_repr}\nState:\n{state_str}"
-
-    def prepare_models(self, denoiser: torch.nn.Module) -> None:
-        """
-        Prepares the models for the guidance technique on a given batch of data. This method should be overridden in
-        subclasses to implement specific model preparation logic.
-        """
-        self._count_prepared += 1
-
-    def cleanup_models(self, denoiser: torch.nn.Module) -> None:
-        """
-        Cleans up the models for the guidance technique after a given batch of data. This method should be overridden
-        in subclasses to implement specific model cleanup logic. It is useful for removing any hooks or other stateful
-        modifications made during `prepare_models`.
-        """
-        pass
-
-    def prepare_inputs(self, data: "BlockState") -> list["BlockState"]:
-        raise NotImplementedError("BaseGuidance::prepare_inputs must be implemented in subclasses.")
-
-    def prepare_inputs_from_block_state(
-        self, data: "BlockState", input_fields: dict[str, str | tuple[str, str]]
-    ) -> list["BlockState"]:
-        raise NotImplementedError("BaseGuidance::prepare_inputs_from_block_state must be implemented in subclasses.")
-
-    def __call__(self, data: list["BlockState"]) -> Any:
-        if not all(hasattr(d, "noise_pred") for d in data):
-            raise ValueError("Expected all data to have `noise_pred` attribute.")
-        if len(data) != self.num_conditions:
-            raise ValueError(
-                f"Expected {self.num_conditions} data items, but got {len(data)}. Please check the input data."
-            )
-        forward_inputs = {getattr(d, self._identifier_key): d.noise_pred for d in data}
-        return self.forward(**forward_inputs)
-
-    def forward(self, *args, **kwargs) -> Any:
-        raise NotImplementedError("BaseGuidance::forward must be implemented in subclasses.")
-
-    @property
-    def is_conditional(self) -> bool:
-        raise NotImplementedError("BaseGuidance::is_conditional must be implemented in subclasses.")
-
-    @property
-    def is_unconditional(self) -> bool:
-        return not self.is_conditional
-
-    @property
-    def num_conditions(self) -> int:
-        raise NotImplementedError("BaseGuidance::num_conditions must be implemented in subclasses.")
-
-    @classmethod
-    def _prepare_batch(
-        cls,
-        data: dict[str, tuple[torch.Tensor, torch.Tensor]],
-        tuple_index: int,
-        identifier: str,
-    ) -> "BlockState":
-        """
-        Prepares a batch of data for the guidance technique. This method is used in the `prepare_inputs` method of the
-        `BaseGuidance` class. It prepares the batch based on the provided tuple index.
-
-        Args:
-            input_fields (`dict[str, str | tuple[str, str]]`):
-                A dictionary where the keys are the names of the fields that will be used to store the data once it is
-                prepared with `prepare_inputs`. The values can be either a string or a tuple of length 2, which is used
-                to look up the required data provided for preparation. If a string is provided, it will be used as the
-                conditional data (or unconditional if used with a guidance method that requires it). If a tuple of
-                length 2 is provided, the first element must be the conditional data identifier and the second element
-                must be the unconditional data identifier or None.
-            data (`BlockState`):
-                The input data to be prepared.
-            tuple_index (`int`):
-                The index to use when accessing input fields that are tuples.
-
-        Returns:
-            `BlockState`: The prepared batch of data.
-        """
-        from ..modular_pipelines.modular_pipeline import BlockState
-
-        data_batch = {}
-        for key, value in data.items():
-            try:
-                if isinstance(value, torch.Tensor):
-                    data_batch[key] = value
-                elif isinstance(value, tuple):
-                    data_batch[key] = value[tuple_index]
-                else:
-                    raise ValueError(f"Invalid value type: {type(value)}")
-            except ValueError:
-                logger.debug(f"`data` does not have attribute(s) {value}, skipping.")
-        data_batch[cls._identifier_key] = identifier
-        return BlockState(**data_batch)
-
-    @classmethod
-    def _prepare_batch_from_block_state(
-        cls,
-        input_fields: dict[str, str | tuple[str, str]],
-        data: "BlockState",
-        tuple_index: int,
-        identifier: str,
-    ) -> "BlockState":
-        """
-        Prepares a batch of data for the guidance technique. This method is used in the `prepare_inputs` method of the
-        `BaseGuidance` class. It prepares the batch based on the provided tuple index.
-
-        Args:
-            input_fields (`dict[str, str | tuple[str, str]]`):
-                A dictionary where the keys are the names of the fields that will be used to store the data once it is
-                prepared with `prepare_inputs`. The values can be either a string or a tuple of length 2, which is used
-                to look up the required data provided for preparation. If a string is provided, it will be used as the
-                conditional data (or unconditional if used with a guidance method that requires it). If a tuple of
-                length 2 is provided, the first element must be the conditional data identifier and the second element
-                must be the unconditional data identifier or None.
-            data (`BlockState`):
-                The input data to be prepared.
-            tuple_index (`int`):
-                The index to use when accessing input fields that are tuples.
-
-        Returns:
-            `BlockState`: The prepared batch of data.
-        """
-        from ..modular_pipelines.modular_pipeline import BlockState
-
-        data_batch = {}
-        for key, value in input_fields.items():
-            try:
-                if isinstance(value, str):
-                    data_batch[key] = getattr(data, value)
-                elif isinstance(value, tuple):
-                    data_batch[key] = getattr(data, value[tuple_index])
-                else:
-                    # We've already checked that value is a string or a tuple of strings with length 2
-                    pass
-            except AttributeError:
-                logger.debug(f"`data` does not have attribute(s) {value}, skipping.")
-        data_batch[cls._identifier_key] = identifier
-        return BlockState(**data_batch)
-
-    @classmethod
-    @validate_hf_hub_args
-    def from_pretrained(
-        cls,
-        pretrained_model_name_or_path: str | os.PathLike | None = None,
-        subfolder: str | None = None,
-        return_unused_kwargs=False,
-        **kwargs,
-    ) -> Self:
-        r"""
-        Instantiate a guider from a pre-defined JSON configuration file in a local directory or Hub repository.
-
-        Parameters:
-            pretrained_model_name_or_path (`str` or `os.PathLike`, *optional*):
-                Can be either:
-
-                    - A string, the *model id* (for example `google/ddpm-celebahq-256`) of a pretrained model hosted on
-                      the Hub.
-                    - A path to a *directory* (for example `./my_model_directory`) containing the guider configuration
-                      saved with [`~BaseGuidance.save_pretrained`].
-            subfolder (`str`, *optional*):
-                The subfolder location of a model file within a larger model repository on the Hub or locally.
-            return_unused_kwargs (`bool`, *optional*, defaults to `False`):
-                Whether kwargs that are not consumed by the Python class should be returned or not.
-            cache_dir (`str | os.PathLike`, *optional*):
-                Path to a directory where a downloaded pretrained model configuration is cached if the standard cache
-                is not used.
-            force_download (`bool`, *optional*, defaults to `False`):
-                Whether or not to force the (re-)download of the model weights and configuration files, overriding the
-                cached versions if they exist.
-
-            proxies (`dict[str, str]`, *optional*):
-                A dictionary of proxy servers to use by protocol or endpoint, for example, `{'http': 'foo.bar:3128',
-                'http://hostname': 'foo.bar:4012'}`. The proxies are used on each request.
-            output_loading_info(`bool`, *optional*, defaults to `False`):
-                Whether or not to also return a dictionary containing missing keys, unexpected keys and error messages.
-            local_files_only(`bool`, *optional*, defaults to `False`):
-                Whether to only load local model weights and configuration files or not. If set to `True`, the model
-                won't be downloaded from the Hub.
-            token (`str` or *bool*, *optional*):
-                The token to use as HTTP bearer authorization for remote files. If `True`, the token generated from
-                `diffusers-cli login` (stored in `~/.huggingface`) is used.
-            revision (`str`, *optional*, defaults to `"main"`):
-                The specific model version to use. It can be a branch name, a tag name, a commit id, or any identifier
-                allowed by Git.
-
-        > [!TIP] > To use private or [gated models](https://huggingface.co/docs/hub/models-gated#gated-models), log-in
-        with `hf > auth login`. You can also activate the special >
-        ["offline-mode"](https://huggingface.co/diffusers/installation.html#offline-mode) to use this method in a >
-        firewalled environment.
-
-        """
-        config, kwargs, commit_hash = cls.load_config(
-            pretrained_model_name_or_path=pretrained_model_name_or_path,
-            subfolder=subfolder,
-            return_unused_kwargs=True,
-            return_commit_hash=True,
-            **kwargs,
-        )
-        return cls.from_config(config, return_unused_kwargs=return_unused_kwargs, **kwargs)
-
-    def save_pretrained(self, save_directory: str | os.PathLike, push_to_hub: bool = False, **kwargs):
-        """
-        Save a guider configuration object to a directory so that it can be reloaded using the
-        [`~BaseGuidance.from_pretrained`] class method.
-
-        Args:
-            save_directory (`str` or `os.PathLike`):
-                Directory where the configuration JSON file will be saved (will be created if it does not exist).
-            push_to_hub (`bool`, *optional*, defaults to `False`):
-                Whether or not to push your model to the Hugging Face Hub after saving it. You can specify the
-                repository you want to push to with `repo_id` (will default to the name of `save_directory` in your
-                namespace).
-            kwargs (`dict[str, Any]`, *optional*):
-                Additional keyword arguments passed along to the [`~utils.PushToHubMixin.push_to_hub`] method.
-        """
-        self.save_config(save_directory=save_directory, push_to_hub=push_to_hub, **kwargs)
-
-
-class GuiderOutput(BaseOutput):
-    pred: torch.Tensor
-    pred_cond: torch.Tensor | None
-    pred_uncond: torch.Tensor | None
-
-
-def rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.0):
-    r"""
-    Rescales `noise_cfg` tensor based on `guidance_rescale` to improve image quality and fix overexposure. Based on
-    Section 3.4 from [Common Diffusion Noise Schedules and Sample Steps are
-    Flawed](https://huggingface.co/papers/2305.08891).
-
-    Args:
-        noise_cfg (`torch.Tensor`):
-            The predicted noise tensor for the guided diffusion process.
-        noise_pred_text (`torch.Tensor`):
-            The predicted noise tensor for the text-guided diffusion process.
-        guidance_rescale (`float`, *optional*, defaults to 0.0):
-            A rescale factor applied to the noise predictions.
-    Returns:
-        noise_cfg (`torch.Tensor`): The rescaled noise prediction tensor.
-    """
-    std_text = noise_pred_text.std(dim=list(range(1, noise_pred_text.ndim)), keepdim=True)
-    std_cfg = noise_cfg.std(dim=list(range(1, noise_cfg.ndim)), keepdim=True)
-    # rescale the results from guidance (fixes overexposure)
-    noise_pred_rescaled = noise_cfg * (std_text / std_cfg)
-    # mix with the original results from guidance by factor guidance_rescale to avoid "plain looking" images
-    noise_cfg = guidance_rescale * noise_pred_rescaled + (1 - guidance_rescale) * noise_cfg
-    return noise_cfg

diffusers/guiders/magnitude_aware_guidance.py

@@ -1,159 +0,0 @@
-# Copyright 2026 The HuggingFace Team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-import math
-from typing import TYPE_CHECKING
-
-import torch
-
-from ..configuration_utils import register_to_config
-from .guider_utils import BaseGuidance, GuiderOutput, rescale_noise_cfg
-
-
-if TYPE_CHECKING:
-    from ..modular_pipelines.modular_pipeline import BlockState
-
-
-class MagnitudeAwareGuidance(BaseGuidance):
-    """
-    Magnitude-Aware Mitigation for Boosted Guidance (MAMBO-G): https://huggingface.co/papers/2508.03442
-
-    Args:
-        guidance_scale (`float`, defaults to `10.0`):
-            The scale parameter for classifier-free guidance. Higher values result in stronger conditioning on the text
-            prompt, while lower values allow for more freedom in generation. Higher values may lead to saturation and
-            deterioration of image quality.
-        alpha (`float`, defaults to `8.0`):
-            The alpha parameter for the magnitude-aware guidance. Higher values cause more aggressive supression of
-            guidance scale when the magnitude of the guidance update is large.
-        guidance_rescale (`float`, defaults to `0.0`):
-            The rescale factor applied to the noise predictions. This is used to improve image quality and fix
-            overexposure. Based on Section 3.4 from [Common Diffusion Noise Schedules and Sample Steps are
-            Flawed](https://huggingface.co/papers/2305.08891).
-        use_original_formulation (`bool`, defaults to `False`):
-            Whether to use the original formulation of classifier-free guidance as proposed in the paper. By default,
-            we use the diffusers-native implementation that has been in the codebase for a long time. See
-            [~guiders.classifier_free_guidance.ClassifierFreeGuidance] for more details.
-        start (`float`, defaults to `0.0`):
-            The fraction of the total number of denoising steps after which guidance starts.
-        stop (`float`, defaults to `1.0`):
-            The fraction of the total number of denoising steps after which guidance stops.
-    """
-
-    _input_predictions = ["pred_cond", "pred_uncond"]
-
-    @register_to_config
-    def __init__(
-        self,
-        guidance_scale: float = 10.0,
-        alpha: float = 8.0,
-        guidance_rescale: float = 0.0,
-        use_original_formulation: bool = False,
-        start: float = 0.0,
-        stop: float = 1.0,
-        enabled: bool = True,
-    ):
-        super().__init__(start, stop, enabled)
-
-        self.guidance_scale = guidance_scale
-        self.alpha = alpha
-        self.guidance_rescale = guidance_rescale
-        self.use_original_formulation = use_original_formulation
-
-    def prepare_inputs(self, data: dict[str, tuple[torch.Tensor, torch.Tensor]]) -> list["BlockState"]:
-        tuple_indices = [0] if self.num_conditions == 1 else [0, 1]
-        data_batches = []
-        for tuple_idx, input_prediction in zip(tuple_indices, self._input_predictions):
-            data_batch = self._prepare_batch(data, tuple_idx, input_prediction)
-            data_batches.append(data_batch)
-        return data_batches
-
-    def prepare_inputs_from_block_state(
-        self, data: "BlockState", input_fields: dict[str, str | tuple[str, str]]
-    ) -> list["BlockState"]:
-        tuple_indices = [0] if self.num_conditions == 1 else [0, 1]
-        data_batches = []
-        for tuple_idx, input_prediction in zip(tuple_indices, self._input_predictions):
-            data_batch = self._prepare_batch_from_block_state(input_fields, data, tuple_idx, input_prediction)
-            data_batches.append(data_batch)
-        return data_batches
-
-    def forward(self, pred_cond: torch.Tensor, pred_uncond: torch.Tensor | None = None) -> GuiderOutput:
-        pred = None
-
-        if not self._is_mambo_g_enabled():
-            pred = pred_cond
-        else:
-            pred = mambo_guidance(
-                pred_cond,
-                pred_uncond,
-                self.guidance_scale,
-                self.alpha,
-                self.use_original_formulation,
-            )
-
-        if self.guidance_rescale > 0.0:
-            pred = rescale_noise_cfg(pred, pred_cond, self.guidance_rescale)
-
-        return GuiderOutput(pred=pred, pred_cond=pred_cond, pred_uncond=pred_uncond)
-
-    @property
-    def is_conditional(self) -> bool:
-        return self._count_prepared == 1
-
-    @property
-    def num_conditions(self) -> int:
-        num_conditions = 1
-        if self._is_mambo_g_enabled():
-            num_conditions += 1
-        return num_conditions
-
-    def _is_mambo_g_enabled(self) -> bool:
-        if not self._enabled:
-            return False
-
-        is_within_range = True
-        if self._num_inference_steps is not None:
-            skip_start_step = int(self._start * self._num_inference_steps)
-            skip_stop_step = int(self._stop * self._num_inference_steps)
-            is_within_range = skip_start_step <= self._step < skip_stop_step
-
-        is_close = False
-        if self.use_original_formulation:
-            is_close = math.isclose(self.guidance_scale, 0.0)
-        else:
-            is_close = math.isclose(self.guidance_scale, 1.0)
-
-        return is_within_range and not is_close
-
-
-def mambo_guidance(
-    pred_cond: torch.Tensor,
-    pred_uncond: torch.Tensor,
-    guidance_scale: float,
-    alpha: float = 8.0,
-    use_original_formulation: bool = False,
-):
-    dim = list(range(1, len(pred_cond.shape)))
-    diff = pred_cond - pred_uncond
-    ratio = torch.norm(diff, dim=dim, keepdim=True) / torch.norm(pred_uncond, dim=dim, keepdim=True)
-    guidance_scale_final = (
-        guidance_scale * torch.exp(-alpha * ratio)
-        if use_original_formulation
-        else 1.0 + (guidance_scale - 1.0) * torch.exp(-alpha * ratio)
-    )
-    pred = pred_cond if use_original_formulation else pred_uncond
-    pred = pred + guidance_scale_final * diff
-
-    return pred

diffusers/guiders/perturbed_attention_guidance.py

@@ -1,289 +0,0 @@
-# Copyright 2026 The HuggingFace Team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-from __future__ import annotations
-
-import math
-from typing import TYPE_CHECKING, Any
-
-import torch
-
-from ..configuration_utils import register_to_config
-from ..hooks import HookRegistry, LayerSkipConfig
-from ..hooks.layer_skip import _apply_layer_skip_hook
-from ..utils import get_logger
-from .guider_utils import BaseGuidance, GuiderOutput, rescale_noise_cfg
-
-
-if TYPE_CHECKING:
-    from ..modular_pipelines.modular_pipeline import BlockState
-
-
-logger = get_logger(__name__)  # pylint: disable=invalid-name
-
-
-class PerturbedAttentionGuidance(BaseGuidance):
-    """
-    Perturbed Attention Guidance (PAG): https://huggingface.co/papers/2403.17377
-
-    The intution behind PAG can be thought of as moving the CFG predicted distribution estimates further away from
-    worse versions of the conditional distribution estimates. PAG was one of the first techniques to introduce the idea
-    of using a worse version of the trained model for better guiding itself in the denoising process. It perturbs the
-    attention scores of the latent stream by replacing the score matrix with an identity matrix for selectively chosen
-    layers.
-
-    Additional reading:
-    - [Guiding a Diffusion Model with a Bad Version of Itself](https://huggingface.co/papers/2406.02507)
-
-    PAG is implemented with similar implementation to SkipLayerGuidance due to overlap in the configuration parameters
-    and implementation details.
-
-    Args:
-        guidance_scale (`float`, defaults to `7.5`):
-            The scale parameter for classifier-free guidance. Higher values result in stronger conditioning on the text
-            prompt, while lower values allow for more freedom in generation. Higher values may lead to saturation and
-            deterioration of image quality.
-        perturbed_guidance_scale (`float`, defaults to `2.8`):
-            The scale parameter for perturbed attention guidance.
-        perturbed_guidance_start (`float`, defaults to `0.01`):
-            The fraction of the total number of denoising steps after which perturbed attention guidance starts.
-        perturbed_guidance_stop (`float`, defaults to `0.2`):
-            The fraction of the total number of denoising steps after which perturbed attention guidance stops.
-        perturbed_guidance_layers (`int` or `list[int]`, *optional*):
-            The layer indices to apply perturbed attention guidance to. Can be a single integer or a list of integers.
-            If not provided, `perturbed_guidance_config` must be provided.
-        perturbed_guidance_config (`LayerSkipConfig` or `list[LayerSkipConfig]`, *optional*):
-            The configuration for the perturbed attention guidance. Can be a single `LayerSkipConfig` or a list of
-            `LayerSkipConfig`. If not provided, `perturbed_guidance_layers` must be provided.
-        guidance_rescale (`float`, defaults to `0.0`):
-            The rescale factor applied to the noise predictions. This is used to improve image quality and fix
-            overexposure. Based on Section 3.4 from [Common Diffusion Noise Schedules and Sample Steps are
-            Flawed](https://huggingface.co/papers/2305.08891).
-        use_original_formulation (`bool`, defaults to `False`):
-            Whether to use the original formulation of classifier-free guidance as proposed in the paper. By default,
-            we use the diffusers-native implementation that has been in the codebase for a long time. See
-            [~guiders.classifier_free_guidance.ClassifierFreeGuidance] for more details.
-        start (`float`, defaults to `0.01`):
-            The fraction of the total number of denoising steps after which guidance starts.
-        stop (`float`, defaults to `0.2`):
-            The fraction of the total number of denoising steps after which guidance stops.
-    """
-
-    # NOTE: The current implementation does not account for joint latent conditioning (text + image/video tokens in
-    # the same latent stream). It assumes the entire latent is a single stream of visual tokens. It would be very
-    # complex to support joint latent conditioning in a model-agnostic manner without specializing the implementation
-    # for each model architecture.
-
-    _input_predictions = ["pred_cond", "pred_uncond", "pred_cond_skip"]
-
-    @register_to_config
-    def __init__(
-        self,
-        guidance_scale: float = 7.5,
-        perturbed_guidance_scale: float = 2.8,
-        perturbed_guidance_start: float = 0.01,
-        perturbed_guidance_stop: float = 0.2,
-        perturbed_guidance_layers: int | list[int] | None = None,
-        perturbed_guidance_config: LayerSkipConfig | list[LayerSkipConfig] | dict[str, Any] = None,
-        guidance_rescale: float = 0.0,
-        use_original_formulation: bool = False,
-        start: float = 0.0,
-        stop: float = 1.0,
-        enabled: bool = True,
-    ):
-        super().__init__(start, stop, enabled)
-
-        self.guidance_scale = guidance_scale
-        self.skip_layer_guidance_scale = perturbed_guidance_scale
-        self.skip_layer_guidance_start = perturbed_guidance_start
-        self.skip_layer_guidance_stop = perturbed_guidance_stop
-        self.guidance_rescale = guidance_rescale
-        self.use_original_formulation = use_original_formulation
-
-        if perturbed_guidance_config is None:
-            if perturbed_guidance_layers is None:
-                raise ValueError(
-                    "`perturbed_guidance_layers` must be provided if `perturbed_guidance_config` is not specified."
-                )
-            perturbed_guidance_config = LayerSkipConfig(
-                indices=perturbed_guidance_layers,
-                fqn="auto",
-                skip_attention=False,
-                skip_attention_scores=True,
-                skip_ff=False,
-            )
-        else:
-            if perturbed_guidance_layers is not None:
-                raise ValueError(
-                    "`perturbed_guidance_layers` should not be provided if `perturbed_guidance_config` is specified."
-                )
-
-        if isinstance(perturbed_guidance_config, dict):
-            perturbed_guidance_config = LayerSkipConfig.from_dict(perturbed_guidance_config)
-
-        if isinstance(perturbed_guidance_config, LayerSkipConfig):
-            perturbed_guidance_config = [perturbed_guidance_config]
-
-        if not isinstance(perturbed_guidance_config, list):
-            raise ValueError(
-                "`perturbed_guidance_config` must be a `LayerSkipConfig`, a list of `LayerSkipConfig`, or a dict that can be converted to a `LayerSkipConfig`."
-            )
-        elif isinstance(next(iter(perturbed_guidance_config), None), dict):
-            perturbed_guidance_config = [LayerSkipConfig.from_dict(config) for config in perturbed_guidance_config]
-
-        for config in perturbed_guidance_config:
-            if config.skip_attention or not config.skip_attention_scores or config.skip_ff:
-                logger.warning(
-                    "Perturbed Attention Guidance is designed to perturb attention scores, so `skip_attention` should be False, `skip_attention_scores` should be True, and `skip_ff` should be False. "
-                    "Please check your configuration. Modifying the config to match the expected values."
-                )
-            config.skip_attention = False
-            config.skip_attention_scores = True
-            config.skip_ff = False
-
-        self.skip_layer_config = perturbed_guidance_config
-        self._skip_layer_hook_names = [f"SkipLayerGuidance_{i}" for i in range(len(self.skip_layer_config))]
-
-    # Copied from diffusers.guiders.skip_layer_guidance.SkipLayerGuidance.prepare_models
-    def prepare_models(self, denoiser: torch.nn.Module) -> None:
-        self._count_prepared += 1
-        if self._is_slg_enabled() and self.is_conditional and self._count_prepared > 1:
-            for name, config in zip(self._skip_layer_hook_names, self.skip_layer_config):
-                _apply_layer_skip_hook(denoiser, config, name=name)
-
-    # Copied from diffusers.guiders.skip_layer_guidance.SkipLayerGuidance.cleanup_models
-    def cleanup_models(self, denoiser: torch.nn.Module) -> None:
-        if self._is_slg_enabled() and self.is_conditional and self._count_prepared > 1:
-            registry = HookRegistry.check_if_exists_or_initialize(denoiser)
-            # Remove the hooks after inference
-            for hook_name in self._skip_layer_hook_names:
-                registry.remove_hook(hook_name, recurse=True)
-
-    # Copied from diffusers.guiders.skip_layer_guidance.SkipLayerGuidance.prepare_inputs
-    def prepare_inputs(self, data: dict[str, tuple[torch.Tensor, torch.Tensor]]) -> list["BlockState"]:
-        if self.num_conditions == 1:
-            tuple_indices = [0]
-            input_predictions = ["pred_cond"]
-        elif self.num_conditions == 2:
-            tuple_indices = [0, 1]
-            input_predictions = (
-                ["pred_cond", "pred_uncond"] if self._is_cfg_enabled() else ["pred_cond", "pred_cond_skip"]
-            )
-        else:
-            tuple_indices = [0, 1, 0]
-            input_predictions = ["pred_cond", "pred_uncond", "pred_cond_skip"]
-        data_batches = []
-        for tuple_idx, input_prediction in zip(tuple_indices, input_predictions):
-            data_batch = self._prepare_batch(data, tuple_idx, input_prediction)
-            data_batches.append(data_batch)
-        return data_batches
-
-    def prepare_inputs_from_block_state(
-        self, data: "BlockState", input_fields: dict[str, str | tuple[str, str]]
-    ) -> list["BlockState"]:
-        if self.num_conditions == 1:
-            tuple_indices = [0]
-            input_predictions = ["pred_cond"]
-        elif self.num_conditions == 2:
-            tuple_indices = [0, 1]
-            input_predictions = (
-                ["pred_cond", "pred_uncond"] if self._is_cfg_enabled() else ["pred_cond", "pred_cond_skip"]
-            )
-        else:
-            tuple_indices = [0, 1, 0]
-            input_predictions = ["pred_cond", "pred_uncond", "pred_cond_skip"]
-        data_batches = []
-        for tuple_idx, input_prediction in zip(tuple_indices, input_predictions):
-            data_batch = self._prepare_batch_from_block_state(input_fields, data, tuple_idx, input_prediction)
-            data_batches.append(data_batch)
-        return data_batches
-
-    # Copied from diffusers.guiders.skip_layer_guidance.SkipLayerGuidance.forward
-    def forward(
-        self,
-        pred_cond: torch.Tensor,
-        pred_uncond: torch.Tensor | None = None,
-        pred_cond_skip: torch.Tensor | None = None,
-    ) -> GuiderOutput:
-        pred = None
-
-        if not self._is_cfg_enabled() and not self._is_slg_enabled():
-            pred = pred_cond
-        elif not self._is_cfg_enabled():
-            shift = pred_cond - pred_cond_skip
-            pred = pred_cond if self.use_original_formulation else pred_cond_skip
-            pred = pred + self.skip_layer_guidance_scale * shift
-        elif not self._is_slg_enabled():
-            shift = pred_cond - pred_uncond
-            pred = pred_cond if self.use_original_formulation else pred_uncond
-            pred = pred + self.guidance_scale * shift
-        else:
-            shift = pred_cond - pred_uncond
-            shift_skip = pred_cond - pred_cond_skip
-            pred = pred_cond if self.use_original_formulation else pred_uncond
-            pred = pred + self.guidance_scale * shift + self.skip_layer_guidance_scale * shift_skip
-
-        if self.guidance_rescale > 0.0:
-            pred = rescale_noise_cfg(pred, pred_cond, self.guidance_rescale)
-
-        return GuiderOutput(pred=pred, pred_cond=pred_cond, pred_uncond=pred_uncond)
-
-    @property
-    # Copied from diffusers.guiders.skip_layer_guidance.SkipLayerGuidance.is_conditional
-    def is_conditional(self) -> bool:
-        return self._count_prepared == 1 or self._count_prepared == 3
-
-    @property
-    # Copied from diffusers.guiders.skip_layer_guidance.SkipLayerGuidance.num_conditions
-    def num_conditions(self) -> int:
-        num_conditions = 1
-        if self._is_cfg_enabled():
-            num_conditions += 1
-        if self._is_slg_enabled():
-            num_conditions += 1
-        return num_conditions
-
-    # Copied from diffusers.guiders.skip_layer_guidance.SkipLayerGuidance._is_cfg_enabled
-    def _is_cfg_enabled(self) -> bool:
-        if not self._enabled:
-            return False
-
-        is_within_range = True
-        if self._num_inference_steps is not None:
-            skip_start_step = int(self._start * self._num_inference_steps)
-            skip_stop_step = int(self._stop * self._num_inference_steps)
-            is_within_range = skip_start_step <= self._step < skip_stop_step
-
-        is_close = False
-        if self.use_original_formulation:
-            is_close = math.isclose(self.guidance_scale, 0.0)
-        else:
-            is_close = math.isclose(self.guidance_scale, 1.0)
-
-        return is_within_range and not is_close
-
-    # Copied from diffusers.guiders.skip_layer_guidance.SkipLayerGuidance._is_slg_enabled
-    def _is_slg_enabled(self) -> bool:
-        if not self._enabled:
-            return False
-
-        is_within_range = True
-        if self._num_inference_steps is not None:
-            skip_start_step = int(self.skip_layer_guidance_start * self._num_inference_steps)
-            skip_stop_step = int(self.skip_layer_guidance_stop * self._num_inference_steps)
-            is_within_range = skip_start_step < self._step < skip_stop_step
-
-        is_zero = math.isclose(self.skip_layer_guidance_scale, 0.0)
-
-        return is_within_range and not is_zero

diffusers/guiders/skip_layer_guidance.py

@@ -1,280 +0,0 @@
-# Copyright 2026 The HuggingFace Team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-from __future__ import annotations
-
-import math
-from typing import TYPE_CHECKING, Any
-
-import torch
-
-from ..configuration_utils import register_to_config
-from ..hooks import HookRegistry, LayerSkipConfig
-from ..hooks.layer_skip import _apply_layer_skip_hook
-from .guider_utils import BaseGuidance, GuiderOutput, rescale_noise_cfg
-
-
-if TYPE_CHECKING:
-    from ..modular_pipelines.modular_pipeline import BlockState
-
-
-class SkipLayerGuidance(BaseGuidance):
-    """
-    Skip Layer Guidance (SLG): https://github.com/Stability-AI/sd3.5
-
-    Spatio-Temporal Guidance (STG): https://huggingface.co/papers/2411.18664
-
-    SLG was introduced by StabilityAI for improving structure and anotomy coherence in generated images. It works by
-    skipping the forward pass of specified transformer blocks during the denoising process on an additional conditional
-    batch of data, apart from the conditional and unconditional batches already used in CFG
-    ([~guiders.classifier_free_guidance.ClassifierFreeGuidance]), and then scaling and shifting the CFG predictions
-    based on the difference between conditional without skipping and conditional with skipping predictions.
-
-    The intution behind SLG can be thought of as moving the CFG predicted distribution estimates further away from
-    worse versions of the conditional distribution estimates (because skipping layers is equivalent to using a worse
-    version of the model for the conditional prediction).
-
-    STG is an improvement and follow-up work combining ideas from SLG, PAG and similar techniques for improving
-    generation quality in video diffusion models.
-
-    Additional reading:
-    - [Guiding a Diffusion Model with a Bad Version of Itself](https://huggingface.co/papers/2406.02507)
-
-    The values for `skip_layer_guidance_scale`, `skip_layer_guidance_start`, and `skip_layer_guidance_stop` are
-    defaulted to the recommendations by StabilityAI for Stable Diffusion 3.5 Medium.
-
-    Args:
-        guidance_scale (`float`, defaults to `7.5`):
-            The scale parameter for classifier-free guidance. Higher values result in stronger conditioning on the text
-            prompt, while lower values allow for more freedom in generation. Higher values may lead to saturation and
-            deterioration of image quality.
-        skip_layer_guidance_scale (`float`, defaults to `2.8`):
-            The scale parameter for skip layer guidance. Anatomy and structure coherence may improve with higher
-            values, but it may also lead to overexposure and saturation.
-        skip_layer_guidance_start (`float`, defaults to `0.01`):
-            The fraction of the total number of denoising steps after which skip layer guidance starts.
-        skip_layer_guidance_stop (`float`, defaults to `0.2`):
-            The fraction of the total number of denoising steps after which skip layer guidance stops.
-        skip_layer_guidance_layers (`int` or `list[int]`, *optional*):
-            The layer indices to apply skip layer guidance to. Can be a single integer or a list of integers. If not
-            provided, `skip_layer_config` must be provided. The recommended values are `[7, 8, 9]` for Stable Diffusion
-            3.5 Medium.
-        skip_layer_config (`LayerSkipConfig` or `list[LayerSkipConfig]`, *optional*):
-            The configuration for the skip layer guidance. Can be a single `LayerSkipConfig` or a list of
-            `LayerSkipConfig`. If not provided, `skip_layer_guidance_layers` must be provided.
-        guidance_rescale (`float`, defaults to `0.0`):
-            The rescale factor applied to the noise predictions. This is used to improve image quality and fix
-            overexposure. Based on Section 3.4 from [Common Diffusion Noise Schedules and Sample Steps are
-            Flawed](https://huggingface.co/papers/2305.08891).
-        use_original_formulation (`bool`, defaults to `False`):
-            Whether to use the original formulation of classifier-free guidance as proposed in the paper. By default,
-            we use the diffusers-native implementation that has been in the codebase for a long time. See
-            [~guiders.classifier_free_guidance.ClassifierFreeGuidance] for more details.
-        start (`float`, defaults to `0.01`):
-            The fraction of the total number of denoising steps after which guidance starts.
-        stop (`float`, defaults to `0.2`):
-            The fraction of the total number of denoising steps after which guidance stops.
-    """
-
-    _input_predictions = ["pred_cond", "pred_uncond", "pred_cond_skip"]
-
-    @register_to_config
-    def __init__(
-        self,
-        guidance_scale: float = 7.5,
-        skip_layer_guidance_scale: float = 2.8,
-        skip_layer_guidance_start: float = 0.01,
-        skip_layer_guidance_stop: float = 0.2,
-        skip_layer_guidance_layers: int | list[int] | None = None,
-        skip_layer_config: LayerSkipConfig | list[LayerSkipConfig] | dict[str, Any] = None,
-        guidance_rescale: float = 0.0,
-        use_original_formulation: bool = False,
-        start: float = 0.0,
-        stop: float = 1.0,
-        enabled: bool = True,
-    ):
-        super().__init__(start, stop, enabled)
-
-        self.guidance_scale = guidance_scale
-        self.skip_layer_guidance_scale = skip_layer_guidance_scale
-        self.skip_layer_guidance_start = skip_layer_guidance_start
-        self.skip_layer_guidance_stop = skip_layer_guidance_stop
-        self.guidance_rescale = guidance_rescale
-        self.use_original_formulation = use_original_formulation
-
-        if not (0.0 <= skip_layer_guidance_start < 1.0):
-            raise ValueError(
-                f"Expected `skip_layer_guidance_start` to be between 0.0 and 1.0, but got {skip_layer_guidance_start}."
-            )
-        if not (skip_layer_guidance_start <= skip_layer_guidance_stop <= 1.0):
-            raise ValueError(
-                f"Expected `skip_layer_guidance_stop` to be between 0.0 and 1.0, but got {skip_layer_guidance_stop}."
-            )
-
-        if skip_layer_guidance_layers is None and skip_layer_config is None:
-            raise ValueError(
-                "Either `skip_layer_guidance_layers` or `skip_layer_config` must be provided to enable Skip Layer Guidance."
-            )
-        if skip_layer_guidance_layers is not None and skip_layer_config is not None:
-            raise ValueError("Only one of `skip_layer_guidance_layers` or `skip_layer_config` can be provided.")
-
-        if skip_layer_guidance_layers is not None:
-            if isinstance(skip_layer_guidance_layers, int):
-                skip_layer_guidance_layers = [skip_layer_guidance_layers]
-            if not isinstance(skip_layer_guidance_layers, list):
-                raise ValueError(
-                    f"Expected `skip_layer_guidance_layers` to be an int or a list of ints, but got {type(skip_layer_guidance_layers)}."
-                )
-            skip_layer_config = [LayerSkipConfig(layer, fqn="auto") for layer in skip_layer_guidance_layers]
-
-        if isinstance(skip_layer_config, dict):
-            skip_layer_config = LayerSkipConfig.from_dict(skip_layer_config)
-
-        if isinstance(skip_layer_config, LayerSkipConfig):
-            skip_layer_config = [skip_layer_config]
-
-        if not isinstance(skip_layer_config, list):
-            raise ValueError(
-                f"Expected `skip_layer_config` to be a LayerSkipConfig or a list of LayerSkipConfig, but got {type(skip_layer_config)}."
-            )
-        elif isinstance(next(iter(skip_layer_config), None), dict):
-            skip_layer_config = [LayerSkipConfig.from_dict(config) for config in skip_layer_config]
-
-        self.skip_layer_config = skip_layer_config
-        self._skip_layer_hook_names = [f"SkipLayerGuidance_{i}" for i in range(len(self.skip_layer_config))]
-
-    def prepare_models(self, denoiser: torch.nn.Module) -> None:
-        self._count_prepared += 1
-        if self._is_slg_enabled() and self.is_conditional and self._count_prepared > 1:
-            for name, config in zip(self._skip_layer_hook_names, self.skip_layer_config):
-                _apply_layer_skip_hook(denoiser, config, name=name)
-
-    def cleanup_models(self, denoiser: torch.nn.Module) -> None:
-        if self._is_slg_enabled() and self.is_conditional and self._count_prepared > 1:
-            registry = HookRegistry.check_if_exists_or_initialize(denoiser)
-            # Remove the hooks after inference
-            for hook_name in self._skip_layer_hook_names:
-                registry.remove_hook(hook_name, recurse=True)
-
-    def prepare_inputs(self, data: dict[str, tuple[torch.Tensor, torch.Tensor]]) -> list["BlockState"]:
-        if self.num_conditions == 1:
-            tuple_indices = [0]
-            input_predictions = ["pred_cond"]
-        elif self.num_conditions == 2:
-            tuple_indices = [0, 1]
-            input_predictions = (
-                ["pred_cond", "pred_uncond"] if self._is_cfg_enabled() else ["pred_cond", "pred_cond_skip"]
-            )
-        else:
-            tuple_indices = [0, 1, 0]
-            input_predictions = ["pred_cond", "pred_uncond", "pred_cond_skip"]
-        data_batches = []
-        for tuple_idx, input_prediction in zip(tuple_indices, input_predictions):
-            data_batch = self._prepare_batch(data, tuple_idx, input_prediction)
-            data_batches.append(data_batch)
-        return data_batches
-
-    def prepare_inputs_from_block_state(
-        self, data: "BlockState", input_fields: dict[str, str | tuple[str, str]]
-    ) -> list["BlockState"]:
-        if self.num_conditions == 1:
-            tuple_indices = [0]
-            input_predictions = ["pred_cond"]
-        elif self.num_conditions == 2:
-            tuple_indices = [0, 1]
-            input_predictions = (
-                ["pred_cond", "pred_uncond"] if self._is_cfg_enabled() else ["pred_cond", "pred_cond_skip"]
-            )
-        else:
-            tuple_indices = [0, 1, 0]
-            input_predictions = ["pred_cond", "pred_uncond", "pred_cond_skip"]
-        data_batches = []
-        for tuple_idx, input_prediction in zip(tuple_indices, input_predictions):
-            data_batch = self._prepare_batch_from_block_state(input_fields, data, tuple_idx, input_prediction)
-            data_batches.append(data_batch)
-        return data_batches
-
-    def forward(
-        self,
-        pred_cond: torch.Tensor,
-        pred_uncond: torch.Tensor | None = None,
-        pred_cond_skip: torch.Tensor | None = None,
-    ) -> GuiderOutput:
-        pred = None
-
-        if not self._is_cfg_enabled() and not self._is_slg_enabled():
-            pred = pred_cond
-        elif not self._is_cfg_enabled():
-            shift = pred_cond - pred_cond_skip
-            pred = pred_cond if self.use_original_formulation else pred_cond_skip
-            pred = pred + self.skip_layer_guidance_scale * shift
-        elif not self._is_slg_enabled():
-            shift = pred_cond - pred_uncond
-            pred = pred_cond if self.use_original_formulation else pred_uncond
-            pred = pred + self.guidance_scale * shift
-        else:
-            shift = pred_cond - pred_uncond
-            shift_skip = pred_cond - pred_cond_skip
-            pred = pred_cond if self.use_original_formulation else pred_uncond
-            pred = pred + self.guidance_scale * shift + self.skip_layer_guidance_scale * shift_skip
-
-        if self.guidance_rescale > 0.0:
-            pred = rescale_noise_cfg(pred, pred_cond, self.guidance_rescale)
-
-        return GuiderOutput(pred=pred, pred_cond=pred_cond, pred_uncond=pred_uncond)
-
-    @property
-    def is_conditional(self) -> bool:
-        return self._count_prepared == 1 or self._count_prepared == 3
-
-    @property
-    def num_conditions(self) -> int:
-        num_conditions = 1
-        if self._is_cfg_enabled():
-            num_conditions += 1
-        if self._is_slg_enabled():
-            num_conditions += 1
-        return num_conditions
-
-    def _is_cfg_enabled(self) -> bool:
-        if not self._enabled:
-            return False
-
-        is_within_range = True
-        if self._num_inference_steps is not None:
-            skip_start_step = int(self._start * self._num_inference_steps)
-            skip_stop_step = int(self._stop * self._num_inference_steps)
-            is_within_range = skip_start_step <= self._step < skip_stop_step
-
-        is_close = False
-        if self.use_original_formulation:
-            is_close = math.isclose(self.guidance_scale, 0.0)
-        else:
-            is_close = math.isclose(self.guidance_scale, 1.0)
-
-        return is_within_range and not is_close
-
-    def _is_slg_enabled(self) -> bool:
-        if not self._enabled:
-            return False
-
-        is_within_range = True
-        if self._num_inference_steps is not None:
-            skip_start_step = int(self.skip_layer_guidance_start * self._num_inference_steps)
-            skip_stop_step = int(self.skip_layer_guidance_stop * self._num_inference_steps)
-            is_within_range = skip_start_step < self._step < skip_stop_step
-
-        is_zero = math.isclose(self.skip_layer_guidance_scale, 0.0)
-
-        return is_within_range and not is_zero

diffusers/guiders/smoothed_energy_guidance.py

@@ -1,269 +0,0 @@
-# Copyright 2026 The HuggingFace Team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-from __future__ import annotations
-
-import math
-from typing import TYPE_CHECKING
-
-import torch
-
-from ..configuration_utils import register_to_config
-from ..hooks import HookRegistry
-from ..hooks.smoothed_energy_guidance_utils import SmoothedEnergyGuidanceConfig, _apply_smoothed_energy_guidance_hook
-from .guider_utils import BaseGuidance, GuiderOutput, rescale_noise_cfg
-
-
-if TYPE_CHECKING:
-    from ..modular_pipelines.modular_pipeline import BlockState
-
-
-class SmoothedEnergyGuidance(BaseGuidance):
-    """
-    Smoothed Energy Guidance (SEG): https://huggingface.co/papers/2408.00760
-
-    SEG is only supported as an experimental prototype feature for now, so the implementation may be modified in the
-    future without warning or guarantee of reproducibility. This implementation assumes:
-    - Generated images are square (height == width)
-    - The model does not combine different modalities together (e.g., text and image latent streams are not combined
-      together such as Flux)
-
-    Args:
-        guidance_scale (`float`, defaults to `7.5`):
-            The scale parameter for classifier-free guidance. Higher values result in stronger conditioning on the text
-            prompt, while lower values allow for more freedom in generation. Higher values may lead to saturation and
-            deterioration of image quality.
-        seg_guidance_scale (`float`, defaults to `3.0`):
-            The scale parameter for smoothed energy guidance. Anatomy and structure coherence may improve with higher
-            values, but it may also lead to overexposure and saturation.
-        seg_blur_sigma (`float`, defaults to `9999999.0`):
-            The amount by which we blur the attention weights. Setting this value greater than 9999.0 results in
-            infinite blur, which means uniform queries. Controlling it exponentially is empirically effective.
-        seg_blur_threshold_inf (`float`, defaults to `9999.0`):
-            The threshold above which the blur is considered infinite.
-        seg_guidance_start (`float`, defaults to `0.0`):
-            The fraction of the total number of denoising steps after which smoothed energy guidance starts.
-        seg_guidance_stop (`float`, defaults to `1.0`):
-            The fraction of the total number of denoising steps after which smoothed energy guidance stops.
-        seg_guidance_layers (`int` or `list[int]`, *optional*):
-            The layer indices to apply smoothed energy guidance to. Can be a single integer or a list of integers. If
-            not provided, `seg_guidance_config` must be provided. The recommended values are `[7, 8, 9]` for Stable
-            Diffusion 3.5 Medium.
-        seg_guidance_config (`SmoothedEnergyGuidanceConfig` or `list[SmoothedEnergyGuidanceConfig]`, *optional*):
-            The configuration for the smoothed energy layer guidance. Can be a single `SmoothedEnergyGuidanceConfig` or
-            a list of `SmoothedEnergyGuidanceConfig`. If not provided, `seg_guidance_layers` must be provided.
-        guidance_rescale (`float`, defaults to `0.0`):
-            The rescale factor applied to the noise predictions. This is used to improve image quality and fix
-            overexposure. Based on Section 3.4 from [Common Diffusion Noise Schedules and Sample Steps are
-            Flawed](https://huggingface.co/papers/2305.08891).
-        use_original_formulation (`bool`, defaults to `False`):
-            Whether to use the original formulation of classifier-free guidance as proposed in the paper. By default,
-            we use the diffusers-native implementation that has been in the codebase for a long time. See
-            [~guiders.classifier_free_guidance.ClassifierFreeGuidance] for more details.
-        start (`float`, defaults to `0.01`):
-            The fraction of the total number of denoising steps after which guidance starts.
-        stop (`float`, defaults to `0.2`):
-            The fraction of the total number of denoising steps after which guidance stops.
-    """
-
-    _input_predictions = ["pred_cond", "pred_uncond", "pred_cond_seg"]
-
-    @register_to_config
-    def __init__(
-        self,
-        guidance_scale: float = 7.5,
-        seg_guidance_scale: float = 2.8,
-        seg_blur_sigma: float = 9999999.0,
-        seg_blur_threshold_inf: float = 9999.0,
-        seg_guidance_start: float = 0.0,
-        seg_guidance_stop: float = 1.0,
-        seg_guidance_layers: int | list[int] | None = None,
-        seg_guidance_config: SmoothedEnergyGuidanceConfig | list[SmoothedEnergyGuidanceConfig] = None,
-        guidance_rescale: float = 0.0,
-        use_original_formulation: bool = False,
-        start: float = 0.0,
-        stop: float = 1.0,
-        enabled: bool = True,
-    ):
-        super().__init__(start, stop, enabled)
-
-        self.guidance_scale = guidance_scale
-        self.seg_guidance_scale = seg_guidance_scale
-        self.seg_blur_sigma = seg_blur_sigma
-        self.seg_blur_threshold_inf = seg_blur_threshold_inf
-        self.seg_guidance_start = seg_guidance_start
-        self.seg_guidance_stop = seg_guidance_stop
-        self.guidance_rescale = guidance_rescale
-        self.use_original_formulation = use_original_formulation
-
-        if not (0.0 <= seg_guidance_start < 1.0):
-            raise ValueError(f"Expected `seg_guidance_start` to be between 0.0 and 1.0, but got {seg_guidance_start}.")
-        if not (seg_guidance_start <= seg_guidance_stop <= 1.0):
-            raise ValueError(f"Expected `seg_guidance_stop` to be between 0.0 and 1.0, but got {seg_guidance_stop}.")
-
-        if seg_guidance_layers is None and seg_guidance_config is None:
-            raise ValueError(
-                "Either `seg_guidance_layers` or `seg_guidance_config` must be provided to enable Smoothed Energy Guidance."
-            )
-        if seg_guidance_layers is not None and seg_guidance_config is not None:
-            raise ValueError("Only one of `seg_guidance_layers` or `seg_guidance_config` can be provided.")
-
-        if seg_guidance_layers is not None:
-            if isinstance(seg_guidance_layers, int):
-                seg_guidance_layers = [seg_guidance_layers]
-            if not isinstance(seg_guidance_layers, list):
-                raise ValueError(
-                    f"Expected `seg_guidance_layers` to be an int or a list of ints, but got {type(seg_guidance_layers)}."
-                )
-            seg_guidance_config = [SmoothedEnergyGuidanceConfig(layer, fqn="auto") for layer in seg_guidance_layers]
-
-        if isinstance(seg_guidance_config, dict):
-            seg_guidance_config = SmoothedEnergyGuidanceConfig.from_dict(seg_guidance_config)
-
-        if isinstance(seg_guidance_config, SmoothedEnergyGuidanceConfig):
-            seg_guidance_config = [seg_guidance_config]
-
-        if not isinstance(seg_guidance_config, list):
-            raise ValueError(
-                f"Expected `seg_guidance_config` to be a SmoothedEnergyGuidanceConfig or a list of SmoothedEnergyGuidanceConfig, but got {type(seg_guidance_config)}."
-            )
-        elif isinstance(next(iter(seg_guidance_config), None), dict):
-            seg_guidance_config = [SmoothedEnergyGuidanceConfig.from_dict(config) for config in seg_guidance_config]
-
-        self.seg_guidance_config = seg_guidance_config
-        self._seg_layer_hook_names = [f"SmoothedEnergyGuidance_{i}" for i in range(len(self.seg_guidance_config))]
-
-    def prepare_models(self, denoiser: torch.nn.Module) -> None:
-        if self._is_seg_enabled() and self.is_conditional and self._count_prepared > 1:
-            for name, config in zip(self._seg_layer_hook_names, self.seg_guidance_config):
-                _apply_smoothed_energy_guidance_hook(denoiser, config, self.seg_blur_sigma, name=name)
-
-    def cleanup_models(self, denoiser: torch.nn.Module):
-        if self._is_seg_enabled() and self.is_conditional and self._count_prepared > 1:
-            registry = HookRegistry.check_if_exists_or_initialize(denoiser)
-            # Remove the hooks after inference
-            for hook_name in self._seg_layer_hook_names:
-                registry.remove_hook(hook_name, recurse=True)
-
-    def prepare_inputs(self, data: dict[str, tuple[torch.Tensor, torch.Tensor]]) -> list["BlockState"]:
-        if self.num_conditions == 1:
-            tuple_indices = [0]
-            input_predictions = ["pred_cond"]
-        elif self.num_conditions == 2:
-            tuple_indices = [0, 1]
-            input_predictions = (
-                ["pred_cond", "pred_uncond"] if self._is_cfg_enabled() else ["pred_cond", "pred_cond_seg"]
-            )
-        else:
-            tuple_indices = [0, 1, 0]
-            input_predictions = ["pred_cond", "pred_uncond", "pred_cond_seg"]
-        data_batches = []
-        for tuple_idx, input_prediction in zip(tuple_indices, input_predictions):
-            data_batch = self._prepare_batch(data, tuple_idx, input_prediction)
-            data_batches.append(data_batch)
-        return data_batches
-
-    def prepare_inputs_from_block_state(
-        self, data: "BlockState", input_fields: dict[str, str | tuple[str, str]]
-    ) -> list["BlockState"]:
-        if self.num_conditions == 1:
-            tuple_indices = [0]
-            input_predictions = ["pred_cond"]
-        elif self.num_conditions == 2:
-            tuple_indices = [0, 1]
-            input_predictions = (
-                ["pred_cond", "pred_uncond"] if self._is_cfg_enabled() else ["pred_cond", "pred_cond_seg"]
-            )
-        else:
-            tuple_indices = [0, 1, 0]
-            input_predictions = ["pred_cond", "pred_uncond", "pred_cond_seg"]
-        data_batches = []
-        for tuple_idx, input_prediction in zip(tuple_indices, input_predictions):
-            data_batch = self._prepare_batch_from_block_state(input_fields, data, tuple_idx, input_prediction)
-            data_batches.append(data_batch)
-        return data_batches
-
-    def forward(
-        self,
-        pred_cond: torch.Tensor,
-        pred_uncond: torch.Tensor | None = None,
-        pred_cond_seg: torch.Tensor | None = None,
-    ) -> GuiderOutput:
-        pred = None
-
-        if not self._is_cfg_enabled() and not self._is_seg_enabled():
-            pred = pred_cond
-        elif not self._is_cfg_enabled():
-            shift = pred_cond - pred_cond_seg
-            pred = pred_cond if self.use_original_formulation else pred_cond_seg
-            pred = pred + self.seg_guidance_scale * shift
-        elif not self._is_seg_enabled():
-            shift = pred_cond - pred_uncond
-            pred = pred_cond if self.use_original_formulation else pred_uncond
-            pred = pred + self.guidance_scale * shift
-        else:
-            shift = pred_cond - pred_uncond
-            shift_seg = pred_cond - pred_cond_seg
-            pred = pred_cond if self.use_original_formulation else pred_uncond
-            pred = pred + self.guidance_scale * shift + self.seg_guidance_scale * shift_seg
-
-        if self.guidance_rescale > 0.0:
-            pred = rescale_noise_cfg(pred, pred_cond, self.guidance_rescale)
-
-        return GuiderOutput(pred=pred, pred_cond=pred_cond, pred_uncond=pred_uncond)
-
-    @property
-    def is_conditional(self) -> bool:
-        return self._count_prepared == 1 or self._count_prepared == 3
-
-    @property
-    def num_conditions(self) -> int:
-        num_conditions = 1
-        if self._is_cfg_enabled():
-            num_conditions += 1
-        if self._is_seg_enabled():
-            num_conditions += 1
-        return num_conditions
-
-    def _is_cfg_enabled(self) -> bool:
-        if not self._enabled:
-            return False
-
-        is_within_range = True
-        if self._num_inference_steps is not None:
-            skip_start_step = int(self._start * self._num_inference_steps)
-            skip_stop_step = int(self._stop * self._num_inference_steps)
-            is_within_range = skip_start_step <= self._step < skip_stop_step
-
-        is_close = False
-        if self.use_original_formulation:
-            is_close = math.isclose(self.guidance_scale, 0.0)
-        else:
-            is_close = math.isclose(self.guidance_scale, 1.0)
-
-        return is_within_range and not is_close
-
-    def _is_seg_enabled(self) -> bool:
-        if not self._enabled:
-            return False
-
-        is_within_range = True
-        if self._num_inference_steps is not None:
-            skip_start_step = int(self.seg_guidance_start * self._num_inference_steps)
-            skip_stop_step = int(self.seg_guidance_stop * self._num_inference_steps)
-            is_within_range = skip_start_step < self._step < skip_stop_step
-
-        is_zero = math.isclose(self.seg_guidance_scale, 0.0)
-
-        return is_within_range and not is_zero

diffusers/guiders/tangential_classifier_free_guidance.py

@@ -1,151 +0,0 @@
-# Copyright 2026 The HuggingFace Team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-from __future__ import annotations
-
-import math
-from typing import TYPE_CHECKING
-
-import torch
-
-from ..configuration_utils import register_to_config
-from .guider_utils import BaseGuidance, GuiderOutput, rescale_noise_cfg
-
-
-if TYPE_CHECKING:
-    from ..modular_pipelines.modular_pipeline import BlockState
-
-
-class TangentialClassifierFreeGuidance(BaseGuidance):
-    """
-    Tangential Classifier Free Guidance (TCFG): https://huggingface.co/papers/2503.18137
-
-    Args:
-        guidance_scale (`float`, defaults to `7.5`):
-            The scale parameter for classifier-free guidance. Higher values result in stronger conditioning on the text
-            prompt, while lower values allow for more freedom in generation. Higher values may lead to saturation and
-            deterioration of image quality.
-        guidance_rescale (`float`, defaults to `0.0`):
-            The rescale factor applied to the noise predictions. This is used to improve image quality and fix
-            overexposure. Based on Section 3.4 from [Common Diffusion Noise Schedules and Sample Steps are
-            Flawed](https://huggingface.co/papers/2305.08891).
-        use_original_formulation (`bool`, defaults to `False`):
-            Whether to use the original formulation of classifier-free guidance as proposed in the paper. By default,
-            we use the diffusers-native implementation that has been in the codebase for a long time. See
-            [~guiders.classifier_free_guidance.ClassifierFreeGuidance] for more details.
-        start (`float`, defaults to `0.0`):
-            The fraction of the total number of denoising steps after which guidance starts.
-        stop (`float`, defaults to `1.0`):
-            The fraction of the total number of denoising steps after which guidance stops.
-    """
-
-    _input_predictions = ["pred_cond", "pred_uncond"]
-
-    @register_to_config
-    def __init__(
-        self,
-        guidance_scale: float = 7.5,
-        guidance_rescale: float = 0.0,
-        use_original_formulation: bool = False,
-        start: float = 0.0,
-        stop: float = 1.0,
-        enabled: bool = True,
-    ):
-        super().__init__(start, stop, enabled)
-
-        self.guidance_scale = guidance_scale
-        self.guidance_rescale = guidance_rescale
-        self.use_original_formulation = use_original_formulation
-
-    def prepare_inputs(self, data: dict[str, tuple[torch.Tensor, torch.Tensor]]) -> list["BlockState"]:
-        tuple_indices = [0] if self.num_conditions == 1 else [0, 1]
-        data_batches = []
-        for tuple_idx, input_prediction in zip(tuple_indices, self._input_predictions):
-            data_batch = self._prepare_batch(data, tuple_idx, input_prediction)
-            data_batches.append(data_batch)
-        return data_batches
-
-    def prepare_inputs_from_block_state(
-        self, data: "BlockState", input_fields: dict[str, str | tuple[str, str]]
-    ) -> list["BlockState"]:
-        tuple_indices = [0] if self.num_conditions == 1 else [0, 1]
-        data_batches = []
-        for tuple_idx, input_prediction in zip(tuple_indices, self._input_predictions):
-            data_batch = self._prepare_batch_from_block_state(input_fields, data, tuple_idx, input_prediction)
-            data_batches.append(data_batch)
-        return data_batches
-
-    def forward(self, pred_cond: torch.Tensor, pred_uncond: torch.Tensor | None = None) -> GuiderOutput:
-        pred = None
-
-        if not self._is_tcfg_enabled():
-            pred = pred_cond
-        else:
-            pred = normalized_guidance(pred_cond, pred_uncond, self.guidance_scale, self.use_original_formulation)
-
-        if self.guidance_rescale > 0.0:
-            pred = rescale_noise_cfg(pred, pred_cond, self.guidance_rescale)
-
-        return GuiderOutput(pred=pred, pred_cond=pred_cond, pred_uncond=pred_uncond)
-
-    @property
-    def is_conditional(self) -> bool:
-        return self._num_outputs_prepared == 1
-
-    @property
-    def num_conditions(self) -> int:
-        num_conditions = 1
-        if self._is_tcfg_enabled():
-            num_conditions += 1
-        return num_conditions
-
-    def _is_tcfg_enabled(self) -> bool:
-        if not self._enabled:
-            return False
-
-        is_within_range = True
-        if self._num_inference_steps is not None:
-            skip_start_step = int(self._start * self._num_inference_steps)
-            skip_stop_step = int(self._stop * self._num_inference_steps)
-            is_within_range = skip_start_step <= self._step < skip_stop_step
-
-        is_close = False
-        if self.use_original_formulation:
-            is_close = math.isclose(self.guidance_scale, 0.0)
-        else:
-            is_close = math.isclose(self.guidance_scale, 1.0)
-
-        return is_within_range and not is_close
-
-
-def normalized_guidance(
-    pred_cond: torch.Tensor, pred_uncond: torch.Tensor, guidance_scale: float, use_original_formulation: bool = False
-) -> torch.Tensor:
-    cond_dtype = pred_cond.dtype
-    preds = torch.stack([pred_cond, pred_uncond], dim=1).float()
-    preds = preds.flatten(2)
-    U, S, Vh = torch.linalg.svd(preds, full_matrices=False)
-    Vh_modified = Vh.clone()
-    Vh_modified[:, 1] = 0
-
-    uncond_flat = pred_uncond.reshape(pred_uncond.size(0), 1, -1).float()
-    x_Vh = torch.matmul(uncond_flat, Vh.transpose(-2, -1))
-    x_Vh_V = torch.matmul(x_Vh, Vh_modified)
-    pred_uncond = x_Vh_V.reshape(pred_uncond.shape).to(cond_dtype)
-
-    pred = pred_cond if use_original_formulation else pred_uncond
-    shift = pred_cond - pred_uncond
-    pred = pred + guidance_scale * shift
-
-    return pred

diffusers/hooks/__init__.py

@@ -1,30 +0,0 @@
-# Copyright 2024 The HuggingFace Team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-from ..utils import is_torch_available
-
-
-if is_torch_available():
-    from .context_parallel import apply_context_parallel
-    from .faster_cache import FasterCacheConfig, apply_faster_cache
-    from .first_block_cache import FirstBlockCacheConfig, apply_first_block_cache
-    from .group_offloading import apply_group_offloading
-    from .hooks import HookRegistry, ModelHook
-    from .layer_skip import LayerSkipConfig, apply_layer_skip
-    from .layerwise_casting import apply_layerwise_casting, apply_layerwise_casting_hook
-    from .mag_cache import MagCacheConfig, apply_mag_cache
-    from .pyramid_attention_broadcast import PyramidAttentionBroadcastConfig, apply_pyramid_attention_broadcast
-    from .smoothed_energy_guidance_utils import SmoothedEnergyGuidanceConfig
-    from .taylorseer_cache import TaylorSeerCacheConfig, apply_taylorseer_cache
-    from .text_kv_cache import TextKVCacheConfig, apply_text_kv_cache

diffusers/hooks/_common.py

@@ -1,61 +0,0 @@
-# Copyright 2026 The HuggingFace Team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-import torch
-
-from ..models.attention import AttentionModuleMixin, FeedForward, LuminaFeedForward
-from ..models.attention_processor import Attention, MochiAttention
-
-
-_ATTENTION_CLASSES = (Attention, MochiAttention, AttentionModuleMixin)
-_FEEDFORWARD_CLASSES = (FeedForward, LuminaFeedForward)
-
-_SPATIAL_TRANSFORMER_BLOCK_IDENTIFIERS = (
-    "blocks",
-    "transformer_blocks",
-    "single_transformer_blocks",
-    "layers",
-    "visual_transformer_blocks",
-)
-_TEMPORAL_TRANSFORMER_BLOCK_IDENTIFIERS = ("temporal_transformer_blocks",)
-_CROSS_TRANSFORMER_BLOCK_IDENTIFIERS = ("blocks", "transformer_blocks", "layers")
-
-_ALL_TRANSFORMER_BLOCK_IDENTIFIERS = tuple(
-    {
-        *_SPATIAL_TRANSFORMER_BLOCK_IDENTIFIERS,
-        *_TEMPORAL_TRANSFORMER_BLOCK_IDENTIFIERS,
-        *_CROSS_TRANSFORMER_BLOCK_IDENTIFIERS,
-    }
-)
-
-# Layers supported for group offloading and layerwise casting
-_GO_LC_SUPPORTED_PYTORCH_LAYERS = (
-    torch.nn.Conv1d,
-    torch.nn.Conv2d,
-    torch.nn.Conv3d,
-    torch.nn.ConvTranspose1d,
-    torch.nn.ConvTranspose2d,
-    torch.nn.ConvTranspose3d,
-    torch.nn.Linear,
-    torch.nn.Embedding,
-    # TODO(aryan): look into torch.nn.LayerNorm, torch.nn.GroupNorm later, seems to be causing some issues with CogVideoX
-    # because of double invocation of the same norm layer in CogVideoXLayerNorm
-)
-
-
-def _get_submodule_from_fqn(module: torch.nn.Module, fqn: str) -> torch.nn.Module | None:
-    for submodule_name, submodule in module.named_modules():
-        if submodule_name == fqn:
-            return submodule
-    return None

diffusers/hooks/_helpers.py

@@ -1,401 +0,0 @@
-# Copyright 2026 The HuggingFace Team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-import inspect
-from dataclasses import dataclass
-from typing import Any, Callable, Type
-
-
-@dataclass
-class AttentionProcessorMetadata:
-    skip_processor_output_fn: Callable[[Any], Any]
-
-
-@dataclass
-class TransformerBlockMetadata:
-    return_hidden_states_index: int = None
-    return_encoder_hidden_states_index: int = None
-    hidden_states_argument_name: str = "hidden_states"
-
-    _cls: Type = None
-    _cached_parameter_indices: dict[str, int] = None
-
-    def _get_parameter_from_args_kwargs(self, identifier: str, args=(), kwargs=None):
-        kwargs = kwargs or {}
-        if identifier in kwargs:
-            return kwargs[identifier]
-        if self._cached_parameter_indices is not None:
-            return args[self._cached_parameter_indices[identifier]]
-        if self._cls is None:
-            raise ValueError("Model class is not set for metadata.")
-        parameters = list(inspect.signature(self._cls.forward).parameters.keys())
-        parameters = parameters[1:]  # skip `self`
-        self._cached_parameter_indices = {param: i for i, param in enumerate(parameters)}
-        if identifier not in self._cached_parameter_indices:
-            raise ValueError(f"Parameter '{identifier}' not found in function signature but was requested.")
-        index = self._cached_parameter_indices[identifier]
-        if index >= len(args):
-            raise ValueError(f"Expected {index} arguments but got {len(args)}.")
-        return args[index]
-
-
-class AttentionProcessorRegistry:
-    _registry = {}
-    # TODO(aryan): this is only required for the time being because we need to do the registrations
-    # for classes. If we do it eagerly, i.e. call the functions in global scope, we will get circular
-    # import errors because of the models imported in this file.
-    _is_registered = False
-
-    @classmethod
-    def register(cls, model_class: Type, metadata: AttentionProcessorMetadata):
-        cls._register()
-        cls._registry[model_class] = metadata
-
-    @classmethod
-    def get(cls, model_class: Type) -> AttentionProcessorMetadata:
-        cls._register()
-        if model_class not in cls._registry:
-            raise ValueError(f"Model class {model_class} not registered.")
-        return cls._registry[model_class]
-
-    @classmethod
-    def _register(cls):
-        if cls._is_registered:
-            return
-        cls._is_registered = True
-        _register_attention_processors_metadata()
-
-
-class TransformerBlockRegistry:
-    _registry = {}
-    # TODO(aryan): this is only required for the time being because we need to do the registrations
-    # for classes. If we do it eagerly, i.e. call the functions in global scope, we will get circular
-    # import errors because of the models imported in this file.
-    _is_registered = False
-
-    @classmethod
-    def register(cls, model_class: Type, metadata: TransformerBlockMetadata):
-        cls._register()
-        metadata._cls = model_class
-        cls._registry[model_class] = metadata
-
-    @classmethod
-    def get(cls, model_class: Type) -> TransformerBlockMetadata:
-        cls._register()
-        if model_class not in cls._registry:
-            raise ValueError(f"Model class {model_class} not registered.")
-        return cls._registry[model_class]
-
-    @classmethod
-    def _register(cls):
-        if cls._is_registered:
-            return
-        cls._is_registered = True
-        _register_transformer_blocks_metadata()
-
-
-def _register_attention_processors_metadata():
-    from ..models.attention_processor import AttnProcessor2_0
-    from ..models.transformers.transformer_cogview4 import CogView4AttnProcessor
-    from ..models.transformers.transformer_flux import FluxAttnProcessor
-    from ..models.transformers.transformer_hunyuanimage import HunyuanImageAttnProcessor
-    from ..models.transformers.transformer_qwenimage import QwenDoubleStreamAttnProcessor2_0
-    from ..models.transformers.transformer_wan import WanAttnProcessor2_0
-    from ..models.transformers.transformer_z_image import ZSingleStreamAttnProcessor
-
-    # AttnProcessor2_0
-    AttentionProcessorRegistry.register(
-        model_class=AttnProcessor2_0,
-        metadata=AttentionProcessorMetadata(
-            skip_processor_output_fn=_skip_proc_output_fn_Attention_AttnProcessor2_0,
-        ),
-    )
-
-    # CogView4AttnProcessor
-    AttentionProcessorRegistry.register(
-        model_class=CogView4AttnProcessor,
-        metadata=AttentionProcessorMetadata(
-            skip_processor_output_fn=_skip_proc_output_fn_Attention_CogView4AttnProcessor,
-        ),
-    )
-
-    # WanAttnProcessor2_0
-    AttentionProcessorRegistry.register(
-        model_class=WanAttnProcessor2_0,
-        metadata=AttentionProcessorMetadata(
-            skip_processor_output_fn=_skip_proc_output_fn_Attention_WanAttnProcessor2_0,
-        ),
-    )
-
-    # FluxAttnProcessor
-    AttentionProcessorRegistry.register(
-        model_class=FluxAttnProcessor,
-        metadata=AttentionProcessorMetadata(skip_processor_output_fn=_skip_proc_output_fn_Attention_FluxAttnProcessor),
-    )
-
-    # QwenDoubleStreamAttnProcessor2
-    AttentionProcessorRegistry.register(
-        model_class=QwenDoubleStreamAttnProcessor2_0,
-        metadata=AttentionProcessorMetadata(
-            skip_processor_output_fn=_skip_proc_output_fn_Attention_QwenDoubleStreamAttnProcessor2_0
-        ),
-    )
-
-    # HunyuanImageAttnProcessor
-    AttentionProcessorRegistry.register(
-        model_class=HunyuanImageAttnProcessor,
-        metadata=AttentionProcessorMetadata(
-            skip_processor_output_fn=_skip_proc_output_fn_Attention_HunyuanImageAttnProcessor,
-        ),
-    )
-
-    # ZSingleStreamAttnProcessor
-    AttentionProcessorRegistry.register(
-        model_class=ZSingleStreamAttnProcessor,
-        metadata=AttentionProcessorMetadata(
-            skip_processor_output_fn=_skip_proc_output_fn_Attention_ZSingleStreamAttnProcessor,
-        ),
-    )
-
-
-def _register_transformer_blocks_metadata():
-    from ..models.attention import BasicTransformerBlock, JointTransformerBlock
-    from ..models.transformers.cogvideox_transformer_3d import CogVideoXBlock
-    from ..models.transformers.transformer_bria import BriaTransformerBlock
-    from ..models.transformers.transformer_cogview4 import CogView4TransformerBlock
-    from ..models.transformers.transformer_flux import FluxSingleTransformerBlock, FluxTransformerBlock
-    from ..models.transformers.transformer_hunyuan_video import (
-        HunyuanVideoSingleTransformerBlock,
-        HunyuanVideoTokenReplaceSingleTransformerBlock,
-        HunyuanVideoTokenReplaceTransformerBlock,
-        HunyuanVideoTransformerBlock,
-    )
-    from ..models.transformers.transformer_hunyuanimage import (
-        HunyuanImageSingleTransformerBlock,
-        HunyuanImageTransformerBlock,
-    )
-    from ..models.transformers.transformer_kandinsky import Kandinsky5TransformerDecoderBlock
-    from ..models.transformers.transformer_ltx import LTXVideoTransformerBlock
-    from ..models.transformers.transformer_mochi import MochiTransformerBlock
-    from ..models.transformers.transformer_motif_video import (
-        MotifVideoSingleTransformerBlock,
-        MotifVideoTransformerBlock,
-    )
-    from ..models.transformers.transformer_qwenimage import QwenImageTransformerBlock
-    from ..models.transformers.transformer_wan import WanTransformerBlock
-    from ..models.transformers.transformer_z_image import ZImageTransformerBlock
-
-    # BasicTransformerBlock
-    TransformerBlockRegistry.register(
-        model_class=BasicTransformerBlock,
-        metadata=TransformerBlockMetadata(
-            return_hidden_states_index=0,
-            return_encoder_hidden_states_index=None,
-        ),
-    )
-    TransformerBlockRegistry.register(
-        model_class=BriaTransformerBlock,
-        metadata=TransformerBlockMetadata(
-            return_hidden_states_index=0,
-            return_encoder_hidden_states_index=None,
-        ),
-    )
-
-    # CogVideoX
-    TransformerBlockRegistry.register(
-        model_class=CogVideoXBlock,
-        metadata=TransformerBlockMetadata(
-            return_hidden_states_index=0,
-            return_encoder_hidden_states_index=1,
-        ),
-    )
-
-    # CogView4
-    TransformerBlockRegistry.register(
-        model_class=CogView4TransformerBlock,
-        metadata=TransformerBlockMetadata(
-            return_hidden_states_index=0,
-            return_encoder_hidden_states_index=1,
-        ),
-    )
-
-    # Flux
-    TransformerBlockRegistry.register(
-        model_class=FluxTransformerBlock,
-        metadata=TransformerBlockMetadata(
-            return_hidden_states_index=1,
-            return_encoder_hidden_states_index=0,
-        ),
-    )
-    TransformerBlockRegistry.register(
-        model_class=FluxSingleTransformerBlock,
-        metadata=TransformerBlockMetadata(
-            return_hidden_states_index=1,
-            return_encoder_hidden_states_index=0,
-        ),
-    )
-
-    # HunyuanVideo
-    TransformerBlockRegistry.register(
-        model_class=HunyuanVideoTransformerBlock,
-        metadata=TransformerBlockMetadata(
-            return_hidden_states_index=0,
-            return_encoder_hidden_states_index=1,
-        ),
-    )
-    TransformerBlockRegistry.register(
-        model_class=HunyuanVideoSingleTransformerBlock,
-        metadata=TransformerBlockMetadata(
-            return_hidden_states_index=0,
-            return_encoder_hidden_states_index=1,
-        ),
-    )
-    TransformerBlockRegistry.register(
-        model_class=HunyuanVideoTokenReplaceTransformerBlock,
-        metadata=TransformerBlockMetadata(
-            return_hidden_states_index=0,
-            return_encoder_hidden_states_index=1,
-        ),
-    )
-    TransformerBlockRegistry.register(
-        model_class=HunyuanVideoTokenReplaceSingleTransformerBlock,
-        metadata=TransformerBlockMetadata(
-            return_hidden_states_index=0,
-            return_encoder_hidden_states_index=1,
-        ),
-    )
-
-    # LTXVideo
-    TransformerBlockRegistry.register(
-        model_class=LTXVideoTransformerBlock,
-        metadata=TransformerBlockMetadata(
-            return_hidden_states_index=0,
-            return_encoder_hidden_states_index=None,
-        ),
-    )
-
-    # Mochi
-    TransformerBlockRegistry.register(
-        model_class=MochiTransformerBlock,
-        metadata=TransformerBlockMetadata(
-            return_hidden_states_index=0,
-            return_encoder_hidden_states_index=1,
-        ),
-    )
-
-    # MotifVideo
-    TransformerBlockRegistry.register(
-        model_class=MotifVideoTransformerBlock,
-        metadata=TransformerBlockMetadata(
-            return_hidden_states_index=0,
-            return_encoder_hidden_states_index=1,
-        ),
-    )
-    TransformerBlockRegistry.register(
-        model_class=MotifVideoSingleTransformerBlock,
-        metadata=TransformerBlockMetadata(
-            return_hidden_states_index=0,
-            return_encoder_hidden_states_index=1,
-        ),
-    )
-
-    # Wan
-    TransformerBlockRegistry.register(
-        model_class=WanTransformerBlock,
-        metadata=TransformerBlockMetadata(
-            return_hidden_states_index=0,
-            return_encoder_hidden_states_index=None,
-        ),
-    )
-
-    # QwenImage
-    TransformerBlockRegistry.register(
-        model_class=QwenImageTransformerBlock,
-        metadata=TransformerBlockMetadata(
-            return_hidden_states_index=1,
-            return_encoder_hidden_states_index=0,
-        ),
-    )
-
-    # HunyuanImage2.1
-    TransformerBlockRegistry.register(
-        model_class=HunyuanImageTransformerBlock,
-        metadata=TransformerBlockMetadata(
-            return_hidden_states_index=0,
-            return_encoder_hidden_states_index=1,
-        ),
-    )
-    TransformerBlockRegistry.register(
-        model_class=HunyuanImageSingleTransformerBlock,
-        metadata=TransformerBlockMetadata(
-            return_hidden_states_index=0,
-            return_encoder_hidden_states_index=1,
-        ),
-    )
-
-    # ZImage
-    TransformerBlockRegistry.register(
-        model_class=ZImageTransformerBlock,
-        metadata=TransformerBlockMetadata(
-            return_hidden_states_index=0,
-            return_encoder_hidden_states_index=None,
-        ),
-    )
-
-    TransformerBlockRegistry.register(
-        model_class=JointTransformerBlock,
-        metadata=TransformerBlockMetadata(
-            return_hidden_states_index=1,
-            return_encoder_hidden_states_index=0,
-        ),
-    )
-
-    # Kandinsky 5.0 (Kandinsky5TransformerDecoderBlock)
-    TransformerBlockRegistry.register(
-        model_class=Kandinsky5TransformerDecoderBlock,
-        metadata=TransformerBlockMetadata(
-            return_hidden_states_index=0,
-            return_encoder_hidden_states_index=None,
-            hidden_states_argument_name="visual_embed",
-        ),
-    )
-
-
-# fmt: off
-def _skip_attention___ret___hidden_states(self, *args, **kwargs):
-    hidden_states = kwargs.get("hidden_states", None)
-    if hidden_states is None and len(args) > 0:
-        hidden_states = args[0]
-    return hidden_states
-
-
-def _skip_attention___ret___hidden_states___encoder_hidden_states(self, *args, **kwargs):
-    hidden_states = kwargs.get("hidden_states", None)
-    encoder_hidden_states = kwargs.get("encoder_hidden_states", None)
-    if hidden_states is None and len(args) > 0:
-        hidden_states = args[0]
-    if encoder_hidden_states is None and len(args) > 1:
-        encoder_hidden_states = args[1]
-    return hidden_states, encoder_hidden_states
-
-
-_skip_proc_output_fn_Attention_AttnProcessor2_0 = _skip_attention___ret___hidden_states
-_skip_proc_output_fn_Attention_CogView4AttnProcessor = _skip_attention___ret___hidden_states___encoder_hidden_states
-_skip_proc_output_fn_Attention_WanAttnProcessor2_0 = _skip_attention___ret___hidden_states
-# not sure what this is yet.
-_skip_proc_output_fn_Attention_FluxAttnProcessor = _skip_attention___ret___hidden_states
-_skip_proc_output_fn_Attention_QwenDoubleStreamAttnProcessor2_0 = _skip_attention___ret___hidden_states
-_skip_proc_output_fn_Attention_HunyuanImageAttnProcessor = _skip_attention___ret___hidden_states
-_skip_proc_output_fn_Attention_ZSingleStreamAttnProcessor = _skip_attention___ret___hidden_states
-# fmt: on

diffusers/hooks/context_parallel.py

@@ -1,382 +0,0 @@
-# Copyright 2026 The HuggingFace Team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-import copy
-import inspect
-from dataclasses import dataclass
-from typing import Type
-
-import torch
-import torch.distributed as dist
-
-
-if torch.distributed.is_available():
-    import torch.distributed._functional_collectives as funcol
-
-from ..models._modeling_parallel import (
-    ContextParallelConfig,
-    ContextParallelInput,
-    ContextParallelModelPlan,
-    ContextParallelOutput,
-    gather_size_by_comm,
-)
-from ..utils import get_logger
-from ..utils.torch_utils import lru_cache_unless_export, maybe_allow_in_graph, unwrap_module
-from .hooks import HookRegistry, ModelHook
-
-
-logger = get_logger(__name__)  # pylint: disable=invalid-name
-
-_CONTEXT_PARALLEL_INPUT_HOOK_TEMPLATE = "cp_input---{}"
-_CONTEXT_PARALLEL_OUTPUT_HOOK_TEMPLATE = "cp_output---{}"
-
-
-# TODO(aryan): consolidate with ._helpers.TransformerBlockMetadata
-@dataclass
-class ModuleForwardMetadata:
-    cached_parameter_indices: dict[str, int] = None
-    _cls: Type = None
-
-    def _get_parameter_from_args_kwargs(self, identifier: str, args=(), kwargs=None):
-        kwargs = kwargs or {}
-
-        if identifier in kwargs:
-            return kwargs[identifier], True, None
-
-        if self.cached_parameter_indices is not None:
-            index = self.cached_parameter_indices.get(identifier, None)
-            if index is None:
-                raise ValueError(f"Parameter '{identifier}' not found in cached indices.")
-            return args[index], False, index
-
-        if self._cls is None:
-            raise ValueError("Model class is not set for metadata.")
-
-        parameters = list(inspect.signature(self._cls.forward).parameters.keys())
-        parameters = parameters[1:]  # skip `self`
-        self.cached_parameter_indices = {param: i for i, param in enumerate(parameters)}
-
-        if identifier not in self.cached_parameter_indices:
-            raise ValueError(f"Parameter '{identifier}' not found in function signature but was requested.")
-
-        index = self.cached_parameter_indices[identifier]
-
-        if index >= len(args):
-            raise ValueError(f"Expected {index} arguments but got {len(args)}.")
-
-        return args[index], False, index
-
-
-def apply_context_parallel(
-    module: torch.nn.Module,
-    parallel_config: ContextParallelConfig,
-    plan: dict[str, ContextParallelModelPlan],
-) -> None:
-    """Apply context parallel on a model."""
-    logger.debug(f"Applying context parallel with CP mesh: {parallel_config._mesh} and plan: {plan}")
-
-    for module_id, cp_model_plan in plan.items():
-        submodule = _get_submodule_by_name(module, module_id)
-        if not isinstance(submodule, list):
-            submodule = [submodule]
-
-        logger.debug(f"Applying ContextParallelHook to {module_id=} identifying a total of {len(submodule)} modules")
-
-        for m in submodule:
-            if isinstance(cp_model_plan, dict):
-                hook = ContextParallelSplitHook(cp_model_plan, parallel_config)
-                hook_name = _CONTEXT_PARALLEL_INPUT_HOOK_TEMPLATE.format(module_id)
-            elif isinstance(cp_model_plan, (ContextParallelOutput, list, tuple)):
-                if isinstance(cp_model_plan, ContextParallelOutput):
-                    cp_model_plan = [cp_model_plan]
-                if not all(isinstance(x, ContextParallelOutput) for x in cp_model_plan):
-                    raise ValueError(f"Expected all elements of cp_model_plan to be CPOutput, but got {cp_model_plan}")
-                hook = ContextParallelGatherHook(cp_model_plan, parallel_config)
-                hook_name = _CONTEXT_PARALLEL_OUTPUT_HOOK_TEMPLATE.format(module_id)
-            else:
-                raise ValueError(f"Unsupported context parallel model plan type: {type(cp_model_plan)}")
-            registry = HookRegistry.check_if_exists_or_initialize(m)
-            registry.register_hook(hook, hook_name)
-
-
-def remove_context_parallel(module: torch.nn.Module, plan: dict[str, ContextParallelModelPlan]) -> None:
-    for module_id, cp_model_plan in plan.items():
-        submodule = _get_submodule_by_name(module, module_id)
-        if not isinstance(submodule, list):
-            submodule = [submodule]
-
-        for m in submodule:
-            registry = HookRegistry.check_if_exists_or_initialize(m)
-            if isinstance(cp_model_plan, dict):
-                hook_name = _CONTEXT_PARALLEL_INPUT_HOOK_TEMPLATE.format(module_id)
-            elif isinstance(cp_model_plan, (ContextParallelOutput, list, tuple)):
-                hook_name = _CONTEXT_PARALLEL_OUTPUT_HOOK_TEMPLATE.format(module_id)
-            else:
-                raise ValueError(f"Unsupported context parallel model plan type: {type(cp_model_plan)}")
-            registry.remove_hook(hook_name)
-
-
-class ContextParallelSplitHook(ModelHook):
-    def __init__(self, metadata: ContextParallelModelPlan, parallel_config: ContextParallelConfig) -> None:
-        super().__init__()
-        self.metadata = metadata
-        self.parallel_config = parallel_config
-        self.module_forward_metadata = None
-
-    def initialize_hook(self, module):
-        cls = unwrap_module(module).__class__
-        self.module_forward_metadata = ModuleForwardMetadata(_cls=cls)
-        return module
-
-    def pre_forward(self, module, *args, **kwargs):
-        args_list = list(args)
-
-        for name, cpm in self.metadata.items():
-            if isinstance(cpm, ContextParallelInput) and cpm.split_output:
-                continue
-
-            # Maybe the parameter was passed as a keyword argument
-            input_val, is_kwarg, index = self.module_forward_metadata._get_parameter_from_args_kwargs(
-                name, args_list, kwargs
-            )
-
-            if input_val is None:
-                continue
-
-            # The input_val may be a tensor or list/tuple of tensors. In certain cases, user may specify to shard
-            # the output instead of input for a particular layer by setting split_output=True
-            if isinstance(input_val, torch.Tensor):
-                input_val = self._prepare_cp_input(input_val, cpm)
-            elif isinstance(input_val, (list, tuple)):
-                if len(input_val) != len(cpm):
-                    raise ValueError(
-                        f"Expected input model plan to have {len(input_val)} elements, but got {len(cpm)}."
-                    )
-                sharded_input_val = []
-                for i, x in enumerate(input_val):
-                    if torch.is_tensor(x) and not cpm[i].split_output:
-                        x = self._prepare_cp_input(x, cpm[i])
-                    sharded_input_val.append(x)
-                input_val = sharded_input_val
-            else:
-                raise ValueError(f"Unsupported input type: {type(input_val)}")
-
-            if is_kwarg:
-                kwargs[name] = input_val
-            elif index is not None and index < len(args_list):
-                args_list[index] = input_val
-            else:
-                raise ValueError(
-                    f"An unexpected error occurred while processing the input '{name}'. Please open an "
-                    f"issue at https://github.com/huggingface/diffusers/issues and provide a minimal reproducible "
-                    f"example along with the full stack trace."
-                )
-
-        return tuple(args_list), kwargs
-
-    def post_forward(self, module, output):
-        is_tensor = isinstance(output, torch.Tensor)
-        is_tensor_list = isinstance(output, (list, tuple)) and all(isinstance(x, torch.Tensor) for x in output)
-
-        if not is_tensor and not is_tensor_list:
-            raise ValueError(f"Expected output to be a tensor or a list/tuple of tensors, but got {type(output)}.")
-
-        output = [output] if is_tensor else list(output)
-        for index, cpm in self.metadata.items():
-            if not isinstance(cpm, ContextParallelInput) or not cpm.split_output:
-                continue
-            if index >= len(output):
-                raise ValueError(f"Index {index} out of bounds for output of length {len(output)}.")
-            current_output = output[index]
-            current_output = self._prepare_cp_input(current_output, cpm)
-            output[index] = current_output
-
-        return output[0] if is_tensor else tuple(output)
-
-    def _prepare_cp_input(self, x: torch.Tensor, cp_input: ContextParallelInput) -> torch.Tensor:
-        if cp_input.expected_dims is not None and x.dim() != cp_input.expected_dims:
-            logger.warning_once(
-                f"Expected input tensor to have {cp_input.expected_dims} dimensions, but got {x.dim()} dimensions, split will not be applied."
-            )
-            return x
-        else:
-            if self.parallel_config.ulysses_anything or self.parallel_config.ring_anything:
-                return PartitionAnythingSharder.shard_anything(
-                    x, cp_input.split_dim, self.parallel_config._flattened_mesh
-                )
-            return EquipartitionSharder.shard(x, cp_input.split_dim, self.parallel_config._flattened_mesh)
-
-
-class ContextParallelGatherHook(ModelHook):
-    def __init__(self, metadata: ContextParallelModelPlan, parallel_config: ContextParallelConfig) -> None:
-        super().__init__()
-        self.metadata = metadata
-        self.parallel_config = parallel_config
-
-    def post_forward(self, module, output):
-        is_tensor = isinstance(output, torch.Tensor)
-
-        if is_tensor:
-            output = [output]
-        elif not (isinstance(output, (list, tuple)) and all(isinstance(x, torch.Tensor) for x in output)):
-            raise ValueError(f"Expected output to be a tensor or a list/tuple of tensors, but got {type(output)}.")
-
-        output = list(output)
-
-        if len(output) != len(self.metadata):
-            raise ValueError(f"Expected output to have {len(self.metadata)} elements, but got {len(output)}.")
-
-        for i, cpm in enumerate(self.metadata):
-            if cpm is None:
-                continue
-            if self.parallel_config.ulysses_anything or self.parallel_config.ring_anything:
-                output[i] = PartitionAnythingSharder.unshard_anything(
-                    output[i], cpm.gather_dim, self.parallel_config._flattened_mesh
-                )
-            else:
-                output[i] = EquipartitionSharder.unshard(
-                    output[i], cpm.gather_dim, self.parallel_config._flattened_mesh
-                )
-
-        return output[0] if is_tensor else tuple(output)
-
-
-class AllGatherFunction(torch.autograd.Function):
-    @staticmethod
-    def forward(ctx, tensor, dim, group):
-        ctx.dim = dim
-        ctx.group = group
-        ctx.world_size = torch.distributed.get_world_size(group)
-        ctx.rank = torch.distributed.get_rank(group)
-        return funcol.all_gather_tensor(tensor, dim, group=group)
-
-    @staticmethod
-    def backward(ctx, grad_output):
-        grad_chunks = torch.chunk(grad_output, ctx.world_size, dim=ctx.dim)
-        return grad_chunks[ctx.rank], None, None
-
-
-class EquipartitionSharder:
-    @classmethod
-    def shard(cls, tensor: torch.Tensor, dim: int, mesh: torch.distributed.device_mesh.DeviceMesh) -> torch.Tensor:
-        # NOTE: the following assertion does not have to be true in general. We simply enforce it for now
-        # because the alternate case has not yet been tested/required for any model.
-        assert tensor.size()[dim] % mesh.size() == 0, (
-            "Tensor size along dimension to be sharded must be divisible by mesh size"
-        )
-
-        # The following is not fullgraph compatible with Dynamo (fails in DeviceMesh.get_rank)
-        # return tensor.chunk(mesh.size(), dim=dim)[mesh.get_rank()]
-
-        return tensor.chunk(mesh.size(), dim=dim)[torch.distributed.get_rank(mesh.get_group())]
-
-    @classmethod
-    def unshard(cls, tensor: torch.Tensor, dim: int, mesh: torch.distributed.device_mesh.DeviceMesh) -> torch.Tensor:
-        tensor = tensor.contiguous()
-        tensor = AllGatherFunction.apply(tensor, dim, mesh.get_group())
-        return tensor
-
-
-class AllGatherAnythingFunction(torch.autograd.Function):
-    @staticmethod
-    def forward(ctx, tensor: torch.Tensor, dim: int, group: dist.device_mesh.DeviceMesh):
-        ctx.dim = dim
-        ctx.group = group
-        ctx.world_size = dist.get_world_size(group)
-        ctx.rank = dist.get_rank(group)
-        gathered_tensor = _all_gather_anything(tensor, dim, group)
-        return gathered_tensor
-
-    @staticmethod
-    def backward(ctx, grad_output):
-        # NOTE: We use `tensor_split` instead of chunk, because the `chunk`
-        # function may return fewer than the specified number of chunks!
-        grad_splits = torch.tensor_split(grad_output, ctx.world_size, dim=ctx.dim)
-        return grad_splits[ctx.rank], None, None
-
-
-class PartitionAnythingSharder:
-    @classmethod
-    def shard_anything(
-        cls, tensor: torch.Tensor, dim: int, mesh: torch.distributed.device_mesh.DeviceMesh
-    ) -> torch.Tensor:
-        assert tensor.size()[dim] >= mesh.size(), (
-            f"Cannot shard tensor of size {tensor.size()} along dim {dim} across mesh of size {mesh.size()}."
-        )
-        # NOTE: We use `tensor_split` instead of chunk, because the `chunk`
-        # function may return fewer than the specified number of chunks!
-        return tensor.tensor_split(mesh.size(), dim=dim)[dist.get_rank(mesh.get_group())]
-
-    @classmethod
-    def unshard_anything(
-        cls, tensor: torch.Tensor, dim: int, mesh: torch.distributed.device_mesh.DeviceMesh
-    ) -> torch.Tensor:
-        tensor = tensor.contiguous()
-        tensor = AllGatherAnythingFunction.apply(tensor, dim, mesh.get_group())
-        return tensor
-
-
-@lru_cache_unless_export(maxsize=64)
-def _fill_gather_shapes(shape: tuple[int], gather_dims: tuple[int], dim: int, world_size: int) -> list[list[int]]:
-    gather_shapes = []
-    for i in range(world_size):
-        rank_shape = list(copy.deepcopy(shape))
-        rank_shape[dim] = gather_dims[i]
-        gather_shapes.append(rank_shape)
-    return gather_shapes
-
-
-@maybe_allow_in_graph
-def _all_gather_anything(tensor: torch.Tensor, dim: int, group: dist.device_mesh.DeviceMesh) -> torch.Tensor:
-    world_size = dist.get_world_size(group=group)
-
-    tensor = tensor.contiguous()
-    shape = tensor.shape
-    rank_dim = shape[dim]
-    gather_dims = gather_size_by_comm(rank_dim, group)
-
-    gather_shapes = _fill_gather_shapes(tuple(shape), tuple(gather_dims), dim, world_size)
-
-    gathered_tensors = [torch.empty(shape, device=tensor.device, dtype=tensor.dtype) for shape in gather_shapes]
-
-    dist.all_gather(gathered_tensors, tensor, group=group)
-    gathered_tensor = torch.cat(gathered_tensors, dim=dim)
-    return gathered_tensor
-
-
-def _get_submodule_by_name(model: torch.nn.Module, name: str) -> torch.nn.Module | list[torch.nn.Module]:
-    if name.count("*") > 1:
-        raise ValueError("Wildcard '*' can only be used once in the name")
-    return _find_submodule_by_name(model, name)
-
-
-def _find_submodule_by_name(model: torch.nn.Module, name: str) -> torch.nn.Module | list[torch.nn.Module]:
-    if name == "":
-        return model
-    first_atom, remaining_name = name.split(".", 1) if "." in name else (name, "")
-    if first_atom == "*":
-        if not isinstance(model, torch.nn.ModuleList):
-            raise ValueError("Wildcard '*' can only be used with ModuleList")
-        submodules = []
-        for submodule in model:
-            subsubmodules = _find_submodule_by_name(submodule, remaining_name)
-            if not isinstance(subsubmodules, list):
-                subsubmodules = [subsubmodules]
-            submodules.extend(subsubmodules)
-        return submodules
-    else:
-        if hasattr(model, first_atom):
-            submodule = getattr(model, first_atom)
-            return _find_submodule_by_name(submodule, remaining_name)
-        else:
-            raise ValueError(f"'{first_atom}' is not a submodule of '{model.__class__.__name__}'")

diffusers/hooks/faster_cache.py

@@ -1,654 +0,0 @@
-# Copyright 2026 The HuggingFace Team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-import re
-from dataclasses import dataclass
-from typing import Any, Callable
-
-import torch
-
-from ..models.attention import AttentionModuleMixin
-from ..models.modeling_outputs import Transformer2DModelOutput
-from ..utils import logging
-from ._common import _ATTENTION_CLASSES
-from .hooks import HookRegistry, ModelHook
-
-
-logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
-
-
-_FASTER_CACHE_DENOISER_HOOK = "faster_cache_denoiser"
-_FASTER_CACHE_BLOCK_HOOK = "faster_cache_block"
-_SPATIAL_ATTENTION_BLOCK_IDENTIFIERS = (
-    "^blocks.*attn",
-    "^transformer_blocks.*attn",
-    "^single_transformer_blocks.*attn",
-)
-_TEMPORAL_ATTENTION_BLOCK_IDENTIFIERS = ("^temporal_transformer_blocks.*attn",)
-_TRANSFORMER_BLOCK_IDENTIFIERS = _SPATIAL_ATTENTION_BLOCK_IDENTIFIERS + _TEMPORAL_ATTENTION_BLOCK_IDENTIFIERS
-_UNCOND_COND_INPUT_KWARGS_IDENTIFIERS = (
-    "hidden_states",
-    "encoder_hidden_states",
-    "timestep",
-    "attention_mask",
-    "encoder_attention_mask",
-)
-
-
-@dataclass
-class FasterCacheConfig:
-    r"""
-    Configuration for [FasterCache](https://huggingface.co/papers/2410.19355).
-
-    Attributes:
-        spatial_attention_block_skip_range (`int`, defaults to `2`):
-            Calculate the attention states every `N` iterations. If this is set to `N`, the attention computation will
-            be skipped `N - 1` times (i.e., cached attention states will be reused) before computing the new attention
-            states again.
-        temporal_attention_block_skip_range (`int`, *optional*, defaults to `None`):
-            Calculate the attention states every `N` iterations. If this is set to `N`, the attention computation will
-            be skipped `N - 1` times (i.e., cached attention states will be reused) before computing the new attention
-            states again.
-        spatial_attention_timestep_skip_range (`tuple[float, float]`, defaults to `(-1, 681)`):
-            The timestep range within which the spatial attention computation can be skipped without a significant loss
-            in quality. This is to be determined by the user based on the underlying model. The first value in the
-            tuple is the lower bound and the second value is the upper bound. Typically, diffusion timesteps for
-            denoising are in the reversed range of 0 to 1000 (i.e. denoising starts at timestep 1000 and ends at
-            timestep 0). For the default values, this would mean that the spatial attention computation skipping will
-            be applicable only after denoising timestep 681 is reached, and continue until the end of the denoising
-            process.
-        temporal_attention_timestep_skip_range (`tuple[float, float]`, *optional*, defaults to `None`):
-            The timestep range within which the temporal attention computation can be skipped without a significant
-            loss in quality. This is to be determined by the user based on the underlying model. The first value in the
-            tuple is the lower bound and the second value is the upper bound. Typically, diffusion timesteps for
-            denoising are in the reversed range of 0 to 1000 (i.e. denoising starts at timestep 1000 and ends at
-            timestep 0).
-        low_frequency_weight_update_timestep_range (`tuple[int, int]`, defaults to `(99, 901)`):
-            The timestep range within which the low frequency weight scaling update is applied. The first value in the
-            tuple is the lower bound and the second value is the upper bound of the timestep range. The callback
-            function for the update is called only within this range.
-        high_frequency_weight_update_timestep_range (`tuple[int, int]`, defaults to `(-1, 301)`):
-            The timestep range within which the high frequency weight scaling update is applied. The first value in the
-            tuple is the lower bound and the second value is the upper bound of the timestep range. The callback
-            function for the update is called only within this range.
-        alpha_low_frequency (`float`, defaults to `1.1`):
-            The weight to scale the low frequency updates by. This is used to approximate the unconditional branch from
-            the conditional branch outputs.
-        alpha_high_frequency (`float`, defaults to `1.1`):
-            The weight to scale the high frequency updates by. This is used to approximate the unconditional branch
-            from the conditional branch outputs.
-        unconditional_batch_skip_range (`int`, defaults to `5`):
-            Process the unconditional branch every `N` iterations. If this is set to `N`, the unconditional branch
-            computation will be skipped `N - 1` times (i.e., cached unconditional branch states will be reused) before
-            computing the new unconditional branch states again.
-        unconditional_batch_timestep_skip_range (`tuple[float, float]`, defaults to `(-1, 641)`):
-            The timestep range within which the unconditional branch computation can be skipped without a significant
-            loss in quality. This is to be determined by the user based on the underlying model. The first value in the
-            tuple is the lower bound and the second value is the upper bound.
-        spatial_attention_block_identifiers (`tuple[str, ...]`, defaults to `("blocks.*attn1", "transformer_blocks.*attn1", "single_transformer_blocks.*attn1")`):
-            The identifiers to match the spatial attention blocks in the model. If the name of the block contains any
-            of these identifiers, FasterCache will be applied to that block. This can either be the full layer names,
-            partial layer names, or regex patterns. Matching will always be done using a regex match.
-        temporal_attention_block_identifiers (`tuple[str, ...]`, defaults to `("temporal_transformer_blocks.*attn1",)`):
-            The identifiers to match the temporal attention blocks in the model. If the name of the block contains any
-            of these identifiers, FasterCache will be applied to that block. This can either be the full layer names,
-            partial layer names, or regex patterns. Matching will always be done using a regex match.
-        attention_weight_callback (`Callable[[torch.nn.Module], float]`, defaults to `None`):
-            The callback function to determine the weight to scale the attention outputs by. This function should take
-            the attention module as input and return a float value. This is used to approximate the unconditional
-            branch from the conditional branch outputs. If not provided, the default weight is 0.5 for all timesteps.
-            Typically, as described in the paper, this weight should gradually increase from 0 to 1 as the inference
-            progresses. Users are encouraged to experiment and provide custom weight schedules that take into account
-            the number of inference steps and underlying model behaviour as denoising progresses.
-        low_frequency_weight_callback (`Callable[[torch.nn.Module], float]`, defaults to `None`):
-            The callback function to determine the weight to scale the low frequency updates by. If not provided, the
-            default weight is 1.1 for timesteps within the range specified (as described in the paper).
-        high_frequency_weight_callback (`Callable[[torch.nn.Module], float]`, defaults to `None`):
-            The callback function to determine the weight to scale the high frequency updates by. If not provided, the
-            default weight is 1.1 for timesteps within the range specified (as described in the paper).
-        tensor_format (`str`, defaults to `"BCFHW"`):
-            The format of the input tensors. This should be one of `"BCFHW"`, `"BFCHW"`, or `"BCHW"`. The format is
-            used to split individual latent frames in order for low and high frequency components to be computed.
-        is_guidance_distilled (`bool`, defaults to `False`):
-            Whether the model is guidance distilled or not. If the model is guidance distilled, FasterCache will not be
-            applied at the denoiser-level to skip the unconditional branch computation (as there is none).
-        _unconditional_conditional_input_kwargs_identifiers (`list[str]`, defaults to `("hidden_states", "encoder_hidden_states", "timestep", "attention_mask", "encoder_attention_mask")`):
-            The identifiers to match the input kwargs that contain the batchwise-concatenated unconditional and
-            conditional inputs. If the name of the input kwargs contains any of these identifiers, FasterCache will
-            split the inputs into unconditional and conditional branches. This must be a list of exact input kwargs
-            names that contain the batchwise-concatenated unconditional and conditional inputs.
-    """
-
-    # In the paper and codebase, they hardcode these values to 2. However, it can be made configurable
-    # after some testing. We default to 2 if these parameters are not provided.
-    spatial_attention_block_skip_range: int = 2
-    temporal_attention_block_skip_range: int | None = None
-
-    spatial_attention_timestep_skip_range: tuple[int, int] = (-1, 681)
-    temporal_attention_timestep_skip_range: tuple[int, int] = (-1, 681)
-
-    # Indicator functions for low/high frequency as mentioned in Equation 11 of the paper
-    low_frequency_weight_update_timestep_range: tuple[int, int] = (99, 901)
-    high_frequency_weight_update_timestep_range: tuple[int, int] = (-1, 301)
-
-    # ⍺1 and ⍺2 as mentioned in Equation 11 of the paper
-    alpha_low_frequency: float = 1.1
-    alpha_high_frequency: float = 1.1
-
-    # n as described in CFG-Cache explanation in the paper - dependent on the model
-    unconditional_batch_skip_range: int = 5
-    unconditional_batch_timestep_skip_range: tuple[int, int] = (-1, 641)
-
-    spatial_attention_block_identifiers: tuple[str, ...] = _SPATIAL_ATTENTION_BLOCK_IDENTIFIERS
-    temporal_attention_block_identifiers: tuple[str, ...] = _TEMPORAL_ATTENTION_BLOCK_IDENTIFIERS
-
-    attention_weight_callback: Callable[[torch.nn.Module], float] = None
-    low_frequency_weight_callback: Callable[[torch.nn.Module], float] = None
-    high_frequency_weight_callback: Callable[[torch.nn.Module], float] = None
-
-    tensor_format: str = "BCFHW"
-    is_guidance_distilled: bool = False
-
-    current_timestep_callback: Callable[[], int] = None
-
-    _unconditional_conditional_input_kwargs_identifiers: list[str] = _UNCOND_COND_INPUT_KWARGS_IDENTIFIERS
-
-    def __repr__(self) -> str:
-        return (
-            f"FasterCacheConfig(\n"
-            f"  spatial_attention_block_skip_range={self.spatial_attention_block_skip_range},\n"
-            f"  temporal_attention_block_skip_range={self.temporal_attention_block_skip_range},\n"
-            f"  spatial_attention_timestep_skip_range={self.spatial_attention_timestep_skip_range},\n"
-            f"  temporal_attention_timestep_skip_range={self.temporal_attention_timestep_skip_range},\n"
-            f"  low_frequency_weight_update_timestep_range={self.low_frequency_weight_update_timestep_range},\n"
-            f"  high_frequency_weight_update_timestep_range={self.high_frequency_weight_update_timestep_range},\n"
-            f"  alpha_low_frequency={self.alpha_low_frequency},\n"
-            f"  alpha_high_frequency={self.alpha_high_frequency},\n"
-            f"  unconditional_batch_skip_range={self.unconditional_batch_skip_range},\n"
-            f"  unconditional_batch_timestep_skip_range={self.unconditional_batch_timestep_skip_range},\n"
-            f"  spatial_attention_block_identifiers={self.spatial_attention_block_identifiers},\n"
-            f"  temporal_attention_block_identifiers={self.temporal_attention_block_identifiers},\n"
-            f"  tensor_format={self.tensor_format},\n"
-            f")"
-        )
-
-
-class FasterCacheDenoiserState:
-    r"""
-    State for [FasterCache](https://huggingface.co/papers/2410.19355) top-level denoiser module.
-    """
-
-    def __init__(self) -> None:
-        self.iteration: int = 0
-        self.low_frequency_delta: torch.Tensor = None
-        self.high_frequency_delta: torch.Tensor = None
-
-    def reset(self):
-        self.iteration = 0
-        self.low_frequency_delta = None
-        self.high_frequency_delta = None
-
-
-class FasterCacheBlockState:
-    r"""
-    State for [FasterCache](https://huggingface.co/papers/2410.19355). Every underlying block that FasterCache is
-    applied to will have an instance of this state.
-    """
-
-    def __init__(self) -> None:
-        self.iteration: int = 0
-        self.batch_size: int = None
-        self.cache: tuple[torch.Tensor, torch.Tensor] = None
-
-    def reset(self):
-        self.iteration = 0
-        self.batch_size = None
-        self.cache = None
-
-
-class FasterCacheDenoiserHook(ModelHook):
-    _is_stateful = True
-
-    def __init__(
-        self,
-        unconditional_batch_skip_range: int,
-        unconditional_batch_timestep_skip_range: tuple[int, int],
-        tensor_format: str,
-        is_guidance_distilled: bool,
-        uncond_cond_input_kwargs_identifiers: list[str],
-        current_timestep_callback: Callable[[], int],
-        low_frequency_weight_callback: Callable[[torch.nn.Module], torch.Tensor],
-        high_frequency_weight_callback: Callable[[torch.nn.Module], torch.Tensor],
-    ) -> None:
-        super().__init__()
-
-        self.unconditional_batch_skip_range = unconditional_batch_skip_range
-        self.unconditional_batch_timestep_skip_range = unconditional_batch_timestep_skip_range
-        # We can't easily detect what args are to be split in unconditional and conditional branches. We
-        # can only do it for kwargs, hence they are the only ones we split. The args are passed as-is.
-        # If a model is to be made compatible with FasterCache, the user must ensure that the inputs that
-        # contain batchwise-concatenated unconditional and conditional inputs are passed as kwargs.
-        self.uncond_cond_input_kwargs_identifiers = uncond_cond_input_kwargs_identifiers
-        self.tensor_format = tensor_format
-        self.is_guidance_distilled = is_guidance_distilled
-
-        self.current_timestep_callback = current_timestep_callback
-        self.low_frequency_weight_callback = low_frequency_weight_callback
-        self.high_frequency_weight_callback = high_frequency_weight_callback
-
-    def initialize_hook(self, module):
-        self.state = FasterCacheDenoiserState()
-        return module
-
-    @staticmethod
-    def _get_cond_input(input: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
-        # Note: this method assumes that the input tensor is batchwise-concatenated with unconditional inputs
-        # followed by conditional inputs.
-        _, cond = input.chunk(2, dim=0)
-        return cond
-
-    def new_forward(self, module: torch.nn.Module, *args, **kwargs) -> Any:
-        # Split the unconditional and conditional inputs. We only want to infer the conditional branch if the
-        # requirements for skipping the unconditional branch are met as described in the paper.
-        # We skip the unconditional branch only if the following conditions are met:
-        #   1. We have completed at least one iteration of the denoiser
-        #   2. The current timestep is within the range specified by the user. This is the optimal timestep range
-        #      where approximating the unconditional branch from the computation of the conditional branch is possible
-        #      without a significant loss in quality.
-        #   3. The current iteration is not a multiple of the unconditional batch skip range. This is done so that
-        #      we compute the unconditional branch at least once every few iterations to ensure minimal quality loss.
-        is_within_timestep_range = (
-            self.unconditional_batch_timestep_skip_range[0]
-            < self.current_timestep_callback()
-            < self.unconditional_batch_timestep_skip_range[1]
-        )
-        should_skip_uncond = (
-            self.state.iteration > 0
-            and is_within_timestep_range
-            and self.state.iteration % self.unconditional_batch_skip_range != 0
-            and not self.is_guidance_distilled
-        )
-
-        if should_skip_uncond:
-            is_any_kwarg_uncond = any(k in self.uncond_cond_input_kwargs_identifiers for k in kwargs.keys())
-            if is_any_kwarg_uncond:
-                logger.debug("FasterCache - Skipping unconditional branch computation")
-                args = tuple([self._get_cond_input(arg) if torch.is_tensor(arg) else arg for arg in args])
-                kwargs = {
-                    k: v if k not in self.uncond_cond_input_kwargs_identifiers else self._get_cond_input(v)
-                    for k, v in kwargs.items()
-                }
-
-        output = self.fn_ref.original_forward(*args, **kwargs)
-
-        if self.is_guidance_distilled:
-            self.state.iteration += 1
-            return output
-
-        if torch.is_tensor(output):
-            hidden_states = output
-        elif isinstance(output, (tuple, Transformer2DModelOutput)):
-            hidden_states = output[0]
-
-        batch_size = hidden_states.size(0)
-
-        if should_skip_uncond:
-            self.state.low_frequency_delta = self.state.low_frequency_delta * self.low_frequency_weight_callback(
-                module
-            )
-            self.state.high_frequency_delta = self.state.high_frequency_delta * self.high_frequency_weight_callback(
-                module
-            )
-
-            if self.tensor_format == "BCFHW":
-                hidden_states = hidden_states.permute(0, 2, 1, 3, 4)
-            if self.tensor_format == "BCFHW" or self.tensor_format == "BFCHW":
-                hidden_states = hidden_states.flatten(0, 1)
-
-            low_freq_cond, high_freq_cond = _split_low_high_freq(hidden_states.float())
-
-            # Approximate/compute the unconditional branch outputs as described in Equation 9 and 10 of the paper
-            low_freq_uncond = self.state.low_frequency_delta + low_freq_cond
-            high_freq_uncond = self.state.high_frequency_delta + high_freq_cond
-            uncond_freq = low_freq_uncond + high_freq_uncond
-
-            uncond_states = torch.fft.ifftshift(uncond_freq)
-            uncond_states = torch.fft.ifft2(uncond_states).real
-
-            if self.tensor_format == "BCFHW" or self.tensor_format == "BFCHW":
-                uncond_states = uncond_states.unflatten(0, (batch_size, -1))
-                hidden_states = hidden_states.unflatten(0, (batch_size, -1))
-            if self.tensor_format == "BCFHW":
-                uncond_states = uncond_states.permute(0, 2, 1, 3, 4)
-                hidden_states = hidden_states.permute(0, 2, 1, 3, 4)
-
-            # Concatenate the approximated unconditional and predicted conditional branches
-            uncond_states = uncond_states.to(hidden_states.dtype)
-            hidden_states = torch.cat([uncond_states, hidden_states], dim=0)
-        else:
-            uncond_states, cond_states = hidden_states.chunk(2, dim=0)
-            if self.tensor_format == "BCFHW":
-                uncond_states = uncond_states.permute(0, 2, 1, 3, 4)
-                cond_states = cond_states.permute(0, 2, 1, 3, 4)
-            if self.tensor_format == "BCFHW" or self.tensor_format == "BFCHW":
-                uncond_states = uncond_states.flatten(0, 1)
-                cond_states = cond_states.flatten(0, 1)
-
-            low_freq_uncond, high_freq_uncond = _split_low_high_freq(uncond_states.float())
-            low_freq_cond, high_freq_cond = _split_low_high_freq(cond_states.float())
-            self.state.low_frequency_delta = low_freq_uncond - low_freq_cond
-            self.state.high_frequency_delta = high_freq_uncond - high_freq_cond
-
-        self.state.iteration += 1
-        if torch.is_tensor(output):
-            output = hidden_states
-        elif isinstance(output, tuple):
-            output = (hidden_states, *output[1:])
-        else:
-            output.sample = hidden_states
-
-        return output
-
-    def reset_state(self, module: torch.nn.Module) -> torch.nn.Module:
-        self.state.reset()
-        return module
-
-
-class FasterCacheBlockHook(ModelHook):
-    _is_stateful = True
-
-    def __init__(
-        self,
-        block_skip_range: int,
-        timestep_skip_range: tuple[int, int],
-        is_guidance_distilled: bool,
-        weight_callback: Callable[[torch.nn.Module], float],
-        current_timestep_callback: Callable[[], int],
-    ) -> None:
-        super().__init__()
-
-        self.block_skip_range = block_skip_range
-        self.timestep_skip_range = timestep_skip_range
-        self.is_guidance_distilled = is_guidance_distilled
-
-        self.weight_callback = weight_callback
-        self.current_timestep_callback = current_timestep_callback
-
-    def initialize_hook(self, module):
-        self.state = FasterCacheBlockState()
-        return module
-
-    def _compute_approximated_attention_output(
-        self, t_2_output: torch.Tensor, t_output: torch.Tensor, weight: float, batch_size: int
-    ) -> torch.Tensor:
-        if t_2_output.size(0) != batch_size:
-            # The cache t_2_output contains both batchwise-concatenated unconditional-conditional branch outputs. Just
-            # take the conditional branch outputs.
-            assert t_2_output.size(0) == 2 * batch_size
-            t_2_output = t_2_output[batch_size:]
-        if t_output.size(0) != batch_size:
-            # The cache t_output contains both batchwise-concatenated unconditional-conditional branch outputs. Just
-            # take the conditional branch outputs.
-            assert t_output.size(0) == 2 * batch_size
-            t_output = t_output[batch_size:]
-        return t_output + (t_output - t_2_output) * weight
-
-    def new_forward(self, module: torch.nn.Module, *args, **kwargs) -> Any:
-        batch_size = [
-            *[arg.size(0) for arg in args if torch.is_tensor(arg)],
-            *[v.size(0) for v in kwargs.values() if torch.is_tensor(v)],
-        ][0]
-        if self.state.batch_size is None:
-            # Will be updated on first forward pass through the denoiser
-            self.state.batch_size = batch_size
-
-        # If we have to skip due to the skip conditions, then let's skip as expected.
-        # But, we can't skip if the denoiser wants to infer both unconditional and conditional branches. This
-        # is because the expected output shapes of attention layer will not match if we only return values from
-        # the cache (which only caches conditional branch outputs). So, if state.batch_size (which is the true
-        # unconditional-conditional batch size) is same as the current batch size, we don't perform the layer
-        # skip. Otherwise, we conditionally skip the layer based on what state.skip_callback returns.
-        is_within_timestep_range = (
-            self.timestep_skip_range[0] < self.current_timestep_callback() < self.timestep_skip_range[1]
-        )
-        if not is_within_timestep_range:
-            should_skip_attention = False
-        else:
-            should_compute_attention = self.state.iteration > 0 and self.state.iteration % self.block_skip_range == 0
-            should_skip_attention = not should_compute_attention
-        if should_skip_attention:
-            should_skip_attention = self.is_guidance_distilled or self.state.batch_size != batch_size
-
-        if should_skip_attention:
-            logger.debug("FasterCache - Skipping attention and using approximation")
-            if torch.is_tensor(self.state.cache[-1]):
-                t_2_output, t_output = self.state.cache
-                weight = self.weight_callback(module)
-                output = self._compute_approximated_attention_output(t_2_output, t_output, weight, batch_size)
-            else:
-                # The cache contains multiple tensors from past N iterations (N=2 for FasterCache). We need to handle all of them.
-                # Diffusers blocks can return multiple tensors - let's call them [A, B, C, ...] for simplicity.
-                # In our cache, we would have [[A_1, B_1, C_1, ...], [A_2, B_2, C_2, ...], ...] where each list is the output from
-                # a forward pass of the block. We need to compute the approximated output for each of these tensors.
-                # The zip(*state.cache) operation will give us [(A_1, A_2, ...), (B_1, B_2, ...), (C_1, C_2, ...), ...] which
-                # allows us to compute the approximated attention output for each tensor in the cache.
-                output = ()
-                for t_2_output, t_output in zip(*self.state.cache):
-                    result = self._compute_approximated_attention_output(
-                        t_2_output, t_output, self.weight_callback(module), batch_size
-                    )
-                    output += (result,)
-        else:
-            logger.debug("FasterCache - Computing attention")
-            output = self.fn_ref.original_forward(*args, **kwargs)
-
-        # Note that the following condition for getting hidden_states should suffice since Diffusers blocks either return
-        # a single hidden_states tensor, or a tuple of (hidden_states, encoder_hidden_states) tensors. We need to handle
-        # both cases.
-        if torch.is_tensor(output):
-            cache_output = output
-            if not self.is_guidance_distilled and cache_output.size(0) == self.state.batch_size:
-                # The output here can be both unconditional-conditional branch outputs or just conditional branch outputs.
-                # This is determined at the higher-level denoiser module. We only want to cache the conditional branch outputs.
-                cache_output = cache_output.chunk(2, dim=0)[1]
-        else:
-            # Cache all return values and perform the same operation as above
-            cache_output = ()
-            for out in output:
-                if not self.is_guidance_distilled and out.size(0) == self.state.batch_size:
-                    out = out.chunk(2, dim=0)[1]
-                cache_output += (out,)
-
-        if self.state.cache is None:
-            self.state.cache = [cache_output, cache_output]
-        else:
-            self.state.cache = [self.state.cache[-1], cache_output]
-
-        self.state.iteration += 1
-        return output
-
-    def reset_state(self, module: torch.nn.Module) -> torch.nn.Module:
-        self.state.reset()
-        return module
-
-
-def apply_faster_cache(module: torch.nn.Module, config: FasterCacheConfig) -> None:
-    r"""
-    Applies [FasterCache](https://huggingface.co/papers/2410.19355) to a given pipeline.
-
-    Args:
-        module (`torch.nn.Module`):
-            The pytorch module to apply FasterCache to. Typically, this should be a transformer architecture supported
-            in Diffusers, such as `CogVideoXTransformer3DModel`, but external implementations may also work.
-        config (`FasterCacheConfig`):
-            The configuration to use for FasterCache.
-
-    Example:
-    ```python
-    >>> import torch
-    >>> from diffusers import CogVideoXPipeline, FasterCacheConfig, apply_faster_cache
-
-    >>> pipe = CogVideoXPipeline.from_pretrained("THUDM/CogVideoX-5b", torch_dtype=torch.bfloat16)
-    >>> pipe.to("cuda")
-
-    >>> config = FasterCacheConfig(
-    ...     spatial_attention_block_skip_range=2,
-    ...     spatial_attention_timestep_skip_range=(-1, 681),
-    ...     low_frequency_weight_update_timestep_range=(99, 641),
-    ...     high_frequency_weight_update_timestep_range=(-1, 301),
-    ...     spatial_attention_block_identifiers=["transformer_blocks"],
-    ...     attention_weight_callback=lambda _: 0.3,
-    ...     tensor_format="BFCHW",
-    ... )
-    >>> apply_faster_cache(pipe.transformer, config)
-    ```
-    """
-
-    logger.warning(
-        "FasterCache is a purely experimental feature and may not work as expected. Not all models support FasterCache. "
-        "The API is subject to change in future releases, with no guarantee of backward compatibility. Please report any issues at "
-        "https://github.com/huggingface/diffusers/issues."
-    )
-
-    if config.attention_weight_callback is None:
-        # If the user has not provided a weight callback, we default to 0.5 for all timesteps.
-        # In the paper, they recommend using a gradually increasing weight from 0 to 1 as the inference progresses, but
-        # this depends from model-to-model. It is required by the user to provide a weight callback if they want to
-        # use a different weight function. Defaulting to 0.5 works well in practice for most cases.
-        logger.warning(
-            "No `attention_weight_callback` provided when enabling FasterCache. Defaulting to using a weight of 0.5 for all timesteps."
-        )
-        config.attention_weight_callback = lambda _: 0.5
-
-    if config.low_frequency_weight_callback is None:
-        logger.debug(
-            "Low frequency weight callback not provided when enabling FasterCache. Defaulting to behaviour described in the paper."
-        )
-
-        def low_frequency_weight_callback(module: torch.nn.Module) -> float:
-            is_within_range = (
-                config.low_frequency_weight_update_timestep_range[0]
-                < config.current_timestep_callback()
-                < config.low_frequency_weight_update_timestep_range[1]
-            )
-            return config.alpha_low_frequency if is_within_range else 1.0
-
-        config.low_frequency_weight_callback = low_frequency_weight_callback
-
-    if config.high_frequency_weight_callback is None:
-        logger.debug(
-            "High frequency weight callback not provided when enabling FasterCache. Defaulting to behaviour described in the paper."
-        )
-
-        def high_frequency_weight_callback(module: torch.nn.Module) -> float:
-            is_within_range = (
-                config.high_frequency_weight_update_timestep_range[0]
-                < config.current_timestep_callback()
-                < config.high_frequency_weight_update_timestep_range[1]
-            )
-            return config.alpha_high_frequency if is_within_range else 1.0
-
-        config.high_frequency_weight_callback = high_frequency_weight_callback
-
-    supported_tensor_formats = ["BCFHW", "BFCHW", "BCHW"]  # TODO(aryan): Support BSC for LTX Video
-    if config.tensor_format not in supported_tensor_formats:
-        raise ValueError(f"`tensor_format` must be one of {supported_tensor_formats}, but got {config.tensor_format}.")
-
-    _apply_faster_cache_on_denoiser(module, config)
-
-    for name, submodule in module.named_modules():
-        if not isinstance(submodule, _ATTENTION_CLASSES):
-            continue
-        if any(re.search(identifier, name) is not None for identifier in _TRANSFORMER_BLOCK_IDENTIFIERS):
-            _apply_faster_cache_on_attention_class(name, submodule, config)
-
-
-def _apply_faster_cache_on_denoiser(module: torch.nn.Module, config: FasterCacheConfig) -> None:
-    hook = FasterCacheDenoiserHook(
-        config.unconditional_batch_skip_range,
-        config.unconditional_batch_timestep_skip_range,
-        config.tensor_format,
-        config.is_guidance_distilled,
-        config._unconditional_conditional_input_kwargs_identifiers,
-        config.current_timestep_callback,
-        config.low_frequency_weight_callback,
-        config.high_frequency_weight_callback,
-    )
-    registry = HookRegistry.check_if_exists_or_initialize(module)
-    registry.register_hook(hook, _FASTER_CACHE_DENOISER_HOOK)
-
-
-def _apply_faster_cache_on_attention_class(name: str, module: AttentionModuleMixin, config: FasterCacheConfig) -> None:
-    is_spatial_self_attention = (
-        any(re.search(identifier, name) is not None for identifier in config.spatial_attention_block_identifiers)
-        and config.spatial_attention_block_skip_range is not None
-        and not getattr(module, "is_cross_attention", False)
-    )
-    is_temporal_self_attention = (
-        any(re.search(identifier, name) is not None for identifier in config.temporal_attention_block_identifiers)
-        and config.temporal_attention_block_skip_range is not None
-        and not module.is_cross_attention
-    )
-
-    block_skip_range, timestep_skip_range, block_type = None, None, None
-    if is_spatial_self_attention:
-        block_skip_range = config.spatial_attention_block_skip_range
-        timestep_skip_range = config.spatial_attention_timestep_skip_range
-        block_type = "spatial"
-    elif is_temporal_self_attention:
-        block_skip_range = config.temporal_attention_block_skip_range
-        timestep_skip_range = config.temporal_attention_timestep_skip_range
-        block_type = "temporal"
-
-    if block_skip_range is None or timestep_skip_range is None:
-        logger.debug(
-            f'Unable to apply FasterCache to the selected layer: "{name}" because it does '
-            f"not match any of the required criteria for spatial or temporal attention layers. Note, "
-            f"however, that this layer may still be valid for applying PAB. Please specify the correct "
-            f"block identifiers in the configuration or use the specialized `apply_faster_cache_on_module` "
-            f"function to apply FasterCache to this layer."
-        )
-        return
-
-    logger.debug(f"Enabling FasterCache ({block_type}) for layer: {name}")
-    hook = FasterCacheBlockHook(
-        block_skip_range,
-        timestep_skip_range,
-        config.is_guidance_distilled,
-        config.attention_weight_callback,
-        config.current_timestep_callback,
-    )
-    registry = HookRegistry.check_if_exists_or_initialize(module)
-    registry.register_hook(hook, _FASTER_CACHE_BLOCK_HOOK)
-
-
-# Reference: https://github.com/Vchitect/FasterCache/blob/fab32c15014636dc854948319c0a9a8d92c7acb4/scripts/latte/faster_cache_sample_latte.py#L127C1-L143C39
-@torch.no_grad()
-def _split_low_high_freq(x):
-    fft = torch.fft.fft2(x)
-    fft_shifted = torch.fft.fftshift(fft)
-    height, width = x.shape[-2:]
-    radius = min(height, width) // 5
-
-    y_grid, x_grid = torch.meshgrid(torch.arange(height), torch.arange(width))
-    center_x, center_y = width // 2, height // 2
-    mask = (x_grid - center_x) ** 2 + (y_grid - center_y) ** 2 <= radius**2
-
-    low_freq_mask = mask.unsqueeze(0).unsqueeze(0).to(x.device)
-    high_freq_mask = ~low_freq_mask
-
-    low_freq_fft = fft_shifted * low_freq_mask
-    high_freq_fft = fft_shifted * high_freq_mask
-
-    return low_freq_fft, high_freq_fft

diffusers/hooks/first_block_cache.py

@@ -1,258 +0,0 @@
-# Copyright 2024 The HuggingFace Team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-from dataclasses import dataclass
-
-import torch
-
-from ..utils import get_logger
-from ..utils.torch_utils import unwrap_module
-from ._common import _ALL_TRANSFORMER_BLOCK_IDENTIFIERS
-from ._helpers import TransformerBlockRegistry
-from .hooks import BaseState, HookRegistry, ModelHook, StateManager
-
-
-logger = get_logger(__name__)  # pylint: disable=invalid-name
-
-_FBC_LEADER_BLOCK_HOOK = "fbc_leader_block_hook"
-_FBC_BLOCK_HOOK = "fbc_block_hook"
-
-
-@dataclass
-class FirstBlockCacheConfig:
-    r"""
-    Configuration for [First Block
-    Cache](https://github.com/chengzeyi/ParaAttention/blob/7a266123671b55e7e5a2fe9af3121f07a36afc78/README.md#first-block-cache-our-dynamic-caching).
-
-    Args:
-        threshold (`float`, defaults to `0.05`):
-            The threshold to determine whether or not a forward pass through all layers of the model is required. A
-            higher threshold usually results in a forward pass through a lower number of layers and faster inference,
-            but might lead to poorer generation quality. A lower threshold may not result in significant generation
-            speedup. The threshold is compared against the absmean difference of the residuals between the current and
-            cached outputs from the first transformer block. If the difference is below the threshold, the forward pass
-            is skipped.
-    """
-
-    threshold: float = 0.05
-
-
-class FBCSharedBlockState(BaseState):
-    def __init__(self) -> None:
-        super().__init__()
-
-        self.head_block_output: torch.Tensor | tuple[torch.Tensor, ...] = None
-        self.head_block_residual: torch.Tensor = None
-        self.tail_block_residuals: torch.Tensor | tuple[torch.Tensor, ...] = None
-        self.should_compute: bool = True
-
-    def reset(self):
-        self.tail_block_residuals = None
-        self.should_compute = True
-
-
-class FBCHeadBlockHook(ModelHook):
-    _is_stateful = True
-
-    def __init__(self, state_manager: StateManager, threshold: float):
-        self.state_manager = state_manager
-        self.threshold = threshold
-        self._metadata = None
-
-    def initialize_hook(self, module):
-        unwrapped_module = unwrap_module(module)
-        self._metadata = TransformerBlockRegistry.get(unwrapped_module.__class__)
-        return module
-
-    def new_forward(self, module: torch.nn.Module, *args, **kwargs):
-        original_hidden_states = self._metadata._get_parameter_from_args_kwargs("hidden_states", args, kwargs)
-
-        output = self.fn_ref.original_forward(*args, **kwargs)
-        is_output_tuple = isinstance(output, tuple)
-
-        if is_output_tuple:
-            hidden_states_residual = output[self._metadata.return_hidden_states_index] - original_hidden_states
-        else:
-            hidden_states_residual = output - original_hidden_states
-
-        shared_state: FBCSharedBlockState = self.state_manager.get_state()
-        hidden_states = encoder_hidden_states = None
-        should_compute = self._should_compute_remaining_blocks(hidden_states_residual)
-        shared_state.should_compute = should_compute
-
-        if not should_compute:
-            # Apply caching
-            if is_output_tuple:
-                hidden_states = (
-                    shared_state.tail_block_residuals[0] + output[self._metadata.return_hidden_states_index]
-                )
-            else:
-                hidden_states = shared_state.tail_block_residuals[0] + output
-
-            if self._metadata.return_encoder_hidden_states_index is not None:
-                assert is_output_tuple
-                encoder_hidden_states = (
-                    shared_state.tail_block_residuals[1] + output[self._metadata.return_encoder_hidden_states_index]
-                )
-
-            if is_output_tuple:
-                return_output = [None] * len(output)
-                return_output[self._metadata.return_hidden_states_index] = hidden_states
-                return_output[self._metadata.return_encoder_hidden_states_index] = encoder_hidden_states
-                return_output = tuple(return_output)
-            else:
-                return_output = hidden_states
-            output = return_output
-        else:
-            if is_output_tuple:
-                head_block_output = [None] * len(output)
-                head_block_output[0] = output[self._metadata.return_hidden_states_index]
-                head_block_output[1] = output[self._metadata.return_encoder_hidden_states_index]
-            else:
-                head_block_output = output
-            shared_state.head_block_output = head_block_output
-            shared_state.head_block_residual = hidden_states_residual
-
-        return output
-
-    def reset_state(self, module):
-        self.state_manager.reset()
-        return module
-
-    @torch.compiler.disable
-    def _should_compute_remaining_blocks(self, hidden_states_residual: torch.Tensor) -> bool:
-        shared_state = self.state_manager.get_state()
-        if shared_state.head_block_residual is None:
-            return True
-        prev_hidden_states_residual = shared_state.head_block_residual
-        absmean = (hidden_states_residual - prev_hidden_states_residual).abs().mean()
-        prev_hidden_states_absmean = prev_hidden_states_residual.abs().mean()
-        diff = (absmean / prev_hidden_states_absmean).item()
-        return diff > self.threshold
-
-
-class FBCBlockHook(ModelHook):
-    def __init__(self, state_manager: StateManager, is_tail: bool = False):
-        super().__init__()
-        self.state_manager = state_manager
-        self.is_tail = is_tail
-        self._metadata = None
-
-    def initialize_hook(self, module):
-        unwrapped_module = unwrap_module(module)
-        self._metadata = TransformerBlockRegistry.get(unwrapped_module.__class__)
-        return module
-
-    def new_forward(self, module: torch.nn.Module, *args, **kwargs):
-        original_hidden_states = self._metadata._get_parameter_from_args_kwargs("hidden_states", args, kwargs)
-        original_encoder_hidden_states = None
-        if self._metadata.return_encoder_hidden_states_index is not None:
-            original_encoder_hidden_states = self._metadata._get_parameter_from_args_kwargs(
-                "encoder_hidden_states", args, kwargs
-            )
-
-        shared_state = self.state_manager.get_state()
-
-        if shared_state.should_compute:
-            output = self.fn_ref.original_forward(*args, **kwargs)
-            if self.is_tail:
-                hidden_states_residual = encoder_hidden_states_residual = None
-                if isinstance(output, tuple):
-                    hidden_states_residual = (
-                        output[self._metadata.return_hidden_states_index] - shared_state.head_block_output[0]
-                    )
-                    encoder_hidden_states_residual = (
-                        output[self._metadata.return_encoder_hidden_states_index] - shared_state.head_block_output[1]
-                    )
-                else:
-                    hidden_states_residual = output - shared_state.head_block_output
-                shared_state.tail_block_residuals = (hidden_states_residual, encoder_hidden_states_residual)
-            return output
-
-        if original_encoder_hidden_states is None:
-            return_output = original_hidden_states
-        else:
-            return_output = [None, None]
-            return_output[self._metadata.return_hidden_states_index] = original_hidden_states
-            return_output[self._metadata.return_encoder_hidden_states_index] = original_encoder_hidden_states
-            return_output = tuple(return_output)
-        return return_output
-
-
-def apply_first_block_cache(module: torch.nn.Module, config: FirstBlockCacheConfig) -> None:
-    """
-    Applies [First Block
-    Cache](https://github.com/chengzeyi/ParaAttention/blob/4de137c5b96416489f06e43e19f2c14a772e28fd/README.md#first-block-cache-our-dynamic-caching)
-    to a given module.
-
-    First Block Cache builds on the ideas of [TeaCache](https://huggingface.co/papers/2411.19108). It is much simpler
-    to implement generically for a wide range of models and has been integrated first for experimental purposes.
-
-    Args:
-        module (`torch.nn.Module`):
-            The pytorch module to apply FBCache to. Typically, this should be a transformer architecture supported in
-            Diffusers, such as `CogVideoXTransformer3DModel`, but external implementations may also work.
-        config (`FirstBlockCacheConfig`):
-            The configuration to use for applying the FBCache method.
-
-    Example:
-        ```python
-        >>> import torch
-        >>> from diffusers import CogView4Pipeline
-        >>> from diffusers.hooks import apply_first_block_cache, FirstBlockCacheConfig
-
-        >>> pipe = CogView4Pipeline.from_pretrained("THUDM/CogView4-6B", torch_dtype=torch.bfloat16)
-        >>> pipe.to("cuda")
-
-        >>> apply_first_block_cache(pipe.transformer, FirstBlockCacheConfig(threshold=0.2))
-
-        >>> prompt = "A photo of an astronaut riding a horse on mars"
-        >>> image = pipe(prompt, generator=torch.Generator().manual_seed(42)).images[0]
-        >>> image.save("output.png")
-        ```
-    """
-
-    state_manager = StateManager(FBCSharedBlockState, (), {})
-    remaining_blocks = []
-
-    for name, submodule in module.named_children():
-        if name not in _ALL_TRANSFORMER_BLOCK_IDENTIFIERS or not isinstance(submodule, torch.nn.ModuleList):
-            continue
-        for index, block in enumerate(submodule):
-            remaining_blocks.append((f"{name}.{index}", block))
-
-    head_block_name, head_block = remaining_blocks.pop(0)
-    tail_block_name, tail_block = remaining_blocks.pop(-1)
-
-    logger.debug(f"Applying FBCHeadBlockHook to '{head_block_name}'")
-    _apply_fbc_head_block_hook(head_block, state_manager, config.threshold)
-
-    for name, block in remaining_blocks:
-        logger.debug(f"Applying FBCBlockHook to '{name}'")
-        _apply_fbc_block_hook(block, state_manager)
-
-    logger.debug(f"Applying FBCBlockHook to tail block '{tail_block_name}'")
-    _apply_fbc_block_hook(tail_block, state_manager, is_tail=True)
-
-
-def _apply_fbc_head_block_hook(block: torch.nn.Module, state_manager: StateManager, threshold: float) -> None:
-    registry = HookRegistry.check_if_exists_or_initialize(block)
-    hook = FBCHeadBlockHook(state_manager, threshold)
-    registry.register_hook(hook, _FBC_LEADER_BLOCK_HOOK)
-
-
-def _apply_fbc_block_hook(block: torch.nn.Module, state_manager: StateManager, is_tail: bool = False) -> None:
-    registry = HookRegistry.check_if_exists_or_initialize(block)
-    hook = FBCBlockHook(state_manager, is_tail)
-    registry.register_hook(hook, _FBC_BLOCK_HOOK)

diffusers/hooks/group_offloading.py

@@ -1,1056 +0,0 @@
-# Copyright 2026 The HuggingFace Team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-import hashlib
-import os
-from contextlib import contextmanager, nullcontext
-from dataclasses import dataclass, replace
-from enum import Enum
-from typing import Set
-
-import safetensors.torch
-import torch
-
-from ..utils import get_logger, is_accelerate_available, is_torchao_available
-from ._common import _GO_LC_SUPPORTED_PYTORCH_LAYERS
-from .hooks import HookRegistry, ModelHook
-
-
-if is_accelerate_available():
-    from accelerate.hooks import AlignDevicesHook, CpuOffload
-    from accelerate.utils import send_to_device
-
-
-logger = get_logger(__name__)  # pylint: disable=invalid-name
-
-
-def _is_torchao_tensor(tensor: torch.Tensor) -> bool:
-    if not is_torchao_available():
-        return False
-    from torchao.utils import TorchAOBaseTensor
-
-    return isinstance(tensor, TorchAOBaseTensor)
-
-
-def _get_torchao_inner_tensor_names(tensor: torch.Tensor) -> list[str]:
-    """Get names of all internal tensor data attributes from a TorchAO tensor."""
-    cls = type(tensor)
-    names = list(getattr(cls, "tensor_data_names", []))
-    for attr_name in getattr(cls, "optional_tensor_data_names", []):
-        if getattr(tensor, attr_name, None) is not None:
-            names.append(attr_name)
-    return names
-
-
-def _swap_torchao_tensor(param: torch.Tensor, source: torch.Tensor) -> None:
-    """Move a TorchAO parameter to the device of `source` via `swap_tensors`.
-
-    `param.data = source` does not work for `_make_wrapper_subclass` tensors because the `.data` setter only replaces
-    the outer wrapper storage while leaving the subclass's internal attributes (e.g. `.qdata`, `.scale`) on the
-    original device. `swap_tensors` swaps the full tensor contents in-place, preserving the parameter's identity so
-    that any dict keyed by `id(param)` remains valid.
-
-    Refer to https://github.com/huggingface/diffusers/pull/13276#discussion_r2944471548 for the full discussion.
-    """
-    torch.utils.swap_tensors(param, source)
-
-
-def _restore_torchao_tensor(param: torch.Tensor, source: torch.Tensor) -> None:
-    """Restore internal tensor data of a TorchAO parameter from `source` without mutating `source`.
-
-    Unlike `_swap_torchao_tensor` this copies attribute references one-by-one via `setattr` so that `source` is **not**
-    modified. Use this when `source` is a cached tensor that must remain unchanged (e.g. a pinned CPU copy in
-    `cpu_param_dict`).
-    """
-    for attr_name in _get_torchao_inner_tensor_names(source):
-        setattr(param, attr_name, getattr(source, attr_name))
-
-
-def _record_stream_torchao_tensor(param: torch.Tensor, stream) -> None:
-    """Record stream for all internal tensors of a TorchAO parameter."""
-    for attr_name in _get_torchao_inner_tensor_names(param):
-        getattr(param, attr_name).record_stream(stream)
-
-
-# fmt: off
-_GROUP_OFFLOADING = "group_offloading"
-_LAYER_EXECUTION_TRACKER = "layer_execution_tracker"
-_LAZY_PREFETCH_GROUP_OFFLOADING = "lazy_prefetch_group_offloading"
-_GROUP_ID_LAZY_LEAF = "lazy_leafs"
-# fmt: on
-
-
-class GroupOffloadingType(str, Enum):
-    BLOCK_LEVEL = "block_level"
-    LEAF_LEVEL = "leaf_level"
-
-
-@dataclass
-class GroupOffloadingConfig:
-    onload_device: torch.device
-    offload_device: torch.device
-    offload_type: GroupOffloadingType
-    non_blocking: bool
-    record_stream: bool
-    low_cpu_mem_usage: bool
-    num_blocks_per_group: int | None = None
-    offload_to_disk_path: str | None = None
-    stream: torch.cuda.Stream | torch.Stream | None = None
-    block_modules: list[str] | None = None
-    exclude_kwargs: list[str] | None = None
-    module_prefix: str = ""
-
-
-class ModuleGroup:
-    def __init__(
-        self,
-        modules: list[torch.nn.Module],
-        offload_device: torch.device,
-        onload_device: torch.device,
-        offload_leader: torch.nn.Module,
-        onload_leader: torch.nn.Module | None = None,
-        parameters: list[torch.nn.Parameter] | None = None,
-        buffers: list[torch.Tensor] | None = None,
-        non_blocking: bool = False,
-        stream: torch.cuda.Stream | torch.Stream | None = None,
-        record_stream: bool | None = False,
-        low_cpu_mem_usage: bool = False,
-        onload_self: bool = True,
-        offload_to_disk_path: str | None = None,
-        group_id: int | str | None = None,
-    ) -> None:
-        self.modules = modules
-        self.offload_device = offload_device
-        self.onload_device = onload_device
-        self.offload_leader = offload_leader
-        self.onload_leader = onload_leader
-        self.parameters = parameters or []
-        self.buffers = buffers or []
-        self.non_blocking = non_blocking or stream is not None
-        self.stream = stream
-        self.record_stream = record_stream
-        self.onload_self = onload_self
-        self.low_cpu_mem_usage = low_cpu_mem_usage
-
-        self.offload_to_disk_path = offload_to_disk_path
-        self._is_offloaded_to_disk = False
-
-        if self.offload_to_disk_path is not None:
-            # Instead of `group_id or str(id(self))` we do this because `group_id` can be "" as well.
-            self.group_id = group_id if group_id is not None else str(id(self))
-            short_hash = _compute_group_hash(self.group_id)
-            self.safetensors_file_path = os.path.join(self.offload_to_disk_path, f"group_{short_hash}.safetensors")
-
-            all_tensors = []
-            for module in self.modules:
-                all_tensors.extend(list(module.parameters()))
-                all_tensors.extend(list(module.buffers()))
-            all_tensors.extend(self.parameters)
-            all_tensors.extend(self.buffers)
-            all_tensors = list(dict.fromkeys(all_tensors))  # Remove duplicates
-
-            self.tensor_to_key = {tensor: f"tensor_{i}" for i, tensor in enumerate(all_tensors)}
-            self.key_to_tensor = {v: k for k, v in self.tensor_to_key.items()}
-            self.cpu_param_dict = {}
-        else:
-            self.cpu_param_dict = self._init_cpu_param_dict()
-
-        self._torch_accelerator_module = (
-            getattr(torch, torch.accelerator.current_accelerator().type)
-            if hasattr(torch, "accelerator")
-            else torch.cuda
-        )
-
-    @staticmethod
-    def _to_cpu(tensor, low_cpu_mem_usage):
-        # For TorchAO tensors, `.data` returns an incomplete wrapper without internal attributes
-        # (e.g. `.qdata`, `.scale`), so we must call `.cpu()` on the tensor directly.
-        t = tensor.cpu() if _is_torchao_tensor(tensor) else tensor.data.cpu()
-        return t if low_cpu_mem_usage else t.pin_memory()
-
-    def _init_cpu_param_dict(self):
-        cpu_param_dict = {}
-        if self.stream is None:
-            return cpu_param_dict
-
-        for module in self.modules:
-            for param in module.parameters():
-                cpu_param_dict[param] = self._to_cpu(param, self.low_cpu_mem_usage)
-            for buffer in module.buffers():
-                cpu_param_dict[buffer] = self._to_cpu(buffer, self.low_cpu_mem_usage)
-
-        for param in self.parameters:
-            cpu_param_dict[param] = self._to_cpu(param, self.low_cpu_mem_usage)
-
-        for buffer in self.buffers:
-            cpu_param_dict[buffer] = self._to_cpu(buffer, self.low_cpu_mem_usage)
-
-        return cpu_param_dict
-
-    @contextmanager
-    def _pinned_memory_tensors(self):
-        try:
-            pinned_dict = {
-                param: tensor.pin_memory() if not tensor.is_pinned() else tensor
-                for param, tensor in self.cpu_param_dict.items()
-            }
-            yield pinned_dict
-        finally:
-            pinned_dict = None
-
-    def _transfer_tensor_to_device(self, tensor, source_tensor, default_stream):
-        moved = source_tensor.to(self.onload_device, non_blocking=self.non_blocking)
-        if _is_torchao_tensor(tensor):
-            _swap_torchao_tensor(tensor, moved)
-        else:
-            tensor.data = moved
-        if self.record_stream:
-            if _is_torchao_tensor(tensor):
-                _record_stream_torchao_tensor(tensor, default_stream)
-            else:
-                tensor.data.record_stream(default_stream)
-
-    def _process_tensors_from_modules(self, pinned_memory=None, default_stream=None):
-        for group_module in self.modules:
-            for param in group_module.parameters():
-                source = pinned_memory[param] if pinned_memory else param.data
-                self._transfer_tensor_to_device(param, source, default_stream)
-            for buffer in group_module.buffers():
-                source = pinned_memory[buffer] if pinned_memory else buffer.data
-                self._transfer_tensor_to_device(buffer, source, default_stream)
-
-        for param in self.parameters:
-            source = pinned_memory[param] if pinned_memory else param.data
-            self._transfer_tensor_to_device(param, source, default_stream)
-
-        for buffer in self.buffers:
-            source = pinned_memory[buffer] if pinned_memory else buffer.data
-            self._transfer_tensor_to_device(buffer, source, default_stream)
-
-    def _check_disk_offload_torchao(self):
-        all_tensors = list(self.tensor_to_key.keys())
-        has_torchao = any(_is_torchao_tensor(t) for t in all_tensors)
-        if has_torchao:
-            raise ValueError(
-                "Disk offloading is not supported for TorchAO quantized tensors because safetensors "
-                "cannot serialize TorchAO subclass tensors. Use memory offloading instead by not "
-                "setting `offload_to_disk_path`."
-            )
-
-    def _onload_from_disk(self):
-        self._check_disk_offload_torchao()
-
-        if self.stream is not None:
-            # Wait for previous Host->Device transfer to complete
-            self.stream.synchronize()
-
-        context = nullcontext() if self.stream is None else self._torch_accelerator_module.stream(self.stream)
-        current_stream = self._torch_accelerator_module.current_stream() if self.record_stream else None
-
-        with context:
-            if self.stream is not None:
-                # Load to CPU first, pin memory, then async copy to the target device
-                loaded_tensors = safetensors.torch.load_file(self.safetensors_file_path, device="cpu")
-                for key, tensor_obj in self.key_to_tensor.items():
-                    pinned_tensor = loaded_tensors[key].pin_memory()
-                    tensor_obj.data = pinned_tensor.to(self.onload_device, non_blocking=self.non_blocking)
-                    if self.record_stream:
-                        tensor_obj.data.record_stream(current_stream)
-            else:
-                # Load directly to the target device
-                onload_device = (
-                    self.onload_device.type if isinstance(self.onload_device, torch.device) else self.onload_device
-                )
-                loaded_tensors = safetensors.torch.load_file(self.safetensors_file_path, device=onload_device)
-                for key, tensor_obj in self.key_to_tensor.items():
-                    tensor_obj.data = loaded_tensors[key]
-
-    def _onload_from_memory(self):
-        if self.stream is not None:
-            # Wait for previous Host->Device transfer to complete
-            self.stream.synchronize()
-
-        context = nullcontext() if self.stream is None else self._torch_accelerator_module.stream(self.stream)
-        default_stream = self._torch_accelerator_module.current_stream() if self.stream is not None else None
-
-        with context:
-            if self.stream is not None:
-                with self._pinned_memory_tensors() as pinned_memory:
-                    self._process_tensors_from_modules(pinned_memory, default_stream=default_stream)
-            else:
-                self._process_tensors_from_modules(None)
-
-    def _offload_to_disk(self):
-        self._check_disk_offload_torchao()
-
-        # TODO: we can potentially optimize this code path by checking if the _all_ the desired
-        # safetensor files exist on the disk and if so, skip this step entirely, reducing IO
-        # overhead. Currently, we just check if the given `safetensors_file_path` exists and if not
-        # we perform a write.
-        # Check if the file has been saved in this session or if it already exists on disk.
-        if not self._is_offloaded_to_disk and not os.path.exists(self.safetensors_file_path):
-            os.makedirs(os.path.dirname(self.safetensors_file_path), exist_ok=True)
-            tensors_to_save = {key: tensor.data.to(self.offload_device) for tensor, key in self.tensor_to_key.items()}
-            safetensors.torch.save_file(tensors_to_save, self.safetensors_file_path)
-
-        # The group is now considered offloaded to disk for the rest of the session.
-        self._is_offloaded_to_disk = True
-
-        # We do this to free up the RAM which is still holding the up tensor data.
-        for tensor_obj in self.tensor_to_key.keys():
-            tensor_obj.data = torch.empty_like(tensor_obj.data, device=self.offload_device)
-
-    def _offload_to_memory(self):
-        if self.stream is not None:
-            if not self.record_stream:
-                self._torch_accelerator_module.current_stream().synchronize()
-
-            for group_module in self.modules:
-                for param in group_module.parameters():
-                    if _is_torchao_tensor(param):
-                        _restore_torchao_tensor(param, self.cpu_param_dict[param])
-                    else:
-                        param.data = self.cpu_param_dict[param]
-            for param in self.parameters:
-                if _is_torchao_tensor(param):
-                    _restore_torchao_tensor(param, self.cpu_param_dict[param])
-                else:
-                    param.data = self.cpu_param_dict[param]
-            for buffer in self.buffers:
-                if _is_torchao_tensor(buffer):
-                    _restore_torchao_tensor(buffer, self.cpu_param_dict[buffer])
-                else:
-                    buffer.data = self.cpu_param_dict[buffer]
-        else:
-            for group_module in self.modules:
-                group_module.to(self.offload_device, non_blocking=False)
-            for param in self.parameters:
-                if _is_torchao_tensor(param):
-                    moved = param.to(self.offload_device, non_blocking=False)
-                    _swap_torchao_tensor(param, moved)
-                else:
-                    param.data = param.data.to(self.offload_device, non_blocking=False)
-            for buffer in self.buffers:
-                if _is_torchao_tensor(buffer):
-                    moved = buffer.to(self.offload_device, non_blocking=False)
-                    _swap_torchao_tensor(buffer, moved)
-                else:
-                    buffer.data = buffer.data.to(self.offload_device, non_blocking=False)
-
-    @torch.compiler.disable()
-    def onload_(self):
-        r"""Onloads the group of parameters to the onload_device."""
-        if self.offload_to_disk_path is not None:
-            self._onload_from_disk()
-        else:
-            self._onload_from_memory()
-
-    @torch.compiler.disable()
-    def offload_(self):
-        r"""Offloads the group of parameters to the offload_device."""
-        if self.offload_to_disk_path:
-            self._offload_to_disk()
-        else:
-            self._offload_to_memory()
-
-
-class GroupOffloadingHook(ModelHook):
-    r"""
-    A hook that offloads groups of torch.nn.Module to the CPU for storage and onloads to accelerator device for
-    computation. Each group has one "onload leader" module that is responsible for onloading, and an "offload leader"
-    module that is responsible for offloading. If prefetching is enabled, the onload leader of the previous module
-    group is responsible for onloading the current module group.
-    """
-
-    _is_stateful = False
-
-    def __init__(self, group: ModuleGroup, *, config: GroupOffloadingConfig) -> None:
-        self.group = group
-        self.next_group: ModuleGroup | None = None
-        self.config = config
-
-    def initialize_hook(self, module: torch.nn.Module) -> torch.nn.Module:
-        if self.group.offload_leader == module:
-            self.group.offload_()
-        return module
-
-    def pre_forward(self, module: torch.nn.Module, *args, **kwargs):
-        # If there wasn't an onload_leader assigned, we assume that the submodule that first called its forward
-        # method is the onload_leader of the group.
-        if self.group.onload_leader is None:
-            self.group.onload_leader = module
-
-        # If the current module is the onload_leader of the group, we onload the group if it is supposed
-        # to onload itself. In the case of using prefetching with streams, we onload the next group if
-        # it is not supposed to onload itself.
-        if self.group.onload_leader == module:
-            if self.group.onload_self:
-                self.group.onload_()
-            else:
-                # onload_self=False means this group relies on prefetching from a previous group.
-                # However, for conditionally-executed modules (e.g. patch_short/patch_mid/patch_long in Helios),
-                # the prefetch chain may not cover them if they were absent during the first forward pass
-                # when the execution order was traced. In that case, their weights remain on offload_device,
-                # so we fall back to a synchronous onload here.
-                params = [p for m in self.group.modules for p in m.parameters()] + list(self.group.parameters)
-                if params and params[0].device == self.group.offload_device:
-                    self.group.onload_()
-                    if self.group.stream is not None:
-                        self.group.stream.synchronize()
-
-            should_onload_next_group = self.next_group is not None and not self.next_group.onload_self
-            if should_onload_next_group:
-                self.next_group.onload_()
-
-            should_synchronize = (
-                not self.group.onload_self and self.group.stream is not None and not should_onload_next_group
-            )
-            if should_synchronize:
-                # If this group didn't onload itself, it means it was asynchronously onloaded by the
-                # previous group. We need to synchronize the side stream to ensure parameters
-                # are completely loaded to proceed with forward pass. Without this, uninitialized
-                # weights will be used in the computation, leading to incorrect results
-                # Also, we should only do this synchronization if we don't already do it from the sync call in
-                # self.next_group.onload_, hence the `not should_onload_next_group` check.
-                self.group.stream.synchronize()
-
-        args = send_to_device(args, self.group.onload_device, non_blocking=self.group.non_blocking)
-
-        # Some Autoencoder models use a feature cache that is passed through submodules
-        # and modified in place. The `send_to_device` call returns a copy of this feature cache object
-        # which breaks the inplace updates. Use `exclude_kwargs` to mark these cache features
-        exclude_kwargs = self.config.exclude_kwargs or []
-        if exclude_kwargs:
-            moved_kwargs = send_to_device(
-                {k: v for k, v in kwargs.items() if k not in exclude_kwargs},
-                self.group.onload_device,
-                non_blocking=self.group.non_blocking,
-            )
-            kwargs.update(moved_kwargs)
-        else:
-            kwargs = send_to_device(kwargs, self.group.onload_device, non_blocking=self.group.non_blocking)
-
-        return args, kwargs
-
-    def post_forward(self, module: torch.nn.Module, output):
-        if self.group.offload_leader == module:
-            self.group.offload_()
-        return output
-
-
-class LazyPrefetchGroupOffloadingHook(ModelHook):
-    r"""
-    A hook, used in conjunction with GroupOffloadingHook, that applies lazy prefetching to groups of torch.nn.Module.
-    This hook is used to determine the order in which the layers are executed during the forward pass. Once the layer
-    invocation order is known, assignments of the next_group attribute for prefetching can be made, which allows
-    prefetching groups in the correct order.
-    """
-
-    _is_stateful = False
-
-    def __init__(self):
-        self.execution_order: list[tuple[str, torch.nn.Module]] = []
-        self._layer_execution_tracker_module_names = set()
-
-    def initialize_hook(self, module):
-        def make_execution_order_update_callback(current_name, current_submodule):
-            def callback():
-                if not torch.compiler.is_compiling():
-                    logger.debug(f"Adding {current_name} to the execution order")
-                self.execution_order.append((current_name, current_submodule))
-
-            return callback
-
-        # To every submodule that contains a group offloading hook (at this point, no prefetching is enabled for any
-        # of the groups), we add a layer execution tracker hook that will be used to determine the order in which the
-        # layers are executed during the forward pass.
-        for name, submodule in module.named_modules():
-            if name == "" or not hasattr(submodule, "_diffusers_hook"):
-                continue
-
-            registry = HookRegistry.check_if_exists_or_initialize(submodule)
-            group_offloading_hook = registry.get_hook(_GROUP_OFFLOADING)
-
-            if group_offloading_hook is not None:
-                # For the first forward pass, we have to load in a blocking manner
-                group_offloading_hook.group.non_blocking = False
-                layer_tracker_hook = LayerExecutionTrackerHook(make_execution_order_update_callback(name, submodule))
-                registry.register_hook(layer_tracker_hook, _LAYER_EXECUTION_TRACKER)
-                self._layer_execution_tracker_module_names.add(name)
-
-        return module
-
-    def post_forward(self, module, output):
-        # At this point, for the current modules' submodules, we know the execution order of the layers. We can now
-        # remove the layer execution tracker hooks and apply prefetching by setting the next_group attribute for each
-        # group offloading hook.
-        num_executed = len(self.execution_order)
-        execution_order_module_names = {name for name, _ in self.execution_order}
-
-        # It may be possible that some layers were not executed during the forward pass. This can happen if the layer
-        # is not used in the forward pass, or if the layer is not executed due to some other reason. In such cases, we
-        # may not be able to apply prefetching in the correct order, which can lead to device-mismatch related errors
-        # if the missing layers end up being executed in the future.
-        if execution_order_module_names != self._layer_execution_tracker_module_names:
-            unexecuted_layers = list(self._layer_execution_tracker_module_names - execution_order_module_names)
-            if not torch.compiler.is_compiling():
-                logger.warning(
-                    "It seems like some layers were not executed during the forward pass. This may lead to problems when "
-                    "applying lazy prefetching with automatic tracing and lead to device-mismatch related errors. Please "
-                    "make sure that all layers are executed during the forward pass. The following layers were not executed:\n"
-                    f"{unexecuted_layers=}"
-                )
-
-        # Remove the layer execution tracker hooks from the submodules
-        base_module_registry = module._diffusers_hook
-        registries = [submodule._diffusers_hook for _, submodule in self.execution_order]
-        group_offloading_hooks = [registry.get_hook(_GROUP_OFFLOADING) for registry in registries]
-
-        for i in range(num_executed):
-            registries[i].remove_hook(_LAYER_EXECUTION_TRACKER, recurse=False)
-
-        # Remove the current lazy prefetch group offloading hook so that it doesn't interfere with the next forward pass
-        base_module_registry.remove_hook(_LAZY_PREFETCH_GROUP_OFFLOADING, recurse=False)
-
-        # LazyPrefetchGroupOffloadingHook is only used with streams, so we know that non_blocking should be True.
-        # We disable non_blocking for the first forward pass, but need to enable it for the subsequent passes to
-        # see the benefits of prefetching.
-        for hook in group_offloading_hooks:
-            hook.group.non_blocking = True
-
-        # Set required attributes for prefetching
-        if num_executed > 0:
-            base_module_group_offloading_hook = base_module_registry.get_hook(_GROUP_OFFLOADING)
-            base_module_group_offloading_hook.next_group = group_offloading_hooks[0].group
-            base_module_group_offloading_hook.next_group.onload_self = False
-
-        for i in range(num_executed - 1):
-            name1, _ = self.execution_order[i]
-            name2, _ = self.execution_order[i + 1]
-            if not torch.compiler.is_compiling():
-                logger.debug(f"Applying lazy prefetch group offloading from {name1} to {name2}")
-            group_offloading_hooks[i].next_group = group_offloading_hooks[i + 1].group
-            group_offloading_hooks[i].next_group.onload_self = False
-
-        return output
-
-
-class LayerExecutionTrackerHook(ModelHook):
-    r"""
-    A hook that tracks the order in which the layers are executed during the forward pass by calling back to the
-    LazyPrefetchGroupOffloadingHook to update the execution order.
-    """
-
-    _is_stateful = False
-
-    def __init__(self, execution_order_update_callback):
-        self.execution_order_update_callback = execution_order_update_callback
-
-    def pre_forward(self, module, *args, **kwargs):
-        self.execution_order_update_callback()
-        return args, kwargs
-
-
-def apply_group_offloading(
-    module: torch.nn.Module,
-    onload_device: str | torch.device,
-    offload_device: str | torch.device = torch.device("cpu"),
-    offload_type: str | GroupOffloadingType = "block_level",
-    num_blocks_per_group: int | None = None,
-    non_blocking: bool = False,
-    use_stream: bool = False,
-    record_stream: bool = False,
-    low_cpu_mem_usage: bool = False,
-    offload_to_disk_path: str | None = None,
-    block_modules: list[str] | None = None,
-    exclude_kwargs: list[str] | None = None,
-) -> None:
-    r"""
-    Applies group offloading to the internal layers of a torch.nn.Module. To understand what group offloading is, and
-    where it is beneficial, we need to first provide some context on how other supported offloading methods work.
-
-    Typically, offloading is done at two levels:
-    - Module-level: In Diffusers, this can be enabled using the `ModelMixin::enable_model_cpu_offload()` method. It
-      works by offloading each component of a pipeline to the CPU for storage, and onloading to the accelerator device
-      when needed for computation. This method is more memory-efficient than keeping all components on the accelerator,
-      but the memory requirements are still quite high. For this method to work, one needs memory equivalent to size of
-      the model in runtime dtype + size of largest intermediate activation tensors to be able to complete the forward
-      pass.
-    - Leaf-level: In Diffusers, this can be enabled using the `ModelMixin::enable_sequential_cpu_offload()` method. It
-      works by offloading the lowest leaf-level parameters of the computation graph to the CPU for storage, and
-      onloading only the leafs to the accelerator device for computation. This uses the lowest amount of accelerator
-      memory, but can be slower due to the excessive number of device synchronizations.
-
-    Group offloading is a middle ground between the two methods. It works by offloading groups of internal layers,
-    (either `torch.nn.ModuleList` or `torch.nn.Sequential`). This method uses lower memory than module-level
-    offloading. It is also faster than leaf-level/sequential offloading, as the number of device synchronizations is
-    reduced.
-
-    Another supported feature (for CUDA devices with support for asynchronous data transfer streams) is the ability to
-    overlap data transfer and computation to reduce the overall execution time compared to sequential offloading. This
-    is enabled using layer prefetching with streams, i.e., the layer that is to be executed next starts onloading to
-    the accelerator device while the current layer is being executed - this increases the memory requirements slightly.
-    Note that this implementation also supports leaf-level offloading but can be made much faster when using streams.
-
-    Args:
-        module (`torch.nn.Module`):
-            The module to which group offloading is applied.
-        onload_device (`torch.device`):
-            The device to which the group of modules are onloaded.
-        offload_device (`torch.device`, defaults to `torch.device("cpu")`):
-            The device to which the group of modules are offloaded. This should typically be the CPU. Default is CPU.
-        offload_type (`str` or `GroupOffloadingType`, defaults to "block_level"):
-            The type of offloading to be applied. Can be one of "block_level" or "leaf_level". Default is
-            "block_level".
-        offload_to_disk_path (`str`, *optional*, defaults to `None`):
-            The path to the directory where parameters will be offloaded. Setting this option can be useful in limited
-            RAM environment settings where a reasonable speed-memory trade-off is desired.
-        num_blocks_per_group (`int`, *optional*):
-            The number of blocks per group when using offload_type="block_level". This is required when using
-            offload_type="block_level".
-        non_blocking (`bool`, defaults to `False`):
-            If True, offloading and onloading is done with non-blocking data transfer.
-        use_stream (`bool`, defaults to `False`):
-            If True, offloading and onloading is done asynchronously using a CUDA stream. This can be useful for
-            overlapping computation and data transfer.
-        record_stream (`bool`, defaults to `False`): When enabled with `use_stream`, it marks the current tensor
-            as having been used by this stream. It is faster at the expense of slightly more memory usage. Refer to the
-            [PyTorch official docs](https://pytorch.org/docs/stable/generated/torch.Tensor.record_stream.html) more
-            details.
-        low_cpu_mem_usage (`bool`, defaults to `False`):
-            If True, the CPU memory usage is minimized by pinning tensors on-the-fly instead of pre-pinning them. This
-            option only matters when using streamed CPU offloading (i.e. `use_stream=True`). This can be useful when
-            the CPU memory is a bottleneck but may counteract the benefits of using streams.
-        block_modules (`list[str]`, *optional*):
-            List of module names that should be treated as blocks for offloading. If provided, only these modules will
-            be considered for block-level offloading. If not provided, the default block detection logic will be used.
-        exclude_kwargs (`list[str]`, *optional*):
-            List of kwarg keys that should not be processed by send_to_device. This is useful for mutable state like
-            caching lists that need to maintain their object identity across forward passes. If not provided, will be
-            inferred from the module's `_skip_keys` attribute if it exists.
-
-    Example:
-        ```python
-        >>> from diffusers import CogVideoXTransformer3DModel
-        >>> from diffusers.hooks import apply_group_offloading
-
-        >>> transformer = CogVideoXTransformer3DModel.from_pretrained(
-        ...     "THUDM/CogVideoX-5b", subfolder="transformer", torch_dtype=torch.bfloat16
-        ... )
-
-        >>> apply_group_offloading(
-        ...     transformer,
-        ...     onload_device=torch.device("cuda"),
-        ...     offload_device=torch.device("cpu"),
-        ...     offload_type="block_level",
-        ...     num_blocks_per_group=2,
-        ...     use_stream=True,
-        ... )
-        ```
-    """
-
-    onload_device = torch.device(onload_device) if isinstance(onload_device, str) else onload_device
-    offload_device = torch.device(offload_device) if isinstance(offload_device, str) else offload_device
-    offload_type = GroupOffloadingType(offload_type)
-
-    stream = None
-    if use_stream:
-        if torch.cuda.is_available():
-            stream = torch.cuda.Stream()
-        elif hasattr(torch, "xpu") and torch.xpu.is_available():
-            stream = torch.Stream()
-        else:
-            raise ValueError("Using streams for data transfer requires a CUDA device, or an Intel XPU device.")
-
-    if not use_stream and record_stream:
-        raise ValueError("`record_stream` cannot be True when `use_stream=False`.")
-    if offload_type == GroupOffloadingType.BLOCK_LEVEL and num_blocks_per_group is None:
-        raise ValueError("`num_blocks_per_group` must be provided when using `offload_type='block_level'.")
-
-    _raise_error_if_accelerate_model_or_sequential_hook_present(module)
-
-    if block_modules is None:
-        block_modules = getattr(module, "_group_offload_block_modules", None)
-
-    if exclude_kwargs is None:
-        exclude_kwargs = getattr(module, "_skip_keys", None)
-
-    config = GroupOffloadingConfig(
-        onload_device=onload_device,
-        offload_device=offload_device,
-        offload_type=offload_type,
-        num_blocks_per_group=num_blocks_per_group,
-        non_blocking=non_blocking,
-        stream=stream,
-        record_stream=record_stream,
-        low_cpu_mem_usage=low_cpu_mem_usage,
-        offload_to_disk_path=offload_to_disk_path,
-        block_modules=block_modules,
-        exclude_kwargs=exclude_kwargs,
-    )
-    _apply_group_offloading(module, config)
-
-
-def _apply_group_offloading(module: torch.nn.Module, config: GroupOffloadingConfig) -> None:
-    if config.offload_type == GroupOffloadingType.BLOCK_LEVEL:
-        _apply_group_offloading_block_level(module, config)
-    elif config.offload_type == GroupOffloadingType.LEAF_LEVEL:
-        _apply_group_offloading_leaf_level(module, config)
-    else:
-        assert False
-
-
-def _apply_group_offloading_block_level(module: torch.nn.Module, config: GroupOffloadingConfig) -> None:
-    r"""
-    This function applies offloading to groups of torch.nn.ModuleList or torch.nn.Sequential blocks, and explicitly
-    defined block modules. In comparison to the "leaf_level" offloading, which is more fine-grained, this offloading is
-    done at the top-level blocks and modules specified in block_modules.
-
-    When block_modules is provided, only those modules will be treated as blocks for offloading. For each specified
-    module, recursively apply block offloading to it.
-    """
-    if config.stream is not None and config.num_blocks_per_group != 1:
-        logger.warning(
-            f"Using streams is only supported for num_blocks_per_group=1. Got {config.num_blocks_per_group=}. Setting it to 1."
-        )
-        config.num_blocks_per_group = 1
-
-    block_modules = set(config.block_modules) if config.block_modules is not None else set()
-
-    # Create module groups for ModuleList and Sequential blocks, and explicitly defined block modules
-    modules_with_group_offloading = set()
-    unmatched_modules = []
-    matched_module_groups = []
-
-    for name, submodule in module.named_children():
-        # Check if this is an explicitly defined block module
-        if name in block_modules:
-            # Track submodule using a prefix to avoid filename collisions during disk offload.
-            # Without this, submodules sharing the same model class would be assigned identical
-            # filenames (derived from the class name).
-            prefix = f"{config.module_prefix}{name}." if config.module_prefix else f"{name}."
-            submodule_config = replace(config, module_prefix=prefix)
-
-            _apply_group_offloading_block_level(submodule, submodule_config)
-            modules_with_group_offloading.add(name)
-
-        elif isinstance(submodule, (torch.nn.ModuleList, torch.nn.Sequential)):
-            # Handle ModuleList and Sequential blocks as before
-            for i in range(0, len(submodule), config.num_blocks_per_group):
-                current_modules = list(submodule[i : i + config.num_blocks_per_group])
-                if len(current_modules) == 0:
-                    continue
-
-                group_id = f"{config.module_prefix}{name}_{i}_{i + len(current_modules) - 1}"
-                group = ModuleGroup(
-                    modules=current_modules,
-                    offload_device=config.offload_device,
-                    onload_device=config.onload_device,
-                    offload_to_disk_path=config.offload_to_disk_path,
-                    offload_leader=current_modules[-1],
-                    onload_leader=current_modules[0],
-                    non_blocking=config.non_blocking,
-                    stream=config.stream,
-                    record_stream=config.record_stream,
-                    low_cpu_mem_usage=config.low_cpu_mem_usage,
-                    onload_self=True,
-                    group_id=group_id,
-                )
-                matched_module_groups.append(group)
-                for j in range(i, i + len(current_modules)):
-                    modules_with_group_offloading.add(f"{name}.{j}")
-        else:
-            # This is an unmatched module
-            unmatched_modules.append((name, submodule))
-
-    # Apply group offloading hooks to the module groups
-    for i, group in enumerate(matched_module_groups):
-        for group_module in group.modules:
-            _apply_group_offloading_hook(group_module, group, config=config)
-
-    # Parameters and Buffers of the top-level module need to be offloaded/onloaded separately
-    # when the forward pass of this module is called. This is because the top-level module is not
-    # part of any group (as doing so would lead to no VRAM savings).
-    parameters = _gather_parameters_with_no_group_offloading_parent(module, modules_with_group_offloading)
-    buffers = _gather_buffers_with_no_group_offloading_parent(module, modules_with_group_offloading)
-    parameters = [param for _, param in parameters]
-    buffers = [buffer for _, buffer in buffers]
-
-    # Create a group for the remaining unmatched submodules of the top-level
-    # module so that they are on the correct device when the forward pass is called.
-    unmatched_modules = [unmatched_module for _, unmatched_module in unmatched_modules]
-    if len(unmatched_modules) > 0 or len(parameters) > 0 or len(buffers) > 0:
-        unmatched_group = ModuleGroup(
-            modules=unmatched_modules,
-            offload_device=config.offload_device,
-            onload_device=config.onload_device,
-            offload_to_disk_path=config.offload_to_disk_path,
-            offload_leader=module,
-            onload_leader=module,
-            parameters=parameters,
-            buffers=buffers,
-            non_blocking=False,
-            stream=None,
-            record_stream=False,
-            onload_self=True,
-            group_id=f"{config.module_prefix}{module.__class__.__name__}_unmatched_group",
-        )
-        if config.stream is None:
-            _apply_group_offloading_hook(module, unmatched_group, config=config)
-        else:
-            _apply_lazy_group_offloading_hook(module, unmatched_group, config=config)
-
-
-def _apply_group_offloading_leaf_level(module: torch.nn.Module, config: GroupOffloadingConfig) -> None:
-    r"""
-    This function applies offloading to groups of leaf modules in a torch.nn.Module. This method has minimal memory
-    requirements. However, it can be slower compared to other offloading methods due to the excessive number of device
-    synchronizations. When using devices that support streams to overlap data transfer and computation, this method can
-    reduce memory usage without any performance degradation.
-    """
-    # Create module groups for leaf modules and apply group offloading hooks
-    modules_with_group_offloading = set()
-    for name, submodule in module.named_modules():
-        if not isinstance(submodule, _GO_LC_SUPPORTED_PYTORCH_LAYERS):
-            continue
-        group = ModuleGroup(
-            modules=[submodule],
-            offload_device=config.offload_device,
-            onload_device=config.onload_device,
-            offload_to_disk_path=config.offload_to_disk_path,
-            offload_leader=submodule,
-            onload_leader=submodule,
-            non_blocking=config.non_blocking,
-            stream=config.stream,
-            record_stream=config.record_stream,
-            low_cpu_mem_usage=config.low_cpu_mem_usage,
-            onload_self=True,
-            group_id=name,
-        )
-        _apply_group_offloading_hook(submodule, group, config=config)
-        modules_with_group_offloading.add(name)
-
-    # Parameters and Buffers at all non-leaf levels need to be offloaded/onloaded separately when the forward pass
-    # of the module is called
-    module_dict = dict(module.named_modules())
-    parameters = _gather_parameters_with_no_group_offloading_parent(module, modules_with_group_offloading)
-    buffers = _gather_buffers_with_no_group_offloading_parent(module, modules_with_group_offloading)
-
-    # Find closest module parent for each parameter and buffer, and attach group hooks
-    parent_to_parameters = {}
-    for name, param in parameters:
-        parent_name = _find_parent_module_in_module_dict(name, module_dict)
-        if parent_name in parent_to_parameters:
-            parent_to_parameters[parent_name].append(param)
-        else:
-            parent_to_parameters[parent_name] = [param]
-
-    parent_to_buffers = {}
-    for name, buffer in buffers:
-        parent_name = _find_parent_module_in_module_dict(name, module_dict)
-        if parent_name in parent_to_buffers:
-            parent_to_buffers[parent_name].append(buffer)
-        else:
-            parent_to_buffers[parent_name] = [buffer]
-
-    parent_names = set(parent_to_parameters.keys()) | set(parent_to_buffers.keys())
-    for name in parent_names:
-        parameters = parent_to_parameters.get(name, [])
-        buffers = parent_to_buffers.get(name, [])
-        parent_module = module_dict[name]
-        group = ModuleGroup(
-            modules=[],
-            offload_device=config.offload_device,
-            onload_device=config.onload_device,
-            offload_leader=parent_module,
-            onload_leader=parent_module,
-            offload_to_disk_path=config.offload_to_disk_path,
-            parameters=parameters,
-            buffers=buffers,
-            non_blocking=config.non_blocking,
-            stream=config.stream,
-            record_stream=config.record_stream,
-            low_cpu_mem_usage=config.low_cpu_mem_usage,
-            onload_self=True,
-            group_id=name,
-        )
-        _apply_group_offloading_hook(parent_module, group, config=config)
-
-    if config.stream is not None:
-        # When using streams, we need to know the layer execution order for applying prefetching (to overlap data transfer
-        # and computation). Since we don't know the order beforehand, we apply a lazy prefetching hook that will find the
-        # execution order and apply prefetching in the correct order.
-        unmatched_group = ModuleGroup(
-            modules=[],
-            offload_device=config.offload_device,
-            onload_device=config.onload_device,
-            offload_to_disk_path=config.offload_to_disk_path,
-            offload_leader=module,
-            onload_leader=module,
-            parameters=None,
-            buffers=None,
-            non_blocking=False,
-            stream=None,
-            record_stream=False,
-            low_cpu_mem_usage=config.low_cpu_mem_usage,
-            onload_self=True,
-            group_id=_GROUP_ID_LAZY_LEAF,
-        )
-        _apply_lazy_group_offloading_hook(module, unmatched_group, config=config)
-
-
-def _apply_group_offloading_hook(
-    module: torch.nn.Module,
-    group: ModuleGroup,
-    *,
-    config: GroupOffloadingConfig,
-) -> None:
-    registry = HookRegistry.check_if_exists_or_initialize(module)
-
-    # We may have already registered a group offloading hook if the module had a torch.nn.Parameter whose parent
-    # is the current module. In such cases, we don't want to overwrite the existing group offloading hook.
-    if registry.get_hook(_GROUP_OFFLOADING) is None:
-        hook = GroupOffloadingHook(group, config=config)
-        registry.register_hook(hook, _GROUP_OFFLOADING)
-
-
-def _apply_lazy_group_offloading_hook(
-    module: torch.nn.Module,
-    group: ModuleGroup,
-    *,
-    config: GroupOffloadingConfig,
-) -> None:
-    registry = HookRegistry.check_if_exists_or_initialize(module)
-
-    # We may have already registered a group offloading hook if the module had a torch.nn.Parameter whose parent
-    # is the current module. In such cases, we don't want to overwrite the existing group offloading hook.
-    if registry.get_hook(_GROUP_OFFLOADING) is None:
-        hook = GroupOffloadingHook(group, config=config)
-        registry.register_hook(hook, _GROUP_OFFLOADING)
-
-    lazy_prefetch_hook = LazyPrefetchGroupOffloadingHook()
-    registry.register_hook(lazy_prefetch_hook, _LAZY_PREFETCH_GROUP_OFFLOADING)
-
-
-def _gather_parameters_with_no_group_offloading_parent(
-    module: torch.nn.Module, modules_with_group_offloading: Set[str]
-) -> list[torch.nn.Parameter]:
-    parameters = []
-    for name, parameter in module.named_parameters():
-        has_parent_with_group_offloading = False
-        atoms = name.split(".")
-        while len(atoms) > 0:
-            parent_name = ".".join(atoms)
-            if parent_name in modules_with_group_offloading:
-                has_parent_with_group_offloading = True
-                break
-            atoms.pop()
-        if not has_parent_with_group_offloading:
-            parameters.append((name, parameter))
-    return parameters
-
-
-def _gather_buffers_with_no_group_offloading_parent(
-    module: torch.nn.Module, modules_with_group_offloading: Set[str]
-) -> list[torch.Tensor]:
-    buffers = []
-    for name, buffer in module.named_buffers():
-        has_parent_with_group_offloading = False
-        atoms = name.split(".")
-        while len(atoms) > 0:
-            parent_name = ".".join(atoms)
-            if parent_name in modules_with_group_offloading:
-                has_parent_with_group_offloading = True
-                break
-            atoms.pop()
-        if not has_parent_with_group_offloading:
-            buffers.append((name, buffer))
-    return buffers
-
-
-def _find_parent_module_in_module_dict(name: str, module_dict: dict[str, torch.nn.Module]) -> str:
-    atoms = name.split(".")
-    while len(atoms) > 0:
-        parent_name = ".".join(atoms)
-        if parent_name in module_dict:
-            return parent_name
-        atoms.pop()
-    return ""
-
-
-def _raise_error_if_accelerate_model_or_sequential_hook_present(module: torch.nn.Module) -> None:
-    if not is_accelerate_available():
-        return
-    for name, submodule in module.named_modules():
-        if not hasattr(submodule, "_hf_hook"):
-            continue
-        if isinstance(submodule._hf_hook, (AlignDevicesHook, CpuOffload)):
-            raise ValueError(
-                f"Cannot apply group offloading to a module that is already applying an alternative "
-                f"offloading strategy from Accelerate. If you want to apply group offloading, please "
-                f"disable the existing offloading strategy first. Offending module: {name} ({type(submodule)})"
-            )
-
-
-def _get_top_level_group_offload_hook(module: torch.nn.Module) -> GroupOffloadingHook | None:
-    for submodule in module.modules():
-        if hasattr(submodule, "_diffusers_hook"):
-            group_offloading_hook = submodule._diffusers_hook.get_hook(_GROUP_OFFLOADING)
-            if group_offloading_hook is not None:
-                return group_offloading_hook
-    return None
-
-
-def _is_group_offload_enabled(module: torch.nn.Module) -> bool:
-    top_level_group_offload_hook = _get_top_level_group_offload_hook(module)
-    return top_level_group_offload_hook is not None
-
-
-def _get_group_onload_device(module: torch.nn.Module) -> torch.device:
-    top_level_group_offload_hook = _get_top_level_group_offload_hook(module)
-    if top_level_group_offload_hook is not None:
-        return top_level_group_offload_hook.config.onload_device
-    raise ValueError("Group offloading is not enabled for the provided module.")
-
-
-def _compute_group_hash(group_id):
-    hashed_id = hashlib.sha256(group_id.encode("utf-8")).hexdigest()
-    # first 16 characters for a reasonably short but unique name
-    return hashed_id[:16]
-
-
-def _maybe_remove_and_reapply_group_offloading(module: torch.nn.Module) -> None:
-    r"""
-    Removes the group offloading hook from the module and re-applies it. This is useful when the module has been
-    modified in-place and the group offloading hook references-to-tensors needs to be updated. The in-place
-    modification can happen in a number of ways, for example, fusing QKV or unloading/loading LoRAs on-the-fly.
-
-    In this implementation, we make an assumption that group offloading has only been applied at the top-level module,
-    and therefore all submodules have the same onload and offload devices. If this assumption is not true, say in the
-    case where user has applied group offloading at multiple levels, this function will not work as expected.
-
-    There is some performance penalty associated with doing this when non-default streams are used, because we need to
-    retrace the execution order of the layers with `LazyPrefetchGroupOffloadingHook`.
-    """
-    top_level_group_offload_hook = _get_top_level_group_offload_hook(module)
-
-    if top_level_group_offload_hook is None:
-        return
-
-    registry = HookRegistry.check_if_exists_or_initialize(module)
-    registry.remove_hook(_GROUP_OFFLOADING, recurse=True)
-    registry.remove_hook(_LAYER_EXECUTION_TRACKER, recurse=True)
-    registry.remove_hook(_LAZY_PREFETCH_GROUP_OFFLOADING, recurse=True)
-
-    _apply_group_offloading(module, top_level_group_offload_hook.config)

diffusers/hooks/hooks.py

@@ -1,312 +0,0 @@
-# Copyright 2026 The HuggingFace Team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-import functools
-from typing import Any
-
-import torch
-
-from ..utils.logging import get_logger
-from ..utils.torch_utils import unwrap_module
-
-
-logger = get_logger(__name__)  # pylint: disable=invalid-name
-
-
-class BaseState:
-    def reset(self, *args, **kwargs) -> None:
-        raise NotImplementedError(
-            "BaseState::reset is not implemented. Please implement this method in the derived class."
-        )
-
-
-class StateManager:
-    def __init__(self, state_cls: BaseState, init_args=None, init_kwargs=None):
-        self._state_cls = state_cls
-        self._init_args = init_args if init_args is not None else ()
-        self._init_kwargs = init_kwargs if init_kwargs is not None else {}
-        self._state_cache = {}
-        self._current_context = None
-
-    def get_state(self):
-        if self._current_context is None:
-            raise ValueError("No context is set. Please set a context before retrieving the state.")
-        if self._current_context not in self._state_cache.keys():
-            self._state_cache[self._current_context] = self._state_cls(*self._init_args, **self._init_kwargs)
-        return self._state_cache[self._current_context]
-
-    def set_context(self, name: str) -> None:
-        self._current_context = name
-
-    def reset(self, *args, **kwargs) -> None:
-        for name, state in list(self._state_cache.items()):
-            state.reset(*args, **kwargs)
-            self._state_cache.pop(name)
-        self._current_context = None
-
-
-class ModelHook:
-    r"""
-    A hook that contains callbacks to be executed just before and after the forward method of a model.
-    """
-
-    _is_stateful = False
-
-    def __init__(self):
-        self.fn_ref: "HookFunctionReference" = None
-
-    def initialize_hook(self, module: torch.nn.Module) -> torch.nn.Module:
-        r"""
-        Hook that is executed when a model is initialized.
-
-        Args:
-            module (`torch.nn.Module`):
-                The module attached to this hook.
-        """
-        return module
-
-    def deinitalize_hook(self, module: torch.nn.Module) -> torch.nn.Module:
-        r"""
-        Hook that is executed when a model is deinitialized.
-
-        Args:
-            module (`torch.nn.Module`):
-                The module attached to this hook.
-        """
-        return module
-
-    def pre_forward(self, module: torch.nn.Module, *args, **kwargs) -> tuple[tuple[Any], dict[str, Any]]:
-        r"""
-        Hook that is executed just before the forward method of the model.
-
-        Args:
-            module (`torch.nn.Module`):
-                The module whose forward pass will be executed just after this event.
-            args (`tuple[Any]`):
-                The positional arguments passed to the module.
-            kwargs (`dict[Str, Any]`):
-                The keyword arguments passed to the module.
-        Returns:
-            `tuple[tuple[Any], dict[Str, Any]]`:
-                A tuple with the treated `args` and `kwargs`.
-        """
-        return args, kwargs
-
-    def post_forward(self, module: torch.nn.Module, output: Any) -> Any:
-        r"""
-        Hook that is executed just after the forward method of the model.
-
-        Args:
-            module (`torch.nn.Module`):
-                The module whose forward pass been executed just before this event.
-            output (`Any`):
-                The output of the module.
-        Returns:
-            `Any`: The processed `output`.
-        """
-        return output
-
-    def detach_hook(self, module: torch.nn.Module) -> torch.nn.Module:
-        r"""
-        Hook that is executed when the hook is detached from a module.
-
-        Args:
-            module (`torch.nn.Module`):
-                The module detached from this hook.
-        """
-        return module
-
-    def reset_state(self, module: torch.nn.Module):
-        if self._is_stateful:
-            raise NotImplementedError("This hook is stateful and needs to implement the `reset_state` method.")
-        return module
-
-    def _set_context(self, module: torch.nn.Module, name: str) -> None:
-        # Iterate over all attributes of the hook to see if any of them have the type `StateManager`. If so, call `set_context` on them.
-        for attr_name in dir(self):
-            attr = getattr(self, attr_name)
-            if isinstance(attr, StateManager):
-                attr.set_context(name)
-        return module
-
-
-class HookFunctionReference:
-    def __init__(self) -> None:
-        """A container class that maintains mutable references to forward pass functions in a hook chain.
-
-        Its mutable nature allows the hook system to modify the execution chain dynamically without rebuilding the
-        entire forward pass structure.
-
-        Attributes:
-            pre_forward: A callable that processes inputs before the main forward pass.
-            post_forward: A callable that processes outputs after the main forward pass.
-            forward: The current forward function in the hook chain.
-            original_forward: The original forward function, stored when a hook provides a custom new_forward.
-
-        The class enables hook removal by allowing updates to the forward chain through reference modification rather
-        than requiring reconstruction of the entire chain. When a hook is removed, only the relevant references need to
-        be updated, preserving the execution order of the remaining hooks.
-        """
-        self.pre_forward = None
-        self.post_forward = None
-        self.forward = None
-        self.original_forward = None
-
-
-class HookRegistry:
-    def __init__(self, module_ref: torch.nn.Module) -> None:
-        super().__init__()
-
-        self.hooks: dict[str, ModelHook] = {}
-
-        self._module_ref = module_ref
-        self._hook_order = []
-        self._fn_refs = []
-
-    def register_hook(self, hook: ModelHook, name: str) -> None:
-        if name in self.hooks.keys():
-            raise ValueError(
-                f"Hook with name {name} already exists in the registry. Please use a different name or "
-                f"first remove the existing hook and then add a new one."
-            )
-
-        self._module_ref = hook.initialize_hook(self._module_ref)
-
-        def create_new_forward(function_reference: HookFunctionReference):
-            def new_forward(module, *args, **kwargs):
-                args, kwargs = function_reference.pre_forward(module, *args, **kwargs)
-                output = function_reference.forward(*args, **kwargs)
-                return function_reference.post_forward(module, output)
-
-            return new_forward
-
-        forward = self._module_ref.forward
-
-        fn_ref = HookFunctionReference()
-        fn_ref.pre_forward = hook.pre_forward
-        fn_ref.post_forward = hook.post_forward
-        fn_ref.forward = forward
-
-        if hasattr(hook, "new_forward"):
-            fn_ref.original_forward = forward
-            fn_ref.forward = functools.update_wrapper(
-                functools.partial(hook.new_forward, self._module_ref), hook.new_forward
-            )
-
-        rewritten_forward = create_new_forward(fn_ref)
-        # Wrap from the original `forward` so `inspect.signature` follows `__wrapped__` to the real
-        # signature instead of the generic `(module, *args, **kwargs)`, which breaks `torch.export`.
-        self._module_ref.forward = functools.update_wrapper(
-            functools.partial(rewritten_forward, self._module_ref), forward
-        )
-
-        hook.fn_ref = fn_ref
-        self.hooks[name] = hook
-        self._hook_order.append(name)
-        self._fn_refs.append(fn_ref)
-
-    def get_hook(self, name: str) -> ModelHook | None:
-        return self.hooks.get(name, None)
-
-    def remove_hook(self, name: str, recurse: bool = True) -> None:
-        if name in self.hooks.keys():
-            num_hooks = len(self._hook_order)
-            hook = self.hooks[name]
-            index = self._hook_order.index(name)
-            fn_ref = self._fn_refs[index]
-
-            old_forward = fn_ref.forward
-            if fn_ref.original_forward is not None:
-                old_forward = fn_ref.original_forward
-
-            if index == num_hooks - 1:
-                self._module_ref.forward = old_forward
-            else:
-                self._fn_refs[index + 1].forward = old_forward
-
-            self._module_ref = hook.deinitalize_hook(self._module_ref)
-            del self.hooks[name]
-            self._hook_order.pop(index)
-            self._fn_refs.pop(index)
-
-        if recurse:
-            for module_name, module in self._module_ref.named_modules():
-                if module_name == "":
-                    continue
-                if hasattr(module, "_diffusers_hook"):
-                    module._diffusers_hook.remove_hook(name, recurse=False)
-
-    def reset_stateful_hooks(self, recurse: bool = True) -> None:
-        for hook_name in reversed(self._hook_order):
-            hook = self.hooks[hook_name]
-            if hook._is_stateful:
-                hook.reset_state(self._module_ref)
-
-        if recurse:
-            for module_name, module in unwrap_module(self._module_ref).named_modules():
-                if module_name == "":
-                    continue
-                module = unwrap_module(module)
-                if hasattr(module, "_diffusers_hook"):
-                    module._diffusers_hook.reset_stateful_hooks(recurse=False)
-
-    @classmethod
-    def check_if_exists_or_initialize(cls, module: torch.nn.Module) -> "HookRegistry":
-        if not hasattr(module, "_diffusers_hook"):
-            module._diffusers_hook = cls(module)
-        return module._diffusers_hook
-
-    def _set_context(self, name: str | None = None) -> None:
-        for hook_name in reversed(self._hook_order):
-            hook = self.hooks[hook_name]
-            if hook._is_stateful:
-                hook._set_context(self._module_ref, name)
-
-        for registry in self._get_child_registries():
-            registry._set_context(name)
-
-    def _get_child_registries(self) -> list["HookRegistry"]:
-        """Return registries of child modules, using a cached list when available.
-
-        The cache is built on first call and reused for subsequent calls. This avoids the cost of walking the full
-        module tree via named_modules() on every _set_context call, which is significant for large models (e.g. ~2.7ms
-        per call on Flux2).
-        """
-        if not hasattr(self, "_child_registries_cache"):
-            self._child_registries_cache = None
-
-        if self._child_registries_cache is not None:
-            return self._child_registries_cache
-
-        registries = []
-        for module_name, module in unwrap_module(self._module_ref).named_modules():
-            if module_name == "":
-                continue
-            module = unwrap_module(module)
-            if hasattr(module, "_diffusers_hook"):
-                registries.append(module._diffusers_hook)
-        self._child_registries_cache = registries
-        return registries
-
-    def __repr__(self) -> str:
-        registry_repr = ""
-        for i, hook_name in enumerate(self._hook_order):
-            if self.hooks[hook_name].__class__.__repr__ is not object.__repr__:
-                hook_repr = self.hooks[hook_name].__repr__()
-            else:
-                hook_repr = self.hooks[hook_name].__class__.__name__
-            registry_repr += f"  ({i}) {hook_name} - {hook_repr}"
-            if i < len(self._hook_order) - 1:
-                registry_repr += "\n"
-        return f"HookRegistry(\n{registry_repr}\n)"

diffusers/hooks/layer_skip.py

@@ -1,263 +0,0 @@
-# Copyright 2026 The HuggingFace Team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-import math
-from dataclasses import asdict, dataclass
-from typing import Callable
-
-import torch
-
-from ..utils import get_logger
-from ..utils.torch_utils import unwrap_module
-from ._common import (
-    _ALL_TRANSFORMER_BLOCK_IDENTIFIERS,
-    _ATTENTION_CLASSES,
-    _FEEDFORWARD_CLASSES,
-    _get_submodule_from_fqn,
-)
-from ._helpers import AttentionProcessorRegistry, TransformerBlockRegistry
-from .hooks import HookRegistry, ModelHook
-
-
-logger = get_logger(__name__)  # pylint: disable=invalid-name
-
-_LAYER_SKIP_HOOK = "layer_skip_hook"
-
-
-# Aryan/YiYi TODO: we need to make guider class a config mixin so I think this is not needed
-# either remove or make it serializable
-@dataclass
-class LayerSkipConfig:
-    r"""
-    Configuration for skipping internal transformer blocks when executing a transformer model.
-
-    Args:
-        indices (`list[int]`):
-            The indices of the layer to skip. This is typically the first layer in the transformer block.
-        fqn (`str`, defaults to `"auto"`):
-            The fully qualified name identifying the stack of transformer blocks. Typically, this is
-            `transformer_blocks`, `single_transformer_blocks`, `blocks`, `layers`, or `temporal_transformer_blocks`.
-            For automatic detection, set this to `"auto"`. "auto" only works on DiT models. For UNet models, you must
-            provide the correct fqn.
-        skip_attention (`bool`, defaults to `True`):
-            Whether to skip attention blocks.
-        skip_ff (`bool`, defaults to `True`):
-            Whether to skip feed-forward blocks.
-        skip_attention_scores (`bool`, defaults to `False`):
-            Whether to skip attention score computation in the attention blocks. This is equivalent to using `value`
-            projections as the output of scaled dot product attention.
-        dropout (`float`, defaults to `1.0`):
-            The dropout probability for dropping the outputs of the skipped layers. By default, this is set to `1.0`,
-            meaning that the outputs of the skipped layers are completely ignored. If set to `0.0`, the outputs of the
-            skipped layers are fully retained, which is equivalent to not skipping any layers.
-    """
-
-    indices: list[int]
-    fqn: str = "auto"
-    skip_attention: bool = True
-    skip_attention_scores: bool = False
-    skip_ff: bool = True
-    dropout: float = 1.0
-
-    def __post_init__(self):
-        if not (0 <= self.dropout <= 1):
-            raise ValueError(f"Expected `dropout` to be between 0.0 and 1.0, but got {self.dropout}.")
-        if not math.isclose(self.dropout, 1.0) and self.skip_attention_scores:
-            raise ValueError(
-                "Cannot set `skip_attention_scores` to True when `dropout` is not 1.0. Please set `dropout` to 1.0."
-            )
-
-    def to_dict(self):
-        return asdict(self)
-
-    @staticmethod
-    def from_dict(data: dict) -> "LayerSkipConfig":
-        return LayerSkipConfig(**data)
-
-
-class AttentionScoreSkipFunctionMode(torch.overrides.TorchFunctionMode):
-    def __torch_function__(self, func, types, args=(), kwargs=None):
-        if kwargs is None:
-            kwargs = {}
-        if func is torch.nn.functional.scaled_dot_product_attention:
-            query = kwargs.get("query", None)
-            key = kwargs.get("key", None)
-            value = kwargs.get("value", None)
-            query = query if query is not None else args[0]
-            key = key if key is not None else args[1]
-            value = value if value is not None else args[2]
-            # If the Q sequence length does not match KV sequence length, methods like
-            # Perturbed Attention Guidance cannot be used (because the caller expects
-            # the same sequence length as Q, but if we return V here, it will not match).
-            # When Q.shape[2] != V.shape[2], PAG will essentially not be applied and
-            # the overall effect would that be of normal CFG with a scale of (guidance_scale + perturbed_guidance_scale).
-            if query.shape[2] == value.shape[2]:
-                return value
-        return func(*args, **kwargs)
-
-
-class AttentionProcessorSkipHook(ModelHook):
-    def __init__(self, skip_processor_output_fn: Callable, skip_attention_scores: bool = False, dropout: float = 1.0):
-        self.skip_processor_output_fn = skip_processor_output_fn
-        self.skip_attention_scores = skip_attention_scores
-        self.dropout = dropout
-
-    def new_forward(self, module: torch.nn.Module, *args, **kwargs):
-        if self.skip_attention_scores:
-            if not math.isclose(self.dropout, 1.0):
-                raise ValueError(
-                    "Cannot set `skip_attention_scores` to True when `dropout` is not 1.0. Please set `dropout` to 1.0."
-                )
-            with AttentionScoreSkipFunctionMode():
-                output = self.fn_ref.original_forward(*args, **kwargs)
-        else:
-            if math.isclose(self.dropout, 1.0):
-                output = self.skip_processor_output_fn(module, *args, **kwargs)
-            else:
-                output = self.fn_ref.original_forward(*args, **kwargs)
-                output = torch.nn.functional.dropout(output, p=self.dropout)
-        return output
-
-
-class FeedForwardSkipHook(ModelHook):
-    def __init__(self, dropout: float):
-        super().__init__()
-        self.dropout = dropout
-
-    def new_forward(self, module: torch.nn.Module, *args, **kwargs):
-        if math.isclose(self.dropout, 1.0):
-            output = kwargs.get("hidden_states", None)
-            if output is None:
-                output = kwargs.get("x", None)
-            if output is None and len(args) > 0:
-                output = args[0]
-        else:
-            output = self.fn_ref.original_forward(*args, **kwargs)
-            output = torch.nn.functional.dropout(output, p=self.dropout)
-        return output
-
-
-class TransformerBlockSkipHook(ModelHook):
-    def __init__(self, dropout: float):
-        super().__init__()
-        self.dropout = dropout
-
-    def initialize_hook(self, module):
-        self._metadata = TransformerBlockRegistry.get(unwrap_module(module).__class__)
-        return module
-
-    def new_forward(self, module: torch.nn.Module, *args, **kwargs):
-        if math.isclose(self.dropout, 1.0):
-            original_hidden_states = self._metadata._get_parameter_from_args_kwargs("hidden_states", args, kwargs)
-            if self._metadata.return_encoder_hidden_states_index is None:
-                output = original_hidden_states
-            else:
-                original_encoder_hidden_states = self._metadata._get_parameter_from_args_kwargs(
-                    "encoder_hidden_states", args, kwargs
-                )
-                output = (original_hidden_states, original_encoder_hidden_states)
-        else:
-            output = self.fn_ref.original_forward(*args, **kwargs)
-            output = torch.nn.functional.dropout(output, p=self.dropout)
-        return output
-
-
-def apply_layer_skip(module: torch.nn.Module, config: LayerSkipConfig) -> None:
-    r"""
-    Apply layer skipping to internal layers of a transformer.
-
-    Args:
-        module (`torch.nn.Module`):
-            The transformer model to which the layer skip hook should be applied.
-        config (`LayerSkipConfig`):
-            The configuration for the layer skip hook.
-
-    Example:
-
-    ```python
-    >>> from diffusers import apply_layer_skip_hook, CogVideoXTransformer3DModel, LayerSkipConfig
-
-    >>> transformer = CogVideoXTransformer3DModel.from_pretrained("THUDM/CogVideoX-5b", torch_dtype=torch.bfloat16)
-    >>> config = LayerSkipConfig(layer_index=[10, 20], fqn="transformer_blocks")
-    >>> apply_layer_skip_hook(transformer, config)
-    ```
-    """
-    _apply_layer_skip_hook(module, config)
-
-
-def _apply_layer_skip_hook(module: torch.nn.Module, config: LayerSkipConfig, name: str | None = None) -> None:
-    name = name or _LAYER_SKIP_HOOK
-
-    if config.skip_attention and config.skip_attention_scores:
-        raise ValueError("Cannot set both `skip_attention` and `skip_attention_scores` to True. Please choose one.")
-    if not math.isclose(config.dropout, 1.0) and config.skip_attention_scores:
-        raise ValueError(
-            "Cannot set `skip_attention_scores` to True when `dropout` is not 1.0. Please set `dropout` to 1.0."
-        )
-
-    if config.fqn == "auto":
-        for identifier in _ALL_TRANSFORMER_BLOCK_IDENTIFIERS:
-            if hasattr(module, identifier):
-                config.fqn = identifier
-                break
-        else:
-            raise ValueError(
-                "Could not find a suitable identifier for the transformer blocks automatically. Please provide a valid "
-                "`fqn` (fully qualified name) that identifies a stack of transformer blocks."
-            )
-
-    transformer_blocks = _get_submodule_from_fqn(module, config.fqn)
-    if transformer_blocks is None or not isinstance(transformer_blocks, torch.nn.ModuleList):
-        raise ValueError(
-            f"Could not find {config.fqn} in the provided module, or configured `fqn` (fully qualified name) does not identify "
-            f"a `torch.nn.ModuleList`. Please provide a valid `fqn` that identifies a stack of transformer blocks."
-        )
-    if len(config.indices) == 0:
-        raise ValueError("Layer index list is empty. Please provide a non-empty list of layer indices to skip.")
-
-    blocks_found = False
-    for i, block in enumerate(transformer_blocks):
-        if i not in config.indices:
-            continue
-
-        blocks_found = True
-
-        if config.skip_attention and config.skip_ff:
-            logger.debug(f"Applying TransformerBlockSkipHook to '{config.fqn}.{i}'")
-            registry = HookRegistry.check_if_exists_or_initialize(block)
-            hook = TransformerBlockSkipHook(config.dropout)
-            registry.register_hook(hook, name)
-
-        elif config.skip_attention or config.skip_attention_scores:
-            for submodule_name, submodule in block.named_modules():
-                if isinstance(submodule, _ATTENTION_CLASSES) and not submodule.is_cross_attention:
-                    logger.debug(f"Applying AttentionProcessorSkipHook to '{config.fqn}.{i}.{submodule_name}'")
-                    output_fn = AttentionProcessorRegistry.get(submodule.processor.__class__).skip_processor_output_fn
-                    registry = HookRegistry.check_if_exists_or_initialize(submodule)
-                    hook = AttentionProcessorSkipHook(output_fn, config.skip_attention_scores, config.dropout)
-                    registry.register_hook(hook, name)
-
-        if config.skip_ff:
-            for submodule_name, submodule in block.named_modules():
-                if isinstance(submodule, _FEEDFORWARD_CLASSES):
-                    logger.debug(f"Applying FeedForwardSkipHook to '{config.fqn}.{i}.{submodule_name}'")
-                    registry = HookRegistry.check_if_exists_or_initialize(submodule)
-                    hook = FeedForwardSkipHook(config.dropout)
-                    registry.register_hook(hook, name)
-
-    if not blocks_found:
-        raise ValueError(
-            f"Could not find any transformer blocks matching the provided indices {config.indices} and "
-            f"fully qualified name '{config.fqn}'. Please check the indices and fqn for correctness."
-        )

diffusers/hooks/layerwise_casting.py

@@ -1,240 +0,0 @@
-# Copyright 2026 The HuggingFace Team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-import re
-from typing import Type
-
-import torch
-
-from ..utils import get_logger, is_peft_available, is_peft_version
-from ._common import _GO_LC_SUPPORTED_PYTORCH_LAYERS
-from .hooks import HookRegistry, ModelHook
-
-
-logger = get_logger(__name__)  # pylint: disable=invalid-name
-
-
-# fmt: off
-_LAYERWISE_CASTING_HOOK = "layerwise_casting"
-_PEFT_AUTOCAST_DISABLE_HOOK = "peft_autocast_disable"
-DEFAULT_SKIP_MODULES_PATTERN = ("pos_embed", "patch_embed", "norm", "^proj_in$", "^proj_out$")
-# fmt: on
-
-_SHOULD_DISABLE_PEFT_INPUT_AUTOCAST = is_peft_available() and is_peft_version(">", "0.14.0")
-if _SHOULD_DISABLE_PEFT_INPUT_AUTOCAST:
-    from peft.helpers import disable_input_dtype_casting
-    from peft.tuners.tuners_utils import BaseTunerLayer
-
-
-class LayerwiseCastingHook(ModelHook):
-    r"""
-    A hook that casts the weights of a module to a high precision dtype for computation, and to a low precision dtype
-    for storage. This process may lead to quality loss in the output, but can significantly reduce the memory
-    footprint.
-    """
-
-    _is_stateful = False
-
-    def __init__(self, storage_dtype: torch.dtype, compute_dtype: torch.dtype, non_blocking: bool) -> None:
-        self.storage_dtype = storage_dtype
-        self.compute_dtype = compute_dtype
-        self.non_blocking = non_blocking
-
-    def initialize_hook(self, module: torch.nn.Module):
-        module.to(dtype=self.storage_dtype, non_blocking=self.non_blocking)
-        return module
-
-    def deinitalize_hook(self, module: torch.nn.Module):
-        raise NotImplementedError(
-            "LayerwiseCastingHook does not support deinitialization. A model once enabled with layerwise casting will "
-            "have casted its weights to a lower precision dtype for storage. Casting this back to the original dtype "
-            "will lead to precision loss, which might have an impact on the model's generation quality. The model should "
-            "be re-initialized and loaded in the original dtype."
-        )
-
-    def pre_forward(self, module: torch.nn.Module, *args, **kwargs):
-        module.to(dtype=self.compute_dtype, non_blocking=self.non_blocking)
-        return args, kwargs
-
-    def post_forward(self, module: torch.nn.Module, output):
-        module.to(dtype=self.storage_dtype, non_blocking=self.non_blocking)
-        return output
-
-
-class PeftInputAutocastDisableHook(ModelHook):
-    r"""
-    A hook that disables the casting of inputs to the module weight dtype during the forward pass. By default, PEFT
-    casts the inputs to the weight dtype of the module, which can lead to precision loss.
-
-    The reasons for needing this are:
-        - If we don't add PEFT layers' weight names to `skip_modules_pattern` when applying layerwise casting, the
-          inputs will be casted to the, possibly lower precision, storage dtype. Reference:
-          https://github.com/huggingface/peft/blob/0facdebf6208139cbd8f3586875acb378813dd97/src/peft/tuners/lora/layer.py#L706
-        - We can, on our end, use something like accelerate's `send_to_device` but for dtypes. This way, we can ensure
-          that the inputs are casted to the computation dtype correctly always. However, there are two goals we are
-          hoping to achieve:
-            1. Making forward implementations independent of device/dtype casting operations as much as possible.
-            2. Performing inference without losing information from casting to different precisions. With the current
-               PEFT implementation (as linked in the reference above), and assuming running layerwise casting inference
-               with storage_dtype=torch.float8_e4m3fn and compute_dtype=torch.bfloat16, inputs are cast to
-               torch.float8_e4m3fn in the lora layer. We will then upcast back to torch.bfloat16 when we continue the
-               forward pass in PEFT linear forward or Diffusers layer forward, with a `send_to_dtype` operation from
-               LayerwiseCastingHook. This will be a lossy operation and result in poorer generation quality.
-    """
-
-    def new_forward(self, module: torch.nn.Module, *args, **kwargs):
-        with disable_input_dtype_casting(module):
-            return self.fn_ref.original_forward(*args, **kwargs)
-
-
-def apply_layerwise_casting(
-    module: torch.nn.Module,
-    storage_dtype: torch.dtype,
-    compute_dtype: torch.dtype,
-    skip_modules_pattern: str | tuple[str, ...] = "auto",
-    skip_modules_classes: tuple[Type[torch.nn.Module], ...] | None = None,
-    non_blocking: bool = False,
-) -> None:
-    r"""
-    Applies layerwise casting to a given module. The module expected here is a Diffusers ModelMixin but it can be any
-    nn.Module using diffusers layers or pytorch primitives.
-
-    Example:
-
-    ```python
-    >>> import torch
-    >>> from diffusers import CogVideoXTransformer3DModel
-
-    >>> transformer = CogVideoXTransformer3DModel.from_pretrained(
-    ...     model_id, subfolder="transformer", torch_dtype=torch.bfloat16
-    ... )
-
-    >>> apply_layerwise_casting(
-    ...     transformer,
-    ...     storage_dtype=torch.float8_e4m3fn,
-    ...     compute_dtype=torch.bfloat16,
-    ...     skip_modules_pattern=["patch_embed", "norm", "proj_out"],
-    ...     non_blocking=True,
-    ... )
-    ```
-
-    Args:
-        module (`torch.nn.Module`):
-            The module whose leaf modules will be cast to a high precision dtype for computation, and to a low
-            precision dtype for storage.
-        storage_dtype (`torch.dtype`):
-            The dtype to cast the module to before/after the forward pass for storage.
-        compute_dtype (`torch.dtype`):
-            The dtype to cast the module to during the forward pass for computation.
-        skip_modules_pattern (`tuple[str, ...]`, defaults to `"auto"`):
-            A list of patterns to match the names of the modules to skip during the layerwise casting process. If set
-            to `"auto"`, the default patterns are used. If set to `None`, no modules are skipped. If set to `None`
-            alongside `skip_modules_classes` being `None`, the layerwise casting is applied directly to the module
-            instead of its internal submodules.
-        skip_modules_classes (`tuple[Type[torch.nn.Module], ...]`, defaults to `None`):
-            A list of module classes to skip during the layerwise casting process.
-        non_blocking (`bool`, defaults to `False`):
-            If `True`, the weight casting operations are non-blocking.
-    """
-    if skip_modules_pattern == "auto":
-        skip_modules_pattern = DEFAULT_SKIP_MODULES_PATTERN
-
-    if skip_modules_classes is None and skip_modules_pattern is None:
-        apply_layerwise_casting_hook(module, storage_dtype, compute_dtype, non_blocking)
-        return
-
-    _apply_layerwise_casting(
-        module,
-        storage_dtype,
-        compute_dtype,
-        skip_modules_pattern,
-        skip_modules_classes,
-        non_blocking,
-    )
-    _disable_peft_input_autocast(module)
-
-
-def _apply_layerwise_casting(
-    module: torch.nn.Module,
-    storage_dtype: torch.dtype,
-    compute_dtype: torch.dtype,
-    skip_modules_pattern: tuple[str, ...] | None = None,
-    skip_modules_classes: tuple[Type[torch.nn.Module], ...] | None = None,
-    non_blocking: bool = False,
-    _prefix: str = "",
-) -> None:
-    should_skip = (skip_modules_classes is not None and isinstance(module, skip_modules_classes)) or (
-        skip_modules_pattern is not None and any(re.search(pattern, _prefix) for pattern in skip_modules_pattern)
-    )
-    if should_skip:
-        logger.debug(f'Skipping layerwise casting for layer "{_prefix}"')
-        return
-
-    if isinstance(module, _GO_LC_SUPPORTED_PYTORCH_LAYERS):
-        logger.debug(f'Applying layerwise casting to layer "{_prefix}"')
-        apply_layerwise_casting_hook(module, storage_dtype, compute_dtype, non_blocking)
-        return
-
-    for name, submodule in module.named_children():
-        layer_name = f"{_prefix}.{name}" if _prefix else name
-        _apply_layerwise_casting(
-            submodule,
-            storage_dtype,
-            compute_dtype,
-            skip_modules_pattern,
-            skip_modules_classes,
-            non_blocking,
-            _prefix=layer_name,
-        )
-
-
-def apply_layerwise_casting_hook(
-    module: torch.nn.Module, storage_dtype: torch.dtype, compute_dtype: torch.dtype, non_blocking: bool
-) -> None:
-    r"""
-    Applies a `LayerwiseCastingHook` to a given module.
-
-    Args:
-        module (`torch.nn.Module`):
-            The module to attach the hook to.
-        storage_dtype (`torch.dtype`):
-            The dtype to cast the module to before the forward pass.
-        compute_dtype (`torch.dtype`):
-            The dtype to cast the module to during the forward pass.
-        non_blocking (`bool`):
-            If `True`, the weight casting operations are non-blocking.
-    """
-    registry = HookRegistry.check_if_exists_or_initialize(module)
-    hook = LayerwiseCastingHook(storage_dtype, compute_dtype, non_blocking)
-    registry.register_hook(hook, _LAYERWISE_CASTING_HOOK)
-
-
-def _is_layerwise_casting_active(module: torch.nn.Module) -> bool:
-    for submodule in module.modules():
-        if (
-            hasattr(submodule, "_diffusers_hook")
-            and submodule._diffusers_hook.get_hook(_LAYERWISE_CASTING_HOOK) is not None
-        ):
-            return True
-    return False
-
-
-def _disable_peft_input_autocast(module: torch.nn.Module) -> None:
-    if not _SHOULD_DISABLE_PEFT_INPUT_AUTOCAST:
-        return
-    for submodule in module.modules():
-        if isinstance(submodule, BaseTunerLayer) and _is_layerwise_casting_active(submodule):
-            registry = HookRegistry.check_if_exists_or_initialize(submodule)
-            hook = PeftInputAutocastDisableHook()
-            registry.register_hook(hook, _PEFT_AUTOCAST_DISABLE_HOOK)

diffusers/hooks/mag_cache.py

@@ -1,468 +0,0 @@
-# Copyright 2026 The HuggingFace Team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-from dataclasses import dataclass
-from typing import List, Optional, Tuple, Union
-
-import torch
-
-from ..utils import get_logger
-from ..utils.torch_utils import unwrap_module
-from ._common import _ALL_TRANSFORMER_BLOCK_IDENTIFIERS
-from ._helpers import TransformerBlockRegistry
-from .hooks import BaseState, HookRegistry, ModelHook, StateManager
-
-
-logger = get_logger(__name__)  # pylint: disable=invalid-name
-
-_MAG_CACHE_LEADER_BLOCK_HOOK = "mag_cache_leader_block_hook"
-_MAG_CACHE_BLOCK_HOOK = "mag_cache_block_hook"
-
-# Default Mag Ratios for Flux models (Dev/Schnell) are provided for convenience.
-# Users must explicitly pass these to the config if using Flux.
-# Reference: https://github.com/Zehong-Ma/MagCache
-FLUX_MAG_RATIOS = torch.tensor(
-    [1.0]
-    + [
-        1.21094,
-        1.11719,
-        1.07812,
-        1.0625,
-        1.03906,
-        1.03125,
-        1.03906,
-        1.02344,
-        1.03125,
-        1.02344,
-        0.98047,
-        1.01562,
-        1.00781,
-        1.0,
-        1.00781,
-        1.0,
-        1.00781,
-        1.0,
-        1.0,
-        0.99609,
-        0.99609,
-        0.98047,
-        0.98828,
-        0.96484,
-        0.95703,
-        0.93359,
-        0.89062,
-    ]
-)
-
-
-def nearest_interp(src_array: torch.Tensor, target_length: int) -> torch.Tensor:
-    """
-    Interpolate the source array to the target length using nearest neighbor interpolation.
-    """
-    src_length = len(src_array)
-    if target_length == 1:
-        return src_array[-1:]
-
-    scale = (src_length - 1) / (target_length - 1)
-    grid = torch.arange(target_length, device=src_array.device, dtype=torch.float32)
-    mapped_indices = torch.round(grid * scale).long()
-    return src_array[mapped_indices]
-
-
-@dataclass
-class MagCacheConfig:
-    r"""
-    Configuration for [MagCache](https://github.com/Zehong-Ma/MagCache).
-
-    Args:
-        threshold (`float`, defaults to `0.06`):
-            The threshold for the accumulated error. If the accumulated error is below this threshold, the block
-            computation is skipped. A higher threshold allows for more aggressive skipping (faster) but may degrade
-            quality.
-        max_skip_steps (`int`, defaults to `3`):
-            The maximum number of consecutive steps that can be skipped (K in the paper).
-        retention_ratio (`float`, defaults to `0.2`):
-            The fraction of initial steps during which skipping is disabled to ensure stability. For example, if
-            `num_inference_steps` is 28 and `retention_ratio` is 0.2, the first 6 steps will never be skipped.
-        num_inference_steps (`int`, defaults to `28`):
-            The number of inference steps used in the pipeline. This is required to interpolate `mag_ratios` correctly.
-        mag_ratios (`torch.Tensor`, *optional*):
-            The pre-computed magnitude ratios for the model. These are checkpoint-dependent. If not provided, you must
-            set `calibrate=True` to calculate them for your specific model. For Flux models, you can use
-            `diffusers.hooks.mag_cache.FLUX_MAG_RATIOS`.
-        calibrate (`bool`, defaults to `False`):
-            If True, enables calibration mode. In this mode, no blocks are skipped. Instead, the hook calculates the
-            magnitude ratios for the current run and logs them at the end. Use this to obtain `mag_ratios` for new
-            models or schedulers.
-    """
-
-    threshold: float = 0.06
-    max_skip_steps: int = 3
-    retention_ratio: float = 0.2
-    num_inference_steps: int = 28
-    mag_ratios: Optional[Union[torch.Tensor, List[float]]] = None
-    calibrate: bool = False
-
-    def __post_init__(self):
-        # User MUST provide ratios OR enable calibration.
-        if self.mag_ratios is None and not self.calibrate:
-            raise ValueError(
-                " `mag_ratios` must be provided for MagCache inference because these ratios are model-dependent.\n"
-                "To get them for your model:\n"
-                "1. Initialize `MagCacheConfig(calibrate=True, ...)`\n"
-                "2. Run inference on your model once.\n"
-                "3. Copy the printed ratios array and pass it to `mag_ratios` in the config.\n"
-                "For Flux models, you can import `FLUX_MAG_RATIOS` from `diffusers.hooks.mag_cache`."
-            )
-
-        if not self.calibrate and self.mag_ratios is not None:
-            if not torch.is_tensor(self.mag_ratios):
-                self.mag_ratios = torch.tensor(self.mag_ratios)
-
-            if len(self.mag_ratios) != self.num_inference_steps:
-                logger.debug(
-                    f"Interpolating mag_ratios from length {len(self.mag_ratios)} to {self.num_inference_steps}"
-                )
-                self.mag_ratios = nearest_interp(self.mag_ratios, self.num_inference_steps)
-
-
-class MagCacheState(BaseState):
-    def __init__(self) -> None:
-        super().__init__()
-        # Cache for the residual (output - input) from the *previous* timestep
-        self.previous_residual: torch.Tensor = None
-
-        # State inputs/outputs for the current forward pass
-        self.head_block_input: Union[torch.Tensor, Tuple[torch.Tensor, ...]] = None
-        self.should_compute: bool = True
-
-        # MagCache accumulators
-        self.accumulated_ratio: float = 1.0
-        self.accumulated_err: float = 0.0
-        self.accumulated_steps: int = 0
-
-        # Current step counter (timestep index)
-        self.step_index: int = 0
-
-        # Calibration storage
-        self.calibration_ratios: List[float] = []
-
-    def reset(self):
-        self.previous_residual = None
-        self.should_compute = True
-        self.accumulated_ratio = 1.0
-        self.accumulated_err = 0.0
-        self.accumulated_steps = 0
-        self.step_index = 0
-        self.calibration_ratios = []
-
-
-class MagCacheHeadHook(ModelHook):
-    _is_stateful = True
-
-    def __init__(self, state_manager: StateManager, config: MagCacheConfig):
-        self.state_manager = state_manager
-        self.config = config
-        self._metadata = None
-
-    def initialize_hook(self, module):
-        unwrapped_module = unwrap_module(module)
-        self._metadata = TransformerBlockRegistry.get(unwrapped_module.__class__)
-        return module
-
-    @torch.compiler.disable
-    def new_forward(self, module: torch.nn.Module, *args, **kwargs):
-        if self.state_manager._current_context is None:
-            self.state_manager.set_context("inference")
-
-        arg_name = self._metadata.hidden_states_argument_name
-        hidden_states = self._metadata._get_parameter_from_args_kwargs(arg_name, args, kwargs)
-
-        state: MagCacheState = self.state_manager.get_state()
-        state.head_block_input = hidden_states
-
-        should_compute = True
-
-        if self.config.calibrate:
-            # Never skip during calibration
-            should_compute = True
-        else:
-            # MagCache Logic
-            current_step = state.step_index
-            if current_step >= len(self.config.mag_ratios):
-                current_scale = 1.0
-            else:
-                current_scale = self.config.mag_ratios[current_step]
-
-            retention_step = int(self.config.retention_ratio * self.config.num_inference_steps + 0.5)
-
-            if current_step >= retention_step:
-                state.accumulated_ratio *= current_scale
-                state.accumulated_steps += 1
-                state.accumulated_err += abs(1.0 - state.accumulated_ratio)
-
-                if (
-                    state.previous_residual is not None
-                    and state.accumulated_err <= self.config.threshold
-                    and state.accumulated_steps <= self.config.max_skip_steps
-                ):
-                    should_compute = False
-                else:
-                    state.accumulated_ratio = 1.0
-                    state.accumulated_steps = 0
-                    state.accumulated_err = 0.0
-
-        state.should_compute = should_compute
-
-        if not should_compute:
-            logger.debug(f"MagCache: Skipping step {state.step_index}")
-            # Apply MagCache: Output = Input + Previous Residual
-
-            output = hidden_states
-            res = state.previous_residual
-
-            if res.device != output.device:
-                res = res.to(output.device)
-
-            # Attempt to apply residual handling shape mismatches (e.g., text+image vs image only)
-            if res.shape == output.shape:
-                output = output + res
-            elif (
-                output.ndim == 3
-                and res.ndim == 3
-                and output.shape[0] == res.shape[0]
-                and output.shape[2] == res.shape[2]
-            ):
-                # Assuming concatenation where image part is at the end (standard in Flux/SD3)
-                diff = output.shape[1] - res.shape[1]
-                if diff > 0:
-                    output = output.clone()
-                    output[:, diff:, :] = output[:, diff:, :] + res
-                else:
-                    logger.warning(
-                        f"MagCache: Dimension mismatch. Input {output.shape}, Residual {res.shape}. "
-                        "Cannot apply residual safely. Returning input without residual."
-                    )
-            else:
-                logger.warning(
-                    f"MagCache: Dimension mismatch. Input {output.shape}, Residual {res.shape}. "
-                    "Cannot apply residual safely. Returning input without residual."
-                )
-
-            if self._metadata.return_encoder_hidden_states_index is not None:
-                original_encoder_hidden_states = self._metadata._get_parameter_from_args_kwargs(
-                    "encoder_hidden_states", args, kwargs
-                )
-                max_idx = max(
-                    self._metadata.return_hidden_states_index, self._metadata.return_encoder_hidden_states_index
-                )
-                ret_list = [None] * (max_idx + 1)
-                ret_list[self._metadata.return_hidden_states_index] = output
-                ret_list[self._metadata.return_encoder_hidden_states_index] = original_encoder_hidden_states
-                return tuple(ret_list)
-            else:
-                return output
-
-        else:
-            # Compute original forward
-            output = self.fn_ref.original_forward(*args, **kwargs)
-            return output
-
-    def reset_state(self, module):
-        self.state_manager.reset()
-        return module
-
-
-class MagCacheBlockHook(ModelHook):
-    def __init__(self, state_manager: StateManager, is_tail: bool = False, config: MagCacheConfig = None):
-        super().__init__()
-        self.state_manager = state_manager
-        self.is_tail = is_tail
-        self.config = config
-        self._metadata = None
-
-    def initialize_hook(self, module):
-        unwrapped_module = unwrap_module(module)
-        self._metadata = TransformerBlockRegistry.get(unwrapped_module.__class__)
-        return module
-
-    @torch.compiler.disable
-    def new_forward(self, module: torch.nn.Module, *args, **kwargs):
-        if self.state_manager._current_context is None:
-            self.state_manager.set_context("inference")
-        state: MagCacheState = self.state_manager.get_state()
-
-        if not state.should_compute:
-            arg_name = self._metadata.hidden_states_argument_name
-            hidden_states = self._metadata._get_parameter_from_args_kwargs(arg_name, args, kwargs)
-
-            if self.is_tail:
-                # Still need to advance step index even if we skip
-                self._advance_step(state)
-
-            if self._metadata.return_encoder_hidden_states_index is not None:
-                encoder_hidden_states = self._metadata._get_parameter_from_args_kwargs(
-                    "encoder_hidden_states", args, kwargs
-                )
-                max_idx = max(
-                    self._metadata.return_hidden_states_index, self._metadata.return_encoder_hidden_states_index
-                )
-                ret_list = [None] * (max_idx + 1)
-                ret_list[self._metadata.return_hidden_states_index] = hidden_states
-                ret_list[self._metadata.return_encoder_hidden_states_index] = encoder_hidden_states
-                return tuple(ret_list)
-
-            return hidden_states
-
-        output = self.fn_ref.original_forward(*args, **kwargs)
-
-        if self.is_tail:
-            # Calculate residual for next steps
-            if isinstance(output, tuple):
-                out_hidden = output[self._metadata.return_hidden_states_index]
-            else:
-                out_hidden = output
-
-            in_hidden = state.head_block_input
-
-            if in_hidden is None:
-                return output
-
-            # Determine residual
-            if out_hidden.shape == in_hidden.shape:
-                residual = out_hidden - in_hidden
-            elif out_hidden.ndim == 3 and in_hidden.ndim == 3 and out_hidden.shape[2] == in_hidden.shape[2]:
-                diff = in_hidden.shape[1] - out_hidden.shape[1]
-                if diff == 0:
-                    residual = out_hidden - in_hidden
-                else:
-                    residual = out_hidden - in_hidden  # Fallback to matching tail
-            else:
-                # Fallback for completely mismatched shapes
-                residual = out_hidden
-
-            if self.config.calibrate:
-                self._perform_calibration_step(state, residual)
-
-            state.previous_residual = residual
-            self._advance_step(state)
-
-        return output
-
-    def _perform_calibration_step(self, state: MagCacheState, current_residual: torch.Tensor):
-        if state.previous_residual is None:
-            # First step has no previous residual to compare against.
-            # log 1.0 as a neutral starting point.
-            ratio = 1.0
-        else:
-            # MagCache Calibration Formula: mean(norm(curr) / norm(prev))
-            # norm(dim=-1) gives magnitude of each token vector
-            curr_norm = torch.linalg.norm(current_residual.float(), dim=-1)
-            prev_norm = torch.linalg.norm(state.previous_residual.float(), dim=-1)
-
-            # Avoid division by zero
-            ratio = (curr_norm / (prev_norm + 1e-8)).mean().item()
-
-        state.calibration_ratios.append(ratio)
-
-    def _advance_step(self, state: MagCacheState):
-        state.step_index += 1
-        if state.step_index >= self.config.num_inference_steps:
-            # End of inference loop
-            if self.config.calibrate:
-                print("\n[MagCache] Calibration Complete. Copy these values to MagCacheConfig(mag_ratios=...):")
-                print(f"{state.calibration_ratios}\n")
-                logger.info(f"MagCache Calibration Results: {state.calibration_ratios}")
-
-            # Reset state
-            state.step_index = 0
-            state.accumulated_ratio = 1.0
-            state.accumulated_steps = 0
-            state.accumulated_err = 0.0
-            state.previous_residual = None
-            state.calibration_ratios = []
-
-
-def apply_mag_cache(module: torch.nn.Module, config: MagCacheConfig) -> None:
-    """
-    Applies MagCache to a given module (typically a Transformer).
-
-    Args:
-        module (`torch.nn.Module`):
-            The module to apply MagCache to.
-        config (`MagCacheConfig`):
-            The configuration for MagCache.
-    """
-    # Initialize registry on the root module so the Pipeline can set context.
-    HookRegistry.check_if_exists_or_initialize(module)
-
-    state_manager = StateManager(MagCacheState, (), {})
-    remaining_blocks = []
-
-    for name, submodule in module.named_children():
-        if name not in _ALL_TRANSFORMER_BLOCK_IDENTIFIERS or not isinstance(submodule, torch.nn.ModuleList):
-            continue
-        for index, block in enumerate(submodule):
-            remaining_blocks.append((f"{name}.{index}", block))
-
-    if not remaining_blocks:
-        logger.warning("MagCache: No transformer blocks found to apply hooks.")
-        return
-
-    # Handle single-block models
-    if len(remaining_blocks) == 1:
-        name, block = remaining_blocks[0]
-        logger.info(f"MagCache: Applying Head+Tail Hooks to single block '{name}'")
-        _apply_mag_cache_block_hook(block, state_manager, config, is_tail=True)
-        _apply_mag_cache_head_hook(block, state_manager, config)
-        return
-
-    head_block_name, head_block = remaining_blocks.pop(0)
-    tail_block_name, tail_block = remaining_blocks.pop(-1)
-
-    logger.info(f"MagCache: Applying Head Hook to {head_block_name}")
-    _apply_mag_cache_head_hook(head_block, state_manager, config)
-
-    for name, block in remaining_blocks:
-        _apply_mag_cache_block_hook(block, state_manager, config)
-
-    logger.info(f"MagCache: Applying Tail Hook to {tail_block_name}")
-    _apply_mag_cache_block_hook(tail_block, state_manager, config, is_tail=True)
-
-
-def _apply_mag_cache_head_hook(block: torch.nn.Module, state_manager: StateManager, config: MagCacheConfig) -> None:
-    registry = HookRegistry.check_if_exists_or_initialize(block)
-
-    # Automatically remove existing hook to allow re-application (e.g. switching modes)
-    if registry.get_hook(_MAG_CACHE_LEADER_BLOCK_HOOK) is not None:
-        registry.remove_hook(_MAG_CACHE_LEADER_BLOCK_HOOK)
-
-    hook = MagCacheHeadHook(state_manager, config)
-    registry.register_hook(hook, _MAG_CACHE_LEADER_BLOCK_HOOK)
-
-
-def _apply_mag_cache_block_hook(
-    block: torch.nn.Module,
-    state_manager: StateManager,
-    config: MagCacheConfig,
-    is_tail: bool = False,
-) -> None:
-    registry = HookRegistry.check_if_exists_or_initialize(block)
-
-    # Automatically remove existing hook to allow re-application
-    if registry.get_hook(_MAG_CACHE_BLOCK_HOOK) is not None:
-        registry.remove_hook(_MAG_CACHE_BLOCK_HOOK)
-
-    hook = MagCacheBlockHook(state_manager, is_tail, config)
-    registry.register_hook(hook, _MAG_CACHE_BLOCK_HOOK)

diffusers/hooks/pyramid_attention_broadcast.py

@@ -1,314 +0,0 @@
-# Copyright 2026 The HuggingFace Team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-import re
-from dataclasses import dataclass
-from typing import Any, Callable
-
-import torch
-
-from ..models.attention import AttentionModuleMixin
-from ..models.attention_processor import Attention, MochiAttention
-from ..utils import logging
-from ._common import (
-    _ATTENTION_CLASSES,
-    _CROSS_TRANSFORMER_BLOCK_IDENTIFIERS,
-    _SPATIAL_TRANSFORMER_BLOCK_IDENTIFIERS,
-    _TEMPORAL_TRANSFORMER_BLOCK_IDENTIFIERS,
-)
-from .hooks import HookRegistry, ModelHook
-
-
-logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
-
-
-_PYRAMID_ATTENTION_BROADCAST_HOOK = "pyramid_attention_broadcast"
-
-
-@dataclass
-class PyramidAttentionBroadcastConfig:
-    r"""
-    Configuration for Pyramid Attention Broadcast.
-
-    Args:
-        spatial_attention_block_skip_range (`int`, *optional*, defaults to `None`):
-            The number of times a specific spatial attention broadcast is skipped before computing the attention states
-            to re-use. If this is set to the value `N`, the attention computation will be skipped `N - 1` times (i.e.,
-            old attention states will be reused) before computing the new attention states again.
-        temporal_attention_block_skip_range (`int`, *optional*, defaults to `None`):
-            The number of times a specific temporal attention broadcast is skipped before computing the attention
-            states to re-use. If this is set to the value `N`, the attention computation will be skipped `N - 1` times
-            (i.e., old attention states will be reused) before computing the new attention states again.
-        cross_attention_block_skip_range (`int`, *optional*, defaults to `None`):
-            The number of times a specific cross-attention broadcast is skipped before computing the attention states
-            to re-use. If this is set to the value `N`, the attention computation will be skipped `N - 1` times (i.e.,
-            old attention states will be reused) before computing the new attention states again.
-        spatial_attention_timestep_skip_range (`tuple[int, int]`, defaults to `(100, 800)`):
-            The range of timesteps to skip in the spatial attention layer. The attention computations will be
-            conditionally skipped if the current timestep is within the specified range.
-        temporal_attention_timestep_skip_range (`tuple[int, int]`, defaults to `(100, 800)`):
-            The range of timesteps to skip in the temporal attention layer. The attention computations will be
-            conditionally skipped if the current timestep is within the specified range.
-        cross_attention_timestep_skip_range (`tuple[int, int]`, defaults to `(100, 800)`):
-            The range of timesteps to skip in the cross-attention layer. The attention computations will be
-            conditionally skipped if the current timestep is within the specified range.
-        spatial_attention_block_identifiers (`tuple[str, ...]`):
-            The identifiers to match against the layer names to determine if the layer is a spatial attention layer.
-        temporal_attention_block_identifiers (`tuple[str, ...]`):
-            The identifiers to match against the layer names to determine if the layer is a temporal attention layer.
-        cross_attention_block_identifiers (`tuple[str, ...]`):
-            The identifiers to match against the layer names to determine if the layer is a cross-attention layer.
-    """
-
-    spatial_attention_block_skip_range: int | None = None
-    temporal_attention_block_skip_range: int | None = None
-    cross_attention_block_skip_range: int | None = None
-
-    spatial_attention_timestep_skip_range: tuple[int, int] = (100, 800)
-    temporal_attention_timestep_skip_range: tuple[int, int] = (100, 800)
-    cross_attention_timestep_skip_range: tuple[int, int] = (100, 800)
-
-    spatial_attention_block_identifiers: tuple[str, ...] = _SPATIAL_TRANSFORMER_BLOCK_IDENTIFIERS
-    temporal_attention_block_identifiers: tuple[str, ...] = _TEMPORAL_TRANSFORMER_BLOCK_IDENTIFIERS
-    cross_attention_block_identifiers: tuple[str, ...] = _CROSS_TRANSFORMER_BLOCK_IDENTIFIERS
-
-    current_timestep_callback: Callable[[], int] = None
-
-    # TODO(aryan): add PAB for MLP layers (very limited speedup from testing with original codebase
-    # so not added for now)
-
-    def __repr__(self) -> str:
-        return (
-            f"PyramidAttentionBroadcastConfig(\n"
-            f"  spatial_attention_block_skip_range={self.spatial_attention_block_skip_range},\n"
-            f"  temporal_attention_block_skip_range={self.temporal_attention_block_skip_range},\n"
-            f"  cross_attention_block_skip_range={self.cross_attention_block_skip_range},\n"
-            f"  spatial_attention_timestep_skip_range={self.spatial_attention_timestep_skip_range},\n"
-            f"  temporal_attention_timestep_skip_range={self.temporal_attention_timestep_skip_range},\n"
-            f"  cross_attention_timestep_skip_range={self.cross_attention_timestep_skip_range},\n"
-            f"  spatial_attention_block_identifiers={self.spatial_attention_block_identifiers},\n"
-            f"  temporal_attention_block_identifiers={self.temporal_attention_block_identifiers},\n"
-            f"  cross_attention_block_identifiers={self.cross_attention_block_identifiers},\n"
-            f"  current_timestep_callback={self.current_timestep_callback}\n"
-            ")"
-        )
-
-
-class PyramidAttentionBroadcastState:
-    r"""
-    State for Pyramid Attention Broadcast.
-
-    Attributes:
-        iteration (`int`):
-            The current iteration of the Pyramid Attention Broadcast. It is necessary to ensure that `reset_state` is
-            called before starting a new inference forward pass for PAB to work correctly.
-        cache (`Any`):
-            The cached output from the previous forward pass. This is used to re-use the attention states when the
-            attention computation is skipped. It is either a tensor or a tuple of tensors, depending on the module.
-    """
-
-    def __init__(self) -> None:
-        self.iteration = 0
-        self.cache = None
-
-    def reset(self):
-        self.iteration = 0
-        self.cache = None
-
-    def __repr__(self):
-        cache_repr = ""
-        if self.cache is None:
-            cache_repr = "None"
-        else:
-            cache_repr = f"Tensor(shape={self.cache.shape}, dtype={self.cache.dtype})"
-        return f"PyramidAttentionBroadcastState(iteration={self.iteration}, cache={cache_repr})"
-
-
-class PyramidAttentionBroadcastHook(ModelHook):
-    r"""A hook that applies Pyramid Attention Broadcast to a given module."""
-
-    _is_stateful = True
-
-    def __init__(
-        self, timestep_skip_range: tuple[int, int], block_skip_range: int, current_timestep_callback: Callable[[], int]
-    ) -> None:
-        super().__init__()
-
-        self.timestep_skip_range = timestep_skip_range
-        self.block_skip_range = block_skip_range
-        self.current_timestep_callback = current_timestep_callback
-
-    def initialize_hook(self, module):
-        self.state = PyramidAttentionBroadcastState()
-        return module
-
-    def new_forward(self, module: torch.nn.Module, *args, **kwargs) -> Any:
-        is_within_timestep_range = (
-            self.timestep_skip_range[0] < self.current_timestep_callback() < self.timestep_skip_range[1]
-        )
-        should_compute_attention = (
-            self.state.cache is None
-            or self.state.iteration == 0
-            or not is_within_timestep_range
-            or self.state.iteration % self.block_skip_range == 0
-        )
-
-        if should_compute_attention:
-            output = self.fn_ref.original_forward(*args, **kwargs)
-        else:
-            output = self.state.cache
-
-        self.state.cache = output
-        self.state.iteration += 1
-        return output
-
-    def reset_state(self, module: torch.nn.Module) -> None:
-        self.state.reset()
-        return module
-
-
-def apply_pyramid_attention_broadcast(module: torch.nn.Module, config: PyramidAttentionBroadcastConfig):
-    r"""
-    Apply [Pyramid Attention Broadcast](https://huggingface.co/papers/2408.12588) to a given pipeline.
-
-    PAB is an attention approximation method that leverages the similarity in attention states between timesteps to
-    reduce the computational cost of attention computation. The key takeaway from the paper is that the attention
-    similarity in the cross-attention layers between timesteps is high, followed by less similarity in the temporal and
-    spatial layers. This allows for the skipping of attention computation in the cross-attention layers more frequently
-    than in the temporal and spatial layers. Applying PAB will, therefore, speedup the inference process.
-
-    Args:
-        module (`torch.nn.Module`):
-            The module to apply Pyramid Attention Broadcast to.
-        config (`PyramidAttentionBroadcastConfig | None`, `optional`, defaults to `None`):
-            The configuration to use for Pyramid Attention Broadcast.
-
-    Example:
-
-    ```python
-    >>> import torch
-    >>> from diffusers import CogVideoXPipeline, PyramidAttentionBroadcastConfig, apply_pyramid_attention_broadcast
-    >>> from diffusers.utils import export_to_video
-
-    >>> pipe = CogVideoXPipeline.from_pretrained("THUDM/CogVideoX-5b", torch_dtype=torch.bfloat16)
-    >>> pipe.to("cuda")
-
-    >>> config = PyramidAttentionBroadcastConfig(
-    ...     spatial_attention_block_skip_range=2,
-    ...     spatial_attention_timestep_skip_range=(100, 800),
-    ...     current_timestep_callback=lambda: pipe.current_timestep,
-    ... )
-    >>> apply_pyramid_attention_broadcast(pipe.transformer, config)
-    ```
-    """
-    if config.current_timestep_callback is None:
-        raise ValueError(
-            "The `current_timestep_callback` function must be provided in the configuration to apply Pyramid Attention Broadcast."
-        )
-
-    if (
-        config.spatial_attention_block_skip_range is None
-        and config.temporal_attention_block_skip_range is None
-        and config.cross_attention_block_skip_range is None
-    ):
-        logger.warning(
-            "Pyramid Attention Broadcast requires one or more of `spatial_attention_block_skip_range`, `temporal_attention_block_skip_range` "
-            "or `cross_attention_block_skip_range` parameters to be set to an integer, not `None`. Defaulting to using `spatial_attention_block_skip_range=2`. "
-            "To avoid this warning, please set one of the above parameters."
-        )
-        config.spatial_attention_block_skip_range = 2
-
-    for name, submodule in module.named_modules():
-        if not isinstance(submodule, (*_ATTENTION_CLASSES, AttentionModuleMixin)):
-            # PAB has been implemented specific to Diffusers' Attention classes. However, this does not mean that PAB
-            # cannot be applied to this layer. For custom layers, users can extend this functionality and implement
-            # their own PAB logic similar to `_apply_pyramid_attention_broadcast_on_attention_class`.
-            continue
-        _apply_pyramid_attention_broadcast_on_attention_class(name, submodule, config)
-
-
-def _apply_pyramid_attention_broadcast_on_attention_class(
-    name: str, module: Attention, config: PyramidAttentionBroadcastConfig
-) -> bool:
-    is_spatial_self_attention = (
-        any(re.search(identifier, name) is not None for identifier in config.spatial_attention_block_identifiers)
-        and config.spatial_attention_block_skip_range is not None
-        and not getattr(module, "is_cross_attention", False)
-    )
-    is_temporal_self_attention = (
-        any(re.search(identifier, name) is not None for identifier in config.temporal_attention_block_identifiers)
-        and config.temporal_attention_block_skip_range is not None
-        and not getattr(module, "is_cross_attention", False)
-    )
-    is_cross_attention = (
-        any(re.search(identifier, name) is not None for identifier in config.cross_attention_block_identifiers)
-        and config.cross_attention_block_skip_range is not None
-        and getattr(module, "is_cross_attention", False)
-    )
-
-    block_skip_range, timestep_skip_range, block_type = None, None, None
-    if is_spatial_self_attention:
-        block_skip_range = config.spatial_attention_block_skip_range
-        timestep_skip_range = config.spatial_attention_timestep_skip_range
-        block_type = "spatial"
-    elif is_temporal_self_attention:
-        block_skip_range = config.temporal_attention_block_skip_range
-        timestep_skip_range = config.temporal_attention_timestep_skip_range
-        block_type = "temporal"
-    elif is_cross_attention:
-        block_skip_range = config.cross_attention_block_skip_range
-        timestep_skip_range = config.cross_attention_timestep_skip_range
-        block_type = "cross"
-
-    if block_skip_range is None or timestep_skip_range is None:
-        logger.info(
-            f'Unable to apply Pyramid Attention Broadcast to the selected layer: "{name}" because it does '
-            f"not match any of the required criteria for spatial, temporal or cross attention layers. Note, "
-            f"however, that this layer may still be valid for applying PAB. Please specify the correct "
-            f"block identifiers in the configuration."
-        )
-        return False
-
-    logger.debug(f"Enabling Pyramid Attention Broadcast ({block_type}) in layer: {name}")
-    _apply_pyramid_attention_broadcast_hook(
-        module, timestep_skip_range, block_skip_range, config.current_timestep_callback
-    )
-    return True
-
-
-def _apply_pyramid_attention_broadcast_hook(
-    module: Attention | MochiAttention,
-    timestep_skip_range: tuple[int, int],
-    block_skip_range: int,
-    current_timestep_callback: Callable[[], int],
-):
-    r"""
-    Apply [Pyramid Attention Broadcast](https://huggingface.co/papers/2408.12588) to a given torch.nn.Module.
-
-    Args:
-        module (`torch.nn.Module`):
-            The module to apply Pyramid Attention Broadcast to.
-        timestep_skip_range (`tuple[int, int]`):
-            The range of timesteps to skip in the attention layer. The attention computations will be conditionally
-            skipped if the current timestep is within the specified range.
-        block_skip_range (`int`):
-            The number of times a specific attention broadcast is skipped before computing the attention states to
-            re-use. If this is set to the value `N`, the attention computation will be skipped `N - 1` times (i.e., old
-            attention states will be reused) before computing the new attention states again.
-        current_timestep_callback (`Callable[[], int]`):
-            A callback function that returns the current inference timestep.
-    """
-    registry = HookRegistry.check_if_exists_or_initialize(module)
-    hook = PyramidAttentionBroadcastHook(timestep_skip_range, block_skip_range, current_timestep_callback)
-    registry.register_hook(hook, _PYRAMID_ATTENTION_BROADCAST_HOOK)

diffusers/hooks/smoothed_energy_guidance_utils.py

@@ -1,166 +0,0 @@
-# Copyright 2026 The HuggingFace Team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-import math
-from dataclasses import asdict, dataclass
-
-import torch
-import torch.nn.functional as F
-
-from ..utils import get_logger
-from ._common import _ALL_TRANSFORMER_BLOCK_IDENTIFIERS, _ATTENTION_CLASSES, _get_submodule_from_fqn
-from .hooks import HookRegistry, ModelHook
-
-
-logger = get_logger(__name__)  # pylint: disable=invalid-name
-
-_SMOOTHED_ENERGY_GUIDANCE_HOOK = "smoothed_energy_guidance_hook"
-
-
-@dataclass
-class SmoothedEnergyGuidanceConfig:
-    r"""
-    Configuration for skipping internal transformer blocks when executing a transformer model.
-
-    Args:
-        indices (`list[int]`):
-            The indices of the layer to skip. This is typically the first layer in the transformer block.
-        fqn (`str`, defaults to `"auto"`):
-            The fully qualified name identifying the stack of transformer blocks. Typically, this is
-            `transformer_blocks`, `single_transformer_blocks`, `blocks`, `layers`, or `temporal_transformer_blocks`.
-            For automatic detection, set this to `"auto"`. "auto" only works on DiT models. For UNet models, you must
-            provide the correct fqn.
-        _query_proj_identifiers (`list[str]`, defaults to `None`):
-            The identifiers for the query projection layers. Typically, these are `to_q`, `query`, or `q_proj`. If
-            `None`, `to_q` is used by default.
-    """
-
-    indices: list[int]
-    fqn: str = "auto"
-    _query_proj_identifiers: list[str] = None
-
-    def to_dict(self):
-        return asdict(self)
-
-    @staticmethod
-    def from_dict(data: dict) -> "SmoothedEnergyGuidanceConfig":
-        return SmoothedEnergyGuidanceConfig(**data)
-
-
-class SmoothedEnergyGuidanceHook(ModelHook):
-    def __init__(self, blur_sigma: float = 1.0, blur_threshold_inf: float = 9999.9) -> None:
-        super().__init__()
-        self.blur_sigma = blur_sigma
-        self.blur_threshold_inf = blur_threshold_inf
-
-    def post_forward(self, module: torch.nn.Module, output: torch.Tensor) -> torch.Tensor:
-        # Copied from https://github.com/SusungHong/SEG-SDXL/blob/cf8256d640d5373541cfea3b3b6caf93272cf986/pipeline_seg.py#L172C31-L172C102
-        kernel_size = math.ceil(6 * self.blur_sigma) + 1 - math.ceil(6 * self.blur_sigma) % 2
-        smoothed_output = _gaussian_blur_2d(output, kernel_size, self.blur_sigma, self.blur_threshold_inf)
-        return smoothed_output
-
-
-def _apply_smoothed_energy_guidance_hook(
-    module: torch.nn.Module, config: SmoothedEnergyGuidanceConfig, blur_sigma: float, name: str | None = None
-) -> None:
-    name = name or _SMOOTHED_ENERGY_GUIDANCE_HOOK
-
-    if config.fqn == "auto":
-        for identifier in _ALL_TRANSFORMER_BLOCK_IDENTIFIERS:
-            if hasattr(module, identifier):
-                config.fqn = identifier
-                break
-        else:
-            raise ValueError(
-                "Could not find a suitable identifier for the transformer blocks automatically. Please provide a valid "
-                "`fqn` (fully qualified name) that identifies a stack of transformer blocks."
-            )
-
-    if config._query_proj_identifiers is None:
-        config._query_proj_identifiers = ["to_q"]
-
-    transformer_blocks = _get_submodule_from_fqn(module, config.fqn)
-    blocks_found = False
-    for i, block in enumerate(transformer_blocks):
-        if i not in config.indices:
-            continue
-
-        blocks_found = True
-
-        for submodule_name, submodule in block.named_modules():
-            if not isinstance(submodule, _ATTENTION_CLASSES) or submodule.is_cross_attention:
-                continue
-            for identifier in config._query_proj_identifiers:
-                query_proj = getattr(submodule, identifier, None)
-                if query_proj is None or not isinstance(query_proj, torch.nn.Linear):
-                    continue
-                logger.debug(
-                    f"Registering smoothed energy guidance hook on {config.fqn}.{i}.{submodule_name}.{identifier}"
-                )
-                registry = HookRegistry.check_if_exists_or_initialize(query_proj)
-                hook = SmoothedEnergyGuidanceHook(blur_sigma)
-                registry.register_hook(hook, name)
-
-    if not blocks_found:
-        raise ValueError(
-            f"Could not find any transformer blocks matching the provided indices {config.indices} and "
-            f"fully qualified name '{config.fqn}'. Please check the indices and fqn for correctness."
-        )
-
-
-# Modified from https://github.com/SusungHong/SEG-SDXL/blob/cf8256d640d5373541cfea3b3b6caf93272cf986/pipeline_seg.py#L71
-def _gaussian_blur_2d(query: torch.Tensor, kernel_size: int, sigma: float, sigma_threshold_inf: float) -> torch.Tensor:
-    """
-    This implementation assumes that the input query is for visual (image/videos) tokens to apply the 2D gaussian blur.
-    However, some models use joint text-visual token attention for which this may not be suitable. Additionally, this
-    implementation also assumes that the visual tokens come from a square image/video. In practice, despite these
-    assumptions, applying the 2D square gaussian blur on the query projections generates reasonable results for
-    Smoothed Energy Guidance.
-
-    SEG is only supported as an experimental prototype feature for now, so the implementation may be modified in the
-    future without warning or guarantee of reproducibility.
-    """
-    assert query.ndim == 3
-
-    is_inf = sigma > sigma_threshold_inf
-    batch_size, seq_len, embed_dim = query.shape
-
-    seq_len_sqrt = int(math.sqrt(seq_len))
-    num_square_tokens = seq_len_sqrt * seq_len_sqrt
-    query_slice = query[:, :num_square_tokens, :]
-    query_slice = query_slice.permute(0, 2, 1)
-    query_slice = query_slice.reshape(batch_size, embed_dim, seq_len_sqrt, seq_len_sqrt)
-
-    if is_inf:
-        kernel_size = min(kernel_size, seq_len_sqrt - (seq_len_sqrt % 2 - 1))
-        kernel_size_half = (kernel_size - 1) / 2
-
-        x = torch.linspace(-kernel_size_half, kernel_size_half, steps=kernel_size)
-        pdf = torch.exp(-0.5 * (x / sigma).pow(2))
-        kernel1d = pdf / pdf.sum()
-        kernel1d = kernel1d.to(query)
-        kernel2d = torch.matmul(kernel1d[:, None], kernel1d[None, :])
-        kernel2d = kernel2d.expand(embed_dim, 1, kernel2d.shape[0], kernel2d.shape[1])
-
-        padding = [kernel_size // 2, kernel_size // 2, kernel_size // 2, kernel_size // 2]
-        query_slice = F.pad(query_slice, padding, mode="reflect")
-        query_slice = F.conv2d(query_slice, kernel2d, groups=embed_dim)
-    else:
-        query_slice[:] = query_slice.mean(dim=(-2, -1), keepdim=True)
-
-    query_slice = query_slice.reshape(batch_size, embed_dim, num_square_tokens)
-    query_slice = query_slice.permute(0, 2, 1)
-    query[:, :num_square_tokens, :] = query_slice.clone()
-
-    return query

diffusers/hooks/taylorseer_cache.py

@@ -1,345 +0,0 @@
-import math
-import re
-from dataclasses import dataclass
-
-import torch
-import torch.nn as nn
-
-from ..utils import logging
-from .hooks import HookRegistry, ModelHook, StateManager
-
-
-logger = logging.get_logger(__name__)
-_TAYLORSEER_CACHE_HOOK = "taylorseer_cache"
-_SPATIAL_ATTENTION_BLOCK_IDENTIFIERS = (
-    "^blocks.*attn",
-    "^transformer_blocks.*attn",
-    "^single_transformer_blocks.*attn",
-)
-_TEMPORAL_ATTENTION_BLOCK_IDENTIFIERS = ("^temporal_transformer_blocks.*attn",)
-_TRANSFORMER_BLOCK_IDENTIFIERS = _SPATIAL_ATTENTION_BLOCK_IDENTIFIERS + _TEMPORAL_ATTENTION_BLOCK_IDENTIFIERS
-_BLOCK_IDENTIFIERS = ("^[^.]*block[^.]*\\.[^.]+$",)
-_PROJ_OUT_IDENTIFIERS = ("^proj_out$",)
-
-
-@dataclass
-class TaylorSeerCacheConfig:
-    """
-    Configuration for TaylorSeer cache. See: https://huggingface.co/papers/2503.06923
-
-    Attributes:
-        cache_interval (`int`, defaults to `5`):
-            The interval between full computation steps. After a full computation, the cached (predicted) outputs are
-            reused for this many subsequent denoising steps before refreshing with a new full forward pass.
-
-        disable_cache_before_step (`int`, defaults to `3`):
-            The denoising step index before which caching is disabled, meaning full computation is performed for the
-            initial steps (0 to disable_cache_before_step - 1) to gather data for Taylor series approximations. During
-            these steps, Taylor factors are updated, but caching/predictions are not applied. Caching begins at this
-            step.
-
-        disable_cache_after_step (`int`, *optional*, defaults to `None`):
-            The denoising step index after which caching is disabled. If set, for steps >= this value, all modules run
-            full computations without predictions or state updates, ensuring accuracy in later stages if needed.
-
-        max_order (`int`, defaults to `1`):
-            The highest order in the Taylor series expansion for approximating module outputs. Higher orders provide
-            better approximations but increase computation and memory usage.
-
-        taylor_factors_dtype (`torch.dtype`, defaults to `torch.bfloat16`):
-            Data type used for storing and computing Taylor series factors. Lower precision reduces memory but may
-            affect stability; higher precision improves accuracy at the cost of more memory.
-
-        skip_predict_identifiers (`list[str]`, *optional*, defaults to `None`):
-            Regex patterns (using `re.fullmatch`) for module names to place as "skip" in "cache" mode. In this mode,
-            the module computes fully during initial or refresh steps but returns a zero tensor (matching recorded
-            shape) during prediction steps to skip computation cheaply.
-
-        cache_identifiers (`list[str]`, *optional*, defaults to `None`):
-            Regex patterns (using `re.fullmatch`) for module names to place in Taylor-series caching mode, where
-            outputs are approximated and cached for reuse.
-
-        use_lite_mode (`bool`, *optional*, defaults to `False`):
-            Enables a lightweight TaylorSeer variant that minimizes memory usage by applying predefined patterns for
-            skipping and caching (e.g., skipping blocks and caching projections). This overrides any custom
-            `inactive_identifiers` or `active_identifiers`.
-
-    Notes:
-        - Patterns are matched using `re.fullmatch` on the module name.
-        - If `skip_predict_identifiers` or `cache_identifiers` are provided, only matching modules are hooked.
-        - If neither is provided, all attention-like modules are hooked by default.
-
-    Example of inactive and active usage:
-
-    ```py
-    def forward(x):
-        x = self.module1(x)  # inactive module: returns zeros tensor based on shape recorded during full compute
-        x = self.module2(x)  # active module: caches output here, avoiding recomputation of prior steps
-        return x
-    ```
-    """
-
-    cache_interval: int = 5
-    disable_cache_before_step: int = 3
-    disable_cache_after_step: int | None = None
-    max_order: int = 1
-    taylor_factors_dtype: torch.dtype | None = torch.bfloat16
-    skip_predict_identifiers: list[str] | None = None
-    cache_identifiers: list[str] | None = None
-    use_lite_mode: bool = False
-
-    def __repr__(self) -> str:
-        return (
-            "TaylorSeerCacheConfig("
-            f"cache_interval={self.cache_interval}, "
-            f"disable_cache_before_step={self.disable_cache_before_step}, "
-            f"disable_cache_after_step={self.disable_cache_after_step}, "
-            f"max_order={self.max_order}, "
-            f"taylor_factors_dtype={self.taylor_factors_dtype}, "
-            f"skip_predict_identifiers={self.skip_predict_identifiers}, "
-            f"cache_identifiers={self.cache_identifiers}, "
-            f"use_lite_mode={self.use_lite_mode})"
-        )
-
-
-class TaylorSeerState:
-    def __init__(
-        self,
-        taylor_factors_dtype: torch.dtype | None = torch.bfloat16,
-        max_order: int = 1,
-        is_inactive: bool = False,
-    ):
-        self.taylor_factors_dtype = taylor_factors_dtype
-        self.max_order = max_order
-        self.is_inactive = is_inactive
-
-        self.module_dtypes: tuple[torch.dtype, ...] = ()
-        self.last_update_step: int | None = None
-        self.taylor_factors: dict[int, dict[int, torch.Tensor]] = {}
-        self.inactive_shapes: tuple[tuple[int, ...], ...] | None = None
-        self.device: torch.device | None = None
-        self.current_step: int = -1
-
-    def reset(self) -> None:
-        self.current_step = -1
-        self.last_update_step = None
-        self.taylor_factors = {}
-        self.inactive_shapes = None
-        self.device = None
-
-    def update(
-        self,
-        outputs: tuple[torch.Tensor, ...],
-    ) -> None:
-        self.module_dtypes = tuple(output.dtype for output in outputs)
-        self.device = outputs[0].device
-
-        if self.is_inactive:
-            self.inactive_shapes = tuple(output.shape for output in outputs)
-        else:
-            for i, features in enumerate(outputs):
-                new_factors: dict[int, torch.Tensor] = {0: features}
-                is_first_update = self.last_update_step is None
-                if not is_first_update:
-                    delta_step = self.current_step - self.last_update_step
-                    if delta_step == 0:
-                        raise ValueError("Delta step cannot be zero for TaylorSeer update.")
-
-                    # Recursive divided differences up to max_order
-                    prev_factors = self.taylor_factors.get(i, {})
-                    for j in range(self.max_order):
-                        prev = prev_factors.get(j)
-                        if prev is None:
-                            break
-                        new_factors[j + 1] = (new_factors[j] - prev.to(features.dtype)) / delta_step
-                self.taylor_factors[i] = {
-                    order: factor.to(self.taylor_factors_dtype) for order, factor in new_factors.items()
-                }
-
-        self.last_update_step = self.current_step
-
-    @torch.compiler.disable
-    def predict(self) -> list[torch.Tensor]:
-        if self.last_update_step is None:
-            raise ValueError("Cannot predict without prior initialization/update.")
-
-        step_offset = self.current_step - self.last_update_step
-
-        outputs = []
-        if self.is_inactive:
-            if self.inactive_shapes is None:
-                raise ValueError("Inactive shapes not set during prediction.")
-            for i in range(len(self.module_dtypes)):
-                outputs.append(
-                    torch.zeros(
-                        self.inactive_shapes[i],
-                        dtype=self.module_dtypes[i],
-                        device=self.device,
-                    )
-                )
-        else:
-            if not self.taylor_factors:
-                raise ValueError("Taylor factors empty during prediction.")
-            num_outputs = len(self.taylor_factors)
-            num_orders = len(self.taylor_factors[0])
-            for i in range(num_outputs):
-                output_dtype = self.module_dtypes[i]
-                taylor_factors = self.taylor_factors[i]
-                output = torch.zeros_like(taylor_factors[0], dtype=output_dtype)
-                for order in range(num_orders):
-                    coeff = (step_offset**order) / math.factorial(order)
-                    factor = taylor_factors[order]
-                    output = output + factor.to(output_dtype) * coeff
-                outputs.append(output)
-        return outputs
-
-
-class TaylorSeerCacheHook(ModelHook):
-    _is_stateful = True
-
-    def __init__(
-        self,
-        cache_interval: int,
-        disable_cache_before_step: int,
-        taylor_factors_dtype: torch.dtype,
-        state_manager: StateManager,
-        disable_cache_after_step: int | None = None,
-    ):
-        super().__init__()
-        self.cache_interval = cache_interval
-        self.disable_cache_before_step = disable_cache_before_step
-        self.disable_cache_after_step = disable_cache_after_step
-        self.taylor_factors_dtype = taylor_factors_dtype
-        self.state_manager = state_manager
-
-    def initialize_hook(self, module: torch.nn.Module):
-        return module
-
-    def reset_state(self, module: torch.nn.Module) -> None:
-        """
-        Reset state between sampling runs.
-        """
-        self.state_manager.reset()
-
-    @torch.compiler.disable
-    def _measure_should_compute(self) -> bool:
-        state: TaylorSeerState = self.state_manager.get_state()
-        state.current_step += 1
-        current_step = state.current_step
-        is_warmup_phase = current_step < self.disable_cache_before_step
-        is_compute_interval = (current_step - self.disable_cache_before_step - 1) % self.cache_interval == 0
-        is_cooldown_phase = self.disable_cache_after_step is not None and current_step >= self.disable_cache_after_step
-        should_compute = is_warmup_phase or is_compute_interval or is_cooldown_phase
-        return should_compute, state
-
-    def new_forward(self, module: torch.nn.Module, *args, **kwargs):
-        should_compute, state = self._measure_should_compute()
-        if should_compute:
-            outputs = self.fn_ref.original_forward(*args, **kwargs)
-            wrapped_outputs = (outputs,) if isinstance(outputs, torch.Tensor) else outputs
-            state.update(wrapped_outputs)
-            return outputs
-
-        outputs_list = state.predict()
-        return outputs_list[0] if len(outputs_list) == 1 else tuple(outputs_list)
-
-
-def _resolve_patterns(config: TaylorSeerCacheConfig) -> tuple[list[str], list[str]]:
-    """
-    Resolve effective inactive and active pattern lists from config + templates.
-    """
-
-    inactive_patterns = config.skip_predict_identifiers if config.skip_predict_identifiers is not None else None
-    active_patterns = config.cache_identifiers if config.cache_identifiers is not None else None
-
-    return inactive_patterns or [], active_patterns or []
-
-
-def apply_taylorseer_cache(module: torch.nn.Module, config: TaylorSeerCacheConfig):
-    """
-    Applies the TaylorSeer cache to a given pipeline (typically the transformer / UNet).
-
-    This function hooks selected modules in the model to enable caching or skipping based on the provided
-    configuration, reducing redundant computations in diffusion denoising loops.
-
-    Args:
-        module (torch.nn.Module): The model subtree to apply the hooks to.
-        config (TaylorSeerCacheConfig): Configuration for the cache.
-
-    Example:
-    ```python
-    >>> import torch
-    >>> from diffusers import FluxPipeline, TaylorSeerCacheConfig
-
-    >>> pipe = FluxPipeline.from_pretrained(
-    ...     "black-forest-labs/FLUX.1-dev",
-    ...     torch_dtype=torch.bfloat16,
-    ... )
-    >>> pipe.to("cuda")
-
-    >>> config = TaylorSeerCacheConfig(
-    ...     cache_interval=5,
-    ...     max_order=1,
-    ...     disable_cache_before_step=3,
-    ...     taylor_factors_dtype=torch.float32,
-    ... )
-    >>> pipe.transformer.enable_cache(config)
-    ```
-    """
-    inactive_patterns, active_patterns = _resolve_patterns(config)
-
-    active_patterns = active_patterns or _TRANSFORMER_BLOCK_IDENTIFIERS
-
-    if config.use_lite_mode:
-        logger.info("Using TaylorSeer Lite variant for cache.")
-        active_patterns = _PROJ_OUT_IDENTIFIERS
-        inactive_patterns = _BLOCK_IDENTIFIERS
-        if config.skip_predict_identifiers or config.cache_identifiers:
-            logger.warning("Lite mode overrides user patterns.")
-
-    for name, submodule in module.named_modules():
-        matches_inactive = any(re.fullmatch(pattern, name) for pattern in inactive_patterns)
-        matches_active = any(re.fullmatch(pattern, name) for pattern in active_patterns)
-        if not (matches_inactive or matches_active):
-            continue
-        _apply_taylorseer_cache_hook(
-            module=submodule,
-            config=config,
-            is_inactive=matches_inactive,
-        )
-
-
-def _apply_taylorseer_cache_hook(
-    module: nn.Module,
-    config: TaylorSeerCacheConfig,
-    is_inactive: bool,
-):
-    """
-    Registers the TaylorSeer hook on the specified nn.Module.
-
-    Args:
-        name: Name of the module.
-        module: The nn.Module to be hooked.
-        config: Cache configuration.
-        is_inactive: Whether this module should operate in "inactive" mode.
-    """
-    state_manager = StateManager(
-        TaylorSeerState,
-        init_kwargs={
-            "taylor_factors_dtype": config.taylor_factors_dtype,
-            "max_order": config.max_order,
-            "is_inactive": is_inactive,
-        },
-    )
-
-    registry = HookRegistry.check_if_exists_or_initialize(module)
-
-    hook = TaylorSeerCacheHook(
-        cache_interval=config.cache_interval,
-        disable_cache_before_step=config.disable_cache_before_step,
-        taylor_factors_dtype=config.taylor_factors_dtype,
-        disable_cache_after_step=config.disable_cache_after_step,
-        state_manager=state_manager,
-    )
-
-    registry.register_hook(hook, _TAYLORSEER_CACHE_HOOK)

diffusers/hooks/text_kv_cache.py

@@ -1,173 +0,0 @@
-# Copyright 2026 The HuggingFace Team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-from dataclasses import dataclass
-
-import torch
-
-from .hooks import BaseState, HookRegistry, ModelHook, StateManager
-
-
-_TEXT_KV_CACHE_TRANSFORMER_HOOK = "text_kv_cache_transformer"
-_TEXT_KV_CACHE_BLOCK_HOOK = "text_kv_cache_block"
-
-
-@dataclass
-class TextKVCacheConfig:
-    """Enable exact (lossless) text K/V caching for transformer models.
-
-    Pre-computes per-block text key and value projections once before the denoising loop and reuses them across all
-    steps. Positive and negative prompts are distinguished via a stable cache key captured by a transformer-level hook
-    before any intermediate tensor allocations.
-    """
-
-    pass
-
-
-class TextKVCacheState(BaseState):
-    """Shared state between the transformer-level and block-level hooks.
-
-    The transformer hook writes the stable ``encoder_hidden_states`` ``data_ptr()`` (captured *before* ``txt_norm``) so
-    that block hooks can use it as a reliable cache key across denoising steps.
-    """
-
-    def __init__(self):
-        self.key: int | None = None
-
-    def reset(self):
-        self.key = None
-
-
-class TextKVCacheBlockState(BaseState):
-    """Per-block state holding cached text key/value projections."""
-
-    def __init__(self):
-        self.kv_cache: dict[int, tuple[torch.Tensor, torch.Tensor]] = {}
-
-    def reset(self):
-        self.kv_cache.clear()
-
-
-class TextKVCacheTransformerHook(ModelHook):
-    """Captures ``encoder_hidden_states.data_ptr()`` before ``txt_norm``
-    and writes it to shared state for the block hooks to read."""
-
-    _is_stateful = True
-
-    def __init__(self, state_manager: StateManager):
-        super().__init__()
-        self.state_manager = state_manager
-
-    def new_forward(self, module: torch.nn.Module, *args, **kwargs):
-        if self.state_manager._current_context is None:
-            self.state_manager.set_context("inference")
-
-        encoder_hidden_states = kwargs.get("encoder_hidden_states")
-        if encoder_hidden_states is not None:
-            state: TextKVCacheState = self.state_manager.get_state()
-            state.key = encoder_hidden_states.data_ptr()
-        return self.fn_ref.original_forward(*args, **kwargs)
-
-    def reset_state(self, module: torch.nn.Module):
-        self.state_manager.reset()
-        return module
-
-
-class TextKVCacheBlockHook(ModelHook):
-    """Caches ``(txt_key, txt_value)`` per block per unique prompt using
-    the stable cache key from the shared state."""
-
-    _is_stateful = True
-
-    def __init__(self, state_manager: StateManager, block_state_manager: StateManager):
-        super().__init__()
-        self.state_manager = state_manager
-        self.block_state_manager = block_state_manager
-
-    def new_forward(self, module: torch.nn.Module, *args, **kwargs):
-        from ..models.transformers.transformer_nucleusmoe_image import _apply_rotary_emb_nucleus
-
-        if self.state_manager._current_context is None:
-            self.state_manager.set_context("inference")
-
-        if self.block_state_manager._current_context is None:
-            self.block_state_manager.set_context("inference")
-
-        if "encoder_hidden_states" in kwargs:
-            encoder_hidden_states = kwargs["encoder_hidden_states"]
-        else:
-            encoder_hidden_states = args[1]
-
-        if "image_rotary_emb" in kwargs:
-            image_rotary_emb = kwargs["image_rotary_emb"]
-        elif len(args) > 3:
-            image_rotary_emb = args[3]
-        else:
-            image_rotary_emb = None
-
-        state: TextKVCacheState = self.state_manager.get_state()
-        cache_key = state.key
-
-        block_state: TextKVCacheBlockState = self.block_state_manager.get_state()
-
-        if cache_key not in block_state.kv_cache:
-            context = module.encoder_proj(encoder_hidden_states)
-
-            attn = module.attn
-            head_dim = attn.inner_dim // attn.heads
-            num_kv_heads = attn.inner_kv_dim // head_dim
-
-            txt_key = attn.add_k_proj(context).unflatten(-1, (num_kv_heads, -1))
-            txt_value = attn.add_v_proj(context).unflatten(-1, (num_kv_heads, -1))
-
-            if attn.norm_added_k is not None:
-                txt_key = attn.norm_added_k(txt_key)
-
-            if image_rotary_emb is not None:
-                _, txt_freqs = image_rotary_emb
-                txt_key = _apply_rotary_emb_nucleus(txt_key, txt_freqs, use_real=False)
-
-            block_state.kv_cache[cache_key] = (txt_key, txt_value)
-
-        txt_key, txt_value = block_state.kv_cache[cache_key]
-
-        attn_kwargs = kwargs.get("attention_kwargs") or {}
-        attn_kwargs["cached_txt_key"] = txt_key
-        attn_kwargs["cached_txt_value"] = txt_value
-        kwargs["attention_kwargs"] = attn_kwargs
-
-        return self.fn_ref.original_forward(*args, **kwargs)
-
-    def reset_state(self, module: torch.nn.Module):
-        self.block_state_manager.reset()
-        return module
-
-
-def apply_text_kv_cache(module: torch.nn.Module, config: TextKVCacheConfig) -> None:
-    from ..models.transformers.transformer_nucleusmoe_image import NucleusMoEImageTransformerBlock
-
-    HookRegistry.check_if_exists_or_initialize(module)
-
-    state_manager = StateManager(TextKVCacheState)
-
-    transformer_hook = TextKVCacheTransformerHook(state_manager)
-    registry = HookRegistry.check_if_exists_or_initialize(module)
-    registry.register_hook(transformer_hook, _TEXT_KV_CACHE_TRANSFORMER_HOOK)
-
-    for _, submodule in module.named_modules():
-        if isinstance(submodule, NucleusMoEImageTransformerBlock):
-            block_state_manager = StateManager(TextKVCacheBlockState)
-            hook = TextKVCacheBlockHook(state_manager, block_state_manager)
-            block_registry = HookRegistry.check_if_exists_or_initialize(submodule)
-            block_registry.register_hook(hook, _TEXT_KV_CACHE_BLOCK_HOOK)

diffusers/hooks/utils.py

@@ -1,43 +0,0 @@
-# Copyright 2026 The HuggingFace Team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-import torch
-
-from ._common import _ALL_TRANSFORMER_BLOCK_IDENTIFIERS, _ATTENTION_CLASSES, _FEEDFORWARD_CLASSES
-
-
-def _get_identifiable_transformer_blocks_in_module(module: torch.nn.Module):
-    module_list_with_transformer_blocks = []
-    for name, submodule in module.named_modules():
-        name_endswith_identifier = any(name.endswith(identifier) for identifier in _ALL_TRANSFORMER_BLOCK_IDENTIFIERS)
-        is_ModuleList = isinstance(submodule, torch.nn.ModuleList)
-        if name_endswith_identifier and is_ModuleList:
-            module_list_with_transformer_blocks.append((name, submodule))
-    return module_list_with_transformer_blocks
-
-
-def _get_identifiable_attention_layers_in_module(module: torch.nn.Module):
-    attention_layers = []
-    for name, submodule in module.named_modules():
-        if isinstance(submodule, _ATTENTION_CLASSES):
-            attention_layers.append((name, submodule))
-    return attention_layers
-
-
-def _get_identifiable_feedforward_layers_in_module(module: torch.nn.Module):
-    feedforward_layers = []
-    for name, submodule in module.named_modules():
-        if isinstance(submodule, _FEEDFORWARD_CLASSES):
-            feedforward_layers.append((name, submodule))
-    return feedforward_layers

diffusers/image_processor.py

@@ -1,1468 +0,0 @@
-# Copyright 2026 The HuggingFace Team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-import math
-import warnings
-
-import numpy as np
-import PIL.Image
-import torch
-import torch.nn.functional as F
-from PIL import Image, ImageFilter, ImageOps
-
-from .configuration_utils import ConfigMixin, register_to_config
-from .utils import CONFIG_NAME, PIL_INTERPOLATION, deprecate
-
-
-PipelineImageInput = (
-    PIL.Image.Image | np.ndarray | torch.Tensor | list[PIL.Image.Image] | list[np.ndarray] | list[torch.Tensor]
-)
-
-PipelineDepthInput = PipelineImageInput
-
-
-def is_valid_image(image) -> bool:
-    r"""
-    Checks if the input is a valid image.
-
-    A valid image can be:
-    - A `PIL.Image.Image`.
-    - A 2D or 3D `np.ndarray` or `torch.Tensor` (grayscale or color image).
-
-    Args:
-        image (`PIL.Image.Image | np.ndarray | torch.Tensor`):
-            The image to validate. It can be a PIL image, a NumPy array, or a torch tensor.
-
-    Returns:
-        `bool`:
-            `True` if the input is a valid image, `False` otherwise.
-    """
-    return isinstance(image, PIL.Image.Image) or isinstance(image, (np.ndarray, torch.Tensor)) and image.ndim in (2, 3)
-
-
-def is_valid_image_imagelist(images):
-    r"""
-    Checks if the input is a valid image or list of images.
-
-    The input can be one of the following formats:
-    - A 4D tensor or numpy array (batch of images).
-    - A valid single image: `PIL.Image.Image`, 2D `np.ndarray` or `torch.Tensor` (grayscale image), 3D `np.ndarray` or
-      `torch.Tensor`.
-    - A list of valid images.
-
-    Args:
-        images (`np.ndarray | torch.Tensor | PIL.Image.Image | list`):
-            The image(s) to check. Can be a batch of images (4D tensor/array), a single image, or a list of valid
-            images.
-
-    Returns:
-        `bool`:
-            `True` if the input is valid, `False` otherwise.
-    """
-    if isinstance(images, (np.ndarray, torch.Tensor)) and images.ndim == 4:
-        return True
-    elif is_valid_image(images):
-        return True
-    elif isinstance(images, list):
-        return all(is_valid_image(image) for image in images)
-    return False
-
-
-class VaeImageProcessor(ConfigMixin):
-    """
-    Image processor for VAE.
-
-    Args:
-        do_resize (`bool`, *optional*, defaults to `True`):
-            Whether to downscale the image's (height, width) dimensions to multiples of `vae_scale_factor`. Can accept
-            `height` and `width` arguments from [`image_processor.VaeImageProcessor.preprocess`] method.
-        vae_scale_factor (`int`, *optional*, defaults to `8`):
-            VAE scale factor. If `do_resize` is `True`, the image is automatically resized to multiples of this factor.
-        resample (`str`, *optional*, defaults to `lanczos`):
-            Resampling filter to use when resizing the image.
-        do_normalize (`bool`, *optional*, defaults to `True`):
-            Whether to normalize the image to [-1,1].
-        do_binarize (`bool`, *optional*, defaults to `False`):
-            Whether to binarize the image to 0/1.
-        do_convert_rgb (`bool`, *optional*, defaults to be `False`):
-            Whether to convert the images to RGB format.
-        do_convert_grayscale (`bool`, *optional*, defaults to be `False`):
-            Whether to convert the images to grayscale format.
-    """
-
-    config_name = CONFIG_NAME
-
-    @register_to_config
-    def __init__(
-        self,
-        do_resize: bool = True,
-        vae_scale_factor: int = 8,
-        vae_latent_channels: int = 4,
-        resample: str = "lanczos",
-        reducing_gap: int | None = None,
-        do_normalize: bool = True,
-        do_binarize: bool = False,
-        do_convert_rgb: bool = False,
-        do_convert_grayscale: bool = False,
-    ):
-        super().__init__()
-        if do_convert_rgb and do_convert_grayscale:
-            raise ValueError(
-                "`do_convert_rgb` and `do_convert_grayscale` can not both be set to `True`,"
-                " if you intended to convert the image into RGB format, please set `do_convert_grayscale = False`.",
-                " if you intended to convert the image into grayscale format, please set `do_convert_rgb = False`",
-            )
-
-    @staticmethod
-    def numpy_to_pil(images: np.ndarray) -> list[PIL.Image.Image]:
-        r"""
-        Convert a numpy image or a batch of images to a PIL image.
-
-        Args:
-            images (`np.ndarray`):
-                The image array to convert to PIL format.
-
-        Returns:
-            `list[PIL.Image.Image]`:
-                A list of PIL images.
-        """
-        if images.ndim == 3:
-            images = images[None, ...]
-        images = (images * 255).round().astype("uint8")
-        if images.shape[-1] == 1:
-            # special case for grayscale (single channel) images
-            pil_images = [Image.fromarray(image.squeeze(), mode="L") for image in images]
-        else:
-            pil_images = [Image.fromarray(image) for image in images]
-
-        return pil_images
-
-    @staticmethod
-    def pil_to_numpy(images: list[PIL.Image.Image] | PIL.Image.Image) -> np.ndarray:
-        r"""
-        Convert a PIL image or a list of PIL images to NumPy arrays.
-
-        Args:
-            images (`PIL.Image.Image` or `list[PIL.Image.Image]`):
-                The PIL image or list of images to convert to NumPy format.
-
-        Returns:
-            `np.ndarray`:
-                A NumPy array representation of the images.
-        """
-        if not isinstance(images, list):
-            images = [images]
-        images = [np.array(image).astype(np.float32) / 255.0 for image in images]
-        images = np.stack(images, axis=0)
-
-        return images
-
-    @staticmethod
-    def numpy_to_pt(images: np.ndarray) -> torch.Tensor:
-        r"""
-        Convert a NumPy image to a PyTorch tensor.
-
-        Args:
-            images (`np.ndarray`):
-                The NumPy image array to convert to PyTorch format.
-
-        Returns:
-            `torch.Tensor`:
-                A PyTorch tensor representation of the images.
-        """
-        if images.ndim == 3:
-            images = images[..., None]
-
-        images = torch.from_numpy(images.transpose(0, 3, 1, 2))
-        return images
-
-    @staticmethod
-    def pt_to_numpy(images: torch.Tensor) -> np.ndarray:
-        r"""
-        Convert a PyTorch tensor to a NumPy image.
-
-        Args:
-            images (`torch.Tensor`):
-                The PyTorch tensor to convert to NumPy format.
-
-        Returns:
-            `np.ndarray`:
-                A NumPy array representation of the images.
-        """
-        images = images.cpu().permute(0, 2, 3, 1).float().numpy()
-        return images
-
-    @staticmethod
-    def normalize(images: np.ndarray | torch.Tensor) -> np.ndarray | torch.Tensor:
-        r"""
-        Normalize an image array to [-1,1].
-
-        Args:
-            images (`np.ndarray` or `torch.Tensor`):
-                The image array to normalize.
-
-        Returns:
-            `np.ndarray` or `torch.Tensor`:
-                The normalized image array.
-        """
-        return 2.0 * images - 1.0
-
-    @staticmethod
-    def denormalize(images: np.ndarray | torch.Tensor) -> np.ndarray | torch.Tensor:
-        r"""
-        Denormalize an image array to [0,1].
-
-        Args:
-            images (`np.ndarray` or `torch.Tensor`):
-                The image array to denormalize.
-
-        Returns:
-            `np.ndarray` or `torch.Tensor`:
-                The denormalized image array.
-        """
-        return (images * 0.5 + 0.5).clamp(0, 1)
-
-    @staticmethod
-    def convert_to_rgb(image: PIL.Image.Image) -> PIL.Image.Image:
-        r"""
-        Converts a PIL image to RGB format.
-
-        Args:
-            image (`PIL.Image.Image`):
-                The PIL image to convert to RGB.
-
-        Returns:
-            `PIL.Image.Image`:
-                The RGB-converted PIL image.
-        """
-        image = image.convert("RGB")
-
-        return image
-
-    @staticmethod
-    def convert_to_grayscale(image: PIL.Image.Image) -> PIL.Image.Image:
-        r"""
-        Converts a given PIL image to grayscale.
-
-        Args:
-            image (`PIL.Image.Image`):
-                The input image to convert.
-
-        Returns:
-            `PIL.Image.Image`:
-                The image converted to grayscale.
-        """
-        image = image.convert("L")
-
-        return image
-
-    @staticmethod
-    def blur(image: PIL.Image.Image, blur_factor: int = 4) -> PIL.Image.Image:
-        r"""
-        Applies Gaussian blur to an image.
-
-        Args:
-            image (`PIL.Image.Image`):
-                The PIL image to convert to grayscale.
-
-        Returns:
-            `PIL.Image.Image`:
-                The grayscale-converted PIL image.
-        """
-        image = image.filter(ImageFilter.GaussianBlur(blur_factor))
-
-        return image
-
-    @staticmethod
-    def get_crop_region(mask_image: PIL.Image.Image, width: int, height: int, pad=0):
-        r"""
-        Finds a rectangular region that contains all masked ares in an image, and expands region to match the aspect
-        ratio of the original image; for example, if user drew mask in a 128x32 region, and the dimensions for
-        processing are 512x512, the region will be expanded to 128x128.
-
-        Args:
-            mask_image (PIL.Image.Image): Mask image.
-            width (int): Width of the image to be processed.
-            height (int): Height of the image to be processed.
-            pad (int, optional): Padding to be added to the crop region. Defaults to 0.
-
-        Returns:
-            tuple: (x1, y1, x2, y2) represent a rectangular region that contains all masked ares in an image and
-            matches the original aspect ratio.
-        """
-
-        mask_image = mask_image.convert("L")
-        mask = np.array(mask_image)
-
-        # 1. find a rectangular region that contains all masked ares in an image
-        h, w = mask.shape
-        crop_left = 0
-        for i in range(w):
-            if not (mask[:, i] == 0).all():
-                break
-            crop_left += 1
-
-        crop_right = 0
-        for i in reversed(range(w)):
-            if not (mask[:, i] == 0).all():
-                break
-            crop_right += 1
-
-        crop_top = 0
-        for i in range(h):
-            if not (mask[i] == 0).all():
-                break
-            crop_top += 1
-
-        crop_bottom = 0
-        for i in reversed(range(h)):
-            if not (mask[i] == 0).all():
-                break
-            crop_bottom += 1
-
-        # 2. add padding to the crop region
-        x1, y1, x2, y2 = (
-            int(max(crop_left - pad, 0)),
-            int(max(crop_top - pad, 0)),
-            int(min(w - crop_right + pad, w)),
-            int(min(h - crop_bottom + pad, h)),
-        )
-
-        # 3. expands crop region to match the aspect ratio of the image to be processed
-        ratio_crop_region = (x2 - x1) / (y2 - y1)
-        ratio_processing = width / height
-
-        if ratio_crop_region > ratio_processing:
-            desired_height = (x2 - x1) / ratio_processing
-            desired_height_diff = int(desired_height - (y2 - y1))
-            y1 -= desired_height_diff // 2
-            y2 += desired_height_diff - desired_height_diff // 2
-            if y2 >= mask_image.height:
-                diff = y2 - mask_image.height
-                y2 -= diff
-                y1 -= diff
-            if y1 < 0:
-                y2 -= y1
-                y1 -= y1
-            if y2 >= mask_image.height:
-                y2 = mask_image.height
-        else:
-            desired_width = (y2 - y1) * ratio_processing
-            desired_width_diff = int(desired_width - (x2 - x1))
-            x1 -= desired_width_diff // 2
-            x2 += desired_width_diff - desired_width_diff // 2
-            if x2 >= mask_image.width:
-                diff = x2 - mask_image.width
-                x2 -= diff
-                x1 -= diff
-            if x1 < 0:
-                x2 -= x1
-                x1 -= x1
-            if x2 >= mask_image.width:
-                x2 = mask_image.width
-
-        return x1, y1, x2, y2
-
-    def _resize_and_fill(
-        self,
-        image: PIL.Image.Image,
-        width: int,
-        height: int,
-    ) -> PIL.Image.Image:
-        r"""
-        Resize the image to fit within the specified width and height, maintaining the aspect ratio, and then center
-        the image within the dimensions, filling empty with data from image.
-
-        Args:
-            image (`PIL.Image.Image`):
-                The image to resize and fill.
-            width (`int`):
-                The width to resize the image to.
-            height (`int`):
-                The height to resize the image to.
-
-        Returns:
-            `PIL.Image.Image`:
-                The resized and filled image.
-        """
-
-        ratio = width / height
-        src_ratio = image.width / image.height
-
-        src_w = width if ratio < src_ratio else image.width * height // image.height
-        src_h = height if ratio >= src_ratio else image.height * width // image.width
-
-        resized = image.resize((src_w, src_h), resample=PIL_INTERPOLATION[self.config.resample])
-        res = Image.new("RGB", (width, height))
-        res.paste(resized, box=(width // 2 - src_w // 2, height // 2 - src_h // 2))
-
-        if ratio < src_ratio:
-            fill_height = height // 2 - src_h // 2
-            if fill_height > 0:
-                res.paste(resized.resize((width, fill_height), box=(0, 0, width, 0)), box=(0, 0))
-                res.paste(
-                    resized.resize((width, fill_height), box=(0, resized.height, width, resized.height)),
-                    box=(0, fill_height + src_h),
-                )
-        elif ratio > src_ratio:
-            fill_width = width // 2 - src_w // 2
-            if fill_width > 0:
-                res.paste(resized.resize((fill_width, height), box=(0, 0, 0, height)), box=(0, 0))
-                res.paste(
-                    resized.resize((fill_width, height), box=(resized.width, 0, resized.width, height)),
-                    box=(fill_width + src_w, 0),
-                )
-
-        return res
-
-    def _resize_and_crop(
-        self,
-        image: PIL.Image.Image,
-        width: int,
-        height: int,
-    ) -> PIL.Image.Image:
-        r"""
-        Resize the image to fit within the specified width and height, maintaining the aspect ratio, and then center
-        the image within the dimensions, cropping the excess.
-
-        Args:
-            image (`PIL.Image.Image`):
-                The image to resize and crop.
-            width (`int`):
-                The width to resize the image to.
-            height (`int`):
-                The height to resize the image to.
-
-        Returns:
-            `PIL.Image.Image`:
-                The resized and cropped image.
-        """
-        ratio = width / height
-        src_ratio = image.width / image.height
-
-        src_w = width if ratio > src_ratio else image.width * height // image.height
-        src_h = height if ratio <= src_ratio else image.height * width // image.width
-
-        resized = image.resize((src_w, src_h), resample=PIL_INTERPOLATION[self.config.resample])
-        res = Image.new("RGB", (width, height))
-        res.paste(resized, box=(width // 2 - src_w // 2, height // 2 - src_h // 2))
-        return res
-
-    def resize(
-        self,
-        image: PIL.Image.Image | np.ndarray | torch.Tensor,
-        height: int,
-        width: int,
-        resize_mode: str = "default",  # "default", "fill", "crop"
-    ) -> PIL.Image.Image | np.ndarray | torch.Tensor:
-        """
-        Resize image.
-
-        Args:
-            image (`PIL.Image.Image`, `np.ndarray` or `torch.Tensor`):
-                The image input, can be a PIL image, numpy array or pytorch tensor.
-            height (`int`):
-                The height to resize to.
-            width (`int`):
-                The width to resize to.
-            resize_mode (`str`, *optional*, defaults to `default`):
-                The resize mode to use, can be one of `default` or `fill`. If `default`, will resize the image to fit
-                within the specified width and height, and it may not maintaining the original aspect ratio. If `fill`,
-                will resize the image to fit within the specified width and height, maintaining the aspect ratio, and
-                then center the image within the dimensions, filling empty with data from image. If `crop`, will resize
-                the image to fit within the specified width and height, maintaining the aspect ratio, and then center
-                the image within the dimensions, cropping the excess. Note that resize_mode `fill` and `crop` are only
-                supported for PIL image input.
-
-        Returns:
-            `PIL.Image.Image`, `np.ndarray` or `torch.Tensor`:
-                The resized image.
-        """
-        if resize_mode != "default" and not isinstance(image, PIL.Image.Image):
-            raise ValueError(f"Only PIL image input is supported for resize_mode {resize_mode}")
-        if isinstance(image, PIL.Image.Image):
-            if resize_mode == "default":
-                image = image.resize(
-                    (width, height),
-                    resample=PIL_INTERPOLATION[self.config.resample],
-                    reducing_gap=self.config.reducing_gap,
-                )
-            elif resize_mode == "fill":
-                image = self._resize_and_fill(image, width, height)
-            elif resize_mode == "crop":
-                image = self._resize_and_crop(image, width, height)
-            else:
-                raise ValueError(f"resize_mode {resize_mode} is not supported")
-
-        elif isinstance(image, torch.Tensor):
-            image = torch.nn.functional.interpolate(
-                image,
-                size=(height, width),
-            )
-        elif isinstance(image, np.ndarray):
-            image = self.numpy_to_pt(image)
-            image = torch.nn.functional.interpolate(
-                image,
-                size=(height, width),
-            )
-            image = self.pt_to_numpy(image)
-
-        return image
-
-    def binarize(self, image: PIL.Image.Image) -> PIL.Image.Image:
-        """
-        Create a mask.
-
-        Args:
-            image (`PIL.Image.Image`):
-                The image input, should be a PIL image.
-
-        Returns:
-            `PIL.Image.Image`:
-                The binarized image. Values less than 0.5 are set to 0, values greater than 0.5 are set to 1.
-        """
-        image[image < 0.5] = 0
-        image[image >= 0.5] = 1
-
-        return image
-
-    def _denormalize_conditionally(
-        self, images: torch.Tensor, do_denormalize: list[bool] | None = None
-    ) -> torch.Tensor:
-        r"""
-        Denormalize a batch of images based on a condition list.
-
-        Args:
-            images (`torch.Tensor`):
-                The input image tensor.
-            do_denormalize (`Optional[list[bool]`, *optional*, defaults to `None`):
-                A list of booleans indicating whether to denormalize each image in the batch. If `None`, will use the
-                value of `do_normalize` in the `VaeImageProcessor` config.
-        """
-        if do_denormalize is None:
-            return self.denormalize(images) if self.config.do_normalize else images
-
-        return torch.stack(
-            [self.denormalize(images[i]) if do_denormalize[i] else images[i] for i in range(images.shape[0])]
-        )
-
-    def get_default_height_width(
-        self,
-        image: PIL.Image.Image | np.ndarray | torch.Tensor,
-        height: int | None = None,
-        width: int | None = None,
-    ) -> tuple[int, int]:
-        r"""
-        Returns the height and width of the image, downscaled to the next integer multiple of `vae_scale_factor`.
-
-        Args:
-            image (`PIL.Image.Image | np.ndarray | torch.Tensor`):
-                The image input, which can be a PIL image, NumPy array, or PyTorch tensor. If it is a NumPy array, it
-                should have shape `[batch, height, width]` or `[batch, height, width, channels]`. If it is a PyTorch
-                tensor, it should have shape `[batch, channels, height, width]`.
-            height (`int | None`, *optional*, defaults to `None`):
-                The height of the preprocessed image. If `None`, the height of the `image` input will be used.
-            width (`int | None`, *optional*, defaults to `None`):
-                The width of the preprocessed image. If `None`, the width of the `image` input will be used.
-
-        Returns:
-            `tuple[int, int]`:
-                A tuple containing the height and width, both resized to the nearest integer multiple of
-                `vae_scale_factor`.
-        """
-
-        if height is None:
-            if isinstance(image, PIL.Image.Image):
-                height = image.height
-            elif isinstance(image, torch.Tensor):
-                height = image.shape[2]
-            else:
-                height = image.shape[1]
-
-        if width is None:
-            if isinstance(image, PIL.Image.Image):
-                width = image.width
-            elif isinstance(image, torch.Tensor):
-                width = image.shape[3]
-            else:
-                width = image.shape[2]
-
-        width, height = (
-            x - x % self.config.vae_scale_factor for x in (width, height)
-        )  # resize to integer multiple of vae_scale_factor
-
-        return height, width
-
-    def preprocess(
-        self,
-        image: PipelineImageInput,
-        height: int | None = None,
-        width: int | None = None,
-        resize_mode: str = "default",  # "default", "fill", "crop"
-        crops_coords: tuple[int, int, int, int] | None = None,
-    ) -> torch.Tensor:
-        """
-        Preprocess the image input.
-
-        Args:
-            image (`PipelineImageInput`):
-                The image input, accepted formats are PIL images, NumPy arrays, PyTorch tensors; Also accept list of
-                supported formats.
-            height (`int`, *optional*):
-                The height in preprocessed image. If `None`, will use the `get_default_height_width()` to get default
-                height.
-            width (`int`, *optional*):
-                The width in preprocessed. If `None`, will use get_default_height_width()` to get the default width.
-            resize_mode (`str`, *optional*, defaults to `default`):
-                The resize mode, can be one of `default` or `fill`. If `default`, will resize the image to fit within
-                the specified width and height, and it may not maintaining the original aspect ratio. If `fill`, will
-                resize the image to fit within the specified width and height, maintaining the aspect ratio, and then
-                center the image within the dimensions, filling empty with data from image. If `crop`, will resize the
-                image to fit within the specified width and height, maintaining the aspect ratio, and then center the
-                image within the dimensions, cropping the excess. Note that resize_mode `fill` and `crop` are only
-                supported for PIL image input.
-            crops_coords (`list[tuple[int, int, int, int]]`, *optional*, defaults to `None`):
-                The crop coordinates for each image in the batch. If `None`, will not crop the image.
-
-        Returns:
-            `torch.Tensor`:
-                The preprocessed image.
-        """
-        supported_formats = (PIL.Image.Image, np.ndarray, torch.Tensor)
-
-        # Expand the missing dimension for 3-dimensional pytorch tensor or numpy array that represents grayscale image
-        if self.config.do_convert_grayscale and isinstance(image, (torch.Tensor, np.ndarray)) and image.ndim == 3:
-            if isinstance(image, torch.Tensor):
-                # if image is a pytorch tensor could have 2 possible shapes:
-                #    1. batch x height x width: we should insert the channel dimension at position 1
-                #    2. channel x height x width: we should insert batch dimension at position 0,
-                #       however, since both channel and batch dimension has same size 1, it is same to insert at position 1
-                #    for simplicity, we insert a dimension of size 1 at position 1 for both cases
-                image = image.unsqueeze(1)
-            else:
-                # if it is a numpy array, it could have 2 possible shapes:
-                #   1. batch x height x width: insert channel dimension on last position
-                #   2. height x width x channel: insert batch dimension on first position
-                if image.shape[-1] == 1:
-                    image = np.expand_dims(image, axis=0)
-                else:
-                    image = np.expand_dims(image, axis=-1)
-
-        if isinstance(image, list) and isinstance(image[0], np.ndarray) and image[0].ndim == 4:
-            warnings.warn(
-                "Passing `image` as a list of 4d np.ndarray is deprecated."
-                "Please concatenate the list along the batch dimension and pass it as a single 4d np.ndarray",
-                FutureWarning,
-            )
-            image = np.concatenate(image, axis=0)
-        if isinstance(image, list) and isinstance(image[0], torch.Tensor) and image[0].ndim == 4:
-            warnings.warn(
-                "Passing `image` as a list of 4d torch.Tensor is deprecated."
-                "Please concatenate the list along the batch dimension and pass it as a single 4d torch.Tensor",
-                FutureWarning,
-            )
-            image = torch.cat(image, axis=0)
-
-        if not is_valid_image_imagelist(image):
-            raise ValueError(
-                f"Input is in incorrect format. Currently, we only support {', '.join(str(x) for x in supported_formats)}"
-            )
-        if not isinstance(image, list):
-            image = [image]
-
-        if isinstance(image[0], PIL.Image.Image):
-            if crops_coords is not None:
-                image = [i.crop(crops_coords) for i in image]
-            if self.config.do_resize:
-                height, width = self.get_default_height_width(image[0], height, width)
-                image = [self.resize(i, height, width, resize_mode=resize_mode) for i in image]
-            if self.config.do_convert_rgb:
-                image = [self.convert_to_rgb(i) for i in image]
-            elif self.config.do_convert_grayscale:
-                image = [self.convert_to_grayscale(i) for i in image]
-            image = self.pil_to_numpy(image)  # to np
-            image = self.numpy_to_pt(image)  # to pt
-
-        elif isinstance(image[0], np.ndarray):
-            image = np.concatenate(image, axis=0) if image[0].ndim == 4 else np.stack(image, axis=0)
-
-            image = self.numpy_to_pt(image)
-
-            height, width = self.get_default_height_width(image, height, width)
-            if self.config.do_resize:
-                image = self.resize(image, height, width)
-
-        elif isinstance(image[0], torch.Tensor):
-            image = torch.cat(image, axis=0) if image[0].ndim == 4 else torch.stack(image, axis=0)
-
-            if self.config.do_convert_grayscale and image.ndim == 3:
-                image = image.unsqueeze(1)
-
-            channel = image.shape[1]
-            # don't need any preprocess if the image is latents
-            if channel == self.config.vae_latent_channels:
-                return image
-
-            height, width = self.get_default_height_width(image, height, width)
-            if self.config.do_resize:
-                image = self.resize(image, height, width)
-
-        # expected range [0,1], normalize to [-1,1]
-        do_normalize = self.config.do_normalize
-        if do_normalize and image.min() < 0:
-            warnings.warn(
-                "Passing `image` as torch tensor with value range in [-1,1] is deprecated. The expected value range for image tensor is [0,1] "
-                f"when passing as pytorch tensor or numpy Array. You passed `image` with value range [{image.min()},{image.max()}]",
-                FutureWarning,
-            )
-            do_normalize = False
-        if do_normalize:
-            image = self.normalize(image)
-
-        if self.config.do_binarize:
-            image = self.binarize(image)
-
-        return image
-
-    def postprocess(
-        self,
-        image: torch.Tensor,
-        output_type: str = "pil",
-        do_denormalize: list[bool] | None = None,
-    ) -> PIL.Image.Image | np.ndarray | torch.Tensor:
-        """
-        Postprocess the image output from tensor to `output_type`.
-
-        Args:
-            image (`torch.Tensor`):
-                The image input, should be a pytorch tensor with shape `B x C x H x W`.
-            output_type (`str`, *optional*, defaults to `pil`):
-                The output type of the image, can be one of `pil`, `np`, `pt`, `latent`.
-            do_denormalize (`list[bool]`, *optional*, defaults to `None`):
-                Whether to denormalize the image to [0,1]. If `None`, will use the value of `do_normalize` in the
-                `VaeImageProcessor` config.
-
-        Returns:
-            `PIL.Image.Image`, `np.ndarray` or `torch.Tensor`:
-                The postprocessed image.
-        """
-        if not isinstance(image, torch.Tensor):
-            raise ValueError(
-                f"Input for postprocessing is in incorrect format: {type(image)}. We only support pytorch tensor"
-            )
-        if output_type not in ["latent", "pt", "np", "pil"]:
-            deprecation_message = (
-                f"the output_type {output_type} is outdated and has been set to `np`. Please make sure to set it to one of these instead: "
-                "`pil`, `np`, `pt`, `latent`"
-            )
-            deprecate("Unsupported output_type", "1.0.0", deprecation_message, standard_warn=False)
-            output_type = "np"
-
-        if output_type == "latent":
-            return image
-
-        image = self._denormalize_conditionally(image, do_denormalize)
-
-        if output_type == "pt":
-            return image
-
-        image = self.pt_to_numpy(image)
-
-        if output_type == "np":
-            return image
-
-        if output_type == "pil":
-            return self.numpy_to_pil(image)
-
-    def apply_overlay(
-        self,
-        mask: PIL.Image.Image,
-        init_image: PIL.Image.Image,
-        image: PIL.Image.Image,
-        crop_coords: tuple[int, int, int, int] | None = None,
-    ) -> PIL.Image.Image:
-        r"""
-        Applies an overlay of the mask and the inpainted image on the original image.
-
-        Args:
-            mask (`PIL.Image.Image`):
-                The mask image that highlights regions to overlay.
-            init_image (`PIL.Image.Image`):
-                The original image to which the overlay is applied.
-            image (`PIL.Image.Image`):
-                The image to overlay onto the original.
-            crop_coords (`tuple[int, int, int, int]`, *optional*):
-                Coordinates to crop the image. If provided, the image will be cropped accordingly.
-
-        Returns:
-            `PIL.Image.Image`:
-                The final image with the overlay applied.
-        """
-
-        width, height = init_image.width, init_image.height
-
-        init_image_masked = PIL.Image.new("RGBa", (width, height))
-        init_image_masked.paste(init_image.convert("RGBA").convert("RGBa"), mask=ImageOps.invert(mask.convert("L")))
-
-        init_image_masked = init_image_masked.convert("RGBA")
-
-        if crop_coords is not None:
-            x, y, x2, y2 = crop_coords
-            w = x2 - x
-            h = y2 - y
-            base_image = PIL.Image.new("RGBA", (width, height))
-            image = self.resize(image, height=h, width=w, resize_mode="crop")
-            base_image.paste(image, (x, y))
-            image = base_image.convert("RGB")
-
-        image = image.convert("RGBA")
-        image.alpha_composite(init_image_masked)
-        image = image.convert("RGB")
-
-        return image
-
-
-class InpaintProcessor(ConfigMixin):
-    """
-    Image processor for inpainting image and mask.
-    """
-
-    config_name = CONFIG_NAME
-
-    @register_to_config
-    def __init__(
-        self,
-        do_resize: bool = True,
-        vae_scale_factor: int = 8,
-        vae_latent_channels: int = 4,
-        resample: str = "lanczos",
-        reducing_gap: int | None = None,
-        do_normalize: bool = True,
-        do_binarize: bool = False,
-        do_convert_grayscale: bool = False,
-        mask_do_normalize: bool = False,
-        mask_do_binarize: bool = True,
-        mask_do_convert_grayscale: bool = True,
-    ):
-        super().__init__()
-
-        self._image_processor = VaeImageProcessor(
-            do_resize=do_resize,
-            vae_scale_factor=vae_scale_factor,
-            vae_latent_channels=vae_latent_channels,
-            resample=resample,
-            reducing_gap=reducing_gap,
-            do_normalize=do_normalize,
-            do_binarize=do_binarize,
-            do_convert_grayscale=do_convert_grayscale,
-        )
-        self._mask_processor = VaeImageProcessor(
-            do_resize=do_resize,
-            vae_scale_factor=vae_scale_factor,
-            vae_latent_channels=vae_latent_channels,
-            resample=resample,
-            reducing_gap=reducing_gap,
-            do_normalize=mask_do_normalize,
-            do_binarize=mask_do_binarize,
-            do_convert_grayscale=mask_do_convert_grayscale,
-        )
-
-    def preprocess(
-        self,
-        image: PIL.Image.Image,
-        mask: PIL.Image.Image | None = None,
-        height: int | None = None,
-        width: int | None = None,
-        padding_mask_crop: int | None = None,
-    ) -> tuple[torch.Tensor, torch.Tensor]:
-        """
-        Preprocess the image and mask.
-        """
-        if mask is None and padding_mask_crop is not None:
-            raise ValueError("mask must be provided if padding_mask_crop is provided")
-
-        # if mask is None, same behavior as regular image processor
-        if mask is None:
-            return self._image_processor.preprocess(image, height=height, width=width)
-
-        if padding_mask_crop is not None:
-            crops_coords = self._image_processor.get_crop_region(mask, width, height, pad=padding_mask_crop)
-            resize_mode = "fill"
-        else:
-            crops_coords = None
-            resize_mode = "default"
-
-        processed_image = self._image_processor.preprocess(
-            image,
-            height=height,
-            width=width,
-            crops_coords=crops_coords,
-            resize_mode=resize_mode,
-        )
-
-        processed_mask = self._mask_processor.preprocess(
-            mask,
-            height=height,
-            width=width,
-            resize_mode=resize_mode,
-            crops_coords=crops_coords,
-        )
-
-        if crops_coords is not None:
-            postprocessing_kwargs = {
-                "crops_coords": crops_coords,
-                "original_image": image,
-                "original_mask": mask,
-            }
-        else:
-            postprocessing_kwargs = {
-                "crops_coords": None,
-                "original_image": None,
-                "original_mask": None,
-            }
-
-        return processed_image, processed_mask, postprocessing_kwargs
-
-    def postprocess(
-        self,
-        image: torch.Tensor,
-        output_type: str = "pil",
-        original_image: PIL.Image.Image | None = None,
-        original_mask: PIL.Image.Image | None = None,
-        crops_coords: tuple[int, int, int, int] | None = None,
-    ) -> tuple[PIL.Image.Image, PIL.Image.Image]:
-        """
-        Postprocess the image, optionally apply mask overlay
-        """
-        image = self._image_processor.postprocess(
-            image,
-            output_type=output_type,
-        )
-        # optionally apply the mask overlay
-        if crops_coords is not None and (original_image is None or original_mask is None):
-            raise ValueError("original_image and original_mask must be provided if crops_coords is provided")
-
-        elif crops_coords is not None and output_type != "pil":
-            raise ValueError("output_type must be 'pil' if crops_coords is provided")
-
-        elif crops_coords is not None:
-            image = [
-                self._image_processor.apply_overlay(original_mask, original_image, i, crops_coords) for i in image
-            ]
-
-        return image
-
-
-class VaeImageProcessorLDM3D(VaeImageProcessor):
-    """
-    Image processor for VAE LDM3D.
-
-    Args:
-        do_resize (`bool`, *optional*, defaults to `True`):
-            Whether to downscale the image's (height, width) dimensions to multiples of `vae_scale_factor`.
-        vae_scale_factor (`int`, *optional*, defaults to `8`):
-            VAE scale factor. If `do_resize` is `True`, the image is automatically resized to multiples of this factor.
-        resample (`str`, *optional*, defaults to `lanczos`):
-            Resampling filter to use when resizing the image.
-        do_normalize (`bool`, *optional*, defaults to `True`):
-            Whether to normalize the image to [-1,1].
-    """
-
-    config_name = CONFIG_NAME
-
-    @register_to_config
-    def __init__(
-        self,
-        do_resize: bool = True,
-        vae_scale_factor: int = 8,
-        resample: str = "lanczos",
-        do_normalize: bool = True,
-    ):
-        super().__init__()
-
-    @staticmethod
-    def numpy_to_pil(images: np.ndarray) -> list[PIL.Image.Image]:
-        r"""
-        Convert a NumPy image or a batch of images to a list of PIL images.
-
-        Args:
-            images (`np.ndarray`):
-                The input NumPy array of images, which can be a single image or a batch.
-
-        Returns:
-            `list[PIL.Image.Image]`:
-                A list of PIL images converted from the input NumPy array.
-        """
-        if images.ndim == 3:
-            images = images[None, ...]
-        images = (images * 255).round().astype("uint8")
-        if images.shape[-1] == 1:
-            # special case for grayscale (single channel) images
-            pil_images = [Image.fromarray(image.squeeze(), mode="L") for image in images]
-        else:
-            pil_images = [Image.fromarray(image[:, :, :3]) for image in images]
-
-        return pil_images
-
-    @staticmethod
-    def depth_pil_to_numpy(images: list[PIL.Image.Image] | PIL.Image.Image) -> np.ndarray:
-        r"""
-        Convert a PIL image or a list of PIL images to NumPy arrays.
-
-        Args:
-            images (`list[PIL.Image.Image, PIL.Image.Image]`):
-                The input image or list of images to be converted.
-
-        Returns:
-            `np.ndarray`:
-                A NumPy array of the converted images.
-        """
-        if not isinstance(images, list):
-            images = [images]
-
-        images = [np.array(image).astype(np.float32) / (2**16 - 1) for image in images]
-        images = np.stack(images, axis=0)
-        return images
-
-    @staticmethod
-    def rgblike_to_depthmap(image: np.ndarray | torch.Tensor) -> np.ndarray | torch.Tensor:
-        r"""
-        Convert an RGB-like depth image to a depth map.
-        """
-        # 1. Cast the tensor to a larger integer type (e.g., int32)
-        #    to safely perform the multiplication by 256.
-        # 2. Perform the 16-bit combination: High-byte * 256 + Low-byte.
-        # 3. Cast the final result to the desired depth map type (uint16) if needed
-        #    before returning, though leaving it as int32/int64 is often safer
-        #    for return value from a library function.
-
-        if isinstance(image, torch.Tensor):
-            # Cast to a safe dtype (e.g., int32 or int64) for the calculation
-            original_dtype = image.dtype
-            image_safe = image.to(torch.int32)
-
-            # Calculate the depth map
-            depth_map = image_safe[:, :, 1] * 256 + image_safe[:, :, 2]
-
-            # You may want to cast the final result to uint16, but casting to a
-            # larger int type (like int32) is sufficient to fix the overflow.
-            # depth_map = depth_map.to(torch.uint16) # Uncomment if uint16 is strictly required
-            return depth_map.to(original_dtype)
-
-        elif isinstance(image, np.ndarray):
-            # NumPy equivalent: Cast to a safe dtype (e.g., np.int32)
-            original_dtype = image.dtype
-            image_safe = image.astype(np.int32)
-
-            # Calculate the depth map
-            depth_map = image_safe[:, :, 1] * 256 + image_safe[:, :, 2]
-
-            # depth_map = depth_map.astype(np.uint16) # Uncomment if uint16 is strictly required
-            return depth_map.astype(original_dtype)
-        else:
-            raise TypeError("Input image must be a torch.Tensor or np.ndarray")
-
-    def numpy_to_depth(self, images: np.ndarray) -> list[PIL.Image.Image]:
-        r"""
-        Convert a NumPy depth image or a batch of images to a list of PIL images.
-
-        Args:
-            images (`np.ndarray`):
-                The input NumPy array of depth images, which can be a single image or a batch.
-
-        Returns:
-            `list[PIL.Image.Image]`:
-                A list of PIL images converted from the input NumPy depth images.
-        """
-        if images.ndim == 3:
-            images = images[None, ...]
-        images_depth = images[:, :, :, 3:]
-        if images.shape[-1] == 6:
-            images_depth = (images_depth * 255).round().astype("uint8")
-            pil_images = [
-                Image.fromarray(self.rgblike_to_depthmap(image_depth), mode="I;16") for image_depth in images_depth
-            ]
-        elif images.shape[-1] == 4:
-            images_depth = (images_depth * 65535.0).astype(np.uint16)
-            pil_images = [Image.fromarray(image_depth, mode="I;16") for image_depth in images_depth]
-        else:
-            raise Exception("Not supported")
-
-        return pil_images
-
-    def postprocess(
-        self,
-        image: torch.Tensor,
-        output_type: str = "pil",
-        do_denormalize: list[bool] | None = None,
-    ) -> PIL.Image.Image | np.ndarray | torch.Tensor:
-        """
-        Postprocess the image output from tensor to `output_type`.
-
-        Args:
-            image (`torch.Tensor`):
-                The image input, should be a pytorch tensor with shape `B x C x H x W`.
-            output_type (`str`, *optional*, defaults to `pil`):
-                The output type of the image, can be one of `pil`, `np`, `pt`, `latent`.
-            do_denormalize (`list[bool]`, *optional*, defaults to `None`):
-                Whether to denormalize the image to [0,1]. If `None`, will use the value of `do_normalize` in the
-                `VaeImageProcessor` config.
-
-        Returns:
-            `PIL.Image.Image`, `np.ndarray` or `torch.Tensor`:
-                The postprocessed image.
-        """
-        if not isinstance(image, torch.Tensor):
-            raise ValueError(
-                f"Input for postprocessing is in incorrect format: {type(image)}. We only support pytorch tensor"
-            )
-        if output_type not in ["latent", "pt", "np", "pil"]:
-            deprecation_message = (
-                f"the output_type {output_type} is outdated and has been set to `np`. Please make sure to set it to one of these instead: "
-                "`pil`, `np`, `pt`, `latent`"
-            )
-            deprecate("Unsupported output_type", "1.0.0", deprecation_message, standard_warn=False)
-            output_type = "np"
-
-        image = self._denormalize_conditionally(image, do_denormalize)
-
-        image = self.pt_to_numpy(image)
-
-        if output_type == "np":
-            if image.shape[-1] == 6:
-                image_depth = np.stack([self.rgblike_to_depthmap(im[:, :, 3:]) for im in image], axis=0)
-            else:
-                image_depth = image[:, :, :, 3:]
-            return image[:, :, :, :3], image_depth
-
-        if output_type == "pil":
-            return self.numpy_to_pil(image), self.numpy_to_depth(image)
-        else:
-            raise Exception(f"This type {output_type} is not supported")
-
-    def preprocess(
-        self,
-        rgb: torch.Tensor | PIL.Image.Image | np.ndarray,
-        depth: torch.Tensor | PIL.Image.Image | np.ndarray,
-        height: int | None = None,
-        width: int | None = None,
-        target_res: int | None = None,
-    ) -> torch.Tensor:
-        r"""
-        Preprocess the image input. Accepted formats are PIL images, NumPy arrays, or PyTorch tensors.
-
-        Args:
-            rgb (`torch.Tensor | PIL.Image.Image | np.ndarray`):
-                The RGB input image, which can be a single image or a batch.
-            depth (`torch.Tensor | PIL.Image.Image | np.ndarray`):
-                The depth input image, which can be a single image or a batch.
-            height (`int | None`, *optional*, defaults to `None`):
-                The desired height of the processed image. If `None`, defaults to the height of the input image.
-            width (`int | None`, *optional*, defaults to `None`):
-                The desired width of the processed image. If `None`, defaults to the width of the input image.
-            target_res (`int | None`, *optional*, defaults to `None`):
-                Target resolution for resizing the images. If specified, overrides height and width.
-
-        Returns:
-            `tuple[torch.Tensor, torch.Tensor]`:
-                A tuple containing the processed RGB and depth images as PyTorch tensors.
-        """
-        supported_formats = (PIL.Image.Image, np.ndarray, torch.Tensor)
-
-        # Expand the missing dimension for 3-dimensional pytorch tensor or numpy array that represents grayscale image
-        if self.config.do_convert_grayscale and isinstance(rgb, (torch.Tensor, np.ndarray)) and rgb.ndim == 3:
-            raise Exception("This is not yet supported")
-
-        if isinstance(rgb, supported_formats):
-            rgb = [rgb]
-            depth = [depth]
-        elif not (isinstance(rgb, list) and all(isinstance(i, supported_formats) for i in rgb)):
-            raise ValueError(
-                f"Input is in incorrect format: {[type(i) for i in rgb]}. Currently, we only support {', '.join(supported_formats)}"
-            )
-
-        if isinstance(rgb[0], PIL.Image.Image):
-            if self.config.do_convert_rgb:
-                raise Exception("This is not yet supported")
-                # rgb = [self.convert_to_rgb(i) for i in rgb]
-                # depth = [self.convert_to_depth(i) for i in depth]  #TODO define convert_to_depth
-            if self.config.do_resize or target_res:
-                height, width = self.get_default_height_width(rgb[0], height, width) if not target_res else target_res
-                rgb = [self.resize(i, height, width) for i in rgb]
-                depth = [self.resize(i, height, width) for i in depth]
-            rgb = self.pil_to_numpy(rgb)  # to np
-            rgb = self.numpy_to_pt(rgb)  # to pt
-
-            depth = self.depth_pil_to_numpy(depth)  # to np
-            depth = self.numpy_to_pt(depth)  # to pt
-
-        elif isinstance(rgb[0], np.ndarray):
-            rgb = np.concatenate(rgb, axis=0) if rgb[0].ndim == 4 else np.stack(rgb, axis=0)
-            rgb = self.numpy_to_pt(rgb)
-            height, width = self.get_default_height_width(rgb, height, width)
-            if self.config.do_resize:
-                rgb = self.resize(rgb, height, width)
-
-            depth = np.concatenate(depth, axis=0) if rgb[0].ndim == 4 else np.stack(depth, axis=0)
-            depth = self.numpy_to_pt(depth)
-            height, width = self.get_default_height_width(depth, height, width)
-            if self.config.do_resize:
-                depth = self.resize(depth, height, width)
-
-        elif isinstance(rgb[0], torch.Tensor):
-            raise Exception("This is not yet supported")
-            # rgb = torch.cat(rgb, axis=0) if rgb[0].ndim == 4 else torch.stack(rgb, axis=0)
-
-            # if self.config.do_convert_grayscale and rgb.ndim == 3:
-            #     rgb = rgb.unsqueeze(1)
-
-            # channel = rgb.shape[1]
-
-            # height, width = self.get_default_height_width(rgb, height, width)
-            # if self.config.do_resize:
-            #     rgb = self.resize(rgb, height, width)
-
-            # depth = torch.cat(depth, axis=0) if depth[0].ndim == 4 else torch.stack(depth, axis=0)
-
-            # if self.config.do_convert_grayscale and depth.ndim == 3:
-            #     depth = depth.unsqueeze(1)
-
-            # channel = depth.shape[1]
-            # # don't need any preprocess if the image is latents
-            # if depth == 4:
-            #     return rgb, depth
-
-            # height, width = self.get_default_height_width(depth, height, width)
-            # if self.config.do_resize:
-            #     depth = self.resize(depth, height, width)
-        # expected range [0,1], normalize to [-1,1]
-        do_normalize = self.config.do_normalize
-        if rgb.min() < 0 and do_normalize:
-            warnings.warn(
-                "Passing `image` as torch tensor with value range in [-1,1] is deprecated. The expected value range for image tensor is [0,1] "
-                f"when passing as pytorch tensor or numpy Array. You passed `image` with value range [{rgb.min()},{rgb.max()}]",
-                FutureWarning,
-            )
-            do_normalize = False
-
-        if do_normalize:
-            rgb = self.normalize(rgb)
-            depth = self.normalize(depth)
-
-        if self.config.do_binarize:
-            rgb = self.binarize(rgb)
-            depth = self.binarize(depth)
-
-        return rgb, depth
-
-
-class IPAdapterMaskProcessor(VaeImageProcessor):
-    """
-    Image processor for IP Adapter image masks.
-
-    Args:
-        do_resize (`bool`, *optional*, defaults to `True`):
-            Whether to downscale the image's (height, width) dimensions to multiples of `vae_scale_factor`.
-        vae_scale_factor (`int`, *optional*, defaults to `8`):
-            VAE scale factor. If `do_resize` is `True`, the image is automatically resized to multiples of this factor.
-        resample (`str`, *optional*, defaults to `lanczos`):
-            Resampling filter to use when resizing the image.
-        do_normalize (`bool`, *optional*, defaults to `False`):
-            Whether to normalize the image to [-1,1].
-        do_binarize (`bool`, *optional*, defaults to `True`):
-            Whether to binarize the image to 0/1.
-        do_convert_grayscale (`bool`, *optional*, defaults to be `True`):
-            Whether to convert the images to grayscale format.
-
-    """
-
-    config_name = CONFIG_NAME
-
-    @register_to_config
-    def __init__(
-        self,
-        do_resize: bool = True,
-        vae_scale_factor: int = 8,
-        resample: str = "lanczos",
-        do_normalize: bool = False,
-        do_binarize: bool = True,
-        do_convert_grayscale: bool = True,
-    ):
-        super().__init__(
-            do_resize=do_resize,
-            vae_scale_factor=vae_scale_factor,
-            resample=resample,
-            do_normalize=do_normalize,
-            do_binarize=do_binarize,
-            do_convert_grayscale=do_convert_grayscale,
-        )
-
-    @staticmethod
-    def downsample(mask: torch.Tensor, batch_size: int, num_queries: int, value_embed_dim: int):
-        """
-        Downsamples the provided mask tensor to match the expected dimensions for scaled dot-product attention. If the
-        aspect ratio of the mask does not match the aspect ratio of the output image, a warning is issued.
-
-        Args:
-            mask (`torch.Tensor`):
-                The input mask tensor generated with `IPAdapterMaskProcessor.preprocess()`.
-            batch_size (`int`):
-                The batch size.
-            num_queries (`int`):
-                The number of queries.
-            value_embed_dim (`int`):
-                The dimensionality of the value embeddings.
-
-        Returns:
-            `torch.Tensor`:
-                The downsampled mask tensor.
-
-        """
-        o_h = mask.shape[1]
-        o_w = mask.shape[2]
-        ratio = o_w / o_h
-        mask_h = int(math.sqrt(num_queries / ratio))
-        mask_h = int(mask_h) + int((num_queries % int(mask_h)) != 0)
-        mask_w = num_queries // mask_h
-
-        mask_downsample = F.interpolate(mask.unsqueeze(0), size=(mask_h, mask_w), mode="bicubic").squeeze(0)
-
-        # Repeat batch_size times
-        if mask_downsample.shape[0] < batch_size:
-            mask_downsample = mask_downsample.repeat(batch_size, 1, 1)
-
-        mask_downsample = mask_downsample.view(mask_downsample.shape[0], -1)
-
-        downsampled_area = mask_h * mask_w
-        # If the output image and the mask do not have the same aspect ratio, tensor shapes will not match
-        # Pad tensor if downsampled_mask.shape[1] is smaller than num_queries
-        if downsampled_area < num_queries:
-            warnings.warn(
-                "The aspect ratio of the mask does not match the aspect ratio of the output image. "
-                "Please update your masks or adjust the output size for optimal performance.",
-                UserWarning,
-            )
-            mask_downsample = F.pad(mask_downsample, (0, num_queries - mask_downsample.shape[1]), value=0.0)
-        # Discard last embeddings if downsampled_mask.shape[1] is bigger than num_queries
-        if downsampled_area > num_queries:
-            warnings.warn(
-                "The aspect ratio of the mask does not match the aspect ratio of the output image. "
-                "Please update your masks or adjust the output size for optimal performance.",
-                UserWarning,
-            )
-            mask_downsample = mask_downsample[:, :num_queries]
-
-        # Repeat last dimension to match SDPA output shape
-        mask_downsample = mask_downsample.view(mask_downsample.shape[0], mask_downsample.shape[1], 1).repeat(
-            1, 1, value_embed_dim
-        )
-
-        return mask_downsample
-
-
-class PixArtImageProcessor(VaeImageProcessor):
-    """
-    Image processor for PixArt image resize and crop.
-
-    Args:
-        do_resize (`bool`, *optional*, defaults to `True`):
-            Whether to downscale the image's (height, width) dimensions to multiples of `vae_scale_factor`. Can accept
-            `height` and `width` arguments from [`image_processor.VaeImageProcessor.preprocess`] method.
-        vae_scale_factor (`int`, *optional*, defaults to `8`):
-            VAE scale factor. If `do_resize` is `True`, the image is automatically resized to multiples of this factor.
-        resample (`str`, *optional*, defaults to `lanczos`):
-            Resampling filter to use when resizing the image.
-        do_normalize (`bool`, *optional*, defaults to `True`):
-            Whether to normalize the image to [-1,1].
-        do_binarize (`bool`, *optional*, defaults to `False`):
-            Whether to binarize the image to 0/1.
-        do_convert_rgb (`bool`, *optional*, defaults to be `False`):
-            Whether to convert the images to RGB format.
-        do_convert_grayscale (`bool`, *optional*, defaults to be `False`):
-            Whether to convert the images to grayscale format.
-    """
-
-    @register_to_config
-    def __init__(
-        self,
-        do_resize: bool = True,
-        vae_scale_factor: int = 8,
-        resample: str = "lanczos",
-        do_normalize: bool = True,
-        do_binarize: bool = False,
-        do_convert_grayscale: bool = False,
-    ):
-        super().__init__(
-            do_resize=do_resize,
-            vae_scale_factor=vae_scale_factor,
-            resample=resample,
-            do_normalize=do_normalize,
-            do_binarize=do_binarize,
-            do_convert_grayscale=do_convert_grayscale,
-        )
-
-    @staticmethod
-    def classify_height_width_bin(height: int, width: int, ratios: dict) -> tuple[int, int]:
-        r"""
-        Returns the binned height and width based on the aspect ratio.
-
-        Args:
-            height (`int`): The height of the image.
-            width (`int`): The width of the image.
-            ratios (`dict`): A dictionary where keys are aspect ratios and values are tuples of (height, width).
-
-        Returns:
-            `tuple[int, int]`: The closest binned height and width.
-        """
-        ar = float(height / width)
-        closest_ratio = min(ratios.keys(), key=lambda ratio: abs(float(ratio) - ar))
-        default_hw = ratios[closest_ratio]
-        return int(default_hw[0]), int(default_hw[1])
-
-    @staticmethod
-    def resize_and_crop_tensor(samples: torch.Tensor, new_width: int, new_height: int) -> torch.Tensor:
-        r"""
-        Resizes and crops a tensor of images to the specified dimensions.
-
-        Args:
-            samples (`torch.Tensor`):
-                A tensor of shape (N, C, H, W) where N is the batch size, C is the number of channels, H is the height,
-                and W is the width.
-            new_width (`int`): The desired width of the output images.
-            new_height (`int`): The desired height of the output images.
-
-        Returns:
-            `torch.Tensor`: A tensor containing the resized and cropped images.
-        """
-        orig_height, orig_width = samples.shape[2], samples.shape[3]
-
-        # Check if resizing is needed
-        if orig_height != new_height or orig_width != new_width:
-            ratio = max(new_height / orig_height, new_width / orig_width)
-            resized_width = int(orig_width * ratio)
-            resized_height = int(orig_height * ratio)
-
-            # Resize
-            samples = F.interpolate(
-                samples, size=(resized_height, resized_width), mode="bilinear", align_corners=False
-            )
-
-            # Center Crop
-            start_x = (resized_width - new_width) // 2
-            end_x = start_x + new_width
-            start_y = (resized_height - new_height) // 2
-            end_y = start_y + new_height
-            samples = samples[:, :, start_y:end_y, start_x:end_x]
-
-        return samples

diffusers/loaders/__init__.py

@@ -1,159 +0,0 @@
-from typing import TYPE_CHECKING
-
-from ..utils import DIFFUSERS_SLOW_IMPORT, _LazyModule, deprecate
-from ..utils.import_utils import is_peft_available, is_torch_available, is_transformers_available
-
-
-def text_encoder_lora_state_dict(text_encoder):
-    deprecate(
-        "text_encoder_load_state_dict in `models`",
-        "0.27.0",
-        "`text_encoder_lora_state_dict` is deprecated and will be removed in 0.27.0. Make sure to retrieve the weights using `get_peft_model`. See https://huggingface.co/docs/peft/v0.6.2/en/quicktour#peftmodel for more information.",
-    )
-    state_dict = {}
-
-    for name, module in text_encoder_attn_modules(text_encoder):
-        for k, v in module.q_proj.lora_linear_layer.state_dict().items():
-            state_dict[f"{name}.q_proj.lora_linear_layer.{k}"] = v
-
-        for k, v in module.k_proj.lora_linear_layer.state_dict().items():
-            state_dict[f"{name}.k_proj.lora_linear_layer.{k}"] = v
-
-        for k, v in module.v_proj.lora_linear_layer.state_dict().items():
-            state_dict[f"{name}.v_proj.lora_linear_layer.{k}"] = v
-
-        for k, v in module.out_proj.lora_linear_layer.state_dict().items():
-            state_dict[f"{name}.out_proj.lora_linear_layer.{k}"] = v
-
-    return state_dict
-
-
-if is_transformers_available():
-
-    def text_encoder_attn_modules(text_encoder):
-        deprecate(
-            "text_encoder_attn_modules in `models`",
-            "0.27.0",
-            "`text_encoder_lora_state_dict` is deprecated and will be removed in 0.27.0. Make sure to retrieve the weights using `get_peft_model`. See https://huggingface.co/docs/peft/v0.6.2/en/quicktour#peftmodel for more information.",
-        )
-        from transformers import CLIPTextModel, CLIPTextModelWithProjection
-
-        attn_modules = []
-
-        if isinstance(text_encoder, (CLIPTextModel, CLIPTextModelWithProjection)):
-            for i, layer in enumerate(text_encoder.text_model.encoder.layers):
-                name = f"text_model.encoder.layers.{i}.self_attn"
-                mod = layer.self_attn
-                attn_modules.append((name, mod))
-        else:
-            raise ValueError(f"do not know how to get attention modules for: {text_encoder.__class__.__name__}")
-
-        return attn_modules
-
-
-_import_structure = {}
-
-if is_torch_available():
-    _import_structure["single_file_model"] = ["FromOriginalModelMixin"]
-    _import_structure["transformer_flux"] = ["FluxTransformer2DLoadersMixin"]
-    _import_structure["transformer_sd3"] = ["SD3Transformer2DLoadersMixin"]
-    _import_structure["unet"] = ["UNet2DConditionLoadersMixin"]
-    _import_structure["utils"] = ["AttnProcsLayers"]
-    if is_transformers_available():
-        _import_structure["single_file"] = ["FromSingleFileMixin"]
-        _import_structure["lora_pipeline"] = [
-            "AceStepLoraLoaderMixin",
-            "AmusedLoraLoaderMixin",
-            "AnimaLoraLoaderMixin",
-            "StableDiffusionLoraLoaderMixin",
-            "SD3LoraLoaderMixin",
-            "AuraFlowLoraLoaderMixin",
-            "StableDiffusionXLLoraLoaderMixin",
-            "LTX2LoraLoaderMixin",
-            "LTXVideoLoraLoaderMixin",
-            "LoraLoaderMixin",
-            "FluxLoraLoaderMixin",
-            "CogVideoXLoraLoaderMixin",
-            "CogView4LoraLoaderMixin",
-            "Mochi1LoraLoaderMixin",
-            "HunyuanVideoLoraLoaderMixin",
-            "SanaLoraLoaderMixin",
-            "Lumina2LoraLoaderMixin",
-            "WanLoraLoaderMixin",
-            "HeliosLoraLoaderMixin",
-            "KandinskyLoraLoaderMixin",
-            "HiDreamImageLoraLoaderMixin",
-            "SkyReelsV2LoraLoaderMixin",
-            "QwenImageLoraLoaderMixin",
-            "Krea2LoraLoaderMixin",
-            "ZImageLoraLoaderMixin",
-            "Flux2LoraLoaderMixin",
-            "Ideogram4LoraLoaderMixin",
-            "ErnieImageLoraLoaderMixin",
-            "CosmosLoraLoaderMixin",
-        ]
-        _import_structure["textual_inversion"] = ["TextualInversionLoaderMixin"]
-        _import_structure["ip_adapter"] = [
-            "IPAdapterMixin",
-            "FluxIPAdapterMixin",
-            "SD3IPAdapterMixin",
-            "ModularIPAdapterMixin",
-        ]
-
-_import_structure["peft"] = ["PeftAdapterMixin"]
-
-
-if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT:
-    if is_torch_available():
-        from .single_file_model import FromOriginalModelMixin
-        from .transformer_flux import FluxTransformer2DLoadersMixin
-        from .transformer_sd3 import SD3Transformer2DLoadersMixin
-        from .unet import UNet2DConditionLoadersMixin
-        from .utils import AttnProcsLayers
-
-        if is_transformers_available():
-            from .ip_adapter import (
-                FluxIPAdapterMixin,
-                IPAdapterMixin,
-                ModularIPAdapterMixin,
-                SD3IPAdapterMixin,
-            )
-            from .lora_pipeline import (
-                AceStepLoraLoaderMixin,
-                AmusedLoraLoaderMixin,
-                AnimaLoraLoaderMixin,
-                AuraFlowLoraLoaderMixin,
-                CogVideoXLoraLoaderMixin,
-                CogView4LoraLoaderMixin,
-                CosmosLoraLoaderMixin,
-                ErnieImageLoraLoaderMixin,
-                Flux2LoraLoaderMixin,
-                FluxLoraLoaderMixin,
-                HeliosLoraLoaderMixin,
-                HiDreamImageLoraLoaderMixin,
-                HunyuanVideoLoraLoaderMixin,
-                Ideogram4LoraLoaderMixin,
-                KandinskyLoraLoaderMixin,
-                Krea2LoraLoaderMixin,
-                LoraLoaderMixin,
-                LTX2LoraLoaderMixin,
-                LTXVideoLoraLoaderMixin,
-                Lumina2LoraLoaderMixin,
-                Mochi1LoraLoaderMixin,
-                QwenImageLoraLoaderMixin,
-                SanaLoraLoaderMixin,
-                SD3LoraLoaderMixin,
-                SkyReelsV2LoraLoaderMixin,
-                StableDiffusionLoraLoaderMixin,
-                StableDiffusionXLLoraLoaderMixin,
-                WanLoraLoaderMixin,
-                ZImageLoraLoaderMixin,
-            )
-            from .single_file import FromSingleFileMixin
-            from .textual_inversion import TextualInversionLoaderMixin
-
-    from .peft import PeftAdapterMixin
-else:
-    import sys
-
-    sys.modules[__name__] = _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)