1
0
mirror of https://github.com/huggingface/diffusers.git synced 2026-10-23 09:09:01 +03:00

[docs] FLUX.3: modular-only notes, lazy loading, callbacks (#15069)

* [docs] FLUX.3: modular-only notes, lazy loading, callbacks

Add a tips callout on what differs from DiffusionPipeline (lazy loading,
outputs, guider, offloading, editing blocks), move the prompting advice
into the upsampling section with the pop-the-step example, and add a
per-step callbacks section with a block that takes callback_on_step_end.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

* [docs] FLUX.3: callback block walks all denoising loops

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

* [docs] modular: add "Coming from DiffusionPipeline" to the quickstart and move per-step callbacks to the callbacks guide

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

* [docs] FLUX.3: show the upsampled caption; clarify the index note in the quickstart

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

* [docs] modular: make "Coming from DiffusionPipeline" its own page

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

---------

Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com>
This commit is contained in:
YiYi Xu
2026-10-16 13:34:11 -10:00
committed by GitHub
parent d42ebf9c6e
commit b800f828f3
7 changed files with 168 additions and 23 deletions

View File

@@ -137,6 +137,8 @@
title: Overview
- local: modular_diffusers/quickstart
title: Quickstart
- local: modular_diffusers/coming_from_diffusion_pipeline
title: Coming from DiffusionPipeline
- local: modular_diffusers/modular_diffusers_states
title: States
- local: modular_diffusers/pipeline_block

View File

@@ -16,7 +16,7 @@
FLUX 3 Image Klein is a 5.6B image generation model from Black Forest Labs: a single transformer ([`Flux3Transformer2DModel`]) generates images from a text prompt, alone or together with one or more reference images. Text conditioning comes from Qwen3-VL, and images use the FLUX.2 packed autoencoder ([`AutoencoderKLFlux2`]).
FLUX 3 Image Klein ships as a modular pipeline. Prompting matters: the model responds best to full scene captions that describe the subject, composition, lighting, and atmosphere, so the pipeline runs the [prompt upsampler](#prompt-upsampling) first on every request, text-to-image and editing alike. It is a separate 9B model from Black Forest Labs ([black-forest-labs/flux-3-image-klein-pup](https://huggingface.co/black-forest-labs/flux-3-image-klein-pup)) that expands `prompt` into such a caption before the text encoder, and a short prompt is all it takes. Get started with a text-to-image generation:
FLUX 3 Image Klein ships as a [modular pipeline](../../modular_diffusers/overview). Get started with a text-to-image generation:
```python
import torch
@@ -26,45 +26,84 @@ pipe = ModularPipeline.from_pretrained("black-forest-labs/flux-3-image-klein-bas
pipe.load_components(dtype={"vae": torch.float32, "default": torch.bfloat16})
pipe.to("cuda")
image = pipe(
prompt="A red fox in a snow-covered forest clearing at golden hour",
generator=torch.Generator().manual_seed(0),
output="images",
)[0]
image.save("fox.png")
```
`prompt` is the only required input: the pipeline carries the recommended sampling settings (50 steps, guidance scale 3.5, scheduler shift 6.93) and generates 1024 x 1024 images by default. It takes one prompt per call; use `num_images_per_prompt` for several images of the same prompt.
As you can see in the example above, loading and running differ a little from a `DiffusionPipeline`. Loading is lazy: [`~ModularPipeline.from_pretrained`] only builds the pipeline, and you call [`~ModularPipeline.load_components`] to load the weights. A call returns the pipeline state, and you ask for what you need with `output`: `output="images"` here, or `output=["images", "prompt"]` to also get the caption the prompt upsampler wrote.
> [!TIP]
> If you're working with `ModularPipeline` for the first time, read the [quickstart](../../modular_diffusers/quickstart), and [Coming from DiffusionPipeline](../../modular_diffusers/coming_from_diffusion_pipeline) if you know the standard pipelines. Besides loading and outputs, a few other things differ: [guidance](#guidance-scale), [callbacks](#per-step-callbacks) and [offloading](#memory).
## Prompt upsampling
Prompting matters: FLUX 3 Image Klein takes full scene captions that describe the subject, composition, lighting, and atmosphere, not short prompts; a request sent as typed gives generic output. That is why the pipeline runs the prompt upsampler first on every request, text-to-image and editing alike. It is a separate 9B model from Black Forest Labs ([black-forest-labs/flux-3-image-klein-pup](https://huggingface.co/black-forest-labs/flux-3-image-klein-pup)) that expands `prompt` into such a caption before the text encoder, and a short prompt is all it takes.
The prompt upsampler turns a short request, together with the reference images when editing, into the full scene caption the model takes: medium and style, scene, lighting and camera, each subject with its colors, materials and position in the frame. One model covers text-to-image, editing with reference images and layout with boxes. The caption it wrote comes back as the `prompt` output. In the example above, ask for it alongside the images to see what the image model received:
```python
out = pipe(
prompt="A red fox in a snow-covered forest clearing at golden hour",
generator=torch.Generator().manual_seed(0),
output=["images", "prompt"],
)
out["images"][0].save("fox.png")
print(out["prompt"]) # the caption the image model received
print(out["prompt"])
```
`prompt` is the only required input: the pipeline carries the recommended sampling settings (50 steps, guidance scale 3.5, scheduler shift 6.93) and generates 1024 x 1024 images by default. It takes one prompt per call; use `num_images_per_prompt` for several images of the same prompt.
```
Full-color outdoor wildlife photograph taken from a low eye-level perspective in a snow-covered forest clearing at golden hour. In the center foreground, a red fox stands in profile facing left, its body turned toward the viewer and its head raised as it looks across the clearing. The fox has dense rust-red winter fur, a white throat and chest, black lower legs, pointed ears with dark tips, and a long bushy tail ending in white. ...
```
## Prompt upsampling
The full caption runs to about 240 words; `max_new_tokens` caps its length (3584 by default).
The prompt upsampler turns a short request, together with the reference images when editing, into the full scene caption the model takes: medium and style, scene, lighting and camera, each subject with its colors, materials and position in the frame. One model covers text-to-image, editing with reference images and layout with boxes; the step picks the task from the inputs, like the rest of the pipeline. Its two components, `prompt_upsampler` and `prompt_upsampler_processor`, load from the upsampler repository with the others. The caption it wrote comes back as the `prompt` output, as above; `max_new_tokens` caps its length (3584 by default).
You can print the step's `doc` to see what it takes and what it writes:
To send a caption as given, one you wrote in full or got from another LLM, build the pipeline without the step:
```python
print(pipe.blocks.sub_blocks["prompt_upsampler"].doc)
```
```
Inputs:
prompt (`str`):
The edit instruction to expand into the full caption (one per call).
image (`Image | list`, *optional*):
The reference image(s) of the edit, at most 10; the same ones the pipeline receives.
boxes (`list`, *optional*):
Layout boxes, each {'id', 'desc', 'from', 'src_bbox', 'tgt_bbox'} with the boxes as [y_min, x_min, y_max, x_max] in
0-1: ...
max_new_tokens (`int`, *optional*, defaults to 3584):
The most tokens the upsampler may generate for the caption.
Outputs:
prompt (`str`):
The full caption, replacing the request.
```
Like other blocks in the pipeline, the `prompt_upsampler` block also runs on its own. The recommended way is to pop it out of the blocks: the popped block becomes a standalone upsampler, and the remaining blocks become a pipeline that takes the expanded caption as `prompt`:
```python
blocks = ModularPipeline.from_pretrained("black-forest-labs/flux-3-image-klein-base").blocks
blocks.sub_blocks.pop("prompt_upsampler")
plain = blocks.init_pipeline("black-forest-labs/flux-3-image-klein-base")
plain.load_components(dtype={"vae": torch.float32, "default": torch.bfloat16})
plain.to("cuda")
pup_block = blocks.sub_blocks.pop("prompt_upsampler")
image = plain(prompt=caption, generator=torch.Generator().manual_seed(0), output="images")[0]
```
Any other LLM can write the caption: give it the system prompt for the task from the upsampler repository's `system_prompts.json` (`t2i`, or `i2i` with the references first in the user message) and the request as the user message, and pass what it returns as `prompt` to the pipeline without the step.
The step also runs on its own, to write captions without generating:
```python
pup = pipe.blocks.sub_blocks["prompt_upsampler"].init_pipeline("black-forest-labs/flux-3-image-klein-base")
pup = pup_block.init_pipeline("black-forest-labs/flux-3-image-klein-base")
pup.load_components(dtype=torch.bfloat16)
pup.to("cuda")
caption = pup(prompt="A red fox in a snow-covered forest clearing at golden hour", output="prompt")
pipe = blocks.init_pipeline("black-forest-labs/flux-3-image-klein-base")
pipe.load_components(dtype={"vae": torch.float32, "default": torch.bfloat16})
pipe.to("cuda")
image = pipe(prompt=caption, generator=torch.Generator().manual_seed(0), output="images")[0]
```
This way it can be served as a separate stage, on its own GPU or behind its own endpoint, and a UI can show the caption it wrote and let the user edit it before generating. Any other LLM can take the upsampler's place, with the same system prompts: give it `system_prompts["t2i"]` from the upsampler repository (or `"i2i"`, with the references first in the user message) as the system prompt and the request as the user message, and pass the caption it returns as `prompt` to `pipe`.
## Image-to-image
Pass one or more reference images with `image`. Every reference keeps its own resolution (center-cropped to the 16-pixel grid; references above 16 megapixels are downscaled first), and the generation size defaults to one megapixel at the first reference's aspect ratio. The prompt the image model takes is the *target* image as a complete caption of the reference with the edit folded in, not a one-line instruction; the upsampler sees the references too and writes it from your instruction: what to keep, what to change.
@@ -168,6 +207,21 @@ t2i = ModularPipeline.from_pretrained("black-forest-labs/flux-3-image-klein-base
print(t2i.blocks.doc)
```
## Per-step callbacks
To use an existing callback function, insert a block into the denoising loop. The [callbacks guide](../../using-diffusers/callback#modular-pipelines) has a simple example with the `CallbackOnStepEnd` block and the `walk_denoise_loops` helper:
```python
blocks = Flux3AutoBlocks() # or Flux3DistilledAutoBlocks
for loop in walk_denoise_loops(blocks):
loop.sub_blocks["callback"] = CallbackOnStepEnd()
pipe = blocks.init_pipeline("black-forest-labs/flux-3-image-klein-base")
image = pipe(prompt=caption, callback_on_step_end=my_callback, output="images")[0]
```
To stop early, raise from the callback; there is no `interrupt` flag.
## Memory
Loading everything in bf16 takes about 39 GB of weights: the transformer is 5.5B parameters, the Qwen3-VL-4B text encoder 4B and the [prompt upsampler](#prompt-upsampling) 9B (about 19 GB on its own), with the image VAE in fp32. A pipeline built without the upsampler step (see there) loads about 20 GB. For GPUs with less memory, enable component offloading; the manager moves components between the CPU and the GPU as each step needs them, the upsampler included:

View File

@@ -0,0 +1,23 @@
<!--Copyright 2025 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.
-->
# Coming from DiffusionPipeline
If you are familiar with [`DiffusionPipeline`] and are working with a [`ModularPipeline`] for the first time, a few things work differently. They are listed here so that you have them in mind before you go through the [quickstart](./quickstart).
- **Loading is lazy and flexible.** [`~ModularPipeline.from_pretrained`] builds the pipeline from the repository's `modular_model_index.json` without loading any weights. [`~ModularPipeline.load_components`] then loads them, and it doesn't have to be all at once: by default it loads everything, but you can load only the components you name with `names`, or the set one workflow needs with `workflow`. [`~ModularPipeline.update_components`] lets you reuse a component you already have, or one that differs from the repository's spec. See [Loading components](./modular_pipeline#loading-components) and [Updating components](./modular_pipeline#updating-components).
- **[`ModularPipeline`] reads a different index, `modular_model_index.json`**, but it also works with a repository that only has a `model_index.json`, as long as the model has a modular pipeline. The reverse is not true: [`DiffusionPipeline`] and the `AutoPipeline` classes don't work with a repository that only has a `modular_model_index.json`.
- **A call returns the pipeline state**, which holds the intermediate values too (`prompt_embeds`, `latents`, ...). Ask for what you need with the `output` argument, `output="images"` or `output=["images", "latents"]`, or read an attribute of the state, as in the [quickstart](./quickstart#run-a-pipeline). See [States](./modular_diffusers_states).
- **Classifier-free guidance is supported in the `guider` component**, not a `guidance_scale` call argument: `pipe.update_components(guider=ClassifierFreeGuidance(guidance_scale=5.0))`. See the [guiders guide](../using-diffusers/guiders).
- **There is no `callback_on_step_end`**; however, you can insert a block into the denoising loop to use your existing callback functions. See [Modular pipelines](../using-diffusers/callback#modular-pipelines) in the callbacks guide.
- **The pipeline is a graph of blocks** that you can change (pop, insert, replace) before building it with [`~ModularPipelineBlocks.init_pipeline`]; `pipe.blocks` returns a copy, so edit the blocks first. See [Compose new workflows](./quickstart#compose-new-workflows).
- **Offloading is handled by a [`ComponentsManager`]**, which can serve several pipelines that share components. See the [ComponentsManager guide](./components_manager#offloading).

View File

@@ -90,3 +90,5 @@ Add more loop blocks to run within each iteration with [`~modular_pipelines.Loop
```py
loop = LoopWrapper.from_blocks_dict({"block1": LoopBlock(), "block2": LoopBlock})
```
To add a loop block to an existing pipeline's denoising loop, for example one that runs a `callback_on_step_end` function at every step, see [Modular pipelines](../using-diffusers/callback#modular-pipelines) in the callbacks guide.

View File

@@ -22,6 +22,7 @@ The Modular Diffusers docs are organized as shown below.
## Quickstart
- The [quickstart](./quickstart) shows you how to run a modular pipeline, understand its structure, and customize it by modifying the blocks that compose it.
- [Coming from DiffusionPipeline](./coming_from_diffusion_pipeline) lists what works differently from a [`DiffusionPipeline`], for readers who know the standard pipelines.
## ModularPipelineBlocks

View File

@@ -14,7 +14,7 @@ specific language governing permissions and limitations under the License.
Modular Diffusers is a framework for quickly building flexible and customizable pipelines. These pipelines can go beyond what standard `DiffusionPipeline`s can do. At the core of Modular Diffusers are [`ModularPipelineBlocks`] that can be combined with other blocks to adapt to new workflows. The blocks are converted into a [`ModularPipeline`], a friendly user-facing interface for running generation tasks.
This guide shows you how to run a modular pipeline, understand its structure, and customize it by modifying the blocks that compose it.
This guide shows you how to run a modular pipeline, understand its structure, and customize it by modifying the blocks that compose it. If you are familiar with [`DiffusionPipeline`], read [Coming from DiffusionPipeline](./coming_from_diffusion_pipeline) first for what works differently.
## Run a pipeline

View File

@@ -14,7 +14,7 @@ specific language governing permissions and limitations under the License.
A callback runs at the end of a denoising step and can change pipeline state or tensors for later steps. Use it to adjust attributes or tensor variables for new behavior without rewriting the pipeline.
These callbacks apply to classic [`DiffusionPipeline`] loops. In [Modular Diffusers](../modular_diffusers/overview), you can build and add custom pipeline blocks instead of `callback_on_step_end`.
These callbacks work with classic [`DiffusionPipeline`] loops out of the box. They also work with [Modular Diffusers](../modular_diffusers/overview) pipelines once a block that runs them is inserted into the denoising loop; see [Modular pipelines](#modular-pipelines).
Diffusers provides several callbacks in the pipeline [overview](../api/pipelines/overview#diffusers.callbacks.PipelineCallback).
@@ -137,3 +137,66 @@ image = pipeline(
callback_on_step_end_tensor_inputs=["latents"],
).images[0]
```
## Modular pipelines
Modular pipelines don't carry the `callback_on_step_end` and `interrupt` mechanism. It will come with first-class streaming support instead: [huggingface/diffusers#14159](https://github.com/huggingface/diffusers/pull/14159), not merged yet, adds `pipe.stream()`, which yields the pipeline state at every denoising step and ends the run when you stop iterating. For now, to reuse your existing callback functions, insert a block into the denoising loop: it receives the loop's `BlockState` with the step index and timestep, and can read or edit `latents`.
> [!TIP]
> Edit the blocks *before* [`~modular_pipelines.ModularPipelineBlocks.init_pipeline`]: `pipe.blocks` returns a copy, so changing it after the pipeline exists has no effect.
The block below takes the callback as a pipeline input and calls it with the same signature as above: the pipeline, the step index, the timestep and a dict of tensors, whose returned values are written back. `walk_denoise_loops` finds every denoising loop in the blocks, whichever workflow it belongs to, so the same snippet works for any modular pipeline.
```py
from typing import Callable
import torch
from diffusers import ModularPipeline
from diffusers.modular_pipelines import InputParam, LoopSequentialPipelineBlocks, ModularPipelineBlocks
class CallbackOnStepEnd(ModularPipelineBlocks):
@property
def inputs(self):
return [
InputParam(
"callback_on_step_end",
type_hint=Callable,
description="Called after each scheduler step as `callback(pipe, i, t, {'latents': ...})`; the dict it returns is written back.",
),
InputParam("latents", required=True),
]
def __call__(self, components, block_state, i, t):
if block_state.callback_on_step_end is not None:
outputs = block_state.callback_on_step_end(components, i, t, {"latents": block_state.latents})
block_state.latents = outputs.get("latents", block_state.latents)
return components, block_state
def walk_denoise_loops(block):
"""Every denoising loop in the blocks, whichever workflow it belongs to."""
if isinstance(block, LoopSequentialPipelineBlocks) and "denoiser" in block.sub_blocks:
yield block
for sub in getattr(block, "sub_blocks", {}).values():
yield from walk_denoise_loops(sub)
repo_id = "stabilityai/stable-diffusion-xl-base-1.0"
blocks = ModularPipeline.from_pretrained(repo_id).blocks
for loop in walk_denoise_loops(blocks):
loop.sub_blocks["callback"] = CallbackOnStepEnd() # appended after the last step of the loop
pipe = blocks.init_pipeline(repo_id)
pipe.load_components(dtype=torch.float16)
pipe.to("cuda")
```
`callback_on_step_end` is now an input of the pipeline, and a callback written for a `DiffusionPipeline`, such as `decode_tensors` above, works as is:
```py
image = pipe(prompt="A croissant shaped like a cute bear.", callback_on_step_end=decode_tensors, output="images")[0]
```
[Early stopping](#early-stopping) works differently: a block cannot skip the loop it runs in, so setting `_interrupt` has no effect. Raise from the callback instead: the exception propagates out of `pipe(...)` (no image is decoded), and the pipeline stays usable afterwards.
This works for every modular pipeline except the inner loop of nested loops (autoregressive video models such as Helios and Wan-Animate-2 denoise each chunk inside a block of the outer loop). There, insert the block into the outer loop instead: it runs once per chunk and receives the chunk index `k` in place of `i` and `t`.