update agent doc to cover more on tests + include it in review CI scope

.ai/AGENTS.md

@@ -30,6 +30,7 @@ Strive to write code as simple and explicit as possible.
 - **Models** — see [models.md](models.md) for model conventions, attention pattern, implementation rules, dependencies, and gotchas. For adding or converting a model, use the [model-integration](./skills/model-integration/SKILL.md) skill.
 - **Pipelines** — see [pipelines.md](pipelines.md) for pipeline conventions, patterns, and gotchas.
 - **Modular pipelines** — see [modular.md](modular.md) for modular pipeline conventions, patterns, and gotchas.
+- **Tests** — see [testing.md](testing.md) for test conventions: required test layers, tester mixins, and dummy-component rules.
 
 ## Skills
 

.ai/modular.md

@@ -6,6 +6,30 @@ Shared reference for modular pipeline conventions, patterns, and gotchas.
 
 When adding a new modular pipeline (or reviewing one), skim `src/diffusers/modular_pipelines/qwenimage/`, `src/diffusers/modular_pipelines/flux2/`, `src/diffusers/modular_pipelines/wan/`, and `src/diffusers/modular_pipelines/helios/` first to establish the pattern. Most conventions (file split between `encoders.py` / `before_denoise.py` / `denoise.py` / `decoders.py`, how `expected_components` / `inputs` / `intermediate_outputs` are declared, the denoise-loop wrapping with `LoopSequentialPipelineBlocks`, top-level assembly via `AutoPipelineBlocks` / `SequentialPipelineBlocks` in `modular_blocks_<model>.py`, the `ModularPipeline` subclass shape, the guider-abstracted denoise body, `kwargs_type="denoiser_input_fields"` plumbing) are easiest to internalize by comparison rather than from a fixed list.
 
+## Running a modular pipeline
+
+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 single block or sub-workflow**: convert it to a pipeline first with `init_pipeline()`. Blocks are never executed directly.
+
+```python
+# one block
+pipe = MyTextEncoderStep().init_pipeline("some-org/tiny-model")  # repo optional if the block needs no pretrained components
+pipe.load_components()                                           # init_pipeline only wires specs; this materializes components
+
+# a chain of blocks
+blocks = SequentialPipelineBlocks.from_blocks_dict({"vision": VisionStep(), "sound": SoundStep()})
+pipe = blocks.init_pipeline()
+
+# run it: declared InputParams are call kwargs; `output=` selects what comes back
+ids = pipe(prompt="a robot", output="cond_input_ids")   # one value (any declared output/intermediate)
+state = pipe(prompt="a robot")                          # or the full state — read values with state.get("name")
+```
+
+- **Swap components and config values** with `pipe.update_components(scheduler=new_scheduler, my_config_flag=False)` — it handles both, keeping the specs and the saved `modular_model_index.json` in sync. Read config via `pipe.config.<name>` (direct attribute access is deprecated).
+- **Don't call a block directly** (`block(components, state)`) and don't hand-build a `PipelineState` to feed it. That is the executor's internal protocol — it only *appears* to work for blocks that never touch `components`, and breaks the moment the block gains a component or config dependency. If you find yourself constructing a `PipelineState`, you want `init_pipeline()` and a normal call instead.
+
 ## File structure
 
 ```
@@ -185,7 +209,7 @@ ComponentSpec(
 
 5. **Using `InputParam.template()` / `OutputParam.template()` when semantics don't match.** Templates carry predefined descriptions — e.g. the `"latents"` output template means "Denoised latents". Don't use it for initial noisy latents from a prepare-latents step. Use a plain `InputParam(...)` / `OutputParam(...)` with an accurate description instead.
 
-6. **Test model paths pointing to contributor repos.** Tiny test models must live under `hf-internal-testing/`, not personal repos like `username/tiny-model`. Move the model before merge.
+6. **Test model paths pointing to contributor repos.** Tiny test models ultimately live under `hf-internal-testing/`, not personal repos like `username/tiny-model`. Developing against a personal repo is fine and not merge-blocking — a maintainer moves the model (before or after merge) and updates the path.
 
 7. **Respect the declared IO system.** Components in `expected_components`, fields in `inputs` / `intermediate_outputs` — once declared, the modular framework guarantees them. So:
     - **Don't read defensively.** Declared components are always set as attributes (possibly `None`); declared upstream outputs are always populated in `block_state` after the upstream block runs. `getattr(components, "vae", None)`, `hasattr(self, "vae")`, `getattr(block_state, "prompt_embeds", None)` are dead code that hides typos. Use `components.vae` / `block_state.prompt_embeds` directly. Check `is not None` only when nullability is meaningful (a component the user might not have loaded).

.ai/pipelines.md

@@ -80,3 +80,11 @@ src/diffusers/pipelines/<model>/
 7. **Don't modify the state of a registered component on the fly.** From inside `__call__` or other helper methods, don't change the state of `self.text_encoder` / `self.transformer` / `self.vae` — no in-place `.to(dtype/device)`, no setting attributes/buffers or swapping submodules. Components are shared and routinely reused across pipelines, so a per-call mutation may silently change another pipeline's outputs. You should pass a component that's already in the right state, and document that expectation explicitly. Only when that's genuinely inconvenient and you must change state for the duration of a call — e.g. swapping in an attention processor — save the original first and restore it before returning, so the component is left exactly as you found it. The PAG pipelines are the reference for this: `pipeline_pag_sd.py` snapshots `original_attn_proc = self.unet.attn_processors`, installs the PAG processors for the denoising loop, then calls `self.unet.set_attn_processor(original_attn_proc)` at the end of `__call__`.
 
 8. **Don't reimplement `DiffusionPipeline`.** A pipeline subclass adds only *pipeline-specific* steps (`__call__`, `check_inputs`, `encode_prompt`, `prepare_latents`, …). Device placement, offloading, and component loading/registration already live on the base class — don't add your own; use what's there.
+
+9. **Build `callback_kwargs` with a loop, never a dict comprehension.** `{k: locals()[k] for k in callback_on_step_end_tensor_inputs}` always raises `KeyError`: inside a comprehension, `locals()` is the comprehension's own scope, not `__call__`'s. Use the standard form (see `pipeline_stable_diffusion.py`):
+   ```python
+   callback_kwargs = {}
+   for k in callback_on_step_end_tensor_inputs:
+       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.

.ai/review-rules.md

@@ -7,6 +7,7 @@ Before reviewing, read and apply the guidelines in:
 - [models.md](models.md) — model conventions, attention pattern, implementation rules, dependencies, gotchas
 - [pipelines.md](pipelines.md) — pipeline conventions, coding style, gotchas
 - [modular.md](modular.md) — modular pipeline conventions, patterns, common mistakes
+- [testing.md](testing.md) — test conventions: required test layers, tester mixins, dummy-component rules. When a PR adds or changes tests, check them against this guide.
 - [skills/model-integration/pitfalls.md](skills/model-integration/pitfalls.md) — known pitfalls causing numerical discrepancies between the reference implementation and the diffusers port (dtype mismatches, config assumptions, etc.)
 
 ## Common mistakes

.ai/skills/model-integration/SKILL.md

@@ -87,41 +87,13 @@ Common conversion patterns to watch for model-level components:
 
 ## Testing
 
-Two test layers must be added for any new pipeline: pipeline-level tests, and (if a new model is introduced) model-level tests. Integration/slow tests and LoRA tests are **not** added in the initial PR — they come later, after discussion with maintainers.
-
-**General rules (apply to both layers):**
-- Keep component sizes tiny so the suite runs fast — small `num_layers`, small hidden/attention dims, low resolution, few frames. Reference `tests/pipelines/wan/test_wan.py` (`get_dummy_components` and `get_dummy_inputs`) for the size scale to target.
-- No LoRA tests in the initial PR (no `LoraTesterMixin`, no `tests/lora/test_lora_layers_<model>.py`).
-- No integration / slow tests in the initial PR — don't add anything gated on `@slow` / `RUN_SLOW=1` yet.
-
-### Pipeline-level tests
-
-- Location: `tests/pipelines/<model>/test_<model>.py` (one file per pipeline variant, e.g. T2V, I2V).
-- Subclass both `PipelineTesterMixin` (from `..test_pipelines_common`) and `unittest.TestCase`.
-- Set `pipeline_class`, `params`, `batch_params`, `image_params` from `..pipeline_params`, and any `required_optional_params` / capability flags (`test_xformers_attention`, `supports_dduf`, etc.) that apply.
-- Implement `get_dummy_components()` (build all sub-modules with tiny configs and a fixed `torch.manual_seed(0)` before each) and `get_dummy_inputs(device, seed=0)`.
-- Skip any inherited tests that don't apply with `@unittest.skip("Test not supported")` rather than deleting them.
-- Reference: `tests/pipelines/wan/test_wan.py`.
-
-### Model-level tests
-
-Only required if the pipeline introduces a new model class (transformer, VAE, etc.). Don't write these by hand — generate them (example command below):
-
-```bash
-python utils/generate_model_tests.py src/diffusers/models/transformers/transformer_<model>.py
-```
-
-- Run with **no `--include` flags** initially. The generator auto-detects mixins/attributes and emits the always-on testers (`ModelTesterMixin`, `MemoryTesterMixin`, `TorchCompileTesterMixin`, plus `AttentionTesterMixin` / `ContextParallelTesterMixin` / `TrainingTesterMixin` as applicable). Optional testers (quantization, caching, single-file, IP adapter, etc.) are added later, after maintainer discussion.
-- The generator writes to `tests/models/transformers/test_models_transformer_<model>.py` (or the matching `unets/` / `autoencoders/` subdir).
-- Fill in the `TODO`s in the generated `<Model>TesterConfig`: `pretrained_model_name_or_path`, `get_init_dict()` (tiny config), `get_dummy_inputs()`, `input_shape`, `output_shape`. Keep init dims small for speed.
-- Do **not** add `LoraTesterMixin` at the start, even if the model subclasses `PeftAdapterMixin` — strip it from the generated file for the initial PR.
-- Reference: `tests/models/transformers/test_models_transformer_flux.py`.
+Two test layers must be added for any new pipeline: pipeline-level tests, and (if a new model is introduced) model-level tests. Conventions for both layers — file locations, tester mixins, dummy-component rules — live in [testing.md](../../testing.md); follow it when writing the tests.
 
 ## Model parity test
 
 Confirm the diffusers implementation matches the reference. Test each component on **CPU/float32** with a strict tolerance (`max_diff < 1e-3`), comparing the **freshly converted** weights against the reference in a single script — both sides side by side, nothing saved to disk in between. See [pitfalls.md](pitfalls.md) for the common sources of numerical discrepancy.
 
-This is an **internal verification tool for integration — it should not be shipped in the PR** (it imports the reference repo). The tests that ship with the PR are the model-level and pipeline-level tests in **Testing**.
+This is an **internal verification tool for integration — it should not be shipped in the PR** (it imports the reference repo). The tests that ship with the PR are the model-level and pipeline-level tests in [testing.md](../../testing.md).
 
 The example below is schematic (placeholder names). `ReferenceModel` is the component **imported from the original repo**, and `convert_my_component` is **the same conversion function you wrote for the conversion script for the component**. You should make sure both load the *same* checkpoint weights and run the *same* input, so any difference is a conversion or implementation bug — not a difference in inputs.
 

.ai/skills/self-review/SKILL.md

@@ -12,8 +12,9 @@ description: >
 
 Runs the same rubric as the `@claude` CI reviewer, so you catch issues before a
 maintainer does — but over your **whole** PR diff. (The CI scopes itself to
-`src/diffusers/` and `.ai/`; for your own PR, also review your tests, docs, and
-scripts.) You're already on the branch with the conventions loaded, so: get the
+`src/diffusers/`, `tests/`, and `.ai/`; for your own PR, also review your docs
+and scripts.) You're already on the branch with the conventions loaded, so: get
+the
 diff → review it against the rubric → report → iterate with the contributor
 until it's ready, then remind them to share the final notes on the PR.
 

.ai/testing.md

@@ -0,0 +1,49 @@
+# Testing
+
+Test conventions for new models and pipelines: what a PR must ship, and what to check existing test files against.
+
+Two test layers must be added for any new pipeline: pipeline-level tests, and (if a new model is introduced) model-level tests. Integration/slow tests and LoRA tests are **not** added in the initial PR — they come later, after discussion with maintainers.
+
+## General rules (apply to all layers)
+
+- Keep component sizes tiny so the suite runs fast — small `num_layers`, small hidden/attention dims, low resolution, few frames. Reference `tests/pipelines/wan/test_wan.py` (`get_dummy_components` and `get_dummy_inputs`) for the size scale to target.
+- Build dummy components from the **real classes** at tiny config — a real VAE with tiny dims, a real tokenizer from an `hf-internal-testing/tiny-random-*` repo. Don't substitute a hand-rolled mock (a bare `nn.Module` with a `SimpleNamespace` config, a fake tokenizer) without a good reason: a mock is written by copying whatever the pipeline reads from the component today, so it can only confirm the pipeline against itself — the test stays green when the component renames a config field or the pipeline starts reading one the component doesn't have, and catching exactly that pipeline↔component contract is what a pipeline test is for. A good reason to stub: the component is impractical to instantiate and only its I/O matters to the pipeline (e.g. `DummyCosmosSafetyChecker` standing in for the huge Cosmos guardrail) — then make it a shared, purpose-built class honoring the real interface.
+- The same applies to test doubles at the call level: don't monkeypatch a component method (e.g. the scheduler's `set_timesteps`) just to capture what the code under test passed to it — that only verifies the caller against itself, not against the real method's contract. Call the real component and assert on its resulting state.
+- No LoRA tests in the initial PR (no `LoraTesterMixin`, no `tests/lora/test_lora_layers_<model>.py`).
+- No integration / slow tests in the initial PR — don't add anything gated on `@slow` / `RUN_SLOW=1` yet.
+
+## Pipeline-level tests 
+
+### Stanard pipelines
+
+- Location: `tests/pipelines/<model>/test_<model>.py` (one file per pipeline variant, e.g. T2V, I2V).
+- Subclass both `PipelineTesterMixin` (from `..test_pipelines_common`) and `unittest.TestCase`.
+- Set `pipeline_class`, `params`, `batch_params`, `image_params` from `..pipeline_params`, and any `required_optional_params` / capability flags (`test_xformers_attention`, `supports_dduf`, etc.) that apply.
+- Implement `get_dummy_components()` (build all sub-modules with tiny configs and a fixed `torch.manual_seed(0)` before each) and `get_dummy_inputs(device, seed=0)`.
+- Skip any inherited tests that don't apply with `@unittest.skip("Test not supported")` rather than deleting them.
+- Reference: `tests/pipelines/wan/test_wan.py`.
+
+### Modular pipelines
+
+- Location: `tests/modular_pipelines/<model>/test_modular_pipeline_<model>.py` (one test class per blocks assembly / pipeline variant).
+- Subclass `ModularPipelineTesterMixin` (from `..test_modular_pipelines_common`) — it runs the pipeline end-to-end (call signature, batch consistency, float16, device placement) against a tiny checkpoint.
+- Set `pipeline_class`, `pipeline_blocks_class`, `pretrained_model_name_or_path`, `params` / `batch_params`, and implement `get_dummy_inputs(seed=0)`. Set `expected_workflow_blocks` to pin the block name → class ordering per workflow.
+- `pretrained_model_name_or_path` is a tiny repo with real components (tiny transformer, real scheduler / VAE / tokenizer configs). Develop against a personal repo; tiny repos ultimately live under `hf-internal-testing/` — not merge-blocking, a maintainer moves it before or after merge.
+- **The tiny repo must mirror the real checkpoint's shape** — same index file type, same pipeline-level config keys, a scheduler configured like the real one. A fixture that doesn't look like the published repos tests a loading/config path no user will ever hit, while the path users *do* hit stays uncovered. If the model ships variants with different configs (base/distilled, different schedules), make one tiny repo and test class per variant — see the flux2 klein base/distilled split.
+- **Bespoke tests go on the tester class as methods**, not as module-level functions — the mixin is pytest-style, so fixtures (`tmp_path`, `pytest.raises`, parametrize) all work in methods.
+- **Test a block's behavior by running it as a pipeline** — `init_pipeline()` → `load_components()` → call it and assert on outputs (see "Running a modular pipeline" in [modular.md](modular.md)). Config-dependent behavior: flip the value with `update_components(...)` and compare real outputs across the two runs. Input validation: `pytest.raises` around a normal `pipe(...)` call. Don't call `block(components, state)` directly or hand-build a `PipelineState`, and don't assert on declared specs (`inputs` / `intermediate_outputs` name lists) — declarations aren't behavior, and `expected_workflow_blocks` already pins the structure.
+- Reference: `tests/modular_pipelines/flux2/test_modular_pipeline_flux2_klein.py` (plus `..._klein_base.py` for the base/distilled variant split).
+
+## Model-level tests
+
+Only required if the pipeline introduces a new model class (transformer, VAE, etc.). Don't write these by hand — generate them (example command below):
+
+```bash
+python utils/generate_model_tests.py src/diffusers/models/transformers/transformer_<model>.py
+```
+
+- Run with **no `--include` flags** initially. The generator auto-detects mixins/attributes and emits the always-on testers (`ModelTesterMixin`, `MemoryTesterMixin`, `TorchCompileTesterMixin`, plus `AttentionTesterMixin` / `ContextParallelTesterMixin` / `TrainingTesterMixin` as applicable). Optional testers (quantization, caching, single-file, IP adapter, etc.) are added later, after maintainer discussion.
+- The generator writes to `tests/models/transformers/test_models_transformer_<model>.py` (or the matching `unets/` / `autoencoders/` subdir).
+- Fill in the `TODO`s in the generated `<Model>TesterConfig`: `pretrained_model_name_or_path`, `get_init_dict()` (tiny config), `get_dummy_inputs()`, `input_shape`, `output_shape`. Keep init dims small for speed.
+- Do **not** add `LoraTesterMixin` at the start, even if the model subclasses `PeftAdapterMixin` — strip it from the generated file for the initial PR.
+- Reference: `tests/models/transformers/test_models_transformer_flux.py`.