huggingface / huggingface/diffusers

chroma model/pipeline review

オープン
#13,619 コメント 0 件 リアクション 0 件 担当者 0 名 GitHub で見る
主要言語
Python
スター
34.5k
フォーク
7.3k
平均マージ
3日 3時間
マージ済み PR(30日)
91

説明

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

コントリビューションガイド

コントリビューションガイドを開く

調査の方向性

Start by running the supplied reproductions and reading the affected Chroma files: the three pipeline files, transformer_chroma.py, transformer_chroma.py, pipeline_output.py, and the existing tests under tests/pipelines/chroma/. Address the reported mask, inpaint, checkpointing, import, test-coverage, and documentation findings, then verify the new behavior with focused fast tests and the requested slow coverage.

索引モデルが issue の本文から書いたものです。

評価

技術スタック
python, pytorch
領域
documentation, machine-learning, testing-qa
issue の種類
バグ
難易度
4/5
見積もり時間
3〜5日
活発さ
静か
明瞭さ
おおむね明確
初心者へのやさしさ
45/100

新しい issue をメールで受け取る

初心者向けの GitHub issue を短くまとめたダイジェスト。