Skip to content
25 changes: 10 additions & 15 deletions .ai/AGENTS.md
Original file line number Diff line number Diff line change
Expand Up @@ -8,36 +8,31 @@ Strive to write code as simple and explicit as possible.
- No defensive code, unused code paths, or legacy stubs — do not add fallback paths, safety checks, or configuration options "just in case"; do not carry unused method parameters "for API consistency", backwards-compatibility aliases for names that never shipped, or deprecation shims for code that was never released. When porting from a research repo, delete training-time code paths, experimental flags, and ablation branches entirely — only keep the inference path you are actually integrating.
- Do not guess user intent and silently correct behavior. Make the expected inputs clear in the docstring, and raise a concise error for unsupported cases rather than adding complex fallback logic.

Before opening the PR, self-review against [review-rules.md](review-rules.md), which collects the most common mistakes we catch in review.

---

## Code formatting

- `make style` and `make fix-copies` should be run as the final step before opening a PR
- `make style` and `make fix-copies` should be run before opening a PR

### Copied Code

- Many classes are kept in sync with a source via a `# Copied from ...` header comment
- Do not edit a `# Copied from` block directly — run `make fix-copies` to propagate changes from the source
- Remove the header to intentionally break the link

### Models

- See [models.md](models.md) for model conventions, attention pattern, implementation rules, dependencies, and gotchas.
- See the [model-integration](./skills/model-integration/SKILL.md) skill for the full integration workflow, file structure, test setup, and other details.

### Pipelines & Schedulers

- See [pipelines.md](pipelines.md) for pipeline conventions, patterns, and gotchas.
## Reference guides

### Modular Pipelines

- See [modular.md](modular.md) for modular pipeline conventions, patterns, and gotchas.
- **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.

## Skills

Task-specific guides live in `.ai/skills/` and are loaded on demand by AI agents. Available skills include:

- [model-integration](./skills/model-integration/SKILL.md) (adding/converting pipelines)
- [parity-testing](./skills/parity-testing/SKILL.md) (debugging numerical parity).
- [self-review](./skills/self-review/SKILL.md) (pre-PR self-review against the project rules)

## Self-review before a PR

Before opening a PR, run self-review against [review-rules.md](review-rules.md). The [self-review skill](skills/self-review/SKILL.md) runs this as the same pass the `@claude` CI reviewer uses.
1 change: 1 addition & 0 deletions .ai/models.md
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@ Linked from `AGENTS.md`, `skills/model-integration/SKILL.md`, and `review-rules.

* Models use `ModelMixin` with `register_to_config` for config serialization.
* When adding a new transformer (or reviewing one), skim `src/diffusers/models/transformers/transformer_flux.py`, `src/diffusers/models/transformers/transformer_flux2.py`, `src/diffusers/models/transformers/transformer_qwenimage.py`, and `src/diffusers/models/transformers/transformer_wan.py` first to establish the pattern. Most conventions (mixin set, file structure, naming, gradient-checkpointing implementation, `_no_split_modules` settings, etc.) are easiest to internalize by comparison rather than from a fixed list.
* **Loading goes through `from_pretrained` / `from_single_file`.** Weights and configs load through the standard paths — never fetched or imported out-of-band at runtime. Don't override or add a custom `from_pretrained`, and don't load weights manually (`load_file(...)`, `hf_hub_download(...)`, or `sys.path.insert(...)` to import a reference repo). For an original-format single checkpoint, add `from_single_file` support (mixin + weight-mapping).

## Attention pattern

Expand Down
12 changes: 12 additions & 0 deletions .ai/pipelines.md
Original file line number Diff line number Diff line change
Expand Up @@ -3,10 +3,22 @@
Shared reference for pipeline-related conventions, patterns, and gotchas.
Linked from `AGENTS.md`, `skills/model-integration/SKILL.md`, and `review-rules.md`.

> **Prefer modular for new pipelines.** [Modular Diffusers](modular.md) is the preferred way to add a new pipeline; the standard `DiffusionPipeline` covered below is still supported but is no longer the default. We prefer modular especially for models that don't fit a fixed task-based structure (e.g. modality baked into the checkpoint) or that are actively evolving. The conventions below apply when you do build or review a standard pipeline.
## Common pipeline conventions

When adding a new pipeline (or reviewing one), skim `pipeline_flux.py`, `pipeline_flux2.py`, `pipeline_qwenimage.py`, `pipeline_wan.py` first to establish the pattern. Most conventions (class structure, mixin set, `__call__` shape — input validation → encode prompt → timesteps → latent prep → denoise loop → decode — `encode_prompt` / `prepare_latents` shape, `output_type` / `generator` / `progress_bar` plumbing, `@torch.no_grad()` on `__call__`, LoRA mixin, `from_single_file` support, etc.) are easiest to internalize by comparison rather than from a fixed list.

## File structure

```
src/diffusers/pipelines/<model>/
__init__.py # Lazy imports
pipeline_<model>.py # Main pipeline (with __call__)
pipeline_<model>_<variant>.py # Variant pipelines (e.g. img2img, inpaint) — one file/class each
pipeline_output.py # Output dataclass
```

## Gotchas

1. **Config-derived static values: prefer `__init__` attributes.** Values that come from a sub-component's config (e.g. `vae_scale_factor`) belong as `self.foo = ...` in `__init__` — not `@property`, not module-level constants. Note the `getattr(...)` fallback — sub-components may not be loaded when the pipeline is constructed (e.g. via `from_pretrained` on a partial config), so don't assume `self.vae` / `self.transformer` exists.
Expand Down
3 changes: 1 addition & 2 deletions .ai/review-rules.md
Original file line number Diff line number Diff line change
Expand Up @@ -7,8 +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
- [skills/parity-testing/SKILL.md](skills/parity-testing/SKILL.md) — testing rules, comparison utilities
- [skills/parity-testing/pitfalls.md](skills/parity-testing/pitfalls.md) — known pitfalls (dtype mismatches, config assumptions, etc.)
- [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

Expand Down
137 changes: 81 additions & 56 deletions .ai/skills/model-integration/SKILL.md
Original file line number Diff line number Diff line change
Expand Up @@ -8,72 +8,82 @@ description: >

## Goal

Integrate a new model into diffusers end-to-end. The overall flow:

1. **Gather info** — ask the user for the reference repo, setup guide, a runnable inference script, and other objectives such as standard vs modular.
2. **Confirm the plan** — once you have everything, tell the user exactly what you'll do: e.g. "I'll integrate model X with pipeline Y into diffusers based on your script. I'll run parity tests (model-level and pipeline-level) using the `parity-testing` skill to verify numerical correctness against the reference."
3. **Implement** — write the diffusers code (model, pipeline, scheduler if needed), convert weights, register in `__init__.py`.
4. **Parity test** — use the `parity-testing` skill to verify component and e2e parity against the reference implementation.
5. **Deliver a unit test** — provide a self-contained test script that runs the diffusers implementation, checks numerical output (np allclose), and saves an image/video for visual verification. This is what the user runs to confirm everything works.

Work one workflow at a time — get it to full parity before moving on.
Integrate a new model into diffusers end-to-end, to full numerical parity with the reference implementation — one workflow at a time.

## Setup — gather before starting

Before writing any code, gather info in this order:

1. **Reference repo** — ask for the github link. If they've already set it up locally, ask for the path. Otherwise, ask what setup steps are needed (install deps, download checkpoints, set env vars, etc.) and run through them before proceeding.
2. **Inference script** — ask for a runnable end-to-end script for a basic workflow first (e.g. T2V). Then ask what other workflows they want to support (I2V, V2V, etc.) and agree on the full implementation order together.
3. **Standard vs modular** — standard pipelines, modular, or both?
3. **Standard vs modular** — **default to modular.** [Modular Diffusers](../../modular.md) is the preferred implementation for new pipelines; the standard `DiffusionPipeline` is still supported but no longer the default. We prefer modular especially for models that don't fit a fixed task-based structure (modality baked into the checkpoint) or that are actively evolving.

Ask step 3 as an `AskUserQuestion`, with modular marked as the recommended default.

Once you have everything, **confirm the plan** with the user before implementing — state exactly what you'll do, e.g. "I'll integrate model X with pipeline Y based on your script, and verify the model matches the reference before considering it done."

Use `AskUserQuestion` with structured choices for step 3 when the options are known.
Then work through the **Integration checklist** below

## Standard Pipeline Integration
## Integration checklist

### File structure for a new model
- [ ] **Transformer model**
Comment thread
yiyixuxu marked this conversation as resolved.
- [ ] Implement the model with `from_pretrained` support (conventions: [models.md](../../models.md))
- [ ] Convert weights (see **Weight / Checkpoint Conversion**)
- [ ] Parity test against the reference (internal, not shipped — see **Model parity test**)
- [ ] Register in the relevant `__init__.py` files (lazy imports)
- [ ] Model-level tests (see **Testing**)
- [ ] **VAE** (if applicable) — reuse an existing `AutoencoderKL*` if possible; if a new one is needed, follow the same sub-steps as the transformer
- [ ] **Scheduler** — reuse an existing scheduler, or add a custom one
- [ ] **Pipeline**
- [ ] Implement the pipeline — see [modular.md](../../modular.md) for modular pipeline, or [pipelines.md](../../pipelines.md) for standard pipeline
- [ ] Add a LoRA mixin if applicable

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Maybe better if this is skipped in the initial implementation. Easier to review and test.

- [ ] Register in the relevant `__init__.py` files (lazy imports)
- [ ] Pipeline-level tests (see **Testing**)
- [ ] **Docs** — see **File structure**
- [ ] **Style** — `make style` and `make quality`

## File structure

A new model PR roughly lands these files (the contents of `pipelines/<model>/` and `modular_pipelines/<model>/` live in their guides):

```
src/diffusers/
models/transformers/transformer_<model>.py # The core model
schedulers/scheduling_<model>.py # If model needs a custom scheduler
pipelines/<model>/
__init__.py
pipeline_<model>.py # Main pipeline
pipeline_<model>_<variant>.py # Variant pipelines (e.g. pyramid, distilled)
pipeline_output.py # Output dataclass
loaders/lora_pipeline.py # LoRA mixin (add to existing file)

models/transformers/transformer_<model>.py # the model (or models/autoencoders/, models/unets/)
schedulers/scheduling_<model>.py # only if a custom scheduler is needed
loaders/lora_pipeline.py # LoRA mixin — add to the existing file
pipelines/<model>/ # standard pipeline — see pipelines.md
modular_pipelines/<model>/ # modular pipeline — see modular.md
tests/
models/transformers/test_models_transformer_<model>.py
pipelines/<model>/test_<model>.py
lora/test_lora_layers_<model>.py

docs/source/en/api/
pipelines/<model>.md
models/<model>_transformer3d.md # or appropriate name
docs/source/en/
_toctree.yml # register the new pages in the docs index
api/models/<model>.md
api/pipelines/<model>.md
```

### Integration checklist
## Model integration specific rules

- [ ] Implement transformer model with `from_pretrained` support
- [ ] Implement or reuse scheduler
- [ ] Implement pipeline(s) with `__call__` method
- [ ] Add LoRA support if applicable
- [ ] Register all classes in `__init__.py` files (lazy imports)
- [ ] Write unit tests (model, pipeline, LoRA)
- [ ] Write docs
- [ ] Run `make style` and `make quality`
- [ ] Test parity with reference implementation (see `parity-testing` skill)
**Match the reference's numerical logic.** Restructuring code to fit diffusers APIs (`ModelMixin`, `ConfigMixin`, blocks for modular, etc.) is expected, and required diffusers conventions (e.g. the attention pattern in [models.md](../../models.md)) take precedence. Beyond those, keep the actual computation as close to the reference as possible — don't reorder operations, change the math, or rename internals for aesthetics, even if it looks unclean. Small deviations make output mismatches very hard to track down.

### Model conventions, attention pattern, and implementation rules
## Weight / Checkpoint Conversion

See [../../models.md](../../models.md) for the attention pattern, implementation rules, common conventions, dependencies, and gotchas. These apply to all model work.
Convert the original checkpoint into diffusers format with a standalone script under `scripts/` (e.g. `scripts/convert_<model>_to_diffusers.py`). The flow:

### Model integration specific rules
1. Map the original state-dict keys to the diffusers module names (renames + any tensor surgery — see patterns below).
2. Instantiate the diffusers model from its config and load the converted state dict.
3. `save_pretrained(...)` to a **persistent** path (never `/tmp/` — see [pitfalls.md](pitfalls.md) #10), then load it back with `from_pretrained` to confirm it round-trips.

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Bug — broken #10 reference. The parenthetical says see [pitfalls.md](pitfalls.md) #10, but in the new pitfalls.md item #10 is "RoPE position dtype", not the storage-location rule. The original pitfall #10 ("NEVER store converted models in /tmp/") was dropped during the consolidation from parity-testing/pitfalls.md.

Suggested fix — inline the advice and drop the dangling anchor:

Suggested change
3. `save_pretrained(...)` to a **persistent** path (never `/tmp/` — see [pitfalls.md](pitfalls.md) #10), then load it back with `from_pretrained` to confirm it round-trips.
3. `save_pretrained(...)` to a **persistent** path (never `/tmp/` — temporary directories get wiped on restart), then load it back with `from_pretrained` to confirm it round-trips.


**Don't combine structural changes with behavioral changes.** Restructuring code to fit diffusers APIs (ModelMixin, ConfigMixin, etc.) is unavoidable. But don't also "improve" the algorithm, refactor computation order, or rename internal variables for aesthetics. Keep numerical logic as close to the reference as possible, even if it looks unclean. For standard → modular, this is stricter: copy loop logic verbatim and only restructure into blocks. Clean up in a separate commit after parity is confirmed.
All weights load through the standard paths — `from_pretrained`, or `from_single_file` (add `FromSingleFileMixin` + a weight-mapping) for an original-format single checkpoint. No custom `from_pretrained`, no manual runtime loading. See the loading rule in [models.md](../../models.md).

### Testing
Common conversion patterns to watch for:
Comment thread
yiyixuxu marked this conversation as resolved.
Outdated
- Fused QKV weights that need splitting into separate Q, K, V
- Scale/shift ordering differences (reference stores `[shift, scale]`, diffusers expects `[scale, shift]`)
- Weight transpositions (linear stored as transposed conv, or vice versa)
- Interleaved head dimensions that need reshaping
- Bias terms absorbed into different layers

## 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.

Expand All @@ -82,7 +92,7 @@ Two test layers must be added for any new pipeline: pipeline-level tests, and (i
- 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
### 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`.
Expand All @@ -91,7 +101,7 @@ Two test layers must be added for any new pipeline: pipeline-level tests, and (i
- 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
### 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):

Expand All @@ -105,20 +115,35 @@ python utils/generate_model_tests.py src/diffusers/models/transformers/transform
- 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`.

---
## Model parity test

## Modular Pipeline Conversion
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.

See [modular.md](../../modular.md) for the full guide on modular pipeline conventions, block types, build order, guider abstraction, gotchas, and conversion checklist.
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**.

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

```python
@torch.inference_mode()
def test_my_component():
# deterministic input — use the same shape & dtype the real model receives at this stage
gen = torch.Generator().manual_seed(42)
x = torch.randn(1, 16, 32, 32, generator=gen, dtype=torch.float32) # adjust to the real input shape

## Weight Conversion Tips
original_state_dict = load_original_weights(...) # the original checkpoint — both sides load these same weights

<!-- TODO: Add concrete examples as we encounter them. Common patterns to watch for:
- Fused QKV weights that need splitting into separate Q, K, V
- Scale/shift ordering differences (reference stores [shift, scale], diffusers expects [scale, shift])
- Weight transpositions (linear stored as transposed conv, or vice versa)
- Interleaved head dimensions that need reshaping
- Bias terms absorbed into different layers
Add each with a before/after code snippet showing the conversion. -->
# reference: the original repo's implementation (load one model at a time to fit in CPU RAM)
ref_model = ReferenceModel(config) # ReferenceModel: imported from the original repo
ref_model.load_state_dict(original_state_dict, strict=True)
ref_model = ref_model.float().eval()
ref_out = ref_model(x).clone() # clone before freeing the model
del ref_model

# diffusers: convert those same weights with your conversion-script function, then run
diff_model = convert_my_component(original_state_dict) # convert_my_component: the fn from convert_<model>_to_diffusers.py
diff_model = diff_model.float().eval()
diff_out = diff_model(x)

max_diff = (ref_out - diff_out).abs().max().item()
assert max_diff < 1e-3, f"FAIL: max_diff={max_diff:.2e}"
```
Loading
Loading