huggingface / huggingface/diffusers
allegro model/pipeline review
- Dominant language
- Python
- Stars
- 34.5k
- Forks
- 7.3k
- Avg merge
- 3d 3h
- Merged PRs (30d)
- 91
Description
# `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.
Contributor guide
Research direction
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.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- api, machine-learning, testing-qa
- Issue type
- Bug
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Quiet
- Clarity
- Mostly clear
- Newbie friendliness
- 35/100