From b9812d6995d0c3e8a721c25b5199b31702cc8167 Mon Sep 17 00:00:00 2001 From: stevhliu Date: Wed, 25 Mar 2026 09:15:40 -0700 Subject: [PATCH 1/3] refactor --- .ai/AGENTS.md | 11 +--- .ai/models.md | 88 +++++++++++++++++++++++++++ .ai/review-rules.md | 4 +- .ai/skills/model-integration/SKILL.md | 78 +----------------------- 4 files changed, 95 insertions(+), 86 deletions(-) create mode 100644 .ai/models.md diff --git a/.ai/AGENTS.md b/.ai/AGENTS.md index a42ca5ede871..61ea4149eb84 100644 --- a/.ai/AGENTS.md +++ b/.ai/AGENTS.md @@ -10,10 +10,6 @@ Strive to write code as simple and explicit as possible. --- -### Dependencies -- No new mandatory dependency without discussion (e.g. `einops`) -- Optional deps guarded with `is_X_available()` and a dummy in `utils/dummy_*.py` - ## Code formatting - `make style` and `make fix-copies` should be run as the final step before opening a PR @@ -23,11 +19,10 @@ Strive to write code as simple and explicit as possible. - Remove the header to intentionally break the link ### Models -- All layer calls should be visible directly in `forward` — avoid helper functions that hide `nn.Module` calls. -- Avoid graph breaks for `torch.compile` compatibility — do not insert NumPy operations in forward implementations and any other patterns that can break `torch.compile` compatibility with `fullgraph=True`. -- See the **model-integration** skill for the attention pattern, pipeline rules, test setup instructions, and other important details. +- See [models.md](models.md) for 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. ## Skills Task-specific guides live in `.ai/skills/` and are loaded on demand by AI agents. -Available skills: **model-integration** (adding/converting pipelines), **parity-testing** (debugging numerical parity). +Available skills: [model-integration](./skills/model-integration/SKILL.md) (adding/converting pipelines), [parity-testing](./skills/parity-testing/SKILL.md) (debugging numerical parity). diff --git a/.ai/models.md b/.ai/models.md new file mode 100644 index 000000000000..1d4ee47158a1 --- /dev/null +++ b/.ai/models.md @@ -0,0 +1,88 @@ +# Models — Conventions & Rules + +Shared reference for model-related conventions, patterns, and gotchas. +Linked from `AGENTS.md`, `skills/model-integration/SKILL.md`, and `review-rules.md`. + +## Coding style + +- All layer calls should be visible directly in `forward` — avoid helper functions that hide `nn.Module` calls. +- Avoid graph breaks for `torch.compile` compatibility — do not insert NumPy operations in forward implementations and any other patterns that can break `torch.compile` compatibility with `fullgraph=True`. +- No new mandatory dependency without discussion (e.g. `einops`). Optional deps guarded with `is_X_available()` and a dummy in `utils/dummy_*.py`. + +## Common diffusers conventions + +- Models use `ModelMixin` with `register_to_config` for config serialization +- Pipelines inherit from `DiffusionPipeline` +- Schedulers use `SchedulerMixin` with `ConfigMixin` +- Use `@torch.no_grad()` on pipeline `__call__` +- Support `output_type="latent"` for skipping VAE decode +- Support `generator` parameter for reproducibility +- Use `self.progress_bar(timesteps)` for progress tracking + +## Attention pattern + +Attention must follow the diffusers pattern: both the `Attention` class and its processor are defined in the model file. The processor's `__call__` handles the actual compute and must use `dispatch_attention_fn` rather than calling `F.scaled_dot_product_attention` directly. The attention class inherits `AttentionModuleMixin` and declares `_default_processor_cls` and `_available_processors`. + +```python +# transformer_mymodel.py + +class MyModelAttnProcessor: + _attention_backend = None + _parallel_config = None + + def __call__(self, attn, hidden_states, attention_mask=None, ...): + query = attn.to_q(hidden_states) + key = attn.to_k(hidden_states) + value = attn.to_v(hidden_states) + # reshape, apply rope, etc. + hidden_states = dispatch_attention_fn( + query, key, value, + attn_mask=attention_mask, + backend=self._attention_backend, + parallel_config=self._parallel_config, + ) + hidden_states = hidden_states.flatten(2, 3) + return attn.to_out[0](hidden_states) + + +class MyModelAttention(nn.Module, AttentionModuleMixin): + _default_processor_cls = MyModelAttnProcessor + _available_processors = [MyModelAttnProcessor] + + def __init__(self, query_dim, heads=8, dim_head=64, ...): + super().__init__() + self.to_q = nn.Linear(query_dim, heads * dim_head, bias=False) + self.to_k = nn.Linear(query_dim, heads * dim_head, bias=False) + self.to_v = nn.Linear(query_dim, heads * dim_head, bias=False) + self.to_out = nn.ModuleList([nn.Linear(heads * dim_head, query_dim), nn.Dropout(0.0)]) + self.set_processor(MyModelAttnProcessor()) + + def forward(self, hidden_states, attention_mask=None, **kwargs): + return self.processor(self, hidden_states, attention_mask, **kwargs) +``` + +Consult the implementations in `src/diffusers/models/transformers/` if you need further references. + +## Implementation rules + +1. **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. +2. **Pipelines must inherit from `DiffusionPipeline`.** Consult implementations in `src/diffusers/pipelines` in case you need references. +3. **Don't subclass an existing pipeline for a variant.** DO NOT use an existing pipeline class (e.g., `FluxPipeline`) to override another pipeline (e.g., `FluxImg2ImgPipeline`) which will be a part of the core codebase (`src`). + +## Gotchas + +1. **Forgetting `__init__.py` lazy imports.** Every new class must be registered in the appropriate `__init__.py` with lazy imports. Missing this causes `ImportError` that only shows up when users try `from diffusers import YourNewClass`. + +2. **Using `einops` or other non-PyTorch deps.** Reference implementations often use `einops.rearrange`. Always rewrite with native PyTorch (`reshape`, `permute`, `unflatten`). Don't add the dependency. If a dependency is truly unavoidable, guard its import: `if is_my_dependency_available(): import my_dependency`. + +3. **Missing `make fix-copies` after `# Copied from`.** If you add `# Copied from` annotations, you must run `make fix-copies` to propagate them. CI will fail otherwise. + +4. **Wrong `_supports_cache_class` / `_no_split_modules`.** These class attributes control KV cache and device placement. Copy from a similar model and verify -- wrong values cause silent correctness bugs or OOM errors. + +5. **Missing `@torch.no_grad()` on pipeline `__call__`.** Forgetting this causes GPU OOM from gradient accumulation during inference. + +6. **Config serialization gaps.** Every `__init__` parameter in a `ModelMixin` subclass must be captured by `register_to_config`. If you add a new param but forget to register it, `from_pretrained` will silently use the default instead of the saved value. + +7. **Forgetting to update `_import_structure` and `_lazy_modules`.** The top-level `src/diffusers/__init__.py` has both -- missing either one causes partial import failures. + +8. **Hardcoded dtype in model forward.** Don't hardcode `torch.float32` or `torch.bfloat16` in the model's forward pass. Use the dtype of the input tensors or `self.dtype` so the model works with any precision. diff --git a/.ai/review-rules.md b/.ai/review-rules.md index 12efc94c4b61..e8fd1c202934 100644 --- a/.ai/review-rules.md +++ b/.ai/review-rules.md @@ -3,8 +3,8 @@ Review-specific rules for Claude. Focus on correctness — style is handled by ruff. Before reviewing, read and apply the guidelines in: -- [AGENTS.md](AGENTS.md) — coding style, dependencies, copied code, model conventions -- [skills/model-integration/SKILL.md](skills/model-integration/SKILL.md) — attention pattern, pipeline rules, implementation checklist, gotchas +- [AGENTS.md](AGENTS.md) — coding style, copied code +- [models.md](models.md) — model conventions, attention pattern, implementation rules, dependencies, gotchas - [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.) diff --git a/.ai/skills/model-integration/SKILL.md b/.ai/skills/model-integration/SKILL.md index 16880e4d4850..1d6501486733 100644 --- a/.ai/skills/model-integration/SKILL.md +++ b/.ai/skills/model-integration/SKILL.md @@ -65,89 +65,15 @@ docs/source/en/api/ - [ ] Run `make style` and `make quality` - [ ] Test parity with reference implementation (see `parity-testing` skill) -### Attention pattern - -Attention must follow the diffusers pattern: both the `Attention` class and its processor are defined in the model file. The processor's `__call__` handles the actual compute and must use `dispatch_attention_fn` rather than calling `F.scaled_dot_product_attention` directly. The attention class inherits `AttentionModuleMixin` and declares `_default_processor_cls` and `_available_processors`. - -```python -# transformer_mymodel.py - -class MyModelAttnProcessor: - _attention_backend = None - _parallel_config = None - - def __call__(self, attn, hidden_states, attention_mask=None, ...): - query = attn.to_q(hidden_states) - key = attn.to_k(hidden_states) - value = attn.to_v(hidden_states) - # reshape, apply rope, etc. - hidden_states = dispatch_attention_fn( - query, key, value, - attn_mask=attention_mask, - backend=self._attention_backend, - parallel_config=self._parallel_config, - ) - hidden_states = hidden_states.flatten(2, 3) - return attn.to_out[0](hidden_states) - - -class MyModelAttention(nn.Module, AttentionModuleMixin): - _default_processor_cls = MyModelAttnProcessor - _available_processors = [MyModelAttnProcessor] - - def __init__(self, query_dim, heads=8, dim_head=64, ...): - super().__init__() - self.to_q = nn.Linear(query_dim, heads * dim_head, bias=False) - self.to_k = nn.Linear(query_dim, heads * dim_head, bias=False) - self.to_v = nn.Linear(query_dim, heads * dim_head, bias=False) - self.to_out = nn.ModuleList([nn.Linear(heads * dim_head, query_dim), nn.Dropout(0.0)]) - self.set_processor(MyModelAttnProcessor()) - - def forward(self, hidden_states, attention_mask=None, **kwargs): - return self.processor(self, hidden_states, attention_mask, **kwargs) -``` - -Consult the implementations in `src/diffusers/models/transformers/` if you need further references. +### Model conventions, attention pattern, and implementation rules -### Implementation rules - -1. **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. -2. **Pipelines must inherit from `DiffusionPipeline`.** Consult implementations in `src/diffusers/pipelines` in case you need references. -3. **Don't subclass an existing pipeline for a variant.** DO NOT use an existing pipeline class (e.g., `FluxPipeline`) to override another pipeline (e.g., `FluxImg2ImgPipeline`) which will be a part of the core codebase (`src`). +See [../../models.md](../../models.md) for the attention pattern, implementation rules, common conventions, dependencies, and gotchas. These apply to all model work. ### Test setup - Slow tests gated with `@slow` and `RUN_SLOW=1` - All model-level tests must use the `BaseModelTesterConfig`, `ModelTesterMixin`, `MemoryTesterMixin`, `AttentionTesterMixin`, `LoraTesterMixin`, and `TrainingTesterMixin` classes initially to write the tests. Any additional tests should be added after discussions with the maintainers. Use `tests/models/transformers/test_models_transformer_flux.py` as a reference. -### Common diffusers conventions - -- Pipelines inherit from `DiffusionPipeline` -- Models use `ModelMixin` with `register_to_config` for config serialization -- Schedulers use `SchedulerMixin` with `ConfigMixin` -- Use `@torch.no_grad()` on pipeline `__call__` -- Support `output_type="latent"` for skipping VAE decode -- Support `generator` parameter for reproducibility -- Use `self.progress_bar(timesteps)` for progress tracking - -## Gotchas - -1. **Forgetting `__init__.py` lazy imports.** Every new class must be registered in the appropriate `__init__.py` with lazy imports. Missing this causes `ImportError` that only shows up when users try `from diffusers import YourNewClass`. - -2. **Using `einops` or other non-PyTorch deps.** Reference implementations often use `einops.rearrange`. Always rewrite with native PyTorch (`reshape`, `permute`, `unflatten`). Don't add the dependency. If a dependency is truly unavoidable, guard its import: `if is_my_dependency_available(): import my_dependency`. - -3. **Missing `make fix-copies` after `# Copied from`.** If you add `# Copied from` annotations, you must run `make fix-copies` to propagate them. CI will fail otherwise. - -4. **Wrong `_supports_cache_class` / `_no_split_modules`.** These class attributes control KV cache and device placement. Copy from a similar model and verify -- wrong values cause silent correctness bugs or OOM errors. - -5. **Missing `@torch.no_grad()` on pipeline `__call__`.** Forgetting this causes GPU OOM from gradient accumulation during inference. - -6. **Config serialization gaps.** Every `__init__` parameter in a `ModelMixin` subclass must be captured by `register_to_config`. If you add a new param but forget to register it, `from_pretrained` will silently use the default instead of the saved value. - -7. **Forgetting to update `_import_structure` and `_lazy_modules`.** The top-level `src/diffusers/__init__.py` has both -- missing either one causes partial import failures. - -8. **Hardcoded dtype in model forward.** Don't hardcode `torch.float32` or `torch.bfloat16` in the model's forward pass. Use the dtype of the input tensors or `self.dtype` so the model works with any precision. - --- ## Modular Pipeline Conversion From aefdaa785ed76359f7433c202e8c00b022f04e47 Mon Sep 17 00:00:00 2001 From: stevhliu Date: Wed, 25 Mar 2026 12:59:48 -0700 Subject: [PATCH 2/3] feedback --- .ai/models.md | 7 +++---- .ai/skills/model-integration/SKILL.md | 4 ++++ 2 files changed, 7 insertions(+), 4 deletions(-) diff --git a/.ai/models.md b/.ai/models.md index 1d4ee47158a1..73eba1a39c08 100644 --- a/.ai/models.md +++ b/.ai/models.md @@ -1,4 +1,4 @@ -# Models — Conventions & Rules +# Model conventions and rules Shared reference for model-related conventions, patterns, and gotchas. Linked from `AGENTS.md`, `skills/model-integration/SKILL.md`, and `review-rules.md`. @@ -65,9 +65,8 @@ Consult the implementations in `src/diffusers/models/transformers/` if you need ## Implementation rules -1. **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. -2. **Pipelines must inherit from `DiffusionPipeline`.** Consult implementations in `src/diffusers/pipelines` in case you need references. -3. **Don't subclass an existing pipeline for a variant.** DO NOT use an existing pipeline class (e.g., `FluxPipeline`) to override another pipeline (e.g., `FluxImg2ImgPipeline`) which will be a part of the core codebase (`src`). +1. **Pipelines must inherit from `DiffusionPipeline`.** Consult implementations in `src/diffusers/pipelines` in case you need references. +2. **Don't subclass an existing pipeline for a variant.** DO NOT use an existing pipeline class (e.g., `FluxPipeline`) to override another pipeline (e.g., `FluxImg2ImgPipeline`) which will be a part of the core codebase (`src`). ## Gotchas diff --git a/.ai/skills/model-integration/SKILL.md b/.ai/skills/model-integration/SKILL.md index 1d6501486733..bc9ab46c6e2f 100644 --- a/.ai/skills/model-integration/SKILL.md +++ b/.ai/skills/model-integration/SKILL.md @@ -69,6 +69,10 @@ docs/source/en/api/ See [../../models.md](../../models.md) for the attention pattern, implementation rules, common conventions, dependencies, and gotchas. These apply to all model work. +### Implementation rule + +**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. + ### Test setup - Slow tests gated with `@slow` and `RUN_SLOW=1` From faa30fa48dde85d1519a05df775f32324f6ba409 Mon Sep 17 00:00:00 2001 From: stevhliu Date: Thu, 26 Mar 2026 13:21:09 -0700 Subject: [PATCH 3/3] feedback --- .ai/AGENTS.md | 21 ++++++++++++++++++--- .ai/models.md | 13 +------------ .ai/skills/model-integration/SKILL.md | 2 +- 3 files changed, 20 insertions(+), 16 deletions(-) diff --git a/.ai/AGENTS.md b/.ai/AGENTS.md index 61ea4149eb84..9f42d26b22cf 100644 --- a/.ai/AGENTS.md +++ b/.ai/AGENTS.md @@ -11,18 +11,33 @@ Strive to write code as simple and explicit as possible. --- ## Code formatting + - `make style` and `make fix-copies` should be run as the final step 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 conventions, attention pattern, implementation rules, dependencies, and gotchas. + +- 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 + +- Pipelines inherit from `DiffusionPipeline` +- Schedulers use `SchedulerMixin` with `ConfigMixin` +- Use `@torch.no_grad()` on pipeline `__call__` +- Support `output_type="latent"` for skipping VAE decode +- Support `generator` parameter for reproducibility +- Use `self.progress_bar(timesteps)` for progress tracking +- Don't subclass an existing pipeline for a variant — DO NOT use an existing pipeline class (e.g., `FluxPipeline`) to override another pipeline (e.g., `FluxImg2ImgPipeline`) which will be a part of the core codebase (`src`) + ## Skills -Task-specific guides live in `.ai/skills/` and are loaded on demand by AI agents. -Available skills: [model-integration](./skills/model-integration/SKILL.md) (adding/converting pipelines), [parity-testing](./skills/parity-testing/SKILL.md) (debugging numerical parity). +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). diff --git a/.ai/models.md b/.ai/models.md index 73eba1a39c08..4821c770f615 100644 --- a/.ai/models.md +++ b/.ai/models.md @@ -9,15 +9,9 @@ Linked from `AGENTS.md`, `skills/model-integration/SKILL.md`, and `review-rules. - Avoid graph breaks for `torch.compile` compatibility — do not insert NumPy operations in forward implementations and any other patterns that can break `torch.compile` compatibility with `fullgraph=True`. - No new mandatory dependency without discussion (e.g. `einops`). Optional deps guarded with `is_X_available()` and a dummy in `utils/dummy_*.py`. -## Common diffusers conventions +## Common model conventions - Models use `ModelMixin` with `register_to_config` for config serialization -- Pipelines inherit from `DiffusionPipeline` -- Schedulers use `SchedulerMixin` with `ConfigMixin` -- Use `@torch.no_grad()` on pipeline `__call__` -- Support `output_type="latent"` for skipping VAE decode -- Support `generator` parameter for reproducibility -- Use `self.progress_bar(timesteps)` for progress tracking ## Attention pattern @@ -63,11 +57,6 @@ class MyModelAttention(nn.Module, AttentionModuleMixin): Consult the implementations in `src/diffusers/models/transformers/` if you need further references. -## Implementation rules - -1. **Pipelines must inherit from `DiffusionPipeline`.** Consult implementations in `src/diffusers/pipelines` in case you need references. -2. **Don't subclass an existing pipeline for a variant.** DO NOT use an existing pipeline class (e.g., `FluxPipeline`) to override another pipeline (e.g., `FluxImg2ImgPipeline`) which will be a part of the core codebase (`src`). - ## Gotchas 1. **Forgetting `__init__.py` lazy imports.** Every new class must be registered in the appropriate `__init__.py` with lazy imports. Missing this causes `ImportError` that only shows up when users try `from diffusers import YourNewClass`. diff --git a/.ai/skills/model-integration/SKILL.md b/.ai/skills/model-integration/SKILL.md index bc9ab46c6e2f..97ed536a0083 100644 --- a/.ai/skills/model-integration/SKILL.md +++ b/.ai/skills/model-integration/SKILL.md @@ -69,7 +69,7 @@ docs/source/en/api/ See [../../models.md](../../models.md) for the attention pattern, implementation rules, common conventions, dependencies, and gotchas. These apply to all model work. -### Implementation rule +### Model integration specific rules **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.