huggingface / huggingface/diffusers

chroma model/pipeline review

Open
#13,619 0 comments 0 reactions 0 assignees View on GitHub
Dominant language
Python
Stars
34.5k
Forks
7.3k
Avg merge
3d 3h
Merged PRs (30d)
91

Description

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

Contributor guide

Open the contributing guide

Research direction

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.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
documentation, machine-learning, testing-qa
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Quiet
Clarity
Mostly clear
Newbie friendliness
45/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.