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:
@@ -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
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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).
|
||||
@@ -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.
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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`.
|
||||
|
||||
Reference in New Issue
Block a user