diff --git a/.ai/models.md b/.ai/models.md index bdc6464c9aec..b30f43f21db3 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 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 +_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..3353b878e0ad 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. 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 @@ -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 @@ -287,6 +315,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 +330,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`.