Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -113,7 +113,8 @@ def __call__(
generation. Can be used to tweak the same generation with different prompts. If not provided, a latents
tensor is generated by sampling using the supplied random `generator`.
output_type (`str`, *optional*, defaults to `"pil"`):
The output format of the generated image. Choose between `PIL.Image` or `np.array`.
The output format of the generated image. Choose between `"pil"` (`PIL.Image`), `"np"` (`np.array`) or
`"pt"` (`torch.Tensor`).
return_dict (`bool`, *optional*, defaults to `True`):
Whether or not to return a [`ImagePipelineOutput`] instead of a plain tuple.

Expand Down Expand Up @@ -220,9 +221,12 @@ def __call__(
image = self.vqvae.decode(latents).sample

image = (image / 2 + 0.5).clamp(0, 1)
image = image.cpu().permute(0, 2, 3, 1).numpy()
if output_type == "pil":
image = self.numpy_to_pil(image)

if output_type != "pt":
image = image.cpu().permute(0, 2, 3, 1).numpy()

if output_type == "pil":
image = self.numpy_to_pil(image)
Comment on lines 223 to +229

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

To support outputting in "pt" so that our testing mixin can handle it.


if not return_dict:
return (image,)
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -97,7 +97,8 @@ def __call__(
A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make
generation deterministic.
output_type (`str`, *optional*, defaults to `"pil"`):
The output format of the generated image. Choose between `PIL.Image` or `np.array`.
The output format of the generated image. Choose between `"pil"` (`PIL.Image`), `"np"` (`np.array`) or
`"pt"` (`torch.Tensor`).
return_dict (`bool`, *optional*, defaults to `True`):
Whether or not to return a [`ImagePipelineOutput`] instead of a plain tuple.

Expand Down Expand Up @@ -185,10 +186,12 @@ def __call__(
image = self.vqvae.decode(latents).sample
image = torch.clamp(image, -1.0, 1.0)
image = image / 2 + 0.5
image = image.cpu().permute(0, 2, 3, 1).numpy()

if output_type == "pil":
image = self.numpy_to_pil(image)
if output_type != "pt":
image = image.cpu().permute(0, 2, 3, 1).numpy()

if output_type == "pil":
image = self.numpy_to_pil(image)

if not return_dict:
return (image,)
Expand Down
Original file line number Diff line number Diff line change
@@ -1,8 +1,22 @@
# coding=utf-8
# Copyright 2026 HuggingFace Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import gc
import inspect
import unittest

import numpy as np
import pytest
import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

Expand All @@ -14,27 +28,43 @@
)

from ...testing_utils import (
assert_tensors_close,
backend_empty_cache,
enable_full_determinism,
require_torch_accelerator,
slow,
torch_device,
)
from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import IPAdapterTesterMixin, PipelineLatentTesterMixin, PipelineTesterMixin
from ..testing_utils import (
BasePipelineTesterConfig,
IPAdapterTesterMixin,
LoraMemoryTesterMixin,
LoraTesterMixin,
MemoryTesterMixin,
PipelineTesterMixin,
UNetLoraTesterMixin,
)


enable_full_determinism()


class LatentConsistencyModelPipelineFastTests(
IPAdapterTesterMixin, PipelineLatentTesterMixin, PipelineTesterMixin, unittest.TestCase
):
class LatentConsistencyModelPipelineTesterConfig(BasePipelineTesterConfig):
pipeline_class = LatentConsistencyModelPipeline
params = TEXT_TO_IMAGE_PARAMS - {"negative_prompt", "negative_prompt_embeds"}
batch_params = TEXT_TO_IMAGE_BATCH_PARAMS - {"negative_prompt"}
image_params = TEXT_TO_IMAGE_IMAGE_PARAMS
image_latents_params = TEXT_TO_IMAGE_IMAGE_PARAMS
# The canonical text-to-image sets minus `negative_prompt` / `negative_prompt_embeds`: LCM is
# guidance-distilled and `__call__` takes no negative prompt.
required_input_params_in_call_signature = frozenset(
[
"prompt",
"height",
"width",
"guidance_scale",
"prompt_embeds",
"cross_attention_kwargs",
]
)
batch_input_params = frozenset(["prompt"])
output_shape = (3, 64, 64)

def get_dummy_components(self):
torch.manual_seed(0)
Expand Down Expand Up @@ -82,7 +112,7 @@ def get_dummy_components(self):
text_encoder = CLIPTextModel(text_encoder_config)
tokenizer = CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip")

components = {
return {
"unet": unet,
"scheduler": scheduler,
"vae": vae,
Expand All @@ -93,155 +123,144 @@ def get_dummy_components(self):
"image_encoder": None,
"requires_safety_checker": False,
}
return components

def get_dummy_inputs(self, device, seed=0):
if str(device).startswith("mps"):
generator = torch.manual_seed(seed)
else:
generator = torch.Generator(device=device).manual_seed(seed)
inputs = {

def get_dummy_inputs(self):
return {
"prompt": "A painting of a squirrel eating a burger",
"generator": generator,
"generator": self.get_generator(0),
"num_inference_steps": 2,
"guidance_scale": 6.0,
"output_type": "np",
# Request torch outputs so tests compare torch tensors directly (see `BasePipelineTesterConfig`).
# Note `"pt"` images are `(batch, channels, height, width)`, unlike `"np"` (`(batch, h, w, c)`).
"output_type": "pt",
}
return inputs

def test_ip_adapter(self):
expected_pipe_slice = None
if torch_device == "cpu":
expected_pipe_slice = np.array([0.1405, 0.5002, 0.5213, 0.1223, 0.3856, 0.4165, 0.5382, 0.3622, 0.3693])
return super().test_ip_adapter(expected_pipe_slice=expected_pipe_slice)

class TestLatentConsistencyModelPipeline(LatentConsistencyModelPipelineTesterConfig, PipelineTesterMixin):
def test_lcm_onestep(self):
device = "cpu" # ensure determinism for the device-dependent torch.Generator

components = self.get_dummy_components()
pipe = LatentConsistencyModelPipeline(**components)
pipe = pipe.to(torch_device)
pipe.set_progress_bar_config(disable=None)
# Run on CPU: the expected slice below is CPU-specific.
pipe = self.get_pipeline()

inputs = self.get_dummy_inputs(device)
inputs = self.get_dummy_inputs()
inputs["num_inference_steps"] = 1
output = pipe(**inputs)
image = output.images
assert image.shape == (1, 64, 64, 3)
image = pipe(**inputs).images
assert image.shape == (1, *self.output_shape)

image_slice = image[0, -3:, -3:, -1]
expected_slice = np.array([0.1444, 0.5229, 0.5344, 0.1384, 0.3999, 0.4320, 0.5345, 0.3559, 0.3685])
assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-3
# fmt: off
expected_slice = torch.tensor([0.1444, 0.5229, 0.5344, 0.1384, 0.3999, 0.4320, 0.5345, 0.3559, 0.3685])
# fmt: on
assert_tensors_close(image[0, -1, -3:, -3:].flatten(), expected_slice, atol=1e-3)

def test_lcm_multistep(self):
device = "cpu" # ensure determinism for the device-dependent torch.Generator

components = self.get_dummy_components()
pipe = LatentConsistencyModelPipeline(**components)
pipe = pipe.to(torch_device)
pipe.set_progress_bar_config(disable=None)
# Run on CPU: the expected slice below is CPU-specific.
pipe = self.get_pipeline()

inputs = self.get_dummy_inputs(device)
output = pipe(**inputs)
image = output.images
assert image.shape == (1, 64, 64, 3)
image = pipe(**self.get_dummy_inputs()).images
assert image.shape == (1, *self.output_shape)

image_slice = image[0, -3:, -3:, -1]
expected_slice = np.array([0.1405, 0.5002, 0.5213, 0.1223, 0.3856, 0.4165, 0.5382, 0.3622, 0.3693])
assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-3
# fmt: off
expected_slice = torch.tensor([0.1405, 0.5002, 0.5213, 0.1223, 0.3856, 0.4165, 0.5382, 0.3622, 0.3693])
# fmt: on
assert_tensors_close(image[0, -1, -3:, -3:].flatten(), expected_slice, atol=1e-3)

def test_lcm_custom_timesteps(self):
device = "cpu" # ensure determinism for the device-dependent torch.Generator
# Run on CPU: the expected slice below is CPU-specific.
pipe = self.get_pipeline()

components = self.get_dummy_components()
pipe = LatentConsistencyModelPipeline(**components)
pipe = pipe.to(torch_device)
pipe.set_progress_bar_config(disable=None)

inputs = self.get_dummy_inputs(device)
inputs = self.get_dummy_inputs()
del inputs["num_inference_steps"]
inputs["timesteps"] = [999, 499]
output = pipe(**inputs)
image = output.images
assert image.shape == (1, 64, 64, 3)
image = pipe(**inputs).images
assert image.shape == (1, *self.output_shape)

image_slice = image[0, -3:, -3:, -1]
expected_slice = np.array([0.1405, 0.5002, 0.5213, 0.1223, 0.3856, 0.4165, 0.5382, 0.3622, 0.3693])
assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-3
# Custom timesteps matching the default 2-step schedule reproduce `test_lcm_multistep`'s output.
# fmt: off
expected_slice = torch.tensor([0.1405, 0.5002, 0.5213, 0.1223, 0.3856, 0.4165, 0.5382, 0.3622, 0.3693])
# fmt: on
assert_tensors_close(image[0, -1, -3:, -3:].flatten(), expected_slice, atol=1e-3)

def test_inference_batch_single_identical(self):
super().test_inference_batch_single_identical(expected_max_diff=5e-4)
def test_inference_batch_single_identical(self, batch_size=3, expected_max_diff=5e-4):
super().test_inference_batch_single_identical(batch_size=batch_size, expected_max_diff=expected_max_diff)

# skip because lcm pipeline apply cfg differently
@pytest.mark.skip("LCM applies classifier-free guidance differently, so the shared CFG callback test cannot run.")
def test_callback_cfg(self):
pass

# override default test because the final latent variable is "denoised" instead of "latents"
def test_callback_inputs(self):
sig = inspect.signature(self.pipeline_class.__call__)
# Overridden because the final latent variable is `denoised` rather than `latents`.
pipe = self.get_pipeline().to(torch_device)

if not ("callback_on_step_end_tensor_inputs" in sig.parameters and "callback_on_step_end" in sig.parameters):
return

components = self.get_dummy_components()
pipe = self.pipeline_class(**components)
pipe = pipe.to(torch_device)
pipe.set_progress_bar_config(disable=None)

self.assertTrue(
hasattr(pipe, "_callback_tensor_inputs"),
f" {self.pipeline_class} should have `_callback_tensor_inputs` that defines a list of tensor variables its callback function can use as inputs",
assert hasattr(pipe, "_callback_tensor_inputs"), (
f"{self.pipeline_class} should have `_callback_tensor_inputs` that defines a list of tensor variables "
"its callback function can use as inputs"
)

def callback_inputs_test(pipe, i, t, callback_kwargs):
missing_callback_inputs = set()
for v in pipe._callback_tensor_inputs:
if v not in callback_kwargs:
missing_callback_inputs.add(v)
self.assertTrue(
len(missing_callback_inputs) == 0, f"Missing callback tensor inputs: {missing_callback_inputs}"
)
last_i = pipe.num_timesteps - 1
if i == last_i:
missing_callback_inputs = {v for v in pipe._callback_tensor_inputs if v not in callback_kwargs}
assert len(missing_callback_inputs) == 0, f"Missing callback tensor inputs: {missing_callback_inputs}"
if i == pipe.num_timesteps - 1:
callback_kwargs["denoised"] = torch.zeros_like(callback_kwargs["denoised"])
return callback_kwargs

inputs = self.get_dummy_inputs(torch_device)
inputs = self.get_dummy_inputs()
inputs["callback_on_step_end"] = callback_inputs_test
inputs["callback_on_step_end_tensor_inputs"] = pipe._callback_tensor_inputs
inputs["output_type"] = "latent"

output = pipe(**inputs)[0]
assert output.abs().sum() == 0

def test_encode_prompt_works_in_isolation(self):
def test_encode_prompt_works_in_isolation(self, extra_required_param_value_dict=None, atol=1e-4, rtol=1e-4):
# `encode_prompt` requires `device` and `do_classifier_free_guidance`, neither of which `__call__`
# exposes with a default for the shared test to pick up.
extra_required_param_value_dict = {
"device": torch.device(torch_device).type,
"do_classifier_free_guidance": self.get_dummy_inputs(device=torch_device).get("guidance_scale", 1.0) > 1.0,
"do_classifier_free_guidance": self.get_dummy_inputs().get("guidance_scale", 1.0) > 1.0,
}
return super().test_encode_prompt_works_in_isolation(extra_required_param_value_dict)
super().test_encode_prompt_works_in_isolation(extra_required_param_value_dict, atol=atol, rtol=rtol)


class TestLatentConsistencyModelPipelineMemory(LatentConsistencyModelPipelineTesterConfig, MemoryTesterMixin):
"""Memory optimization tests (CPU offload, group offload, layerwise casting) for the LCM pipeline."""


class TestLatentConsistencyModelPipelineIPAdapter(LatentConsistencyModelPipelineTesterConfig, IPAdapterTesterMixin):
"""IP-Adapter tests for the LCM pipeline."""


class TestLatentConsistencyModelPipelineLoRA(
LatentConsistencyModelPipelineTesterConfig, LoraTesterMixin, UNetLoraTesterMixin
):
"""LoRA tests for the LCM pipeline."""


class TestLatentConsistencyModelPipelineLoRAMemory(LatentConsistencyModelPipelineTesterConfig, LoraMemoryTesterMixin):
"""LoRA x memory-optimization tests (group offload, CPU offload) for the LCM pipeline."""


@slow
@require_torch_accelerator
class LatentConsistencyModelPipelineSlowTests(unittest.TestCase):
def setUp(self):
class TestLatentConsistencyModelPipelineIntegration:
@pytest.fixture(autouse=True)
def cleanup(self):
gc.collect()
backend_empty_cache(torch_device)
yield
gc.collect()
backend_empty_cache(torch_device)

def get_inputs(self, device, generator_device="cpu", dtype=torch.float32, seed=0):
generator = torch.Generator(device=generator_device).manual_seed(seed)
latents = np.random.RandomState(seed).standard_normal((1, 4, 64, 64))
latents = torch.from_numpy(latents).to(device=device, dtype=dtype)
inputs = {
return {
"prompt": "a photograph of an astronaut riding a horse",
"latents": latents,
"generator": generator,
"num_inference_steps": 3,
"guidance_scale": 7.5,
"output_type": "np",
}
return inputs

def test_lcm_onestep(self):
pipe = LatentConsistencyModelPipeline.from_pretrained("SimianLuo/LCM_Dreamshaper_v7", safety_checker=None)
Expand All @@ -264,8 +283,7 @@ def test_lcm_multistep(self):
pipe = pipe.to(torch_device)
pipe.set_progress_bar_config(disable=None)

inputs = self.get_inputs(torch_device)
image = pipe(**inputs).images
image = pipe(**self.get_inputs(torch_device)).images
assert image.shape == (1, 512, 512, 3)

image_slice = image[0, -3:, -3:, -1].flatten()
Expand Down
Loading
Loading