huggingface / huggingface/diffusers

pixart_alpha model/pipeline review

Đang mở
#13,631 0 bình luận 0 reaction 0 người được giao Xem trên GitHub
Ngôn ngữ chính
Python
Star
34.5k
Fork
7.3k
Merge trung bình
3 ngày 3 giờ
Pull request đã merge (30 ngày)
91

Mô tả

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

Hướng dẫn đóng góp

Mở hướng dẫn đóng góp

Hướng nghiên cứu

Bắt đầu với các tệp PixArt bị ảnh hưởng: pipeline_pixart_alpha.py, pipeline_pixart_sigma.py, pixart_alpha/__init__.py và pixart_transformer_2d.py, sau đó chạy các snippet tái hiện cho từng phát hiện. So sánh Alpha với Sigma và kiểm tra các API scheduler và attention được tham chiếu. Công việc được hoàn tất khi năm trường hợp đã báo cáo không còn thất bại và phạm vi kiểm thử hồi quy bao gồm suy luận một bước, xác thực kích thước, các kiểu dữ liệu hỗn hợp, lazy export và trạng thái fuse/unfuse của QKV.

Do mô hình lập chỉ mục viết ra từ nội dung của issue.

Đánh giá

Công nghệ
python, pytorch
Lĩnh vực
machine-learning, testing-qa
Loại issue
Lỗi
Độ khó
4/5
Thời gian dự kiến
3-5 ngày
Mức độ hoạt động
Ít trao đổi
Độ rõ ràng
Đặc tả rõ ràng
Mức phù hợp với người mới
52/100

Nhận issue mới trong hộp thư của bạn

Bản tóm tắt ngắn những issue GitHub phù hợp với người mới.