huggingface / huggingface/diffusers

consistency_models model/pipeline review

Aperta
#13,643 0 commenti 0 reazioni 0 assegnatari Vedi su GitHub
Lingua principale
Python
Stelle
34.5k
Fork
7.3k
Merge medio
3g 3h
PR unite (30g)
91

Descrizione

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

Guida per i contributori

Apri la guida per i contributori

Direzione di ricerca

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.

Scritto dal modello di indicizzazione a partire dal testo della issue.

Valutazione

Stack tecnologico
python, pytorch
Ambito
machine-learning
Tipo di issue
Bug
Difficoltà
4/5
Tempo stimato
3-5 giorni
Stato di attività
Tranquilla
Chiarezza
Specificata chiaramente
Idoneità per principianti
62/100

Ricevi le nuove issue nella tua casella

Un breve riepilogo di issue GitHub adatte ai principianti.