huggingface / huggingface/diffusers

pixart_alpha model/pipeline review

オープン
#13,631 コメント 0 件 リアクション 0 件 担当者 0 名 GitHub で見る
主要言語
Python
スター
34.5k
フォーク
7.3k
平均マージ
3日 3時間
マージ済み PR(30日)
91

説明

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

コントリビューションガイド

コントリビューションガイドを開く

調査の方向性

影響を受ける PixArt ファイル: pipeline_pixart_alpha.py、pipeline_pixart_sigma.py、pixart_alpha/__init__.py、pixart_transformer_2d.py から始め、その後、各指摘について再現スニペットを実行してください。Alpha と Sigma を比較し、参照されている scheduler API と attention API を調査してください。5 件の報告されたケースで失敗しなくなり、回帰テストで 1 ステップ推論、次元検証、混在する dtype、lazy export、QKV の fuse/unfuse 状態がカバーされれば完了です。

索引モデルが issue の本文から書いたものです。

評価

技術スタック
python, pytorch
領域
machine-learning, testing-qa
issue の種類
バグ
難易度
4/5
見積もり時間
3〜5日
活発さ
静か
明瞭さ
明確に書かれている
初心者へのやさしさ
52/100

新しい issue をメールで受け取る

初心者向けの GitHub issue を短くまとめたダイジェスト。