huggingface / huggingface/diffusers

allegro model/pipeline review

Aperta
#13,647 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

# `allegro` model/pipeline review

Commit tested: `0f1abc4ae8b0eb2a3b40e82a310507281144c423`

Review performed against the repository review rules.

## Issue 1: VAE decode/encode fail unless tiling is enabled

Affected code:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/autoencoders/autoencoder_kl_allegro.py#L798-L841
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/allegro/pipeline_allegro.py#L972-L976

Problem:
`AutoencoderKLAllegro.encode()` and `.decode()` raise `NotImplementedError` unless `vae.enable_tiling()` was called. The pipeline also fails by default when `output_type != "latent"` unless the user remembered to enable tiling. This also breaks save/load parity because `use_tiling` is runtime state, not serialized config.

Impact:
Basic public VAE APIs and default pipeline inference are fragile. Docs include a quantized `AllegroPipeline` example that does not enable tiling, so that path can fail at decode time.

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

vae = AutoencoderKLAllegro(
block_out_channels=(8, 8, 8, 8),
latent_channels=4,
layers_per_block=1,
norm_num_groups=2,
)
vae.decode(torch.randn(1, 4, 2, 2, 2))
```

Relevant precedent:
Related mitigation, not a duplicate full fix: https://github.com/huggingface/diffusers/pull/10212

Suggested fix:
```python
def _encode(self, x):
if self.use_tiling:
return self.tiled_encode(x)
batch_size = x.shape[0]
h = self.encoder(x)
h = h.permute(0, 2, 1, 3, 4).flatten(0, 1)
h = self.quant_conv(h)
return h.unflatten(0, (batch_size, -1)).permute(0, 2, 1, 3, 4)

def _decode(self, z):
if self.use_tiling:
return self.tiled_decode(z)
batch_size = z.shape[0]
z = z.permute(0, 2, 1, 3, 4).flatten(0, 1)
z = self.post_quant_conv(z)
z = z.unflatten(0, (batch_size, -1)).permute(0, 2, 1, 3, 4)
return self.decoder(z)
```

## Issue 2: `num_videos_per_prompt` is silently ignored

Affected code:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/allegro/pipeline_allegro.py#L824
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/allegro/pipeline_allegro.py#L868-L900

Problem:
`__call__` accepts and documents `num_videos_per_prompt`, but line 824 overwrites any user value with `1`.

Impact:
Users requesting multiple videos per prompt always receive one video, with no warning or error.

Reproduction:
```python
import torch
from diffusers import AllegroPipeline, AllegroTransformer3DModel, AutoencoderKLAllegro, DDIMScheduler

transformer = AllegroTransformer3DModel(
num_attention_heads=2, attention_head_dim=12, in_channels=4, out_channels=4,
num_layers=1, cross_attention_dim=24, sample_width=8, sample_height=8,
sample_frames=8, caption_channels=24,
)
vae = AutoencoderKLAllegro(block_out_channels=(8, 8, 8, 8), latent_channels=4, layers_per_block=1, norm_num_groups=2)
pipe = AllegroPipeline(None, None, vae, transformer, DDIMScheduler())
pipe.set_progress_bar_config(disable=True)

out = pipe(
prompt_embeds=torch.randn(1, 16, 24),
prompt_attention_mask=torch.ones(1, 16, dtype=torch.long),
guidance_scale=1.0,
num_videos_per_prompt=3,
num_inference_steps=1,
height=16, width=16, num_frames=8,
output_type="latent",
).frames
print(out.shape) # torch.Size([1, 4, 2, 2, 2]), expected batch 3
```

Relevant precedent:
`HunyuanVideoPipeline` preserves `num_videos_per_prompt` through prompt encoding and latent preparation.

Suggested fix:
```python
# Remove this line from __call__:
num_videos_per_prompt = 1
```

## Issue 3: Custom timestep scheduler state is overwritten

Affected code:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/allegro/pipeline_allegro.py#L892-L895

Problem:
`retrieve_timesteps()` correctly applies custom `timesteps`, but the next line calls `self.scheduler.set_timesteps(num_inference_steps, device=device)` again, replacing scheduler internal state with default timesteps while the loop still iterates over the custom local `timesteps`.

Impact:
Schedulers whose `step()` depends on `scheduler.timesteps` can use mismatched sigma/step indices for custom timesteps.

Reproduction:
```python
import torch
from diffusers import AllegroPipeline, AllegroTransformer3DModel, AutoencoderKLAllegro, EulerDiscreteScheduler

transformer = AllegroTransformer3DModel(
num_attention_heads=2, attention_head_dim=12, in_channels=4, out_channels=4,
num_layers=1, cross_attention_dim=24, sample_width=8, sample_height=8,
sample_frames=8, caption_channels=24,
)
vae = AutoencoderKLAllegro(block_out_channels=(8, 8, 8, 8), latent_channels=4, layers_per_block=1, norm_num_groups=2)
pipe = AllegroPipeline(None, None, vae, transformer, EulerDiscreteScheduler(num_train_timesteps=1000))
pipe.set_progress_bar_config(disable=True)

seen = []
def cb(pipe, i, t, kwargs):
seen.append((int(t), [int(x) for x in pipe.scheduler.timesteps[:2]]))
return kwargs

pipe(
prompt_embeds=torch.randn(1, 16, 24),
prompt_attention_mask=torch.ones(1, 16, dtype=torch.long),
guidance_scale=1.0,
timesteps=[999, 500],
height=16, width=16, num_frames=8,
output_type="latent",
callback_on_step_end=cb,
)
print(seen) # [(999, [999, 0]), (500, [999, 0])]
```

Relevant precedent:
QwenImage and Hunyuan pipelines call `retrieve_timesteps()` once and do not reset the scheduler afterward.

Suggested fix:
```python
timesteps, num_inference_steps = retrieve_timesteps(
self.scheduler, num_inference_steps, timestep_device, timesteps
)
# Delete the second set_timesteps call.
```

## Issue 4: Allegro attention ignores the attention backend dispatcher

Affected code:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/attention_processor.py#L1993-L2066
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/transformers/transformer_allegro.py#L81-L100

Problem:
`AllegroAttnProcessor2_0` calls `F.scaled_dot_product_attention` directly and has no `_attention_backend` attribute. `model.set_attention_backend()` therefore cannot route Allegro attention through `dispatch_attention_fn`.

Impact:
Allegro cannot reliably use newer diffusers attention backends, backend-specific validation, or context-parallel-compatible attention paths.

Reproduction:
```python
from diffusers import AllegroTransformer3DModel
from diffusers.models.attention_processor import Attention

model = AllegroTransformer3DModel(
num_attention_heads=2, attention_head_dim=12, in_channels=4, out_channels=4,
num_layers=1, cross_attention_dim=24, sample_width=8, sample_height=8,
sample_frames=8, caption_channels=24,
)
model.set_attention_backend("native")
print([(n, hasattr(m.processor, "_attention_backend")) for n, m in model.named_modules() if isinstance(m, Attention)])
# [('transformer_blocks.0.attn1', False), ('transformer_blocks.0.attn2', False)]
```

Relevant precedent:
The current review rules point to Flux/Wan/Qwen-style processors that call `dispatch_attention_fn`.

Suggested fix:
Move Allegro attention to the modern model-local attention pattern, add `_attention_backend` and `_parallel_config` on the processor, and replace the direct SDPA call with `dispatch_attention_fn(..., backend=self._attention_backend, parallel_config=self._parallel_config)`.

## Issue 5: `AllegroPipelineOutput` is not exported from the Allegro package

Affected code:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/allegro/__init__.py#L13-L35
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/allegro/pipeline_output.py#L10-L22

Problem:
The output class exists and is documented, but `from diffusers.pipelines.allegro import AllegroPipelineOutput` fails because `pipeline_output` is absent from the lazy import structure.

Impact:
Public import behavior is inconsistent with many pipeline packages that expose their output dataclasses.

Reproduction:
```python
from diffusers.pipelines.allegro import AllegroPipelineOutput
```

Relevant precedent:
`src/diffusers/pipelines/qwenimage/__init__.py` exports `QwenImagePipelineOutput` through `_import_structure["pipeline_output"]`.

Suggested fix:
```python
_import_structure = {"pipeline_output": ["AllegroPipelineOutput"]}

# In the TYPE_CHECKING / slow import branch:
from .pipeline_output import AllegroPipelineOutput
```

## Issue 6: Test coverage does not exercise meaningful Allegro VAE decode behavior

Affected code:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/tests/pipelines/allegro/test_allegro.py#L153-L180
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/tests/pipelines/allegro/test_allegro.py#L279-L338
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/tests/pipelines/allegro/test_allegro.py#L356-L377

Problem:
There is no standalone `AutoencoderKLAllegro` model test. Several important pipeline tests are skipped because non-tiled decoding is unimplemented, and the fast/slow inference tests compare against random tensors rather than fixed expected outputs.

Impact:
The current tests can pass while decode is unimplemented or returns meaningless output, so regressions in the VAE and pipeline output quality are not caught.

Reproduction:
```python
from pathlib import Path

print(Path("tests/models/autoencoders/test_models_autoencoder_kl_allegro.py").exists())
text = Path("tests/pipelines/allegro/test_allegro.py").read_text()
print("Decoding without tiling is not yet implemented" in text)
print("expected_video = torch.randn" in text)
```

Relevant precedent:
Other video VAEs such as CogVideoX/Wan have model-level VAE coverage and pipeline tests with deterministic expected slices.

Suggested fix:
Add a dedicated `AutoencoderKLAllegro` fast test using `sample_size=16` and tiny channels, cover `encode`, `decode`, save/load, slicing, and tiling. Replace random expected tensors in pipeline tests with deterministic slices from known-good outputs.

Duplicate-search status: searched GitHub Issues and PRs for `Allegro`, `pipeline_allegro.py`, `autoencoder_kl_allegro.py`, `transformer_allegro.py`, `AutoencoderKLAllegro Decoding without tiling`, `Allegro num_videos_per_prompt`, `Allegro timesteps scheduler.set_timesteps`, `AllegroPipelineOutput`, and `AllegroAttnProcessor2_0 attention backend`. I found related historical PRs, especially https://github.com/huggingface/diffusers/pull/10212, but no open duplicate for the actionable issues above.

Test execution note: targeted `.venv` snippets were run successfully. A direct pytest collection of `tests/pipelines/allegro/test_allegro.py::AllegroPipelineFastTests::test_inference` failed in this Windows environment because the installed PyTorch lacks `torch._C._distributed_c10d`, matching the known class of issue in https://github.com/huggingface/diffusers/issues/12409.

Guida per i contributori

Apri la guida per i contributori

Direzione di ricerca

Start with the affected Allegro files: autoencoder_kl_allegro.py, pipeline_allegro.py, attention_processor.py, transformer_allegro.py, the Allegro package exports, and tests/pipelines/allegro/test_allegro.py. Run the supplied reproductions and inspect related Qwen, HunyuanVideo, Flux, Wan, and CogVideoX implementations. Done means the six reported behaviors are corrected or covered by meaningful deterministic tests, with targeted tests passing in a supported environment.

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

Valutazione

Stack tecnologico
python, pytorch
Ambito
api, machine-learning, testing-qa
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.