huggingface / huggingface/diffusers

consistency_models model/pipeline review

Đang mở
#13,643 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ả

# `consistency_models` model/pipeline review

Commit tested: `0f1abc4ae8b0eb2a3b40e82a310507281144c423`

Review performed against the repository review rules.

## Issue 1: Random class labels ignore the supplied generator

Affected code:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/consistency_models/pipeline_consistency_models.py#L132-L143

Problem:
For class-conditional UNets, omitting `class_labels` makes the pipeline sample random labels with `torch.randint(...)`, but `prepare_class_labels` does not receive or use the user-supplied `generator`. Seeded inference is therefore not fully controlled by `generator`.

Impact:
Two calls with identical explicit generators can produce different images when `class_labels=None`, which breaks the usual diffusers reproducibility contract.

Reproduction:
```python
import torch
from diffusers import CMStochasticIterativeScheduler, ConsistencyModelPipeline, UNet2DModel

unet = UNet2DModel(
sample_size=8,
in_channels=3,
out_channels=3,
layers_per_block=1,
block_out_channels=(8,),
down_block_types=("DownBlock2D",),
up_block_types=("UpBlock2D",),
norm_num_groups=1,
num_class_embeds=10,
)
pipe = ConsistencyModelPipeline(
unet=unet,
scheduler=CMStochasticIterativeScheduler(num_train_timesteps=40, sigma_min=0.002, sigma_max=80.0),
).to("cpu")
pipe.set_progress_bar_config(disable=True)

latents = torch.zeros((1, 3, 8, 8))
torch.manual_seed(0)

out_a = pipe(
latents=latents,
generator=torch.Generator(device="cpu").manual_seed(123),
class_labels=None,
num_inference_steps=1,
output_type="pt",
).images
out_b = pipe(
latents=latents,
generator=torch.Generator(device="cpu").manual_seed(123),
class_labels=None,
num_inference_steps=1,
output_type="pt",
).images

print(torch.equal(out_a, out_b), (out_a - out_b).abs().max().item())
```

Relevant precedent:
`randn_tensor` handles explicit generators, including CPU generators that later move tensors to the execution device:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/utils/torch_utils.py#L152-L199

Suggested fix:
```python
def prepare_class_labels(self, batch_size, device, generator=None, class_labels=None):
if self.unet.config.num_class_embeds is None:
return None

if isinstance(class_labels, list):
class_labels = torch.tensor(class_labels, dtype=torch.long)
elif isinstance(class_labels, int):
class_labels = torch.tensor([class_labels] * batch_size, dtype=torch.long)
elif class_labels is None:
if isinstance(generator, list):
class_labels = torch.cat(
[
torch.randint(
0,
self.unet.config.num_class_embeds,
size=(1,),
generator=g,
device=g.device,
).cpu()
for g in generator
]
)
else:
rand_device = generator.device if generator is not None else torch.device("cpu")
class_labels = torch.randint(
0,
self.unet.config.num_class_embeds,
size=(batch_size,),
generator=generator,
device=rand_device,
)

return class_labels.to(device=device, dtype=torch.long)
```

Also pass `generator=generator` from `__call__`.

Duplicate check:
No matching existing issue or PR found for `ConsistencyModelPipeline class_labels generator` or `consistency_models class_labels random`.

## Issue 2: Latent shape handling hardcodes RGB square images

Affected code:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/consistency_models/pipeline_consistency_models.py#L158-L161
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/consistency_models/pipeline_consistency_models.py#L223-L241

Problem:
`check_inputs` expects latent shape `(batch_size, 3, img_size, img_size)`, even though `prepare_latents` uses `self.unet.config.in_channels`. The same `img_size` value is passed as both height and width, so tuple `sample_size` configs also fail before inference.

Impact:
Valid `UNet2DModel` configs with non-3 channel counts or non-square `sample_size` cannot be used through this pipeline, and valid user-supplied latents are rejected.

Reproduction:
```python
import torch
from diffusers import CMStochasticIterativeScheduler, ConsistencyModelPipeline, UNet2DModel

scheduler = CMStochasticIterativeScheduler(num_train_timesteps=40, sigma_min=0.002, sigma_max=80.0)

one_channel_unet = UNet2DModel(
sample_size=8,
in_channels=1,
out_channels=1,
layers_per_block=1,
block_out_channels=(8,),
down_block_types=("DownBlock2D",),
up_block_types=("UpBlock2D",),
norm_num_groups=1,
)
pipe = ConsistencyModelPipeline(unet=one_channel_unet, scheduler=scheduler).to("cpu")
pipe.set_progress_bar_config(disable=True)

try:
pipe(latents=torch.zeros((1, 1, 8, 8)), num_inference_steps=1, output_type="pt")
except Exception as e:
print(type(e).__name__, e)

rect_unet = UNet2DModel(
sample_size=(8, 10),
in_channels=3,
out_channels=3,
layers_per_block=1,
block_out_channels=(8,),
down_block_types=("DownBlock2D",),
up_block_types=("UpBlock2D",),
norm_num_groups=1,
)
pipe = ConsistencyModelPipeline(unet=rect_unet, scheduler=scheduler).to("cpu")
pipe.set_progress_bar_config(disable=True)

try:
pipe(num_inference_steps=1, output_type="pt")
except Exception as e:
print(type(e).__name__, e)
```

Relevant precedent:
`DDPMPipeline` derives shape from both `unet.config.in_channels` and tuple `sample_size`:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/ddpm/pipeline_ddpm.py#L100-L109

Suggested fix:
```python
sample_size = self.unet.config.sample_size
if isinstance(sample_size, int):
height = width = sample_size
else:
height, width = sample_size

self.check_inputs(num_inference_steps, timesteps, latents, batch_size, height, width, callback_steps)

sample = self.prepare_latents(
batch_size=batch_size,
num_channels=self.unet.config.in_channels,
height=height,
width=width,
dtype=self.unet.dtype,
device=device,
generator=generator,
latents=latents,
)
```

and update validation to:
```python
expected_shape = (batch_size, self.unet.config.in_channels, height, width)
```

Duplicate check:
No matching existing issue or PR found for `ConsistencyModelPipeline sample_size tuple` or `pipeline_consistency_models in_channels latents`.

## Coverage / Search Status

Reviewed: public lazy imports and top-level exports, pipeline config/loading surface, runtime dtype/device/offload path, scheduler interaction, class-conditioning path, docs, examples, deprecation status, and tests under `tests/pipelines/consistency_models`.

Fast tests exist at `tests/pipelines/consistency_models/test_consistency_models.py`. Slow tests also exist under `ConsistencyModelPipelineSlowTests`, so slow coverage is not missing. Current tests do not cover omitted class labels on class-conditional models, non-RGB latent shapes, or tuple `sample_size`.

I attempted `.venv` pytest collection for the target test file, but this environment's Torch build is missing `torch._C._distributed_c10d`, which breaks collection through shared test utilities. The standalone reproductions above were run with `.venv`.

Duplicate searches were run with `gh search issues --include-prs` for `consistency_models`, `ConsistencyModelPipeline`, `pipeline_consistency_models`, `CMStochasticIterativeScheduler ConsistencyModelPipeline`, and the specific class-label / latent-shape failure modes. Existing results included older consistency-model sampling/training items, but none matched the two findings above.

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

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

Hướng nghiên cứu

Start with src/diffusers/pipelines/consistency_models/pipeline_consistency_models.py, especially prepare_class_labels, check_inputs, prepare_latents, and __call__. Read the comparable shape handling in pipeline_ddpm.py, then run tests/pipelines/consistency_models/test_consistency_models.py. Done means seeded omitted class labels and valid non-RGB or rectangular latent shapes work, with regression coverage for both findings.

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
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
62/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.