AI/diffusers
1
0
Fork 0
mirror of https://github.com/huggingface/diffusers.git synced 2026-10-24 20:59:27 +03:00
Code Activity

[tests] refactor l-series pipeline tests (#14637)

Browse Source 
* refactor l-series pipeline tests

* define sets explicitly.
This commit is contained in:
Sayak Paul
2026-08-31 22:34:40 +05:30
committed by GitHub
parent 45ff6a6c18
commit 45b32f4321
22 changed files with 1334 additions and 2329 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
src/diffusers/pipelines/latent_diffusion/pipeline_latent_diffusion.py
View File

@@ -113,7 +113,8 @@ class LDMTextToImagePipeline(DiffusionPipeline):
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.
@@ -220,9 +221,12 @@ class LDMTextToImagePipeline(DiffusionPipeline):
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)
if not return_dict:
return (image,)

11
src/diffusers/pipelines/latent_diffusion/pipeline_latent_diffusion_superresolution.py
View File

@@ -97,7 +97,8 @@ class LDMSuperResolutionPipeline(DiffusionPipeline):
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.
@@ -185,10 +186,12 @@ class LDMSuperResolutionPipeline(DiffusionPipeline):
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,)

220
tests/pipelines/latent_consistency_models/test_latent_consistency_models.py
View File

@@ -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
@@ -14,27 +28,43 @@ from diffusers import (
)
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)
@@ -82,7 +112,7 @@ class LatentConsistencyModelPipelineFastTests(
text_encoder = CLIPTextModel(text_encoder_config)
tokenizer = CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip")
components = {
return {
"unet": unet,
"scheduler": scheduler,
"vae": vae,
@@ -93,120 +123,86 @@ class LatentConsistencyModelPipelineFastTests(
"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
# 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()
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
# 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)
image = pipe(**self.get_dummy_inputs()).images
assert image.shape == (1, *self.output_shape)
inputs = self.get_dummy_inputs(device)
output = pipe(**inputs)
image = output.images
assert image.shape == (1, 64, 64, 3)
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"
@@ -214,18 +210,42 @@ class LatentConsistencyModelPipelineFastTests(
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)
@@ -233,7 +253,7 @@ class LatentConsistencyModelPipelineSlowTests(unittest.TestCase):
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,
@@ -241,7 +261,6 @@ class LatentConsistencyModelPipelineSlowTests(unittest.TestCase):
"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)
@@ -264,8 +283,7 @@ class LatentConsistencyModelPipelineSlowTests(unittest.TestCase):
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()

233
tests/pipelines/latent_consistency_models/test_latent_consistency_models_img2img.py
View File

@@ -1,9 +1,23 @@
# 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 random
import unittest
import numpy as np
import pytest
import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer
@@ -15,6 +29,7 @@ from diffusers import (
)
from ...testing_utils import (
assert_tensors_close,
backend_empty_cache,
enable_full_determinism,
floats_tensor,
@@ -23,26 +38,38 @@ from ...testing_utils import (
slow,
torch_device,
)
from ..pipeline_params import (
IMAGE_TO_IMAGE_IMAGE_PARAMS,
TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS,
TEXT_GUIDED_IMAGE_VARIATION_PARAMS,
from ..testing_utils import (
BasePipelineTesterConfig,
IPAdapterTesterMixin,
LoraMemoryTesterMixin,
LoraTesterMixin,
MemoryTesterMixin,
PipelineTesterMixin,
UNetLoraTesterMixin,
)
from ..test_pipelines_common import IPAdapterTesterMixin, PipelineLatentTesterMixin, PipelineTesterMixin
enable_full_determinism()
class LatentConsistencyModelImg2ImgPipelineFastTests(
IPAdapterTesterMixin, PipelineLatentTesterMixin, PipelineTesterMixin, unittest.TestCase
):
class LatentConsistencyModelImg2ImgPipelineTesterConfig(BasePipelineTesterConfig):
pipeline_class = LatentConsistencyModelImg2ImgPipeline
params = TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {"height", "width", "negative_prompt", "negative_prompt_embeds"}
required_optional_params = PipelineTesterMixin.required_optional_params - {"latents", "negative_prompt"}
batch_params = TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
image_params = IMAGE_TO_IMAGE_IMAGE_PARAMS
image_latents_params = IMAGE_TO_IMAGE_IMAGE_PARAMS
# The canonical text-guided image-variation sets minus `height` / `width`, which this pipeline derives from
# `image`, and minus `negative_prompt` / `negative_prompt_embeds`, which guidance-distilled LCM does not take.
required_input_params_in_call_signature = frozenset(
[
"prompt",
"image",
"guidance_scale",
"prompt_embeds",
]
)
batch_input_params = frozenset(["prompt", "image"])
output_shape = (3, 32, 32)
# `__call__` starts from the supplied image, so it takes no `latents`.
optional_input_params = frozenset(
["num_inference_steps", "num_images_per_prompt", "generator", "output_type", "return_dict"]
)
def get_dummy_components(self):
torch.manual_seed(0)
@@ -90,7 +117,7 @@ class LatentConsistencyModelImg2ImgPipelineFastTests(
text_encoder = CLIPTextModel(text_encoder_config)
tokenizer = CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip")
components = {
return {
"unet": unet,
"scheduler": scheduler,
"vae": vae,
@@ -101,119 +128,87 @@ class LatentConsistencyModelImg2ImgPipelineFastTests(
"image_encoder": None,
"requires_safety_checker": False,
}
return components
def get_dummy_inputs(self, device, seed=0):
image = floats_tensor((1, 3, 32, 32), rng=random.Random(seed)).to(device)
def get_dummy_inputs(self):
image = floats_tensor((1, 3, 32, 32), rng=random.Random(0))
image = image / 2 + 0.5
if str(device).startswith("mps"):
generator = torch.manual_seed(seed)
else:
generator = torch.Generator(device=device).manual_seed(seed)
inputs = {
return {
"prompt": "A painting of a squirrel eating a burger",
"image": image,
"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.3952, 0.3674, 0.2839, 0.5461, 0.5532, 0.3737, 0.4652, 0.4986, 0.4422])
return super().test_ip_adapter(expected_pipe_slice=expected_pipe_slice)
class TestLatentConsistencyModelImg2ImgPipeline(
LatentConsistencyModelImg2ImgPipelineTesterConfig, PipelineTesterMixin
):
def test_lcm_onestep(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 = self.pipeline_class(**components)
pipe = pipe.to(torch_device)
pipe.set_progress_bar_config(disable=None)
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, 32, 32, 3)
image = pipe(**inputs).images
assert image.shape == (1, *self.output_shape)
image_slice = image[0, -3:, -3:, -1]
expected_slice = np.array([0.4317, 0.3653, 0.2190, 0.7136, 0.6321, 0.3634, 0.5846, 0.6095, 0.4969])
assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-3
# fmt: off
expected_slice = torch.tensor([0.4317, 0.3653, 0.2190, 0.7136, 0.6321, 0.3634, 0.5846, 0.6095, 0.4969])
# 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
# Run on CPU: the expected slice below is CPU-specific.
pipe = self.get_pipeline()
components = self.get_dummy_components()
pipe = self.pipeline_class(**components)
pipe = pipe.to(torch_device)
pipe.set_progress_bar_config(disable=None)
image = pipe(**self.get_dummy_inputs()).images
assert image.shape == (1, *self.output_shape)
inputs = self.get_dummy_inputs(device)
output = pipe(**inputs)
image = output.images
assert image.shape == (1, 32, 32, 3)
image_slice = image[0, -3:, -3:, -1]
expected_slice = np.array([0.4083, 0.3668, 0.2467, 0.6268, 0.5976, 0.3750, 0.5071, 0.5439, 0.4677])
assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-3
# fmt: off
expected_slice = torch.tensor([0.4083, 0.3668, 0.2467, 0.6268, 0.5976, 0.3750, 0.5071, 0.5439, 0.4677])
# 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 = self.pipeline_class(**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, 32, 32, 3)
image = pipe(**inputs).images
assert image.shape == (1, *self.output_shape)
image_slice = image[0, -3:, -3:, -1]
expected_slice = np.array([0.3985, 0.3444, 0.2534, 0.6969, 0.6167, 0.3622, 0.5754, 0.5850, 0.4959])
assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-3
# fmt: off
expected_slice = torch.tensor([0.3985, 0.3444, 0.2534, 0.6969, 0.6167, 0.3622, 0.5754, 0.5850, 0.4959])
# 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)
# 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"
@@ -221,18 +216,48 @@ class LatentConsistencyModelImg2ImgPipelineFastTests(
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 TestLatentConsistencyModelImg2ImgPipelineMemory(
LatentConsistencyModelImg2ImgPipelineTesterConfig, MemoryTesterMixin
):
"""Memory optimization tests (CPU offload, group offload, layerwise casting) for the LCM img2img pipeline."""
class TestLatentConsistencyModelImg2ImgPipelineIPAdapter(
LatentConsistencyModelImg2ImgPipelineTesterConfig, IPAdapterTesterMixin
):
"""IP-Adapter tests for the LCM img2img pipeline."""
class TestLatentConsistencyModelImg2ImgPipelineLoRA(
LatentConsistencyModelImg2ImgPipelineTesterConfig, LoraTesterMixin, UNetLoraTesterMixin
):
"""LoRA tests for the LCM img2img pipeline."""
class TestLatentConsistencyModelImg2ImgPipelineLoRAMemory(
LatentConsistencyModelImg2ImgPipelineTesterConfig, LoraMemoryTesterMixin
):
"""LoRA x memory-optimization tests (group offload, CPU offload) for the LCM img2img pipeline."""
@slow
@require_torch_accelerator
class LatentConsistencyModelImg2ImgPipelineSlowTests(unittest.TestCase):
def setUp(self):
class TestLatentConsistencyModelImg2ImgPipelineIntegration:
@pytest.fixture(autouse=True)
def cleanup(self):
gc.collect()
backend_empty_cache(torch_device)
yield
gc.collect()
backend_empty_cache(torch_device)
@@ -246,7 +271,7 @@ class LatentConsistencyModelImg2ImgPipelineSlowTests(unittest.TestCase):
)
init_image = init_image.resize((512, 512))
inputs = {
return {
"prompt": "a photograph of an astronaut riding a horse",
"latents": latents,
"generator": generator,
@@ -255,7 +280,6 @@ class LatentConsistencyModelImg2ImgPipelineSlowTests(unittest.TestCase):
"output_type": "np",
"image": init_image,
}
return inputs
def test_lcm_onestep(self):
pipe = LatentConsistencyModelImg2ImgPipeline.from_pretrained(
@@ -282,8 +306,7 @@ class LatentConsistencyModelImg2ImgPipelineSlowTests(unittest.TestCase):
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()

124
tests/pipelines/latent_diffusion/test_latent_diffusion.py
View File

@@ -14,15 +14,16 @@
# limitations under the License.
import gc
import unittest
import numpy as np
import pytest
import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer
from diffusers import AutoencoderKL, DDIMScheduler, LDMTextToImagePipeline, UNet2DConditionModel
from ...testing_utils import (
assert_tensors_close,
backend_empty_cache,
enable_full_determinism,
load_numpy,
@@ -31,26 +32,25 @@ from ...testing_utils import (
torch_device,
)
from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import PipelineTesterMixin
from ..testing_utils import BasePipelineTesterConfig, MemoryTesterMixin, PipelineTesterMixin
enable_full_determinism()
class LDMTextToImagePipelineFastTests(PipelineTesterMixin, unittest.TestCase):
class LDMTextToImagePipelineTesterConfig(BasePipelineTesterConfig):
pipeline_class = LDMTextToImagePipeline
params = TEXT_TO_IMAGE_PARAMS - {
# This pipeline predates prompt embeddings and negative prompting; `__call__` only takes a raw `prompt`.
required_input_params_in_call_signature = TEXT_TO_IMAGE_PARAMS - {
"negative_prompt",
"negative_prompt_embeds",
"cross_attention_kwargs",
"prompt_embeds",
}
required_optional_params = PipelineTesterMixin.required_optional_params - {
"num_images_per_prompt",
"callback",
"callback_steps",
}
batch_params = TEXT_TO_IMAGE_BATCH_PARAMS
batch_input_params = TEXT_TO_IMAGE_BATCH_PARAMS
output_shape = (3, 16, 16)
# `__call__` has no `num_images_per_prompt`.
optional_input_params = frozenset(["num_inference_steps", "generator", "latents", "output_type", "return_dict"])
def get_dummy_components(self):
torch.manual_seed(0)
@@ -95,124 +95,86 @@ class LDMTextToImagePipelineFastTests(PipelineTesterMixin, unittest.TestCase):
text_encoder = CLIPTextModel(text_encoder_config)
tokenizer = CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip")
components = {
return {
"unet": unet,
"scheduler": scheduler,
"vqvae": vae,
"bert": text_encoder,
"tokenizer": tokenizer,
}
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
class TestLDMTextToImagePipeline(LDMTextToImagePipelineTesterConfig, PipelineTesterMixin):
def test_inference_text2img(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 = LDMTextToImagePipeline(**components)
pipe.to(device)
pipe.set_progress_bar_config(disable=None)
image = pipe(**self.get_dummy_inputs()).images
assert image.shape == (1, *self.output_shape)
inputs = self.get_dummy_inputs(device)
image = pipe(**inputs).images
image_slice = image[0, -3:, -3:, -1]
# fmt: off
expected_slice = torch.tensor([0.6511, 0.5873, 0.5183, 0.5046, 0.6663, 0.3908, 0.5729, 0.5979, 0.5058])
# fmt: on
assert_tensors_close(image[0, -1, -3:, -3:].flatten(), expected_slice, atol=1e-3)
assert image.shape == (1, 16, 16, 3)
expected_slice = np.array([0.6511, 0.5873, 0.5183, 0.5046, 0.6663, 0.3908, 0.5729, 0.5979, 0.5058])
assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-3
class TestLDMTextToImagePipelineMemory(LDMTextToImagePipelineTesterConfig, MemoryTesterMixin):
"""Memory optimization tests (CPU offload, group offload, layerwise casting) for the LDM text2img pipeline."""
@nightly
@require_torch_accelerator
class LDMTextToImagePipelineSlowTests(unittest.TestCase):
def setUp(self):
super().setUp()
class TestLDMTextToImagePipelineIntegration:
@pytest.fixture(autouse=True)
def cleanup(self):
gc.collect()
backend_empty_cache(torch_device)
yield
gc.collect()
backend_empty_cache(torch_device)
def tearDown(self):
super().tearDown()
gc.collect()
backend_empty_cache(torch_device)
def get_inputs(self, device, dtype=torch.float32, seed=0):
def get_inputs(self, device, dtype=torch.float32, seed=0, num_inference_steps=3):
generator = torch.manual_seed(seed)
latents = np.random.RandomState(seed).standard_normal((1, 4, 32, 32))
latents = torch.from_numpy(latents).to(device=device, dtype=dtype)
inputs = {
return {
"prompt": "A painting of a squirrel eating a burger",
"latents": latents,
"generator": generator,
"num_inference_steps": 3,
"num_inference_steps": num_inference_steps,
"guidance_scale": 6.0,
"output_type": "np",
}
return inputs
def test_ldm_default_ddim(self):
pipe = LDMTextToImagePipeline.from_pretrained("CompVis/ldm-text2im-large-256").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
image_slice = image[0, -3:, -3:, -1].flatten()
assert image.shape == (1, 256, 256, 3)
expected_slice = np.array([0.51825, 0.52850, 0.52543, 0.54258, 0.52304, 0.52569, 0.54363, 0.55276, 0.56878])
max_diff = np.abs(expected_slice - image_slice).max()
assert max_diff < 1e-3
assert np.abs(expected_slice - image_slice).max() < 1e-3
@nightly
@require_torch_accelerator
class LDMTextToImagePipelineNightlyTests(unittest.TestCase):
def setUp(self):
super().setUp()
gc.collect()
backend_empty_cache(torch_device)
def tearDown(self):
super().tearDown()
gc.collect()
backend_empty_cache(torch_device)
def get_inputs(self, device, dtype=torch.float32, seed=0):
generator = torch.manual_seed(seed)
latents = np.random.RandomState(seed).standard_normal((1, 4, 32, 32))
latents = torch.from_numpy(latents).to(device=device, dtype=dtype)
inputs = {
"prompt": "A painting of a squirrel eating a burger",
"latents": latents,
"generator": generator,
"num_inference_steps": 50,
"guidance_scale": 6.0,
"output_type": "np",
}
return inputs
def test_ldm_default_ddim(self):
def test_ldm_default_ddim_full_schedule(self):
pipe = LDMTextToImagePipeline.from_pretrained("CompVis/ldm-text2im-large-256").to(torch_device)
pipe.set_progress_bar_config(disable=None)
inputs = self.get_inputs(torch_device)
image = pipe(**inputs).images[0]
image = pipe(**self.get_inputs(torch_device, num_inference_steps=50)).images[0]
expected_image = load_numpy(
"https://huggingface.co/datasets/diffusers/test-arrays/resolve/main/ldm_text2img/ldm_large_256_ddim.npy"
)
max_diff = np.abs(expected_image - image).max()
assert max_diff < 1e-3
assert np.abs(expected_image - image).max() < 1e-3

110
tests/pipelines/latent_diffusion/test_latent_diffusion_superresolution.py
View File

@@ -14,42 +14,44 @@
# limitations under the License.
import random
import unittest
import numpy as np
import pytest
import torch
from diffusers import DDIMScheduler, LDMSuperResolutionPipeline, UNet2DModel, VQModel
from diffusers.utils import PIL_INTERPOLATION
from ...testing_utils import (
assert_tensors_close,
enable_full_determinism,
floats_tensor,
load_image,
nightly,
require_accelerator,
require_torch,
torch_device,
)
from ..testing_utils import BasePipelineTesterConfig, MemoryTesterMixin, PipelineTesterMixin
enable_full_determinism()
class LDMSuperResolutionPipelineFastTests(unittest.TestCase):
@property
def dummy_image(self):
batch_size = 1
num_channels = 3
sizes = (32, 32)
class LDMSuperResolutionPipelineTesterConfig(BasePipelineTesterConfig):
pipeline_class = LDMSuperResolutionPipeline
required_input_params_in_call_signature = frozenset(["image", "batch_size"])
# `__call__` takes exactly one `image` — a `PIL.Image` or a single tensor whose leading dim *is* the batch —
# and rejects the list of images the shared batching helpers build, so there is nothing for them to batch.
# `test_batched_image_input` below covers batching through the API the pipeline actually offers.
batch_input_params = frozenset()
output_shape = (3, 64, 64)
# An unconditional upscaler: no prompt, so no `num_images_per_prompt`, and the noise is always sampled
# internally rather than passed in as `latents`.
optional_input_params = frozenset(["num_inference_steps", "generator", "output_type", "return_dict"])
image = floats_tensor((batch_size, num_channels) + sizes, rng=random.Random(0)).to(torch_device)
return image
@property
def dummy_uncond_unet(self):
def get_dummy_components(self):
torch.manual_seed(0)
model = UNet2DModel(
unet = UNet2DModel(
block_out_channels=(32, 64),
layers_per_block=2,
sample_size=32,
@@ -58,12 +60,8 @@ class LDMSuperResolutionPipelineFastTests(unittest.TestCase):
down_block_types=("DownBlock2D", "AttnDownBlock2D"),
up_block_types=("AttnUpBlock2D", "UpBlock2D"),
)
return model
@property
def dummy_vq_model(self):
torch.manual_seed(0)
model = VQModel(
vqvae = VQModel(
block_out_channels=[32, 64],
in_channels=3,
out_channels=3,
@@ -71,54 +69,61 @@ class LDMSuperResolutionPipelineFastTests(unittest.TestCase):
up_block_types=["UpDecoderBlock2D", "UpDecoderBlock2D"],
latent_channels=3,
)
return model
scheduler = DDIMScheduler()
return {"unet": unet, "vqvae": vqvae, "scheduler": scheduler}
def get_dummy_inputs(self):
return {
"image": floats_tensor((1, 3, 32, 32), rng=random.Random(0)),
"generator": self.get_generator(0),
"num_inference_steps": 2,
# 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",
}
class TestLDMSuperResolutionPipeline(LDMSuperResolutionPipelineTesterConfig, PipelineTesterMixin):
def test_inference_superresolution(self):
device = "cpu"
unet = self.dummy_uncond_unet
scheduler = DDIMScheduler()
vqvae = self.dummy_vq_model
# Run on CPU: the expected slice below is CPU-specific.
pipe = self.get_pipeline()
ldm = LDMSuperResolutionPipeline(unet=unet, vqvae=vqvae, scheduler=scheduler)
ldm.to(device)
ldm.set_progress_bar_config(disable=None)
image = pipe(**self.get_dummy_inputs()).images
assert image.shape == (1, *self.output_shape)
init_image = self.dummy_image.to(device)
# fmt: off
expected_slice = torch.tensor([0.8678, 0.8245, 0.6381, 0.6830, 0.4385, 0.5599, 0.4641, 0.6201, 0.5150])
# fmt: on
assert_tensors_close(image[0, -1, -3:, -3:].flatten(), expected_slice, atol=1e-2)
generator = torch.Generator(device=device).manual_seed(0)
image = ldm(image=init_image, generator=generator, num_inference_steps=2, output_type="np").images
def test_batched_image_input(self):
# The `batch_size` argument is ignored for tensor input (`__call__` derives it from `image.shape[0]`), so
# batching means stacking frames into `image` itself.
pipe = self.get_pipeline().to(torch_device)
image_slice = image[0, -3:, -3:, -1]
inputs = self.get_dummy_inputs()
inputs["image"] = inputs["image"].repeat(3, 1, 1, 1)
images = pipe(**inputs).images
assert image.shape == (1, 64, 64, 3)
expected_slice = np.array([0.8678, 0.8245, 0.6381, 0.6830, 0.4385, 0.5599, 0.4641, 0.6201, 0.5150])
assert images.shape == (3, *self.output_shape)
assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2
@pytest.mark.skip("`__call__` rejects a list of images; see `test_batched_image_input`.")
def test_inference_batch_consistent(self):
pass
@require_accelerator
def test_inference_superresolution_fp16(self):
unet = self.dummy_uncond_unet
scheduler = DDIMScheduler()
vqvae = self.dummy_vq_model
@pytest.mark.skip("`__call__` rejects a list of images; see `test_batched_image_input`.")
def test_inference_batch_single_identical(self):
pass
# put models in fp16
unet = unet.half()
vqvae = vqvae.half()
ldm = LDMSuperResolutionPipeline(unet=unet, vqvae=vqvae, scheduler=scheduler)
ldm.to(torch_device)
ldm.set_progress_bar_config(disable=None)
init_image = self.dummy_image.to(torch_device)
image = ldm(init_image, num_inference_steps=2, output_type="np").images
assert image.shape == (1, 64, 64, 3)
class TestLDMSuperResolutionPipelineMemory(LDMSuperResolutionPipelineTesterConfig, MemoryTesterMixin):
"""Memory optimization tests (CPU offload, group offload, layerwise casting) for the LDM upscaler pipeline."""
@nightly
@require_torch
class LDMSuperResolutionPipelineIntegrationTests(unittest.TestCase):
class TestLDMSuperResolutionPipelineIntegration:
def test_inference_superresolution(self):
init_image = load_image(
"https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
@@ -136,5 +141,4 @@ class LDMSuperResolutionPipelineIntegrationTests(unittest.TestCase):
assert image.shape == (1, 256, 256, 3)
expected_slice = np.array([0.7644, 0.7679, 0.7642, 0.7633, 0.7666, 0.7560, 0.7425, 0.7257, 0.6907])
assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2

294
tests/pipelines/latte/test_latte.py
View File

@@ -14,25 +14,20 @@
# limitations under the License.
import gc
import inspect
import tempfile
import unittest
import numpy as np
import pytest
import torch
from transformers import AutoConfig, AutoTokenizer, T5EncoderModel
from diffusers import (
AutoencoderKL,
DDIMScheduler,
FasterCacheConfig,
LattePipeline,
LatteTransformer3DModel,
PyramidAttentionBroadcastConfig,
)
from diffusers.utils.import_utils import is_xformers_available
from ...testing_utils import (
assert_tensors_close,
backend_empty_cache,
enable_full_determinism,
numpy_cosine_similarity_distance,
@@ -40,51 +35,25 @@ from ...testing_utils import (
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 (
from ..testing_utils import (
BasePipelineTesterConfig,
FasterCacheTesterMixin,
MemoryTesterMixin,
PipelineTesterMixin,
PyramidAttentionBroadcastTesterMixin,
to_np,
)
enable_full_determinism()
class LattePipelineFastTests(
PipelineTesterMixin, PyramidAttentionBroadcastTesterMixin, FasterCacheTesterMixin, unittest.TestCase
):
class LattePipelineTesterConfig(BasePipelineTesterConfig):
pipeline_class = LattePipeline
params = TEXT_TO_IMAGE_PARAMS - {"cross_attention_kwargs"}
batch_params = TEXT_TO_IMAGE_BATCH_PARAMS
image_params = TEXT_TO_IMAGE_IMAGE_PARAMS
image_latents_params = TEXT_TO_IMAGE_IMAGE_PARAMS
required_optional_params = PipelineTesterMixin.required_optional_params
test_layerwise_casting = True
test_group_offloading = True
pab_config = PyramidAttentionBroadcastConfig(
spatial_attention_block_skip_range=2,
temporal_attention_block_skip_range=2,
cross_attention_block_skip_range=2,
spatial_attention_timestep_skip_range=(100, 700),
temporal_attention_timestep_skip_range=(100, 800),
cross_attention_timestep_skip_range=(100, 800),
spatial_attention_block_identifiers=["transformer_blocks"],
temporal_attention_block_identifiers=["temporal_transformer_blocks"],
cross_attention_block_identifiers=["transformer_blocks"],
)
faster_cache_config = FasterCacheConfig(
spatial_attention_block_skip_range=2,
temporal_attention_block_skip_range=2,
spatial_attention_timestep_skip_range=(-1, 901),
temporal_attention_timestep_skip_range=(-1, 901),
unconditional_batch_skip_range=2,
attention_weight_callback=lambda _: 0.5,
required_input_params_in_call_signature = frozenset(
["prompt", "height", "width", "guidance_scale", "negative_prompt", "prompt_embeds", "negative_prompt_embeds"]
)
batch_input_params = frozenset(["prompt", "negative_prompt"])
output_shape = (1, 3, 8, 8)
def get_dummy_components(self, num_layers: int = 1):
torch.manual_seed(0)
@@ -111,211 +80,127 @@ class LattePipelineFastTests(
scheduler = DDIMScheduler()
config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5")
text_encoder = T5EncoderModel(config)
tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5")
components = {
"transformer": transformer.eval(),
"vae": vae.eval(),
return {
"transformer": transformer,
"vae": vae,
"scheduler": scheduler,
"text_encoder": text_encoder.eval(),
"text_encoder": text_encoder,
"tokenizer": tokenizer,
}
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",
"negative_prompt": "low quality",
"generator": generator,
"generator": self.get_generator(0),
"num_inference_steps": 2,
"guidance_scale": 5.0,
"height": 8,
"width": 8,
"video_length": 1,
"output_type": "pt",
"clean_caption": False,
# Request torch outputs so tests compare torch tensors directly (see `BasePipelineTesterConfig`).
"output_type": "pt",
}
return inputs
def test_inference(self):
device = "cpu"
components = self.get_dummy_components()
pipe = self.pipeline_class(**components)
pipe.to(device)
pipe.set_progress_bar_config(disable=None)
class TestLattePipeline(LattePipelineTesterConfig, PipelineTesterMixin):
def test_inference_batch_single_identical(self, batch_size=3, expected_max_diff=1e-3):
super().test_inference_batch_single_identical(batch_size=batch_size, expected_max_diff=expected_max_diff)
inputs = self.get_dummy_inputs(device)
video = pipe(**inputs).frames
generated_video = video[0]
self.assertEqual(generated_video.shape, (1, 3, 8, 8))
expected_video = torch.randn(1, 3, 8, 8)
max_diff = np.abs(generated_video - expected_video).max()
self.assertLessEqual(max_diff, 1e10)
def test_callback_inputs(self):
sig = inspect.signature(self.pipeline_class.__call__)
has_callback_tensor_inputs = "callback_on_step_end_tensor_inputs" in sig.parameters
has_callback_step_end = "callback_on_step_end" in sig.parameters
if not (has_callback_tensor_inputs and has_callback_step_end):
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",
)
def callback_inputs_subset(pipe, i, t, callback_kwargs):
# iterate over callback args
for tensor_name, tensor_value in callback_kwargs.items():
# check that we're only passing in allowed tensor inputs
assert tensor_name in pipe._callback_tensor_inputs
return callback_kwargs
def callback_inputs_all(pipe, i, t, callback_kwargs):
for tensor_name in pipe._callback_tensor_inputs:
assert tensor_name in callback_kwargs
# iterate over callback args
for tensor_name, tensor_value in callback_kwargs.items():
# check that we're only passing in allowed tensor inputs
assert tensor_name in pipe._callback_tensor_inputs
return callback_kwargs
inputs = self.get_dummy_inputs(torch_device)
# Test passing in a subset
inputs["callback_on_step_end"] = callback_inputs_subset
inputs["callback_on_step_end_tensor_inputs"] = ["latents"]
output = pipe(**inputs)[0]
# Test passing in a everything
inputs["callback_on_step_end"] = callback_inputs_all
inputs["callback_on_step_end_tensor_inputs"] = pipe._callback_tensor_inputs
output = pipe(**inputs)[0]
def callback_inputs_change_tensor(pipe, i, t, callback_kwargs):
is_last = i == (pipe.num_timesteps - 1)
if is_last:
callback_kwargs["latents"] = torch.zeros_like(callback_kwargs["latents"])
return callback_kwargs
inputs["callback_on_step_end"] = callback_inputs_change_tensor
inputs["callback_on_step_end_tensor_inputs"] = pipe._callback_tensor_inputs
output = pipe(**inputs)[0]
assert output.abs().sum() < 1e10
def test_inference_batch_single_identical(self):
self._test_inference_batch_single_identical(batch_size=3, expected_max_diff=1e-3)
@unittest.skip("Not supported.")
def test_attention_slicing_forward_pass(self):
pass
@unittest.skipIf(
torch_device != "cuda" or not is_xformers_available(),
reason="XFormers attention is only available with CUDA and `xformers` installed",
)
def test_xformers_attention_forwardGenerator_pass(self):
super()._test_xformers_attention_forwardGenerator_pass(test_mean_pixel_difference=False)
@unittest.skip("Test not supported because `encode_prompt()` has multiple returns.")
@pytest.mark.skip("`encode_prompt()` has multiple returns, which the shared test cannot unpack.")
def test_encode_prompt_works_in_isolation(self):
pass
def test_save_load_optional_components(self):
if not hasattr(self.pipeline_class, "_optional_components"):
return
def test_save_load_optional_components(self, tmp_path, expected_max_difference=1.0):
# `tokenizer` and `text_encoder` are the optional components here, so the pipeline can no longer turn a
# prompt into embeddings once they are dropped. Encode up front and drive the reloaded pipeline with the
# embeddings instead of the shared implementation's plain `get_dummy_inputs()`.
pipe = self.get_pipeline().to(torch_device)
components = self.get_dummy_components()
pipe = self.pipeline_class(**components)
prompt_embeds, negative_prompt_embeds = pipe.encode_prompt(self.get_dummy_inputs()["prompt"])
for component in pipe.components.values():
if hasattr(component, "set_default_attn_processor"):
component.set_default_attn_processor()
pipe.to(torch_device)
pipe.set_progress_bar_config(disable=None)
inputs = self.get_dummy_inputs()
inputs.pop("prompt")
inputs.update(
{
"prompt_embeds": prompt_embeds,
"negative_prompt": None,
"negative_prompt_embeds": negative_prompt_embeds,
"mask_feature": False,
}
)
inputs = self.get_dummy_inputs(torch_device)
prompt = inputs["prompt"]
generator = inputs["generator"]
(
prompt_embeds,
negative_prompt_embeds,
) = pipe.encode_prompt(prompt)
# inputs with prompt converted to embeddings
inputs = {
"prompt_embeds": prompt_embeds,
"negative_prompt": None,
"negative_prompt_embeds": negative_prompt_embeds,
"generator": generator,
"num_inference_steps": 2,
"guidance_scale": 5.0,
"height": 8,
"width": 8,
"video_length": 1,
"mask_feature": False,
"output_type": "pt",
"clean_caption": False,
}
# set all optional components to None
for optional_component in pipe._optional_components:
setattr(pipe, optional_component, None)
output = pipe(**inputs)[0]
with tempfile.TemporaryDirectory() as tmpdir:
pipe.save_pretrained(tmpdir, safe_serialization=False)
pipe_loaded = self.pipeline_class.from_pretrained(tmpdir)
pipe_loaded.to(torch_device)
for component in pipe_loaded.components.values():
if hasattr(component, "set_default_attn_processor"):
component.set_default_attn_processor()
pipe_loaded.set_progress_bar_config(disable=None)
pipe.save_pretrained(tmp_path, safe_serialization=False)
pipe_loaded = self.pipeline_class.from_pretrained(tmp_path)
pipe_loaded.to(torch_device)
pipe_loaded.set_progress_bar_config(disable=None)
for optional_component in pipe._optional_components:
self.assertTrue(
getattr(pipe_loaded, optional_component) is None,
f"`{optional_component}` did not stay set to None after loading.",
assert getattr(pipe_loaded, optional_component) is None, (
f"`{optional_component}` did not stay set to None after loading."
)
output_loaded = pipe_loaded(**inputs)[0]
max_diff = np.abs(to_np(output) - to_np(output_loaded)).max()
self.assertLess(max_diff, 1.0)
assert_tensors_close(
output_loaded,
output,
atol=expected_max_difference,
msg="Output changed after dropping optional components.",
)
class TestLattePipelineMemory(LattePipelineTesterConfig, MemoryTesterMixin):
"""Memory optimization tests (CPU offload, group offload, layerwise casting) for the Latte pipeline."""
class TestLattePipelinePyramidAttentionBroadcast(LattePipelineTesterConfig, PyramidAttentionBroadcastTesterMixin):
"""Pyramid Attention Broadcast tests for the Latte pipeline."""
PAB_CONFIG = {
"spatial_attention_block_skip_range": 2,
"temporal_attention_block_skip_range": 2,
"cross_attention_block_skip_range": 2,
"spatial_attention_timestep_skip_range": (100, 700),
"temporal_attention_timestep_skip_range": (100, 800),
"cross_attention_timestep_skip_range": (100, 800),
"spatial_attention_block_identifiers": ["transformer_blocks"],
"temporal_attention_block_identifiers": ["temporal_transformer_blocks"],
"cross_attention_block_identifiers": ["transformer_blocks"],
}
class TestLattePipelineFasterCache(LattePipelineTesterConfig, FasterCacheTesterMixin):
"""FasterCache tests for the Latte pipeline."""
FASTER_CACHE_CONFIG = {
"spatial_attention_block_skip_range": 2,
"temporal_attention_block_skip_range": 2,
"spatial_attention_timestep_skip_range": (-1, 901),
"temporal_attention_timestep_skip_range": (-1, 901),
"unconditional_batch_skip_range": 2,
"attention_weight_callback": lambda _: 0.5,
}
@slow
@require_torch_accelerator
class LattePipelineIntegrationTests(unittest.TestCase):
class TestLattePipelineIntegration:
prompt = "A painting of a squirrel eating a burger."
def setUp(self):
super().setUp()
@pytest.fixture(autouse=True)
def cleanup(self):
gc.collect()
backend_empty_cache(torch_device)
def tearDown(self):
super().tearDown()
yield
gc.collect()
backend_empty_cache(torch_device)
@@ -324,10 +209,9 @@ class LattePipelineIntegrationTests(unittest.TestCase):
pipe = LattePipeline.from_pretrained("maxin-cn/Latte-1", torch_dtype=torch.float16)
pipe.enable_model_cpu_offload(device=torch_device)
prompt = self.prompt
videos = pipe(
prompt=prompt,
prompt=self.prompt,
height=512,
width=512,
generator=generator,

134
tests/pipelines/ledits_pp/test_ledits_pp_stable_diffusion.py
View File

@@ -15,9 +15,9 @@
import gc
import random
import unittest
import numpy as np
import pytest
import torch
from PIL import Image
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer
@@ -31,6 +31,7 @@ from diffusers import (
from ...testing_utils import (
Expectations,
assert_tensors_close,
backend_empty_cache,
enable_full_determinism,
floats_tensor,
@@ -45,8 +46,10 @@ from ...testing_utils import (
enable_full_determinism()
# LEdits++ does not compose the shared `PipelineTesterMixin`: `__call__` takes neither a prompt nor an image and
# only works once `invert()` has populated the pipeline's state, so the shared dummy-input contract does not apply.
@skip_mps
class LEditsPPPipelineStableDiffusionFastTests(unittest.TestCase):
class TestLEditsPPPipelineStableDiffusion:
pipeline_class = LEditsPPPipelineStableDiffusion
def get_dummy_components(self):
@@ -97,48 +100,42 @@ class LEditsPPPipelineStableDiffusionFastTests(unittest.TestCase):
}
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 = {
"generator": generator,
def get_generator(self, seed=0):
return torch.Generator("cpu").manual_seed(seed)
def get_pipeline(self):
pipe = self.pipeline_class(**self.get_dummy_components()).to(torch_device)
pipe.set_progress_bar_config(disable=None)
return pipe
def get_dummy_inputs(self):
return {
"generator": self.get_generator(0),
"editing_prompt": ["wearing glasses", "sunshine"],
"reverse_editing_direction": [False, True],
"edit_guidance_scale": [10.0, 5.0],
}
return inputs
def get_dummy_inversion_inputs(self, device, seed=0):
def get_dummy_inversion_inputs(self):
images = floats_tensor((2, 3, 32, 32), rng=random.Random(0)).cpu().permute(0, 2, 3, 1)
images = 255 * images
image_1 = Image.fromarray(np.uint8(images[0])).convert("RGB")
image_2 = Image.fromarray(np.uint8(images[1])).convert("RGB")
if str(device).startswith("mps"):
generator = torch.manual_seed(seed)
else:
generator = torch.Generator(device=device).manual_seed(seed)
inputs = {
return {
"image": [image_1, image_2],
"source_prompt": "",
"source_guidance_scale": 3.5,
"num_inversion_steps": 20,
"skip": 0.15,
"generator": generator,
"generator": self.get_generator(0),
}
return inputs
def test_ledits_pp_inversion(self):
device = "cpu" # ensure determinism for the device-dependent torch.Generator
components = self.get_dummy_components()
sd_pipe = LEditsPPPipelineStableDiffusion(**components)
sd_pipe = sd_pipe.to(torch_device)
sd_pipe.set_progress_bar_config(disable=None)
# The expected slices below are CPU-specific.
sd_pipe = self.get_pipeline()
inputs = self.get_dummy_inversion_inputs(device)
inputs = self.get_dummy_inversion_inputs()
inputs["image"] = inputs["image"][0]
sd_pipe.invert(**inputs)
assert sd_pipe.init_latents.shape == (
@@ -148,20 +145,18 @@ class LEditsPPPipelineStableDiffusionFastTests(unittest.TestCase):
int(32 / sd_pipe.vae_scale_factor),
)
latent_slice = sd_pipe.init_latents[0, -1, -3:, -3:].to(device)
latent_slice = sd_pipe.init_latents[0, -1, -3:, -3:].cpu()
expected_slice = np.array([-0.9084, -0.0367, 0.2940, 0.0839, 0.6890, 0.2651, -0.7104, 2.1090, -0.7822])
assert np.abs(latent_slice.flatten() - expected_slice).max() < 1e-3
# fmt: off
expected_slice = torch.tensor([-0.9084, -0.0367, 0.2940, 0.0839, 0.6890, 0.2651, -0.7104, 2.1090, -0.7822])
# fmt: on
assert_tensors_close(latent_slice.flatten(), expected_slice, atol=1e-3)
def test_ledits_pp_inversion_batch(self):
device = "cpu" # ensure determinism for the device-dependent torch.Generator
components = self.get_dummy_components()
sd_pipe = LEditsPPPipelineStableDiffusion(**components)
sd_pipe = sd_pipe.to(torch_device)
sd_pipe.set_progress_bar_config(disable=None)
# The expected slices below are CPU-specific.
sd_pipe = self.get_pipeline()
inputs = self.get_dummy_inversion_inputs(device)
sd_pipe.invert(**inputs)
sd_pipe.invert(**self.get_dummy_inversion_inputs())
assert sd_pipe.init_latents.shape == (
2,
4,
@@ -169,27 +164,23 @@ class LEditsPPPipelineStableDiffusionFastTests(unittest.TestCase):
int(32 / sd_pipe.vae_scale_factor),
)
latent_slice = sd_pipe.init_latents[0, -1, -3:, -3:].to(device)
expected_slice = np.array([0.2528, 0.1458, -0.2166, 0.4565, -0.5657, -1.0286, -0.9961, 0.5933, 1.1173])
assert np.abs(latent_slice.flatten() - expected_slice).max() < 1e-3
latent_slice = sd_pipe.init_latents[1, -1, -3:, -3:].to(device)
expected_slice = np.array([-0.0796, 2.0583, 0.5501, 0.5358, 0.0282, -0.2803, -1.0470, 0.7023, -0.0072])
assert np.abs(latent_slice.flatten() - expected_slice).max() < 1e-3
# fmt: off
expected_slices = torch.tensor(
[
[0.2528, 0.1458, -0.2166, 0.4565, -0.5657, -1.0286, -0.9961, 0.5933, 1.1173],
[-0.0796, 2.0583, 0.5501, 0.5358, 0.0282, -0.2803, -1.0470, 0.7023, -0.0072],
]
)
# fmt: on
for i, expected_slice in enumerate(expected_slices):
latent_slice = sd_pipe.init_latents[i, -1, -3:, -3:].cpu()
assert_tensors_close(latent_slice.flatten(), expected_slice, atol=1e-3)
def test_ledits_pp_warmup_steps(self):
device = "cpu" # ensure determinism for the device-dependent torch.Generator
components = self.get_dummy_components()
pipe = LEditsPPPipelineStableDiffusion(**components)
pipe = pipe.to(torch_device)
pipe.set_progress_bar_config(disable=None)
pipe = self.get_pipeline()
pipe.invert(**self.get_dummy_inversion_inputs())
inversion_inputs = self.get_dummy_inversion_inputs(device)
pipe.invert(**inversion_inputs)
inputs = self.get_dummy_inputs(device)
inputs = self.get_dummy_inputs()
inputs["edit_warmup_steps"] = [0, 5]
pipe(**inputs).images
@@ -204,13 +195,9 @@ class LEditsPPPipelineStableDiffusionFastTests(unittest.TestCase):
pipe(**inputs).images
def test_callback_inputs(self):
device = "cpu" # ensure determinism for the device-dependent torch.Generator
components = self.get_dummy_components()
pipe = LEditsPPPipelineStableDiffusion(**components)
pipe = pipe.to(torch_device)
pipe.set_progress_bar_config(disable=None)
pipe = self.get_pipeline()
inversion_inputs = self.get_dummy_inversion_inputs(device)
inversion_inputs = self.get_dummy_inversion_inputs()
inversion_inputs["image"] = inversion_inputs["image"][0]
pipe.invert(**inversion_inputs)
@@ -220,7 +207,7 @@ class LEditsPPPipelineStableDiffusionFastTests(unittest.TestCase):
return callback_kwargs
inputs = self.get_dummy_inputs(device)
inputs = self.get_dummy_inputs()
inputs["callback_on_step_end"] = callback_inputs_all
inputs["callback_on_step_end_tensor_inputs"] = pipe._callback_tensor_inputs
pipe(**inputs)
@@ -228,26 +215,23 @@ class LEditsPPPipelineStableDiffusionFastTests(unittest.TestCase):
@slow
@require_torch_accelerator
class LEditsPPPipelineStableDiffusionSlowTests(unittest.TestCase):
def setUp(self):
super().setUp()
class TestLEditsPPPipelineStableDiffusionIntegration:
@pytest.fixture(autouse=True)
def cleanup(self):
gc.collect()
backend_empty_cache(torch_device)
yield
gc.collect()
backend_empty_cache(torch_device)
def tearDown(self):
super().tearDown()
gc.collect()
backend_empty_cache(torch_device)
@classmethod
def setUpClass(cls):
raw_image = load_image(
@pytest.fixture(scope="class")
def raw_image(self):
image = load_image(
"https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/pix2pix/cat_6.png"
)
raw_image = raw_image.convert("RGB").resize((512, 512))
cls.raw_image = raw_image
return image.convert("RGB").resize((512, 512))
def test_ledits_pp_editing(self):
def test_ledits_pp_editing(self, raw_image):
pipe = LEditsPPPipelineStableDiffusion.from_pretrained(
"stable-diffusion-v1-5/stable-diffusion-v1-5", safety_checker=None, torch_dtype=torch.float16
)
@@ -255,7 +239,7 @@ class LEditsPPPipelineStableDiffusionSlowTests(unittest.TestCase):
pipe.set_progress_bar_config(disable=None)
generator = torch.manual_seed(0)
_ = pipe.invert(image=self.raw_image, generator=generator)
_ = pipe.invert(image=raw_image, generator=generator)
generator = torch.manual_seed(0)
inputs = {
"generator": generator,

122
tests/pipelines/ledits_pp/test_ledits_pp_stable_diffusion_xl.py
View File

@@ -14,9 +14,9 @@
# limitations under the License.
import random
import unittest
import numpy as np
import pytest
import torch
from PIL import Image
from transformers import (
@@ -38,6 +38,7 @@ from diffusers import (
# from diffusers.image_processor import VaeImageProcessor
from ...testing_utils import (
assert_tensors_close,
enable_full_determinism,
floats_tensor,
load_image,
@@ -51,8 +52,10 @@ from ...testing_utils import (
enable_full_determinism()
# LEdits++ does not compose the shared `PipelineTesterMixin`: `__call__` takes neither a prompt nor an image and
# only works once `invert()` has populated the pipeline's state, so the shared dummy-input contract does not apply.
@skip_mps
class LEditsPPPipelineStableDiffusionXLFastTests(unittest.TestCase):
class TestLEditsPPPipelineStableDiffusionXL:
pipeline_class = LEditsPPPipelineStableDiffusionXL
def get_dummy_components(self, skip_first_text_encoder=False, time_cond_proj_dim=None):
@@ -144,48 +147,42 @@ class LEditsPPPipelineStableDiffusionXLFastTests(unittest.TestCase):
}
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 = {
"generator": generator,
def get_generator(self, seed=0):
return torch.Generator("cpu").manual_seed(seed)
def get_pipeline(self):
pipe = self.pipeline_class(**self.get_dummy_components()).to(torch_device)
pipe.set_progress_bar_config(disable=None)
return pipe
def get_dummy_inputs(self):
return {
"generator": self.get_generator(0),
"editing_prompt": ["wearing glasses", "sunshine"],
"reverse_editing_direction": [False, True],
"edit_guidance_scale": [10.0, 5.0],
}
return inputs
def get_dummy_inversion_inputs(self, device, seed=0):
def get_dummy_inversion_inputs(self):
images = floats_tensor((2, 3, 32, 32), rng=random.Random(0)).cpu().permute(0, 2, 3, 1)
images = 255 * images
image_1 = Image.fromarray(np.uint8(images[0])).convert("RGB")
image_2 = Image.fromarray(np.uint8(images[1])).convert("RGB")
if str(device).startswith("mps"):
generator = torch.manual_seed(seed)
else:
generator = torch.Generator(device=device).manual_seed(seed)
inputs = {
return {
"image": [image_1, image_2],
"source_prompt": "",
"source_guidance_scale": 3.5,
"num_inversion_steps": 20,
"skip": 0.15,
"generator": generator,
"generator": self.get_generator(0),
}
return inputs
def test_ledits_pp_inversion(self):
device = "cpu" # ensure determinism for the device-dependent torch.Generator
components = self.get_dummy_components()
sd_pipe = LEditsPPPipelineStableDiffusionXL(**components)
sd_pipe = sd_pipe.to(torch_device)
sd_pipe.set_progress_bar_config(disable=None)
# The expected slice below is CPU-specific.
sd_pipe = self.get_pipeline()
inputs = self.get_dummy_inversion_inputs(device)
inputs = self.get_dummy_inversion_inputs()
inputs["image"] = inputs["image"][0]
sd_pipe.invert(**inputs)
assert sd_pipe.init_latents.shape == (
@@ -195,19 +192,18 @@ class LEditsPPPipelineStableDiffusionXLFastTests(unittest.TestCase):
int(32 / sd_pipe.vae_scale_factor),
)
latent_slice = sd_pipe.init_latents[0, -1, -3:, -3:].to(device)
expected_slice = np.array([-0.9084, -0.0367, 0.2940, 0.0839, 0.6890, 0.2651, -0.7103, 2.1090, -0.7821])
assert np.abs(latent_slice.flatten() - expected_slice).max() < 1e-3
latent_slice = sd_pipe.init_latents[0, -1, -3:, -3:].cpu()
# fmt: off
expected_slice = torch.tensor([-0.9084, -0.0367, 0.2940, 0.0839, 0.6890, 0.2651, -0.7103, 2.1090, -0.7821])
# fmt: on
assert_tensors_close(latent_slice.flatten(), expected_slice, atol=1e-3)
def test_ledits_pp_inversion_batch(self):
device = "cpu" # ensure determinism for the device-dependent torch.Generator
components = self.get_dummy_components()
sd_pipe = LEditsPPPipelineStableDiffusionXL(**components)
sd_pipe = sd_pipe.to(torch_device)
sd_pipe.set_progress_bar_config(disable=None)
# The expected slices below are CPU-specific.
sd_pipe = self.get_pipeline()
inputs = self.get_dummy_inversion_inputs(device)
sd_pipe.invert(**inputs)
sd_pipe.invert(**self.get_dummy_inversion_inputs())
assert sd_pipe.init_latents.shape == (
2,
4,
@@ -215,29 +211,26 @@ class LEditsPPPipelineStableDiffusionXLFastTests(unittest.TestCase):
int(32 / sd_pipe.vae_scale_factor),
)
latent_slice = sd_pipe.init_latents[0, -1, -3:, -3:].to(device)
expected_slice = np.array([0.2528, 0.1458, -0.2166, 0.4565, -0.5656, -1.0286, -0.9961, 0.5933, 1.1172])
assert np.abs(latent_slice.flatten() - expected_slice).max() < 1e-3
latent_slice = sd_pipe.init_latents[1, -1, -3:, -3:].to(device)
expected_slice = np.array([-0.0796, 2.0583, 0.5500, 0.5358, 0.0282, -0.2803, -1.0470, 0.7024, -0.0072])
assert np.abs(latent_slice.flatten() - expected_slice).max() < 1e-3
# fmt: off
expected_slices = torch.tensor(
[
[0.2528, 0.1458, -0.2166, 0.4565, -0.5656, -1.0286, -0.9961, 0.5933, 1.1172],
[-0.0796, 2.0583, 0.5500, 0.5358, 0.0282, -0.2803, -1.0470, 0.7024, -0.0072],
]
)
# fmt: on
for i, expected_slice in enumerate(expected_slices):
latent_slice = sd_pipe.init_latents[i, -1, -3:, -3:].cpu()
assert_tensors_close(latent_slice.flatten(), expected_slice, atol=1e-3)
def test_ledits_pp_warmup_steps(self):
device = "cpu" # ensure determinism for the device-dependent torch.Generator
components = self.get_dummy_components()
pipe = LEditsPPPipelineStableDiffusionXL(**components)
pipe = pipe.to(torch_device)
pipe.set_progress_bar_config(disable=None)
pipe = self.get_pipeline()
inversion_inputs = self.get_dummy_inversion_inputs(device)
inversion_inputs = self.get_dummy_inversion_inputs()
inversion_inputs["image"] = inversion_inputs["image"][0]
pipe.invert(**inversion_inputs)
inputs = self.get_dummy_inputs(device)
inputs = self.get_dummy_inputs()
inputs["edit_warmup_steps"] = [0, 5]
pipe(**inputs).images
@@ -252,13 +245,9 @@ class LEditsPPPipelineStableDiffusionXLFastTests(unittest.TestCase):
pipe(**inputs).images
def test_callback_inputs(self):
device = "cpu" # ensure determinism for the device-dependent torch.Generator
components = self.get_dummy_components()
pipe = LEditsPPPipelineStableDiffusionXL(**components)
pipe = pipe.to(torch_device)
pipe.set_progress_bar_config(disable=None)
pipe = self.get_pipeline()
inversion_inputs = self.get_dummy_inversion_inputs(device)
inversion_inputs = self.get_dummy_inversion_inputs()
inversion_inputs["image"] = inversion_inputs["image"][0]
pipe.invert(**inversion_inputs)
@@ -268,7 +257,7 @@ class LEditsPPPipelineStableDiffusionXLFastTests(unittest.TestCase):
return callback_kwargs
inputs = self.get_dummy_inputs(device)
inputs = self.get_dummy_inputs()
inputs["callback_on_step_end"] = callback_inputs_all
inputs["callback_on_step_end_tensor_inputs"] = pipe._callback_tensor_inputs
pipe(**inputs)
@@ -276,16 +265,15 @@ class LEditsPPPipelineStableDiffusionXLFastTests(unittest.TestCase):
@slow
@require_torch_accelerator
class LEditsPPPipelineStableDiffusionXLSlowTests(unittest.TestCase):
@classmethod
def setUpClass(cls):
raw_image = load_image(
class TestLEditsPPPipelineStableDiffusionXLIntegration:
@pytest.fixture(scope="class")
def raw_image(self):
image = load_image(
"https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/pix2pix/cat_6.png"
)
raw_image = raw_image.convert("RGB").resize((512, 512))
cls.raw_image = raw_image
return image.convert("RGB").resize((512, 512))
def test_ledits_pp_edit(self):
def test_ledits_pp_edit(self, raw_image):
pipe = LEditsPPPipelineStableDiffusionXL.from_pretrained(
"stabilityai/stable-diffusion-xl-base-1.0", safety_checker=None, add_watermarker=None
)
@@ -293,7 +281,7 @@ class LEditsPPPipelineStableDiffusionXLSlowTests(unittest.TestCase):
pipe.set_progress_bar_config(disable=None)
generator = torch.manual_seed(0)
_ = pipe.invert(image=self.raw_image, generator=generator, num_zero_noise_steps=0)
_ = pipe.invert(image=raw_image, generator=generator, num_zero_noise_steps=0)
inputs = {
"generator": generator,
"editing_prompt": ["cat", "dog"],

83
tests/pipelines/llada2/test_llada2.py
View File

@@ -1,5 +1,4 @@
import unittest
import pytest
import torch
from diffusers import BlockRefinementScheduler, LLaDA2Pipeline
@@ -41,7 +40,7 @@ def _make_pipeline(tokenizer=None):
return LLaDA2Pipeline(model=model, scheduler=scheduler, tokenizer=tokenizer)
class LLaDA2PipelineTest(unittest.TestCase):
class TestLLaDA2Pipeline:
def test_pipeline_runs(self):
pipe = _make_pipeline().to("cpu")
@@ -61,8 +60,8 @@ class LLaDA2PipelineTest(unittest.TestCase):
output_type="seq",
)
self.assertEqual(out.sequences.shape, (2, 24))
self.assertFalse((out.sequences == 31).any().item())
assert out.sequences.shape == (2, 24)
assert not (out.sequences == 31).any().item()
def test_pipeline_return_tuple(self):
pipe = _make_pipeline().to("cpu")
@@ -83,8 +82,8 @@ class LLaDA2PipelineTest(unittest.TestCase):
return_dict=False,
)
self.assertEqual(sequences.shape, (1, 16))
self.assertIsNone(texts)
assert sequences.shape == (1, 16)
assert texts is None
def test_output_type_seq(self):
"""output_type='seq' should return sequences but no texts."""
@@ -104,9 +103,9 @@ class LLaDA2PipelineTest(unittest.TestCase):
output_type="seq",
)
self.assertIsNotNone(out.sequences)
self.assertEqual(out.sequences.shape, (1, 16))
self.assertIsNone(out.texts)
assert out.sequences is not None
assert out.sequences.shape == (1, 16)
assert out.texts is None
def test_output_type_text_without_tokenizer(self):
"""output_type='text' without a tokenizer should return texts=None."""
@@ -126,8 +125,8 @@ class LLaDA2PipelineTest(unittest.TestCase):
output_type="text",
)
self.assertIsNotNone(out.sequences)
self.assertIsNone(out.texts)
assert out.sequences is not None
assert out.texts is None
def test_output_type_text_with_tokenizer(self):
"""output_type='text' with a tokenizer should return decoded texts."""
@@ -155,16 +154,16 @@ class LLaDA2PipelineTest(unittest.TestCase):
output_type="text",
)
self.assertIsNotNone(out.sequences)
self.assertIsNotNone(out.texts)
self.assertEqual(len(out.texts), 1)
self.assertTrue(out.texts[0].startswith("decoded_"))
assert out.sequences is not None
assert out.texts is not None
assert len(out.texts) == 1
assert out.texts[0].startswith("decoded_")
def test_output_type_invalid_raises(self):
"""Invalid output_type should raise ValueError."""
pipe = _make_pipeline().to("cpu")
with self.assertRaises(ValueError):
with pytest.raises(ValueError):
pipe(
input_ids=torch.tensor([[5, 6, 7, 8]], dtype=torch.long),
use_chat_template=False,
@@ -186,9 +185,9 @@ class LLaDA2PipelineTest(unittest.TestCase):
add_generation_prompt=False,
chat_template_kwargs=None,
)
self.assertTrue(torch.equal(result_ids, ids))
self.assertEqual(result_mask.shape, ids.shape)
self.assertTrue((result_mask == 1).all().item())
assert torch.equal(result_ids, ids)
assert result_mask.shape == ids.shape
assert (result_mask == 1).all().item()
def test_prepare_input_ids_from_1d_tensor(self):
pipe = _make_pipeline()
@@ -201,12 +200,12 @@ class LLaDA2PipelineTest(unittest.TestCase):
add_generation_prompt=False,
chat_template_kwargs=None,
)
self.assertEqual(result_ids.shape, (1, 3))
self.assertEqual(result_mask.shape, (1, 3))
assert result_ids.shape == (1, 3)
assert result_mask.shape == (1, 3)
def test_prepare_input_ids_no_tokenizer_raises(self):
pipe = _make_pipeline(tokenizer=None)
with self.assertRaises(ValueError):
with pytest.raises(ValueError):
pipe._prepare_input_ids(
prompt="hello",
messages=None,
@@ -220,7 +219,7 @@ class LLaDA2PipelineTest(unittest.TestCase):
pipe = _make_pipeline()
# Manually set tokenizer to a simple object so _prepare_input_ids doesn't short-circuit
pipe.tokenizer = type("Tok", (), {"eos_token_id": None, "mask_token_id": None})()
with self.assertRaises(ValueError):
with pytest.raises(ValueError):
pipe._prepare_input_ids(
prompt="hello",
messages=[{"role": "user", "content": "hi"}],
@@ -233,7 +232,7 @@ class LLaDA2PipelineTest(unittest.TestCase):
def test_prepare_input_ids_neither_raises(self):
pipe = _make_pipeline()
pipe.tokenizer = type("Tok", (), {"eos_token_id": None, "mask_token_id": None})()
with self.assertRaises(ValueError):
with pytest.raises(ValueError):
pipe._prepare_input_ids(
prompt=None,
messages=None,
@@ -244,7 +243,7 @@ class LLaDA2PipelineTest(unittest.TestCase):
)
class LLaDA2RegressionTest(unittest.TestCase):
class TestLLaDA2Regression:
"""Pin the regressions identified in https://github.com/huggingface/diffusers/issues/13598."""
def test_attention_mask_carried_through_for_pre_tokenized_input(self):
@@ -278,16 +277,16 @@ class LLaDA2RegressionTest(unittest.TestCase):
output_type="seq",
)
self.assertGreater(len(captured), 0)
assert len(captured) > 0
first_mask = captured[0]
# Padded prompt positions stay zero in the runtime mask (Issue #1).
self.assertEqual(first_mask[0, 3].item(), 0)
self.assertEqual(first_mask[1, 1].item(), 0)
self.assertEqual(first_mask[1, 2].item(), 0)
self.assertEqual(first_mask[1, 3].item(), 0)
assert first_mask[0, 3].item() == 0
assert first_mask[1, 1].item() == 0
assert first_mask[1, 2].item() == 0
assert first_mask[1, 3].item() == 0
# Real prompt positions stay one.
self.assertEqual(first_mask[0, 0].item(), 1)
self.assertEqual(first_mask[1, 0].item(), 1)
assert first_mask[0, 0].item() == 1
assert first_mask[1, 0].item() == 1
def test_block_length_routes_into_scheduler_transfer_schedule(self):
"""Issue #2: the per-call `block_length` must drive the scheduler's `_transfer_schedule`."""
@@ -313,9 +312,9 @@ class LLaDA2RegressionTest(unittest.TestCase):
callback_on_step_end_tensor_inputs=["transfer_index"],
)
# With block_length=num_inference_steps=8 the schedule commits exactly one token per step.
self.assertEqual(commits[0], 1)
self.assertEqual(commits[1], 1)
self.assertEqual(commits[2], 1)
assert commits[0] == 1
assert commits[1] == 1
assert commits[2] == 1
def test_callback_tensor_inputs_advertised_keys_resolve(self):
"""Issue #3: every advertised callback key must be a bound local at callback time."""
@@ -341,7 +340,7 @@ class LLaDA2RegressionTest(unittest.TestCase):
callback_on_step_end=cb,
callback_on_step_end_tensor_inputs=keys,
)
self.assertEqual(set(observed), set(keys))
assert set(observed) == set(keys)
def test_eos_at_first_generated_position_triggers_finished(self):
"""Issue #4: EOS exactly at index `prompt_length` must mark the row finished."""
@@ -357,7 +356,7 @@ class LLaDA2RegressionTest(unittest.TestCase):
mask_token_id=99,
prompt_length=1,
)
self.assertTrue(bool(finished[0].item()))
assert bool(finished[0].item())
def test_finished_rows_are_frozen_for_subsequent_blocks(self):
"""Issue #5: once a row emits EOS, later blocks must not overwrite its committed tokens."""
@@ -393,7 +392,7 @@ class LLaDA2RegressionTest(unittest.TestCase):
output_type="seq",
)
# Row 0's first generated tokens must not be overwritten by later-block sampling (token 7).
self.assertNotIn(7, out.sequences[0].tolist()[:2])
assert 7 not in out.sequences[0].tolist()[:2]
def test_progress_bar_disable_is_preserved_after_call(self):
"""Issue #6: calling the pipeline must not mutate `_progress_bar_config`."""
@@ -412,8 +411,4 @@ class LLaDA2RegressionTest(unittest.TestCase):
eos_early_stop=False,
output_type="seq",
)
self.assertEqual(pipe._progress_bar_config, before)
if __name__ == "__main__":
unittest.main()
assert pipe._progress_bar_config == before

144
tests/pipelines/longcat_audio_dit/test_longcat_audio_dit.py
View File

@@ -13,9 +13,9 @@
# limitations under the License.
import os
import unittest
from pathlib import Path
import pytest
import torch
from transformers import AutoTokenizer, UMT5Config, UMT5EncoderModel
@@ -26,24 +26,33 @@ from diffusers import (
LongCatAudioDiTVae,
)
from ...testing_utils import enable_full_determinism, require_torch_accelerator, slow, torch_device
from ...testing_utils import (
assert_tensors_close,
enable_full_determinism,
require_torch_accelerator,
slow,
torch_device,
)
from ..pipeline_params import TEXT_TO_AUDIO_BATCH_PARAMS, TEXT_TO_AUDIO_PARAMS
from ..test_pipelines_common import PipelineTesterMixin
from ..testing_utils import BasePipelineTesterConfig, MemoryTesterMixin, PipelineTesterMixin
enable_full_determinism()
class LongCatAudioDiTPipelineFastTests(PipelineTesterMixin, unittest.TestCase):
class LongCatAudioDiTPipelineTesterConfig(BasePipelineTesterConfig):
pipeline_class = LongCatAudioDiTPipeline
params = (
# This pipeline sizes its output with `audio_duration_s` and takes no precomputed prompt embeddings.
required_input_params_in_call_signature = (
TEXT_TO_AUDIO_PARAMS
- {"audio_length_in_s", "prompt_embeds", "negative_prompt_embeds", "cross_attention_kwargs"}
) | {"audio_duration_s"}
batch_params = TEXT_TO_AUDIO_BATCH_PARAMS
required_optional_params = PipelineTesterMixin.required_optional_params - {"num_images_per_prompt"}
test_attention_slicing = False
test_xformers_attention = False
batch_input_params = TEXT_TO_AUDIO_BATCH_PARAMS
# Waveform length for `audio_duration_s=0.1` at the tiny VAE's 24 kHz sample rate, as one mono channel.
output_shape = (1, 4800)
# An audio pipeline: `__call__` has no `num_images_per_prompt`, and the noise is always sampled internally
# rather than passed in as `latents`.
optional_input_params = frozenset(["num_inference_steps", "generator", "output_type", "return_dict"])
def get_dummy_components(self):
torch.manual_seed(0)
@@ -77,85 +86,32 @@ class LongCatAudioDiTPipelineFastTests(PipelineTesterMixin, unittest.TestCase):
"transformer": transformer,
}
def get_dummy_inputs(self, device, seed=0, prompt="soft ocean ambience"):
if str(device).startswith("mps"):
generator = torch.manual_seed(seed)
else:
generator = torch.Generator(device=device).manual_seed(seed)
def get_dummy_inputs(self):
return {
"prompt": prompt,
"prompt": "soft ocean ambience",
"audio_duration_s": 0.1,
"num_inference_steps": 2,
"guidance_scale": 1.0,
"generator": generator,
"generator": self.get_generator(0),
# Request torch outputs so tests compare torch tensors directly (see `BasePipelineTesterConfig`).
"output_type": "pt",
}
class TestLongCatAudioDiTPipeline(LongCatAudioDiTPipelineTesterConfig, PipelineTesterMixin):
def test_inference(self):
device = "cpu"
pipe = self.pipeline_class(**self.get_dummy_components())
pipe.to(device)
pipe.set_progress_bar_config(disable=None)
pipe = self.get_pipeline().to(torch_device)
output = pipe(**self.get_dummy_inputs(device)).audios
audios = pipe(**self.get_dummy_inputs()).audios
self.assertEqual(output.ndim, 3)
self.assertEqual(output.shape[0], 1)
self.assertEqual(output.shape[1], 1)
self.assertGreater(output.shape[-1], 0)
assert audios.shape == (1, *self.output_shape)
def test_save_load_local(self):
import tempfile
device = "cpu"
pipe = self.pipeline_class(**self.get_dummy_components())
pipe.to(device)
with tempfile.TemporaryDirectory() as tmp_dir:
pipe.save_pretrained(tmp_dir)
reloaded = self.pipeline_class.from_pretrained(tmp_dir, local_files_only=True)
output = reloaded(**self.get_dummy_inputs(device, seed=0)).audios
self.assertIsInstance(reloaded, LongCatAudioDiTPipeline)
self.assertEqual(output.ndim, 3)
self.assertGreater(output.shape[-1], 0)
def test_inference_batch_single_identical(self):
self._test_inference_batch_single_identical(expected_max_diff=2e-3)
def test_model_cpu_offload_forward_pass(self):
self.skipTest(
"LongCatAudioDiTPipeline offload coverage is not ready for the standard PipelineTesterMixin test."
)
def test_cpu_offload_forward_pass_twice(self):
self.skipTest(
"LongCatAudioDiTPipeline offload coverage is not ready for the standard PipelineTesterMixin test."
)
def test_sequential_cpu_offload_forward_pass(self):
self.skipTest(
"LongCatAudioDiTPipeline uses `torch.nn.utils.weight_norm`, which is not compatible with "
"sequential offloading."
)
def test_sequential_offload_forward_pass_twice(self):
self.skipTest(
"LongCatAudioDiTPipeline uses `torch.nn.utils.weight_norm`, which is not compatible with "
"sequential offloading."
)
def test_pipeline_level_group_offloading_inference(self):
self.skipTest(
"LongCatAudioDiTPipeline group offloading coverage is not ready for the standard PipelineTesterMixin test."
)
def test_num_images_per_prompt(self):
self.skipTest("LongCatAudioDiTPipeline does not support num_images_per_prompt.")
def test_inference_batch_single_identical(self, batch_size=3, expected_max_diff=2e-3):
super().test_inference_batch_single_identical(batch_size=batch_size, expected_max_diff=expected_max_diff)
@pytest.mark.skip("`LongCatAudioDiTPipeline.encode_prompt` has a custom signature.")
def test_encode_prompt_works_in_isolation(self):
self.skipTest("LongCatAudioDiTPipeline.encode_prompt has a custom signature.")
pass
def test_uniform_flow_match_scheduler_grid_matches_manual_updates(self):
num_inference_steps = 6
@@ -165,7 +121,7 @@ class LongCatAudioDiTPipelineFastTests(PipelineTesterMixin, unittest.TestCase):
expected_grid = torch.linspace(0, 1, num_inference_steps + 1, dtype=torch.float32)
actual_timesteps = scheduler.timesteps / scheduler.config.num_train_timesteps
self.assertTrue(torch.allclose(actual_timesteps, expected_grid[:-1], atol=1e-6, rtol=0))
assert_tensors_close(actual_timesteps, expected_grid[:-1], atol=1e-6, rtol=0)
sample = torch.zeros(1, 2, 3)
model_output = torch.ones_like(sample)
@@ -174,7 +130,31 @@ class LongCatAudioDiTPipelineFastTests(PipelineTesterMixin, unittest.TestCase):
expected = expected + model_output * (t1 - t0)
sample = scheduler.step(model_output, scheduler_t, sample, return_dict=False)[0]
self.assertTrue(torch.allclose(sample, expected, atol=1e-6, rtol=0))
assert_tensors_close(sample, expected, atol=1e-6, rtol=0)
class TestLongCatAudioDiTPipelineMemory(LongCatAudioDiTPipelineTesterConfig, MemoryTesterMixin):
"""Memory optimization tests (CPU offload, group offload, layerwise casting) for the LongCat AudioDiT pipeline."""
@pytest.mark.skip("Offload coverage is not ready for this pipeline.")
def test_model_cpu_offload_forward_pass(self):
pass
@pytest.mark.skip("Offload coverage is not ready for this pipeline.")
def test_cpu_offload_forward_pass_twice(self):
pass
@pytest.mark.skip("The pipeline uses `torch.nn.utils.weight_norm`, incompatible with sequential offloading.")
def test_sequential_cpu_offload_forward_pass(self):
pass
@pytest.mark.skip("The pipeline uses `torch.nn.utils.weight_norm`, incompatible with sequential offloading.")
def test_sequential_offload_forward_pass_twice(self):
pass
@pytest.mark.skip("Group offloading coverage is not ready for this pipeline.")
def test_pipeline_level_group_offloading_inference(self):
pass
def test_longcat_audio_top_level_imports():
@@ -185,22 +165,20 @@ def test_longcat_audio_top_level_imports():
@slow
@require_torch_accelerator
class LongCatAudioDiTPipelineSlowTests(unittest.TestCase):
pipeline_class = LongCatAudioDiTPipeline
class TestLongCatAudioDiTPipelineIntegration:
def test_longcat_audio_pipeline_from_pretrained_real_local_weights(self):
model_path = Path(
os.getenv("LONGCAT_AUDIO_DIT_MODEL_PATH", "/data/models/meituan-longcat/LongCat-AudioDiT-1B")
)
tokenizer_path_env = os.getenv("LONGCAT_AUDIO_DIT_TOKENIZER_PATH")
if tokenizer_path_env is None:
raise unittest.SkipTest("LONGCAT_AUDIO_DIT_TOKENIZER_PATH is not set")
pytest.skip("LONGCAT_AUDIO_DIT_TOKENIZER_PATH is not set")
tokenizer_path = Path(tokenizer_path_env)
if not model_path.exists():
raise unittest.SkipTest(f"LongCat-AudioDiT model path not found: {model_path}")
pytest.skip(f"LongCat-AudioDiT model path not found: {model_path}")
if not tokenizer_path.exists():
raise unittest.SkipTest(f"LongCat-AudioDiT tokenizer path not found: {tokenizer_path}")
pytest.skip(f"LongCat-AudioDiT tokenizer path not found: {tokenizer_path}")
pipe = LongCatAudioDiTPipeline.from_pretrained(
model_path,

225
tests/pipelines/ltx/test_ltx_condition.py
View File

@@ -12,10 +12,6 @@
# See the License for the specific language governing permissions and
# limitations under the License.
import inspect
import unittest
import numpy as np
import torch
from transformers import AutoConfig, AutoTokenizer, T5EncoderModel
@@ -27,31 +23,40 @@ from diffusers import (
)
from diffusers.pipelines.ltx.pipeline_ltx_condition import LTXVideoCondition
from ...testing_utils import enable_full_determinism, 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 PipelineTesterMixin, to_np
from ...testing_utils import assert_tensors_close, enable_full_determinism, torch_device
from ..testing_utils import (
BasePipelineTesterConfig,
LoraMemoryTesterMixin,
LoraTesterMixin,
MemoryTesterMixin,
PipelineTesterMixin,
)
enable_full_determinism()
class LTXConditionPipelineFastTests(PipelineTesterMixin, unittest.TestCase):
class LTXConditionPipelineTesterConfig(BasePipelineTesterConfig):
pipeline_class = LTXConditionPipeline
params = TEXT_TO_IMAGE_PARAMS - {"cross_attention_kwargs"}
batch_params = TEXT_TO_IMAGE_BATCH_PARAMS.union({"image"})
image_params = TEXT_TO_IMAGE_IMAGE_PARAMS
image_latents_params = TEXT_TO_IMAGE_IMAGE_PARAMS
required_optional_params = frozenset(
required_input_params_in_call_signature = frozenset(
[
"num_inference_steps",
"generator",
"latents",
"return_dict",
"callback_on_step_end",
"callback_on_step_end_tensor_inputs",
"conditions",
"image",
"prompt",
"height",
"width",
"guidance_scale",
"negative_prompt",
"prompt_embeds",
"negative_prompt_embeds",
]
)
test_xformers_attention = False
batch_input_params = frozenset(["prompt", "negative_prompt", "image"])
output_shape = (9, 3, 32, 32)
# LTX is a video pipeline: it exposes `num_videos_per_prompt`, not the base default `num_images_per_prompt`.
optional_input_params = frozenset(
["num_inference_steps", "num_videos_per_prompt", "generator", "latents", "output_type", "return_dict"]
)
def get_dummy_components(self):
torch.manual_seed(0)
@@ -98,31 +103,25 @@ class LTXConditionPipelineFastTests(PipelineTesterMixin, unittest.TestCase):
text_encoder = T5EncoderModel(config).eval()
tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5")
components = {
return {
"transformer": transformer,
"vae": vae,
"scheduler": scheduler,
"text_encoder": text_encoder,
"tokenizer": tokenizer,
}
return components
def get_dummy_inputs(self, device, seed=0, use_conditions=False):
if str(device).startswith("mps"):
generator = torch.manual_seed(seed)
else:
generator = torch.Generator(device=device).manual_seed(seed)
def get_dummy_inputs(self, use_conditions: bool = False):
"""Dummy inputs for the pipeline.
image = torch.randn((1, 3, 32, 32), generator=generator, device=device)
if use_conditions:
conditions = LTXVideoCondition(
image=image,
)
else:
conditions = None
The conditioning image can be passed either directly as `image` or wrapped in an `LTXVideoCondition`;
`use_conditions` picks the latter so `test_inference` can check the two routes agree.
"""
generator = self.get_generator(0)
image = torch.randn((1, 3, 32, 32), generator=generator)
inputs = {
"conditions": conditions,
return {
"conditions": LTXVideoCondition(image=image) if use_conditions else None,
"image": None if use_conditions else image,
"prompt": "dance monkey",
"negative_prompt": "",
@@ -134,138 +133,34 @@ class LTXConditionPipelineFastTests(PipelineTesterMixin, unittest.TestCase):
# 8 * k + 1 is the recommendation
"num_frames": 9,
"max_sequence_length": 16,
# Request torch outputs so tests compare torch tensors directly (see `BasePipelineTesterConfig`).
"output_type": "pt",
}
return inputs
class TestLTXConditionPipeline(LTXConditionPipelineTesterConfig, PipelineTesterMixin):
def test_inference(self):
device = "cpu"
pipe = self.get_pipeline().to(torch_device)
components = self.get_dummy_components()
pipe = self.pipeline_class(**components)
pipe.to(device)
pipe.set_progress_bar_config(disable=None)
generated_video = pipe(**self.get_dummy_inputs()).frames[0]
generated_video_with_conditions = pipe(**self.get_dummy_inputs(use_conditions=True)).frames[0]
inputs = self.get_dummy_inputs(device)
inputs2 = self.get_dummy_inputs(device, use_conditions=True)
video = pipe(**inputs).frames
generated_video = video[0]
video2 = pipe(**inputs2).frames
generated_video2 = video2[0]
self.assertEqual(generated_video.shape, (9, 3, 32, 32))
max_diff = np.abs(generated_video - generated_video2).max()
self.assertLessEqual(max_diff, 1e-3)
def test_callback_inputs(self):
sig = inspect.signature(self.pipeline_class.__call__)
has_callback_tensor_inputs = "callback_on_step_end_tensor_inputs" in sig.parameters
has_callback_step_end = "callback_on_step_end" in sig.parameters
if not (has_callback_tensor_inputs and has_callback_step_end):
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 generated_video.shape == self.output_shape
assert_tensors_close(
generated_video_with_conditions,
generated_video,
atol=1e-3,
msg="Passing the image through `conditions` should match passing it as `image`.",
)
def callback_inputs_subset(pipe, i, t, callback_kwargs):
# iterate over callback args
for tensor_name, tensor_value in callback_kwargs.items():
# check that we're only passing in allowed tensor inputs
assert tensor_name in pipe._callback_tensor_inputs
return callback_kwargs
def callback_inputs_all(pipe, i, t, callback_kwargs):
for tensor_name in pipe._callback_tensor_inputs:
assert tensor_name in callback_kwargs
# iterate over callback args
for tensor_name, tensor_value in callback_kwargs.items():
# check that we're only passing in allowed tensor inputs
assert tensor_name in pipe._callback_tensor_inputs
return callback_kwargs
inputs = self.get_dummy_inputs(torch_device)
# Test passing in a subset
inputs["callback_on_step_end"] = callback_inputs_subset
inputs["callback_on_step_end_tensor_inputs"] = ["latents"]
output = pipe(**inputs)[0]
# Test passing in a everything
inputs["callback_on_step_end"] = callback_inputs_all
inputs["callback_on_step_end_tensor_inputs"] = pipe._callback_tensor_inputs
output = pipe(**inputs)[0]
def callback_inputs_change_tensor(pipe, i, t, callback_kwargs):
is_last = i == (pipe.num_timesteps - 1)
if is_last:
callback_kwargs["latents"] = torch.zeros_like(callback_kwargs["latents"])
return callback_kwargs
inputs["callback_on_step_end"] = callback_inputs_change_tensor
inputs["callback_on_step_end_tensor_inputs"] = pipe._callback_tensor_inputs
output = pipe(**inputs)[0]
assert output.abs().sum() < 1e10
def test_inference_batch_single_identical(self):
self._test_inference_batch_single_identical(batch_size=3, expected_max_diff=1e-3)
def test_attention_slicing_forward_pass(
self, test_max_difference=True, test_mean_pixel_difference=True, expected_max_diff=1e-3
):
if not self.test_attention_slicing:
return
components = self.get_dummy_components()
pipe = self.pipeline_class(**components)
for component in pipe.components.values():
if hasattr(component, "set_default_attn_processor"):
component.set_default_attn_processor()
pipe.to(torch_device)
pipe.set_progress_bar_config(disable=None)
generator_device = "cpu"
inputs = self.get_dummy_inputs(generator_device)
output_without_slicing = pipe(**inputs)[0]
pipe.enable_attention_slicing(slice_size=1)
inputs = self.get_dummy_inputs(generator_device)
output_with_slicing1 = pipe(**inputs)[0]
pipe.enable_attention_slicing(slice_size=2)
inputs = self.get_dummy_inputs(generator_device)
output_with_slicing2 = pipe(**inputs)[0]
if test_max_difference:
max_diff1 = np.abs(to_np(output_with_slicing1) - to_np(output_without_slicing)).max()
max_diff2 = np.abs(to_np(output_with_slicing2) - to_np(output_without_slicing)).max()
self.assertLess(
max(max_diff1, max_diff2),
expected_max_diff,
"Attention slicing should not affect the inference results",
)
def test_inference_batch_single_identical(self, batch_size=3, expected_max_diff=1e-3):
super().test_inference_batch_single_identical(batch_size=batch_size, expected_max_diff=expected_max_diff)
def test_vae_tiling(self, expected_diff_max: float = 0.2):
generator_device = "cpu"
components = self.get_dummy_components()
pipe = self.pipeline_class(**components)
pipe.to("cpu")
pipe.set_progress_bar_config(disable=None)
pipe = self.get_pipeline().to(torch_device)
# Without tiling
inputs = self.get_dummy_inputs(generator_device)
inputs = self.get_dummy_inputs()
inputs["height"] = inputs["width"] = 128
output_without_tiling = pipe(**inputs)[0]
@@ -276,12 +171,22 @@ class LTXConditionPipelineFastTests(PipelineTesterMixin, unittest.TestCase):
tile_sample_stride_height=64,
tile_sample_stride_width=64,
)
inputs = self.get_dummy_inputs(generator_device)
inputs = self.get_dummy_inputs()
inputs["height"] = inputs["width"] = 128
output_with_tiling = pipe(**inputs)[0]
self.assertLess(
(to_np(output_without_tiling) - to_np(output_with_tiling)).max(),
expected_diff_max,
"VAE tiling should not affect the inference results",
assert (output_without_tiling - output_with_tiling).abs().max() < expected_diff_max, (
"VAE tiling should not affect the inference results."
)
class TestLTXConditionPipelineMemory(LTXConditionPipelineTesterConfig, MemoryTesterMixin):
"""Memory optimization tests (CPU offload, group offload, layerwise casting) for the LTX condition pipeline."""
class TestLTXConditionPipelineLoRA(LTXConditionPipelineTesterConfig, LoraTesterMixin):
"""LoRA tests for the LTX condition pipeline."""
class TestLTXConditionPipelineLoRAMemory(LTXConditionPipelineTesterConfig, LoraMemoryTesterMixin):
"""LoRA x memory-optimization tests (group offload, CPU offload) for the LTX condition pipeline."""

204
tests/pipelines/ltx/test_ltx_image2video.py
View File

@@ -12,10 +12,6 @@
# See the License for the specific language governing permissions and
# limitations under the License.
import inspect
import unittest
import numpy as np
import torch
from transformers import AutoConfig, AutoTokenizer, T5EncoderModel
@@ -27,30 +23,38 @@ from diffusers import (
)
from ...testing_utils import enable_full_determinism, 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 PipelineTesterMixin, to_np
from ..testing_utils import (
BasePipelineTesterConfig,
LoraMemoryTesterMixin,
LoraTesterMixin,
MemoryTesterMixin,
PipelineTesterMixin,
)
enable_full_determinism()
class LTXImageToVideoPipelineFastTests(PipelineTesterMixin, unittest.TestCase):
class LTXImageToVideoPipelineTesterConfig(BasePipelineTesterConfig):
pipeline_class = LTXImageToVideoPipeline
params = TEXT_TO_IMAGE_PARAMS - {"cross_attention_kwargs"}
batch_params = TEXT_TO_IMAGE_BATCH_PARAMS.union({"image"})
image_params = TEXT_TO_IMAGE_IMAGE_PARAMS
image_latents_params = TEXT_TO_IMAGE_IMAGE_PARAMS
required_optional_params = frozenset(
required_input_params_in_call_signature = frozenset(
[
"num_inference_steps",
"generator",
"latents",
"return_dict",
"callback_on_step_end",
"callback_on_step_end_tensor_inputs",
"image",
"prompt",
"height",
"width",
"guidance_scale",
"negative_prompt",
"prompt_embeds",
"negative_prompt_embeds",
]
)
test_xformers_attention = False
batch_input_params = frozenset(["prompt", "negative_prompt", "image"])
output_shape = (9, 3, 32, 32)
# LTX is a video pipeline: it exposes `num_videos_per_prompt`, not the base default `num_images_per_prompt`.
optional_input_params = frozenset(
["num_inference_steps", "num_videos_per_prompt", "generator", "latents", "output_type", "return_dict"]
)
def get_dummy_components(self):
torch.manual_seed(0)
@@ -97,24 +101,19 @@ class LTXImageToVideoPipelineFastTests(PipelineTesterMixin, unittest.TestCase):
text_encoder = T5EncoderModel(config).eval()
tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5")
components = {
return {
"transformer": transformer,
"vae": vae,
"scheduler": scheduler,
"text_encoder": text_encoder,
"tokenizer": tokenizer,
}
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)
def get_dummy_inputs(self):
generator = self.get_generator(0)
image = torch.rand((1, 3, 32, 32), generator=generator)
image = torch.rand((1, 3, 32, 32), generator=generator, device=device)
inputs = {
return {
"image": image,
"prompt": "dance monkey",
"negative_prompt": "",
@@ -126,135 +125,20 @@ class LTXImageToVideoPipelineFastTests(PipelineTesterMixin, unittest.TestCase):
# 8 * k + 1 is the recommendation
"num_frames": 9,
"max_sequence_length": 16,
# Request torch outputs so tests compare torch tensors directly (see `BasePipelineTesterConfig`).
"output_type": "pt",
}
return inputs
def test_inference(self):
device = "cpu"
components = self.get_dummy_components()
pipe = self.pipeline_class(**components)
pipe.to(device)
pipe.set_progress_bar_config(disable=None)
inputs = self.get_dummy_inputs(device)
video = pipe(**inputs).frames
generated_video = video[0]
self.assertEqual(generated_video.shape, (9, 3, 32, 32))
expected_video = torch.randn(9, 3, 32, 32)
max_diff = torch.amax(torch.abs(generated_video - expected_video))
self.assertLessEqual(max_diff, 1e10)
def test_callback_inputs(self):
sig = inspect.signature(self.pipeline_class.__call__)
has_callback_tensor_inputs = "callback_on_step_end_tensor_inputs" in sig.parameters
has_callback_step_end = "callback_on_step_end" in sig.parameters
if not (has_callback_tensor_inputs and has_callback_step_end):
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",
)
def callback_inputs_subset(pipe, i, t, callback_kwargs):
# iterate over callback args
for tensor_name, tensor_value in callback_kwargs.items():
# check that we're only passing in allowed tensor inputs
assert tensor_name in pipe._callback_tensor_inputs
return callback_kwargs
def callback_inputs_all(pipe, i, t, callback_kwargs):
for tensor_name in pipe._callback_tensor_inputs:
assert tensor_name in callback_kwargs
# iterate over callback args
for tensor_name, tensor_value in callback_kwargs.items():
# check that we're only passing in allowed tensor inputs
assert tensor_name in pipe._callback_tensor_inputs
return callback_kwargs
inputs = self.get_dummy_inputs(torch_device)
# Test passing in a subset
inputs["callback_on_step_end"] = callback_inputs_subset
inputs["callback_on_step_end_tensor_inputs"] = ["latents"]
output = pipe(**inputs)[0]
# Test passing in a everything
inputs["callback_on_step_end"] = callback_inputs_all
inputs["callback_on_step_end_tensor_inputs"] = pipe._callback_tensor_inputs
output = pipe(**inputs)[0]
def callback_inputs_change_tensor(pipe, i, t, callback_kwargs):
is_last = i == (pipe.num_timesteps - 1)
if is_last:
callback_kwargs["latents"] = torch.zeros_like(callback_kwargs["latents"])
return callback_kwargs
inputs["callback_on_step_end"] = callback_inputs_change_tensor
inputs["callback_on_step_end_tensor_inputs"] = pipe._callback_tensor_inputs
output = pipe(**inputs)[0]
assert output.abs().sum() < 1e10
def test_inference_batch_single_identical(self):
self._test_inference_batch_single_identical(batch_size=3, expected_max_diff=1e-3)
def test_attention_slicing_forward_pass(
self, test_max_difference=True, test_mean_pixel_difference=True, expected_max_diff=1e-3
):
if not self.test_attention_slicing:
return
components = self.get_dummy_components()
pipe = self.pipeline_class(**components)
for component in pipe.components.values():
if hasattr(component, "set_default_attn_processor"):
component.set_default_attn_processor()
pipe.to(torch_device)
pipe.set_progress_bar_config(disable=None)
generator_device = "cpu"
inputs = self.get_dummy_inputs(generator_device)
output_without_slicing = pipe(**inputs)[0]
pipe.enable_attention_slicing(slice_size=1)
inputs = self.get_dummy_inputs(generator_device)
output_with_slicing1 = pipe(**inputs)[0]
pipe.enable_attention_slicing(slice_size=2)
inputs = self.get_dummy_inputs(generator_device)
output_with_slicing2 = pipe(**inputs)[0]
if test_max_difference:
max_diff1 = np.abs(to_np(output_with_slicing1) - to_np(output_without_slicing)).max()
max_diff2 = np.abs(to_np(output_with_slicing2) - to_np(output_without_slicing)).max()
self.assertLess(
max(max_diff1, max_diff2),
expected_max_diff,
"Attention slicing should not affect the inference results",
)
class TestLTXImageToVideoPipeline(LTXImageToVideoPipelineTesterConfig, PipelineTesterMixin):
def test_inference_batch_single_identical(self, batch_size=3, expected_max_diff=1e-3):
super().test_inference_batch_single_identical(batch_size=batch_size, expected_max_diff=expected_max_diff)
def test_vae_tiling(self, expected_diff_max: float = 0.2):
generator_device = "cpu"
components = self.get_dummy_components()
pipe = self.pipeline_class(**components)
pipe.to("cpu")
pipe.set_progress_bar_config(disable=None)
pipe = self.get_pipeline().to(torch_device)
# Without tiling
inputs = self.get_dummy_inputs(generator_device)
inputs = self.get_dummy_inputs()
inputs["height"] = inputs["width"] = 128
output_without_tiling = pipe(**inputs)[0]
@@ -265,12 +149,22 @@ class LTXImageToVideoPipelineFastTests(PipelineTesterMixin, unittest.TestCase):
tile_sample_stride_height=64,
tile_sample_stride_width=64,
)
inputs = self.get_dummy_inputs(generator_device)
inputs = self.get_dummy_inputs()
inputs["height"] = inputs["width"] = 128
output_with_tiling = pipe(**inputs)[0]
self.assertLess(
(to_np(output_without_tiling) - to_np(output_with_tiling)).max(),
expected_diff_max,
"VAE tiling should not affect the inference results",
assert (output_without_tiling - output_with_tiling).abs().max() < expected_diff_max, (
"VAE tiling should not affect the inference results."
)
class TestLTXImageToVideoPipelineMemory(LTXImageToVideoPipelineTesterConfig, MemoryTesterMixin):
"""Memory optimization tests (CPU offload, group offload, layerwise casting) for the LTX I2V pipeline."""
class TestLTXImageToVideoPipelineLoRA(LTXImageToVideoPipelineTesterConfig, LoraTesterMixin):
"""LoRA tests for the LTX I2V pipeline."""
class TestLTXImageToVideoPipelineLoRAMemory(LTXImageToVideoPipelineTesterConfig, LoraMemoryTesterMixin):
"""LoRA x memory-optimization tests (group offload, CPU offload) for the LTX I2V pipeline."""

92
tests/pipelines/ltx/test_ltx_latent_upsample.py
View File

@@ -12,27 +12,27 @@
# See the License for the specific language governing permissions and
# limitations under the License.
import unittest
import numpy as np
import pytest
import torch
from diffusers import AutoencoderKLLTXVideo, LTXLatentUpsamplePipeline
from diffusers.pipelines.ltx.modeling_latent_upsampler import LTXLatentUpsamplerModel
from ...testing_utils import enable_full_determinism
from ..test_pipelines_common import PipelineTesterMixin, to_np
from ...testing_utils import enable_full_determinism, torch_device
from ..testing_utils import BasePipelineTesterConfig, MemoryTesterMixin, PipelineTesterMixin
enable_full_determinism()
class LTXLatentUpsamplePipelineFastTests(PipelineTesterMixin, unittest.TestCase):
class LTXLatentUpsamplePipelineTesterConfig(BasePipelineTesterConfig):
pipeline_class = LTXLatentUpsamplePipeline
params = {"video", "generator"}
batch_params = {"video", "generator"}
required_optional_params = frozenset(["generator", "latents", "return_dict"])
test_xformers_attention = False
required_input_params_in_call_signature = frozenset(["video", "height", "width", "latents"])
batch_input_params = frozenset(["video"])
output_shape = (5, 3, 32, 32)
# This pipeline takes a video rather than a prompt, so it has neither `num_images_per_prompt` nor
# `num_inference_steps` — upsampling is a single forward pass through the latent upsampler.
optional_input_params = frozenset(["generator", "latents", "output_type", "return_dict"])
def get_dummy_components(self):
torch.manual_seed(0)
@@ -68,57 +68,31 @@ class LTXLatentUpsamplePipelineFastTests(PipelineTesterMixin, unittest.TestCase)
temporal_upsample=False,
)
components = {
return {
"vae": vae,
"latent_upsampler": latent_upsampler,
}
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)
def get_dummy_inputs(self):
generator = self.get_generator(0)
video = torch.randn((5, 3, 32, 32), generator=generator)
video = torch.randn((5, 3, 32, 32), generator=generator, device=device)
inputs = {
return {
"video": video,
"generator": generator,
"height": 16,
"width": 16,
# Request torch outputs so tests compare torch tensors directly (see `BasePipelineTesterConfig`).
"output_type": "pt",
}
return inputs
def test_inference(self):
device = "cpu"
components = self.get_dummy_components()
pipe = self.pipeline_class(**components)
pipe.to(device)
pipe.set_progress_bar_config(disable=None)
inputs = self.get_dummy_inputs(device)
video = pipe(**inputs).frames
generated_video = video[0]
self.assertEqual(generated_video.shape, (5, 3, 32, 32))
expected_video = torch.randn(5, 3, 32, 32)
max_diff = np.abs(generated_video - expected_video).max()
self.assertLessEqual(max_diff, 1e10)
class TestLTXLatentUpsamplePipeline(LTXLatentUpsamplePipelineTesterConfig, PipelineTesterMixin):
def test_vae_tiling(self, expected_diff_max: float = 0.25):
generator_device = "cpu"
components = self.get_dummy_components()
pipe = self.pipeline_class(**components)
pipe.to("cpu")
pipe.set_progress_bar_config(disable=None)
pipe = self.get_pipeline().to(torch_device)
# Without tiling
inputs = self.get_dummy_inputs(generator_device)
inputs = self.get_dummy_inputs()
inputs["height"] = inputs["width"] = 128
output_without_tiling = pipe(**inputs)[0]
@@ -129,30 +103,24 @@ class LTXLatentUpsamplePipelineFastTests(PipelineTesterMixin, unittest.TestCase)
tile_sample_stride_height=64,
tile_sample_stride_width=64,
)
inputs = self.get_dummy_inputs(generator_device)
inputs = self.get_dummy_inputs()
inputs["height"] = inputs["width"] = 128
output_with_tiling = pipe(**inputs)[0]
self.assertLess(
(to_np(output_without_tiling) - to_np(output_with_tiling)).max(),
expected_diff_max,
"VAE tiling should not affect the inference results",
assert (output_without_tiling - output_with_tiling).abs().max() < expected_diff_max, (
"VAE tiling should not affect the inference results."
)
@unittest.skip("Test is not applicable.")
def test_callback_inputs(self):
pass
@unittest.skip("Test is not applicable.")
def test_attention_slicing_forward_pass(
self, test_max_difference=True, test_mean_pixel_difference=True, expected_max_diff=1e-3
):
pass
@unittest.skip("Test is not applicable.")
# `__call__` documents batched video input as unsupported (`batch_size` is pinned to 1), so the batching
# tests below have nothing to assert against.
@pytest.mark.skip("Batched video input is not supported by this pipeline.")
def test_inference_batch_consistent(self):
pass
@unittest.skip("Test is not applicable.")
@pytest.mark.skip("Batched video input is not supported by this pipeline.")
def test_inference_batch_single_identical(self):
pass
class TestLTXLatentUpsamplePipelineMemory(LTXLatentUpsamplePipelineTesterConfig, MemoryTesterMixin):
"""Memory optimization tests (CPU offload, group offload, layerwise casting) for the LTX upsampler pipeline."""

199
tests/pipelines/ltx2/test_ltx2.py
View File

@@ -14,170 +14,28 @@
import pytest
import torch
from transformers import AutoTokenizer, Gemma3ForConditionalGeneration
from diffusers import (
AutoencoderKLLTX2Audio,
AutoencoderKLLTX2Video,
FlowMatchEulerDiscreteScheduler,
LTX2Pipeline,
LTX2VideoTransformer3DModel,
)
from diffusers.pipelines.ltx2 import LTX2DurationHead, LTX2TextConnectors
from diffusers.pipelines.ltx2.vocoder import LTX2Vocoder
from diffusers import LTX2Pipeline
from ...testing_utils import assert_tensors_close, enable_full_determinism, require_torch_accelerator, torch_device
from ..testing_utils import (
BasePipelineTesterConfig,
LoraMemoryTesterMixin,
LoraTesterMixin,
MemoryTesterMixin,
PipelineTesterMixin,
from ...testing_utils import assert_tensors_close, enable_full_determinism, torch_device
from ..testing_utils import PipelineTesterMixin
from .testing_utils import (
LTX2BaseTesterConfig,
LTX2LoraMemoryTesterMixin,
LTX2LoraTesterMixin,
LTX2MemoryTesterMixin,
)
enable_full_determinism()
class LTX2PipelineTesterConfig(BasePipelineTesterConfig):
class LTX2PipelineTesterConfig(LTX2BaseTesterConfig):
pipeline_class = LTX2Pipeline
required_input_params_in_call_signature = frozenset(
["prompt", "height", "width", "guidance_scale", "negative_prompt", "prompt_embeds", "negative_prompt_embeds"]
)
batch_input_params = frozenset(["prompt", "negative_prompt"])
output_shape = (5, 3, 32, 32)
# LTX2 is a video pipeline (`num_videos_per_prompt`, not `num_images_per_prompt`) and takes a second latent
# input for the audio stream.
optional_input_params = frozenset(
[
"num_inference_steps",
"num_videos_per_prompt",
"generator",
"latents",
"audio_latents",
"output_type",
"return_dict",
]
)
base_text_encoder_ckpt_id = "hf-internal-testing/tiny-gemma3"
def get_dummy_components(self):
tokenizer = AutoTokenizer.from_pretrained(self.base_text_encoder_ckpt_id)
text_encoder = Gemma3ForConditionalGeneration.from_pretrained(self.base_text_encoder_ckpt_id)
torch.manual_seed(0)
transformer = LTX2VideoTransformer3DModel(
in_channels=4,
out_channels=4,
patch_size=1,
patch_size_t=1,
num_attention_heads=2,
attention_head_dim=8,
cross_attention_dim=16,
audio_in_channels=4,
audio_out_channels=4,
audio_num_attention_heads=2,
audio_attention_head_dim=4,
audio_cross_attention_dim=8,
num_layers=2,
qk_norm="rms_norm_across_heads",
caption_channels=text_encoder.config.text_config.hidden_size,
rope_double_precision=False,
rope_type="split",
)
torch.manual_seed(0)
connectors = LTX2TextConnectors(
caption_channels=text_encoder.config.text_config.hidden_size,
text_proj_in_factor=text_encoder.config.text_config.num_hidden_layers + 1,
video_connector_num_attention_heads=4,
video_connector_attention_head_dim=8,
video_connector_num_layers=1,
video_connector_num_learnable_registers=None,
audio_connector_num_attention_heads=4,
audio_connector_attention_head_dim=8,
audio_connector_num_layers=1,
audio_connector_num_learnable_registers=None,
connector_rope_base_seq_len=32,
rope_theta=10000.0,
rope_double_precision=False,
causal_temporal_positioning=False,
rope_type="split",
)
torch.manual_seed(0)
vae = AutoencoderKLLTX2Video(
in_channels=3,
out_channels=3,
latent_channels=4,
block_out_channels=(8,),
decoder_block_out_channels=(8,),
layers_per_block=(1,),
decoder_layers_per_block=(1, 1),
spatio_temporal_scaling=(True,),
decoder_spatio_temporal_scaling=(True,),
decoder_inject_noise=(False, False),
downsample_type=("spatial",),
upsample_residual=(False,),
upsample_factor=(1,),
timestep_conditioning=False,
patch_size=1,
patch_size_t=1,
encoder_causal=True,
decoder_causal=False,
)
vae.use_framewise_encoding = False
vae.use_framewise_decoding = False
torch.manual_seed(0)
audio_vae = AutoencoderKLLTX2Audio(
base_channels=4,
output_channels=2,
ch_mult=(1,),
num_res_blocks=1,
attn_resolutions=None,
in_channels=2,
resolution=32,
latent_channels=2,
norm_type="pixel",
causality_axis="height",
dropout=0.0,
mid_block_add_attention=False,
sample_rate=16000,
mel_hop_length=160,
is_causal=True,
mel_bins=8,
)
torch.manual_seed(0)
vocoder = LTX2Vocoder(
in_channels=audio_vae.config.output_channels * audio_vae.config.mel_bins,
hidden_channels=32,
out_channels=2,
upsample_kernel_sizes=[4, 4],
upsample_factors=[2, 2],
resnet_kernel_sizes=[3],
resnet_dilations=[[1, 3, 5]],
leaky_relu_negative_slope=0.1,
output_sampling_rate=16000,
)
scheduler = FlowMatchEulerDiscreteScheduler()
return {
"transformer": transformer,
"vae": vae,
"audio_vae": audio_vae,
"scheduler": scheduler,
"text_encoder": text_encoder,
"tokenizer": tokenizer,
"connectors": connectors,
"vocoder": vocoder,
"processor": None,
"prompt_enhancer": None,
"duration_head": None,
}
def get_dummy_inputs(self):
return {
@@ -206,23 +64,6 @@ class LTX2PipelineTesterConfig(BasePipelineTesterConfig):
"output_type": "pt",
}
def get_dummy_duration_head(self):
torch.manual_seed(0)
# The dummy connectors emit 4 heads * 8 head_dim = 32 wide output for both streams.
return LTX2DurationHead(
video_cross_attention_dim=32,
audio_cross_attention_dim=32,
pooler_hidden_dim=8,
num_queries=1,
num_pooler_heads=2,
mlp_hidden_dim=8,
)
def get_pipeline_with_duration_head(self):
components = self.get_dummy_components()
components["duration_head"] = self.get_dummy_duration_head()
return self.get_pipeline(**components).to(torch_device)
class TestLTX2Pipeline(LTX2PipelineTesterConfig, PipelineTesterMixin):
def test_inference(self):
@@ -404,29 +245,13 @@ class TestLTX2Pipeline(LTX2PipelineTesterConfig, PipelineTesterMixin):
pipe(**inputs)
class TestLTX2PipelineMemory(LTX2PipelineTesterConfig, MemoryTesterMixin):
class TestLTX2PipelineMemory(LTX2PipelineTesterConfig, LTX2MemoryTesterMixin):
"""Memory optimization tests (CPU offload, group offload, layerwise casting) for the LTX2 pipeline."""
@require_torch_accelerator
def test_group_offloading_inference(self):
# The shared helper only offloads a fixed set of component names and leaves LTX2's extra module
# components (`connectors`, `audio_vae`, `vocoder`) on CPU, so the forward pass mixes devices.
# Pipeline-level offloading, which walks every component, is exercised by
# `test_pipeline_level_group_offloading_inference`.
pytest.skip("Using test_pipeline_level_group_offloading_inference instead")
class TestLTX2PipelineLoRA(LTX2PipelineTesterConfig, LoraTesterMixin):
class TestLTX2PipelineLoRA(LTX2PipelineTesterConfig, LTX2LoraTesterMixin):
"""LoRA tests for the LTX2 pipeline."""
# `LTX2Pipeline` advertises `connectors` as LoRA-loadable and `load_lora_weights` does handle connector
# LoRAs, but `save_lora_weights` only accepts `transformer_lora_layers` — so connector adapters cannot
# round-trip through the public API these tests drive. Scope the tests to the transformer until that closes.
lora_loadable_components = ["transformer"]
class TestLTX2PipelineLoRAMemory(LTX2PipelineTesterConfig, LoraMemoryTesterMixin):
class TestLTX2PipelineLoRAMemory(LTX2PipelineTesterConfig, LTX2LoraMemoryTesterMixin):
"""LoRA x memory-optimization tests (group offload, CPU offload) for the LTX2 pipeline."""
# See `TestLTX2PipelineLoRA`.
lora_loadable_components = ["transformer"]

212
tests/pipelines/ltx2/test_ltx2_condition.py
View File

@@ -12,193 +12,48 @@
# See the License for the specific language governing permissions and
# limitations under the License.
import unittest
import torch
from transformers import AutoTokenizer, Gemma3ForConditionalGeneration
from diffusers import (
AutoencoderKLLTX2Audio,
AutoencoderKLLTX2Video,
FlowMatchEulerDiscreteScheduler,
LTX2ConditionPipeline,
LTX2VideoTransformer3DModel,
)
from diffusers.pipelines.ltx2 import LTX2TextConnectors
from diffusers.pipelines.ltx2.latent_upsampler import LTX2LatentUpsamplerModel
from diffusers import LTX2ConditionPipeline
from diffusers.pipelines.ltx2.pipeline_ltx2_condition import LTX2VideoCondition
from diffusers.pipelines.ltx2.vocoder import LTX2Vocoder
from ...testing_utils import enable_full_determinism
from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import PipelineTesterMixin
from ..testing_utils import PipelineTesterMixin
from .testing_utils import (
LTX2BaseTesterConfig,
LTX2LoraMemoryTesterMixin,
LTX2LoraTesterMixin,
LTX2MemoryTesterMixin,
)
enable_full_determinism()
class LTX2ConditionPipelineFastTests(PipelineTesterMixin, unittest.TestCase):
class LTX2ConditionPipelineTesterConfig(LTX2BaseTesterConfig):
pipeline_class = LTX2ConditionPipeline
params = TEXT_TO_IMAGE_PARAMS - {"cross_attention_kwargs"}
batch_params = TEXT_TO_IMAGE_BATCH_PARAMS
image_params = TEXT_TO_IMAGE_IMAGE_PARAMS
image_latents_params = TEXT_TO_IMAGE_IMAGE_PARAMS
required_optional_params = frozenset(
required_input_params_in_call_signature = frozenset(
[
"num_inference_steps",
"generator",
"latents",
"audio_latents",
"return_dict",
"callback_on_step_end",
"callback_on_step_end_tensor_inputs",
"conditions",
"prompt",
"height",
"width",
"guidance_scale",
"negative_prompt",
"prompt_embeds",
"negative_prompt_embeds",
]
)
test_attention_slicing = False
test_xformers_attention = False
batch_input_params = frozenset(["prompt", "negative_prompt"])
unset_components = ("audio_scheduler", "processor", "prompt_enhancer", "duration_head")
base_text_encoder_ckpt_id = "hf-internal-testing/tiny-gemma3"
def get_dummy_components(self):
tokenizer = AutoTokenizer.from_pretrained(self.base_text_encoder_ckpt_id)
text_encoder = Gemma3ForConditionalGeneration.from_pretrained(self.base_text_encoder_ckpt_id)
torch.manual_seed(0)
transformer = LTX2VideoTransformer3DModel(
in_channels=4,
out_channels=4,
patch_size=1,
patch_size_t=1,
num_attention_heads=2,
attention_head_dim=8,
cross_attention_dim=16,
audio_in_channels=4,
audio_out_channels=4,
audio_num_attention_heads=2,
audio_attention_head_dim=4,
audio_cross_attention_dim=8,
num_layers=2,
qk_norm="rms_norm_across_heads",
caption_channels=text_encoder.config.text_config.hidden_size,
rope_double_precision=False,
rope_type="split",
)
torch.manual_seed(0)
connectors = LTX2TextConnectors(
caption_channels=text_encoder.config.text_config.hidden_size,
text_proj_in_factor=text_encoder.config.text_config.num_hidden_layers + 1,
video_connector_num_attention_heads=4,
video_connector_attention_head_dim=8,
video_connector_num_layers=1,
video_connector_num_learnable_registers=None,
audio_connector_num_attention_heads=4,
audio_connector_attention_head_dim=8,
audio_connector_num_layers=1,
audio_connector_num_learnable_registers=None,
connector_rope_base_seq_len=32,
rope_theta=10000.0,
rope_double_precision=False,
causal_temporal_positioning=False,
rope_type="split",
)
torch.manual_seed(0)
vae = AutoencoderKLLTX2Video(
in_channels=3,
out_channels=3,
latent_channels=4,
block_out_channels=(8,),
decoder_block_out_channels=(8,),
layers_per_block=(1,),
decoder_layers_per_block=(1, 1),
spatio_temporal_scaling=(True,),
decoder_spatio_temporal_scaling=(True,),
decoder_inject_noise=(False, False),
downsample_type=("spatial",),
upsample_residual=(False,),
upsample_factor=(1,),
timestep_conditioning=False,
patch_size=1,
patch_size_t=1,
encoder_causal=True,
decoder_causal=False,
)
vae.use_framewise_encoding = False
vae.use_framewise_decoding = False
torch.manual_seed(0)
audio_vae = AutoencoderKLLTX2Audio(
base_channels=4,
output_channels=2,
ch_mult=(1,),
num_res_blocks=1,
attn_resolutions=None,
in_channels=2,
resolution=32,
latent_channels=2,
norm_type="pixel",
causality_axis="height",
dropout=0.0,
mid_block_add_attention=False,
sample_rate=16000,
mel_hop_length=160,
is_causal=True,
mel_bins=8,
)
torch.manual_seed(0)
vocoder = LTX2Vocoder(
in_channels=audio_vae.config.output_channels * audio_vae.config.mel_bins,
hidden_channels=32,
out_channels=2,
upsample_kernel_sizes=[4, 4],
upsample_factors=[2, 2],
resnet_kernel_sizes=[3],
resnet_dilations=[[1, 3, 5]],
leaky_relu_negative_slope=0.1,
output_sampling_rate=16000,
)
scheduler = FlowMatchEulerDiscreteScheduler()
components = {
"transformer": transformer,
"vae": vae,
"audio_vae": audio_vae,
"scheduler": scheduler,
"text_encoder": text_encoder,
"tokenizer": tokenizer,
"connectors": connectors,
"vocoder": vocoder,
"audio_scheduler": None,
"processor": None,
"prompt_enhancer": None,
"duration_head": None,
}
return components
def get_dummy_upsample_component(self, in_channels=4, mid_channels=32, num_blocks_per_stage=1):
upsampler = LTX2LatentUpsamplerModel(
in_channels=in_channels,
mid_channels=mid_channels,
num_blocks_per_stage=num_blocks_per_stage,
)
return upsampler
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)
image = torch.rand((1, 3, 32, 32), generator=generator, device=device)
def get_dummy_inputs(self):
generator = self.get_generator(0)
image = torch.rand((1, 3, 32, 32), generator=generator)
# Synthetic float tensors skip H.264 CRF re-compression (training path uses PIL/uint8).
img_cond = LTX2VideoCondition(frames=image, index=0, strength=1.0, crf=0)
inputs = {
return {
"conditions": img_cond,
"prompt": "a robot dancing",
"negative_prompt": "",
@@ -221,10 +76,23 @@ class LTX2ConditionPipelineFastTests(PipelineTesterMixin, unittest.TestCase):
"num_frames": 5,
"frame_rate": 25.0,
"max_sequence_length": 16,
# Request torch outputs so tests compare torch tensors directly (see `BasePipelineTesterConfig`).
"output_type": "pt",
}
return inputs
def test_inference_batch_single_identical(self):
self._test_inference_batch_single_identical(batch_size=2, expected_max_diff=1e-3)
class TestLTX2ConditionPipeline(LTX2ConditionPipelineTesterConfig, PipelineTesterMixin):
def test_inference_batch_single_identical(self, batch_size=2, expected_max_diff=1e-3):
super().test_inference_batch_single_identical(batch_size=batch_size, expected_max_diff=expected_max_diff)
class TestLTX2ConditionPipelineMemory(LTX2ConditionPipelineTesterConfig, LTX2MemoryTesterMixin):
"""Memory optimization tests (CPU offload, group offload, layerwise casting) for the LTX2 condition pipeline."""
class TestLTX2ConditionPipelineLoRA(LTX2ConditionPipelineTesterConfig, LTX2LoraTesterMixin):
"""LoRA tests for the LTX2 condition pipeline."""
class TestLTX2ConditionPipelineLoRAMemory(LTX2ConditionPipelineTesterConfig, LTX2LoraMemoryTesterMixin):
"""LoRA x memory-optimization tests (group offload, CPU offload) for the LTX2 condition pipeline."""

8
tests/pipelines/ltx2/test_ltx2_connectors.py
View File

@@ -12,19 +12,17 @@
# See the License for the specific language governing permissions and
# limitations under the License.
import unittest
import torch
from diffusers.pipelines.ltx2.connectors import LTX2ConnectorTransformer1d
from ...testing_utils import enable_full_determinism
from ...testing_utils import assert_tensors_close, enable_full_determinism
enable_full_determinism()
class LTX2ConnectorRegisterLayoutTests(unittest.TestCase):
class TestLTX2ConnectorRegisterLayout:
"""The connector must lay out its sequence exactly like the original LTX
implementation (``ltx_core`` ``_replace_padded_with_learnable_registers``,
also matched by ComfyUI): the valid tokens move to the front *in their
@@ -82,7 +80,7 @@ class LTX2ConnectorRegisterLayoutTests(unittest.TestCase):
with torch.no_grad():
output, _ = connector(hidden_states, additive_mask)
expected = self.reference_layout(connector, hidden_states, binary_mask)
self.assertTrue(torch.allclose(output, expected, atol=1e-5))
assert_tensors_close(output, expected, atol=1e-5)
def test_register_layout_left_padded(self):
self.check_layout([5])

365
tests/pipelines/ltx2/test_ltx2_hdr.py
View File

@@ -12,191 +12,57 @@
# See the License for the specific language governing permissions and
# limitations under the License.
import tempfile
import unittest
from contextlib import contextmanager
import numpy as np
import torch
from transformers import AutoTokenizer, Gemma3ForConditionalGeneration
import diffusers
from diffusers import (
AutoencoderKLLTX2Audio,
AutoencoderKLLTX2Video,
FlowMatchEulerDiscreteScheduler,
LTX2HDRPipeline,
LTX2VideoTransformer3DModel,
from diffusers import LTX2HDRPipeline
from diffusers.pipelines.ltx2 import LTX2HDRReferenceCondition
from ...testing_utils import enable_full_determinism
from ..testing_utils import PipelineTesterMixin
from .testing_utils import (
LTX2BaseTesterConfig,
LTX2LoraMemoryTesterMixin,
LTX2LoraTesterMixin,
LTX2MemoryTesterMixin,
)
from diffusers.pipelines.ltx2 import LTX2HDRReferenceCondition, LTX2TextConnectors
from diffusers.pipelines.ltx2.latent_upsampler import LTX2LatentUpsamplerModel
from diffusers.pipelines.ltx2.vocoder import LTX2Vocoder
from diffusers.utils import logging
from ...testing_utils import enable_full_determinism, require_accelerator, 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 PipelineTesterMixin, to_np
enable_full_determinism()
class LTX2HDRPipelineFastTests(PipelineTesterMixin, unittest.TestCase):
class LTX2HDRPipelineTesterConfig(LTX2BaseTesterConfig):
pipeline_class = LTX2HDRPipeline
params = TEXT_TO_IMAGE_PARAMS - {"cross_attention_kwargs"}
batch_params = TEXT_TO_IMAGE_BATCH_PARAMS
image_params = TEXT_TO_IMAGE_IMAGE_PARAMS
image_latents_params = TEXT_TO_IMAGE_IMAGE_PARAMS
required_optional_params = frozenset(
required_input_params_in_call_signature = frozenset(
[
"num_inference_steps",
"generator",
"latents",
"return_dict",
"callback_on_step_end",
"callback_on_step_end_tensor_inputs",
"reference_conditions",
"prompt",
"height",
"width",
"guidance_scale",
"negative_prompt",
"prompt_embeds",
"negative_prompt_embeds",
]
)
test_attention_slicing = False
test_xformers_attention = False
batch_input_params = frozenset(["prompt", "negative_prompt"])
# `postprocess_hdr_video` permutes to channels-last for both `"pt"` and `"np"`, so this pipeline's frames come
# back as (num_frames, height, width, channels) rather than the channels-first layout `"pt"` usually implies.
output_shape = (5, 32, 32, 3)
# LTX2 is a video pipeline: it exposes `num_videos_per_prompt`, not the base default `num_images_per_prompt`.
# Unlike the other LTX2 pipelines this one renders video only, so there is no `audio_latents`.
optional_input_params = frozenset(
["num_inference_steps", "num_videos_per_prompt", "generator", "latents", "output_type", "return_dict"]
)
unset_components = ("audio_scheduler",)
base_text_encoder_ckpt_id = "hf-internal-testing/tiny-gemma3"
def get_dummy_components(self):
tokenizer = AutoTokenizer.from_pretrained(self.base_text_encoder_ckpt_id)
text_encoder = Gemma3ForConditionalGeneration.from_pretrained(self.base_text_encoder_ckpt_id)
torch.manual_seed(0)
transformer = LTX2VideoTransformer3DModel(
in_channels=4,
out_channels=4,
patch_size=1,
patch_size_t=1,
num_attention_heads=2,
attention_head_dim=8,
cross_attention_dim=16,
audio_in_channels=4,
audio_out_channels=4,
audio_num_attention_heads=2,
audio_attention_head_dim=4,
audio_cross_attention_dim=8,
num_layers=2,
qk_norm="rms_norm_across_heads",
caption_channels=text_encoder.config.text_config.hidden_size,
rope_double_precision=False,
rope_type="split",
)
torch.manual_seed(0)
connectors = LTX2TextConnectors(
caption_channels=text_encoder.config.text_config.hidden_size,
text_proj_in_factor=text_encoder.config.text_config.num_hidden_layers + 1,
video_connector_num_attention_heads=4,
video_connector_attention_head_dim=8,
video_connector_num_layers=1,
video_connector_num_learnable_registers=None,
audio_connector_num_attention_heads=4,
audio_connector_attention_head_dim=8,
audio_connector_num_layers=1,
audio_connector_num_learnable_registers=None,
connector_rope_base_seq_len=32,
rope_theta=10000.0,
rope_double_precision=False,
causal_temporal_positioning=False,
rope_type="split",
)
torch.manual_seed(0)
vae = AutoencoderKLLTX2Video(
in_channels=3,
out_channels=3,
latent_channels=4,
block_out_channels=(8,),
decoder_block_out_channels=(8,),
layers_per_block=(1,),
decoder_layers_per_block=(1, 1),
spatio_temporal_scaling=(True,),
decoder_spatio_temporal_scaling=(True,),
decoder_inject_noise=(False, False),
downsample_type=("spatial",),
upsample_residual=(False,),
upsample_factor=(1,),
timestep_conditioning=False,
patch_size=1,
patch_size_t=1,
encoder_causal=True,
decoder_causal=False,
)
vae.use_framewise_encoding = False
vae.use_framewise_decoding = False
torch.manual_seed(0)
audio_vae = AutoencoderKLLTX2Audio(
base_channels=4,
output_channels=2,
ch_mult=(1,),
num_res_blocks=1,
attn_resolutions=None,
in_channels=2,
resolution=32,
latent_channels=2,
norm_type="pixel",
causality_axis="height",
dropout=0.0,
mid_block_add_attention=False,
sample_rate=16000,
mel_hop_length=160,
is_causal=True,
mel_bins=8,
)
torch.manual_seed(0)
vocoder = LTX2Vocoder(
in_channels=audio_vae.config.output_channels * audio_vae.config.mel_bins,
hidden_channels=32,
out_channels=2,
upsample_kernel_sizes=[4, 4],
upsample_factors=[2, 2],
resnet_kernel_sizes=[3],
resnet_dilations=[[1, 3, 5]],
leaky_relu_negative_slope=0.1,
output_sampling_rate=16000,
)
scheduler = FlowMatchEulerDiscreteScheduler()
components = {
"transformer": transformer,
"vae": vae,
"audio_vae": audio_vae,
"scheduler": scheduler,
"text_encoder": text_encoder,
"tokenizer": tokenizer,
"connectors": connectors,
"vocoder": vocoder,
"audio_scheduler": None,
}
return components
def get_dummy_upsample_component(self, in_channels=4, mid_channels=32, num_blocks_per_stage=1):
upsampler = LTX2LatentUpsamplerModel(
in_channels=in_channels,
mid_channels=mid_channels,
num_blocks_per_stage=num_blocks_per_stage,
)
return upsampler
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)
image = torch.rand((1, 3, 32, 32), generator=generator, device=device)
def get_dummy_inputs(self):
generator = self.get_generator(0)
image = torch.rand((1, 3, 32, 32), generator=generator)
img_cond = LTX2HDRReferenceCondition(frames=image, strength=1.0)
inputs = {
return {
"reference_conditions": img_cond,
"prompt": "a robot dancing",
"negative_prompt": "",
@@ -208,145 +74,48 @@ class LTX2HDRPipelineFastTests(PipelineTesterMixin, unittest.TestCase):
"num_frames": 5,
"frame_rate": 25.0,
"max_sequence_length": 16,
# Request torch outputs so tests compare torch tensors directly (see `BasePipelineTesterConfig`).
"output_type": "pt",
}
class TestLTX2HDRPipeline(LTX2HDRPipelineTesterConfig, PipelineTesterMixin):
# The HDR video processor applies the inverse LogC3 transfer function, whose exponential blows tiny numerical
# differences up into large pixel ones. Tests that compare two decoded runs against each other therefore
# compare latents instead: `latent_outputs()` flips `get_dummy_inputs` over to `output_type="latent"` for the
# duration of the base implementation.
_force_latent_output = False
def get_dummy_inputs(self):
inputs = super().get_dummy_inputs()
if self._force_latent_output:
inputs["output_type"] = "latent"
return inputs
# Override to set the dummy inputs `output_type` to "latent" for this test, as the HDR video processor appears to
# amplify small numerical differences due to applying the exponential inverse LogC3 inverse transfer function
def test_inference_batch_single_identical(
self,
batch_size=2,
expected_max_diff=1e-4,
additional_params_copy_to_batched_inputs=["num_inference_steps"],
):
components = self.get_dummy_components()
for key in components:
if "text_encoder" in key and hasattr(components[key], "eval"):
components[key].eval()
pipe = self.pipeline_class(**components)
for components in pipe.components.values():
if hasattr(components, "set_default_attn_processor"):
components.set_default_attn_processor()
@contextmanager
def latent_outputs(self):
self._force_latent_output = True
try:
yield
finally:
self._force_latent_output = False
pipe.to(torch_device)
pipe.set_progress_bar_config(disable=None)
inputs = self.get_dummy_inputs(torch_device)
# NOTE: explicitly set output_type="latent" for this test to avoid postprocessor issues
inputs["output_type"] = "latent"
# Reset generator in case it is has been used in self.get_dummy_inputs
inputs["generator"] = self.get_generator(0)
def test_inference_batch_single_identical(self, batch_size=2, expected_max_diff=1e-4):
with self.latent_outputs():
super().test_inference_batch_single_identical(batch_size=batch_size, expected_max_diff=expected_max_diff)
logger = logging.get_logger(pipe.__module__)
logger.setLevel(level=diffusers.logging.FATAL)
def test_save_load_float16(self, tmp_path, expected_max_diff=1e-2):
with self.latent_outputs():
super().test_save_load_float16(tmp_path, expected_max_diff=expected_max_diff)
# batchify inputs
batched_inputs = {}
batched_inputs.update(inputs)
for name in self.batch_params:
if name not in inputs:
continue
class TestLTX2HDRPipelineMemory(LTX2HDRPipelineTesterConfig, LTX2MemoryTesterMixin):
"""Memory optimization tests (CPU offload, group offload, layerwise casting) for the LTX2 HDR pipeline."""
value = inputs[name]
if name == "prompt":
print(f"prompt value type: {type(value)}")
len_prompt = len(value)
batched_inputs[name] = [value[: len_prompt // i] for i in range(1, batch_size + 1)]
batched_inputs[name][-1] = 100 * "very long"
else:
batched_inputs[name] = batch_size * [value]
print(f"Prompt input: {inputs['prompt']}")
print(f"Prompt batched input {batched_inputs['prompt']}")
print(f"Batch size: {batch_size}")
class TestLTX2HDRPipelineLoRA(LTX2HDRPipelineTesterConfig, LTX2LoraTesterMixin):
"""LoRA tests for the LTX2 HDR pipeline."""
if "generator" in inputs:
batched_inputs["generator"] = [self.get_generator(i) for i in range(batch_size)]
if "batch_size" in inputs:
batched_inputs["batch_size"] = batch_size
for arg in additional_params_copy_to_batched_inputs:
batched_inputs[arg] = inputs[arg]
output = pipe(**inputs)
output_batch = pipe(**batched_inputs)
assert output_batch[0].shape[0] == batch_size
max_diff = np.abs(to_np(output_batch[0][0]) - to_np(output[0][0])).max()
assert max_diff < expected_max_diff
# Override to set the dummy inputs `output_type` to "latent" for this test, as the HDR video processor appears to
# amplify small numerical differences due to applying the exponential inverse LogC3 inverse transfer function
@unittest.skipIf(torch_device not in ["cuda", "xpu"], reason="float16 requires CUDA or XPU")
@require_accelerator
def test_save_load_float16(self, expected_max_diff=1e-2):
components = self.get_dummy_components()
for name, module in components.items():
# Account for components with _keep_in_fp32_modules
if hasattr(module, "_keep_in_fp32_modules") and module._keep_in_fp32_modules is not None:
for name, param in module.named_parameters():
if any(
module_to_keep_in_fp32 in name.split(".")
for module_to_keep_in_fp32 in module._keep_in_fp32_modules
):
param.data = param.data.to(torch_device).to(torch.float32)
else:
param.data = param.data.to(torch_device).to(torch.float16)
for name, buf in module.named_buffers():
if not buf.is_floating_point():
buf.data = buf.data.to(torch_device)
elif any(
module_to_keep_in_fp32 in name.split(".")
for module_to_keep_in_fp32 in module._keep_in_fp32_modules
):
buf.data = buf.data.to(torch_device).to(torch.float32)
else:
buf.data = buf.data.to(torch_device).to(torch.float16)
elif hasattr(module, "half"):
components[name] = module.to(torch_device).half()
for key, component in components.items():
if hasattr(component, "eval"):
component.eval()
pipe = self.pipeline_class(**components)
for component in pipe.components.values():
if hasattr(component, "set_default_attn_processor"):
component.set_default_attn_processor()
pipe.to(torch_device)
pipe.set_progress_bar_config(disable=None)
inputs = self.get_dummy_inputs(torch_device)
# NOTE: explicitly set output_type="latent" for this test to avoid postprocessor issues
inputs["output_type"] = "latent"
output = pipe(**inputs)[0]
with tempfile.TemporaryDirectory() as tmpdir:
pipe.save_pretrained(tmpdir)
pipe_loaded = self.pipeline_class.from_pretrained(tmpdir, torch_dtype=torch.float16)
for component in pipe_loaded.components.values():
if hasattr(component, "set_default_attn_processor"):
component.set_default_attn_processor()
pipe_loaded.to(torch_device)
pipe_loaded.set_progress_bar_config(disable=None)
for name, component in pipe_loaded.components.items():
if hasattr(component, "dtype"):
self.assertTrue(
component.dtype == torch.float16,
f"`{name}.dtype` switched from `float16` to {component.dtype} after loading.",
)
inputs = self.get_dummy_inputs(torch_device)
# NOTE: explicitly set output_type="latent" for this test to avoid postprocessor issues
inputs["output_type"] = "latent"
output_loaded = pipe_loaded(**inputs)[0]
max_diff = np.abs(to_np(output) - to_np(output_loaded)).max()
self.assertLess(
max_diff, expected_max_diff, "The output of the fp16 pipeline changed after saving and loading."
)
class TestLTX2HDRPipelineLoRAMemory(LTX2HDRPipelineTesterConfig, LTX2LoraMemoryTesterMixin):
"""LoRA x memory-optimization tests (group offload, CPU offload) for the LTX2 HDR pipeline."""

342
tests/pipelines/ltx2/test_ltx2_image2video.py
View File

@@ -12,193 +12,54 @@
# See the License for the specific language governing permissions and
# limitations under the License.
import unittest
import pytest
import torch
from transformers import AutoTokenizer, Gemma3ForConditionalGeneration
from diffusers import (
AutoencoderKLLTX2Audio,
AutoencoderKLLTX2Video,
FlowMatchEulerDiscreteScheduler,
LTX2ImageToVideoPipeline,
LTX2VideoTransformer3DModel,
)
from diffusers.pipelines.ltx2 import (
LTX2DurationHead,
LTX2LatentUpsamplePipeline,
LTX2TextConnectors,
)
from diffusers import LTX2ImageToVideoPipeline
from diffusers.pipelines.ltx2 import LTX2LatentUpsamplePipeline
from diffusers.pipelines.ltx2.latent_upsampler import LTX2LatentUpsamplerModel
from diffusers.pipelines.ltx2.vocoder import LTX2Vocoder
from ...testing_utils import enable_full_determinism, 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 PipelineTesterMixin
from ...testing_utils import assert_tensors_close, enable_full_determinism, torch_device
from ..testing_utils import PipelineTesterMixin
from .testing_utils import (
LTX2BaseTesterConfig,
LTX2LoraMemoryTesterMixin,
LTX2LoraTesterMixin,
LTX2MemoryTesterMixin,
)
enable_full_determinism()
class LTX2ImageToVideoPipelineFastTests(PipelineTesterMixin, unittest.TestCase):
class LTX2ImageToVideoPipelineTesterConfig(LTX2BaseTesterConfig):
pipeline_class = LTX2ImageToVideoPipeline
params = TEXT_TO_IMAGE_PARAMS - {"cross_attention_kwargs"}
batch_params = TEXT_TO_IMAGE_BATCH_PARAMS.union({"image"})
image_params = TEXT_TO_IMAGE_IMAGE_PARAMS
image_latents_params = TEXT_TO_IMAGE_IMAGE_PARAMS
required_optional_params = frozenset(
required_input_params_in_call_signature = frozenset(
[
"num_inference_steps",
"generator",
"latents",
"audio_latents",
"return_dict",
"callback_on_step_end",
"callback_on_step_end_tensor_inputs",
"image",
"prompt",
"height",
"width",
"guidance_scale",
"negative_prompt",
"prompt_embeds",
"negative_prompt_embeds",
]
)
test_attention_slicing = False
test_xformers_attention = False
base_text_encoder_ckpt_id = "hf-internal-testing/tiny-gemma3"
def get_dummy_components(self):
tokenizer = AutoTokenizer.from_pretrained(self.base_text_encoder_ckpt_id)
text_encoder = Gemma3ForConditionalGeneration.from_pretrained(self.base_text_encoder_ckpt_id)
torch.manual_seed(0)
transformer = LTX2VideoTransformer3DModel(
in_channels=4,
out_channels=4,
patch_size=1,
patch_size_t=1,
num_attention_heads=2,
attention_head_dim=8,
cross_attention_dim=16,
audio_in_channels=4,
audio_out_channels=4,
audio_num_attention_heads=2,
audio_attention_head_dim=4,
audio_cross_attention_dim=8,
num_layers=2,
qk_norm="rms_norm_across_heads",
caption_channels=text_encoder.config.text_config.hidden_size,
rope_double_precision=False,
rope_type="split",
)
torch.manual_seed(0)
connectors = LTX2TextConnectors(
caption_channels=text_encoder.config.text_config.hidden_size,
text_proj_in_factor=text_encoder.config.text_config.num_hidden_layers + 1,
video_connector_num_attention_heads=4,
video_connector_attention_head_dim=8,
video_connector_num_layers=1,
video_connector_num_learnable_registers=None,
audio_connector_num_attention_heads=4,
audio_connector_attention_head_dim=8,
audio_connector_num_layers=1,
audio_connector_num_learnable_registers=None,
connector_rope_base_seq_len=32,
rope_theta=10000.0,
rope_double_precision=False,
causal_temporal_positioning=False,
rope_type="split",
)
torch.manual_seed(0)
vae = AutoencoderKLLTX2Video(
in_channels=3,
out_channels=3,
latent_channels=4,
block_out_channels=(8,),
decoder_block_out_channels=(8,),
layers_per_block=(1,),
decoder_layers_per_block=(1, 1),
spatio_temporal_scaling=(True,),
decoder_spatio_temporal_scaling=(True,),
decoder_inject_noise=(False, False),
downsample_type=("spatial",),
upsample_residual=(False,),
upsample_factor=(1,),
timestep_conditioning=False,
patch_size=1,
patch_size_t=1,
encoder_causal=True,
decoder_causal=False,
)
vae.use_framewise_encoding = False
vae.use_framewise_decoding = False
torch.manual_seed(0)
audio_vae = AutoencoderKLLTX2Audio(
base_channels=4,
output_channels=2,
ch_mult=(1,),
num_res_blocks=1,
attn_resolutions=None,
in_channels=2,
resolution=32,
latent_channels=2,
norm_type="pixel",
causality_axis="height",
dropout=0.0,
mid_block_add_attention=False,
sample_rate=16000,
mel_hop_length=160,
is_causal=True,
mel_bins=8,
)
torch.manual_seed(0)
vocoder = LTX2Vocoder(
in_channels=audio_vae.config.output_channels * audio_vae.config.mel_bins,
hidden_channels=32,
out_channels=2,
upsample_kernel_sizes=[4, 4],
upsample_factors=[2, 2],
resnet_kernel_sizes=[3],
resnet_dilations=[[1, 3, 5]],
leaky_relu_negative_slope=0.1,
output_sampling_rate=16000,
)
scheduler = FlowMatchEulerDiscreteScheduler()
components = {
"transformer": transformer,
"vae": vae,
"audio_vae": audio_vae,
"scheduler": scheduler,
"text_encoder": text_encoder,
"tokenizer": tokenizer,
"connectors": connectors,
"vocoder": vocoder,
"processor": None,
"prompt_enhancer": None,
"duration_head": None,
}
return components
batch_input_params = frozenset(["prompt", "negative_prompt", "image"])
def get_dummy_upsample_component(self, in_channels=4, mid_channels=32, num_blocks_per_stage=1):
upsampler = LTX2LatentUpsamplerModel(
return LTX2LatentUpsamplerModel(
in_channels=in_channels,
mid_channels=mid_channels,
num_blocks_per_stage=num_blocks_per_stage,
)
return upsampler
def get_dummy_inputs(self):
generator = self.get_generator(0)
image = torch.rand((1, 3, 32, 32), generator=generator)
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)
image = torch.rand((1, 3, 32, 32), generator=generator, device=device)
inputs = {
return {
"image": image,
"prompt": "a robot dancing",
"negative_prompt": "",
@@ -223,27 +84,23 @@ class LTX2ImageToVideoPipelineFastTests(PipelineTesterMixin, unittest.TestCase):
"max_sequence_length": 16,
# Synthetic float tensors skip H.264 CRF re-compression (training path uses PIL).
"image_crf": 0,
# Request torch outputs so tests compare torch tensors directly (see `BasePipelineTesterConfig`).
"output_type": "pt",
}
return inputs
class TestLTX2ImageToVideoPipeline(LTX2ImageToVideoPipelineTesterConfig, PipelineTesterMixin):
def test_inference(self):
device = "cpu"
# Run on CPU: the expected slices below are CPU-specific.
pipe = self.get_pipeline()
components = self.get_dummy_components()
pipe = self.pipeline_class(**components)
pipe.to(device)
pipe.set_progress_bar_config(disable=None)
inputs = self.get_dummy_inputs(device)
output = pipe(**inputs)
output = pipe(**self.get_dummy_inputs())
video = output.frames
audio = output.audio
self.assertEqual(video.shape, (1, 5, 3, 32, 32))
self.assertEqual(audio.shape[0], 1)
self.assertEqual(audio.shape[1], components["vocoder"].config.out_channels)
assert video.shape == (1, *self.output_shape)
assert audio.shape[0] == 1
assert audio.shape[1] == pipe.vocoder.config.out_channels
# fmt: off
expected_video_slice = torch.tensor(
@@ -263,26 +120,22 @@ class LTX2ImageToVideoPipelineFastTests(PipelineTesterMixin, unittest.TestCase):
generated_video_slice = torch.cat([video[:8], video[-8:]])
generated_audio_slice = torch.cat([audio[:8], audio[-8:]])
assert torch.allclose(expected_video_slice, generated_video_slice, atol=1e-4, rtol=1e-4)
assert torch.allclose(expected_audio_slice, generated_audio_slice, atol=1e-4, rtol=1e-4)
assert_tensors_close(generated_video_slice, expected_video_slice, atol=1e-4, rtol=1e-4)
assert_tensors_close(generated_audio_slice, expected_audio_slice, atol=1e-4, rtol=1e-4)
def test_two_stages_inference(self):
device = "cpu"
# Run on CPU: the expected slices below are CPU-specific.
pipe = self.get_pipeline()
components = self.get_dummy_components()
pipe = self.pipeline_class(**components)
pipe.to(device)
pipe.set_progress_bar_config(disable=None)
inputs = self.get_dummy_inputs(device)
inputs = self.get_dummy_inputs()
inputs["output_type"] = "latent"
first_stage_output = pipe(**inputs)
video_latent = first_stage_output.frames
audio_latent = first_stage_output.audio
self.assertEqual(video_latent.shape, (1, 4, 3, 16, 16))
self.assertEqual(audio_latent.shape, (1, 2, 5, 2))
self.assertEqual(audio_latent.shape[1], components["vocoder"].config.out_channels)
assert video_latent.shape == (1, 4, 3, 16, 16)
assert audio_latent.shape == (1, 2, 5, 2)
assert audio_latent.shape[1] == pipe.vocoder.config.out_channels
inputs["latents"] = video_latent
inputs["audio_latents"] = audio_latent
@@ -291,9 +144,9 @@ class LTX2ImageToVideoPipelineFastTests(PipelineTesterMixin, unittest.TestCase):
video = second_stage_output.frames
audio = second_stage_output.audio
self.assertEqual(video.shape, (1, 5, 3, 32, 32))
self.assertEqual(audio.shape[0], 1)
self.assertEqual(audio.shape[1], components["vocoder"].config.out_channels)
assert video.shape == (1, *self.output_shape)
assert audio.shape[0] == 1
assert audio.shape[1] == pipe.vocoder.config.out_channels
# fmt: off
expected_video_slice = torch.tensor(
@@ -313,31 +166,27 @@ class LTX2ImageToVideoPipelineFastTests(PipelineTesterMixin, unittest.TestCase):
generated_video_slice = torch.cat([video[:8], video[-8:]])
generated_audio_slice = torch.cat([audio[:8], audio[-8:]])
assert torch.allclose(expected_video_slice, generated_video_slice, atol=1e-4, rtol=1e-4)
assert torch.allclose(expected_audio_slice, generated_audio_slice, atol=1e-4, rtol=1e-4)
assert_tensors_close(generated_video_slice, expected_video_slice, atol=1e-4, rtol=1e-4)
assert_tensors_close(generated_audio_slice, expected_audio_slice, atol=1e-4, rtol=1e-4)
def test_two_stages_inference_with_upsampler(self):
device = "cpu"
# Run on CPU: the expected slices below are CPU-specific.
pipe = self.get_pipeline()
components = self.get_dummy_components()
pipe = self.pipeline_class(**components)
pipe.to(device)
pipe.set_progress_bar_config(disable=None)
inputs = self.get_dummy_inputs(device)
inputs = self.get_dummy_inputs()
inputs["output_type"] = "latent"
first_stage_output = pipe(**inputs)
video_latent = first_stage_output.frames
audio_latent = first_stage_output.audio
self.assertEqual(video_latent.shape, (1, 4, 3, 16, 16))
self.assertEqual(audio_latent.shape, (1, 2, 5, 2))
self.assertEqual(audio_latent.shape[1], components["vocoder"].config.out_channels)
assert video_latent.shape == (1, 4, 3, 16, 16)
assert audio_latent.shape == (1, 2, 5, 2)
assert audio_latent.shape[1] == pipe.vocoder.config.out_channels
upsampler = self.get_dummy_upsample_component(in_channels=video_latent.shape[1])
upsample_pipe = LTX2LatentUpsamplePipeline(vae=pipe.vae, latent_upsampler=upsampler)
upscaled_video_latent = upsample_pipe(latents=video_latent, output_type="latent", return_dict=False)[0]
self.assertEqual(upscaled_video_latent.shape, (1, 4, 3, 32, 32))
assert upscaled_video_latent.shape == (1, 4, 3, 32, 32)
inputs["latents"] = upscaled_video_latent
inputs["audio_latents"] = audio_latent
@@ -346,9 +195,9 @@ class LTX2ImageToVideoPipelineFastTests(PipelineTesterMixin, unittest.TestCase):
video = second_stage_output.frames
audio = second_stage_output.audio
self.assertEqual(video.shape, (1, 5, 3, 64, 64))
self.assertEqual(audio.shape[0], 1)
self.assertEqual(audio.shape[1], components["vocoder"].config.out_channels)
assert video.shape == (1, 5, 3, 64, 64)
assert audio.shape[0] == 1
assert audio.shape[1] == pipe.vocoder.config.out_channels
# fmt: off
expected_video_slice = torch.tensor(
@@ -368,31 +217,16 @@ class LTX2ImageToVideoPipelineFastTests(PipelineTesterMixin, unittest.TestCase):
generated_video_slice = torch.cat([video[:8], video[-8:]])
generated_audio_slice = torch.cat([audio[:8], audio[-8:]])
assert torch.allclose(expected_video_slice, generated_video_slice, atol=1e-4, rtol=1e-4)
assert torch.allclose(expected_audio_slice, generated_audio_slice, atol=1e-4, rtol=1e-4)
assert_tensors_close(generated_video_slice, expected_video_slice, atol=1e-4, rtol=1e-4)
assert_tensors_close(generated_audio_slice, expected_audio_slice, atol=1e-4, rtol=1e-4)
def test_inference_batch_single_identical(self):
self._test_inference_batch_single_identical(batch_size=2, expected_max_diff=2e-2)
def get_dummy_duration_head(self):
torch.manual_seed(0)
# The dummy connectors emit 4 heads * 8 head_dim = 32 wide output for both streams.
return LTX2DurationHead(
video_cross_attention_dim=32,
audio_cross_attention_dim=32,
pooler_hidden_dim=8,
num_queries=1,
num_pooler_heads=2,
mlp_hidden_dim=8,
)
def test_inference_batch_single_identical(self, batch_size=2, expected_max_diff=2e-2):
super().test_inference_batch_single_identical(batch_size=batch_size, expected_max_diff=expected_max_diff)
def test_auto_duration_produces_a_grid_valid_frame_count(self):
components = self.get_dummy_components()
components["duration_head"] = self.get_dummy_duration_head()
pipe = self.pipeline_class(**components).to(torch_device)
pipe.set_progress_bar_config(disable=None)
pipe = self.get_pipeline_with_duration_head()
inputs = self.get_dummy_inputs(torch_device)
inputs = self.get_dummy_inputs()
inputs.pop("num_frames")
inputs["min_seconds"] = 0.5
inputs["max_seconds"] = 2.0
@@ -403,12 +237,9 @@ class LTX2ImageToVideoPipelineFastTests(PipelineTesterMixin, unittest.TestCase):
assert 0 < len(frames) <= round(2.0 * inputs["frame_rate"])
def test_omitting_num_frames_auto_predicts_when_a_head_is_present(self):
components = self.get_dummy_components()
components["duration_head"] = self.get_dummy_duration_head()
pipe = self.pipeline_class(**components).to(torch_device)
pipe.set_progress_bar_config(disable=None)
pipe = self.get_pipeline_with_duration_head()
inputs = self.get_dummy_inputs(torch_device)
inputs = self.get_dummy_inputs()
inputs.pop("num_frames")
inputs["output_type"] = "latent"
latents = pipe(**inputs).frames
@@ -422,10 +253,9 @@ class LTX2ImageToVideoPipelineFastTests(PipelineTesterMixin, unittest.TestCase):
# Guards backwards compatibility: a pre-2.5 pipeline has no duration_head and must keep 121.
components = self.get_dummy_components()
assert components.get("duration_head") is None
pipe = self.pipeline_class(**components).to(torch_device)
pipe.set_progress_bar_config(disable=None)
pipe = self.get_pipeline(**components).to(torch_device)
inputs = self.get_dummy_inputs(torch_device)
inputs = self.get_dummy_inputs()
inputs.pop("num_frames")
# Decoding 121 frames is needlessly slow here; the latent frame count already pins num_frames down.
inputs["output_type"] = "latent"
@@ -436,12 +266,9 @@ class LTX2ImageToVideoPipelineFastTests(PipelineTesterMixin, unittest.TestCase):
assert latents.shape[2] == expected_latent_frames
def test_explicit_num_frames_wins_over_a_present_head(self):
components = self.get_dummy_components()
components["duration_head"] = self.get_dummy_duration_head()
pipe = self.pipeline_class(**components).to(torch_device)
pipe.set_progress_bar_config(disable=None)
pipe = self.get_pipeline_with_duration_head()
inputs = self.get_dummy_inputs(torch_device)
inputs = self.get_dummy_inputs()
inputs["num_frames"] = 9
frames = pipe(**inputs).frames[0]
@@ -450,28 +277,21 @@ class LTX2ImageToVideoPipelineFastTests(PipelineTesterMixin, unittest.TestCase):
def test_auto_duration_with_multiple_prompts_raises(self):
# The head predicts one duration, so it cannot serve prompts with different natural lengths.
# Without this guard the pipeline silently applied the first prompt's length to all of them.
components = self.get_dummy_components()
components["duration_head"] = self.get_dummy_duration_head()
pipe = self.pipeline_class(**components).to(torch_device)
pipe.set_progress_bar_config(disable=None)
pipe = self.get_pipeline_with_duration_head()
inputs = self.get_dummy_inputs(torch_device)
inputs = self.get_dummy_inputs()
inputs["prompt"] = ["a robot dancing", "a much longer and quite different scene"]
inputs["negative_prompt"] = ["", ""]
inputs.pop("num_frames")
with self.assertRaises(ValueError) as ctx:
with pytest.raises(ValueError, match="2 prompts were supplied"):
pipe(**inputs)
assert "2 prompts were supplied" in str(ctx.exception)
def test_multiple_prompts_still_work_with_an_explicit_num_frames(self):
# The guard must be scoped to the auto path -- batched prompts with an integer are unaffected.
components = self.get_dummy_components()
components["duration_head"] = self.get_dummy_duration_head()
pipe = self.pipeline_class(**components).to(torch_device)
pipe.set_progress_bar_config(disable=None)
pipe = self.get_pipeline_with_duration_head()
inputs = self.get_dummy_inputs(torch_device)
inputs = self.get_dummy_inputs()
inputs["prompt"] = ["a robot dancing", "a much longer and quite different scene"]
inputs["negative_prompt"] = ["", ""]
inputs["num_frames"] = 5
@@ -480,3 +300,15 @@ class LTX2ImageToVideoPipelineFastTests(PipelineTesterMixin, unittest.TestCase):
latents = pipe(**inputs).frames
assert latents.shape[0] == 2
class TestLTX2ImageToVideoPipelineMemory(LTX2ImageToVideoPipelineTesterConfig, LTX2MemoryTesterMixin):
"""Memory optimization tests (CPU offload, group offload, layerwise casting) for the LTX2 I2V pipeline."""
class TestLTX2ImageToVideoPipelineLoRA(LTX2ImageToVideoPipelineTesterConfig, LTX2LoraTesterMixin):
"""LoRA tests for the LTX2 I2V pipeline."""
class TestLTX2ImageToVideoPipelineLoRAMemory(LTX2ImageToVideoPipelineTesterConfig, LTX2LoraMemoryTesterMixin):
"""LoRA x memory-optimization tests (group offload, CPU offload) for the LTX2 I2V pipeline."""

212
tests/pipelines/ltx2/test_ltx2_in_context.py
View File

@@ -12,192 +12,49 @@
# See the License for the specific language governing permissions and
# limitations under the License.
import unittest
import torch
from transformers import AutoTokenizer, Gemma3ForConditionalGeneration
from diffusers import (
AutoencoderKLLTX2Audio,
AutoencoderKLLTX2Video,
FlowMatchEulerDiscreteScheduler,
LTX2InContextPipeline,
LTX2VideoTransformer3DModel,
)
from diffusers.pipelines.ltx2 import LTX2TextConnectors
from diffusers.pipelines.ltx2.latent_upsampler import LTX2LatentUpsamplerModel
from diffusers import LTX2InContextPipeline
from diffusers.pipelines.ltx2.pipeline_ltx2_condition import LTX2VideoCondition
from diffusers.pipelines.ltx2.vocoder import LTX2Vocoder
from ...testing_utils import enable_full_determinism
from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import PipelineTesterMixin
from ..testing_utils import PipelineTesterMixin
from .testing_utils import (
LTX2BaseTesterConfig,
LTX2LoraMemoryTesterMixin,
LTX2LoraTesterMixin,
LTX2MemoryTesterMixin,
)
enable_full_determinism()
class LTX2InContextPipelineFastTests(PipelineTesterMixin, unittest.TestCase):
class LTX2InContextPipelineTesterConfig(LTX2BaseTesterConfig):
pipeline_class = LTX2InContextPipeline
params = TEXT_TO_IMAGE_PARAMS - {"cross_attention_kwargs"}
batch_params = TEXT_TO_IMAGE_BATCH_PARAMS
image_params = TEXT_TO_IMAGE_IMAGE_PARAMS
image_latents_params = TEXT_TO_IMAGE_IMAGE_PARAMS
required_optional_params = frozenset(
required_input_params_in_call_signature = frozenset(
[
"num_inference_steps",
"generator",
"latents",
"audio_latents",
"return_dict",
"callback_on_step_end",
"callback_on_step_end_tensor_inputs",
"conditions",
"reference_conditions",
"prompt",
"height",
"width",
"guidance_scale",
"negative_prompt",
"prompt_embeds",
"negative_prompt_embeds",
]
)
test_attention_slicing = False
test_xformers_attention = False
batch_input_params = frozenset(["prompt", "negative_prompt"])
unset_components = ("audio_scheduler", "processor", "prompt_enhancer")
base_text_encoder_ckpt_id = "hf-internal-testing/tiny-gemma3"
def get_dummy_components(self):
tokenizer = AutoTokenizer.from_pretrained(self.base_text_encoder_ckpt_id)
text_encoder = Gemma3ForConditionalGeneration.from_pretrained(self.base_text_encoder_ckpt_id)
torch.manual_seed(0)
transformer = LTX2VideoTransformer3DModel(
in_channels=4,
out_channels=4,
patch_size=1,
patch_size_t=1,
num_attention_heads=2,
attention_head_dim=8,
cross_attention_dim=16,
audio_in_channels=4,
audio_out_channels=4,
audio_num_attention_heads=2,
audio_attention_head_dim=4,
audio_cross_attention_dim=8,
num_layers=2,
qk_norm="rms_norm_across_heads",
caption_channels=text_encoder.config.text_config.hidden_size,
rope_double_precision=False,
rope_type="split",
)
torch.manual_seed(0)
connectors = LTX2TextConnectors(
caption_channels=text_encoder.config.text_config.hidden_size,
text_proj_in_factor=text_encoder.config.text_config.num_hidden_layers + 1,
video_connector_num_attention_heads=4,
video_connector_attention_head_dim=8,
video_connector_num_layers=1,
video_connector_num_learnable_registers=None,
audio_connector_num_attention_heads=4,
audio_connector_attention_head_dim=8,
audio_connector_num_layers=1,
audio_connector_num_learnable_registers=None,
connector_rope_base_seq_len=32,
rope_theta=10000.0,
rope_double_precision=False,
causal_temporal_positioning=False,
rope_type="split",
)
torch.manual_seed(0)
vae = AutoencoderKLLTX2Video(
in_channels=3,
out_channels=3,
latent_channels=4,
block_out_channels=(8,),
decoder_block_out_channels=(8,),
layers_per_block=(1,),
decoder_layers_per_block=(1, 1),
spatio_temporal_scaling=(True,),
decoder_spatio_temporal_scaling=(True,),
decoder_inject_noise=(False, False),
downsample_type=("spatial",),
upsample_residual=(False,),
upsample_factor=(1,),
timestep_conditioning=False,
patch_size=1,
patch_size_t=1,
encoder_causal=True,
decoder_causal=False,
)
vae.use_framewise_encoding = False
vae.use_framewise_decoding = False
torch.manual_seed(0)
audio_vae = AutoencoderKLLTX2Audio(
base_channels=4,
output_channels=2,
ch_mult=(1,),
num_res_blocks=1,
attn_resolutions=None,
in_channels=2,
resolution=32,
latent_channels=2,
norm_type="pixel",
causality_axis="height",
dropout=0.0,
mid_block_add_attention=False,
sample_rate=16000,
mel_hop_length=160,
is_causal=True,
mel_bins=8,
)
torch.manual_seed(0)
vocoder = LTX2Vocoder(
in_channels=audio_vae.config.output_channels * audio_vae.config.mel_bins,
hidden_channels=32,
out_channels=2,
upsample_kernel_sizes=[4, 4],
upsample_factors=[2, 2],
resnet_kernel_sizes=[3],
resnet_dilations=[[1, 3, 5]],
leaky_relu_negative_slope=0.1,
output_sampling_rate=16000,
)
scheduler = FlowMatchEulerDiscreteScheduler()
components = {
"transformer": transformer,
"vae": vae,
"audio_vae": audio_vae,
"scheduler": scheduler,
"text_encoder": text_encoder,
"tokenizer": tokenizer,
"connectors": connectors,
"vocoder": vocoder,
"audio_scheduler": None,
"processor": None,
"prompt_enhancer": None,
}
return components
def get_dummy_upsample_component(self, in_channels=4, mid_channels=32, num_blocks_per_stage=1):
upsampler = LTX2LatentUpsamplerModel(
in_channels=in_channels,
mid_channels=mid_channels,
num_blocks_per_stage=num_blocks_per_stage,
)
return upsampler
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)
image = torch.rand((1, 3, 32, 32), generator=generator, device=device)
def get_dummy_inputs(self):
generator = self.get_generator(0)
image = torch.rand((1, 3, 32, 32), generator=generator)
# Synthetic float tensors skip H.264 CRF re-compression (training path uses PIL/uint8).
img_cond = LTX2VideoCondition(frames=image, index=0, strength=1.0, crf=0)
inputs = {
return {
"conditions": img_cond,
"prompt": "a robot dancing",
"negative_prompt": "",
@@ -209,10 +66,23 @@ class LTX2InContextPipelineFastTests(PipelineTesterMixin, unittest.TestCase):
"num_frames": 5,
"frame_rate": 25.0,
"max_sequence_length": 16,
# Request torch outputs so tests compare torch tensors directly (see `BasePipelineTesterConfig`).
"output_type": "pt",
}
return inputs
def test_inference_batch_single_identical(self):
self._test_inference_batch_single_identical(batch_size=2, expected_max_diff=1e-3)
class TestLTX2InContextPipeline(LTX2InContextPipelineTesterConfig, PipelineTesterMixin):
def test_inference_batch_single_identical(self, batch_size=2, expected_max_diff=1e-3):
super().test_inference_batch_single_identical(batch_size=batch_size, expected_max_diff=expected_max_diff)
class TestLTX2InContextPipelineMemory(LTX2InContextPipelineTesterConfig, LTX2MemoryTesterMixin):
"""Memory optimization tests (CPU offload, group offload, layerwise casting) for the LTX2 IC pipeline."""
class TestLTX2InContextPipelineLoRA(LTX2InContextPipelineTesterConfig, LTX2LoraTesterMixin):
"""LoRA tests for the LTX2 IC pipeline."""
class TestLTX2InContextPipelineLoRAMemory(LTX2InContextPipelineTesterConfig, LTX2LoraMemoryTesterMixin):
"""LoRA x memory-optimization tests (group offload, CPU offload) for the LTX2 IC pipeline."""

228
tests/pipelines/ltx2/testing_utils.py Normal file
View File

@@ -0,0 +1,228 @@
# Copyright 2026 The HuggingFace Team.
#
# 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.
"""Shared test fixtures for the LTX2 pipelines.
Every LTX2 pipeline in this directory takes the same set of dummy sub-modules — they differ only in which
optional components they accept and in what `__call__` takes — so the component builder lives here and the
per-pipeline configs subclass `LTX2BaseTesterConfig`. The two scoped tester mixins below carry the
pipeline-family-wide skips so each test file does not restate them.
"""
import pytest
import torch
from transformers import AutoTokenizer, Gemma3ForConditionalGeneration
from diffusers import (
AutoencoderKLLTX2Audio,
AutoencoderKLLTX2Video,
FlowMatchEulerDiscreteScheduler,
LTX2VideoTransformer3DModel,
)
from diffusers.pipelines.ltx2 import LTX2DurationHead, LTX2TextConnectors
from diffusers.pipelines.ltx2.vocoder import LTX2Vocoder
from ...testing_utils import torch_device
from ..testing_utils import (
BasePipelineTesterConfig,
LoraMemoryTesterMixin,
LoraTesterMixin,
MemoryTesterMixin,
)
class LTX2BaseTesterConfig(BasePipelineTesterConfig):
"""Dummy component set shared by every LTX2 pipeline in this directory."""
# LTX2 is a video pipeline (`num_videos_per_prompt`, not `num_images_per_prompt`) and takes a second latent
# input for the audio stream. `LTX2HDRPipeline` renders video only and overrides this without `audio_latents`.
optional_input_params = frozenset(
[
"num_inference_steps",
"num_videos_per_prompt",
"generator",
"latents",
"audio_latents",
"output_type",
"return_dict",
]
)
output_shape = (5, 3, 32, 32)
base_text_encoder_ckpt_id = "hf-internal-testing/tiny-gemma3"
# Components the pipeline accepts but that these fast tests leave unset. Which ones exist differs per
# pipeline — only the pipelines that can predict a duration take a `duration_head`, and only some take an
# `audio_scheduler` — so each config lists its own set and `get_dummy_components` fills them with `None`.
unset_components = ("processor", "prompt_enhancer", "duration_head")
def get_dummy_components(self):
tokenizer = AutoTokenizer.from_pretrained(self.base_text_encoder_ckpt_id)
text_encoder = Gemma3ForConditionalGeneration.from_pretrained(self.base_text_encoder_ckpt_id)
torch.manual_seed(0)
transformer = LTX2VideoTransformer3DModel(
in_channels=4,
out_channels=4,
patch_size=1,
patch_size_t=1,
num_attention_heads=2,
attention_head_dim=8,
cross_attention_dim=16,
audio_in_channels=4,
audio_out_channels=4,
audio_num_attention_heads=2,
audio_attention_head_dim=4,
audio_cross_attention_dim=8,
num_layers=2,
qk_norm="rms_norm_across_heads",
caption_channels=text_encoder.config.text_config.hidden_size,
rope_double_precision=False,
rope_type="split",
)
torch.manual_seed(0)
connectors = LTX2TextConnectors(
caption_channels=text_encoder.config.text_config.hidden_size,
text_proj_in_factor=text_encoder.config.text_config.num_hidden_layers + 1,
video_connector_num_attention_heads=4,
video_connector_attention_head_dim=8,
video_connector_num_layers=1,
video_connector_num_learnable_registers=None,
audio_connector_num_attention_heads=4,
audio_connector_attention_head_dim=8,
audio_connector_num_layers=1,
audio_connector_num_learnable_registers=None,
connector_rope_base_seq_len=32,
rope_theta=10000.0,
rope_double_precision=False,
causal_temporal_positioning=False,
rope_type="split",
)
torch.manual_seed(0)
vae = AutoencoderKLLTX2Video(
in_channels=3,
out_channels=3,
latent_channels=4,
block_out_channels=(8,),
decoder_block_out_channels=(8,),
layers_per_block=(1,),
decoder_layers_per_block=(1, 1),
spatio_temporal_scaling=(True,),
decoder_spatio_temporal_scaling=(True,),
decoder_inject_noise=(False, False),
downsample_type=("spatial",),
upsample_residual=(False,),
upsample_factor=(1,),
timestep_conditioning=False,
patch_size=1,
patch_size_t=1,
encoder_causal=True,
decoder_causal=False,
)
vae.use_framewise_encoding = False
vae.use_framewise_decoding = False
torch.manual_seed(0)
audio_vae = AutoencoderKLLTX2Audio(
base_channels=4,
output_channels=2,
ch_mult=(1,),
num_res_blocks=1,
attn_resolutions=None,
in_channels=2,
resolution=32,
latent_channels=2,
norm_type="pixel",
causality_axis="height",
dropout=0.0,
mid_block_add_attention=False,
sample_rate=16000,
mel_hop_length=160,
is_causal=True,
mel_bins=8,
)
torch.manual_seed(0)
vocoder = LTX2Vocoder(
in_channels=audio_vae.config.output_channels * audio_vae.config.mel_bins,
hidden_channels=32,
out_channels=2,
upsample_kernel_sizes=[4, 4],
upsample_factors=[2, 2],
resnet_kernel_sizes=[3],
resnet_dilations=[[1, 3, 5]],
leaky_relu_negative_slope=0.1,
output_sampling_rate=16000,
)
scheduler = FlowMatchEulerDiscreteScheduler()
return {
"transformer": transformer,
"vae": vae,
"audio_vae": audio_vae,
"scheduler": scheduler,
"text_encoder": text_encoder,
"tokenizer": tokenizer,
"connectors": connectors,
"vocoder": vocoder,
**dict.fromkeys(self.unset_components),
}
def get_dummy_duration_head(self):
"""A tiny `LTX2DurationHead`, for the pipelines that accept one (`duration_head` in `unset_components`)."""
torch.manual_seed(0)
# The dummy connectors emit 4 heads * 8 head_dim = 32 wide output for both streams.
return LTX2DurationHead(
video_cross_attention_dim=32,
audio_cross_attention_dim=32,
pooler_hidden_dim=8,
num_queries=1,
num_pooler_heads=2,
mlp_hidden_dim=8,
)
def get_pipeline_with_duration_head(self):
components = self.get_dummy_components()
components["duration_head"] = self.get_dummy_duration_head()
return self.get_pipeline(**components).to(torch_device)
class LTX2MemoryTesterMixin(MemoryTesterMixin):
"""`MemoryTesterMixin` for the LTX2 pipelines in this directory."""
# The shared helper only offloads a fixed set of component names and leaves LTX2's extra module components
# (`connectors`, `audio_vae`, `vocoder`) on CPU, so the forward pass mixes devices. Pipeline-level
# offloading, which walks every component, is exercised by `test_pipeline_level_group_offloading_inference`.
@pytest.mark.skip("Using test_pipeline_level_group_offloading_inference instead")
def test_group_offloading_inference(self):
pass
class LTX2LoraTesterMixin(LoraTesterMixin):
"""`LoraTesterMixin` for the LTX2 pipelines in this directory.
Every LTX2 pipeline advertises `connectors` as LoRA-loadable and `load_lora_weights` does handle connector
LoRAs, but `save_lora_weights` only accepts `transformer_lora_layers` — so connector adapters cannot
round-trip through the public API these tests drive. Scope the tests to the transformer until that closes.
"""
lora_loadable_components = ["transformer"]
class LTX2LoraMemoryTesterMixin(LoraMemoryTesterMixin):
"""`LoraMemoryTesterMixin` for the LTX2 pipelines, scoped to the transformer — see `LTX2LoraTesterMixin`."""
lora_loadable_components = ["transformer"]

89
tests/pipelines/lumina/test_lumina_nextdit.py
View File

@@ -1,7 +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 unittest
import numpy as np
import pytest
import torch
from transformers import AutoTokenizer, GemmaConfig, GemmaForCausalLM
@@ -19,12 +34,12 @@ from ...testing_utils import (
slow,
torch_device,
)
from ..test_pipelines_common import PipelineTesterMixin
from ..testing_utils import BasePipelineTesterConfig, MemoryTesterMixin, PipelineTesterMixin
class LuminaPipelineFastTests(unittest.TestCase, PipelineTesterMixin):
class LuminaPipelineTesterConfig(BasePipelineTesterConfig):
pipeline_class = LuminaPipeline
params = frozenset(
required_input_params_in_call_signature = frozenset(
[
"prompt",
"height",
@@ -35,10 +50,10 @@ class LuminaPipelineFastTests(unittest.TestCase, PipelineTesterMixin):
"negative_prompt_embeds",
]
)
batch_params = frozenset(["prompt", "negative_prompt"])
test_layerwise_casting = True
test_group_offloading = True
batch_input_params = frozenset(["prompt", "negative_prompt"])
# The dummy `AutoencoderKL` has a single block and so does not up/downsample: latents at
# `32 // vae_scale_factor` decode to a 4x4 image rather than back to the requested 32x32.
output_shape = (3, 4, 4)
def get_dummy_components(self):
torch.manual_seed(0)
@@ -75,72 +90,63 @@ class LuminaPipelineFastTests(unittest.TestCase, PipelineTesterMixin):
)
text_encoder = GemmaForCausalLM(config)
components = {
"transformer": transformer.eval(),
"vae": vae.eval(),
return {
"transformer": transformer,
"vae": vae,
"scheduler": scheduler,
"text_encoder": text_encoder.eval(),
"text_encoder": text_encoder,
"tokenizer": tokenizer,
}
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="cpu").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": 5.0,
"output_type": "np",
"height": 32,
"width": 32,
# Request torch outputs so tests compare torch tensors directly (see `BasePipelineTesterConfig`).
"output_type": "pt",
}
return inputs
@unittest.skip("xformers attention processor does not exist for Lumina")
def test_xformers_attention_forwardGenerator_pass(self):
pass
class TestLuminaPipeline(LuminaPipelineTesterConfig, PipelineTesterMixin):
pass
class TestLuminaPipelineMemory(LuminaPipelineTesterConfig, MemoryTesterMixin):
"""Memory optimization tests (CPU offload, group offload, layerwise casting) for the Lumina pipeline."""
@slow
@require_torch_accelerator
class LuminaPipelineSlowTests(unittest.TestCase):
class TestLuminaPipelineIntegration:
pipeline_class = LuminaPipeline
repo_id = "Alpha-VLLM/Lumina-Next-SFT-diffusers"
def setUp(self):
super().setUp()
@pytest.fixture(autouse=True)
def cleanup(self):
gc.collect()
backend_empty_cache(torch_device)
def tearDown(self):
super().tearDown()
yield
gc.collect()
backend_empty_cache(torch_device)
def get_inputs(self, device, seed=0):
if str(device).startswith("mps"):
generator = torch.manual_seed(seed)
else:
generator = torch.Generator(device="cpu").manual_seed(seed)
return {
"prompt": "A photo of a cat",
"num_inference_steps": 2,
"guidance_scale": 5.0,
"output_type": "np",
"generator": generator,
"generator": torch.Generator(device="cpu").manual_seed(seed),
}
def test_lumina_inference(self):
pipe = self.pipeline_class.from_pretrained(self.repo_id, torch_dtype=torch.bfloat16)
pipe.enable_model_cpu_offload(device=torch_device)
inputs = self.get_inputs(torch_device)
image = pipe(**inputs).images[0]
image = pipe(**self.get_inputs(torch_device)).images[0]
image_slice = image[0, :10, :10]
expected_slice = np.array(
[
@@ -159,5 +165,4 @@ class LuminaPipelineSlowTests(unittest.TestCase):
)
max_diff = numpy_cosine_similarity_distance(expected_slice.flatten(), image_slice.flatten())
assert max_diff < 1e-4