From ea36a0ffb95c869c243b0fd1e92b951b9d080259 Mon Sep 17 00:00:00 2001 From: yiyixuxu Date: Wed, 12 Aug 2026 09:03:43 +0000 Subject: [PATCH 1/3] Add agent-doc lessons from the Wan-Animate-2 modular integration - models.md: document _skip_keys (group offloading exclude_kwargs + device_map dispatch skip_keys) with the in-tree cache examples - modular.md: canonical flat-blockset packing (standalone children, InsertableDict groups, no cross-preset-file imports); variant presets carry their own model_name + mapping entry (#14451); composed-blockset inputs/outputs overrides; guider ownership and requires_unconditional_embeds; always ship modular_model_index.json; auto-docstring drift + check_forward_call_docstrings checklist items; randn_tensor gotcha - pipelines.md: the same randn_tensor gotcha Co-Authored-By: Claude Fable 5 --- .ai/models.md | 14 ++++++++++++++ .ai/modular.md | 36 +++++++++++++++++++++++++++++++++++- .ai/pipelines.md | 2 ++ 3 files changed, 51 insertions(+), 1 deletion(-) diff --git a/.ai/models.md b/.ai/models.md index bdc6464c9aec..1051071d8c74 100644 --- a/.ai/models.md +++ b/.ai/models.md @@ -141,6 +141,20 @@ _keep_in_fp32_modules = ["time_embedder", "scale_shift_table", "norm1", "norm2", If `None` (default), all modules follow the requested `torch_dtype`. +### `_skip_keys` + +**API:** every device-placement hook that moves call inputs — `apply_group_offloading(model, ...)` reads it as `exclude_kwargs` (`hooks/group_offloading.py`), and `from_pretrained(..., device_map=...)` (including the `device_map={"": "cpu"}` offload placement) forwards it to accelerate's `dispatch_model(skip_keys=...)` (`pipelines/pipeline_loading_utils.py`). + +These hooks move every tensor in the call's args/kwargs to the execution device before `forward` runs. List the forward kwargs that carry non-tensor runtime objects holding large tensors — a KV-cache, an encoder feature cache — so the hooks leave them alone instead of transferring (or repeatedly re-transferring) their contents on every call. + +```python +# in-tree examples +_skip_keys = ["kv_cache"] # transformer_wan_animate_2.py, transformer_flux2.py +_skip_keys = ["feat_cache", "feat_idx"] # autoencoder_kl_wan.py +``` + +If unset, offloading treats every kwarg as movable input, which at best wastes transfers and at worst breaks the object's device placement mid-inference. + ### `_cp_plan` **API:** `model.enable_parallelism(config=parallel_config)` — when the config includes `context_parallel_config`, this plan is used by `apply_context_parallel()` to shard tensors across GPUs for sequence parallelism (`modeling_utils.py:1665`). diff --git a/.ai/modular.md b/.ai/modular.md index b68129d08536..91d189565533 100644 --- a/.ai/modular.md +++ b/.ai/modular.md @@ -10,7 +10,7 @@ When adding a new modular pipeline (or reviewing one), skim `src/diffusers/modul This section provides guidance on how to execute pipelines and blocks — in scripts, debugging sessions, and tests alike. -- **Full pipeline from a repo**: `ModularPipeline.from_pretrained(repo_id)` — the base class, not the model subclass; it resolves the right class from the repo's `modular_model_index.json` (falling back to a standard `model_index.json`). Then `pipe.load_components()` and call it. +- **Full pipeline from a repo**: `ModularPipeline.from_pretrained(repo_id)` — the base class, not the model subclass; it resolves the right class from the repo's `modular_model_index.json` (falling back to a standard `model_index.json`). Then `pipe.load_components()` and call it. A repo meant for `ModularPipeline.from_pretrained` should always ship a `modular_model_index.json` (its `_class_name` routes to the right pipeline class); the `model_index.json` fallback is a last resort with limited support. - **A single block or sub-workflow**: convert it to a pipeline first with `init_pipeline()`. Blocks are never executed directly. ```python @@ -185,6 +185,8 @@ class AutoDenoise(ConditionalPipelineBlocks): A different checkpoint (distilled / turbo / a variant with its own schedule) can have its own blockset mapped to it: give the variant a `ModularPipeline` subclass carrying its `default_blocks_name`, and checkpoints route to it automatically — via `_class_name` in `modular_model_index.json`, or, for repos that only ship a standard `model_index.json`, a config-keyed map fn in `MODULAR_PIPELINE_MAPPING` (see `_flux2_klein_map_fn`). +A variant blockset also declares its **own `model_name`** (every block class carries one), with its own `_create_default_map_fn` entry in `MODULAR_PIPELINE_MAPPING` — e.g. `wan-animate-2` / `wan-animate-2-distilled`. Sharing the base's name means `blocks.init_pipeline()` resolves the *base* pipeline class (the map fn gets no config on that path) and `save_pretrained` then round-trips the wrong `_class_name` — see huggingface/diffusers#14451. + Default to taking that option. The only reason not to split is when the variant behaves literally the same. If the split buys anything at all — the distilled variant doesn't have to declare `negative_prompt`, doesn't carry a guider, and its docs describe exactly what the checkpoint does — make the separate blockset. It costs almost nothing: blocksets compose the same shared leaf blocks, and only the steps that truly differ need new block classes. See `modular_blocks_flux2_klein.py`, which reuses the base flux2 leaf blocks and swaps in just a `negative_prompt`-free text encoder and a guider-free denoise step. Don't fall back to the standard-pipeline habit of a config flag branching inside a shared block (`ConfigSpec(name="is_distilled")` + `if components.config.is_distilled:`). That keeps both variants' behavior bundled in one blockset — and the input surface is the one thing it can never fix: a repo can override components and config values per checkpoint, but never which inputs the blocks declare, so the distilled checkpoint would still accept `negative_prompt` and silently ignore it. @@ -209,6 +211,20 @@ Standard pipelines accept `prompt_embeds` / `image_latents` as `__call__` inputs Prefer flat sequences over nested compositions. Put the `Auto` / `Conditional` selection at the top level and make each workflow variant a flat `InsertableDict` of leaf blocks. Try not to nest `AutoPipelineBlocks` inside `SequentialPipelineBlocks` inside `AutoPipelineBlocks` — debugging which workflow was selected, and which block inside which sub-block touched which state, becomes painful. See `flux2/modular_blocks_flux2_klein.py` for the canonical shape. +The default blockset's top-level children are exactly the steps worth running standalone — `text_encoder` / `image_encoder` / `vae_encoder` / `denoise` / `decode` — each poppable and usable on its own. Multi-step children (a preprocess + encode pair, a prepare + loop pair) are assembled as a module-level `InsertableDict` plus a `SequentialPipelineBlocks` reading it: + +```python +MyImageEncoderBlocks = InsertableDict([("preprocess", MyProcessImagesInputStep()), ("encode", MyImageClipEncoderStep())]) + +# auto_docstring +class MyImageEncodeStep(SequentialPipelineBlocks): + model_name = "my-model" + block_classes = MyImageEncoderBlocks.values() + block_names = MyImageEncoderBlocks.keys() +``` + +Preset files never import from each other: each `modular_blocks_*.py` self-assembles its groups from the leaf files (`encoders.py`, `denoise.py`, ...), even when a group is identical to the sibling preset's — see `modular_blocks_wan_animate_2_distilled.py`, `modular_blocks_flux2_klein.py`. + ## InputParam / OutputParam Use `.template("")` for params with a canonical meaning (`prompt`, `negative_prompt`, `image`, `generator`, `num_inference_steps`, `latents`, `prompt_embeds`, `images`, `videos`, etc.) — the template carries a vetted description and type hint. The full registry lives in [`src/diffusers/modular_pipelines/modular_pipeline_utils.py`](../src/diffusers/modular_pipelines/modular_pipeline_utils.py) (`INPUT_PARAM_TEMPLATES`, `OUTPUT_PARAM_TEMPLATES`); read that file rather than relying on a hardcoded list here, since names get added. @@ -248,6 +264,18 @@ if block_state.num_frames is None: A declared default is part of the block's contract, so the assembled pipeline is aware of it: the generated docstring shows it and `default_call_parameters` reports it. Resolved inside the body instead, the input renders as `*optional*` with no default, and nothing at the pipeline level can report what the block will actually do. Don't worry about branches of a conditional blockset declaring different defaults for the same input — each branch resolves its own at runtime. Resolve inside `__call__` only when the default is *computed* — derived from other inputs or component config (`height = components.default_sample_size * components.vae_scale_factor`). And when several blocks in a sequence share an input, declare the same default on each (or only on the first block that reads it): in a sequence the input is one shared value, so disagreeing declarations are silently resolved first-block-wins. +**A composed `SequentialPipelineBlocks` can override `inputs` / `outputs`.** Two uses: narrow `outputs` to what downstream actually consumes, so the docstring shows the step's product instead of every internal intermediate (Wan-Animate-2's core denoise exposes only `segment_frames`); and change one input default per preset by mapping over `super().inputs` (the distilled core denoise turns `num_inference_steps` into `default=10`): + +```python +@property +def inputs(self): + # The distilled checkpoint samples in few steps. + return [ + InputParam.template("num_inference_steps", default=10) if param.name == "num_inference_steps" else param + for param in super().inputs + ] +``` + ## ComponentSpec patterns ```python @@ -264,6 +292,8 @@ ComponentSpec( ) ``` +The guider spec is declared **only by the denoise blocks** — each preset's denoise step pins its scale there (base `FrozenDict({"guidance_scale": 3.0})`, distilled `1.0`), which is how one repo's checkpoint gets the right guidance with no `guidance_scale` input anywhere. The text encoder never declares a guider: it asks `components.requires_unconditional_embeds` (a pipeline-class property that consults the guider when one exists) to decide whether to encode the negative prompt, defaulting it to `""`; an explicitly passed `negative_prompt` is the standalone escape hatch. See `wan_animate_2/encoders.py` / `denoise.py`. + ## Gotchas 1. **Importing from standard pipelines.** The modular and standard pipeline systems are parallel — modular blocks must not import from `diffusers.pipelines.*`. For shared utility methods (e.g. `_pack_latents`, `retrieve_timesteps`), either redefine as standalone functions or use `# Copied from diffusers.pipelines....` headers. See `wan/before_denoise.py` and `helios/before_denoise.py` for examples. @@ -287,6 +317,8 @@ ComponentSpec( 9. **Serving a checkpoint variant through a config flag in a shared block.** `ConfigSpec(name="is_distilled")` plus `if components.config.is_distilled:` bundles two checkpoints' behavior into one blockset — and it can't change the input surface at all (the distilled variant would still accept `negative_prompt`). Suggest a separate blockset for the variant instead (see Key pattern: Checkpoint variants). +10. **Raw `torch.randn(device=...)` for noise.** Use `randn_tensor(...)` from `utils/torch_utils`: it draws on the generator's device and moves the result, so CPU generators (what the test mixins pass) work, and the CUDA-generator path is bit-identical to `torch.randn`. + ## Conversion checklist - [ ] Read original pipeline's `__call__` end-to-end, map stages @@ -300,5 +332,7 @@ ComponentSpec( - [ ] Assemble blocks in `modular_blocks_.py` - [ ] Wire up `__init__.py` with lazy imports - [ ] Add `# auto_docstring` above all assembled blocks (SequentialPipelineBlocks, AutoPipelineBlocks, etc.), run `python utils/modular_auto_docstring.py --fix_and_overwrite`, and verify the generated docstrings — all parameters should have proper descriptions with no "TODO" placeholders indicating missing definitions +- [ ] `--fix_and_overwrite` regenerates **every** modular family — revert the files outside your model's folder before committing (careful with path filters: `grep -v pipelines/wan/` also matches `modular_pipelines/wan/`) +- [ ] `python utils/check_forward_call_docstrings.py` must pass (CI gates it): every `forward` / `__call__` parameter needs its own docstring entry (no `a (…), b (…):` fused entries) and a `Returns:` section - [ ] Run `make style` and `make quality` - [ ] Test all workflows for parity with reference diff --git a/.ai/pipelines.md b/.ai/pipelines.md index 7384808a1cda..f0a773962c37 100644 --- a/.ai/pipelines.md +++ b/.ai/pipelines.md @@ -88,3 +88,5 @@ src/diffusers/pipelines// callback_kwargs[k] = locals()[k] ``` The bug is invisible until someone actually passes `callback_on_step_end` — the `PipelineTesterMixin` callback tests are what catch it. + +10. **Raw `torch.randn(device=...)` for noise.** Use `randn_tensor(...)` from `utils/torch_utils`: it draws on the generator's device and moves the result, so CPU generators work, and the CUDA-generator path is bit-identical to `torch.randn`. From 6d9541132319097249e05607f6091a5c907d2cdd Mon Sep 17 00:00:00 2001 From: YiYi Xu Date: Thu, 13 Aug 2026 10:34:46 -1000 Subject: [PATCH 2/3] Apply suggestions from code review Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com> --- .ai/modular.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.ai/modular.md b/.ai/modular.md index 91d189565533..463f171b9e29 100644 --- a/.ai/modular.md +++ b/.ai/modular.md @@ -10,7 +10,7 @@ When adding a new modular pipeline (or reviewing one), skim `src/diffusers/modul This section provides guidance on how to execute pipelines and blocks — in scripts, debugging sessions, and tests alike. -- **Full pipeline from a repo**: `ModularPipeline.from_pretrained(repo_id)` — the base class, not the model subclass; it resolves the right class from the repo's `modular_model_index.json` (falling back to a standard `model_index.json`). Then `pipe.load_components()` and call it. A repo meant for `ModularPipeline.from_pretrained` should always ship a `modular_model_index.json` (its `_class_name` routes to the right pipeline class); the `model_index.json` fallback is a last resort with limited support. +- **Full pipeline from a repo**: `ModularPipeline.from_pretrained(repo_id)` — the base class, not the model subclass; it resolves the right class from the repo's `modular_model_index.json` (falling back to a standard `model_index.json`). Then `pipe.load_components()` and call it. New modular repositories should include `modular_model_index.json` because it records modular block and component metadata. `model_index.json` remains supported for compatibility with standard repositories, but it cannot express all modular metadata. - **A single block or sub-workflow**: convert it to a pipeline first with `init_pipeline()`. Blocks are never executed directly. ```python From 49827631518d358086b198c9d6f41bbce3407043 Mon Sep 17 00:00:00 2001 From: yiyixuxu Date: Thu, 13 Aug 2026 20:50:48 +0000 Subject: [PATCH 3/3] Address review comments MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Reword the `_skip_keys` guidance: "non-tensor runtime objects holding large tensors" contradicted itself, say "state the model places itself" instead. Drop the paragraph claiming the guider spec is declared only by the denoise blocks — `WanTextEncoderStep` declares one too, and the codebase is inconsistent enough that no rule can be stated yet (see #14469). Co-Authored-By: Claude Opus 5 --- .ai/models.md | 2 +- .ai/modular.md | 2 -- 2 files changed, 1 insertion(+), 3 deletions(-) diff --git a/.ai/models.md b/.ai/models.md index 1051071d8c74..b30f43f21db3 100644 --- a/.ai/models.md +++ b/.ai/models.md @@ -145,7 +145,7 @@ If `None` (default), all modules follow the requested `torch_dtype`. **API:** every device-placement hook that moves call inputs — `apply_group_offloading(model, ...)` reads it as `exclude_kwargs` (`hooks/group_offloading.py`), and `from_pretrained(..., device_map=...)` (including the `device_map={"": "cpu"}` offload placement) forwards it to accelerate's `dispatch_model(skip_keys=...)` (`pipelines/pipeline_loading_utils.py`). -These hooks move every tensor in the call's args/kwargs to the execution device before `forward` runs. List the forward kwargs that carry non-tensor runtime objects holding large tensors — a KV-cache, an encoder feature cache — so the hooks leave them alone instead of transferring (or repeatedly re-transferring) their contents on every call. +These hooks move every tensor in the call's args/kwargs to the execution device before `forward` runs. List the forward kwargs that carry state the model places itself — a KV-cache, an encoder feature cache — so the hooks leave them alone instead of transferring (or repeatedly re-transferring) their contents on every call. ```python # in-tree examples diff --git a/.ai/modular.md b/.ai/modular.md index 463f171b9e29..3353b878e0ad 100644 --- a/.ai/modular.md +++ b/.ai/modular.md @@ -292,8 +292,6 @@ ComponentSpec( ) ``` -The guider spec is declared **only by the denoise blocks** — each preset's denoise step pins its scale there (base `FrozenDict({"guidance_scale": 3.0})`, distilled `1.0`), which is how one repo's checkpoint gets the right guidance with no `guidance_scale` input anywhere. The text encoder never declares a guider: it asks `components.requires_unconditional_embeds` (a pipeline-class property that consults the guider when one exists) to decide whether to encode the negative prompt, defaulting it to `""`; an explicitly passed `negative_prompt` is the standalone escape hatch. See `wan_animate_2/encoders.py` / `denoise.py`. - ## Gotchas 1. **Importing from standard pipelines.** The modular and standard pipeline systems are parallel — modular blocks must not import from `diffusers.pipelines.*`. For shared utility methods (e.g. `_pack_latents`, `retrieve_timesteps`), either redefine as standalone functions or use `# Copied from diffusers.pipelines....` headers. See `wan/before_denoise.py` and `helios/before_denoise.py` for examples.