Skip to content

Fix consistency_models class-label sampling and latent shapes - #13679

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

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

Conversation

@hlky

@hlkyhlky commented May 5, 2026

Copy link
Copy Markdown
Contributor

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 freeto join this conversation on GitHub. Already have an account? Sign in to comment

Labels

fixes-issuepipelinessize/LPR with diff > 200 LOCtests

Projects

None yet

Development

Successfully merging this pull request may close these issues.

consistency_models model/pipeline review

1 participant

@hlky