huggingface / huggingface/diffusers
consistency_models model/pipeline review
- 主要言語
- Python
- スター
- 34.5k
- フォーク
- 7.3k
- 平均マージ
- 3日 3時間
- マージ済み PR(30日)
- 91
説明
# `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.
コントリビューションガイド
調査の方向性
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.
索引モデルが issue の本文から書いたものです。
評価
- 技術スタック
- python, pytorch
- 領域
- machine-learning
- issue の種類
- バグ
- 難易度
- 4/5
- 見積もり時間
- 3〜5日
- 活発さ
- 静か
- 明瞭さ
- 明確に書かれている
- 初心者へのやさしさ
- 62/100