huggingface / huggingface/diffusers

lumina2 model/pipeline review

Aperta
#13,613 0 commenti 0 reazioni 0 assegnatari Vedi su GitHub
Lingua principale
Python
Stelle
34.5k
Fork
7.3k
Merge medio
3g 3h
PR unite (30g)
91

Descrizione

# `lumina2` model/pipeline review

Commit tested: `0f1abc4ae8b0eb2a3b40e82a310507281144c423`

Review performed against the repository review rules.

Reviewed target pipeline/model/init files, top-level lazy exports, dummy exports, fast/model/single-file/LoRA tests, docs, and DreamBooth example coverage. Duplicate search was run with `gh search issues/prs` for `lumina2`, affected class/function names, and each specific failure mode. Only the scheduler `image_seq_len` item had an existing duplicate.

Targeted pytest command using `.venv` was attempted, but collection failed in this local torch build because `torch._C._distributed_c10d` is missing.

## Issue 1: Deprecated alias is exported but cannot be constructed

Affected code:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/lumina2/pipeline_lumina2.py#L801-L818
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/__init__.py#L622-L623

Problem:
`Lumina2Text2ImgPipeline` remains publicly exported, but its constructor calls `deprecate(..., "0.34", ...)`. The current package version is `0.38.0.dev0`, so `deprecate` raises a `ValueError` instead of warning.

Impact:
Users can import the backwards-compatible alias, but any construction or config load path that instantiates it fails immediately.

Reproduction:
```python
from diffusers import Lumina2Text2ImgPipeline

Lumina2Text2ImgPipeline(
transformer=None,
scheduler=None,
vae=None,
text_encoder=None,
tokenizer=None,
)
```

Relevant precedent:
Related rename PR, but not a duplicate for the current failure: https://github.com/huggingface/diffusers/pull/10827

Suggested fix:
```python
# If keeping the alias:
deprecate(
"diffusers.pipelines.lumina2.pipeline_lumina2.Lumina2Text2ImgPipeline",
"1.0.0",
deprecation_message,
)

# Or remove the alias from pipeline_lumina2.py, lazy exports, top-level exports, and dummy objects.
```

## Issue 2: Precomputed negative prompt embeds are not repeated for `num_images_per_prompt`

Affected code:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/lumina2/pipeline_lumina2.py#L297-L339

Problem:
`encode_prompt` repeats `prompt_embeds` and `prompt_attention_mask` unconditionally, but repeats `negative_prompt_embeds` and `negative_prompt_attention_mask` only when the pipeline encoded them itself. If the caller supplies precomputed negative embeddings and `num_images_per_prompt > 1`, the positive and negative batches diverge.

Impact:
Classifier-free guidance can broadcast incorrectly for batch size 1, or fail with a shape mismatch for larger batches.

Reproduction:
```python
import torch
from diffusers import Lumina2Pipeline

pipe = Lumina2Pipeline(transformer=None, scheduler=None, vae=None, text_encoder=None, tokenizer=None)
pe, pm, ne, nm = pipe.encode_prompt(
prompt=None,
do_classifier_free_guidance=True,
num_images_per_prompt=2,
prompt_embeds=torch.randn(2, 3, 4),
negative_prompt_embeds=torch.randn(2, 3, 4),
prompt_attention_mask=torch.ones(2, 3, dtype=torch.bool),
negative_prompt_attention_mask=torch.ones(2, 3, dtype=torch.bool),
)
print(pe.shape, pm.shape, ne.shape, nm.shape)
# positive batch is 4, negative batch is still 2
```

Relevant precedent:
`LuminaPipeline` repeats generated negative embeddings and masks with the positive batch.
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/lumina/pipeline_lumina.py#L360-L365

Suggested fix:
```python
if do_classifier_free_guidance:
if negative_prompt_embeds is None:
negative_prompt_embeds, negative_prompt_attention_mask = self._get_gemma_prompt_embeds(...)

negative_batch_size, neg_seq_len, _ = negative_prompt_embeds.shape
negative_prompt_embeds = negative_prompt_embeds.repeat(1, num_images_per_prompt, 1)
negative_prompt_embeds = negative_prompt_embeds.view(
negative_batch_size * num_images_per_prompt, neg_seq_len, -1
)
negative_prompt_attention_mask = negative_prompt_attention_mask.repeat(num_images_per_prompt, 1)
negative_prompt_attention_mask = negative_prompt_attention_mask.view(
negative_batch_size * num_images_per_prompt, -1
)
```

## Issue 3: Precomputed prompt embeds are not cast to transformer dtype

Affected code:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/lumina2/pipeline_lumina2.py#L290-L302
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/lumina2/pipeline_lumina2.py#L684-L692

Problem:
When `prompt_embeds` are supplied directly, `encode_prompt` does not cast them to the transformer dtype. `__call__` then creates latents with `prompt_embeds.dtype`, so float32 prompt embeds plus a bf16 transformer produce float32 latents fed into bf16 linear layers.

Impact:
Common precomputed-embedding workflows fail for bf16/quantized transformer usage.

Reproduction:
```python
import torch
from diffusers import FlowMatchEulerDiscreteScheduler, Lumina2Pipeline, Lumina2Transformer2DModel

transformer = Lumina2Transformer2DModel(
sample_size=4, patch_size=2, in_channels=4, hidden_size=8, num_layers=1,
num_refiner_layers=1, num_attention_heads=1, num_kv_heads=1,
multiple_of=16, axes_dim_rope=(4, 2, 2), axes_lens=(32, 32, 32), cap_feat_dim=8,
).to(torch.bfloat16)
pipe = Lumina2Pipeline(transformer, FlowMatchEulerDiscreteScheduler(), None, None, None)
pipe.set_progress_bar_config(disable=True)

pipe(
prompt=None,
prompt_embeds=torch.randn(1, 4, 8, dtype=torch.float32),
prompt_attention_mask=torch.ones(1, 4, dtype=torch.bool),
guidance_scale=1.0,
num_inference_steps=1,
height=32,
width=32,
output_type="latent",
)
```

Relevant precedent:
Flux/Qwen-style prompt paths normalize prompt tensors before transformer use.
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/qwenimage/pipeline_qwenimage.py#L256-L264

Suggested fix:
```python
dtype = self.transformer.dtype if self.transformer is not None else prompt_embeds.dtype
prompt_embeds = prompt_embeds.to(device=device, dtype=dtype)
if negative_prompt_embeds is not None:
negative_prompt_embeds = negative_prompt_embeds.to(device=device, dtype=dtype)
```

## Issue 4: Scheduler shift uses latent channels as image sequence length

Affected code:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/lumina2/pipeline_lumina2.py#L697-L716

Problem:
`image_seq_len = latents.shape[1]` reads the channel dimension. Lumina2 latents are BCHW at this point, so the scheduler `mu` should be based on the number of image tokens after patching.

Impact:
Resolution-dependent timestep shifting is wrong. For a 1024x1024 image with default 16-channel latents and patch size 2, the code uses `16` instead of `4096`.

Reproduction:
```python
import torch
from diffusers import Lumina2Transformer2DModel

def calculate_shift(image_seq_len, base_seq_len=256, max_seq_len=4096, base_shift=0.5, max_shift=1.15):
m = (max_shift - base_shift) / (max_seq_len - base_seq_len)
return image_seq_len * m + (base_shift - m * base_seq_len)

transformer = Lumina2Transformer2DModel(
sample_size=128, patch_size=2, in_channels=16, hidden_size=8, num_layers=0,
num_refiner_layers=0, num_attention_heads=1, num_kv_heads=1,
multiple_of=16, axes_dim_rope=(4, 2, 2), cap_feat_dim=8,
)
latents = torch.zeros(1, transformer.config.in_channels, 128, 128)
wrong = latents.shape[1]
correct = (latents.shape[2] // transformer.config.patch_size) * (latents.shape[3] // transformer.config.patch_size)
print(wrong, calculate_shift(wrong))
print(correct, calculate_shift(correct))
```

Relevant precedent:
Duplicate/related existing items: https://github.com/huggingface/diffusers/issues/12913 and https://github.com/huggingface/diffusers/pull/13272

Suggested fix:
```python
patch_size = self.transformer.config.patch_size
image_seq_len = (latents.shape[2] // patch_size) * (latents.shape[3] // patch_size)
```

## Issue 5: Lumina2 attention bypasses diffusers attention dispatch

Affected code:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/transformers/transformer_lumina2.py#L69-L145

Problem:
`Lumina2AttnProcessor2_0` calls `F.scaled_dot_product_attention` directly and has no `_attention_backend` / `_parallel_config`. `ModelMixin.set_attention_backend()` cannot configure these processors.

Impact:
Lumina2 misses the current attention backend system, including backend selection and context-parallel-compatible dispatch.

Reproduction:
```python
from diffusers import Lumina2Transformer2DModel

model = Lumina2Transformer2DModel(
sample_size=16, patch_size=2, in_channels=4, hidden_size=24,
num_layers=1, num_refiner_layers=1, num_attention_heads=3,
num_kv_heads=1, multiple_of=2, axes_dim_rope=(4, 2, 2),
axes_lens=(128, 128, 128), cap_feat_dim=32,
)
print([(m.processor.__class__.__name__, hasattr(m.processor, "_attention_backend")) for m in model.modules() if hasattr(m, "processor")])
model.set_attention_backend("native")
print([getattr(m.processor, "_attention_backend", "MISSING") for m in model.modules() if hasattr(m, "processor")])
```

Relevant precedent:
Flux processors declare `_attention_backend` and call `dispatch_attention_fn`.
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/transformers/transformer_flux.py#L75-L125

Suggested fix:
Port Lumina2 attention to the current pattern: define a Lumina2 attention module with `AttentionModuleMixin`, give the processor `_attention_backend` and `_parallel_config`, and call `dispatch_attention_fn` on `(batch, sequence, heads, head_dim)` query/key/value tensors.

## Issue 6: RoPE precomputes complex128 tensors by default

Affected code:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/transformers/transformer_lumina2.py#L243-L249

Problem:
RoPE precompute uses `torch.float64` unless MPS is globally available, which produces `torch.complex128` frequency tensors. The review rules disallow unconditional float64 and call out NPU/MPS compatibility.

Impact:
This adds unnecessary memory/cast overhead and can break unsupported float64 backends.

Reproduction:
```python
from diffusers import Lumina2Transformer2DModel

model = Lumina2Transformer2DModel(
sample_size=16, patch_size=2, in_channels=4, hidden_size=24,
num_layers=1, num_refiner_layers=1, num_attention_heads=3,
num_kv_heads=1, multiple_of=2, axes_dim_rope=(4, 2, 2),
axes_lens=(128, 128, 128), cap_feat_dim=32,
)
print([freqs.dtype for freqs in model.rope_embedder.freqs_cis])
# [torch.complex128, torch.complex128, torch.complex128]
```

Relevant precedent:
Flux gates MPS and NPU explicitly for RoPE dtype.
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/transformers/transformer_flux.py#L505-L518

Suggested fix:
```python
freqs_dtype = torch.float32
```

## Issue 7: RoPE position construction breaks `torch.compile(fullgraph=True)`

Affected code:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/transformers/transformer_lumina2.py#L263-L291

Problem:
`attention_mask.sum(dim=1).tolist()` moves tensor values into Python, then Python `max()` and per-sample loops use those values for shapes/slices. This violates the model review rule to avoid graph breaks in forward implementations.

Impact:
`torch.compile(..., fullgraph=True)` cannot compile the model, and device execution pays synchronization overhead.

Reproduction:
```python
import torch
from diffusers import Lumina2Transformer2DModel

model = Lumina2Transformer2DModel(
sample_size=16, patch_size=2, in_channels=4, hidden_size=24,
num_layers=1, num_refiner_layers=1, num_attention_heads=3,
num_kv_heads=1, multiple_of=2, axes_dim_rope=(4, 2, 2),
axes_lens=(128, 128, 128), cap_feat_dim=32,
).eval()

compiled = torch.compile(model, fullgraph=True, backend="eager")
compiled(
hidden_states=torch.randn(1, 4, 16, 16),
timestep=torch.rand(1),
encoder_hidden_states=torch.randn(1, 16, 32),
encoder_attention_mask=torch.ones(1, 16, dtype=torch.bool),
)
```

Relevant precedent:
The model review rules explicitly require avoiding graph breaks in forward implementations.

Suggested fix:
Vectorize position-id and joint-sequence assembly so mask lengths stay as tensors. If variable effective caption lengths are required, use tensor masks/scatter operations rather than `.tolist()` and Python-derived allocation sizes.

## Issue 8: Slow end-to-end Lumina2 pipeline coverage is missing

Affected code:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/tests/pipelines/lumina2/test_pipeline_lumina2.py#L16-L117
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/tests/single_file/test_lumina2_transformer.py#L30-L38

Problem:
Lumina2 has fast dummy pipeline/model tests, LoRA tests, DreamBooth example tests, and a single-file transformer loading test, but no slow end-to-end pipeline test with the real `Alpha-VLLM/Lumina-Image-2.0` components and an expected output slice.

Impact:
Scheduler/parity bugs such as the `mu` issue, dtype behavior, and real-checkpoint prompt/decoder behavior can ship without a slow regression signal.

Reproduction:
```python
from pathlib import Path

matches = []
for path in Path("tests").rglob("*lumina2*.py"):
text = path.read_text(encoding="utf-8")
if "@slow" in text or "@nightly" in text:
matches.append(str(path))
print(matches)
# []
```

Relevant precedent:
Flux has a real-checkpoint slow pipeline test class.
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/tests/pipelines/flux/test_pipeline_flux.py#L238-L275

Suggested fix:
Add a `Lumina2PipelineSlowTests` class under `tests/pipelines/lumina2/test_pipeline_lumina2.py` that loads the public checkpoint or cached test slices, runs a tiny deterministic inference, and asserts an expected image slice.

Guida per i contributori

Apri la guida per i contributori

Direzione di ricerca

Inizia dalle sezioni interessate di src/diffusers/pipelines/lumina2/pipeline_lumina2.py e src/diffusers/models/transformers/transformer_lumina2.py, quindi confronta le implementazioni citate di Lumina e Flux. Esegui i test mirati di Lumina2 dopo aver risolto l’errore locale di raccolta di torch segnalato. Il lavoro è completato quando i comportamenti segnalati relativi a costruzione, embedding, scheduler, attention-dispatch e RoPE sono coperti da regressioni superate.

Scritto dal modello di indicizzazione a partire dal testo della issue.

Valutazione

Stack tecnologico
python, pytorch
Ambito
machine-learning
Tipo di issue
Bug
Difficoltà
5/5
Tempo stimato
Più di una settimana
Stato di attività
Tranquilla
Chiarezza
Abbastanza chiara
Idoneità per principianti
35/100

Ricevi le nuove issue nella tua casella

Un breve riepilogo di issue GitHub adatte ai principianti.