huggingface / huggingface/diffusers

pixart_alpha model/pipeline review

Open
#13,631 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

# `pixart_alpha` model/pipeline review

Commit tested: `0f1abc4ae8b0eb2a3b40e82a310507281144c423`

Review performed against the repository review rules.

Duplicate search: searched GitHub Issues and PRs for `pixart_alpha`, affected class/function/file names, and failure modes. Existing duplicates/related items found for Issue 1 and Issue 5; no duplicates found for the other report items.

Coverage status: fast and slow tests exist for PixArt Alpha, PixArt Sigma, and `PixArtTransformer2DModel`. No missing slow tests found for the listed target files. Current coverage misses the DPM one-step path, non-patch-aligned image sizes, mixed transformer/VAE dtype decode, lazy constant exports, and QKV unfuse edge cases.

## Issue 1: Existing duplicate: PixArtAlpha one-step branch still breaks scheduler outputs

Affected code:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/pixart_alpha/pipeline_pixart_alpha.py#L949-L952
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/schedulers/scheduling_dpmsolver_multistep.py#L1277-L1278

Problem:
`PixArtAlphaPipeline` indexes `scheduler.step(..., return_dict=False)[1]` whenever `num_inference_steps == 1`. `DPMSolverMultistepScheduler`, the pipeline’s declared scheduler type, returns a one-item tuple, so one-step inference crashes. Existing open duplicate/related issue: https://github.com/huggingface/diffusers/issues/8689 covers this same one-step special case for another scheduler.

Impact:
One-step PixArt Alpha inference is broken for standard scheduler outputs and for scheduler swaps.

Reproduction:
```python
import torch
from diffusers import AutoencoderKL, DPMSolverMultistepScheduler, PixArtAlphaPipeline, PixArtTransformer2DModel

transformer = PixArtTransformer2DModel(
sample_size=8, num_layers=1, patch_size=2, attention_head_dim=2, num_attention_heads=2,
in_channels=4, cross_attention_dim=8, out_channels=8, use_additional_conditions=False,
).eval()

pipe = PixArtAlphaPipeline(None, None, AutoencoderKL().eval(), transformer, DPMSolverMultistepScheduler())
embeds = torch.randn(1, 8, 8)
mask = torch.ones(1, 8, dtype=torch.long)

pipe(prompt_embeds=embeds, prompt_attention_mask=mask, guidance_scale=1.0,
num_inference_steps=1, use_resolution_binning=False, output_type="latent")
```

Relevant precedent:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/pixart_alpha/pipeline_pixart_sigma.py#L881-L882

Suggested fix:
```python
latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0]
```

## Issue 2: Height/width validation allows sizes that cannot be patchified

Affected code:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/pixart_alpha/pipeline_pixart_alpha.py#L468-L469
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/pixart_alpha/pipeline_pixart_sigma.py#L413-L414

Problem:
Both pipelines only require `height` and `width` to be divisible by 8. Real PixArt latents are divided by `vae_scale_factor` and then by `transformer.config.patch_size`, so dimensions must be divisible by `vae_scale_factor * patch_size` unless resolution binning changes them first.

Impact:
`use_resolution_binning=False` accepts documented-valid sizes and then crashes inside denoising.

Reproduction:
```python
import torch
from diffusers import AutoencoderKL, DDIMScheduler, PixArtAlphaPipeline, PixArtTransformer2DModel

transformer = PixArtTransformer2DModel(
sample_size=8, num_layers=1, patch_size=2, attention_head_dim=2, num_attention_heads=2,
in_channels=4, cross_attention_dim=8, out_channels=8, use_additional_conditions=False,
).eval()
vae = AutoencoderKL(
block_out_channels=(8, 8, 8, 8),
down_block_types=("DownEncoderBlock2D",) * 4,
up_block_types=("UpDecoderBlock2D",) * 4,
norm_num_groups=4,
).eval()

pipe = PixArtAlphaPipeline(None, None, vae, transformer, DDIMScheduler())
embeds = torch.randn(1, 8, 8)
mask = torch.ones(1, 8, dtype=torch.long)

pipe(prompt_embeds=embeds, prompt_attention_mask=mask, guidance_scale=1.0,
num_inference_steps=1, height=24, width=24,
use_resolution_binning=False, output_type="latent")
```

Relevant precedent:
Sana casts provided latents to the requested dtype and validates around its latent geometry:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/sana/pipeline_sana.py#L688-L705

Suggested fix:
```python
divisor = self.vae_scale_factor * self.transformer.config.patch_size
if height % divisor != 0 or width % divisor != 0:
raise ValueError(f"`height` and `width` have to be divisible by {divisor} but are {height} and {width}.")
```

## Issue 3: PixArtAlpha decode does not cast latents to VAE dtype

Affected code:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/pixart_alpha/pipeline_pixart_alpha.py#L964-L966

Problem:
Alpha decodes `latents / scaling_factor` directly. If the transformer runs in `float16` while the VAE remains `float32`, decode fails with an input/bias dtype mismatch. Sigma already casts latents to `self.vae.dtype`.

Impact:
Mixed precision component loading, partial offload workflows, or manually supplied components can fail at the final decode step.

Reproduction:
```python
import torch
from diffusers import AutoencoderKL, DDIMScheduler, PixArtAlphaPipeline, PixArtTransformer2DModel

transformer = PixArtTransformer2DModel(
sample_size=8, num_layers=1, patch_size=2, attention_head_dim=2, num_attention_heads=2,
in_channels=4, cross_attention_dim=8, out_channels=8, use_additional_conditions=False,
).eval().to(dtype=torch.float16)

pipe = PixArtAlphaPipeline(None, None, AutoencoderKL().eval(), transformer, DDIMScheduler())
embeds = torch.randn(1, 8, 8, dtype=torch.float16)
mask = torch.ones(1, 8, dtype=torch.long)

pipe(prompt_embeds=embeds, prompt_attention_mask=mask, guidance_scale=1.0,
num_inference_steps=1, use_resolution_binning=False, output_type="np")
```

Relevant precedent:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/pixart_alpha/pipeline_pixart_sigma.py#L894-L895

Suggested fix:
```python
image = self.vae.decode(latents.to(self.vae.dtype) / self.vae.config.scaling_factor, return_dict=False)[0]
```

## Issue 4: PixArt aspect-ratio constants are imported only in the non-lazy branch

Affected code:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/pixart_alpha/__init__.py#L25-L42

Problem:
The `TYPE_CHECKING`/slow-import branch imports `ASPECT_RATIO_256_BIN`, `ASPECT_RATIO_512_BIN`, `ASPECT_RATIO_1024_BIN`, and `ASPECT_RATIO_2048_BIN`, but `_import_structure` exposes only the two pipeline classes. Normal lazy imports fail for names that the non-lazy branch advertises.

Impact:
Public subpackage imports are inconsistent across lazy and slow import modes.

Reproduction:
```python
from diffusers.pipelines.pixart_alpha import ASPECT_RATIO_1024_BIN
```

Relevant precedent:
The same file already imports these constants in the eager branch.

Suggested fix:
```python
_import_structure["pipeline_pixart_alpha"] = [
"ASPECT_RATIO_256_BIN",
"ASPECT_RATIO_512_BIN",
"ASPECT_RATIO_1024_BIN",
"PixArtAlphaPipeline",
]
_import_structure["pipeline_pixart_sigma"] = ["ASPECT_RATIO_2048_BIN", "PixArtSigmaPipeline"]
```

## Issue 5: Existing related issue: PixArtTransformer QKV unfuse state is not safe

Affected code:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/transformers/pixart_transformer_2d.py#L203-L225

Problem:
`original_attn_processors` is only created inside `fuse_qkv_projections()`. Calling `unfuse_qkv_projections()` before `fuse_qkv_projections()` raises `AttributeError`, and calling fuse twice overwrites the saved original processors with fused processors. Related existing issue: https://github.com/huggingface/diffusers/issues/13592 reports the same copied QKV state bug for UNet.

Impact:
The public attention optimization API is not idempotent and can leave PixArt fused after an enable-twice-disable flow.

Reproduction:
```python
from diffusers import PixArtTransformer2DModel

model = PixArtTransformer2DModel(
sample_size=8, num_layers=1, attention_head_dim=2, num_attention_heads=2,
cross_attention_dim=8, num_embeds_ada_norm=8, use_additional_conditions=False,
)

model.unfuse_qkv_projections() # AttributeError

model.fuse_qkv_projections()
model.fuse_qkv_projections()
model.unfuse_qkv_projections()
print({p.__class__.__name__ for p in model.attn_processors.values()}) # still fused
```

Relevant precedent:
The method doc says it disables fused projection “if enabled”.

Suggested fix:
```python
# in __init__
self.original_attn_processors = None

# in fuse_qkv_projections
if self.original_attn_processors is None:
self.original_attn_processors = self.attn_processors

# in unfuse_qkv_projections
if self.original_attn_processors is not None:
self.set_attn_processor(self.original_attn_processors)
self.original_attn_processors = None
```

Contributor guide

Open the contributing guide

Research direction

Start with the affected PixArt files: pipeline_pixart_alpha.py, pipeline_pixart_sigma.py, pixart_alpha/__init__.py, and pixart_transformer_2d.py, then run the reproduction snippets for each finding. Compare Alpha with Sigma and inspect the referenced scheduler and attention APIs. Done means the five reported cases no longer fail and regression coverage exercises one-step inference, dimension validation, mixed dtypes, lazy exports, and QKV fuse/unfuse state.

Written by the indexing model from the issue text.

Assessment

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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.