From b800f828f3951c4d206e154411bcbfa693cd322e Mon Sep 17 00:00:00 2001 From: YiYi Xu Date: Fri, 16 Oct 2026 13:34:11 -1000 Subject: [PATCH] [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 * [docs] FLUX.3: callback block walks all denoising loops Co-Authored-By: Claude Fable 5.1 * [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 * [docs] FLUX.3: show the upsampled caption; clarify the index note in the quickstart Co-Authored-By: Claude Fable 5.1 * [docs] modular: make "Coming from DiffusionPipeline" its own page Co-Authored-By: Claude Fable 5.1 --------- Co-authored-by: Claude Fable 5.1 --- docs/source/en/_toctree.yml | 2 + docs/source/en/api/pipelines/flux3.md | 96 +++++++++++++++---- .../coming_from_diffusion_pipeline.md | 23 +++++ .../loop_sequential_pipeline_blocks.md | 2 + docs/source/en/modular_diffusers/overview.md | 1 + .../source/en/modular_diffusers/quickstart.md | 2 +- docs/source/en/using-diffusers/callback.md | 65 ++++++++++++- 7 files changed, 168 insertions(+), 23 deletions(-) create mode 100644 docs/source/en/modular_diffusers/coming_from_diffusion_pipeline.md diff --git a/docs/source/en/_toctree.yml b/docs/source/en/_toctree.yml index 67854093f2..82a15e66d5 100644 --- a/docs/source/en/_toctree.yml +++ b/docs/source/en/_toctree.yml @@ -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 diff --git a/docs/source/en/api/pipelines/flux3.md b/docs/source/en/api/pipelines/flux3.md index 35eb070d08..b99b4eea4d 100644 --- a/docs/source/en/api/pipelines/flux3.md +++ b/docs/source/en/api/pipelines/flux3.md @@ -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: diff --git a/docs/source/en/modular_diffusers/coming_from_diffusion_pipeline.md b/docs/source/en/modular_diffusers/coming_from_diffusion_pipeline.md new file mode 100644 index 0000000000..b7a991d2d6 --- /dev/null +++ b/docs/source/en/modular_diffusers/coming_from_diffusion_pipeline.md @@ -0,0 +1,23 @@ + + +# 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). diff --git a/docs/source/en/modular_diffusers/loop_sequential_pipeline_blocks.md b/docs/source/en/modular_diffusers/loop_sequential_pipeline_blocks.md index 74a8689227..0a82ad9a4c 100644 --- a/docs/source/en/modular_diffusers/loop_sequential_pipeline_blocks.md +++ b/docs/source/en/modular_diffusers/loop_sequential_pipeline_blocks.md @@ -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. diff --git a/docs/source/en/modular_diffusers/overview.md b/docs/source/en/modular_diffusers/overview.md index f80fff3061..d27542d2fe 100644 --- a/docs/source/en/modular_diffusers/overview.md +++ b/docs/source/en/modular_diffusers/overview.md @@ -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 diff --git a/docs/source/en/modular_diffusers/quickstart.md b/docs/source/en/modular_diffusers/quickstart.md index 2961d6c6ec..8ff3bd8985 100644 --- a/docs/source/en/modular_diffusers/quickstart.md +++ b/docs/source/en/modular_diffusers/quickstart.md @@ -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 diff --git a/docs/source/en/using-diffusers/callback.md b/docs/source/en/using-diffusers/callback.md index 8a44f5e825..1cf4970e53 100644 --- a/docs/source/en/using-diffusers/callback.md +++ b/docs/source/en/using-diffusers/callback.md @@ -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`.