huggingface / huggingface/diffusers
chroma model/pipeline review
- Ngôn ngữ chính
- Python
- Star
- 34.5k
- Fork
- 7.3k
- Merge trung bình
- 3 ngày 3 giờ
- Pull request đã merge (30 ngày)
- 91
Mô tả
# `chroma` model/pipeline review
Commit tested: `0f1abc4ae8b0eb2a3b40e82a310507281144c423`
Review performed against the repository review rules.
Reviewed: public/lazy imports, model config/loading hooks, runtime dtype/device paths, attention masks/processors, offload paths, docs, examples, fast/slow tests, and related Flux/Qwen precedents. Top-level imports for `ChromaPipeline`, `ChromaImg2ImgPipeline`, `ChromaInpaintPipeline`, and `ChromaTransformer2DModel` work.
## Issue 1: Existing Chroma float-mask bug is still present
Affected code:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/chroma/pipeline_chroma.py#L249-L252
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/chroma/pipeline_chroma_img2img.py#L262-L264
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/chroma/pipeline_chroma_inpainting.py#L278-L280
Problem:
The Chroma pipelines convert padding masks to `prompt_embeds` dtype. SDPA interprets float masks as additive bias, not keep/drop masks, so `0.0` does not mask padding. This is an exact duplicate of closed issue https://github.com/huggingface/diffusers/issues/12116 and related earlier issue https://github.com/huggingface/diffusers/issues/11724, but it is still reproducible at this commit.
Impact:
Masked T5 padding can still influence image tokens, especially for short prompts, causing quality/parity regressions.
Reproduction:
```python
import torch
from diffusers import ChromaTransformer2DModel
torch.manual_seed(0)
model = ChromaTransformer2DModel(
in_channels=4, out_channels=4, num_layers=1, num_single_layers=1,
attention_head_dim=4, num_attention_heads=1, joint_attention_dim=8,
axes_dims_rope=(0, 2, 2), approximator_num_channels=16,
approximator_hidden_dim=8, approximator_layers=1,
).eval()
base_encoder = torch.randn(1, 3, 8)
changed_encoder = base_encoder.clone()
changed_encoder[:, 2] += 1000 # token 2 is masked
common = dict(
hidden_states=torch.randn(1, 2, 4),
timestep=torch.ones(1),
txt_ids=torch.zeros(3, 3),
img_ids=torch.zeros(2, 3),
)
for dtype in (torch.bool, torch.float32):
mask = torch.tensor([[1, 1, 0, 1, 1]], dtype=dtype)
with torch.no_grad():
a = model(encoder_hidden_states=base_encoder, attention_mask=mask, **common).sample
b = model(encoder_hidden_states=changed_encoder, attention_mask=mask, **common).sample
print(dtype, (a - b).abs().max().item())
```
Relevant precedent:
`attention_dispatch._normalize_attn_mask` requires bool masks for mask-derived sequence lengths:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/attention_dispatch.py#L639-L647
Suggested fix:
```python
# In all Chroma _get_t5_prompt_embeds methods:
attention_mask = mask_indices <= seq_lengths.unsqueeze(1)
# In all Chroma _prepare_attention_mask methods:
attention_mask = attention_mask.to(dtype=torch.bool)
image_attention_mask = torch.ones(
batch_size, sequence_length, device=attention_mask.device, dtype=torch.bool
)
attention_mask = torch.cat([attention_mask, image_attention_mask], dim=1)
```
## Issue 2: `ChromaInpaintPipeline` crashes on missing `guidance_embeds`
Affected code:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/chroma/pipeline_chroma_inpainting.py#L1042-L1047
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/transformers/transformer_chroma.py#L411-L426
Problem:
The inpaint pipeline copied Flux guidance-embedding handling, but `ChromaTransformer2DModel` does not register `guidance_embeds` and does not accept a `guidance` forward argument. Accessing `self.transformer.config.guidance_embeds` raises `AttributeError`.
Impact:
`ChromaInpaintPipeline.__call__` fails before denoising with the standard Chroma transformer.
Reproduction:
```python
from diffusers import ChromaTransformer2DModel
transformer = ChromaTransformer2DModel(
in_channels=4, out_channels=4, num_layers=0, num_single_layers=0,
attention_head_dim=4, num_attention_heads=1, joint_attention_dim=8,
axes_dims_rope=(0, 2, 2), approximator_num_channels=16,
approximator_hidden_dim=8,
)
print("guidance_embeds" in transformer.config)
print(transformer.config.guidance_embeds) # AttributeError
```
Relevant precedent:
Flux only does this because `FluxTransformer2DModel` registers `guidance_embeds`:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/transformers/transformer_flux.py#L579-L603
Suggested fix:
```python
# Remove the inpaint-only guidance_embeds block entirely.
# ChromaTransformer2DModel has no guidance input and the local `guidance` value is unused.
```
## Issue 3: Gradient checkpointing drops Chroma attention masks in single blocks
Affected code:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/transformers/transformer_chroma.py#L558-L571
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/transformers/transformer_chroma.py#L586-L604
Problem:
The gradient-checkpointed dual-block path omits `joint_attention_kwargs`, and the single-block path omits both `attention_mask` and `joint_attention_kwargs`. With checkpointing enabled, the model no longer computes the same function.
Impact:
Training/fine-tuning with gradient checkpointing can silently ignore padding masks in single-stream blocks and diverge from non-checkpointed training.
Reproduction:
```python
import torch
from diffusers import ChromaTransformer2DModel
torch.manual_seed(0)
model = ChromaTransformer2DModel(
in_channels=4, out_channels=4, num_layers=0, num_single_layers=1,
attention_head_dim=4, num_attention_heads=1, joint_attention_dim=8,
axes_dims_rope=(0, 2, 2), approximator_num_channels=16,
approximator_hidden_dim=8, approximator_layers=1,
).eval()
inputs = dict(
hidden_states=torch.randn(1, 2, 4, requires_grad=True),
encoder_hidden_states=torch.randn(1, 2, 8, requires_grad=True),
timestep=torch.ones(1),
txt_ids=torch.zeros(2, 3),
img_ids=torch.zeros(2, 3),
attention_mask=torch.tensor([[1, 0, 1, 1]], dtype=torch.bool),
)
out_no_ckpt = model(**inputs).sample.detach()
model.enable_gradient_checkpointing()
out_ckpt = model(**inputs).sample.detach()
print((out_no_ckpt - out_ckpt).abs().max().item())
```
Relevant precedent:
No exact duplicate found for `ChromaTransformer2DModel attention_mask gradient_checkpointing`.
Suggested fix:
```python
encoder_hidden_states, hidden_states = self._gradient_checkpointing_func(
block, hidden_states, encoder_hidden_states, temb, image_rotary_emb, attention_mask, joint_attention_kwargs
)
hidden_states = self._gradient_checkpointing_func(
block, hidden_states, temb, image_rotary_emb, attention_mask, joint_attention_kwargs
)
```
## Issue 4: Chroma inpaint has no fast tests, and Chroma has no slow tests
Affected code:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/tests/pipelines/chroma/test_pipeline_chroma.py#L13-L19
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/tests/pipelines/chroma/test_pipeline_chroma_img2img.py#L14-L20
Problem:
Only text2img and img2img fast pipeline tests exist. There is no `ChromaInpaintPipeline` fast test, and no Chroma slow tests anywhere under `tests/`.
Impact:
The inpaint runtime crash above is not covered, and there is no slow coverage against published Chroma checkpoints.
Reproduction:
```python
from pathlib import Path
files = sorted(str(p).replace("\\", "/") for p in Path("tests").rglob("*chroma*.py"))
print("\n".join(files))
print("has_inpaint_fast_test", any("inpaint" in f for f in files))
print("has_slow_marker", any("@slow" in Path(f).read_text(encoding="utf-8") for f in files))
```
Relevant precedent:
Flux and QwenImage have inpaint fast-test classes:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/tests/pipelines/flux/test_pipeline_flux_inpaint.py#L21-L27
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/tests/pipelines/qwenimage/test_qwenimage_inpaint.py#L37-L56
Suggested fix:
```python
# Add tests/pipelines/chroma/test_pipeline_chroma_inpainting.py using the existing
# Chroma tiny components plus image/mask tensors, and add at least one @slow Chroma
# pipeline test against a published Chroma checkpoint or saved test slices.
```
## Issue 5: `ChromaTransformer2DModel` uses deprecated `FluxPosEmbed` import path
Affected code:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/transformers/transformer_chroma.py#L29-L33
Problem:
`FluxPosEmbed` is imported from `diffusers.models.embeddings`, whose shim emits a deprecation warning and asks callers to import from `diffusers.models.transformers.transformer_flux`.
Impact:
Every Chroma transformer construction emits a user-visible `FutureWarning`.
Reproduction:
```python
import warnings
from diffusers import ChromaTransformer2DModel
with warnings.catch_warnings(record=True) as caught:
warnings.simplefilter("always")
ChromaTransformer2DModel(
in_channels=4, out_channels=4, num_layers=0, num_single_layers=0,
attention_head_dim=4, num_attention_heads=1, joint_attention_dim=8,
axes_dims_rope=(0, 2, 2), approximator_num_channels=16,
approximator_hidden_dim=8,
)
print([str(w.message) for w in caught if "FluxPosEmbed" in str(w.message)])
```
Relevant precedent:
The non-deprecated class lives here:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/transformers/transformer_flux.py#L494-L498
Suggested fix:
```python
from ..embeddings import PixArtAlphaTextProjection, Timesteps, get_timestep_embedding
from .transformer_flux import FluxAttention, FluxAttnProcessor, FluxPosEmbed
```
## Issue 6: Chroma inpaint/output docs contain copied or AI-artifact text
Affected code:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/chroma/pipeline_chroma_inpainting.py#L1-L4
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/chroma/pipeline_chroma_inpainting.py#L172-L195
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/chroma/pipeline_output.py#L11-L17
Problem:
The inpaint file includes `contentReference[oaicite:*]` artifacts and its class docstring describes Flux, Black Forest Labs, DDIM, CLIP, `text_encoder_2`, and `tokenizer_2`, none of which match this Chroma pipeline signature. `ChromaPipelineOutput` also says “Stable Diffusion pipelines.”
Impact:
Generated API docs are misleading and violate the repo review rule against ephemeral context/artifacts.
Reproduction:
```python
from pathlib import Path
inpaint = Path("src/diffusers/pipelines/chroma/pipeline_chroma_inpainting.py").read_text()
output = Path("src/diffusers/pipelines/chroma/pipeline_output.py").read_text()
for token in ["contentReference[oaicite", "The Flux pipeline", "[`DDIMScheduler`]", "[`CLIPTextModel`]", "text_encoder_2"]:
print(token, token in inpaint)
print("Stable Diffusion output doc", "Output class for Stable Diffusion pipelines." in output)
```
Relevant precedent:
No duplicate issue or PR found for the doc artifacts.
Suggested fix:
```python
# Replace the copied inpaint docstring with Chroma-specific text:
# - Chroma image inpainting
# - FlowMatchEulerDiscreteScheduler
# - T5EncoderModel / T5TokenizerFast
# - no text_encoder_2/tokenizer_2/CLIP text encoder
# Also change ChromaPipelineOutput to "Output class for Chroma pipelines."
```
Hướng dẫn đóng góp
Hướng nghiên cứu
Bắt đầu bằng cách chạy các bản tái hiện được cung cấp và đọc các tệp Chroma bị ảnh hưởng: ba tệp pipeline, transformer_chroma.py, transformer_chroma.py, pipeline_output.py và các test hiện có trong tests/pipelines/chroma/. Xử lý các vấn đề đã báo cáo liên quan đến mask, inpaint, checkpointing, import, độ bao phủ test và tài liệu, sau đó xác minh hành vi mới bằng các test nhanh tập trung và phần bao phủ chậm được yêu cầu.
Do mô hình lập chỉ mục viết ra từ nội dung của issue.
Đánh giá
- Công nghệ
- python, pytorch
- Lĩnh vực
- documentation, machine-learning, testing-qa
- Loại issue
- Lỗi
- Độ khó
- 4/5
- Thời gian dự kiến
- 3-5 ngày
- Mức độ hoạt động
- Ít trao đổi
- Độ rõ ràng
- Khá rõ ràng
- Mức phù hợp với người mới
- 45/100