Skip to content

Fix consistency_models class-label sampling and latent shapes#13679

Open
hlky wants to merge 1 commit intohuggingface:mainfrom
hlky:codex/consistency-models-review-fixes
Open

Fix consistency_models class-label sampling and latent shapes#13679
hlky wants to merge 1 commit intohuggingface:mainfrom
hlky:codex/consistency-models-review-fixes

Conversation

@hlky
Copy link
Copy Markdown
Contributor

@hlky hlky commented May 5, 2026

What does this PR do?

Fixes #13643.

Selected target: ConsistencyModelPipeline (consistency_models).

No duplicate fix was found.

  • Uses the supplied generator when sampling random class labels for class-conditional consistency models.
  • Preserves existing explicit class-label dtype and scalar-label behavior.
  • Derives latent validation and preparation shapes from unet.config.in_channels and unet.config.sample_size.
  • Supports tuple sample_size configs.
  • Replaces the local image postprocess helper with VaeImageProcessor.
  • Adds focused fast tests for the fixed behavior.
  • Adds type hints and updates docstrings for the touched pipeline methods.

AGENTS.md suggestions

  • Prefer shared processor utilities for standard pipeline output conversion. For image outputs, use VaeImageProcessor.postprocess(...) instead of hand-rolling denormalization, tensor-to-NumPy conversion, or PIL conversion logic inside the pipeline.
  • Keep public and helper type hints aligned with runtime behavior. If an argument accepts None, generator lists, tensor/list/int unions, or returns either an output dataclass or tuple, reflect that in the annotation.
  • Keep docstrings synchronized with signatures and behavior. Public docs should mention supported output_type values, generator-list support, and behavior that changes when optional inputs such as class_labels are omitted.

Tests

.\.venv\Scripts\python.exe -m pytest tests/pipelines/consistency_models/test_consistency_models.py -q -k "random_class_labels_use_generator or latents_use_unet_in_channels or tuple_sample_size"
.\.venv\Scripts\python.exe -m ruff check src/diffusers/pipelines/consistency_models/pipeline_consistency_models.py tests/pipelines/consistency_models/test_consistency_models.py
.\.venv\Scripts\python.exe -m ruff format --check src/diffusers/pipelines/consistency_models/pipeline_consistency_models.py tests/pipelines/consistency_models/test_consistency_models.py
git diff --check upstream/main...HEAD

Result: tests and checks passed.

@hlky hlky marked this pull request as ready for review May 5, 2026 00:50
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Projects

None yet

Development

Successfully merging this pull request may close these issues.

consistency_models model/pipeline review

1 participant