[Agent docs] some updated based on recent integration

.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`).

.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("<name>")` 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_<model>.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

.ai/pipelines.md

@@ -88,3 +88,5 @@ src/diffusers/pipelines/<model>/
        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`.