Canvas: switch to step-distilled model + AOTI (PR#29 revendor, aokit.load + eager fallback, 16 steps no guider)

README.md

@@ -14,7 +14,7 @@ short_description: FLUX.3 image-small bounding-box layout canvas (ZeroGPU)
 # FLUX.3 image-small — layout-control canvas (private)
 
 A custom-HTML canvas on top of the **diffusers-native FLUX.3 modular pipeline**
-([`diffusers-new-model-addition-f3` PR #23](https://github.com/huggingface/diffusers-new-model-addition-f3/pull/23)),
+([`diffusers-new-model-addition-f3` PR #29](https://github.com/huggingface/diffusers-new-model-addition-f3/pull/29)), running the **step-distilled** model AOT-compiled with aokit,
 on ZeroGPU.
 
 Drag a bounding box for every element that matters, give each box an id and a
@@ -23,11 +23,11 @@ every element placed inside its box. Toggle the **Bounding boxes** overlay on th
 result, or **Replay** to regenerate.
 
 > **Private.** Weights are read at runtime from the private repo
-> `diffusers-internal-dev/f3-small-img` (never copied here). Keep this Space private.
+> `diffusers-internal-dev/f3-small-image-distilled` (never copied here). Keep this Space private.
 
 Layout control uses FLUX.3's region format: the final prompt is the scene caption
 followed by a JSON list of regions, each `{"id", "bbox": [y0, x0, y1, x1] (0-1000), "desc"}`.
 
 ## Setup
 
-Needs an `HF_TOKEN` secret with read access to `diffusers-internal-dev/f3-small-img`.
+Needs an `HF_TOKEN` secret with read access to `diffusers-internal-dev/f3-small-image-distilled` and the private AOTI artifact `diffusers-modular/f3-small-image-distilled-aoti-sm120`.

app.py

@@ -1,9 +1,14 @@
 """FLUX.3 image-small — layout-control canvas (private ZeroGPU demo, diffusers modular).
 
 A custom HTML/JS canvas on top of the diffusers-native FLUX.3 modular pipeline
-(PR #23, vendored `diffusers/`). You drag bounding boxes on an aspect-ratio
-artboard, give each box an id + a description, write a one-line scene prompt, and
-FLUX.3 renders the scene with every element placed inside its box.
+(PR #29, vendored `diffusers/`), running the **step-distilled** flow-map model
+(`Flux3DistilledAutoBlocks`, 16 steps, no guider / negative prompt), with its
+transformer block stacks **AOT-compiled via aokit** for ZeroGPU (`aokit.load` +
+a per-block eager fallback; set `F3_AOTI=0` to run eager).
+
+You drag bounding boxes on an aspect-ratio artboard, give each box an id + a
+description, write a one-line scene prompt, and FLUX.3 renders the scene with
+every element placed inside its box.
 
 Layout control uses FLUX.3's region format (see `integration-notes/flux3_examples.py`):
 the final prompt is the scene caption followed by a JSON list of regions, each
@@ -23,18 +28,22 @@ import base64
 import io
 import json
 import random
+from pathlib import Path
 
 import gradio as gr
 import torch
+from huggingface_hub import hf_hub_download
 
 from diffusers import ModularPipeline
 
-REPO = os.environ.get("F3_REPO", "diffusers-internal-dev/f3-small-img")
+# Step-distilled FLUX.3 image model + its private AOTI artifact.
+REPO = os.environ.get("F3_REPO", "diffusers-internal-dev/f3-small-image-distilled")
+AOTI_REPO = os.environ.get("F3_AOTI_REPO", "diffusers-modular/f3-small-image-distilled-aoti-sm120")
+AOTI_ENABLED = os.environ.get("F3_AOTI", "1") != "0"
 HF_TOKEN = os.environ.get("HF_TOKEN")
 
 MAX_SEED = 2**63 - 1
-DEFAULT_STEPS = int(os.environ.get("F3_STEPS", "50"))
-DEFAULT_GUIDANCE = 3.5
+DEFAULT_STEPS = int(os.environ.get("F3_STEPS", "16"))  # distilled default
 
 torch.backends.cuda.matmul.allow_tf32 = True
 torch.backends.cudnn.allow_tf32 = True
@@ -50,14 +59,66 @@ RATIOS = {
     "9:16": (768, 1344),
 }
 
-# ---- load the pipeline on CPU at import; place on cuda under ZeroGPU ----
-print(f"[app] loading modular pipeline from {REPO} ...", flush=True)
+
+def _eager_fallback(compiled, eager):
+    """Call the AOTI graph; fall back to the eager forward on any shape it can't
+    handle (keeps odd resolutions / layouts working instead of crashing)."""
+    def forward(*args, **kwargs):
+        try:
+            return compiled(*args, **kwargs)
+        except Exception:
+            return eager(*args, **kwargs)
+    return forward
+
+
+def _load_aoti(transformer, repo, token):
+    """aokit.load the compiled stacks, then wrap each block with an eager fallback.
+    The artifact's torch version is asserted against the runtime."""
+    import aokit
+
+    cfg = json.loads(Path(hf_hub_download(repo, "config.json", token=token)).read_text())
+    if cfg.get("torch") != torch.__version__:
+        raise RuntimeError(
+            f"AOTI artifact built for torch {cfg.get('torch')}, runtime has {torch.__version__}"
+        )
+    stack_names = list(cfg.get("stacks", {})) or [
+        "mode_blocks.image", "mode_blocks.image_cond", "txt_mode_blocks", "single_transformer_blocks",
+    ]
+
+    def _blocks(name):
+        try:
+            stack = transformer.get_submodule(name)
+        except AttributeError:
+            return []
+        return list(stack) if isinstance(stack, torch.nn.ModuleList) else [stack]
+
+    eager = {name: [blk.forward for blk in _blocks(name)] for name in stack_names}
+    aokit.load(transformer, repo)  # patches block.forward -> compiled (weights stay live)
+    wrapped = 0
+    for name in stack_names:
+        for blk, eager_fwd in zip(_blocks(name), eager[name]):
+            blk.forward = _eager_fallback(blk.forward, eager_fwd)
+            wrapped += 1
+    print(f"[app] AOTI loaded from {repo}: {len(stack_names)} stacks, {wrapped} blocks wrapped "
+          f"(torch {cfg.get('torch')}, {cfg.get('arch')})", flush=True)
+    return cfg
+
+
+# ---- load the pipeline on CPU at import; place on cuda + AOTI under ZeroGPU ----
+print(f"[app] loading distilled modular pipeline from {REPO} ...", flush=True)
 PIPE = ModularPipeline.from_pretrained(REPO, token=HF_TOKEN)
 PIPE.load_components(dtype={"vae": torch.float32, "default": torch.bfloat16}, token=HF_TOKEN)
 print("[app] components loaded on CPU.", flush=True)
+AOTI_LOADED = False
 if spaces is not None:
     PIPE.to("cuda")
-    print("[app] pipeline on cuda; ready.", flush=True)
+    if AOTI_ENABLED:
+        try:
+            _load_aoti(PIPE.transformer, AOTI_REPO, HF_TOKEN)
+            AOTI_LOADED = True
+        except Exception as exc:  # fall back to eager if the artifact is unavailable
+            print(f"[app] AOTI unavailable, running eager: {exc!r}", flush=True)
+    print(f"[app] pipeline on cuda; ready (aoti={AOTI_LOADED}).", flush=True)
 
 
 def _build_prompt(scene: str, elements: list) -> str:
@@ -88,7 +149,8 @@ def _duration(state_json, *_):
     except Exception:
         w, h = 1024, 1024
     px = w * h
-    return int(min(300, 45 + DEFAULT_STEPS * 2.6 * max(1.0, px / (1024 * 1024))))
+    # Distilled flow-map model: single forward per step (no CFG), so cheaper/step.
+    return int(min(300, 40 + DEFAULT_STEPS * 1.6 * max(1.0, px / (1024 * 1024))))
 
 
 def _generate(state_json):
@@ -109,7 +171,7 @@ def _generate(state_json):
     seed = st.get("seed")
     seed = int(seed) if isinstance(seed, (int, float)) and seed >= 0 else random.randint(0, MAX_SEED)
 
-    PIPE.guider = PIPE.guider.new(guidance_scale=DEFAULT_GUIDANCE)
+    # Distilled flow-map model has no guider / guidance.
     image = PIPE(
         prompt=full_prompt,
         height=h,

diffusers/__init__.py

@@ -527,6 +527,8 @@ else:
             "Flux2KleinModularPipeline",
             "Flux2ModularPipeline",
             "Flux3AutoBlocks",
+            "Flux3DistilledAutoBlocks",
+            "Flux3DistilledModularPipeline",
             "Flux3ModularPipeline",
             "FluxAutoBlocks",
             "FluxKontextAutoBlocks",
@@ -1386,6 +1388,8 @@ if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT:
             Flux2KleinModularPipeline,
             Flux2ModularPipeline,
             Flux3AutoBlocks,
+            Flux3DistilledAutoBlocks,
+            Flux3DistilledModularPipeline,
             Flux3ModularPipeline,
             FluxAutoBlocks,
             FluxKontextAutoBlocks,

diffusers/experimental/README.md

@@ -0,0 +1,5 @@
+# 🧨 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/models/README.md

@@ -0,0 +1,3 @@
+# Models
+
+For more detail on the models, please refer to the [docs](https://huggingface.co/docs/diffusers/api/models/overview).

diffusers/models/transformers/transformer_flux3.py

@@ -346,7 +346,9 @@ class Flux3Transformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, Attenti
     passes through its own per-mode transformer blocks (`mode_blocks` for the image modes, `txt_mode_blocks` for the
     text context), then the text context and the active image modes are concatenated for the joint single blocks, and
     only the image tokens pass through the output layers. The modulation is per mode: the image tokens are modulated by
-    the flow timestep, the reference image tokens and the text context by timestep `0`.
+    the flow timestep, the reference image tokens and the text context by timestep `0`. A flow-map checkpoint
+    (`flow_map=True`, the step-distilled models) additionally embeds the timestep each denoising step jumps to and adds
+    it to every mode's modulation.
 
     Args:
         in_channels (`int`, defaults to `128`):
@@ -373,6 +375,9 @@ class Flux3Transformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, Attenti
             The base frequency of the rotary positional embeddings.
         eps (`float`, defaults to `1e-6`):
             The epsilon of the layer norms and the QK norms.
+        flow_map (`bool`, defaults to `False`):
+            Whether the model is a flow map (step-distilled): adds the `target_timestep_embedder` and requires
+            `target_timestep` in the forward pass.
     """
 
     _supports_gradient_checkpointing = True
@@ -395,6 +400,7 @@ class Flux3Transformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, Attenti
         axes_dims_rope: tuple[int, ...] = (32, 32, 32, 32),
         rope_theta: int = 10000,
         eps: float = 1e-6,
+        flow_map: bool = False,
     ):
         super().__init__()
 
@@ -410,6 +416,12 @@ class Flux3Transformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, Attenti
         self.timestep_embedder = TimestepEmbedding(
             in_channels=timestep_channels, time_embed_dim=self.inner_dim, sample_proj_bias=False
         )
+        # flow map (step-distilled checkpoints): a second embedder for the timestep each step jumps to
+        self.target_timestep_embedder = (
+            TimestepEmbedding(in_channels=timestep_channels, time_embed_dim=self.inner_dim, sample_proj_bias=False)
+            if flow_map
+            else None
+        )
 
         # 3. Modulation: one head per mode (+ text) for the per-mode blocks and for the joint blocks
         self.early_stream_modulations = nn.ModuleDict(
@@ -486,6 +498,7 @@ class Flux3Transformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, Attenti
         txt_ids: torch.Tensor,
         image_cond_hidden_states: torch.Tensor | None = None,
         image_cond_ids: torch.Tensor | None = None,
+        target_timestep: torch.Tensor | None = None,
         attention_kwargs: dict[str, Any] | None = None,
         return_dict: bool = True,
     ) -> Transformer2DModelOutput | tuple[torch.Tensor]:
@@ -509,6 +522,9 @@ class Flux3Transformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, Attenti
                 VAE-encoded reference image tokens (clean conditioning content, for image-to-image generation).
             image_cond_ids (`torch.Tensor` of shape `(cond_seq_len, 4)`, *optional*):
                 Position ids of the reference image tokens. Required with `image_cond_hidden_states`.
+            target_timestep (`torch.Tensor` of shape `(batch_size,)`, *optional*):
+                Flow timestep in `[0, 1]` that the current denoising step jumps to. Required by, and only by, a flow-map
+                model (`flow_map=True`).
             attention_kwargs (`dict`, *optional*):
                 A kwargs dictionary whose `scale` entry, if specified, scales the LoRA layers for this forward pass.
             return_dict (`bool`, *optional*, defaults to `True`):
@@ -524,6 +540,10 @@ class Flux3Transformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, Attenti
             raise ValueError("`image_cond_hidden_states` requires `image_cond_ids`.")
         if image_cond_hidden_states is None and image_cond_ids is not None:
             raise ValueError("`image_cond_ids` was passed without `image_cond_hidden_states`.")
+        if (target_timestep is not None) != bool(self.config.flow_map):
+            raise ValueError(
+                "`target_timestep` is required by, and only by, a flow-map transformer (`flow_map=True`)."
+            )
 
         batch_size, image_seq_len = hidden_states.shape[:2]
         dtype = hidden_states.dtype
@@ -531,9 +551,17 @@ class Flux3Transformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, Attenti
         # 1. Timestep embeddings: the image tokens at `timestep`, the text context and the reference tokens at 0.
         #    Each is a per-token embedding (batch_size, seq_len, dim) since the modulation is per token.
         image_temb = self.time_proj(timestep.to(dtype) * 1000)
-        image_temb = self.timestep_embedder(image_temb.to(dtype))[:, None, :].expand(batch_size, image_seq_len, -1)
+        image_temb = self.timestep_embedder(image_temb.to(dtype))[:, None, :]
         zero_temb = self.time_proj(torch.zeros(batch_size, device=hidden_states.device, dtype=dtype))
         zero_temb = self.timestep_embedder(zero_temb.to(dtype))[:, None, :]
+        if target_timestep is not None:
+            # flow map: the embedding of the timestep this step jumps to is added to every mode's embedding, so it
+            # reaches the per-mode blocks, the joint blocks and the output layer
+            target_temb = self.time_proj(target_timestep.to(dtype) * 1000)
+            target_temb = self.target_timestep_embedder(target_temb.to(dtype))[:, None, :]
+            image_temb = image_temb + target_temb
+            zero_temb = zero_temb + target_temb
+        image_temb = image_temb.expand(batch_size, image_seq_len, -1)
         zero_temb_act = F.silu(zero_temb)
 
         # 2. Text context: the two modulation sets, input projection and RoPE

diffusers/modular_pipelines/__init__.py

@@ -83,6 +83,8 @@ else:
     ]
     _import_structure["flux3"] = [
         "Flux3AutoBlocks",
+        "Flux3DistilledAutoBlocks",
+        "Flux3DistilledModularPipeline",
         "Flux3ModularPipeline",
     ]
     _import_structure["ideogram4"] = [
@@ -172,7 +174,12 @@ if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT:
             Flux2KleinModularPipeline,
             Flux2ModularPipeline,
         )
-        from .flux3 import Flux3AutoBlocks, Flux3ModularPipeline
+        from .flux3 import (
+            Flux3AutoBlocks,
+            Flux3DistilledAutoBlocks,
+            Flux3DistilledModularPipeline,
+            Flux3ModularPipeline,
+        )
         from .helios import (
             HeliosAutoBlocks,
             HeliosModularPipeline,

diffusers/modular_pipelines/flux3/__init__.py

@@ -22,7 +22,8 @@ except OptionalDependencyNotAvailable:
     _dummy_objects.update(get_objects_from_module(dummy_torch_and_transformers_objects))
 else:
     _import_structure["modular_blocks_flux3"] = ["Flux3AutoBlocks"]
-    _import_structure["modular_pipeline"] = ["Flux3ModularPipeline"]
+    _import_structure["modular_blocks_flux3_distilled"] = ["Flux3DistilledAutoBlocks"]
+    _import_structure["modular_pipeline"] = ["Flux3DistilledModularPipeline", "Flux3ModularPipeline"]
 
 if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT:
     try:
@@ -32,7 +33,8 @@ if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT:
         from ...utils.dummy_torch_and_transformers_objects import *  # noqa F403
     else:
         from .modular_blocks_flux3 import Flux3AutoBlocks
-        from .modular_pipeline import Flux3ModularPipeline
+        from .modular_blocks_flux3_distilled import Flux3DistilledAutoBlocks
+        from .modular_pipeline import Flux3DistilledModularPipeline, Flux3ModularPipeline
 else:
     import sys
 

diffusers/modular_pipelines/flux3/before_denoise.py

@@ -92,9 +92,11 @@ def retrieve_timesteps(
 class Flux3SetTimestepsStep(ModularPipelineBlocks):
     model_name = "flux3"
 
-    def __init__(self, shift: float = 6.93):
-        # the constant timestep shift ("alpha") of the FLUX.3 sampling recipe, configured on the scheduler
+    def __init__(self, shift: float = 6.93, num_inference_steps: int = 50):
+        # the constant timestep shift ("alpha") of the FLUX.3 sampling recipe, configured on the scheduler, and the
+        # default number of steps (50; 16 for the step-distilled checkpoints)
         self.shift = shift
+        self.num_inference_steps = num_inference_steps
         super().__init__()
 
     @property
@@ -118,7 +120,7 @@ class Flux3SetTimestepsStep(ModularPipelineBlocks):
     @property
     def inputs(self) -> list[InputParam]:
         return [
-            InputParam.template("num_inference_steps"),
+            InputParam.template("num_inference_steps", default=self.num_inference_steps),
             InputParam.template("timesteps", note="overrides `num_inference_steps` when provided"),
             InputParam.template("sigmas", note="overrides `num_inference_steps` when provided"),
         ]
@@ -395,3 +397,40 @@ class Flux3TextRoPEInputsStep(ModularPipelineBlocks):
 
         self.set_block_state(state, block_state)
         return components, state
+
+
+class Flux3DistilledTextRoPEInputsStep(Flux3TextRoPEInputsStep):
+    model_name = "flux3-distilled"
+
+    @property
+    def description(self) -> str:
+        return (
+            "Step that prepares the 4D RoPE position IDs of the text tokens for the step-distilled FLUX.3 model (no "
+            "negative prompt). Place it after the text encoder step."
+        )
+
+    @property
+    def inputs(self) -> list[InputParam]:
+        return [
+            InputParam.template(
+                "prompt_embeds",
+                description="Text embeddings used to guide the generation. Can be generated from text_encoder step.",
+            ),
+        ]
+
+    @property
+    def intermediate_outputs(self) -> list[OutputParam]:
+        return [
+            OutputParam(
+                name="txt_ids",
+                kwargs_type="denoiser_input_fields",
+                type_hint=torch.Tensor,
+                description="4D position IDs (T, H, W, L) for text tokens, used for RoPE calculation.",
+            ),
+        ]
+
+    def __call__(self, components: Flux3ModularPipeline, state: PipelineState) -> PipelineState:
+        block_state = self.get_block_state(state)
+        block_state.txt_ids = self._prepare_text_ids(block_state.prompt_embeds).to(block_state.prompt_embeds.device)
+        self.set_block_state(state, block_state)
+        return components, state

diffusers/modular_pipelines/flux3/denoise.py

@@ -138,6 +138,76 @@ class Flux3LoopDenoiser(ModularPipelineBlocks):
         return components, block_state
 
 
+class Flux3DistilledLoopDenoiser(ModularPipelineBlocks):
+    model_name = "flux3-distilled"
+
+    @property
+    def expected_components(self) -> list[ComponentSpec]:
+        return [ComponentSpec("transformer", Flux3Transformer2DModel)]
+
+    @property
+    def description(self) -> str:
+        return (
+            "Step within the denoising loop that predicts the velocity of the image tokens with a step-distilled "
+            "(flow-map) FLUX.3 transformer: no classifier-free guidance, and each step also passes the transformer the "
+            "timestep it jumps to. The static conditioning fields (`prompt_embeds`, `txt_ids`, `img_ids`, the "
+            "reference image tokens and their ids) arrive via `denoiser_input_fields`. "
+            "This block should be used to compose the `sub_blocks` attribute of a `LoopSequentialPipelineBlocks` "
+            "object (e.g. `Flux3DistilledDenoiseStep`)"
+        )
+
+    @property
+    def inputs(self) -> list[tuple[str, Any]]:
+        return [
+            InputParam(
+                "latents",
+                required=True,
+                type_hint=torch.Tensor,
+                description="The image latent tokens to denoise. Shape: (B, seq_len, C)",
+            ),
+            InputParam.template("denoiser_input_fields"),
+        ]
+
+    @property
+    def intermediate_outputs(self) -> list[OutputParam]:
+        return [OutputParam("noise_pred", type_hint=torch.Tensor, description="The predicted velocity for this step")]
+
+    @torch.no_grad()
+    def __call__(
+        self, components: Flux3ModularPipeline, block_state: BlockState, i: int, t: torch.Tensor
+    ) -> PipelineState:
+        denoiser_input_fields = block_state.denoiser_input_fields
+
+        # transformer kwargs = every collected static field matching the forward signature (`img_ids`, `txt_ids`, and
+        # the `image_cond_*` fields of image-to-image generation)
+        transformer_args = set(inspect.signature(components.transformer.forward).parameters.keys())
+        transformer_kwargs = {
+            field_name: field_value
+            for field_name, field_value in denoiser_input_fields.items()
+            if field_name in transformer_args
+        }
+
+        latents = block_state.latents.to(components.transformer.dtype)
+        batch_size = latents.shape[0]
+        # the flow timestep in [0, 1] (scheduler timesteps are sigmas * 1000), and the one this step jumps to: the next
+        # scheduler timestep, 0 after the last
+        timesteps = block_state.timesteps
+        next_t = timesteps[i + 1] if i + 1 < len(timesteps) else torch.zeros_like(t)
+        timestep = (t / 1000).expand(batch_size).to(latents.dtype)
+        target_timestep = (next_t / 1000).expand(batch_size).to(latents.dtype)
+
+        block_state.noise_pred = components.transformer(
+            hidden_states=latents,
+            encoder_hidden_states=denoiser_input_fields["prompt_embeds"],
+            timestep=timestep,
+            target_timestep=target_timestep,
+            return_dict=False,
+            **transformer_kwargs,
+        )[0]
+
+        return components, block_state
+
+
 class Flux3LoopAfterDenoiser(ModularPipelineBlocks):
     model_name = "flux3"
 
@@ -276,3 +346,73 @@ class Flux3DenoiseStep(Flux3DenoiseLoopWrapper):
             " - `Flux3LoopAfterDenoiser`\n"
             "This block supports text-to-image and image-to-image generation."
         )
+
+
+# auto_docstring
+class Flux3DistilledDenoiseStep(Flux3DenoiseLoopWrapper):
+    """
+    Denoise step that iteratively denoises the latents for the step-distilled (flow-map) FLUX.3 model.
+      Its loop logic is defined in `Flux3DenoiseLoopWrapper.__call__` method At each iteration, it runs blocks defined
+      in `sub_blocks` sequentially:
+       - `Flux3DistilledLoopDenoiser`
+       - `Flux3LoopAfterDenoiser`
+      This block supports text-to-image and image-to-image generation, without classifier-free guidance.
+
+      Components:
+          transformer (`Flux3Transformer2DModel`) scheduler (`FlowMatchEulerDiscreteScheduler`)
+
+      Inputs:
+          timesteps (`Tensor`):
+              The timesteps to use for the denoising process.
+          num_inference_steps (`int`):
+              The number of inference steps to use for the denoising process.
+          latents (`Tensor`):
+              The image latent tokens to denoise. Shape: (B, seq_len, C)
+          **denoiser_input_fields (`None`, *optional*):
+              conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc.
+
+      Outputs:
+          latents (`Tensor`):
+              The denoised latents
+    """
+
+    model_name = "flux3-distilled"
+    """
+    Denoise step that iteratively denoises the latents for the step-distilled (flow-map) FLUX.3 model.
+      Its loop logic is defined in `Flux3DenoiseLoopWrapper.__call__` method At each iteration, it runs blocks defined
+      in `sub_blocks` sequentially:
+       - `Flux3DistilledLoopDenoiser`
+       - `Flux3LoopAfterDenoiser`
+      This block supports text-to-image and image-to-image generation, without classifier-free guidance.
+
+      Components:
+          transformer (`Flux3Transformer2DModel`) scheduler (`FlowMatchEulerDiscreteScheduler`)
+
+      Inputs:
+          timesteps (`Tensor`):
+              The timesteps to use for the denoising process.
+          num_inference_steps (`int`):
+              The number of inference steps to use for the denoising process.
+          latents (`Tensor`):
+              The image latent tokens to denoise. Shape: (B, seq_len, C)
+          **denoiser_input_fields (`None`, *optional*):
+              conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc.
+
+      Outputs:
+          latents (`Tensor`):
+              The denoised latents
+    """
+
+    block_classes = [Flux3DistilledLoopDenoiser(), Flux3LoopAfterDenoiser()]
+    block_names = ["denoiser", "after_denoiser"]
+
+    @property
+    def description(self) -> str:
+        return (
+            "Denoise step that iteratively denoises the latents for the step-distilled (flow-map) FLUX.3 model. \n"
+            "Its loop logic is defined in `Flux3DenoiseLoopWrapper.__call__` method \n"
+            "At each iteration, it runs blocks defined in `sub_blocks` sequentially:\n"
+            " - `Flux3DistilledLoopDenoiser`\n"
+            " - `Flux3LoopAfterDenoiser`\n"
+            "This block supports text-to-image and image-to-image generation, without classifier-free guidance."
+        )

diffusers/modular_pipelines/flux3/encoders.py

@@ -326,3 +326,70 @@ class Flux3TextEncoderStep(ModularPipelineBlocks):
 
         self.set_block_state(state, block_state)
         return components, state
+
+
+class Flux3DistilledTextEncoderStep(Flux3TextEncoderStep):
+    model_name = "flux3-distilled"
+
+    @property
+    def description(self) -> str:
+        return (
+            "Text Encoder step that generates text embeddings using Qwen3-VL to guide the generation with the "
+            "step-distilled FLUX.3 model, which samples without classifier-free guidance (no negative prompt). Takes "
+            f"one prompt per call; it goes through the chat template and is right-padded to a multiple of "
+            f"{self.pad_multiple} tokens."
+        )
+
+    @property
+    def inputs(self) -> list[InputParam]:
+        return [
+            InputParam.template("prompt", description="The prompt to guide the generation (one prompt per call)"),
+            InputParam.template(
+                "max_sequence_length",
+                default=8192,
+                description="Maximum number of prompt tokens; longer prompts are truncated.",
+            ),
+            InputParam(
+                "text_encoder_out_layers",
+                type_hint=tuple[int],
+                default=(4, 8, 12, 16, 20, 24, 28, 32),
+                required=False,
+                description="Indices of the text-encoder hidden layers whose states are stacked to form the text embeddings",
+            ),
+        ]
+
+    @property
+    def intermediate_outputs(self) -> list[OutputParam]:
+        return [
+            OutputParam(
+                "prompt_embeds",
+                kwargs_type="denoiser_input_fields",
+                type_hint=torch.Tensor,
+                description="Text embeddings (stacked Qwen3-VL hidden states) used to guide the image generation",
+            ),
+        ]
+
+    @staticmethod
+    def check_inputs(block_state):
+        if not isinstance(block_state.prompt, str):
+            raise ValueError(
+                f"`prompt` must be a str; batched prompts are not supported, got {type(block_state.prompt).__name__}."
+            )
+
+    @torch.no_grad()
+    def __call__(self, components: Flux3ModularPipeline, state: PipelineState) -> PipelineState:
+        block_state = self.get_block_state(state)
+        self.check_inputs(block_state)
+
+        block_state.prompt_embeds = self._get_qwen3_vl_prompt_embeds(
+            text_encoder=components.text_encoder,
+            tokenizer=components.tokenizer,
+            prompt=block_state.prompt,
+            device=components._execution_device,
+            max_sequence_length=block_state.max_sequence_length,
+            pad_multiple=self.pad_multiple,
+            hidden_states_layers=block_state.text_encoder_out_layers,
+        )
+
+        self.set_block_state(state, block_state)
+        return components, state

diffusers/modular_pipelines/flux3/inputs.py

@@ -173,6 +173,73 @@ class Flux3TextInputStep(ModularPipelineBlocks):
         return components, state
 
 
+class Flux3DistilledTextInputStep(ModularPipelineBlocks):
+    model_name = "flux3-distilled"
+
+    @property
+    def description(self) -> str:
+        return (
+            "This step:\n"
+            "  1. Determines `batch_size` from `prompt_embeds` and the compute `dtype` from the transformer\n"
+            "  2. Repeats the text embeddings to the batch size (batch_size * num_images_per_prompt)\n"
+            "For the step-distilled FLUX.3 model, which samples without classifier-free guidance (no negative "
+            "embeddings)."
+        )
+
+    @property
+    def expected_components(self) -> list[ComponentSpec]:
+        return [ComponentSpec("transformer", Flux3Transformer2DModel)]
+
+    @property
+    def inputs(self) -> list[InputParam]:
+        return [
+            InputParam.template("num_images_per_prompt"),
+            InputParam(
+                "prompt_embeds",
+                required=True,
+                kwargs_type="denoiser_input_fields",
+                type_hint=torch.Tensor,
+                description="Pre-generated text embeddings. Can be generated from text_encoder step.",
+            ),
+        ]
+
+    @property
+    def intermediate_outputs(self) -> list[OutputParam]:
+        return [
+            OutputParam(
+                "batch_size",
+                type_hint=int,
+                description="Number of prompts, the final batch size of model inputs should be batch_size * num_images_per_prompt",
+            ),
+            OutputParam(
+                "dtype",
+                type_hint=torch.dtype,
+                description="Data type of model tensor inputs (the transformer's dtype)",
+            ),
+            OutputParam(
+                "prompt_embeds",
+                type_hint=torch.Tensor,
+                kwargs_type="denoiser_input_fields",
+                description="Text embeddings used to guide the image generation",
+            ),
+        ]
+
+    @torch.no_grad()
+    def __call__(self, components: Flux3ModularPipeline, state: PipelineState) -> PipelineState:
+        block_state = self.get_block_state(state)
+
+        block_state.batch_size = block_state.prompt_embeds.shape[0]
+        # the denoiser consumes every tensor this chain prepares, so its dtype is the compute dtype
+        block_state.dtype = components.transformer.dtype
+
+        block_state.prompt_embeds = repeat_tensor_to_batch_size(
+            "prompt_embeds", block_state.prompt_embeds, block_state.batch_size, block_state.num_images_per_prompt
+        ).to(block_state.dtype)
+
+        self.set_block_state(state, block_state)
+        return components, state
+
+
 class Flux3ImageCondInputStep(ModularPipelineBlocks):
     model_name = "flux3"
 

diffusers/modular_pipelines/flux3/modular_blocks_flux3_distilled.py

@@ -0,0 +1,333 @@
+# Copyright 2026 Black Forest Labs and 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 logging
+from ..modular_pipeline import AutoPipelineBlocks, SequentialPipelineBlocks
+from ..modular_pipeline_utils import InsertableDict, OutputParam
+from .before_denoise import (
+    Flux3DistilledTextRoPEInputsStep,
+    Flux3PrepareImageCondLatentsStep,
+    Flux3PrepareLatentsStep,
+    Flux3SetTimestepsStep,
+)
+from .decoders import Flux3DecodeStep, Flux3UnpackLatentsStep
+from .denoise import Flux3DistilledDenoiseStep
+from .encoders import Flux3DistilledTextEncoderStep
+from .inputs import Flux3DistilledTextInputStep, Flux3ImageCondInputStep
+from .modular_blocks_flux3 import Flux3AutoVaeEncoderStep
+
+
+logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
+
+
+# The step-distilled (flow-map) FLUX.3 checkpoints: 16 steps and no classifier-free guidance. The text encoder,
+# text input, text RoPE and denoise steps have distilled versions without the negative prompt; the denoise loop
+# passes the transformer the timestep each step jumps to and runs a single pass. Everything else is shared with the
+# base model.
+
+###
+### Core denoise
+###
+
+Flux3DistilledCoreDenoiseBlocks = InsertableDict(
+    [
+        ("input", Flux3DistilledTextInputStep()),
+        ("prepare_latents", Flux3PrepareLatentsStep()),
+        ("set_timesteps", Flux3SetTimestepsStep(num_inference_steps=16)),
+        ("prepare_rope_inputs", Flux3DistilledTextRoPEInputsStep()),
+        ("denoise", Flux3DistilledDenoiseStep()),
+        ("after_denoise", Flux3UnpackLatentsStep()),
+    ]
+)
+
+
+# auto_docstring
+class Flux3DistilledCoreDenoiseStep(SequentialPipelineBlocks):
+    """
+    Core denoise step that performs the denoising process for the step-distilled FLUX.3 model, for text-to-image
+    generation.
+
+      Components:
+          transformer (`Flux3Transformer2DModel`) scheduler (`FlowMatchEulerDiscreteScheduler`)
+
+      Inputs:
+          num_images_per_prompt (`int`, *optional*, defaults to 1):
+              The number of images to generate per prompt.
+          prompt_embeds (`Tensor`):
+              Pre-generated text embeddings. Can be generated from text_encoder step.
+          height (`int`, *optional*, defaults to 1024):
+              The height in pixels of the generated image.
+          width (`int`, *optional*, defaults to 1024):
+              The width in pixels of the generated image.
+          latents (`Tensor`, *optional*):
+              Pre-generated noisy latents for image generation.
+          generator (`Generator`, *optional*):
+              Torch generator for deterministic generation.
+          num_inference_steps (`int`, *optional*, defaults to 16):
+              The number of denoising steps.
+          timesteps (`Tensor`, *optional*):
+              Timesteps for the denoising process. (overrides `num_inference_steps` when provided)
+          sigmas (`list`, *optional*):
+              Custom sigmas for the denoising process. (overrides `num_inference_steps` when provided)
+          **denoiser_input_fields (`None`, *optional*):
+              conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc.
+
+      Outputs:
+          latents (`Tensor`):
+              Denoised latents.
+    """
+
+    model_name = "flux3-distilled"
+    block_classes = Flux3DistilledCoreDenoiseBlocks.values()
+    block_names = Flux3DistilledCoreDenoiseBlocks.keys()
+
+    @property
+    def description(self):
+        return (
+            "Core denoise step that performs the denoising process for the step-distilled FLUX.3 model, for "
+            "text-to-image generation."
+        )
+
+    @property
+    def outputs(self):
+        return [
+            OutputParam.template("latents"),
+        ]
+
+
+Flux3DistilledImageConditionedCoreDenoiseBlocks = InsertableDict(
+    [
+        ("input", Flux3DistilledTextInputStep()),
+        ("cond_input", Flux3ImageCondInputStep()),
+        ("prepare_latents", Flux3PrepareLatentsStep()),
+        ("prepare_image_cond_latents", Flux3PrepareImageCondLatentsStep()),
+        ("set_timesteps", Flux3SetTimestepsStep(num_inference_steps=16)),
+        ("prepare_rope_inputs", Flux3DistilledTextRoPEInputsStep()),
+        ("denoise", Flux3DistilledDenoiseStep()),
+        ("after_denoise", Flux3UnpackLatentsStep()),
+    ]
+)
+
+
+# auto_docstring
+class Flux3DistilledImageConditionedCoreDenoiseStep(SequentialPipelineBlocks):
+    """
+    Core denoise step that performs the denoising process for the step-distilled FLUX.3 model with reference images
+    (image-to-image).
+
+      Components:
+          transformer (`Flux3Transformer2DModel`) scheduler (`FlowMatchEulerDiscreteScheduler`)
+
+      Inputs:
+          num_images_per_prompt (`int`, *optional*, defaults to 1):
+              The number of images to generate per prompt.
+          prompt_embeds (`Tensor`):
+              Pre-generated text embeddings. Can be generated from text_encoder step.
+          image_cond_latents (`list`):
+              Normalized packed latents of the reference images, one tensor of shape (B, C, height_i//16, width_i//16)
+              per image
+          height (`int`, *optional*):
+              The height in pixels of the generated image; defaults to the first reference image's aspect ratio at the
+              default area
+          width (`int`, *optional*):
+              The width in pixels of the generated image; defaults to the first reference image's aspect ratio at the
+              default area
+          latents (`Tensor`, *optional*):
+              Pre-generated noisy latents for image generation.
+          generator (`Generator`, *optional*):
+              Torch generator for deterministic generation.
+          num_inference_steps (`int`, *optional*, defaults to 16):
+              The number of denoising steps.
+          timesteps (`Tensor`, *optional*):
+              Timesteps for the denoising process. (overrides `num_inference_steps` when provided)
+          sigmas (`list`, *optional*):
+              Custom sigmas for the denoising process. (overrides `num_inference_steps` when provided)
+          **denoiser_input_fields (`None`, *optional*):
+              conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc.
+
+      Outputs:
+          latents (`Tensor`):
+              Denoised latents.
+    """
+
+    model_name = "flux3-distilled"
+    block_classes = Flux3DistilledImageConditionedCoreDenoiseBlocks.values()
+    block_names = Flux3DistilledImageConditionedCoreDenoiseBlocks.keys()
+
+    @property
+    def description(self):
+        return (
+            "Core denoise step that performs the denoising process for the step-distilled FLUX.3 model with reference "
+            "images (image-to-image)."
+        )
+
+    @property
+    def outputs(self):
+        return [
+            OutputParam.template("latents"),
+        ]
+
+
+# auto_docstring
+class Flux3DistilledAutoCoreDenoiseStep(AutoPipelineBlocks):
+    """
+    Core denoise step that selects the step-distilled FLUX.3 denoising workflow by the conditioning inputs.
+       - `Flux3DistilledImageConditionedCoreDenoiseStep` is used when `image_cond_latents` is provided
+         (image-to-image).
+       - `Flux3DistilledCoreDenoiseStep` is used otherwise (text-to-image).
+
+      Components:
+          transformer (`Flux3Transformer2DModel`) scheduler (`FlowMatchEulerDiscreteScheduler`)
+
+      Inputs:
+          num_images_per_prompt (`int`, *optional*, defaults to 1):
+              The number of images to generate per prompt.
+          prompt_embeds (`Tensor`):
+              Pre-generated text embeddings. Can be generated from text_encoder step.
+          image_cond_latents (`list`, *optional*):
+              Normalized packed latents of the reference images, one tensor of shape (B, C, height_i//16, width_i//16)
+              per image
+          height (`int`, *optional*, defaults to None or 1024, depending on the workflow):
+              The height in pixels of the generated image; defaults to the first reference image's aspect ratio at the
+              default area
+          width (`int`, *optional*, defaults to None or 1024, depending on the workflow):
+              The width in pixels of the generated image; defaults to the first reference image's aspect ratio at the
+              default area
+          latents (`Tensor`):
+              Pre-generated noisy latents for image generation.
+          generator (`Generator`, *optional*):
+              Torch generator for deterministic generation.
+          num_inference_steps (`int`):
+              The number of denoising steps.
+          timesteps (`Tensor`):
+              Timesteps for the denoising process. (overrides `num_inference_steps` when provided)
+          sigmas (`list`, *optional*):
+              Custom sigmas for the denoising process. (overrides `num_inference_steps` when provided)
+          **denoiser_input_fields (`None`, *optional*):
+              conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc.
+
+      Outputs:
+          latents (`Tensor`):
+              Denoised latents.
+    """
+
+    model_name = "flux3-distilled"
+    block_classes = [Flux3DistilledImageConditionedCoreDenoiseStep, Flux3DistilledCoreDenoiseStep]
+    block_names = ["image_conditioned", "text2image"]
+    block_trigger_inputs = ["image_cond_latents", None]
+
+    @property
+    def description(self):
+        return (
+            "Core denoise step that selects the step-distilled FLUX.3 denoising workflow by the conditioning inputs.\n"
+            " - `Flux3DistilledImageConditionedCoreDenoiseStep` is used when `image_cond_latents` is provided "
+            "(image-to-image).\n"
+            " - `Flux3DistilledCoreDenoiseStep` is used otherwise (text-to-image)."
+        )
+
+
+###
+### Auto blocks
+###
+
+DISTILLED_AUTO_BLOCKS = InsertableDict(
+    [
+        ("text_encoder", Flux3DistilledTextEncoderStep()),
+        ("vae_encoder", Flux3AutoVaeEncoderStep()),
+        ("denoise", Flux3DistilledAutoCoreDenoiseStep()),
+        ("decode", Flux3DecodeStep()),
+    ]
+)
+
+
+# auto_docstring
+class Flux3DistilledAutoBlocks(SequentialPipelineBlocks):
+    """
+    Auto blocks that perform text-to-image and image-to-image generation using the step-distilled FLUX.3 model (16
+    steps, no classifier-free guidance).
+
+      Supported workflows:
+        - `text2image`: requires `prompt`
+        - `image_conditioned`: requires `prompt`, `image`
+
+      Components:
+          text_encoder (`Qwen3VLForConditionalGeneration`) tokenizer (`AutoProcessor`) image_processor
+          (`Flux3ImageProcessor`) vae (`AutoencoderKLFlux2`) transformer (`Flux3Transformer2DModel`) scheduler
+          (`FlowMatchEulerDiscreteScheduler`)
+
+      Inputs:
+          prompt (`str`):
+              The prompt to guide the generation (one prompt per call)
+          max_sequence_length (`int`, *optional*, defaults to 8192):
+              Maximum number of prompt tokens; longer prompts are truncated.
+          text_encoder_out_layers (`tuple`, *optional*, defaults to (4, 8, 12, 16, 20, 24, 28, 32)):
+              Indices of the text-encoder hidden layers whose states are stacked to form the text embeddings
+          image (`Image | list`, *optional*):
+              The reference image or images to condition the generation on; each keeps its own resolution
+              (center-cropped to the 16-pixel grid)
+          num_images_per_prompt (`int`, *optional*, defaults to 1):
+              The number of images to generate per prompt.
+          image_cond_latents (`list`, *optional*):
+              Normalized packed latents of the reference images, one tensor of shape (B, C, height_i//16, width_i//16)
+              per image
+          height (`int`, *optional*, defaults to None or 1024, depending on the workflow):
+              The height in pixels of the generated image; defaults to the first reference image's aspect ratio at the
+              default area
+          width (`int`, *optional*, defaults to None or 1024, depending on the workflow):
+              The width in pixels of the generated image; defaults to the first reference image's aspect ratio at the
+              default area
+          latents (`Tensor`):
+              Pre-generated noisy latents for image generation.
+          generator (`Generator`, *optional*):
+              Torch generator for deterministic generation.
+          num_inference_steps (`int`):
+              The number of denoising steps.
+          timesteps (`Tensor`):
+              Timesteps for the denoising process. (overrides `num_inference_steps` when provided)
+          sigmas (`list`, *optional*):
+              Custom sigmas for the denoising process. (overrides `num_inference_steps` when provided)
+          **denoiser_input_fields (`None`, *optional*):
+              conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc.
+          output_type (`str`, *optional*, defaults to pil):
+              Output format: 'pil', 'np', 'pt'.
+
+      Outputs:
+          images (`list`):
+              Generated images.
+          latents (`Tensor`):
+              Denoised latents.
+    """
+
+    model_name = "flux3-distilled"
+    block_classes = DISTILLED_AUTO_BLOCKS.values()
+    block_names = DISTILLED_AUTO_BLOCKS.keys()
+    _workflow_map = {
+        "text2image": {"prompt": True},
+        "image_conditioned": {"prompt": True, "image": True},
+    }
+
+    @property
+    def description(self):
+        return (
+            "Auto blocks that perform text-to-image and image-to-image generation using the step-distilled FLUX.3 "
+            "model (16 steps, no classifier-free guidance)."
+        )
+
+    @property
+    def outputs(self):
+        return [
+            OutputParam.template("images"),
+            OutputParam.template("latents"),
+        ]

diffusers/modular_pipelines/flux3/modular_pipeline.py

@@ -46,3 +46,11 @@ class Flux3ModularPipeline(ModularPipeline):
         if getattr(self, "guider", None) is None:
             return False
         return self.guider._enabled and self.guider.num_conditions > 1
+
+
+class Flux3DistilledModularPipeline(Flux3ModularPipeline):
+    """
+    A ModularPipeline for the step-distilled FLUX.3 model (16 steps, no classifier-free guidance).
+    """
+
+    default_blocks_name = "Flux3DistilledAutoBlocks"

diffusers/modular_pipelines/modular_pipeline.py

@@ -141,6 +141,7 @@ MODULAR_PIPELINE_MAPPING = OrderedDict(
         ("flux2", _create_default_map_fn("Flux2ModularPipeline")),
         ("flux2-klein", _flux2_klein_map_fn),
         ("flux3", _create_default_map_fn("Flux3ModularPipeline")),
+        ("flux3-distilled", _create_default_map_fn("Flux3DistilledModularPipeline")),
         ("ideogram4", _create_default_map_fn("Ideogram4ModularPipeline")),
         ("krea2", _krea2_map_fn),
         ("qwenimage", _create_default_map_fn("QwenImageModularPipeline")),

diffusers/pipelines/README.md

@@ -0,0 +1,171 @@
+# 🧨 Diffusers Pipelines
+
+Pipelines provide a simple way to run state-of-the-art diffusion models in inference.
+Most diffusion systems consist of multiple independently-trained models and highly adaptable scheduler
+components - all of which are needed to have a functioning end-to-end diffusion system.
+
+As an example, [Stable Diffusion](https://huggingface.co/blog/stable_diffusion) has three independently trained models:
+- [Autoencoder](https://github.com/huggingface/diffusers/blob/5cbed8e0d157f65d3ddc2420dfd09f2df630e978/src/diffusers/models/vae.py#L392)
+- [Conditional Unet](https://github.com/huggingface/diffusers/blob/5cbed8e0d157f65d3ddc2420dfd09f2df630e978/src/diffusers/models/unet_2d_condition.py#L12)
+- [CLIP text encoder](https://huggingface.co/docs/transformers/main/en/model_doc/clip#transformers.CLIPTextModel)
+- a scheduler component, [scheduler](https://github.com/huggingface/diffusers/blob/main/src/diffusers/schedulers/scheduling_pndm.py),
+- a [CLIPImageProcessor](https://huggingface.co/docs/transformers/main/en/model_doc/clip#transformers.CLIPImageProcessor),
+- as well as a [safety checker](https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/stable_diffusion/safety_checker.py).
+All of these components are necessary to run stable diffusion in inference even though they were trained
+or created independently from each other.
+
+To that end, we strive to offer all open-sourced, state-of-the-art diffusion system under a unified API.
+More specifically, we strive to provide pipelines that
+- 1. can load the officially published weights and yield 1-to-1 the same outputs as the original implementation according to the corresponding paper (*e.g.* [LDMTextToImagePipeline](https://github.com/huggingface/diffusers/tree/main/src/diffusers/pipelines/latent_diffusion), uses the officially released weights of [High-Resolution Image Synthesis with Latent Diffusion Models](https://huggingface.co/papers/2112.10752)),
+- 2. have a simple user interface to run the model in inference (see the [Pipelines API](#pipelines-api) section),
+- 3. are easy to understand with code that is self-explanatory and can be read along-side the official paper (see [Pipelines summary](#pipelines-summary)),
+- 4. can easily be contributed by the community (see the [Contribution](#contribution) section).
+
+**Note** that pipelines do not (and should not) offer any training functionality.
+If you are looking for *official* training examples, please have a look at [examples](https://github.com/huggingface/diffusers/tree/main/examples).
+
+
+## Pipelines Summary
+
+The following table summarizes all officially supported pipelines, their corresponding paper, and if
+available a colab notebook to directly try them out.
+
+| Pipeline                                                                                                                      | Source                                                                                                                       | Tasks | Colab
+|-------------------------------------------------------------------------------------------------------------------------------|------------------------------------------------------------------------------------------------------------------------------|:---:|:---:|
+| [dance diffusion](https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/dance_diffusion)                 | [**Dance Diffusion**](https://github.com/Harmonai-org/sample-generator)                                                      | *Unconditional Audio Generation* |
+| [ddpm](https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/ddpm)                                       | [**Denoising Diffusion Probabilistic Models**](https://huggingface.co/papers/2006.11239)                                             | *Unconditional Image Generation* |
+| [ddim](https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/ddim)                                       | [**Denoising Diffusion Implicit Models**](https://huggingface.co/papers/2010.02502)                                                  | *Unconditional Image Generation* | [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/huggingface/notebooks/blob/main/diffusers/training_example.ipynb)
+| [latent_diffusion](https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/latent_diffusion)               | [**High-Resolution Image Synthesis with Latent Diffusion Models**](https://huggingface.co/papers/2112.10752)                         | *Text-to-Image Generation* |
+| [latent_diffusion_uncond](https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/latent_diffusion_uncond) | [**High-Resolution Image Synthesis with Latent Diffusion Models**](https://huggingface.co/papers/2112.10752)                         | *Unconditional Image Generation* |
+| [pndm](https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/pndm)                                       | [**Pseudo Numerical Methods for Diffusion Models on Manifolds**](https://huggingface.co/papers/2202.09778)                           | *Unconditional Image Generation* |
+| [score_sde_ve](https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/score_sde_ve)                       | [**Score-Based Generative Modeling through Stochastic Differential Equations**](https://openreview.net/forum?id=PxTIG12RRHS) | *Unconditional Image Generation* |
+| [score_sde_vp](https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/score_sde_vp)                       | [**Score-Based Generative Modeling through Stochastic Differential Equations**](https://openreview.net/forum?id=PxTIG12RRHS) | *Unconditional Image Generation* |
+| [stable_diffusion](https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/stable_diffusion)               | [**Stable Diffusion**](https://stability.ai/blog/stable-diffusion-public-release)                                            | *Text-to-Image Generation* | [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/huggingface/notebooks/blob/main/diffusers/stable_diffusion.ipynb)
+| [stable_diffusion](https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/stable_diffusion)               | [**Stable Diffusion**](https://stability.ai/blog/stable-diffusion-public-release)                                            | *Image-to-Image Text-Guided Generation* | [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/huggingface/notebooks/blob/main/diffusers/image_2_image_using_diffusers.ipynb)
+| [stable_diffusion](https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/stable_diffusion)               | [**Stable Diffusion**](https://stability.ai/blog/stable-diffusion-public-release)                                            | *Text-Guided Image Inpainting* | [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/huggingface/notebooks/blob/main/diffusers/in_painting_with_stable_diffusion_using_diffusers.ipynb)
+| [stochastic_karras_ve](https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/stochastic_karras_ve)       | [**Elucidating the Design Space of Diffusion-Based Generative Models**](https://huggingface.co/papers/2206.00364)                    | *Unconditional Image Generation* |
+
+**Note**: Pipelines are simple examples of how to play around with the diffusion systems as described in the corresponding papers.
+However, most of them can be adapted to use different scheduler components or even different model components. Some pipeline examples are shown in the [Examples](#examples) below.
+
+## Pipelines API
+
+Diffusion models often consist of multiple independently-trained models or other previously existing components.
+
+
+Each model has been trained independently on a different task and the scheduler can easily be swapped out and replaced with a different one.
+During inference, we however want to be able to easily load all components and use them in inference - even if one component, *e.g.* CLIP's text encoder, originates from a different library, such as [Transformers](https://github.com/huggingface/transformers). To that end, all pipelines provide the following functionality:
+
+- [`from_pretrained` method](https://github.com/huggingface/diffusers/blob/5cbed8e0d157f65d3ddc2420dfd09f2df630e978/src/diffusers/pipeline_utils.py#L139) that accepts a Hugging Face Hub repository id, *e.g.* [stable-diffusion-v1-5/stable-diffusion-v1-5](https://huggingface.co/stable-diffusion-v1-5/stable-diffusion-v1-5) or a path to a local directory, *e.g.*
+"./stable-diffusion". To correctly retrieve which models and components should be loaded, one has to provide a `model_index.json` file, *e.g.* [stable-diffusion-v1-5/stable-diffusion-v1-5/model_index.json](https://huggingface.co/stable-diffusion-v1-5/stable-diffusion-v1-5/blob/main/model_index.json), which defines all components that should be
+loaded into the pipelines. More specifically, for each model/component one needs to define the format `<name>: ["<library>", "<class name>"]`. `<name>` is the attribute name given to the loaded instance of `<class name>` which can be found in the library or pipeline folder called `"<library>"`.
+- [`save_pretrained`](https://github.com/huggingface/diffusers/blob/5cbed8e0d157f65d3ddc2420dfd09f2df630e978/src/diffusers/pipeline_utils.py#L90) that accepts a local path, *e.g.* `./stable-diffusion` under which all models/components of the pipeline will be saved. For each component/model a folder is created inside the local path that is named after the given attribute name, *e.g.* `./stable_diffusion/unet`.
+In addition, a `model_index.json` file is created at the root of the local path, *e.g.* `./stable_diffusion/model_index.json` so that the complete pipeline can again be instantiated
+from the local path.
+- [`to`](https://github.com/huggingface/diffusers/blob/5cbed8e0d157f65d3ddc2420dfd09f2df630e978/src/diffusers/pipeline_utils.py#L118) which accepts a `string` or `torch.device` to move all models that are of type `torch.nn.Module` to the passed device. The behavior is fully analogous to [PyTorch's `to` method](https://pytorch.org/docs/stable/generated/torch.nn.Module.html#torch.nn.Module.to).
+- [`__call__`] method to use the pipeline in inference. `__call__` defines inference logic of the pipeline and should ideally encompass all aspects of it, from pre-processing to forwarding tensors to the different models and schedulers, as well as post-processing. The API of the `__call__` method can strongly vary from pipeline to pipeline. *E.g.* a text-to-image pipeline, such as [`StableDiffusionPipeline`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py) should accept among other things the text prompt to generate the image. A pure image generation pipeline, such as [DDPMPipeline](https://github.com/huggingface/diffusers/tree/main/src/diffusers/pipelines/ddpm) on the other hand can be run without providing any inputs. To better understand what inputs can be adapted for
+each pipeline, one should look directly into the respective pipeline.
+
+**Note**: All pipelines have PyTorch's autograd disabled by decorating the `__call__` method with a [`torch.no_grad`](https://pytorch.org/docs/stable/generated/torch.no_grad.html) decorator because pipelines should
+not be used for training. If you want to store the gradients during the forward pass, we recommend writing your own pipeline, see also our [community-examples](https://github.com/huggingface/diffusers/tree/main/examples/community)
+
+## Contribution
+
+We are more than happy about any contribution to the officially supported pipelines 🤗. We aspire
+all of our pipelines to be  **self-contained**, **easy-to-tweak**, **beginner-friendly** and for **one-purpose-only**.
+
+- **Self-contained**: A pipeline shall be as self-contained as possible. More specifically, this means that all functionality should be either directly defined in the pipeline file itself, should be inherited from (and only from) the [`DiffusionPipeline` class](https://github.com/huggingface/diffusers/blob/5cbed8e0d157f65d3ddc2420dfd09f2df630e978/src/diffusers/pipeline_utils.py#L56) or be directly attached to the model and scheduler components of the pipeline.
+- **Easy-to-use**: Pipelines should be extremely easy to use - one should be able to load the pipeline and
+use it for its designated task, *e.g.* text-to-image generation, in just a couple of lines of code. Most
+logic including pre-processing, an unrolled diffusion loop, and post-processing should all happen inside the `__call__` method.
+- **Easy-to-tweak**: Certain pipelines will not be able to handle all use cases and tasks that you might like them to. If you want to use a certain pipeline for a specific use case that is not yet supported, you might have to copy the pipeline file and tweak the code to your needs. We try to make the pipeline code as readable as possible so that each part –from pre-processing to diffusing to post-processing– can easily be adapted. If you would like the community to benefit from your customized pipeline, we would love to see a contribution to our [community-examples](https://github.com/huggingface/diffusers/tree/main/examples/community). If you feel that an important pipeline should be part of the official pipelines but isn't, a contribution to the [official pipelines](https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines) would be even better.
+- **One-purpose-only**: Pipelines should be used for one task and one task only. Even if two tasks are very similar from a modeling point of view, *e.g.* image2image translation and in-painting, pipelines shall be used for one task only to keep them *easy-to-tweak* and *readable*.
+
+## Examples
+
+### Text-to-Image generation with Stable Diffusion
+
+```python
+# make sure you're logged in with `hf auth login`
+from diffusers import StableDiffusionPipeline, LMSDiscreteScheduler
+
+pipe = StableDiffusionPipeline.from_pretrained("stable-diffusion-v1-5/stable-diffusion-v1-5")
+pipe = pipe.to("cuda")
+
+prompt = "a photo of an astronaut riding a horse on mars"
+image = pipe(prompt).images[0]
+
+image.save("astronaut_rides_horse.png")
+```
+
+### Image-to-Image text-guided generation with Stable Diffusion
+
+The `StableDiffusionImg2ImgPipeline` lets you pass a text prompt and an initial image to condition the generation of new images.
+
+```python
+import requests
+from PIL import Image
+from io import BytesIO
+
+from diffusers import StableDiffusionImg2ImgPipeline
+
+# load the pipeline
+device = "cuda"
+pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
+    "stable-diffusion-v1-5/stable-diffusion-v1-5",
+    torch_dtype=torch.float16,
+).to(device)
+
+# let's download an initial image
+url = "https://raw.githubusercontent.com/CompVis/stable-diffusion/main/assets/stable-samples/img2img/sketch-mountains-input.jpg"
+
+response = requests.get(url)
+init_image = Image.open(BytesIO(response.content)).convert("RGB")
+init_image = init_image.resize((768, 512))
+
+prompt = "A fantasy landscape, trending on artstation"
+
+images = pipe(prompt=prompt, image=init_image, strength=0.75, guidance_scale=7.5).images
+
+images[0].save("fantasy_landscape.png")
+```
+You can also run this example on colab [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/huggingface/notebooks/blob/main/diffusers/image_2_image_using_diffusers.ipynb)
+
+### Tweak prompts reusing seeds and latents
+
+You can generate your own latents to reproduce results, or tweak your prompt on a specific result you liked. [This notebook](https://github.com/pcuenca/diffusers-examples/blob/main/notebooks/stable-diffusion-seeds.ipynb) shows how to do it step by step. You can also run it in Google Colab [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/pcuenca/diffusers-examples/blob/main/notebooks/stable-diffusion-seeds.ipynb).
+
+
+### In-painting using Stable Diffusion
+
+The `StableDiffusionInpaintPipeline` lets you edit specific parts of an image by providing a mask and text prompt.
+
+```python
+import PIL
+import requests
+import torch
+from io import BytesIO
+
+from diffusers import StableDiffusionInpaintPipeline
+
+def download_image(url):
+    response = requests.get(url)
+    return PIL.Image.open(BytesIO(response.content)).convert("RGB")
+
+img_url = "https://raw.githubusercontent.com/CompVis/latent-diffusion/main/data/inpainting_examples/overture-creations-5sI6fQgYIuo.png"
+mask_url = "https://raw.githubusercontent.com/CompVis/latent-diffusion/main/data/inpainting_examples/overture-creations-5sI6fQgYIuo_mask.png"
+
+init_image = download_image(img_url).resize((512, 512))
+mask_image = download_image(mask_url).resize((512, 512))
+
+pipe = StableDiffusionInpaintPipeline.from_pretrained(
+    "stable-diffusion-v1-5/stable-diffusion-inpainting",
+    torch_dtype=torch.float16,
+)
+pipe = pipe.to("cuda")
+
+prompt = "Face of a yellow cat, high resolution, sitting on a park bench"
+image = pipe(prompt=prompt, image=init_image, mask_image=mask_image).images[0]
+```
+
+You can also run this example on colab [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/huggingface/notebooks/blob/main/diffusers/in_painting_with_stable_diffusion_using_diffusers.ipynb)

diffusers/pipelines/deprecated/README.md

@@ -0,0 +1,3 @@
+# Deprecated Pipelines
+
+This folder contains pipelines that have very low usage as measured by model downloads, issues and PRs. While you can still use the pipelines just as before, we will stop testing the pipelines and will not accept any changes to existing files.

diffusers/pipelines/stable_diffusion/README.md

@@ -0,0 +1,176 @@
+# Stable Diffusion
+
+## Overview
+
+Stable Diffusion was proposed in [Stable Diffusion Announcement](https://stability.ai/blog/stable-diffusion-announcement) by Patrick Esser and Robin Rombach and the Stability AI team.
+
+The summary of the model is the following:
+
+*Stable Diffusion is a text-to-image model that will empower billions of people to create stunning art within seconds. It is a breakthrough in speed and quality meaning that it can run on consumer GPUs. You can see some of the amazing output that has been created by this model without pre or post-processing on this page. The model itself builds upon the work of the team at CompVis and Runway in their widely used latent diffusion model combined with insights from the conditional diffusion models by our lead generative AI developer Katherine Crowson, Dall-E 2 by Open AI, Imagen by Google Brain and many others. We are delighted that AI media generation is a cooperative field and hope it can continue this way to bring the gift of creativity to all.*
+
+## Tips:
+
+- Stable Diffusion has the same architecture as [Latent Diffusion](https://huggingface.co/papers/2112.10752) but uses a frozen CLIP Text Encoder instead of training the text encoder jointly with the diffusion model.
+- An in-detail explanation of the Stable Diffusion model can be found under [Stable Diffusion with 🧨 Diffusers](https://huggingface.co/blog/stable_diffusion).
+- If you don't want to rely on the Hugging Face Hub and having to pass a authentication token, you can
+download the weights with `git lfs install; git clone https://huggingface.co/stable-diffusion-v1-5/stable-diffusion-v1-5` and instead pass the local path to the cloned folder to `from_pretrained` as shown below.
+- Stable Diffusion can work with a variety of different samplers as is shown below.
+
+## Available Pipelines:
+
+| Pipeline | Tasks | Colab
+|---|---|:---:|
+| [pipeline_stable_diffusion.py](https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py) | *Text-to-Image Generation* | [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/huggingface/notebooks/blob/main/diffusers/training_example.ipynb)
+| [pipeline_stable_diffusion_img2img](https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_img2img.py) | *Image-to-Image Text-Guided Generation* | [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/huggingface/notebooks/blob/main/diffusers/image_2_image_using_diffusers.ipynb)
+| [pipeline_stable_diffusion_inpaint](https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_inpaint.py) | *Text-Guided Image Inpainting* | [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/huggingface/notebooks/blob/main/diffusers/in_painting_with_stable_diffusion_using_diffusers.ipynb)
+
+## Examples:
+
+### Using Stable Diffusion without being logged into the Hub.
+
+If you want to download the model weights using a single Python line, you need to be logged in via `hf auth login`.
+
+```python
+from diffusers import DiffusionPipeline
+
+pipeline = DiffusionPipeline.from_pretrained("stable-diffusion-v1-5/stable-diffusion-v1-5")
+```
+
+This however can make it difficult to build applications on top of `diffusers` as you will always have to pass the token around. A potential way to solve this issue is by downloading the weights to a local path `"./stable-diffusion-v1-5"`:
+
+```
+git lfs install
+git clone https://huggingface.co/stable-diffusion-v1-5/stable-diffusion-v1-5
+```
+
+and simply passing the local path to `from_pretrained`:
+
+```python
+from diffusers import StableDiffusionPipeline
+
+pipe = StableDiffusionPipeline.from_pretrained("./stable-diffusion-v1-5")
+```
+
+### Text-to-Image with default PLMS scheduler
+
+```python
+# make sure you're logged in with `hf auth login`
+from diffusers import StableDiffusionPipeline
+
+pipe = StableDiffusionPipeline.from_pretrained("stable-diffusion-v1-5/stable-diffusion-v1-5")
+pipe = pipe.to("cuda")
+
+prompt = "a photo of an astronaut riding a horse on mars"
+image = pipe(prompt).images[0]
+
+image.save("astronaut_rides_horse.png")
+```
+
+### Text-to-Image with DDIM scheduler
+
+```python
+# make sure you're logged in with `hf auth login`
+from diffusers import StableDiffusionPipeline, DDIMScheduler
+
+scheduler =  DDIMScheduler.from_pretrained("CompVis/stable-diffusion-v1-4", subfolder="scheduler")
+
+pipe = StableDiffusionPipeline.from_pretrained(
+    "stable-diffusion-v1-5/stable-diffusion-v1-5",
+    scheduler=scheduler,
+).to("cuda")
+
+prompt = "a photo of an astronaut riding a horse on mars"
+image = pipe(prompt).images[0]
+
+image.save("astronaut_rides_horse.png")
+```
+
+### Text-to-Image with K-LMS scheduler
+
+```python
+# make sure you're logged in with `hf auth login`
+from diffusers import StableDiffusionPipeline, LMSDiscreteScheduler
+
+lms = LMSDiscreteScheduler.from_pretrained("CompVis/stable-diffusion-v1-4", subfolder="scheduler")
+
+pipe = StableDiffusionPipeline.from_pretrained(
+    "stable-diffusion-v1-5/stable-diffusion-v1-5",
+    scheduler=lms,
+).to("cuda")
+
+prompt = "a photo of an astronaut riding a horse on mars"
+image = pipe(prompt).images[0]
+
+image.save("astronaut_rides_horse.png")
+```
+
+### CycleDiffusion using Stable Diffusion and DDIM scheduler
+
+```python
+import requests
+import torch
+from PIL import Image
+from io import BytesIO
+
+from diffusers import CycleDiffusionPipeline, DDIMScheduler
+
+
+# load the scheduler. CycleDiffusion only supports stochastic schedulers.
+
+# load the pipeline
+# make sure you're logged in with `hf auth login`
+model_id_or_path = "CompVis/stable-diffusion-v1-4"
+scheduler = DDIMScheduler.from_pretrained(model_id_or_path, subfolder="scheduler")
+pipe = CycleDiffusionPipeline.from_pretrained(model_id_or_path, scheduler=scheduler).to("cuda")
+
+# let's download an initial image
+url = "https://raw.githubusercontent.com/ChenWu98/cycle-diffusion/main/data/dalle2/An%20astronaut%20riding%20a%20horse.png"
+response = requests.get(url)
+init_image = Image.open(BytesIO(response.content)).convert("RGB")
+init_image = init_image.resize((512, 512))
+init_image.save("horse.png")
+
+# let's specify a prompt
+source_prompt = "An astronaut riding a horse"
+prompt = "An astronaut riding an elephant"
+
+# call the pipeline
+image = pipe(
+    prompt=prompt,
+    source_prompt=source_prompt,
+    image=init_image,
+    num_inference_steps=100,
+    eta=0.1,
+    strength=0.8,
+    guidance_scale=2,
+    source_guidance_scale=1,
+).images[0]
+
+image.save("horse_to_elephant.png")
+
+# let's try another example
+# See more samples at the original repo: https://github.com/ChenWu98/cycle-diffusion
+url = "https://raw.githubusercontent.com/ChenWu98/cycle-diffusion/main/data/dalle2/A%20black%20colored%20car.png"
+response = requests.get(url)
+init_image = Image.open(BytesIO(response.content)).convert("RGB")
+init_image = init_image.resize((512, 512))
+init_image.save("black.png")
+
+source_prompt = "A black colored car"
+prompt = "A blue colored car"
+
+# call the pipeline
+torch.manual_seed(0)
+image = pipe(
+    prompt=prompt,
+    source_prompt=source_prompt,
+    image=init_image,
+    num_inference_steps=100,
+    eta=0.1,
+    strength=0.85,
+    guidance_scale=3,
+    source_guidance_scale=1,
+).images[0]
+
+image.save("black_to_blue.png")
+```

diffusers/schedulers/README.md

@@ -0,0 +1,3 @@
+# Schedulers
+
+For more information on the schedulers, please refer to the [docs](https://huggingface.co/docs/diffusers/api/schedulers/overview).

diffusers/utils/dummy_torch_and_transformers_objects.py

@@ -227,6 +227,36 @@ class Flux3AutoBlocks(metaclass=DummyObject):
         requires_backends(cls, ["torch", "transformers"])
 
 
+class Flux3DistilledAutoBlocks(metaclass=DummyObject):
+    _backends = ["torch", "transformers"]
+
+    def __init__(self, *args, **kwargs):
+        requires_backends(self, ["torch", "transformers"])
+
+    @classmethod
+    def from_config(cls, *args, **kwargs):
+        requires_backends(cls, ["torch", "transformers"])
+
+    @classmethod
+    def from_pretrained(cls, *args, **kwargs):
+        requires_backends(cls, ["torch", "transformers"])
+
+
+class Flux3DistilledModularPipeline(metaclass=DummyObject):
+    _backends = ["torch", "transformers"]
+
+    def __init__(self, *args, **kwargs):
+        requires_backends(self, ["torch", "transformers"])
+
+    @classmethod
+    def from_config(cls, *args, **kwargs):
+        requires_backends(cls, ["torch", "transformers"])
+
+    @classmethod
+    def from_pretrained(cls, *args, **kwargs):
+        requires_backends(cls, ["torch", "transformers"])
+
+
 class Flux3ModularPipeline(metaclass=DummyObject):
     _backends = ["torch", "transformers"]
 

requirements.txt

@@ -13,3 +13,4 @@ pillow
 huggingface_hub>=1.31.0,<2.0
 gradio==5.49.1
 spaces
+aokit