Skip to content

Commit bfe91ff

Browse files
stevhliustevhliu
committed
refactor
1 parent 762ae05 commit bfe91ff

4 files changed

Lines changed: 104 additions & 84 deletions

File tree

‎.ai/AGENTS.md‎

Lines changed: 3 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -10,10 +10,6 @@ Strive to write code as simple and explicit as possible.
1010

1111
---
1212

13-
### Dependencies
14-
- No new mandatory dependency without discussion (e.g. `einops`)
15-
- Optional deps guarded with `is_X_available()` and a dummy in `utils/dummy_*.py`
16-
1713
## Code formatting
1814
- `make style` and `make fix-copies` should be run as the final step before opening a PR
1915

@@ -23,11 +19,10 @@ Strive to write code as simple and explicit as possible.
2319
- Remove the header to intentionally break the link
2420

2521
### Models
26-
- All layer calls should be visible directly in `forward` — avoid helper functions that hide `nn.Module` calls.
27-
- 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`.
28-
- See the **model-integration** skill for the attention pattern, pipeline rules, test setup instructions, and other important details.
22+
- See [models.md](models.md) for conventions, attention pattern, implementation rules, dependencies, and gotchas.
23+
- See the [model-integration](./skills/model-integration/SKILL.md) skill for the full integration workflow, file structure, test setup, and other details.
2924

3025
## Skills
3126

3227
Task-specific guides live in `.ai/skills/` and are loaded on demand by AI agents.
33-
Available skills: **model-integration** (adding/converting pipelines), **parity-testing** (debugging numerical parity).
28+
Available skills: [model-integration](./skills/model-integration/SKILL.md) (adding/converting pipelines), [parity-testing](./skills/parity-testing/SKILL.md) (debugging numerical parity).

‎.ai/models.md‎

Lines changed: 88 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,88 @@
1+
# Models — Conventions & Rules
2+
3+
Shared reference for model-related conventions, patterns, and gotchas.
4+
Linked from `AGENTS.md`, `skills/model-integration/SKILL.md`, and `review-rules.md`.
5+
6+
## Coding style
7+
8+
- All layer calls should be visible directly in `forward` — avoid helper functions that hide `nn.Module` calls.
9+
- 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`.
10+
- No new mandatory dependency without discussion (e.g. `einops`). Optional deps guarded with `is_X_available()` and a dummy in `utils/dummy_*.py`.
11+
12+
## Common diffusers conventions
13+
14+
- Models use `ModelMixin` with `register_to_config` for config serialization
15+
- Pipelines inherit from `DiffusionPipeline`
16+
- Schedulers use `SchedulerMixin` with `ConfigMixin`
17+
- Use `@torch.no_grad()` on pipeline `__call__`
18+
- Support `output_type="latent"` for skipping VAE decode
19+
- Support `generator` parameter for reproducibility
20+
- Use `self.progress_bar(timesteps)` for progress tracking
21+
22+
## Attention pattern
23+
24+
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`.
25+
26+
```python
27+
# transformer_mymodel.py
28+
29+
class MyModelAttnProcessor:
30+
_attention_backend = None
31+
_parallel_config = None
32+
33+
def __call__(self, attn, hidden_states, attention_mask=None, ...):
34+
query = attn.to_q(hidden_states)
35+
key = attn.to_k(hidden_states)
36+
value = attn.to_v(hidden_states)
37+
# reshape, apply rope, etc.
38+
hidden_states = dispatch_attention_fn(
39+
query, key, value,
40+
attn_mask=attention_mask,
41+
backend=self._attention_backend,
42+
parallel_config=self._parallel_config,
43+
)
44+
hidden_states = hidden_states.flatten(2, 3)
45+
return attn.to_out[0](hidden_states)
46+
47+
48+
class MyModelAttention(nn.Module, AttentionModuleMixin):
49+
_default_processor_cls = MyModelAttnProcessor
50+
_available_processors = [MyModelAttnProcessor]
51+
52+
def __init__(self, query_dim, heads=8, dim_head=64, ...):
53+
super().__init__()
54+
self.to_q = nn.Linear(query_dim, heads * dim_head, bias=False)
55+
self.to_k = nn.Linear(query_dim, heads * dim_head, bias=False)
56+
self.to_v = nn.Linear(query_dim, heads * dim_head, bias=False)
57+
self.to_out = nn.ModuleList([nn.Linear(heads * dim_head, query_dim), nn.Dropout(0.0)])
58+
self.set_processor(MyModelAttnProcessor())
59+
60+
def forward(self, hidden_states, attention_mask=None, **kwargs):
61+
return self.processor(self, hidden_states, attention_mask, **kwargs)
62+
```
63+
64+
Consult the implementations in `src/diffusers/models/transformers/` if you need further references.
65+
66+
## Implementation rules
67+
68+
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.
69+
2. **Pipelines must inherit from `DiffusionPipeline`.** Consult implementations in `src/diffusers/pipelines` in case you need references.
70+
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`).
71+
72+
## Gotchas
73+
74+
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`.
75+
76+
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`.
77+
78+
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.
79+
80+
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.
81+
82+
5. **Missing `@torch.no_grad()` on pipeline `__call__`.** Forgetting this causes GPU OOM from gradient accumulation during inference.
83+
84+
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.
85+
86+
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.
87+
88+
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.

‎.ai/review-rules.md‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,11 @@
1+
# PR Review Rules
2+
3+
Review-specific rules for Claude. Focus on correctness — style is handled by ruff.
4+
5+
Before reviewing, read and apply the guidelines in:
6+
- [AGENTS.md](AGENTS.md) — coding style, copied code
7+
- [models.md](models.md) — model conventions, attention pattern, implementation rules, dependencies, gotchas
8+
- [skills/parity-testing/SKILL.md](skills/parity-testing/SKILL.md) — testing rules, comparison utilities
9+
- [skills/parity-testing/pitfalls.md](skills/parity-testing/pitfalls.md) — known pitfalls (dtype mismatches, config assumptions, etc.)
10+
11+
## Common mistakes (add new rules below this line)

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

Lines changed: 2 additions & 76 deletions
Original file line numberDiff line numberDiff line change
@@ -65,89 +65,15 @@ docs/source/en/api/
6565
- [ ] Run `make style` and `make quality`
6666
- [ ] Test parity with reference implementation (see `parity-testing` skill)
6767

68-
### Attention pattern
69-
70-
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`.
71-
72-
```python
73-
# transformer_mymodel.py
74-
75-
class MyModelAttnProcessor:
76-
_attention_backend = None
77-
_parallel_config = None
78-
79-
def __call__(self, attn, hidden_states, attention_mask=None, ...):
80-
query = attn.to_q(hidden_states)
81-
key = attn.to_k(hidden_states)
82-
value = attn.to_v(hidden_states)
83-
# reshape, apply rope, etc.
84-
hidden_states = dispatch_attention_fn(
85-
query, key, value,
86-
attn_mask=attention_mask,
87-
backend=self._attention_backend,
88-
parallel_config=self._parallel_config,
89-
)
90-
hidden_states = hidden_states.flatten(2, 3)
91-
return attn.to_out[0](hidden_states)
92-
93-
94-
class MyModelAttention(nn.Module, AttentionModuleMixin):
95-
_default_processor_cls = MyModelAttnProcessor
96-
_available_processors = [MyModelAttnProcessor]
97-
98-
def __init__(self, query_dim, heads=8, dim_head=64, ...):
99-
super().__init__()
100-
self.to_q = nn.Linear(query_dim, heads * dim_head, bias=False)
101-
self.to_k = nn.Linear(query_dim, heads * dim_head, bias=False)
102-
self.to_v = nn.Linear(query_dim, heads * dim_head, bias=False)
103-
self.to_out = nn.ModuleList([nn.Linear(heads * dim_head, query_dim), nn.Dropout(0.0)])
104-
self.set_processor(MyModelAttnProcessor())
105-
106-
def forward(self, hidden_states, attention_mask=None, **kwargs):
107-
return self.processor(self, hidden_states, attention_mask, **kwargs)
108-
```
109-
110-
Consult the implementations in `src/diffusers/models/transformers/` if you need further references.
68+
### Model conventions, attention pattern, and implementation rules
11169

112-
### Implementation rules
113-
114-
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.
115-
2. **Pipelines must inherit from `DiffusionPipeline`.** Consult implementations in `src/diffusers/pipelines` in case you need references.
116-
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`).
70+
See [../../models.md](../../models.md) for the attention pattern, implementation rules, common conventions, dependencies, and gotchas. These apply to all model work.
11771

11872
### Test setup
11973

12074
- Slow tests gated with `@slow` and `RUN_SLOW=1`
12175
- 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.
12276

123-
### Common diffusers conventions
124-
125-
- Pipelines inherit from `DiffusionPipeline`
126-
- Models use `ModelMixin` with `register_to_config` for config serialization
127-
- Schedulers use `SchedulerMixin` with `ConfigMixin`
128-
- Use `@torch.no_grad()` on pipeline `__call__`
129-
- Support `output_type="latent"` for skipping VAE decode
130-
- Support `generator` parameter for reproducibility
131-
- Use `self.progress_bar(timesteps)` for progress tracking
132-
133-
## Gotchas
134-
135-
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`.
136-
137-
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`.
138-
139-
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.
140-
141-
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.
142-
143-
5. **Missing `@torch.no_grad()` on pipeline `__call__`.** Forgetting this causes GPU OOM from gradient accumulation during inference.
144-
145-
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.
146-
147-
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.
148-
149-
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.
150-
15177
---
15278

15379
## Modular Pipeline Conversion

0 commit comments

Comments
 (0)