diff --git a/tests/pipelines/chroma/test_pipeline_chroma.py b/tests/pipelines/chroma/test_pipeline_chroma.py index e5b2a27b87bf..ecaf056bf716 100644 --- a/tests/pipelines/chroma/test_pipeline_chroma.py +++ b/tests/pipelines/chroma/test_pipeline_chroma.py @@ -1,28 +1,25 @@ -import unittest - -import numpy as np import torch from transformers import AutoConfig, AutoTokenizer, T5EncoderModel from diffusers import AutoencoderKL, ChromaPipeline, ChromaTransformer2DModel, FlowMatchEulerDiscreteScheduler -from ...testing_utils import torch_device -from ..test_pipelines_common import FluxIPAdapterTesterMixin, PipelineTesterMixin, check_qkv_fused_layers_exist +from ...testing_utils import assert_tensors_close, torch_device +from ..flux.testing_utils import FluxIPAdapterTesterMixin +from ..testing_utils import ( + BasePipelineTesterConfig, + MemoryTesterMixin, + PipelineTesterMixin, + check_qkv_fused_layers_exist, +) -class ChromaPipelineFastTests( - unittest.TestCase, - PipelineTesterMixin, - FluxIPAdapterTesterMixin, -): +class ChromaPipelineTesterConfig(BasePipelineTesterConfig): pipeline_class = ChromaPipeline - params = frozenset(["prompt", "height", "width", "guidance_scale", "prompt_embeds"]) - batch_params = frozenset(["prompt"]) - - # there is no xformers processor for Flux - test_xformers_attention = False - test_layerwise_casting = True - test_group_offloading = True + required_input_params_in_call_signature = frozenset( + ["prompt", "height", "width", "guidance_scale", "prompt_embeds"] + ) + batch_input_params = frozenset(["prompt"]) + output_shape = (3, 8, 8) def get_dummy_components(self, num_layers: int = 1, num_single_layers: int = 1): torch.manual_seed(0) @@ -75,82 +72,88 @@ def get_dummy_components(self, num_layers: int = 1, num_single_layers: int = 1): "feature_extractor": None, } - 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", "negative_prompt": "bad, ugly", - "generator": generator, + "generator": self.get_generator(0), "num_inference_steps": 2, "guidance_scale": 5.0, "height": 8, "width": 8, "max_sequence_length": 48, - "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 TestChromaPipeline(ChromaPipelineTesterConfig, PipelineTesterMixin): def test_chroma_different_prompts(self): - pipe = self.pipeline_class(**self.get_dummy_components()).to(torch_device) + pipe = self.get_pipeline().to(torch_device) - inputs = self.get_dummy_inputs(torch_device) + inputs = self.get_dummy_inputs() output_same_prompt = pipe(**inputs).images[0] - inputs = self.get_dummy_inputs(torch_device) + inputs = self.get_dummy_inputs() inputs["prompt"] = "a different prompt" output_different_prompts = pipe(**inputs).images[0] - max_diff = np.abs(output_same_prompt - output_different_prompts).max() + max_diff = (output_same_prompt - output_different_prompts).abs().max() # Outputs should be different here # For some reasons, they don't show large differences - assert max_diff > 1e-6 + assert max_diff > 1e-6, "Outputs should be different for different prompts." def test_fused_qkv_projections(self): - device = "cpu" # ensure determinism for the device-dependent torch.Generator - components = self.get_dummy_components() - pipe = self.pipeline_class(**components) - pipe = pipe.to(device) - pipe.set_progress_bar_config(disable=None) + # Run on CPU to keep the seeded generator deterministic across the three forward passes. + pipe = self.get_pipeline() - inputs = self.get_dummy_inputs(device) + inputs = self.get_dummy_inputs() image = pipe(**inputs).images - original_image_slice = image[0, -3:, -3:, -1] + original_image_slice = image[0, -1, -3:, -3:] # TODO (sayakpaul): will refactor this once `fuse_qkv_projections()` has been added # to the pipeline level. pipe.transformer.fuse_qkv_projections() - self.assertTrue( - check_qkv_fused_layers_exist(pipe.transformer, ["to_qkv"]), - ("Something wrong with the fused attention layers. Expected all the attention projections to be fused."), + assert check_qkv_fused_layers_exist(pipe.transformer, ["to_qkv"]), ( + "Something wrong with the fused attention layers. Expected all the attention projections to be fused." ) - inputs = self.get_dummy_inputs(device) + inputs = self.get_dummy_inputs() image = pipe(**inputs).images - image_slice_fused = image[0, -3:, -3:, -1] + image_slice_fused = image[0, -1, -3:, -3:] pipe.transformer.unfuse_qkv_projections() - inputs = self.get_dummy_inputs(device) + inputs = self.get_dummy_inputs() image = pipe(**inputs).images - image_slice_disabled = image[0, -3:, -3:, -1] - - assert np.allclose(original_image_slice, image_slice_fused, atol=1e-3, rtol=1e-3), ( - "Fusion of QKV projections shouldn't affect the outputs." + image_slice_disabled = image[0, -1, -3:, -3:] + + assert_tensors_close( + original_image_slice, + image_slice_fused, + atol=1e-3, + rtol=1e-3, + msg="Fusion of QKV projections shouldn't affect the outputs.", ) - assert np.allclose(image_slice_fused, image_slice_disabled, atol=1e-3, rtol=1e-3), ( - "Outputs, with QKV projection fusion enabled, shouldn't change when fused QKV projections are disabled." + assert_tensors_close( + image_slice_fused, + image_slice_disabled, + atol=1e-3, + rtol=1e-3, + msg="Outputs, with QKV projection fusion enabled, shouldn't change when fused QKV projections are disabled.", ) - assert np.allclose(original_image_slice, image_slice_disabled, atol=1e-2, rtol=1e-2), ( - "Original outputs should match when fused QKV projections are disabled." + assert_tensors_close( + original_image_slice, + image_slice_disabled, + atol=1e-2, + rtol=1e-2, + msg="Original outputs should match when fused QKV projections are disabled.", ) def test_chroma_image_output_shape(self): - pipe = self.pipeline_class(**self.get_dummy_components()).to(torch_device) - inputs = self.get_dummy_inputs(torch_device) + pipe = self.get_pipeline().to(torch_device) + inputs = self.get_dummy_inputs() height_width_pairs = [(32, 32), (72, 57)] for height, width in height_width_pairs: @@ -159,5 +162,15 @@ def test_chroma_image_output_shape(self): inputs.update({"height": height, "width": width}) image = pipe(**inputs).images[0] - output_height, output_width, _ = image.shape - assert (output_height, output_width) == (expected_height, expected_width) + _, output_height, output_width = image.shape + assert (output_height, output_width) == (expected_height, expected_width), ( + f"Output shape {image.shape} does not match expected shape {(expected_height, expected_width)}" + ) + + +class TestChromaPipelineIPAdapter(ChromaPipelineTesterConfig, FluxIPAdapterTesterMixin): + """IP-Adapter tests for the Chroma pipeline.""" + + +class TestChromaPipelineMemory(ChromaPipelineTesterConfig, MemoryTesterMixin): + """Memory optimization tests (CPU offload, group offload, layerwise casting) for the Chroma pipeline.""" diff --git a/tests/pipelines/chroma/test_pipeline_chroma_img2img.py b/tests/pipelines/chroma/test_pipeline_chroma_img2img.py index 2b415d303a4b..224d2bd6c98f 100644 --- a/tests/pipelines/chroma/test_pipeline_chroma_img2img.py +++ b/tests/pipelines/chroma/test_pipeline_chroma_img2img.py @@ -1,29 +1,27 @@ import random -import unittest -import numpy as np import torch from transformers import AutoConfig, AutoTokenizer, T5EncoderModel from diffusers import AutoencoderKL, ChromaImg2ImgPipeline, ChromaTransformer2DModel, FlowMatchEulerDiscreteScheduler -from ...testing_utils import floats_tensor, torch_device -from ..test_pipelines_common import FluxIPAdapterTesterMixin, PipelineTesterMixin, check_qkv_fused_layers_exist +from ...testing_utils import assert_tensors_close, floats_tensor, torch_device +from ..flux.testing_utils import FluxIPAdapterTesterMixin +from ..testing_utils import ( + BasePipelineTesterConfig, + MemoryTesterMixin, + PipelineTesterMixin, + check_qkv_fused_layers_exist, +) -class ChromaImg2ImgPipelineFastTests( - unittest.TestCase, - PipelineTesterMixin, - FluxIPAdapterTesterMixin, -): +class ChromaImg2ImgPipelineTesterConfig(BasePipelineTesterConfig): pipeline_class = ChromaImg2ImgPipeline - params = frozenset(["prompt", "height", "width", "guidance_scale", "prompt_embeds"]) - batch_params = frozenset(["prompt"]) - - # there is no xformers processor for Flux - test_xformers_attention = False - test_layerwise_casting = True - test_group_offloading = True + required_input_params_in_call_signature = frozenset( + ["prompt", "height", "width", "guidance_scale", "prompt_embeds"] + ) + batch_input_params = frozenset(["prompt"]) + output_shape = (3, 8, 8) def get_dummy_components(self, num_layers: int = 1, num_single_layers: int = 1): torch.manual_seed(0) @@ -76,84 +74,91 @@ def get_dummy_components(self, num_layers: int = 1, num_single_layers: int = 1): "feature_extractor": None, } - def get_dummy_inputs(self, device, seed=0): - image = floats_tensor((1, 3, 32, 32), rng=random.Random(seed)).to(device) - if str(device).startswith("mps"): - generator = torch.manual_seed(seed) - else: - generator = torch.Generator(device="cpu").manual_seed(seed) + def get_dummy_inputs(self): + image = floats_tensor((1, 3, 32, 32), rng=random.Random(0)).to(torch_device) - 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": 5.0, "height": 8, "width": 8, "max_sequence_length": 48, "strength": 0.8, - "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 TestChromaImg2ImgPipeline(ChromaImg2ImgPipelineTesterConfig, PipelineTesterMixin): def test_chroma_different_prompts(self): - pipe = self.pipeline_class(**self.get_dummy_components()).to(torch_device) + pipe = self.get_pipeline().to(torch_device) - inputs = self.get_dummy_inputs(torch_device) + inputs = self.get_dummy_inputs() output_same_prompt = pipe(**inputs).images[0] - inputs = self.get_dummy_inputs(torch_device) + inputs = self.get_dummy_inputs() inputs["prompt"] = "a different prompt" output_different_prompts = pipe(**inputs).images[0] - max_diff = np.abs(output_same_prompt - output_different_prompts).max() + max_diff = (output_same_prompt - output_different_prompts).abs().max() # Outputs should be different here # For some reasons, they don't show large differences - assert max_diff > 1e-6 + assert max_diff > 1e-6, "Outputs should be different for different prompts." def test_fused_qkv_projections(self): - device = "cpu" # ensure determinism for the device-dependent torch.Generator - components = self.get_dummy_components() - pipe = self.pipeline_class(**components) - pipe = pipe.to(device) - pipe.set_progress_bar_config(disable=None) + # Run on CPU to keep the seeded generator deterministic across the three forward passes. + pipe = self.get_pipeline() - inputs = self.get_dummy_inputs(device) + inputs = self.get_dummy_inputs() image = pipe(**inputs).images - original_image_slice = image[0, -3:, -3:, -1] + original_image_slice = image[0, -1, -3:, -3:] # TODO (sayakpaul): will refactor this once `fuse_qkv_projections()` has been added # to the pipeline level. pipe.transformer.fuse_qkv_projections() - self.assertTrue( - check_qkv_fused_layers_exist(pipe.transformer, ["to_qkv"]), - ("Something wrong with the fused attention layers. Expected all the attention projections to be fused."), + assert check_qkv_fused_layers_exist(pipe.transformer, ["to_qkv"]), ( + "Something wrong with the fused attention layers. Expected all the attention projections to be fused." ) - inputs = self.get_dummy_inputs(device) + inputs = self.get_dummy_inputs() image = pipe(**inputs).images - image_slice_fused = image[0, -3:, -3:, -1] + image_slice_fused = image[0, -1, -3:, -3:] pipe.transformer.unfuse_qkv_projections() - inputs = self.get_dummy_inputs(device) + inputs = self.get_dummy_inputs() image = pipe(**inputs).images - image_slice_disabled = image[0, -3:, -3:, -1] - - assert np.allclose(original_image_slice, image_slice_fused, atol=1e-3, rtol=1e-3), ( - "Fusion of QKV projections shouldn't affect the outputs." + image_slice_disabled = image[0, -1, -3:, -3:] + + assert_tensors_close( + original_image_slice, + image_slice_fused, + atol=1e-3, + rtol=1e-3, + msg="Fusion of QKV projections shouldn't affect the outputs.", ) - assert np.allclose(image_slice_fused, image_slice_disabled, atol=1e-3, rtol=1e-3), ( - "Outputs, with QKV projection fusion enabled, shouldn't change when fused QKV projections are disabled." + assert_tensors_close( + image_slice_fused, + image_slice_disabled, + atol=1e-3, + rtol=1e-3, + msg="Outputs, with QKV projection fusion enabled, shouldn't change when fused QKV projections are disabled.", ) - assert np.allclose(original_image_slice, image_slice_disabled, atol=1e-2, rtol=1e-2), ( - "Original outputs should match when fused QKV projections are disabled." + assert_tensors_close( + original_image_slice, + image_slice_disabled, + atol=1e-2, + rtol=1e-2, + msg="Original outputs should match when fused QKV projections are disabled.", ) def test_chroma_image_output_shape(self): - pipe = self.pipeline_class(**self.get_dummy_components()).to(torch_device) - inputs = self.get_dummy_inputs(torch_device) + pipe = self.get_pipeline().to(torch_device) + inputs = self.get_dummy_inputs() height_width_pairs = [(32, 32), (72, 57)] for height, width in height_width_pairs: @@ -162,5 +167,15 @@ def test_chroma_image_output_shape(self): inputs.update({"height": height, "width": width}) image = pipe(**inputs).images[0] - output_height, output_width, _ = image.shape - assert (output_height, output_width) == (expected_height, expected_width) + _, output_height, output_width = image.shape + assert (output_height, output_width) == (expected_height, expected_width), ( + f"Output shape {image.shape} does not match expected shape {(expected_height, expected_width)}" + ) + + +class TestChromaImg2ImgPipelineIPAdapter(ChromaImg2ImgPipelineTesterConfig, FluxIPAdapterTesterMixin): + """IP-Adapter tests for the Chroma img2img pipeline.""" + + +class TestChromaImg2ImgPipelineMemory(ChromaImg2ImgPipelineTesterConfig, MemoryTesterMixin): + """Memory optimization tests (CPU offload, group offload, layerwise casting) for the Chroma img2img pipeline.""" diff --git a/tests/pipelines/chronoedit/test_chronoedit.py b/tests/pipelines/chronoedit/test_chronoedit.py index 374243d5c9ed..2447492ece10 100644 --- a/tests/pipelines/chronoedit/test_chronoedit.py +++ b/tests/pipelines/chronoedit/test_chronoedit.py @@ -12,8 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -import unittest - +import pytest import torch from PIL import Image from transformers import ( @@ -32,31 +31,21 @@ FlowMatchEulerDiscreteScheduler, ) -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 - - -enable_full_determinism() +from ...testing_utils import assert_tensors_close +from ..testing_utils import BasePipelineTesterConfig, MemoryTesterMixin, PipelineTesterMixin -class ChronoEditPipelineFastTests(PipelineTesterMixin, unittest.TestCase): +class ChronoEditPipelineTesterConfig(BasePipelineTesterConfig): pipeline_class = ChronoEditPipeline - params = TEXT_TO_IMAGE_PARAMS - {"cross_attention_kwargs", "height", "width"} - 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( - [ - "num_inference_steps", - "generator", - "latents", - "return_dict", - "callback_on_step_end", - "callback_on_step_end_tensor_inputs", - ] + required_input_params_in_call_signature = frozenset( + ["image", "prompt", "negative_prompt", "guidance_scale", "prompt_embeds", "negative_prompt_embeds"] + ) + batch_input_params = frozenset(["prompt"]) + output_shape = (5, 3, 16, 16) + # ChronoEdit 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"] ) - test_xformers_attention = False def get_dummy_components(self): torch.manual_seed(0) @@ -72,7 +61,9 @@ def get_dummy_components(self): # TODO: impl FlowDPMSolverMultistepScheduler scheduler = FlowMatchEulerDiscreteScheduler(shift=7.0) config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5") - text_encoder = T5EncoderModel(config) + # `eval()` because a directly constructed model stays in training mode, which leaves T5's + # dropout active and makes the pipeline outputs non-deterministic across calls. + text_encoder = T5EncoderModel(config).eval() tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5") torch.manual_seed(0) @@ -107,7 +98,7 @@ def get_dummy_components(self): torch.manual_seed(0) image_processor = CLIPImageProcessor(crop_size=32, size=32) - components = { + return { "transformer": transformer, "vae": vae, "scheduler": scheduler, @@ -116,43 +107,36 @@ def get_dummy_components(self): "image_encoder": image_encoder, "image_processor": image_processor, } - 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): image_height = 16 image_width = 16 image = Image.new("RGB", (image_width, image_height)) - inputs = { + return { "image": image, "prompt": "dance monkey", "negative_prompt": "negative", # TODO "height": image_height, "width": image_width, - "generator": generator, + "generator": self.get_generator(0), "num_inference_steps": 2, "guidance_scale": 6.0, "num_frames": 5, "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) +class TestChronoEditPipeline(ChronoEditPipelineTesterConfig, PipelineTesterMixin): + def test_inference(self): + # Run on CPU: the expected slice below is CPU-specific. + pipe = self.get_pipeline() - inputs = self.get_dummy_inputs(device) + inputs = self.get_dummy_inputs() video = pipe(**inputs).frames generated_video = video[0] - self.assertEqual(generated_video.shape, (5, 3, 16, 16)) + assert generated_video.shape == self.output_shape # fmt: off expected_slice = torch.tensor([0.4525, 0.4520, 0.4485, 0.4533, 0.4522, 0.4522, 0.4529, 0.4528, 0.5023, 0.5067, 0.5023, 0.5061, 0.5024, 0.4977, 0.5118, 0.5190]) @@ -160,18 +144,18 @@ def test_inference(self): generated_slice = generated_video.flatten() generated_slice = torch.cat([generated_slice[:8], generated_slice[-8:]]) - self.assertTrue(torch.allclose(generated_slice, expected_slice, atol=1e-3)) - - @unittest.skip("Test not supported") - def test_attention_slicing_forward_pass(self): - pass + assert_tensors_close(generated_slice, expected_slice, atol=1e-3) - @unittest.skip("TODO: revisit failing as it requires a very high threshold to pass") + @pytest.mark.skip("TODO: revisit failing as it requires a very high threshold to pass") def test_inference_batch_single_identical(self): pass - @unittest.skip( + @pytest.mark.skip( "ChronoEditPipeline has to run in mixed precision. Save/Load the entire pipeline in FP16 will result in errors" ) def test_save_load_float16(self): pass + + +class TestChronoEditPipelineMemory(ChronoEditPipelineTesterConfig, MemoryTesterMixin): + """Memory optimization tests (CPU offload, group offload, layerwise casting) for the ChronoEdit pipeline.""" diff --git a/tests/pipelines/cogview3/test_cogview3plus.py b/tests/pipelines/cogview3/test_cogview3plus.py index 374cb6a2a295..014298e12927 100644 --- a/tests/pipelines/cogview3/test_cogview3plus.py +++ b/tests/pipelines/cogview3/test_cogview3plus.py @@ -13,10 +13,8 @@ # limitations under the License. import gc -import inspect -import unittest -import numpy as np +import pytest import torch from transformers import AutoConfig, AutoTokenizer, T5EncoderModel @@ -30,35 +28,19 @@ 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 ( - PipelineTesterMixin, - to_np, -) +from ..testing_utils import BasePipelineTesterConfig, MemoryTesterMixin, PipelineTesterMixin enable_full_determinism() -class CogView3PlusPipelineFastTests(PipelineTesterMixin, unittest.TestCase): +class CogView3PlusPipelineTesterConfig(BasePipelineTesterConfig): pipeline_class = CogView3PlusPipeline - 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( - [ - "num_inference_steps", - "generator", - "latents", - "return_dict", - "callback_on_step_end", - "callback_on_step_end_tensor_inputs", - ] + required_input_params_in_call_signature = frozenset( + ["prompt", "height", "width", "guidance_scale", "negative_prompt", "prompt_embeds", "negative_prompt_embeds"] ) - test_xformers_attention = False - test_layerwise_casting = True - test_group_offloading = True + batch_input_params = frozenset(["prompt", "negative_prompt"]) + output_shape = (3, 16, 16) def get_dummy_components(self): torch.manual_seed(0) @@ -90,166 +72,64 @@ def get_dummy_components(self): torch.manual_seed(0) scheduler = CogVideoXDDIMScheduler() config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5") - text_encoder = T5EncoderModel(config) + # `eval()` because a directly constructed model stays in training mode, which leaves T5's + # dropout active and makes the pipeline outputs non-deterministic across calls. + 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) - inputs = { + def get_dummy_inputs(self): + return { "prompt": "dance monkey", "negative_prompt": "", - "generator": generator, + "generator": self.get_generator(0), "num_inference_steps": 2, "guidance_scale": 6.0, "height": 16, "width": 16, "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) +class TestCogView3PlusPipeline(CogView3PlusPipelineTesterConfig, PipelineTesterMixin): + def test_inference(self): + pipe = self.get_pipeline() - inputs = self.get_dummy_inputs(device) - image = pipe(**inputs)[0] + image = pipe(**self.get_dummy_inputs())[0] generated_image = image[0] - self.assertEqual(generated_image.shape, (3, 16, 16)) - expected_image = torch.randn(3, 16, 16) - max_diff = np.abs(generated_image - expected_image).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 + assert generated_image.shape == self.output_shape - inputs = self.get_dummy_inputs(torch_device) + 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) - # 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] + def test_encode_prompt_works_in_isolation(self): + super().test_encode_prompt_works_in_isolation(atol=1e-3, rtol=1e-3) - 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_encode_prompt_works_in_isolation(self): - return super().test_encode_prompt_works_in_isolation(atol=1e-3, rtol=1e-3) +class TestCogView3PlusPipelineMemory(CogView3PlusPipelineTesterConfig, MemoryTesterMixin): + """Memory optimization tests (CPU offload, group offload, layerwise casting) for the CogView3Plus pipeline.""" @slow @require_torch_accelerator -class CogView3PlusPipelineIntegrationTests(unittest.TestCase): +class TestCogView3PlusPipelineIntegration: 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) diff --git a/tests/pipelines/consisid/test_consisid.py b/tests/pipelines/consisid/test_consisid.py index b427eeea1d8c..1d7b11a9e8b2 100644 --- a/tests/pipelines/consisid/test_consisid.py +++ b/tests/pipelines/consisid/test_consisid.py @@ -13,10 +13,8 @@ # limitations under the License. import gc -import inspect -import unittest -import numpy as np +import pytest import torch from PIL import Image from transformers import AutoConfig, AutoTokenizer, T5EncoderModel @@ -32,35 +30,23 @@ 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 ( - PipelineTesterMixin, - to_np, -) +from ..testing_utils import BasePipelineTesterConfig, MemoryTesterMixin, PipelineTesterMixin enable_full_determinism() -class ConsisIDPipelineFastTests(PipelineTesterMixin, unittest.TestCase): +class ConsisIDPipelineTesterConfig(BasePipelineTesterConfig): pipeline_class = ConsisIDPipeline - 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( - [ - "num_inference_steps", - "generator", - "latents", - "return_dict", - "callback_on_step_end", - "callback_on_step_end_tensor_inputs", - ] + 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", "image"]) + output_shape = (8, 3, 16, 16) + # ConsisID 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"] ) - test_xformers_attention = False - test_layerwise_casting = True - test_group_offloading = True def get_dummy_components(self): torch.manual_seed(0) @@ -123,34 +109,30 @@ def get_dummy_components(self): torch.manual_seed(0) scheduler = DDIMScheduler() config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5") - text_encoder = T5EncoderModel(config) + # `eval()` because a directly constructed model stays in training mode, which leaves T5's + # dropout active and makes the pipeline outputs non-deterministic across calls. + 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): image_height = 16 image_width = 16 image = Image.new("RGB", (image_width, image_height)) id_vit_hidden = [torch.ones([1, 2, 2])] * 1 id_cond = torch.ones(1, 2) - inputs = { + return { "image": image, "prompt": "dance monkey", "negative_prompt": "", - "generator": generator, + "generator": self.get_generator(0), "num_inference_steps": 2, "guidance_scale": 6.0, "height": image_height, @@ -159,129 +141,24 @@ def get_dummy_inputs(self, device, seed=0): "max_sequence_length": 16, "id_vit_hidden": id_vit_hidden, "id_cond": id_cond, + # 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 TestConsisIDPipeline(ConsisIDPipelineTesterConfig, PipelineTesterMixin): + def test_inference(self): + pipe = self.get_pipeline() - inputs = self.get_dummy_inputs(device) - video = pipe(**inputs).frames + video = pipe(**self.get_dummy_inputs()).frames generated_video = video[0] - self.assertEqual(generated_video.shape, (8, 3, 16, 16)) - expected_video = torch.randn(8, 3, 16, 16) - 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 + assert generated_video.shape == self.output_shape - # 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() - for key in components: - if "text_encoder" in key and hasattr(components[key], "eval"): - components[key].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) - - 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.4): - generator_device = "cpu" components = self.get_dummy_components() # The reason to modify it this way is because ConsisID Transformer limits the generation to resolutions used during initialization. @@ -293,12 +170,10 @@ def test_vae_tiling(self, expected_diff_max: float = 0.4): sample_width=16, ) - pipe = self.pipeline_class(**components) - pipe.to("cpu") - pipe.set_progress_bar_config(disable=None) + pipe = self.get_pipeline(**components) # 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] @@ -309,29 +184,29 @@ def test_vae_tiling(self, expected_diff_max: float = 0.4): tile_overlap_factor_height=1 / 12, tile_overlap_factor_width=1 / 12, ) - 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).max() < expected_diff_max, ( + "VAE tiling should not affect the inference results" ) +class TestConsisIDPipelineMemory(ConsisIDPipelineTesterConfig, MemoryTesterMixin): + """Memory optimization tests (CPU offload, group offload, layerwise casting) for the ConsisID pipeline.""" + + @slow @require_torch_accelerator -class ConsisIDPipelineIntegrationTests(unittest.TestCase): +class TestConsisIDPipelineIntegration: 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) diff --git a/tests/pipelines/consistency_models/test_consistency_models.py b/tests/pipelines/consistency_models/test_consistency_models.py index 0ab0c0af2588..2b4faf2595ab 100644 --- a/tests/pipelines/consistency_models/test_consistency_models.py +++ b/tests/pipelines/consistency_models/test_consistency_models.py @@ -1,7 +1,7 @@ import gc -import unittest import numpy as np +import pytest import torch from torch.backends.cuda import sdp_kernel @@ -14,6 +14,7 @@ from ...testing_utils import ( Expectations, + assert_tensors_close, backend_empty_cache, enable_full_determinism, nightly, @@ -22,19 +23,20 @@ torch_device, ) from ..pipeline_params import UNCONDITIONAL_IMAGE_GENERATION_BATCH_PARAMS, UNCONDITIONAL_IMAGE_GENERATION_PARAMS -from ..test_pipelines_common import PipelineTesterMixin +from ..testing_utils import BasePipelineTesterConfig, MemoryTesterMixin, PipelineTesterMixin enable_full_determinism() -class ConsistencyModelPipelineFastTests(PipelineTesterMixin, unittest.TestCase): +class ConsistencyModelPipelineTesterConfig(BasePipelineTesterConfig): pipeline_class = ConsistencyModelPipeline - params = UNCONDITIONAL_IMAGE_GENERATION_PARAMS - batch_params = UNCONDITIONAL_IMAGE_GENERATION_BATCH_PARAMS - - # Override required_optional_params to remove num_images_per_prompt - required_optional_params = frozenset( + required_input_params_in_call_signature = UNCONDITIONAL_IMAGE_GENERATION_PARAMS + batch_input_params = UNCONDITIONAL_IMAGE_GENERATION_BATCH_PARAMS + output_shape = (3, 32, 32) + # Unconditional generation: the pipeline takes a `batch_size` instead of `num_images_per_prompt`, and still + # exposes the legacy `callback` / `callback_steps` arguments. + optional_input_params = frozenset( [ "num_inference_steps", "generator", @@ -75,110 +77,102 @@ def get_dummy_components(self, class_cond=False): sigma_max=80.0, ) - components = { + return { "unet": unet, "scheduler": scheduler, } - 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 { "batch_size": 1, "num_inference_steps": None, "timesteps": [22, 0], - "generator": generator, - "output_type": "np", + "generator": self.get_generator(0), + # 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 TestConsistencyModelPipeline(ConsistencyModelPipelineTesterConfig, PipelineTesterMixin): def test_consistency_model_pipeline_multistep(self): - device = "cpu" # ensure determinism for the device-dependent torch.Generator - components = self.get_dummy_components() - pipe = ConsistencyModelPipeline(**components) - pipe = pipe.to(device) - pipe.set_progress_bar_config(disable=None) + # Run on CPU: the expected slice below is CPU-specific. + pipe = self.get_pipeline() - inputs = self.get_dummy_inputs(device) - image = pipe(**inputs).images - assert image.shape == (1, 32, 32, 3) + image = pipe(**self.get_dummy_inputs()).images + assert image.shape == (1, *self.output_shape) - image_slice = image[0, -3:, -3:, -1] - expected_slice = np.array([0.3572, 0.6273, 0.4031, 0.3961, 0.4321, 0.5730, 0.5266, 0.4780, 0.5004]) + image_slice = image[0, -1, -3:, -3:] + # fmt: off + expected_slice = torch.tensor([0.3572, 0.6273, 0.4031, 0.3961, 0.4321, 0.5730, 0.5266, 0.4780, 0.5004]) + # fmt: on - assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-3 + assert_tensors_close(image_slice.flatten(), expected_slice, atol=1e-3) def test_consistency_model_pipeline_multistep_class_cond(self): - device = "cpu" # ensure determinism for the device-dependent torch.Generator - components = self.get_dummy_components(class_cond=True) - pipe = ConsistencyModelPipeline(**components) - pipe = pipe.to(device) - pipe.set_progress_bar_config(disable=None) + # Run on CPU: the expected slice below is CPU-specific. + pipe = self.get_pipeline(**self.get_dummy_components(class_cond=True)) - inputs = self.get_dummy_inputs(device) + inputs = self.get_dummy_inputs() inputs["class_labels"] = 0 image = pipe(**inputs).images - assert image.shape == (1, 32, 32, 3) + assert image.shape == (1, *self.output_shape) - image_slice = image[0, -3:, -3:, -1] - expected_slice = np.array([0.3572, 0.6273, 0.4031, 0.3961, 0.4321, 0.5730, 0.5266, 0.4780, 0.5004]) + image_slice = image[0, -1, -3:, -3:] + # fmt: off + expected_slice = torch.tensor([0.3572, 0.6273, 0.4031, 0.3961, 0.4321, 0.5730, 0.5266, 0.4780, 0.5004]) + # fmt: on - assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-3 + assert_tensors_close(image_slice.flatten(), expected_slice, atol=1e-3) def test_consistency_model_pipeline_onestep(self): - device = "cpu" # ensure determinism for the device-dependent torch.Generator - components = self.get_dummy_components() - pipe = ConsistencyModelPipeline(**components) - pipe = pipe.to(device) - pipe.set_progress_bar_config(disable=None) + # Run on CPU: the expected slice below is CPU-specific. + pipe = self.get_pipeline() - inputs = self.get_dummy_inputs(device) + inputs = self.get_dummy_inputs() inputs["num_inference_steps"] = 1 inputs["timesteps"] = None image = pipe(**inputs).images - assert image.shape == (1, 32, 32, 3) + assert image.shape == (1, *self.output_shape) - image_slice = image[0, -3:, -3:, -1] - expected_slice = np.array([0.5004, 0.5004, 0.4994, 0.5008, 0.4976, 0.5018, 0.4990, 0.4982, 0.4987]) + image_slice = image[0, -1, -3:, -3:] + # fmt: off + expected_slice = torch.tensor([0.5004, 0.5004, 0.4994, 0.5008, 0.4976, 0.5018, 0.4990, 0.4982, 0.4987]) + # fmt: on - assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-3 + assert_tensors_close(image_slice.flatten(), expected_slice, atol=1e-3) def test_consistency_model_pipeline_onestep_class_cond(self): - device = "cpu" # ensure determinism for the device-dependent torch.Generator - components = self.get_dummy_components(class_cond=True) - pipe = ConsistencyModelPipeline(**components) - pipe = pipe.to(device) - pipe.set_progress_bar_config(disable=None) + # Run on CPU: the expected slice below is CPU-specific. + pipe = self.get_pipeline(**self.get_dummy_components(class_cond=True)) - inputs = self.get_dummy_inputs(device) + inputs = self.get_dummy_inputs() inputs["num_inference_steps"] = 1 inputs["timesteps"] = None inputs["class_labels"] = 0 image = pipe(**inputs).images - assert image.shape == (1, 32, 32, 3) + assert image.shape == (1, *self.output_shape) - image_slice = image[0, -3:, -3:, -1] - expected_slice = np.array([0.5004, 0.5004, 0.4994, 0.5008, 0.4976, 0.5018, 0.4990, 0.4982, 0.4987]) + image_slice = image[0, -1, -3:, -3:] + # fmt: off + expected_slice = torch.tensor([0.5004, 0.5004, 0.4994, 0.5008, 0.4976, 0.5018, 0.4990, 0.4982, 0.4987]) + # fmt: on - assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-3 + assert_tensors_close(image_slice.flatten(), expected_slice, atol=1e-3) + + +class TestConsistencyModelPipelineMemory(ConsistencyModelPipelineTesterConfig, MemoryTesterMixin): + """Memory optimization tests (CPU offload, group offload, layerwise casting) for the consistency model pipeline.""" @nightly @require_torch_accelerator -class ConsistencyModelPipelineSlowTests(unittest.TestCase): - def setUp(self): - super().setUp() +class TestConsistencyModelPipelineSlow: + @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) diff --git a/tests/pipelines/controlnet/test_controlnet.py b/tests/pipelines/controlnet/test_controlnet.py index 2e002d62aeac..b93d746ea9d5 100644 --- a/tests/pipelines/controlnet/test_controlnet.py +++ b/tests/pipelines/controlnet/test_controlnet.py @@ -14,10 +14,9 @@ # limitations under the License. import gc -import tempfile -import unittest import numpy as np +import pytest import torch from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer @@ -31,10 +30,10 @@ UNet2DConditionModel, ) from diffusers.pipelines.controlnet.pipeline_controlnet import MultiControlNetModel -from diffusers.utils.import_utils import is_xformers_available from diffusers.utils.torch_utils import randn_tensor from ...testing_utils import ( + assert_tensors_close, backend_empty_cache, backend_max_memory_allocated, backend_reset_max_memory_allocated, @@ -46,37 +45,19 @@ slow, torch_device, ) -from ..pipeline_params import ( - IMAGE_TO_IMAGE_IMAGE_PARAMS, - TEXT_TO_IMAGE_BATCH_PARAMS, - TEXT_TO_IMAGE_IMAGE_PARAMS, - TEXT_TO_IMAGE_PARAMS, -) -from ..test_pipelines_common import ( - IPAdapterTesterMixin, - PipelineKarrasSchedulerTesterMixin, - PipelineLatentTesterMixin, - PipelineTesterMixin, -) +from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_PARAMS +from ..stable_diffusion.ip_adapter_tester import IPAdapterTesterMixin +from ..testing_utils import BasePipelineTesterConfig, MemoryTesterMixin, PipelineTesterMixin enable_full_determinism() -class ControlNetPipelineFastTests( - IPAdapterTesterMixin, - PipelineLatentTesterMixin, - PipelineKarrasSchedulerTesterMixin, - PipelineTesterMixin, - unittest.TestCase, -): +class ControlNetPipelineTesterConfig(BasePipelineTesterConfig): pipeline_class = StableDiffusionControlNetPipeline - params = TEXT_TO_IMAGE_PARAMS - batch_params = TEXT_TO_IMAGE_BATCH_PARAMS - image_params = IMAGE_TO_IMAGE_IMAGE_PARAMS - image_latents_params = TEXT_TO_IMAGE_IMAGE_PARAMS - test_layerwise_casting = True - test_group_offloading = True + required_input_params_in_call_signature = TEXT_TO_IMAGE_PARAMS + batch_input_params = TEXT_TO_IMAGE_BATCH_PARAMS + output_shape = (3, 64, 64) def get_dummy_components(self, time_cond_proj_dim=None): torch.manual_seed(0) @@ -135,7 +116,7 @@ def get_dummy_components(self, time_cond_proj_dim=None): text_encoder = CLIPTextModel(text_encoder_config) tokenizer = CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip") - components = { + return { "unet": unet, "controlnet": controlnet, "scheduler": scheduler, @@ -146,112 +127,81 @@ def get_dummy_components(self, time_cond_proj_dim=None): "feature_extractor": None, "image_encoder": None, } - 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): + # The conditioning image is drawn from the same generator that is handed to the pipeline, so the pipeline + # sees an already-advanced generator state — keep the order to stay comparable with the expected slices. + generator = self.get_generator(0) controlnet_embedder_scale_factor = 2 image = randn_tensor( (1, 3, 32 * controlnet_embedder_scale_factor, 32 * controlnet_embedder_scale_factor), generator=generator, - device=torch.device(device), + device=torch.device(torch_device), ) - inputs = { + return { "prompt": "A painting of a squirrel eating a burger", "generator": generator, "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", "image": image, } - return inputs - - def test_attention_slicing_forward_pass(self): - return self._test_attention_slicing_forward_pass(expected_max_diff=2e-3) - def test_ip_adapter(self): - expected_pipe_slice = None - if torch_device == "cpu": - expected_pipe_slice = np.array([0.5245, 0.3353, 0.1784, 0.7556, 0.6239, 0.4673, 0.6931, 0.7487, 0.4644]) - return super().test_ip_adapter(expected_pipe_slice=expected_pipe_slice) +class TestControlNetPipeline(ControlNetPipelineTesterConfig, PipelineTesterMixin): + 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) - @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): - self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=2e-3) + def _run_lcm_pipeline(self, **extra_inputs): + pipe = self.get_pipeline(**self.get_dummy_components(time_cond_proj_dim=256)) + pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config) + pipe = pipe.to(torch_device) - def test_inference_batch_single_identical(self): - self._test_inference_batch_single_identical(expected_max_diff=2e-3) + inputs = self.get_dummy_inputs() + inputs.update(extra_inputs) + image = pipe(**inputs).images - def test_controlnet_lcm(self): - device = "cpu" # ensure determinism for the device-dependent torch.Generator + assert image.shape == (1, *self.output_shape) - components = self.get_dummy_components(time_cond_proj_dim=256) - sd_pipe = StableDiffusionControlNetPipeline(**components) - sd_pipe.scheduler = LCMScheduler.from_config(sd_pipe.scheduler.config) - sd_pipe = sd_pipe.to(torch_device) - sd_pipe.set_progress_bar_config(disable=None) + image_slice = image[0, -1, -3:, -3:] + # fmt: off + expected_slice = torch.tensor([0.52700454, 0.3930534, 0.25509018, 0.7132304, 0.53696585, 0.46568912, 0.7095368, 0.7059624, 0.4744786]) + # fmt: on - inputs = self.get_dummy_inputs(device) - output = sd_pipe(**inputs) - image = output.images + assert_tensors_close(image_slice.flatten().cpu(), expected_slice, atol=1e-2) - image_slice = image[0, -3:, -3:, -1] - - assert image.shape == (1, 64, 64, 3) - expected_slice = np.array( - [0.52700454, 0.3930534, 0.25509018, 0.7132304, 0.53696585, 0.46568912, 0.7095368, 0.7059624, 0.4744786] - ) - - assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2 + def test_controlnet_lcm(self): + self._run_lcm_pipeline() def test_controlnet_lcm_custom_timesteps(self): - device = "cpu" # ensure determinism for the device-dependent torch.Generator - - components = self.get_dummy_components(time_cond_proj_dim=256) - sd_pipe = StableDiffusionControlNetPipeline(**components) - sd_pipe.scheduler = LCMScheduler.from_config(sd_pipe.scheduler.config) - sd_pipe = sd_pipe.to(torch_device) - sd_pipe.set_progress_bar_config(disable=None) - - inputs = self.get_dummy_inputs(device) - del inputs["num_inference_steps"] - inputs["timesteps"] = [999, 499] - output = sd_pipe(**inputs) - image = output.images - - image_slice = image[0, -3:, -3:, -1] - - assert image.shape == (1, 64, 64, 3) - expected_slice = np.array( - [0.52700454, 0.3930534, 0.25509018, 0.7132304, 0.53696585, 0.46568912, 0.7095368, 0.7059624, 0.4744786] - ) - - assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2 + # `[999, 499]` is the schedule `num_inference_steps=2` produces, so this lands on the same slice as + # `test_controlnet_lcm`. `num_inference_steps` has to be `None` when `timesteps` is passed explicitly. + self._run_lcm_pipeline(num_inference_steps=None, timesteps=[999, 499]) def test_encode_prompt_works_in_isolation(self): 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) -class StableDiffusionMultiControlNetPipelineFastTests( - IPAdapterTesterMixin, PipelineTesterMixin, PipelineKarrasSchedulerTesterMixin, unittest.TestCase -): +class TestControlNetPipelineIPAdapter(ControlNetPipelineTesterConfig, IPAdapterTesterMixin): + """IP-Adapter tests for the Stable Diffusion ControlNet pipeline.""" + + +class TestControlNetPipelineMemory(ControlNetPipelineTesterConfig, MemoryTesterMixin): + """Memory optimization tests (CPU offload, group offload, layerwise casting) for the ControlNet pipeline.""" + + +class StableDiffusionMultiControlNetPipelineTesterConfig(BasePipelineTesterConfig): pipeline_class = StableDiffusionControlNetPipeline - params = TEXT_TO_IMAGE_PARAMS - batch_params = TEXT_TO_IMAGE_BATCH_PARAMS - image_params = frozenset([]) # TO_DO: add image_params once refactored VaeImageProcessor.preprocess + required_input_params_in_call_signature = TEXT_TO_IMAGE_PARAMS + batch_input_params = TEXT_TO_IMAGE_BATCH_PARAMS + output_shape = (3, 64, 64) def get_dummy_components(self): torch.manual_seed(0) @@ -331,7 +281,7 @@ def init_weights(m): controlnet = MultiControlNetModel([controlnet1, controlnet2]) - components = { + return { "unet": unet, "controlnet": controlnet, "scheduler": scheduler, @@ -342,155 +292,121 @@ def init_weights(m): "feature_extractor": None, "image_encoder": None, } - 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): + # The conditioning images are drawn from the same generator that is handed to the pipeline, so the pipeline + # sees an already-advanced generator state — keep the order to stay comparable across runs. + generator = self.get_generator(0) controlnet_embedder_scale_factor = 2 - images = [ randn_tensor( (1, 3, 32 * controlnet_embedder_scale_factor, 32 * controlnet_embedder_scale_factor), generator=generator, - device=torch.device(device), + device=torch.device(torch_device), ), randn_tensor( (1, 3, 32 * controlnet_embedder_scale_factor, 32 * controlnet_embedder_scale_factor), generator=generator, - device=torch.device(device), + device=torch.device(torch_device), ), ] - inputs = { + return { "prompt": "A painting of a squirrel eating a burger", "generator": generator, "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", "image": images, } - return inputs +class TestStableDiffusionMultiControlNetPipeline( + StableDiffusionMultiControlNetPipelineTesterConfig, PipelineTesterMixin +): def test_control_guidance_switch(self): - components = self.get_dummy_components() - pipe = self.pipeline_class(**components) - pipe.to(torch_device) + pipe = self.get_pipeline().to(torch_device) scale = 10.0 steps = 4 - inputs = self.get_dummy_inputs(torch_device) - inputs["num_inference_steps"] = steps - inputs["controlnet_conditioning_scale"] = scale - output_1 = pipe(**inputs)[0] + def run(**extra): + inputs = self.get_dummy_inputs() + inputs["num_inference_steps"] = steps + inputs["controlnet_conditioning_scale"] = scale + return pipe(**inputs, **extra)[0] - inputs = self.get_dummy_inputs(torch_device) - inputs["num_inference_steps"] = steps - inputs["controlnet_conditioning_scale"] = scale - output_2 = pipe(**inputs, control_guidance_start=0.1, control_guidance_end=0.2)[0] + output_1 = run() + output_2 = run(control_guidance_start=0.1, control_guidance_end=0.2) + output_3 = run(control_guidance_start=[0.1, 0.3], control_guidance_end=[0.2, 0.7]) + output_4 = run(control_guidance_start=0.4, control_guidance_end=[0.5, 0.8]) - inputs = self.get_dummy_inputs(torch_device) - inputs["num_inference_steps"] = steps - inputs["controlnet_conditioning_scale"] = scale - output_3 = pipe(**inputs, control_guidance_start=[0.1, 0.3], control_guidance_end=[0.2, 0.7])[0] + # make sure that all outputs are different + assert (output_1 - output_2).abs().sum() > 1e-3 + assert (output_1 - output_3).abs().sum() > 1e-3 + assert (output_1 - output_4).abs().sum() > 1e-3 - inputs = self.get_dummy_inputs(torch_device) - inputs["num_inference_steps"] = steps - inputs["controlnet_conditioning_scale"] = scale - output_4 = pipe(**inputs, control_guidance_start=0.4, control_guidance_end=[0.5, 0.8])[0] + 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) - # make sure that all outputs are different - assert np.sum(np.abs(output_1 - output_2)) > 1e-3 - assert np.sum(np.abs(output_1 - output_3)) > 1e-3 - assert np.sum(np.abs(output_1 - output_4)) > 1e-3 - - def test_attention_slicing_forward_pass(self): - return self._test_attention_slicing_forward_pass(expected_max_diff=2e-3) - - @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): - self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=2e-3) - - def test_inference_batch_single_identical(self): - self._test_inference_batch_single_identical(expected_max_diff=2e-3) - - def test_ip_adapter(self): - expected_pipe_slice = None - if torch_device == "cpu": - expected_pipe_slice = np.array([0.2395, 0.3421, 0.4023, 0.5345, 0.3496, 0.2402, 0.4645, 0.4563, 0.3786]) - return super().test_ip_adapter(expected_pipe_slice=expected_pipe_slice) - - def test_save_pretrained_raise_not_implemented_exception(self): - components = self.get_dummy_components() - pipe = self.pipeline_class(**components) - pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) - with tempfile.TemporaryDirectory() as tmpdir: - try: - # save_pretrained is not implemented for Multi-ControlNet - pipe.save_pretrained(tmpdir) - except NotImplementedError: - pass + def test_save_pretrained_raise_not_implemented_exception(self, tmp_path): + pipe = self.get_pipeline().to(torch_device) + try: + # save_pretrained is not implemented for Multi-ControlNet + pipe.save_pretrained(tmp_path) + except NotImplementedError: + pass def test_inference_multiple_prompt_input(self): - device = "cpu" - - components = self.get_dummy_components() - sd_pipe = StableDiffusionControlNetPipeline(**components) - sd_pipe = sd_pipe.to(torch_device) - sd_pipe.set_progress_bar_config(disable=None) + pipe = self.get_pipeline().to(torch_device) - inputs = self.get_dummy_inputs(device) + inputs = self.get_dummy_inputs() inputs["prompt"] = [inputs["prompt"], inputs["prompt"]] inputs["image"] = [inputs["image"], inputs["image"]] - output = sd_pipe(**inputs) - image = output.images + image = pipe(**inputs).images - assert image.shape == (2, 64, 64, 3) + assert image.shape == (2, *self.output_shape) image_1, image_2 = image # make sure that the outputs are different - assert np.sum(np.abs(image_1 - image_2)) > 1e-3 + assert (image_1 - image_2).abs().sum() > 1e-3 # multiple prompts, single image conditioning - inputs = self.get_dummy_inputs(device) + inputs = self.get_dummy_inputs() inputs["prompt"] = [inputs["prompt"], inputs["prompt"]] - output_1 = sd_pipe(**inputs) + output_1 = pipe(**inputs) - assert np.abs(image - output_1.images).max() < 1e-3 + assert (image - output_1.images).abs().max() < 1e-3 # multiple prompts, multiple image conditioning - inputs = self.get_dummy_inputs(device) - inputs["prompt"] = [inputs["prompt"], inputs["prompt"], inputs["prompt"], inputs["prompt"]] - inputs["image"] = [inputs["image"], inputs["image"], inputs["image"], inputs["image"]] - output_2 = sd_pipe(**inputs) - image = output_2.images + inputs = self.get_dummy_inputs() + inputs["prompt"] = [inputs["prompt"]] * 4 + inputs["image"] = [inputs["image"]] * 4 + image = pipe(**inputs).images - assert image.shape == (4, 64, 64, 3) + assert image.shape == (4, *self.output_shape) def test_encode_prompt_works_in_isolation(self): 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) -class StableDiffusionMultiControlNetOneModelPipelineFastTests( - IPAdapterTesterMixin, PipelineTesterMixin, PipelineKarrasSchedulerTesterMixin, unittest.TestCase +class TestStableDiffusionMultiControlNetPipelineIPAdapter( + StableDiffusionMultiControlNetPipelineTesterConfig, IPAdapterTesterMixin ): + """IP-Adapter tests for the Stable Diffusion Multi-ControlNet pipeline.""" + + +class StableDiffusionMultiControlNetOneModelPipelineTesterConfig(BasePipelineTesterConfig): pipeline_class = StableDiffusionControlNetPipeline - params = TEXT_TO_IMAGE_PARAMS - batch_params = TEXT_TO_IMAGE_BATCH_PARAMS - image_params = frozenset([]) # TO_DO: add image_params once refactored VaeImageProcessor.preprocess + required_input_params_in_call_signature = TEXT_TO_IMAGE_PARAMS + batch_input_params = TEXT_TO_IMAGE_BATCH_PARAMS + output_shape = (3, 64, 64) def get_dummy_components(self): torch.manual_seed(0) @@ -558,7 +474,7 @@ def init_weights(m): controlnet = MultiControlNetModel([controlnet]) - components = { + return { "unet": unet, "controlnet": controlnet, "scheduler": scheduler, @@ -569,121 +485,90 @@ def init_weights(m): "feature_extractor": None, "image_encoder": None, } - 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): + # The conditioning image is drawn from the same generator that is handed to the pipeline, so the pipeline + # sees an already-advanced generator state — keep the order to stay comparable across runs. + generator = self.get_generator(0) controlnet_embedder_scale_factor = 2 - images = [ randn_tensor( (1, 3, 32 * controlnet_embedder_scale_factor, 32 * controlnet_embedder_scale_factor), generator=generator, - device=torch.device(device), + device=torch.device(torch_device), ), ] - inputs = { + return { "prompt": "A painting of a squirrel eating a burger", "generator": generator, "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", "image": images, } - return inputs +class TestStableDiffusionMultiControlNetOneModelPipeline( + StableDiffusionMultiControlNetOneModelPipelineTesterConfig, PipelineTesterMixin +): def test_control_guidance_switch(self): - components = self.get_dummy_components() - pipe = self.pipeline_class(**components) - pipe.to(torch_device) + pipe = self.get_pipeline().to(torch_device) scale = 10.0 steps = 4 - inputs = self.get_dummy_inputs(torch_device) - inputs["num_inference_steps"] = steps - inputs["controlnet_conditioning_scale"] = scale - output_1 = pipe(**inputs)[0] - - inputs = self.get_dummy_inputs(torch_device) - inputs["num_inference_steps"] = steps - inputs["controlnet_conditioning_scale"] = scale - output_2 = pipe(**inputs, control_guidance_start=0.1, control_guidance_end=0.2)[0] - - inputs = self.get_dummy_inputs(torch_device) - inputs["num_inference_steps"] = steps - inputs["controlnet_conditioning_scale"] = scale - output_3 = pipe( - **inputs, - control_guidance_start=[0.1], - control_guidance_end=[0.2], - )[0] - - inputs = self.get_dummy_inputs(torch_device) - inputs["num_inference_steps"] = steps - inputs["controlnet_conditioning_scale"] = scale - output_4 = pipe(**inputs, control_guidance_start=0.4, control_guidance_end=[0.5])[0] + def run(**extra): + inputs = self.get_dummy_inputs() + inputs["num_inference_steps"] = steps + inputs["controlnet_conditioning_scale"] = scale + return pipe(**inputs, **extra)[0] + + output_1 = run() + output_2 = run(control_guidance_start=0.1, control_guidance_end=0.2) + output_3 = run(control_guidance_start=[0.1], control_guidance_end=[0.2]) + output_4 = run(control_guidance_start=0.4, control_guidance_end=[0.5]) # make sure that all outputs are different - assert np.sum(np.abs(output_1 - output_2)) > 1e-3 - assert np.sum(np.abs(output_1 - output_3)) > 1e-3 - assert np.sum(np.abs(output_1 - output_4)) > 1e-3 - - def test_attention_slicing_forward_pass(self): - return self._test_attention_slicing_forward_pass(expected_max_diff=2e-3) - - @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): - self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=2e-3) - - def test_inference_batch_single_identical(self): - self._test_inference_batch_single_identical(expected_max_diff=2e-3) - - def test_ip_adapter(self): - expected_pipe_slice = None - if torch_device == "cpu": - expected_pipe_slice = np.array([0.5234, 0.3198, 0.1596, 0.8201, 0.6316, 0.4566, 0.7209, 0.7761, 0.4757]) - return super().test_ip_adapter(expected_pipe_slice=expected_pipe_slice) - - def test_save_pretrained_raise_not_implemented_exception(self): - components = self.get_dummy_components() - pipe = self.pipeline_class(**components) - pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) - with tempfile.TemporaryDirectory() as tmpdir: - try: - # save_pretrained is not implemented for Multi-ControlNet - pipe.save_pretrained(tmpdir) - except NotImplementedError: - pass + assert (output_1 - output_2).abs().sum() > 1e-3 + assert (output_1 - output_3).abs().sum() > 1e-3 + assert (output_1 - output_4).abs().sum() > 1e-3 + + 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) + + def test_save_pretrained_raise_not_implemented_exception(self, tmp_path): + pipe = self.get_pipeline().to(torch_device) + try: + # save_pretrained is not implemented for Multi-ControlNet + pipe.save_pretrained(tmp_path) + except NotImplementedError: + pass def test_encode_prompt_works_in_isolation(self): 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) + + +class TestStableDiffusionMultiControlNetOneModelPipelineIPAdapter( + StableDiffusionMultiControlNetOneModelPipelineTesterConfig, IPAdapterTesterMixin +): + """IP-Adapter tests for the Stable Diffusion Multi-ControlNet (single model) pipeline.""" @slow @require_torch_accelerator -class ControlNetPipelineSlowTests(unittest.TestCase): - def setUp(self): - super().setUp() +class TestControlNetPipelineSlow: + @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) @@ -1032,14 +917,12 @@ def test_v11_shuffle_global_pool_conditions(self): @slow @require_torch_accelerator -class StableDiffusionMultiControlNetPipelineSlowTests(unittest.TestCase): - def setUp(self): - super().setUp() +class TestStableDiffusionMultiControlNetPipelineSlow: + @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) diff --git a/tests/pipelines/controlnet/test_controlnet_img2img.py b/tests/pipelines/controlnet/test_controlnet_img2img.py index b42730d73bb3..86c7f2b88bba 100644 --- a/tests/pipelines/controlnet/test_controlnet_img2img.py +++ b/tests/pipelines/controlnet/test_controlnet_img2img.py @@ -17,10 +17,9 @@ import gc import random -import tempfile -import unittest import numpy as np +import pytest import torch from PIL import Image from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer @@ -34,7 +33,6 @@ ) from diffusers.pipelines.controlnet.pipeline_controlnet import MultiControlNetModel from diffusers.utils import load_image -from diffusers.utils.import_utils import is_xformers_available from diffusers.utils.torch_utils import randn_tensor from ...testing_utils import ( @@ -46,34 +44,19 @@ slow, torch_device, ) -from ..pipeline_params import ( - IMAGE_TO_IMAGE_IMAGE_PARAMS, - TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS, - TEXT_GUIDED_IMAGE_VARIATION_PARAMS, -) -from ..test_pipelines_common import ( - IPAdapterTesterMixin, - PipelineKarrasSchedulerTesterMixin, - PipelineLatentTesterMixin, - PipelineTesterMixin, -) +from ..pipeline_params import TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS, TEXT_GUIDED_IMAGE_VARIATION_PARAMS +from ..stable_diffusion.ip_adapter_tester import IPAdapterTesterMixin +from ..testing_utils import BasePipelineTesterConfig, MemoryTesterMixin, PipelineTesterMixin enable_full_determinism() -class ControlNetImg2ImgPipelineFastTests( - IPAdapterTesterMixin, - PipelineLatentTesterMixin, - PipelineKarrasSchedulerTesterMixin, - PipelineTesterMixin, - unittest.TestCase, -): +class ControlNetImg2ImgPipelineTesterConfig(BasePipelineTesterConfig): pipeline_class = StableDiffusionControlNetImg2ImgPipeline - params = TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {"height", "width"} - batch_params = TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS - image_params = IMAGE_TO_IMAGE_IMAGE_PARAMS.union({"control_image"}) - image_latents_params = IMAGE_TO_IMAGE_IMAGE_PARAMS + required_input_params_in_call_signature = TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {"height", "width"} + batch_input_params = TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS + output_shape = (3, 64, 64) def get_dummy_components(self): torch.manual_seed(0) @@ -131,7 +114,7 @@ def get_dummy_components(self): text_encoder = CLIPTextModel(text_encoder_config) tokenizer = CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip") - components = { + return { "unet": unet, "controlnet": controlnet, "scheduler": scheduler, @@ -142,69 +125,59 @@ def get_dummy_components(self): "feature_extractor": None, "image_encoder": None, } - 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): + # The control image is drawn from the same generator that is handed to the pipeline, so the pipeline sees an + # already-advanced generator state — keep the order to stay comparable across runs. + generator = self.get_generator(0) controlnet_embedder_scale_factor = 2 control_image = randn_tensor( (1, 3, 32 * controlnet_embedder_scale_factor, 32 * controlnet_embedder_scale_factor), generator=generator, - device=torch.device(device), + device=torch.device(torch_device), ) - image = floats_tensor(control_image.shape, rng=random.Random(seed)).to(device) + image = floats_tensor(control_image.shape, rng=random.Random(0)).to(torch_device) image = image.cpu().permute(0, 2, 3, 1)[0] image = Image.fromarray(np.uint8(image)).convert("RGB").resize((64, 64)) - inputs = { + + return { "prompt": "A painting of a squirrel eating a burger", "generator": generator, "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", "image": image, "control_image": control_image, } - return inputs - - def test_attention_slicing_forward_pass(self): - return self._test_attention_slicing_forward_pass(expected_max_diff=2e-3) - def test_ip_adapter(self): - expected_pipe_slice = None - if torch_device == "cpu": - expected_pipe_slice = np.array([0.7051, 0.5142, 0.3632, 0.5896, 0.4731, 0.4075, 0.6059, 0.6865, 0.4203]) - return super().test_ip_adapter(expected_pipe_slice=expected_pipe_slice) - - @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): - self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=2e-3) - - def test_inference_batch_single_identical(self): - self._test_inference_batch_single_identical(expected_max_diff=2e-3) +class TestControlNetImg2ImgPipeline(ControlNetImg2ImgPipelineTesterConfig, PipelineTesterMixin): + 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) def test_encode_prompt_works_in_isolation(self): 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) -class StableDiffusionMultiControlNetPipelineFastTests( - IPAdapterTesterMixin, PipelineTesterMixin, PipelineKarrasSchedulerTesterMixin, unittest.TestCase -): +class TestControlNetImg2ImgPipelineIPAdapter(ControlNetImg2ImgPipelineTesterConfig, IPAdapterTesterMixin): + """IP-Adapter tests for the Stable Diffusion ControlNet img2img pipeline.""" + + +class TestControlNetImg2ImgPipelineMemory(ControlNetImg2ImgPipelineTesterConfig, MemoryTesterMixin): + """Memory optimization tests (CPU offload, group offload, layerwise casting) for the ControlNet img2img pipeline.""" + + +class StableDiffusionMultiControlNetImg2ImgPipelineTesterConfig(BasePipelineTesterConfig): pipeline_class = StableDiffusionControlNetImg2ImgPipeline - params = TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {"height", "width"} - batch_params = TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS - image_params = frozenset([]) # TO_DO: add image_params once refactored VaeImageProcessor.preprocess + required_input_params_in_call_signature = TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {"height", "width"} + batch_input_params = TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS + output_shape = (3, 64, 64) def get_dummy_components(self): torch.manual_seed(0) @@ -284,7 +257,7 @@ def init_weights(m): controlnet = MultiControlNetModel([controlnet1, controlnet2]) - components = { + return { "unet": unet, "controlnet": controlnet, "scheduler": scheduler, @@ -295,126 +268,100 @@ def init_weights(m): "feature_extractor": None, "image_encoder": None, } - 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): + # The control images are drawn from the same generator that is handed to the pipeline, so the pipeline sees + # an already-advanced generator state — keep the order to stay comparable across runs. + generator = self.get_generator(0) controlnet_embedder_scale_factor = 2 - control_image = [ randn_tensor( (1, 3, 32 * controlnet_embedder_scale_factor, 32 * controlnet_embedder_scale_factor), generator=generator, - device=torch.device(device), + device=torch.device(torch_device), ), randn_tensor( (1, 3, 32 * controlnet_embedder_scale_factor, 32 * controlnet_embedder_scale_factor), generator=generator, - device=torch.device(device), + device=torch.device(torch_device), ), ] - image = floats_tensor(control_image[0].shape, rng=random.Random(seed)).to(device) + image = floats_tensor(control_image[0].shape, rng=random.Random(0)).to(torch_device) image = image.cpu().permute(0, 2, 3, 1)[0] image = Image.fromarray(np.uint8(image)).convert("RGB").resize((64, 64)) - inputs = { + + return { "prompt": "A painting of a squirrel eating a burger", "generator": generator, "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", "image": image, "control_image": control_image, } - return inputs +class TestStableDiffusionMultiControlNetImg2ImgPipeline( + StableDiffusionMultiControlNetImg2ImgPipelineTesterConfig, PipelineTesterMixin +): def test_control_guidance_switch(self): - components = self.get_dummy_components() - pipe = self.pipeline_class(**components) - pipe.to(torch_device) + pipe = self.get_pipeline().to(torch_device) scale = 10.0 steps = 4 - inputs = self.get_dummy_inputs(torch_device) - inputs["num_inference_steps"] = steps - inputs["controlnet_conditioning_scale"] = scale - output_1 = pipe(**inputs)[0] + def run(**extra): + inputs = self.get_dummy_inputs() + inputs["num_inference_steps"] = steps + inputs["controlnet_conditioning_scale"] = scale + return pipe(**inputs, **extra)[0] - inputs = self.get_dummy_inputs(torch_device) - inputs["num_inference_steps"] = steps - inputs["controlnet_conditioning_scale"] = scale - output_2 = pipe(**inputs, control_guidance_start=0.1, control_guidance_end=0.2)[0] + output_1 = run() + output_2 = run(control_guidance_start=0.1, control_guidance_end=0.2) + output_3 = run(control_guidance_start=[0.1, 0.3], control_guidance_end=[0.2, 0.7]) + output_4 = run(control_guidance_start=0.4, control_guidance_end=[0.5, 0.8]) - inputs = self.get_dummy_inputs(torch_device) - inputs["num_inference_steps"] = steps - inputs["controlnet_conditioning_scale"] = scale - output_3 = pipe(**inputs, control_guidance_start=[0.1, 0.3], control_guidance_end=[0.2, 0.7])[0] + # make sure that all outputs are different + assert (output_1 - output_2).abs().sum() > 1e-3 + assert (output_1 - output_3).abs().sum() > 1e-3 + assert (output_1 - output_4).abs().sum() > 1e-3 - inputs = self.get_dummy_inputs(torch_device) - inputs["num_inference_steps"] = steps - inputs["controlnet_conditioning_scale"] = scale - output_4 = pipe(**inputs, control_guidance_start=0.4, control_guidance_end=[0.5, 0.8])[0] + 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) - # make sure that all outputs are different - assert np.sum(np.abs(output_1 - output_2)) > 1e-3 - assert np.sum(np.abs(output_1 - output_3)) > 1e-3 - assert np.sum(np.abs(output_1 - output_4)) > 1e-3 - - def test_attention_slicing_forward_pass(self): - return self._test_attention_slicing_forward_pass(expected_max_diff=2e-3) - - @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): - self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=2e-3) - - def test_inference_batch_single_identical(self): - self._test_inference_batch_single_identical(expected_max_diff=2e-3) - - def test_ip_adapter(self): - expected_pipe_slice = None - if torch_device == "cpu": - expected_pipe_slice = np.array([0.5241, 0.7318, 0.6612, 0.3910, 0.5169, 0.5151, 0.6959, 0.5856, 0.5151]) - return super().test_ip_adapter(expected_pipe_slice=expected_pipe_slice) - - def test_save_pretrained_raise_not_implemented_exception(self): - components = self.get_dummy_components() - pipe = self.pipeline_class(**components) - pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) - with tempfile.TemporaryDirectory() as tmpdir: - try: - # save_pretrained is not implemented for Multi-ControlNet - pipe.save_pretrained(tmpdir) - except NotImplementedError: - pass + def test_save_pretrained_raise_not_implemented_exception(self, tmp_path): + pipe = self.get_pipeline().to(torch_device) + try: + # save_pretrained is not implemented for Multi-ControlNet + pipe.save_pretrained(tmp_path) + except NotImplementedError: + pass def test_encode_prompt_works_in_isolation(self): 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) + + +class TestStableDiffusionMultiControlNetImg2ImgPipelineIPAdapter( + StableDiffusionMultiControlNetImg2ImgPipelineTesterConfig, IPAdapterTesterMixin +): + """IP-Adapter tests for the Stable Diffusion Multi-ControlNet img2img pipeline.""" @slow @require_torch_accelerator -class ControlNetImg2ImgPipelineSlowTests(unittest.TestCase): - def setUp(self): - super().setUp() +class TestControlNetImg2ImgPipelineSlow: + @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) diff --git a/tests/pipelines/controlnet/test_controlnet_inpaint.py b/tests/pipelines/controlnet/test_controlnet_inpaint.py index b0420409e86d..297b4d033865 100644 --- a/tests/pipelines/controlnet/test_controlnet_inpaint.py +++ b/tests/pipelines/controlnet/test_controlnet_inpaint.py @@ -17,10 +17,9 @@ import gc import random -import tempfile -import unittest import numpy as np +import pytest import torch from PIL import Image from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer @@ -34,7 +33,6 @@ ) from diffusers.pipelines.controlnet.pipeline_controlnet import MultiControlNetModel from diffusers.utils import load_image -from diffusers.utils.import_utils import is_xformers_available from diffusers.utils.torch_utils import randn_tensor from ...testing_utils import ( @@ -47,25 +45,18 @@ slow, torch_device, ) -from ..pipeline_params import ( - TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS, - TEXT_GUIDED_IMAGE_INPAINTING_PARAMS, - TEXT_TO_IMAGE_IMAGE_PARAMS, -) -from ..test_pipelines_common import PipelineKarrasSchedulerTesterMixin, PipelineLatentTesterMixin, PipelineTesterMixin +from ..pipeline_params import TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS, TEXT_GUIDED_IMAGE_INPAINTING_PARAMS +from ..testing_utils import BasePipelineTesterConfig, MemoryTesterMixin, PipelineTesterMixin enable_full_determinism() -class ControlNetInpaintPipelineFastTests( - PipelineLatentTesterMixin, PipelineKarrasSchedulerTesterMixin, PipelineTesterMixin, unittest.TestCase -): +class ControlNetInpaintPipelineTesterConfig(BasePipelineTesterConfig): pipeline_class = StableDiffusionControlNetInpaintPipeline - params = TEXT_GUIDED_IMAGE_INPAINTING_PARAMS - batch_params = TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS - image_params = frozenset({"control_image"}) # skip `image` and `mask` for now, only test for control_image - image_latents_params = TEXT_TO_IMAGE_IMAGE_PARAMS + required_input_params_in_call_signature = TEXT_GUIDED_IMAGE_INPAINTING_PARAMS + batch_input_params = TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS + output_shape = (3, 64, 64) def get_dummy_components(self): torch.manual_seed(0) @@ -120,7 +111,7 @@ def get_dummy_components(self): text_encoder = CLIPTextModel(text_encoder_config) tokenizer = CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip") - components = { + return { "unet": unet, "controlnet": controlnet, "scheduler": scheduler, @@ -131,65 +122,63 @@ def get_dummy_components(self): "feature_extractor": None, "image_encoder": None, } - 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): + # The control image is drawn from the same generator that is handed to the pipeline, so the pipeline sees an + # already-advanced generator state — keep the order to stay comparable across runs. + generator = self.get_generator(0) controlnet_embedder_scale_factor = 2 control_image = randn_tensor( (1, 3, 32 * controlnet_embedder_scale_factor, 32 * controlnet_embedder_scale_factor), generator=generator, - device=torch.device(device), + device=torch.device(torch_device), ) - init_image = floats_tensor((1, 3, 32, 32), rng=random.Random(seed)).to(device) + init_image = floats_tensor((1, 3, 32, 32), rng=random.Random(0)).to(torch_device) init_image = init_image.cpu().permute(0, 2, 3, 1)[0] image = Image.fromarray(np.uint8(init_image)).convert("RGB").resize((64, 64)) mask_image = Image.fromarray(np.uint8(init_image + 4)).convert("RGB").resize((64, 64)) - inputs = { + return { "prompt": "A painting of a squirrel eating a burger", "generator": generator, "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", "image": image, "mask_image": mask_image, "control_image": control_image, } - return inputs - def test_attention_slicing_forward_pass(self): - return self._test_attention_slicing_forward_pass(expected_max_diff=2e-3) +class ControlNetInpaintPipelineTests: + """Tests shared by the inpaint-UNet and the plain-UNet ("simple") ControlNet inpainting configs.""" - @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): - self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=2e-3) - - def test_inference_batch_single_identical(self): - self._test_inference_batch_single_identical(expected_max_diff=2e-3) + 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) def test_encode_prompt_works_in_isolation(self): 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) -class ControlNetSimpleInpaintPipelineFastTests(ControlNetInpaintPipelineFastTests): - pipeline_class = StableDiffusionControlNetInpaintPipeline - params = TEXT_GUIDED_IMAGE_INPAINTING_PARAMS - batch_params = TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS - image_params = frozenset([]) +class TestControlNetInpaintPipeline( + ControlNetInpaintPipelineTesterConfig, ControlNetInpaintPipelineTests, PipelineTesterMixin +): + pass + + +class TestControlNetInpaintPipelineMemory(ControlNetInpaintPipelineTesterConfig, MemoryTesterMixin): + """Memory optimization tests (CPU offload, group offload, layerwise casting) for the ControlNet inpaint pipeline.""" + + +class ControlNetSimpleInpaintPipelineTesterConfig(ControlNetInpaintPipelineTesterConfig): + """Same contract as `ControlNetInpaintPipelineTesterConfig`, but with a plain (non-inpainting) UNet.""" def get_dummy_components(self): torch.manual_seed(0) @@ -244,7 +233,7 @@ def get_dummy_components(self): text_encoder = CLIPTextModel(text_encoder_config) tokenizer = CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip") - components = { + return { "unet": unet, "controlnet": controlnet, "scheduler": scheduler, @@ -255,15 +244,19 @@ def get_dummy_components(self): "feature_extractor": None, "image_encoder": None, } - return components -class MultiControlNetInpaintPipelineFastTests( - PipelineTesterMixin, PipelineKarrasSchedulerTesterMixin, unittest.TestCase +class TestControlNetSimpleInpaintPipeline( + ControlNetSimpleInpaintPipelineTesterConfig, ControlNetInpaintPipelineTests, PipelineTesterMixin ): + pass + + +class MultiControlNetInpaintPipelineTesterConfig(BasePipelineTesterConfig): pipeline_class = StableDiffusionControlNetInpaintPipeline - params = TEXT_GUIDED_IMAGE_INPAINTING_PARAMS - batch_params = TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS + required_input_params_in_call_signature = TEXT_GUIDED_IMAGE_INPAINTING_PARAMS + batch_input_params = TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS + output_shape = (3, 64, 64) def get_dummy_components(self): torch.manual_seed(0) @@ -339,7 +332,7 @@ def init_weights(m): controlnet = MultiControlNetModel([controlnet1, controlnet2]) - components = { + return { "unet": unet, "controlnet": controlnet, "scheduler": scheduler, @@ -350,123 +343,94 @@ def init_weights(m): "feature_extractor": None, "image_encoder": None, } - 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): + # The control images are drawn from the same generator that is handed to the pipeline, so the pipeline sees + # an already-advanced generator state — keep the order to stay comparable across runs. + generator = self.get_generator(0) controlnet_embedder_scale_factor = 2 - control_image = [ randn_tensor( (1, 3, 32 * controlnet_embedder_scale_factor, 32 * controlnet_embedder_scale_factor), generator=generator, - device=torch.device(device), + device=torch.device(torch_device), ), randn_tensor( (1, 3, 32 * controlnet_embedder_scale_factor, 32 * controlnet_embedder_scale_factor), generator=generator, - device=torch.device(device), + device=torch.device(torch_device), ), ] - init_image = floats_tensor((1, 3, 32, 32), rng=random.Random(seed)).to(device) + init_image = floats_tensor((1, 3, 32, 32), rng=random.Random(0)).to(torch_device) init_image = init_image.cpu().permute(0, 2, 3, 1)[0] image = Image.fromarray(np.uint8(init_image)).convert("RGB").resize((64, 64)) mask_image = Image.fromarray(np.uint8(init_image + 4)).convert("RGB").resize((64, 64)) - inputs = { + return { "prompt": "A painting of a squirrel eating a burger", "generator": generator, "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", "image": image, "mask_image": mask_image, "control_image": control_image, } - return inputs +class TestMultiControlNetInpaintPipeline(MultiControlNetInpaintPipelineTesterConfig, PipelineTesterMixin): def test_control_guidance_switch(self): - components = self.get_dummy_components() - pipe = self.pipeline_class(**components) - pipe.to(torch_device) + pipe = self.get_pipeline().to(torch_device) scale = 10.0 steps = 4 - inputs = self.get_dummy_inputs(torch_device) - inputs["num_inference_steps"] = steps - inputs["controlnet_conditioning_scale"] = scale - output_1 = pipe(**inputs)[0] + def run(**extra): + inputs = self.get_dummy_inputs() + inputs["num_inference_steps"] = steps + inputs["controlnet_conditioning_scale"] = scale + return pipe(**inputs, **extra)[0] - inputs = self.get_dummy_inputs(torch_device) - inputs["num_inference_steps"] = steps - inputs["controlnet_conditioning_scale"] = scale - output_2 = pipe(**inputs, control_guidance_start=0.1, control_guidance_end=0.2)[0] + output_1 = run() + output_2 = run(control_guidance_start=0.1, control_guidance_end=0.2) + output_3 = run(control_guidance_start=[0.1, 0.3], control_guidance_end=[0.2, 0.7]) + output_4 = run(control_guidance_start=0.4, control_guidance_end=[0.5, 0.8]) - inputs = self.get_dummy_inputs(torch_device) - inputs["num_inference_steps"] = steps - inputs["controlnet_conditioning_scale"] = scale - output_3 = pipe(**inputs, control_guidance_start=[0.1, 0.3], control_guidance_end=[0.2, 0.7])[0] + # make sure that all outputs are different + assert (output_1 - output_2).abs().sum() > 1e-3 + assert (output_1 - output_3).abs().sum() > 1e-3 + assert (output_1 - output_4).abs().sum() > 1e-3 - inputs = self.get_dummy_inputs(torch_device) - inputs["num_inference_steps"] = steps - inputs["controlnet_conditioning_scale"] = scale - output_4 = pipe(**inputs, control_guidance_start=0.4, control_guidance_end=[0.5, 0.8])[0] + 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) - # make sure that all outputs are different - assert np.sum(np.abs(output_1 - output_2)) > 1e-3 - assert np.sum(np.abs(output_1 - output_3)) > 1e-3 - assert np.sum(np.abs(output_1 - output_4)) > 1e-3 - - def test_attention_slicing_forward_pass(self): - return self._test_attention_slicing_forward_pass(expected_max_diff=2e-3) - - @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): - self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=2e-3) - - def test_inference_batch_single_identical(self): - self._test_inference_batch_single_identical(expected_max_diff=2e-3) - - def test_save_pretrained_raise_not_implemented_exception(self): - components = self.get_dummy_components() - pipe = self.pipeline_class(**components) - pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) - with tempfile.TemporaryDirectory() as tmpdir: - try: - # save_pretrained is not implemented for Multi-ControlNet - pipe.save_pretrained(tmpdir) - except NotImplementedError: - pass + def test_save_pretrained_raise_not_implemented_exception(self, tmp_path): + pipe = self.get_pipeline().to(torch_device) + try: + # save_pretrained is not implemented for Multi-ControlNet + pipe.save_pretrained(tmp_path) + except NotImplementedError: + pass def test_encode_prompt_works_in_isolation(self): 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) @slow @require_torch_accelerator -class ControlNetInpaintPipelineSlowTests(unittest.TestCase): - def setUp(self): - super().setUp() +class TestControlNetInpaintPipelineSlow: + @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) diff --git a/tests/pipelines/controlnet/test_controlnet_inpaint_sdxl.py b/tests/pipelines/controlnet/test_controlnet_inpaint_sdxl.py index e5a5410a011a..ddc0216a7767 100644 --- a/tests/pipelines/controlnet/test_controlnet_inpaint_sdxl.py +++ b/tests/pipelines/controlnet/test_controlnet_inpaint_sdxl.py @@ -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 ( @@ -36,39 +36,24 @@ StableDiffusionXLControlNetInpaintPipeline, UNet2DConditionModel, ) -from diffusers.utils.import_utils import is_xformers_available -from ...testing_utils import ( - enable_full_determinism, - floats_tensor, - require_torch_accelerator, - torch_device, -) +from ...testing_utils import assert_tensors_close, enable_full_determinism, floats_tensor, torch_device from ..pipeline_params import ( - IMAGE_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_CALLBACK_CFG_PARAMS, - TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS, ) -from ..test_pipelines_common import ( - PipelineKarrasSchedulerTesterMixin, - PipelineLatentTesterMixin, - PipelineTesterMixin, -) +from ..testing_utils import BasePipelineTesterConfig, MemoryTesterMixin, PipelineTesterMixin enable_full_determinism() -class ControlNetPipelineSDXLFastTests( - PipelineLatentTesterMixin, PipelineKarrasSchedulerTesterMixin, PipelineTesterMixin, unittest.TestCase -): +class StableDiffusionXLControlNetInpaintPipelineTesterConfig(BasePipelineTesterConfig): pipeline_class = StableDiffusionXLControlNetInpaintPipeline - params = TEXT_TO_IMAGE_PARAMS - batch_params = TEXT_TO_IMAGE_BATCH_PARAMS - image_params = frozenset(IMAGE_TO_IMAGE_IMAGE_PARAMS.union({"mask_image", "control_image"})) - image_latents_params = TEXT_TO_IMAGE_IMAGE_PARAMS + required_input_params_in_call_signature = TEXT_TO_IMAGE_PARAMS + batch_input_params = TEXT_TO_IMAGE_BATCH_PARAMS + output_shape = (3, 64, 64) callback_cfg_params = TEXT_TO_IMAGE_CALLBACK_CFG_PARAMS.union( { "add_text_embeds", @@ -175,7 +160,7 @@ def get_dummy_components(self): size=224, ) - components = { + return { "unet": unet, "controlnet": controlnet, "scheduler": scheduler, @@ -187,25 +172,19 @@ def get_dummy_components(self): "image_encoder": image_encoder, "feature_extractor": feature_extractor, } - return components - - def get_dummy_inputs(self, device, seed=0, img_res=64): - if str(device).startswith("mps"): - generator = torch.manual_seed(seed) - else: - generator = torch.Generator(device=device).manual_seed(seed) + def get_dummy_inputs(self, img_res=64): # Get random floats in [0, 1] as image - image = floats_tensor((1, 3, 32, 32), rng=random.Random(seed)).to(device) + image = floats_tensor((1, 3, 32, 32), rng=random.Random(0)).to(torch_device) image = image.cpu().permute(0, 2, 3, 1)[0] mask_image = torch.ones_like(image) controlnet_embedder_scale_factor = 2 control_image = ( floats_tensor( (1, 3, 32 * controlnet_embedder_scale_factor, 32 * controlnet_embedder_scale_factor), - rng=random.Random(seed), + rng=random.Random(0), ) - .to(device) + .to(torch_device) .cpu() ) control_image = control_image.cpu().permute(0, 2, 3, 1)[0] @@ -218,139 +197,93 @@ def get_dummy_inputs(self, device, seed=0, img_res=64): mask_image = Image.fromarray(np.uint8(mask_image)).convert("L").resize((img_res, img_res)) control_image = Image.fromarray(np.uint8(control_image)).convert("RGB").resize((img_res, img_res)) - inputs = { + 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", "image": init_image, "mask_image": mask_image, "control_image": control_image, } - return inputs - - def test_attention_slicing_forward_pass(self): - return self._test_attention_slicing_forward_pass(expected_max_diff=2e-3) - - @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): - self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=2e-3) - - def test_inference_batch_single_identical(self): - self._test_inference_batch_single_identical(expected_max_diff=2e-3) - - @require_torch_accelerator - def test_stable_diffusion_xl_offloads(self): - pipes = [] - components = self.get_dummy_components() - sd_pipe = self.pipeline_class(**components).to(torch_device) - pipes.append(sd_pipe) - - components = self.get_dummy_components() - sd_pipe = self.pipeline_class(**components) - sd_pipe.enable_model_cpu_offload(device=torch_device) - pipes.append(sd_pipe) - - components = self.get_dummy_components() - sd_pipe = self.pipeline_class(**components) - sd_pipe.enable_sequential_cpu_offload(device=torch_device) - pipes.append(sd_pipe) - image_slices = [] - for pipe in pipes: - pipe.unet.set_default_attn_processor() - inputs = self.get_dummy_inputs(torch_device) - image = pipe(**inputs).images - - image_slices.append(image[0, -3:, -3:, -1].flatten()) +class TestStableDiffusionXLControlNetInpaintPipeline( + StableDiffusionXLControlNetInpaintPipelineTesterConfig, PipelineTesterMixin +): + 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) - assert np.abs(image_slices[0] - image_slices[1]).max() < 1e-3 - assert np.abs(image_slices[0] - image_slices[2]).max() < 1e-3 + @pytest.mark.skip("TODO(Patrick, Sayak) - skip for now as this requires more refiner tests") + def test_save_load_optional_components(self): + pass def test_stable_diffusion_xl_multi_prompts(self): - components = self.get_dummy_components() - sd_pipe = self.pipeline_class(**components).to(torch_device) + pipe = self.get_pipeline().to(torch_device) # forward with single prompt - inputs = self.get_dummy_inputs(torch_device) - output = sd_pipe(**inputs) - image_slice_1 = output.images[0, -3:, -3:, -1] + image_slice_1 = pipe(**self.get_dummy_inputs()).images[0, -1, -3:, -3:] # forward with same prompt duplicated - inputs = self.get_dummy_inputs(torch_device) + inputs = self.get_dummy_inputs() inputs["prompt_2"] = inputs["prompt"] - output = sd_pipe(**inputs) - image_slice_2 = output.images[0, -3:, -3:, -1] + image_slice_2 = pipe(**inputs).images[0, -1, -3:, -3:] # ensure the results are equal - assert np.abs(image_slice_1.flatten() - image_slice_2.flatten()).max() < 1e-4 + assert (image_slice_1 - image_slice_2).abs().max() < 1e-4 # forward with different prompt - inputs = self.get_dummy_inputs(torch_device) + inputs = self.get_dummy_inputs() inputs["prompt_2"] = "different prompt" - output = sd_pipe(**inputs) - image_slice_3 = output.images[0, -3:, -3:, -1] + image_slice_3 = pipe(**inputs).images[0, -1, -3:, -3:] # ensure the results are not equal - assert np.abs(image_slice_1.flatten() - image_slice_3.flatten()).max() > 1e-4 + assert (image_slice_1 - image_slice_3).abs().max() > 1e-4 # manually set a negative_prompt - inputs = self.get_dummy_inputs(torch_device) + inputs = self.get_dummy_inputs() inputs["negative_prompt"] = "negative prompt" - output = sd_pipe(**inputs) - image_slice_1 = output.images[0, -3:, -3:, -1] + image_slice_1 = pipe(**inputs).images[0, -1, -3:, -3:] # forward with same negative_prompt duplicated - inputs = self.get_dummy_inputs(torch_device) + inputs = self.get_dummy_inputs() inputs["negative_prompt"] = "negative prompt" inputs["negative_prompt_2"] = inputs["negative_prompt"] - output = sd_pipe(**inputs) - image_slice_2 = output.images[0, -3:, -3:, -1] + image_slice_2 = pipe(**inputs).images[0, -1, -3:, -3:] # ensure the results are equal - assert np.abs(image_slice_1.flatten() - image_slice_2.flatten()).max() < 1e-4 + assert (image_slice_1 - image_slice_2).abs().max() < 1e-4 # forward with different negative_prompt - inputs = self.get_dummy_inputs(torch_device) + inputs = self.get_dummy_inputs() inputs["negative_prompt"] = "negative prompt" inputs["negative_prompt_2"] = "different negative prompt" - output = sd_pipe(**inputs) - image_slice_3 = output.images[0, -3:, -3:, -1] + image_slice_3 = pipe(**inputs).images[0, -1, -3:, -3:] # ensure the results are not equal - assert np.abs(image_slice_1.flatten() - image_slice_3.flatten()).max() > 1e-4 + assert (image_slice_1 - image_slice_3).abs().max() > 1e-4 def test_controlnet_sdxl_guess(self): - device = "cpu" - - components = self.get_dummy_components() - - sd_pipe = self.pipeline_class(**components) - sd_pipe = sd_pipe.to(device) + # Run on CPU: the expected slice is CPU-specific. + pipe = self.get_pipeline() - sd_pipe.set_progress_bar_config(disable=None) - - inputs = self.get_dummy_inputs(device) + inputs = self.get_dummy_inputs() inputs["guess_mode"] = True - output = sd_pipe(**inputs) - image_slice = output.images[0, -3:, -3:, -1] + image_slice = pipe(**inputs).images[0, -1, -3:, -3:] - expected_slice = np.array( - [0.559845, 0.506214, 0.467439, 0.587851, 0.536138, 0.480375, 0.598892, 0.571167, 0.434700] - ) + # fmt: off + expected_slice = torch.tensor([0.559845, 0.506214, 0.467439, 0.587851, 0.536138, 0.480375, 0.598892, 0.571167, 0.434700]) + # fmt: on # make sure that it's equal - assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-4 + assert_tensors_close(image_slice.flatten().cpu(), expected_slice, atol=1e-4) - # TODO(Patrick, Sayak) - skip for now as this requires more refiner tests - def test_save_load_optional_components(self): - pass - def test_float16_inference(self): - super().test_float16_inference(expected_max_diff=5e-1) +class TestStableDiffusionXLControlNetInpaintPipelineMemory( + StableDiffusionXLControlNetInpaintPipelineTesterConfig, MemoryTesterMixin +): + """Memory optimization tests (CPU offload, group offload, layerwise casting) for the SDXL ControlNet inpaint pipeline.""" diff --git a/tests/pipelines/controlnet/test_controlnet_sdxl.py b/tests/pipelines/controlnet/test_controlnet_sdxl.py index 926a55f679c5..ae6ad2e54224 100644 --- a/tests/pipelines/controlnet/test_controlnet_sdxl.py +++ b/tests/pipelines/controlnet/test_controlnet_sdxl.py @@ -15,9 +15,9 @@ import copy import gc -import unittest import numpy as np +import pytest import torch from transformers import CLIPTextConfig, CLIPTextModel, CLIPTextModelWithProjection, CLIPTokenizer @@ -33,10 +33,10 @@ ) from diffusers.models.unets.unet_2d_blocks import UNetMidBlock2D from diffusers.pipelines.controlnet.pipeline_controlnet import MultiControlNetModel -from diffusers.utils.import_utils import is_xformers_available from diffusers.utils.torch_utils import randn_tensor from ...testing_utils import ( + assert_tensors_close, backend_empty_cache, enable_full_determinism, load_image, @@ -44,37 +44,19 @@ slow, torch_device, ) -from ..pipeline_params import ( - IMAGE_TO_IMAGE_IMAGE_PARAMS, - TEXT_TO_IMAGE_BATCH_PARAMS, - TEXT_TO_IMAGE_IMAGE_PARAMS, - TEXT_TO_IMAGE_PARAMS, -) -from ..test_pipelines_common import ( - IPAdapterTesterMixin, - PipelineKarrasSchedulerTesterMixin, - PipelineLatentTesterMixin, - PipelineTesterMixin, -) +from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_PARAMS +from ..stable_diffusion.ip_adapter_tester import IPAdapterTesterMixin +from ..testing_utils import BasePipelineTesterConfig, MemoryTesterMixin, PipelineTesterMixin enable_full_determinism() -class StableDiffusionXLControlNetPipelineFastTests( - IPAdapterTesterMixin, - PipelineLatentTesterMixin, - PipelineKarrasSchedulerTesterMixin, - PipelineTesterMixin, - unittest.TestCase, -): +class StableDiffusionXLControlNetPipelineTesterConfig(BasePipelineTesterConfig): pipeline_class = StableDiffusionXLControlNetPipeline - params = TEXT_TO_IMAGE_PARAMS - batch_params = TEXT_TO_IMAGE_BATCH_PARAMS - image_params = IMAGE_TO_IMAGE_IMAGE_PARAMS - image_latents_params = TEXT_TO_IMAGE_IMAGE_PARAMS - test_layerwise_casting = True - test_group_offloading = True + required_input_params_in_call_signature = TEXT_TO_IMAGE_PARAMS + batch_input_params = TEXT_TO_IMAGE_BATCH_PARAMS + output_shape = (3, 64, 64) def get_dummy_components(self, time_cond_proj_dim=None): torch.manual_seed(0) @@ -150,7 +132,7 @@ def get_dummy_components(self, time_cond_proj_dim=None): text_encoder_2 = CLIPTextModelWithProjection(text_encoder_config) tokenizer_2 = CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip") - components = { + return { "unet": unet, "controlnet": controlnet, "scheduler": scheduler, @@ -162,182 +144,123 @@ def get_dummy_components(self, time_cond_proj_dim=None): "feature_extractor": None, "image_encoder": None, } - 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): + # The conditioning image is drawn from the same generator that is handed to the pipeline, so the pipeline + # sees an already-advanced generator state — keep the order to stay comparable with the expected slices. + generator = self.get_generator(0) controlnet_embedder_scale_factor = 2 image = randn_tensor( (1, 3, 32 * controlnet_embedder_scale_factor, 32 * controlnet_embedder_scale_factor), generator=generator, - device=torch.device(device), + device=torch.device(torch_device), ) - inputs = { + return { "prompt": "A painting of a squirrel eating a burger", "generator": generator, "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", "image": image, } - return inputs - def test_attention_slicing_forward_pass(self): - return self._test_attention_slicing_forward_pass(expected_max_diff=2e-3) +class StableDiffusionXLControlNetPipelineTests: + """Tests shared by the SDXL ControlNet config and the SSD-1B one, which only differ in their components. - def test_ip_adapter(self, from_ssd1b=False, expected_pipe_slice=None): - if not from_ssd1b: - expected_pipe_slice = None - if torch_device == "cpu": - expected_pipe_slice = np.array( - [0.7153, 0.5634, 0.5697, 0.6209, 0.5700, 0.6044, 0.4283, 0.4552, 0.5271] - ) - return super().test_ip_adapter(expected_pipe_slice=expected_pipe_slice) + The two expected slices below are the only per-config values, so subclasses set them instead of overriding the + tests. + """ - @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): - self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=2e-3) + expected_guess_slice: torch.Tensor + expected_lcm_slice: torch.Tensor - def test_inference_batch_single_identical(self): - self._test_inference_batch_single_identical(expected_max_diff=2e-3) + 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) - @unittest.skip("We test this functionality elsewhere already.") + @pytest.mark.skip("We test this functionality elsewhere already.") def test_save_load_optional_components(self): pass - @require_torch_accelerator - def test_stable_diffusion_xl_offloads(self): - pipes = [] - components = self.get_dummy_components() - sd_pipe = self.pipeline_class(**components).to(torch_device) - pipes.append(sd_pipe) - - components = self.get_dummy_components() - sd_pipe = self.pipeline_class(**components) - sd_pipe.enable_model_cpu_offload(device=torch_device) - pipes.append(sd_pipe) - - components = self.get_dummy_components() - sd_pipe = self.pipeline_class(**components) - sd_pipe.enable_sequential_cpu_offload(device=torch_device) - pipes.append(sd_pipe) - - image_slices = [] - for pipe in pipes: - pipe.unet.set_default_attn_processor() - - inputs = self.get_dummy_inputs(torch_device) - image = pipe(**inputs).images - - image_slices.append(image[0, -3:, -3:, -1].flatten()) - - assert np.abs(image_slices[0] - image_slices[1]).max() < 1e-3 - assert np.abs(image_slices[0] - image_slices[2]).max() < 1e-3 - def test_stable_diffusion_xl_multi_prompts(self): - components = self.get_dummy_components() - sd_pipe = self.pipeline_class(**components).to(torch_device) + pipe = self.get_pipeline().to(torch_device) # forward with single prompt - inputs = self.get_dummy_inputs(torch_device) - output = sd_pipe(**inputs) - image_slice_1 = output.images[0, -3:, -3:, -1] + image_slice_1 = pipe(**self.get_dummy_inputs()).images[0, -1, -3:, -3:] # forward with same prompt duplicated - inputs = self.get_dummy_inputs(torch_device) + inputs = self.get_dummy_inputs() inputs["prompt_2"] = inputs["prompt"] - output = sd_pipe(**inputs) - image_slice_2 = output.images[0, -3:, -3:, -1] + image_slice_2 = pipe(**inputs).images[0, -1, -3:, -3:] # ensure the results are equal - assert np.abs(image_slice_1.flatten() - image_slice_2.flatten()).max() < 1e-4 + assert (image_slice_1 - image_slice_2).abs().max() < 1e-4 # forward with different prompt - inputs = self.get_dummy_inputs(torch_device) + inputs = self.get_dummy_inputs() inputs["prompt_2"] = "different prompt" - output = sd_pipe(**inputs) - image_slice_3 = output.images[0, -3:, -3:, -1] + image_slice_3 = pipe(**inputs).images[0, -1, -3:, -3:] # ensure the results are not equal - assert np.abs(image_slice_1.flatten() - image_slice_3.flatten()).max() > 1e-4 + assert (image_slice_1 - image_slice_3).abs().max() > 1e-4 # manually set a negative_prompt - inputs = self.get_dummy_inputs(torch_device) + inputs = self.get_dummy_inputs() inputs["negative_prompt"] = "negative prompt" - output = sd_pipe(**inputs) - image_slice_1 = output.images[0, -3:, -3:, -1] + image_slice_1 = pipe(**inputs).images[0, -1, -3:, -3:] # forward with same negative_prompt duplicated - inputs = self.get_dummy_inputs(torch_device) + inputs = self.get_dummy_inputs() inputs["negative_prompt"] = "negative prompt" inputs["negative_prompt_2"] = inputs["negative_prompt"] - output = sd_pipe(**inputs) - image_slice_2 = output.images[0, -3:, -3:, -1] + image_slice_2 = pipe(**inputs).images[0, -1, -3:, -3:] # ensure the results are equal - assert np.abs(image_slice_1.flatten() - image_slice_2.flatten()).max() < 1e-4 + assert (image_slice_1 - image_slice_2).abs().max() < 1e-4 # forward with different negative_prompt - inputs = self.get_dummy_inputs(torch_device) + inputs = self.get_dummy_inputs() inputs["negative_prompt"] = "negative prompt" inputs["negative_prompt_2"] = "different negative prompt" - output = sd_pipe(**inputs) - image_slice_3 = output.images[0, -3:, -3:, -1] + image_slice_3 = pipe(**inputs).images[0, -1, -3:, -3:] # ensure the results are not equal - assert np.abs(image_slice_1.flatten() - image_slice_3.flatten()).max() > 1e-4 + assert (image_slice_1 - image_slice_3).abs().max() > 1e-4 def test_controlnet_sdxl_guess(self): - device = "cpu" - - components = self.get_dummy_components() - - sd_pipe = self.pipeline_class(**components) - sd_pipe = sd_pipe.to(device) + # Run on CPU: the expected slice is CPU-specific. + pipe = self.get_pipeline() - sd_pipe.set_progress_bar_config(disable=None) - - inputs = self.get_dummy_inputs(device) + inputs = self.get_dummy_inputs() inputs["guess_mode"] = True - output = sd_pipe(**inputs) - image_slice = output.images[0, -3:, -3:, -1] - - expected_slice = np.array( - [0.715316, 0.563373, 0.569716, 0.620860, 0.569999, 0.604369, 0.428272, 0.455195, 0.527119] - ) + image_slice = pipe(**inputs).images[0, -1, -3:, -3:] # make sure that it's equal - assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-4 + assert_tensors_close(image_slice.flatten().cpu(), self.expected_guess_slice, atol=1e-4) def test_controlnet_sdxl_lcm(self): - device = "cpu" # ensure determinism for the device-dependent torch.Generator - - components = self.get_dummy_components(time_cond_proj_dim=256) - sd_pipe = StableDiffusionXLControlNetPipeline(**components) - sd_pipe.scheduler = LCMScheduler.from_config(sd_pipe.scheduler.config) - sd_pipe = sd_pipe.to(torch_device) - sd_pipe.set_progress_bar_config(disable=None) + pipe = self.get_pipeline(**self.get_dummy_components(time_cond_proj_dim=256)) + pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config) + pipe = pipe.to(torch_device) - inputs = self.get_dummy_inputs(device) - output = sd_pipe(**inputs) - image = output.images + image = pipe(**self.get_dummy_inputs()).images + assert image.shape == (1, *self.output_shape) - image_slice = image[0, -3:, -3:, -1] + image_slice = image[0, -1, -3:, -3:] + assert_tensors_close(image_slice.flatten().cpu(), self.expected_lcm_slice, atol=1e-2) - assert image.shape == (1, 64, 64, 3) - expected_slice = np.array([0.7820, 0.6195, 0.6193, 0.7045, 0.6706, 0.5837, 0.4147, 0.5232, 0.4868]) - assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2 +class TestStableDiffusionXLControlNetPipeline( + StableDiffusionXLControlNetPipelineTesterConfig, StableDiffusionXLControlNetPipelineTests, PipelineTesterMixin +): + # fmt: off + expected_guess_slice = torch.tensor([0.715316, 0.563373, 0.569716, 0.620860, 0.569999, 0.604369, 0.428272, 0.455195, 0.527119]) + expected_lcm_slice = torch.tensor([0.7820, 0.6195, 0.6193, 0.7045, 0.6706, 0.5837, 0.4147, 0.5232, 0.4868]) + # fmt: on # Copied from test_stable_diffusion_xl.py:test_stable_diffusion_two_xl_mixture_of_denoiser_fast # with `StableDiffusionXLControlNetPipeline` instead of `StableDiffusionXLPipeline` @@ -357,7 +280,7 @@ def assert_run_mixture( expected_tss, num_train_timesteps=pipe_1.scheduler.config.num_train_timesteps, ): - inputs = self.get_dummy_inputs(torch_device) + inputs = self.get_dummy_inputs() inputs["num_inference_steps"] = num_steps class scheduler_cls(scheduler_cls_orig): @@ -444,13 +367,23 @@ def new_step(self, *args, **kwargs): assert_run_mixture(steps, split, scheduler_cls_timesteps[0], scheduler_cls_timesteps[1]) -class StableDiffusionXLMultiControlNetPipelineFastTests( - PipelineTesterMixin, PipelineKarrasSchedulerTesterMixin, unittest.TestCase +class TestStableDiffusionXLControlNetPipelineIPAdapter( + StableDiffusionXLControlNetPipelineTesterConfig, IPAdapterTesterMixin +): + """IP-Adapter tests for the SDXL ControlNet pipeline.""" + + +class TestStableDiffusionXLControlNetPipelineMemory( + StableDiffusionXLControlNetPipelineTesterConfig, MemoryTesterMixin ): + """Memory optimization tests (CPU offload, group offload, layerwise casting) for the SDXL ControlNet pipeline.""" + + +class StableDiffusionXLMultiControlNetPipelineTesterConfig(BasePipelineTesterConfig): pipeline_class = StableDiffusionXLControlNetPipeline - params = TEXT_TO_IMAGE_PARAMS - batch_params = TEXT_TO_IMAGE_BATCH_PARAMS - image_params = frozenset([]) # TO_DO: add image_params once refactored VaeImageProcessor.preprocess + required_input_params_in_call_signature = TEXT_TO_IMAGE_PARAMS + batch_input_params = TEXT_TO_IMAGE_BATCH_PARAMS + output_shape = (3, 64, 64) def get_dummy_components(self): torch.manual_seed(0) @@ -553,7 +486,7 @@ def init_weights(m): controlnet = MultiControlNetModel([controlnet1, controlnet2]) - components = { + return { "unet": unet, "controlnet": controlnet, "scheduler": scheduler, @@ -565,98 +498,75 @@ def init_weights(m): "feature_extractor": None, "image_encoder": None, } - 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): + # The conditioning images are drawn from the same generator that is handed to the pipeline, so the pipeline + # sees an already-advanced generator state — keep the order to stay comparable across runs. + generator = self.get_generator(0) controlnet_embedder_scale_factor = 2 - images = [ randn_tensor( (1, 3, 32 * controlnet_embedder_scale_factor, 32 * controlnet_embedder_scale_factor), generator=generator, - device=torch.device(device), + device=torch.device(torch_device), ), randn_tensor( (1, 3, 32 * controlnet_embedder_scale_factor, 32 * controlnet_embedder_scale_factor), generator=generator, - device=torch.device(device), + device=torch.device(torch_device), ), ] - inputs = { + return { "prompt": "A painting of a squirrel eating a burger", "generator": generator, "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", "image": images, } - return inputs +class TestStableDiffusionXLMultiControlNetPipeline( + StableDiffusionXLMultiControlNetPipelineTesterConfig, PipelineTesterMixin +): def test_control_guidance_switch(self): - components = self.get_dummy_components() - pipe = self.pipeline_class(**components) - pipe.to(torch_device) + pipe = self.get_pipeline().to(torch_device) scale = 10.0 steps = 4 - inputs = self.get_dummy_inputs(torch_device) - inputs["num_inference_steps"] = steps - inputs["controlnet_conditioning_scale"] = scale - output_1 = pipe(**inputs)[0] + def run(**extra): + inputs = self.get_dummy_inputs() + inputs["num_inference_steps"] = steps + inputs["controlnet_conditioning_scale"] = scale + return pipe(**inputs, **extra)[0] - inputs = self.get_dummy_inputs(torch_device) - inputs["num_inference_steps"] = steps - inputs["controlnet_conditioning_scale"] = scale - output_2 = pipe(**inputs, control_guidance_start=0.1, control_guidance_end=0.2)[0] - - inputs = self.get_dummy_inputs(torch_device) - inputs["num_inference_steps"] = steps - inputs["controlnet_conditioning_scale"] = scale - output_3 = pipe(**inputs, control_guidance_start=[0.1, 0.3], control_guidance_end=[0.2, 0.7])[0] - - inputs = self.get_dummy_inputs(torch_device) - inputs["num_inference_steps"] = steps - inputs["controlnet_conditioning_scale"] = scale - output_4 = pipe(**inputs, control_guidance_start=0.4, control_guidance_end=[0.5, 0.8])[0] + output_1 = run() + output_2 = run(control_guidance_start=0.1, control_guidance_end=0.2) + output_3 = run(control_guidance_start=[0.1, 0.3], control_guidance_end=[0.2, 0.7]) + output_4 = run(control_guidance_start=0.4, control_guidance_end=[0.5, 0.8]) # make sure that all outputs are different - assert np.sum(np.abs(output_1 - output_2)) > 1e-3 - assert np.sum(np.abs(output_1 - output_3)) > 1e-3 - assert np.sum(np.abs(output_1 - output_4)) > 1e-3 - - def test_attention_slicing_forward_pass(self): - return self._test_attention_slicing_forward_pass(expected_max_diff=2e-3) - - @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): - self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=2e-3) + assert (output_1 - output_2).abs().sum() > 1e-3 + assert (output_1 - output_3).abs().sum() > 1e-3 + assert (output_1 - output_4).abs().sum() > 1e-3 - def test_inference_batch_single_identical(self): - self._test_inference_batch_single_identical(expected_max_diff=2e-3) + 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) - @unittest.skip("We test this functionality elsewhere already.") + @pytest.mark.skip("We test this functionality elsewhere already.") def test_save_load_optional_components(self): pass -class StableDiffusionXLMultiControlNetOneModelPipelineFastTests( - PipelineKarrasSchedulerTesterMixin, PipelineTesterMixin, unittest.TestCase -): +class StableDiffusionXLMultiControlNetOneModelPipelineTesterConfig(BasePipelineTesterConfig): pipeline_class = StableDiffusionXLControlNetPipeline - params = TEXT_TO_IMAGE_PARAMS - batch_params = TEXT_TO_IMAGE_BATCH_PARAMS - image_params = frozenset([]) # TO_DO: add image_params once refactored VaeImageProcessor.preprocess + required_input_params_in_call_signature = TEXT_TO_IMAGE_PARAMS + batch_input_params = TEXT_TO_IMAGE_BATCH_PARAMS + output_shape = (3, 64, 64) def get_dummy_components(self): torch.manual_seed(0) @@ -741,7 +651,7 @@ def init_weights(m): controlnet = MultiControlNetModel([controlnet]) - components = { + return { "unet": unet, "controlnet": controlnet, "scheduler": scheduler, @@ -753,118 +663,88 @@ def init_weights(m): "feature_extractor": None, "image_encoder": None, } - 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): + # The conditioning image is drawn from the same generator that is handed to the pipeline, so the pipeline + # sees an already-advanced generator state — keep the order to stay comparable across runs. + generator = self.get_generator(0) controlnet_embedder_scale_factor = 2 images = [ randn_tensor( (1, 3, 32 * controlnet_embedder_scale_factor, 32 * controlnet_embedder_scale_factor), generator=generator, - device=torch.device(device), + device=torch.device(torch_device), ), ] - inputs = { + return { "prompt": "A painting of a squirrel eating a burger", "generator": generator, "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", "image": images, } - return inputs +class TestStableDiffusionXLMultiControlNetOneModelPipeline( + StableDiffusionXLMultiControlNetOneModelPipelineTesterConfig, PipelineTesterMixin +): def test_control_guidance_switch(self): - components = self.get_dummy_components() - pipe = self.pipeline_class(**components) - pipe.to(torch_device) + pipe = self.get_pipeline().to(torch_device) scale = 10.0 steps = 4 - inputs = self.get_dummy_inputs(torch_device) - inputs["num_inference_steps"] = steps - inputs["controlnet_conditioning_scale"] = scale - output_1 = pipe(**inputs)[0] + def run(**extra): + inputs = self.get_dummy_inputs() + inputs["num_inference_steps"] = steps + inputs["controlnet_conditioning_scale"] = scale + return pipe(**inputs, **extra)[0] - inputs = self.get_dummy_inputs(torch_device) - inputs["num_inference_steps"] = steps - inputs["controlnet_conditioning_scale"] = scale - output_2 = pipe(**inputs, control_guidance_start=0.1, control_guidance_end=0.2)[0] - - inputs = self.get_dummy_inputs(torch_device) - inputs["num_inference_steps"] = steps - inputs["controlnet_conditioning_scale"] = scale - output_3 = pipe( - **inputs, - control_guidance_start=[0.1], - control_guidance_end=[0.2], - )[0] - - inputs = self.get_dummy_inputs(torch_device) - inputs["num_inference_steps"] = steps - inputs["controlnet_conditioning_scale"] = scale - output_4 = pipe(**inputs, control_guidance_start=0.4, control_guidance_end=[0.5])[0] + output_1 = run() + output_2 = run(control_guidance_start=0.1, control_guidance_end=0.2) + output_3 = run(control_guidance_start=[0.1], control_guidance_end=[0.2]) + output_4 = run(control_guidance_start=0.4, control_guidance_end=[0.5]) # make sure that all outputs are different - assert np.sum(np.abs(output_1 - output_2)) > 1e-3 - assert np.sum(np.abs(output_1 - output_3)) > 1e-3 - assert np.sum(np.abs(output_1 - output_4)) > 1e-3 + assert (output_1 - output_2).abs().sum() > 1e-3 + assert (output_1 - output_3).abs().sum() > 1e-3 + assert (output_1 - output_4).abs().sum() > 1e-3 - def test_attention_slicing_forward_pass(self): - return self._test_attention_slicing_forward_pass(expected_max_diff=2e-3) + 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) - @unittest.skip("We test this functionality elsewhere already.") + @pytest.mark.skip("We test this functionality elsewhere already.") def test_save_load_optional_components(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): - self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=2e-3) - - def test_inference_batch_single_identical(self): - self._test_inference_batch_single_identical(expected_max_diff=2e-3) - def test_negative_conditions(self): - components = self.get_dummy_components() - pipe = self.pipeline_class(**components) - pipe.to(torch_device) + pipe = self.get_pipeline().to(torch_device) - inputs = self.get_dummy_inputs(torch_device) - image = pipe(**inputs).images - image_slice_without_neg_cond = image[0, -3:, -3:, -1] + inputs = self.get_dummy_inputs() + image_slice_without_neg_cond = pipe(**inputs).images[0, -1, -3:, -3:] - image = pipe( + image_slice_with_neg_cond = pipe( **inputs, negative_original_size=(512, 512), negative_crops_coords_top_left=(0, 0), negative_target_size=(1024, 1024), - ).images - image_slice_with_neg_cond = image[0, -3:, -3:, -1] + ).images[0, -1, -3:, -3:] - self.assertTrue(np.abs(image_slice_without_neg_cond - image_slice_with_neg_cond).max() > 1e-2) + assert (image_slice_without_neg_cond - image_slice_with_neg_cond).abs().max() > 1e-2 @slow @require_torch_accelerator -class ControlNetSDXLPipelineSlowTests(unittest.TestCase): - def setUp(self): - super().setUp() +class TestControlNetSDXLPipelineSlow: + @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) @@ -915,81 +795,8 @@ def test_depth(self): assert np.allclose(original_image, expected_image, atol=1e-04) -class StableDiffusionSSD1BControlNetPipelineFastTests(StableDiffusionXLControlNetPipelineFastTests): - def test_controlnet_sdxl_guess(self): - device = "cpu" - - components = self.get_dummy_components() - - sd_pipe = self.pipeline_class(**components) - sd_pipe = sd_pipe.to(device) - - sd_pipe.set_progress_bar_config(disable=None) - - inputs = self.get_dummy_inputs(device) - inputs["guess_mode"] = True - - output = sd_pipe(**inputs) - image_slice = output.images[0, -3:, -3:, -1] - - expected_slice = np.array( - [0.669912, 0.557802, 0.523260, 0.596366, 0.552897, 0.576922, 0.433411, 0.450273, 0.491615] - ) - - # make sure that it's equal - assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-4 - - def test_ip_adapter(self): - expected_pipe_slice = None - if torch_device == "cpu": - expected_pipe_slice = np.array([0.6699, 0.5578, 0.5233, 0.5964, 0.5529, 0.5769, 0.4334, 0.4503, 0.4916]) - - return super().test_ip_adapter(from_ssd1b=True, expected_pipe_slice=expected_pipe_slice) - - def test_controlnet_sdxl_lcm(self): - device = "cpu" # ensure determinism for the device-dependent torch.Generator - - components = self.get_dummy_components(time_cond_proj_dim=256) - sd_pipe = StableDiffusionXLControlNetPipeline(**components) - sd_pipe.scheduler = LCMScheduler.from_config(sd_pipe.scheduler.config) - sd_pipe = sd_pipe.to(torch_device) - sd_pipe.set_progress_bar_config(disable=None) - - inputs = self.get_dummy_inputs(device) - output = sd_pipe(**inputs) - image = output.images - - image_slice = image[0, -3:, -3:, -1] - - assert image.shape == (1, 64, 64, 3) - expected_slice = np.array([0.6787, 0.5117, 0.5558, 0.6963, 0.6571, 0.5928, 0.4121, 0.5468, 0.5057]) - - assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2 - - def test_conditioning_channels(self): - unet = UNet2DConditionModel( - block_out_channels=(32, 64), - layers_per_block=2, - sample_size=32, - in_channels=4, - out_channels=4, - down_block_types=("DownBlock2D", "CrossAttnDownBlock2D"), - up_block_types=("CrossAttnUpBlock2D", "UpBlock2D"), - mid_block_type="UNetMidBlock2D", - # SD2-specific config below - attention_head_dim=(2, 4), - use_linear_projection=True, - addition_embed_type="text_time", - addition_time_embed_dim=8, - transformer_layers_per_block=(1, 2), - projection_class_embeddings_input_dim=80, # 6 * 8 + 32 - cross_attention_dim=64, - time_cond_proj_dim=None, - ) - - controlnet = ControlNetModel.from_unet(unet, conditioning_channels=4) - assert type(controlnet.mid_block) is UNetMidBlock2D - assert controlnet.conditioning_channels == 4 +class StableDiffusionSSD1BControlNetPipelineTesterConfig(StableDiffusionXLControlNetPipelineTesterConfig): + """Same contract as the SDXL ControlNet config, but built on an SSD-1B-shaped UNet / ControlNet.""" def get_dummy_components(self, time_cond_proj_dim=None): torch.manual_seed(0) @@ -1067,7 +874,7 @@ def get_dummy_components(self, time_cond_proj_dim=None): text_encoder_2 = CLIPTextModelWithProjection(text_encoder_config) tokenizer_2 = CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip") - components = { + return { "unet": unet, "controlnet": controlnet, "scheduler": scheduler, @@ -1079,4 +886,49 @@ def get_dummy_components(self, time_cond_proj_dim=None): "feature_extractor": None, "image_encoder": None, } - return components + + +class TestStableDiffusionSSD1BControlNetPipeline( + StableDiffusionSSD1BControlNetPipelineTesterConfig, StableDiffusionXLControlNetPipelineTests, PipelineTesterMixin +): + # fmt: off + expected_guess_slice = torch.tensor([0.669912, 0.557802, 0.523260, 0.596366, 0.552897, 0.576922, 0.433411, 0.450273, 0.491615]) + expected_lcm_slice = torch.tensor([0.6787, 0.5117, 0.5558, 0.6963, 0.6571, 0.5928, 0.4121, 0.5468, 0.5057]) + # fmt: on + + def test_conditioning_channels(self): + unet = UNet2DConditionModel( + block_out_channels=(32, 64), + layers_per_block=2, + sample_size=32, + in_channels=4, + out_channels=4, + down_block_types=("DownBlock2D", "CrossAttnDownBlock2D"), + up_block_types=("CrossAttnUpBlock2D", "UpBlock2D"), + mid_block_type="UNetMidBlock2D", + # SD2-specific config below + attention_head_dim=(2, 4), + use_linear_projection=True, + addition_embed_type="text_time", + addition_time_embed_dim=8, + transformer_layers_per_block=(1, 2), + projection_class_embeddings_input_dim=80, # 6 * 8 + 32 + cross_attention_dim=64, + time_cond_proj_dim=None, + ) + + controlnet = ControlNetModel.from_unet(unet, conditioning_channels=4) + assert type(controlnet.mid_block) is UNetMidBlock2D + assert controlnet.conditioning_channels == 4 + + +class TestStableDiffusionSSD1BControlNetPipelineIPAdapter( + StableDiffusionSSD1BControlNetPipelineTesterConfig, IPAdapterTesterMixin +): + """IP-Adapter tests for the SSD-1B ControlNet pipeline.""" + + +class TestStableDiffusionSSD1BControlNetPipelineMemory( + StableDiffusionSSD1BControlNetPipelineTesterConfig, MemoryTesterMixin +): + """Memory optimization tests (CPU offload, group offload, layerwise casting) for the SSD-1B ControlNet pipeline.""" diff --git a/tests/pipelines/controlnet/test_controlnet_sdxl_img2img.py b/tests/pipelines/controlnet/test_controlnet_sdxl_img2img.py index d242a1bacb92..9b9caa248ba1 100644 --- a/tests/pipelines/controlnet/test_controlnet_sdxl_img2img.py +++ b/tests/pipelines/controlnet/test_controlnet_sdxl_img2img.py @@ -14,9 +14,8 @@ # limitations under the License. import random -import unittest -import numpy as np +import pytest import torch from transformers import CLIPTextConfig, CLIPTextModel, CLIPTextModelWithProjection, CLIPTokenizer @@ -27,44 +26,27 @@ StableDiffusionXLControlNetImg2ImgPipeline, UNet2DConditionModel, ) -from diffusers.utils.import_utils import is_xformers_available -from ...testing_utils import ( - enable_full_determinism, - floats_tensor, - require_torch_accelerator, - torch_device, -) +from ...testing_utils import assert_tensors_close, enable_full_determinism, floats_tensor, torch_device from ..pipeline_params import ( - IMAGE_TO_IMAGE_IMAGE_PARAMS, TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS, TEXT_GUIDED_IMAGE_VARIATION_PARAMS, TEXT_TO_IMAGE_CALLBACK_CFG_PARAMS, ) -from ..test_pipelines_common import ( - IPAdapterTesterMixin, - PipelineKarrasSchedulerTesterMixin, - PipelineLatentTesterMixin, - PipelineTesterMixin, -) +from ..stable_diffusion.ip_adapter_tester import IPAdapterTesterMixin +from ..testing_utils import BasePipelineTesterConfig, MemoryTesterMixin, PipelineTesterMixin enable_full_determinism() -class ControlNetPipelineSDXLImg2ImgFastTests( - IPAdapterTesterMixin, - PipelineLatentTesterMixin, - PipelineKarrasSchedulerTesterMixin, - PipelineTesterMixin, - unittest.TestCase, -): +class StableDiffusionXLControlNetImg2ImgPipelineTesterConfig(BasePipelineTesterConfig): pipeline_class = StableDiffusionXLControlNetImg2ImgPipeline - params = TEXT_GUIDED_IMAGE_VARIATION_PARAMS - required_optional_params = PipelineTesterMixin.required_optional_params - {"latents"} - batch_params = TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS - image_params = IMAGE_TO_IMAGE_IMAGE_PARAMS - image_latents_params = IMAGE_TO_IMAGE_IMAGE_PARAMS + required_input_params_in_call_signature = TEXT_GUIDED_IMAGE_VARIATION_PARAMS + batch_input_params = TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS + output_shape = (3, 64, 64) + # The img2img pipeline derives its starting latents from `image`, so it takes no `latents` argument. + optional_input_params = BasePipelineTesterConfig.optional_input_params - {"latents"} callback_cfg_params = TEXT_TO_IMAGE_CALLBACK_CFG_PARAMS.union( {"add_text_embeds", "add_time_ids", "add_neg_time_ids"} ) @@ -142,7 +124,7 @@ def get_dummy_components(self, skip_first_text_encoder=False): text_encoder_2 = CLIPTextModelWithProjection(text_encoder_config) tokenizer_2 = CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip") - components = { + return { "unet": unet, "controlnet": controlnet, "scheduler": scheduler, @@ -154,176 +136,119 @@ def get_dummy_components(self, skip_first_text_encoder=False): "image_encoder": None, "feature_extractor": None, } - return components - def get_dummy_inputs(self, device, seed=0): + def get_dummy_inputs(self): controlnet_embedder_scale_factor = 2 image = floats_tensor( (1, 3, 32 * controlnet_embedder_scale_factor, 32 * controlnet_embedder_scale_factor), - rng=random.Random(seed), - ).to(device) - - if str(device).startswith("mps"): - generator = torch.manual_seed(seed) - else: - generator = torch.Generator(device=device).manual_seed(seed) + rng=random.Random(0), + ).to(torch_device) - inputs = { + 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", "image": image, "control_image": image, } - return inputs - def test_ip_adapter(self): - expected_pipe_slice = None - if torch_device == "cpu": - expected_pipe_slice = np.array([0.6710, 0.5497, 0.5469, 0.5758, 0.5990, 0.5996, 0.5583, 0.5506, 0.5368]) - # TODO: update after slices.p - return super().test_ip_adapter(expected_pipe_slice=expected_pipe_slice) +class TestStableDiffusionXLControlNetImg2ImgPipeline( + StableDiffusionXLControlNetImg2ImgPipelineTesterConfig, PipelineTesterMixin +): + # Guess mode is expected to land on the same slice at this tolerance, so both tests below share it. + # fmt: off + expected_slice = torch.tensor([0.55813384, 0.4668495, 0.46676695, 0.6121852, 0.55514586, 0.49157068, 0.5960574, 0.56897247, 0.43931544]) + # fmt: on def test_stable_diffusion_xl_controlnet_img2img(self): - device = "cpu" # ensure determinism for the device-dependent torch.Generator - components = self.get_dummy_components() - sd_pipe = self.pipeline_class(**components) - sd_pipe = sd_pipe.to(device) - sd_pipe.set_progress_bar_config(disable=None) - - inputs = self.get_dummy_inputs(device) - image = sd_pipe(**inputs).images - image_slice = image[0, -3:, -3:, -1] - assert image.shape == (1, 64, 64, 3) - - expected_slice = np.array( - [0.55813384, 0.4668495, 0.46676695, 0.6121852, 0.55514586, 0.49157068, 0.5960574, 0.56897247, 0.43931544] - ) - - assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2 + # Run on CPU: the expected slice is CPU-specific. + pipe = self.get_pipeline() - def test_stable_diffusion_xl_controlnet_img2img_guess(self): - device = "cpu" - - components = self.get_dummy_components() + image = pipe(**self.get_dummy_inputs()).images + assert image.shape == (1, *self.output_shape) - sd_pipe = self.pipeline_class(**components) - sd_pipe = sd_pipe.to(device) + image_slice = image[0, -1, -3:, -3:] + assert_tensors_close(image_slice.flatten().cpu(), self.expected_slice, atol=1e-2) - sd_pipe.set_progress_bar_config(disable=None) + def test_stable_diffusion_xl_controlnet_img2img_guess(self): + # Run on CPU: the expected slice is CPU-specific. + pipe = self.get_pipeline() - inputs = self.get_dummy_inputs(device) + inputs = self.get_dummy_inputs() inputs["guess_mode"] = True - output = sd_pipe(**inputs) - image_slice = output.images[0, -3:, -3:, -1] - assert output.images.shape == (1, 64, 64, 3) - - expected_slice = np.array( - [0.55813384, 0.4668495, 0.46676695, 0.6121852, 0.55514586, 0.49157068, 0.5960574, 0.56897247, 0.43931544] - ) + image = pipe(**inputs).images + assert image.shape == (1, *self.output_shape) + image_slice = image[0, -1, -3:, -3:] # make sure that it's equal - assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2 - - def test_attention_slicing_forward_pass(self): - return self._test_attention_slicing_forward_pass(expected_max_diff=2e-3) - - @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): - self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=2e-3) + assert_tensors_close(image_slice.flatten().cpu(), self.expected_slice, atol=1e-2) - def test_inference_batch_single_identical(self): - self._test_inference_batch_single_identical(expected_max_diff=2e-3) + 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) - # TODO(Patrick, Sayak) - skip for now as this requires more refiner tests + @pytest.mark.skip("TODO(Patrick, Sayak) - skip for now as this requires more refiner tests") def test_save_load_optional_components(self): pass - @require_torch_accelerator - def test_stable_diffusion_xl_offloads(self): - pipes = [] - components = self.get_dummy_components() - sd_pipe = self.pipeline_class(**components).to(torch_device) - pipes.append(sd_pipe) - - components = self.get_dummy_components() - sd_pipe = self.pipeline_class(**components) - sd_pipe.enable_model_cpu_offload(device=torch_device) - pipes.append(sd_pipe) - - components = self.get_dummy_components() - sd_pipe = self.pipeline_class(**components) - sd_pipe.enable_sequential_cpu_offload(device=torch_device) - pipes.append(sd_pipe) - - image_slices = [] - for pipe in pipes: - pipe.unet.set_default_attn_processor() - - inputs = self.get_dummy_inputs(torch_device) - image = pipe(**inputs).images - - image_slices.append(image[0, -3:, -3:, -1].flatten()) - - assert np.abs(image_slices[0] - image_slices[1]).max() < 1e-3 - assert np.abs(image_slices[0] - image_slices[2]).max() < 1e-3 - def test_stable_diffusion_xl_multi_prompts(self): - components = self.get_dummy_components() - sd_pipe = self.pipeline_class(**components).to(torch_device) + pipe = self.get_pipeline().to(torch_device) # forward with single prompt - inputs = self.get_dummy_inputs(torch_device) - output = sd_pipe(**inputs) - image_slice_1 = output.images[0, -3:, -3:, -1] + image_slice_1 = pipe(**self.get_dummy_inputs()).images[0, -1, -3:, -3:] # forward with same prompt duplicated - inputs = self.get_dummy_inputs(torch_device) + inputs = self.get_dummy_inputs() inputs["prompt_2"] = inputs["prompt"] - output = sd_pipe(**inputs) - image_slice_2 = output.images[0, -3:, -3:, -1] + image_slice_2 = pipe(**inputs).images[0, -1, -3:, -3:] # ensure the results are equal - assert np.abs(image_slice_1.flatten() - image_slice_2.flatten()).max() < 1e-4 + assert (image_slice_1 - image_slice_2).abs().max() < 1e-4 # forward with different prompt - inputs = self.get_dummy_inputs(torch_device) + inputs = self.get_dummy_inputs() inputs["prompt_2"] = "different prompt" - output = sd_pipe(**inputs) - image_slice_3 = output.images[0, -3:, -3:, -1] + image_slice_3 = pipe(**inputs).images[0, -1, -3:, -3:] # ensure the results are not equal - assert np.abs(image_slice_1.flatten() - image_slice_3.flatten()).max() > 1e-4 + assert (image_slice_1 - image_slice_3).abs().max() > 1e-4 # manually set a negative_prompt - inputs = self.get_dummy_inputs(torch_device) + inputs = self.get_dummy_inputs() inputs["negative_prompt"] = "negative prompt" - output = sd_pipe(**inputs) - image_slice_1 = output.images[0, -3:, -3:, -1] + image_slice_1 = pipe(**inputs).images[0, -1, -3:, -3:] # forward with same negative_prompt duplicated - inputs = self.get_dummy_inputs(torch_device) + inputs = self.get_dummy_inputs() inputs["negative_prompt"] = "negative prompt" inputs["negative_prompt_2"] = inputs["negative_prompt"] - output = sd_pipe(**inputs) - image_slice_2 = output.images[0, -3:, -3:, -1] + image_slice_2 = pipe(**inputs).images[0, -1, -3:, -3:] # ensure the results are equal - assert np.abs(image_slice_1.flatten() - image_slice_2.flatten()).max() < 1e-4 + assert (image_slice_1 - image_slice_2).abs().max() < 1e-4 # forward with different negative_prompt - inputs = self.get_dummy_inputs(torch_device) + inputs = self.get_dummy_inputs() inputs["negative_prompt"] = "negative prompt" inputs["negative_prompt_2"] = "different negative prompt" - output = sd_pipe(**inputs) - image_slice_3 = output.images[0, -3:, -3:, -1] + image_slice_3 = pipe(**inputs).images[0, -1, -3:, -3:] # ensure the results are not equal - assert np.abs(image_slice_1.flatten() - image_slice_3.flatten()).max() > 1e-4 + assert (image_slice_1 - image_slice_3).abs().max() > 1e-4 + + +class TestStableDiffusionXLControlNetImg2ImgPipelineIPAdapter( + StableDiffusionXLControlNetImg2ImgPipelineTesterConfig, IPAdapterTesterMixin +): + """IP-Adapter tests for the SDXL ControlNet img2img pipeline.""" + + +class TestStableDiffusionXLControlNetImg2ImgPipelineMemory( + StableDiffusionXLControlNetImg2ImgPipelineTesterConfig, MemoryTesterMixin +): + """Memory optimization tests (CPU offload, group offload, layerwise casting) for the SDXL ControlNet img2img pipeline.""" diff --git a/tests/pipelines/controlnet_flux/test_controlnet_flux.py b/tests/pipelines/controlnet_flux/test_controlnet_flux.py index 6e223973eee0..90c5e617f61b 100644 --- a/tests/pipelines/controlnet_flux/test_controlnet_flux.py +++ b/tests/pipelines/controlnet_flux/test_controlnet_flux.py @@ -14,9 +14,9 @@ # limitations under the License. import gc -import unittest import numpy as np +import pytest import torch from huggingface_hub import hf_hub_download from transformers import AutoConfig, CLIPTextConfig, CLIPTextModel, CLIPTokenizer, T5EncoderModel, T5TokenizerFast @@ -32,6 +32,7 @@ from diffusers.utils.torch_utils import randn_tensor from ...testing_utils import ( + assert_tensors_close, backend_empty_cache, enable_full_determinism, nightly, @@ -39,19 +40,20 @@ require_big_accelerator, torch_device, ) -from ..test_pipelines_common import FluxIPAdapterTesterMixin, PipelineTesterMixin +from ..flux.testing_utils import FluxIPAdapterTesterMixin +from ..testing_utils import BasePipelineTesterConfig, MemoryTesterMixin, PipelineTesterMixin enable_full_determinism() -class FluxControlNetPipelineFastTests(unittest.TestCase, PipelineTesterMixin, FluxIPAdapterTesterMixin): +class FluxControlNetPipelineTesterConfig(BasePipelineTesterConfig): pipeline_class = FluxControlNetPipeline - - params = frozenset(["prompt", "height", "width", "guidance_scale", "prompt_embeds", "pooled_prompt_embeds"]) - batch_params = frozenset(["prompt"]) - test_layerwise_casting = True - test_group_offloading = True + required_input_params_in_call_signature = frozenset( + ["prompt", "height", "width", "guidance_scale", "prompt_embeds", "pooled_prompt_embeds"] + ) + batch_input_params = frozenset(["prompt"]) + output_shape = (3, 32, 32) def get_dummy_components(self): torch.manual_seed(0) @@ -98,7 +100,9 @@ def get_dummy_components(self): torch.manual_seed(0) config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5") - text_encoder_2 = T5EncoderModel(config) + # `eval()` because a directly constructed model stays in training mode, which leaves T5's + # dropout active and makes the pipeline outputs non-deterministic across calls. + text_encoder_2 = T5EncoderModel(config).eval() tokenizer = CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip") tokenizer_2 = T5TokenizerFast.from_pretrained("hf-internal-testing/tiny-random-t5") @@ -133,60 +137,47 @@ def get_dummy_components(self): "feature_extractor": None, } - 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) - + def get_dummy_inputs(self): + # The control image is drawn from the same generator that is handed to the pipeline, so the pipeline sees an + # already-advanced generator state — keep the order to stay comparable with the expected slices below. + generator = self.get_generator(0) control_image = randn_tensor( (1, 3, 32, 32), generator=generator, - device=torch.device(device), + device=torch.device(torch_device), dtype=torch.float32, ) - controlnet_conditioning_scale = 0.5 - - inputs = { + return { "prompt": "A painting of a squirrel eating a burger", "generator": generator, "num_inference_steps": 2, "guidance_scale": 3.5, - "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", "control_image": control_image, - "controlnet_conditioning_scale": controlnet_conditioning_scale, + "controlnet_conditioning_scale": 0.5, } - return inputs +class TestFluxControlNetPipeline(FluxControlNetPipelineTesterConfig, PipelineTesterMixin): def test_controlnet_flux(self): - components = self.get_dummy_components() - flux_pipe = FluxControlNetPipeline(**components) - flux_pipe = flux_pipe.to(torch_device, dtype=torch.float32) - flux_pipe.set_progress_bar_config(disable=None) - - inputs = self.get_dummy_inputs(torch_device) - output = flux_pipe(**inputs) - image = output.images - - image_slice = image[0, -3:, -3:, -1] - - assert image.shape == (1, 32, 32, 3) + pipe = self.get_pipeline().to(torch_device, dtype=torch.float32) - expected_slice = np.array([0.6751, 0.6115, 0.5290, 0.6200, 0.5736, 0.6409, 0.5588, 0.6046, 0.5594]) + image = pipe(**self.get_dummy_inputs()).images + assert image.shape == (1, *self.output_shape) - assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2, ( - f"Expected: {expected_slice}, got: {image_slice.flatten()}" - ) + image_slice = image[0, -1, -3:, -3:] + # fmt: off + expected_slice = torch.tensor([0.6751, 0.6115, 0.5290, 0.6200, 0.5736, 0.6409, 0.5588, 0.6046, 0.5594]) + # fmt: on - @unittest.skip("xFormersAttnProcessor does not work with SD3 Joint Attention") - def test_xformers_attention_forwardGenerator_pass(self): - pass + assert_tensors_close(image_slice.flatten().cpu(), expected_slice, atol=1e-2) def test_flux_image_output_shape(self): - pipe = self.pipeline_class(**self.get_dummy_components()).to(torch_device) - inputs = self.get_dummy_inputs(torch_device) + pipe = self.get_pipeline().to(torch_device) + inputs = self.get_dummy_inputs() height_width_pairs = [(32, 32), (72, 56)] for height, width in height_width_pairs: @@ -203,22 +194,30 @@ def test_flux_image_output_shape(self): } ) image = pipe(**inputs).images[0] - output_height, output_width, _ = image.shape - assert (output_height, output_width) == (expected_height, expected_width) + _, output_height, output_width = image.shape + assert (output_height, output_width) == (expected_height, expected_width), ( + f"Output shape {image.shape} does not match expected shape {(expected_height, expected_width)}" + ) + + +class TestFluxControlNetPipelineIPAdapter(FluxControlNetPipelineTesterConfig, FluxIPAdapterTesterMixin): + """IP-Adapter tests for the Flux ControlNet pipeline.""" + + +class TestFluxControlNetPipelineMemory(FluxControlNetPipelineTesterConfig, MemoryTesterMixin): + """Memory optimization tests (CPU offload, group offload, layerwise casting) for the Flux ControlNet pipeline.""" @nightly @require_big_accelerator -class FluxControlNetPipelineSlowTests(unittest.TestCase): +class TestFluxControlNetPipelineSlow: pipeline_class = FluxControlNetPipeline - 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) diff --git a/tests/pipelines/controlnet_flux/test_controlnet_flux_img2img.py b/tests/pipelines/controlnet_flux/test_controlnet_flux_img2img.py index a4749188dfd8..165e7bda2dc9 100644 --- a/tests/pipelines/controlnet_flux/test_controlnet_flux_img2img.py +++ b/tests/pipelines/controlnet_flux/test_controlnet_flux_img2img.py @@ -1,6 +1,3 @@ -import unittest - -import numpy as np import torch from transformers import AutoConfig, AutoTokenizer, CLIPTextConfig, CLIPTextModel, CLIPTokenizer, T5EncoderModel @@ -13,13 +10,18 @@ ) from diffusers.utils.torch_utils import randn_tensor -from ...testing_utils import torch_device -from ..test_pipelines_common import PipelineTesterMixin, check_qkv_fused_layers_exist +from ...testing_utils import assert_tensors_close, torch_device +from ..testing_utils import ( + BasePipelineTesterConfig, + MemoryTesterMixin, + PipelineTesterMixin, + check_qkv_fused_layers_exist, +) -class FluxControlNetImg2ImgPipelineFastTests(unittest.TestCase, PipelineTesterMixin): +class FluxControlNetImg2ImgPipelineTesterConfig(BasePipelineTesterConfig): pipeline_class = FluxControlNetImg2ImgPipeline - params = frozenset( + required_input_params_in_call_signature = frozenset( [ "prompt", "image", @@ -33,9 +35,8 @@ class FluxControlNetImg2ImgPipelineFastTests(unittest.TestCase, PipelineTesterMi "pooled_prompt_embeds", ] ) - batch_params = frozenset(["prompt", "image", "control_image"]) - - test_xformers_attention = False + batch_input_params = frozenset(["prompt", "image", "control_image"]) + output_shape = (3, 32, 32) def get_dummy_components(self): torch.manual_seed(0) @@ -69,7 +70,9 @@ def get_dummy_components(self): torch.manual_seed(0) config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5") - text_encoder_2 = T5EncoderModel(config) + # `eval()` because a directly constructed model stays in training mode, which leaves T5's + # dropout active and makes the pipeline outputs non-deterministic across calls. + text_encoder_2 = T5EncoderModel(config).eval() tokenizer = CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip") tokenizer_2 = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5") @@ -114,20 +117,19 @@ def get_dummy_components(self): "controlnet": controlnet, } - 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) - - image = torch.randn(1, 3, 32, 32).to(device) - control_image = torch.randn(1, 3, 32, 32).to(device) + def get_dummy_inputs(self): + # Seeded so that repeated `get_dummy_inputs()` calls hand the pipeline the same images — the shared tests + # compare two runs against each other. + image = randn_tensor((1, 3, 32, 32), generator=self.get_generator(0), device=torch.device(torch_device)) + control_image = randn_tensor( + (1, 3, 32, 32), generator=self.get_generator(1), device=torch.device(torch_device) + ) - inputs = { + return { "prompt": "A painting of a squirrel eating a burger", "image": image, "control_image": control_image, - "generator": generator, + "generator": self.get_generator(0), "num_inference_steps": 2, "guidance_scale": 5.0, "controlnet_conditioning_scale": 1.0, @@ -135,63 +137,74 @@ def get_dummy_inputs(self, device, seed=0): "height": 32, "width": 32, "max_sequence_length": 48, - "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 TestFluxControlNetImg2ImgPipeline(FluxControlNetImg2ImgPipelineTesterConfig, PipelineTesterMixin): def test_flux_controlnet_different_prompts(self): - pipe = self.pipeline_class(**self.get_dummy_components()).to(torch_device) + pipe = self.get_pipeline().to(torch_device) - inputs = self.get_dummy_inputs(torch_device) + inputs = self.get_dummy_inputs() output_same_prompt = pipe(**inputs).images[0] - inputs = self.get_dummy_inputs(torch_device) + inputs = self.get_dummy_inputs() inputs["prompt_2"] = "a different prompt" output_different_prompts = pipe(**inputs).images[0] - max_diff = np.abs(output_same_prompt - output_different_prompts).max() + max_diff = (output_same_prompt - output_different_prompts).abs().max() - assert max_diff > 1e-6 + assert max_diff > 1e-6, "Outputs should be different for different prompts." def test_fused_qkv_projections(self): - device = "cpu" # ensure determinism for the device-dependent torch.Generator - components = self.get_dummy_components() - pipe = self.pipeline_class(**components) - pipe = pipe.to(device) - pipe.set_progress_bar_config(disable=None) + # Run on CPU to keep the seeded generator deterministic across the three forward passes. + pipe = self.get_pipeline() - inputs = self.get_dummy_inputs(device) + inputs = self.get_dummy_inputs() image = pipe(**inputs).images - original_image_slice = image[0, -3:, -3:, -1] + original_image_slice = image[0, -1, -3:, -3:] pipe.transformer.fuse_qkv_projections() - self.assertTrue( - check_qkv_fused_layers_exist(pipe.transformer, ["to_qkv"]), - ("Something wrong with the fused attention layers. Expected all the attention projections to be fused."), + assert check_qkv_fused_layers_exist(pipe.transformer, ["to_qkv"]), ( + "Something wrong with the fused attention layers. Expected all the attention projections to be fused." ) - inputs = self.get_dummy_inputs(device) + inputs = self.get_dummy_inputs() image = pipe(**inputs).images - image_slice_fused = image[0, -3:, -3:, -1] + image_slice_fused = image[0, -1, -3:, -3:] pipe.transformer.unfuse_qkv_projections() - inputs = self.get_dummy_inputs(device) + inputs = self.get_dummy_inputs() image = pipe(**inputs).images - image_slice_disabled = image[0, -3:, -3:, -1] - - assert np.allclose(original_image_slice, image_slice_fused, atol=1e-3, rtol=1e-3), ( - "Fusion of QKV projections shouldn't affect the outputs." + image_slice_disabled = image[0, -1, -3:, -3:] + + assert_tensors_close( + original_image_slice, + image_slice_fused, + atol=1e-3, + rtol=1e-3, + msg="Fusion of QKV projections shouldn't affect the outputs.", ) - assert np.allclose(image_slice_fused, image_slice_disabled, atol=1e-3, rtol=1e-3), ( - "Outputs, with QKV projection fusion enabled, shouldn't change when fused QKV projections are disabled." + assert_tensors_close( + image_slice_fused, + image_slice_disabled, + atol=1e-3, + rtol=1e-3, + msg="Outputs, with QKV projection fusion enabled, shouldn't change when fused QKV projections are disabled.", ) - assert np.allclose(original_image_slice, image_slice_disabled, atol=1e-2, rtol=1e-2), ( - "Original outputs should match when fused QKV projections are disabled." + assert_tensors_close( + original_image_slice, + image_slice_disabled, + atol=1e-2, + rtol=1e-2, + msg="Original outputs should match when fused QKV projections are disabled.", ) def test_flux_image_output_shape(self): - pipe = self.pipeline_class(**self.get_dummy_components()).to(torch_device) - inputs = self.get_dummy_inputs(torch_device) + pipe = self.get_pipeline().to(torch_device) + inputs = self.get_dummy_inputs() height_width_pairs = [(32, 32), (72, 56)] for height, width in height_width_pairs: @@ -214,5 +227,11 @@ def test_flux_image_output_shape(self): } ) image = pipe(**inputs).images[0] - output_height, output_width, _ = image.shape - assert (output_height, output_width) == (expected_height, expected_width) + _, output_height, output_width = image.shape + assert (output_height, output_width) == (expected_height, expected_width), ( + f"Output shape {image.shape} does not match expected shape {(expected_height, expected_width)}" + ) + + +class TestFluxControlNetImg2ImgPipelineMemory(FluxControlNetImg2ImgPipelineTesterConfig, MemoryTesterMixin): + """Memory optimization tests (CPU offload, group offload, layerwise casting) for the Flux ControlNet img2img pipeline.""" diff --git a/tests/pipelines/controlnet_flux/test_controlnet_flux_inpaint.py b/tests/pipelines/controlnet_flux/test_controlnet_flux_inpaint.py index 6eb560d90848..eeacc636e89d 100644 --- a/tests/pipelines/controlnet_flux/test_controlnet_flux_inpaint.py +++ b/tests/pipelines/controlnet_flux/test_controlnet_flux_inpaint.py @@ -1,7 +1,5 @@ import random -import unittest -import numpy as np import torch from transformers import AutoConfig, AutoTokenizer, CLIPTextConfig, CLIPTextModel, CLIPTokenizer, T5EncoderModel @@ -15,15 +13,15 @@ from diffusers.utils.torch_utils import randn_tensor from ...testing_utils import enable_full_determinism, floats_tensor, torch_device -from ..test_pipelines_common import PipelineTesterMixin +from ..testing_utils import BasePipelineTesterConfig, MemoryTesterMixin, PipelineTesterMixin enable_full_determinism() -class FluxControlNetInpaintPipelineTests(unittest.TestCase, PipelineTesterMixin): +class FluxControlNetInpaintPipelineTesterConfig(BasePipelineTesterConfig): pipeline_class = FluxControlNetInpaintPipeline - params = frozenset( + required_input_params_in_call_signature = frozenset( [ "prompt", "height", @@ -39,8 +37,8 @@ class FluxControlNetInpaintPipelineTests(unittest.TestCase, PipelineTesterMixin) "controlnet_conditioning_scale", ] ) - batch_params = frozenset(["prompt", "image", "mask_image", "control_image"]) - test_xformers_attention = False + batch_input_params = frozenset(["prompt", "image", "mask_image", "control_image"]) + output_shape = (3, 32, 32) def get_dummy_components(self): torch.manual_seed(0) @@ -74,7 +72,9 @@ def get_dummy_components(self): torch.manual_seed(0) config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5") - text_encoder_2 = T5EncoderModel(config) + # `eval()` because a directly constructed model stays in training mode, which leaves T5's + # dropout active and makes the pipeline outputs non-deterministic across calls. + text_encoder_2 = T5EncoderModel(config).eval() tokenizer = CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip") tokenizer_2 = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5") @@ -120,73 +120,59 @@ def get_dummy_components(self): "controlnet": controlnet, } - 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): + image = floats_tensor((1, 3, 32, 32), rng=random.Random(0)).to(torch_device) + mask_image = torch.ones((1, 1, 32, 32)).to(torch_device) + control_image = floats_tensor((1, 3, 32, 32), rng=random.Random(0)).to(torch_device) - image = floats_tensor((1, 3, 32, 32), rng=random.Random(seed)).to(device) - mask_image = torch.ones((1, 1, 32, 32)).to(device) - control_image = floats_tensor((1, 3, 32, 32), rng=random.Random(seed)).to(device) - - inputs = { + return { "prompt": "A painting of a squirrel eating a burger", "image": image, "mask_image": mask_image, "control_image": control_image, - "generator": generator, + "generator": self.get_generator(0), "num_inference_steps": 2, "guidance_scale": 5.0, "height": 32, "width": 32, "max_sequence_length": 48, "strength": 0.8, - "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 TestFluxControlNetInpaintPipeline(FluxControlNetInpaintPipelineTesterConfig, PipelineTesterMixin): def test_flux_controlnet_inpaint_with_num_images_per_prompt(self): - device = "cpu" # ensure determinism for the device-dependent torch.Generator - components = self.get_dummy_components() - pipe = self.pipeline_class(**components) - pipe = pipe.to(device) - pipe.set_progress_bar_config(disable=None) + pipe = self.get_pipeline().to(torch_device) - inputs = self.get_dummy_inputs(device) + inputs = self.get_dummy_inputs() inputs["num_images_per_prompt"] = 2 - output = pipe(**inputs) - images = output.images + images = pipe(**inputs).images - assert images.shape == (2, 32, 32, 3) + assert images.shape == (2, *self.output_shape) def test_flux_controlnet_inpaint_with_controlnet_conditioning_scale(self): - device = "cpu" # ensure determinism for the device-dependent torch.Generator - components = self.get_dummy_components() - pipe = self.pipeline_class(**components) - pipe = pipe.to(device) - pipe.set_progress_bar_config(disable=None) + pipe = self.get_pipeline().to(torch_device) - inputs = self.get_dummy_inputs(device) - output_default = pipe(**inputs) - image_default = output_default.images + inputs = self.get_dummy_inputs() + image_default = pipe(**inputs).images inputs["controlnet_conditioning_scale"] = 0.5 - output_scaled = pipe(**inputs) - image_scaled = output_scaled.images + image_scaled = pipe(**inputs).images # Ensure that changing the controlnet_conditioning_scale produces a different output - assert not np.allclose(image_default, image_scaled, atol=0.01) - - def test_attention_slicing_forward_pass(self): - super().test_attention_slicing_forward_pass(expected_max_diff=3e-3) + assert not torch.allclose(image_default, image_scaled, atol=0.01), ( + "Changing `controlnet_conditioning_scale` should change the output." + ) - def test_inference_batch_single_identical(self): - super().test_inference_batch_single_identical(expected_max_diff=3e-3) + def test_inference_batch_single_identical(self, batch_size=3, expected_max_diff=3e-3): + super().test_inference_batch_single_identical(batch_size=batch_size, expected_max_diff=expected_max_diff) def test_flux_image_output_shape(self): - pipe = self.pipeline_class(**self.get_dummy_components()).to(torch_device) - inputs = self.get_dummy_inputs(torch_device) + pipe = self.get_pipeline().to(torch_device) + inputs = self.get_dummy_inputs() height_width_pairs = [(32, 32), (72, 56)] for height, width in height_width_pairs: @@ -211,5 +197,11 @@ def test_flux_image_output_shape(self): } ) image = pipe(**inputs).images[0] - output_height, output_width, _ = image.shape - assert (output_height, output_width) == (expected_height, expected_width) + _, output_height, output_width = image.shape + assert (output_height, output_width) == (expected_height, expected_width), ( + f"Output shape {image.shape} does not match expected shape {(expected_height, expected_width)}" + ) + + +class TestFluxControlNetInpaintPipelineMemory(FluxControlNetInpaintPipelineTesterConfig, MemoryTesterMixin): + """Memory optimization tests (CPU offload, group offload, layerwise casting) for the Flux ControlNet inpaint pipeline.""" diff --git a/tests/pipelines/controlnet_hunyuandit/test_controlnet_hunyuandit.py b/tests/pipelines/controlnet_hunyuandit/test_controlnet_hunyuandit.py index 1f765a1675ae..1ce485968654 100644 --- a/tests/pipelines/controlnet_hunyuandit/test_controlnet_hunyuandit.py +++ b/tests/pipelines/controlnet_hunyuandit/test_controlnet_hunyuandit.py @@ -14,9 +14,9 @@ # limitations under the License. import gc -import unittest import numpy as np +import pytest import torch from transformers import AutoConfig, AutoTokenizer, BertModel, T5EncoderModel @@ -31,21 +31,22 @@ from diffusers.utils.torch_utils import randn_tensor from ...testing_utils import ( + assert_tensors_close, backend_empty_cache, enable_full_determinism, require_torch_accelerator, slow, torch_device, ) -from ..test_pipelines_common import PipelineTesterMixin +from ..testing_utils import BasePipelineTesterConfig, MemoryTesterMixin, PipelineTesterMixin enable_full_determinism() -class HunyuanDiTControlNetPipelineFastTests(unittest.TestCase, PipelineTesterMixin): +class HunyuanDiTControlNetPipelineTesterConfig(BasePipelineTesterConfig): pipeline_class = HunyuanDiTControlNetPipeline - params = frozenset( + required_input_params_in_call_signature = frozenset( [ "prompt", "height", @@ -56,8 +57,8 @@ class HunyuanDiTControlNetPipelineFastTests(unittest.TestCase, PipelineTesterMix "negative_prompt_embeds", ] ) - batch_params = frozenset(["prompt", "negative_prompt"]) - test_layerwise_casting = True + batch_input_params = frozenset(["prompt", "negative_prompt"]) + output_shape = (3, 16, 16) def get_dummy_components(self): torch.manual_seed(0) @@ -99,10 +100,12 @@ def get_dummy_components(self): torch.manual_seed(0) config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5") - text_encoder_2 = T5EncoderModel(config) + # `eval()` because a directly constructed model stays in training mode, which leaves T5's + # dropout active and makes the pipeline outputs non-deterministic across calls. + text_encoder_2 = T5EncoderModel(config).eval() tokenizer_2 = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5") - components = { + return { "transformer": transformer.eval(), "vae": vae.eval(), "scheduler": scheduler, @@ -114,90 +117,81 @@ def get_dummy_components(self): "feature_extractor": None, "controlnet": controlnet, } - 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) + def get_dummy_inputs(self): + # The control image is drawn from the same generator that is handed to the pipeline, so the pipeline sees an + # already-advanced generator state — keep the order to stay comparable with the expected slice below. + generator = self.get_generator(0) control_image = randn_tensor( (1, 3, 16, 16), generator=generator, - device=torch.device(device), + device=torch.device(torch_device), dtype=torch.float32, ) - controlnet_conditioning_scale = 0.5 - - inputs = { + return { "prompt": "A painting of a squirrel eating a burger", "generator": generator, "num_inference_steps": 2, "guidance_scale": 5.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", "control_image": control_image, - "controlnet_conditioning_scale": controlnet_conditioning_scale, + "controlnet_conditioning_scale": 0.5, } - return inputs +class TestHunyuanDiTControlNetPipeline(HunyuanDiTControlNetPipelineTesterConfig, PipelineTesterMixin): def test_controlnet_hunyuandit(self): - components = self.get_dummy_components() - pipe = HunyuanDiTControlNetPipeline(**components) - pipe = pipe.to(torch_device, dtype=torch.float32) - pipe.set_progress_bar_config(disable=None) + pipe = self.get_pipeline().to(torch_device, dtype=torch.float32) - inputs = self.get_dummy_inputs(torch_device) - output = pipe(**inputs) - image = output.images + image = pipe(**self.get_dummy_inputs()).images + assert image.shape == (1, *self.output_shape) - image_slice = image[0, -3:, -3:, -1] - assert image.shape == (1, 16, 16, 3) + image_slice = image[0, -1, -3:, -3:] + # fmt: off + expected_slice = torch.tensor([0.5925, 0.5392, 0.4450, 0.7140, 0.3954, 0.3553, 0.3842, 0.5994, 0.3765]) + # fmt: on - expected_slice = np.array([0.5925, 0.5392, 0.4450, 0.7140, 0.3954, 0.3553, 0.3842, 0.5994, 0.3765]) + assert_tensors_close(image_slice.flatten().cpu(), expected_slice, atol=1e-2) - assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2, ( - f"Expected: {expected_slice}, got: {image_slice.flatten()}" - ) - - def test_inference_batch_single_identical(self): - self._test_inference_batch_single_identical( - expected_max_diff=1e-3, - ) - - def test_sequential_cpu_offload_forward_pass(self): - # TODO(YiYi) need to fix later - pass - - def test_sequential_offload_forward_pass_twice(self): - # TODO(YiYi) need to fix later - pass + 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) + @pytest.mark.skip("TODO(YiYi) need to fix later") def test_save_load_optional_components(self): - # TODO(YiYi) need to fix later pass - @unittest.skip( + @pytest.mark.skip( "Test not supported as `encode_prompt` is called two times separately which deivates from about 99% of the pipelines we have." ) def test_encode_prompt_works_in_isolation(self): pass +class TestHunyuanDiTControlNetPipelineMemory(HunyuanDiTControlNetPipelineTesterConfig, MemoryTesterMixin): + """Memory optimization tests (CPU offload, group offload, layerwise casting) for the HunyuanDiT ControlNet pipeline.""" + + @pytest.mark.skip("TODO(YiYi) need to fix later") + def test_sequential_cpu_offload_forward_pass(self): + pass + + @pytest.mark.skip("TODO(YiYi) need to fix later") + def test_sequential_offload_forward_pass_twice(self): + pass + + @slow @require_torch_accelerator -class HunyuanDiTControlNetPipelineSlowTests(unittest.TestCase): +class TestHunyuanDiTControlNetPipelineSlow: pipeline_class = HunyuanDiTControlNetPipeline - 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) diff --git a/tests/pipelines/controlnet_sd3/test_controlnet_inpaint_sd3.py b/tests/pipelines/controlnet_sd3/test_controlnet_inpaint_sd3.py index 09bce003379b..554fbc150e09 100644 --- a/tests/pipelines/controlnet_sd3/test_controlnet_inpaint_sd3.py +++ b/tests/pipelines/controlnet_sd3/test_controlnet_inpaint_sd3.py @@ -1,206 +1,202 @@ -# 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 unittest - -import numpy as np -import torch -from transformers import ( - AutoConfig, - AutoTokenizer, - CLIPTextConfig, - CLIPTextModelWithProjection, - CLIPTokenizer, - T5EncoderModel, -) - -from diffusers import ( - AutoencoderKL, - FlowMatchEulerDiscreteScheduler, - SD3Transformer2DModel, - StableDiffusion3ControlNetInpaintingPipeline, -) -from diffusers.models import SD3ControlNetModel -from diffusers.utils.torch_utils import randn_tensor - -from ...testing_utils import enable_full_determinism, torch_device -from ..test_pipelines_common import PipelineTesterMixin - - -enable_full_determinism() - - -class StableDiffusion3ControlInpaintNetPipelineFastTests(unittest.TestCase, PipelineTesterMixin): - pipeline_class = StableDiffusion3ControlNetInpaintingPipeline - params = frozenset( - [ - "prompt", - "height", - "width", - "guidance_scale", - "negative_prompt", - "prompt_embeds", - "negative_prompt_embeds", - ] - ) - batch_params = frozenset(["prompt", "negative_prompt"]) - - def get_dummy_components(self): - torch.manual_seed(0) - transformer = SD3Transformer2DModel( - sample_size=32, - patch_size=1, - in_channels=8, - num_layers=4, - attention_head_dim=8, - num_attention_heads=4, - joint_attention_dim=32, - caption_projection_dim=32, - pooled_projection_dim=64, - out_channels=8, - ) - - torch.manual_seed(0) - controlnet = SD3ControlNetModel( - sample_size=32, - patch_size=1, - in_channels=8, - num_layers=1, - attention_head_dim=8, - num_attention_heads=4, - joint_attention_dim=32, - caption_projection_dim=32, - pooled_projection_dim=64, - out_channels=8, - extra_conditioning_channels=1, - ) - clip_text_encoder_config = CLIPTextConfig( - bos_token_id=0, - eos_token_id=2, - hidden_size=32, - intermediate_size=37, - layer_norm_eps=1e-05, - num_attention_heads=4, - num_hidden_layers=5, - pad_token_id=1, - vocab_size=1000, - hidden_act="gelu", - projection_dim=32, - ) - - torch.manual_seed(0) - text_encoder = CLIPTextModelWithProjection(clip_text_encoder_config) - - torch.manual_seed(0) - text_encoder_2 = CLIPTextModelWithProjection(clip_text_encoder_config) - - torch.manual_seed(0) - config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5") - text_encoder_3 = T5EncoderModel(config) - - tokenizer = CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip") - tokenizer_2 = CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip") - tokenizer_3 = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5") - - torch.manual_seed(0) - vae = AutoencoderKL( - sample_size=32, - in_channels=3, - out_channels=3, - block_out_channels=(4,), - layers_per_block=1, - latent_channels=8, - norm_num_groups=1, - use_quant_conv=False, - use_post_quant_conv=False, - shift_factor=0.0609, - scaling_factor=1.5035, - ) - - scheduler = FlowMatchEulerDiscreteScheduler() - - return { - "scheduler": scheduler, - "text_encoder": text_encoder, - "text_encoder_2": text_encoder_2, - "text_encoder_3": text_encoder_3, - "tokenizer": tokenizer, - "tokenizer_2": tokenizer_2, - "tokenizer_3": tokenizer_3, - "transformer": transformer, - "vae": vae, - "controlnet": controlnet, - "image_encoder": None, - "feature_extractor": None, - } - - 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) - - control_image = randn_tensor( - (1, 3, 32, 32), - generator=generator, - device=torch.device(device), - dtype=torch.float32, - ) - - control_mask = randn_tensor( - (1, 1, 32, 32), - generator=generator, - device=torch.device(device), - dtype=torch.float32, - ) - - controlnet_conditioning_scale = 0.95 - - inputs = { - "prompt": "A painting of a squirrel eating a burger", - "generator": generator, - "num_inference_steps": 2, - "guidance_scale": 7.0, - "output_type": "np", - "control_image": control_image, - "control_mask": control_mask, - "controlnet_conditioning_scale": controlnet_conditioning_scale, - } - - return inputs - - def test_controlnet_inpaint_sd3(self): - components = self.get_dummy_components() - sd_pipe = StableDiffusion3ControlNetInpaintingPipeline(**components) - sd_pipe = sd_pipe.to(torch_device, dtype=torch.float32) - sd_pipe.set_progress_bar_config(disable=None) - - inputs = self.get_dummy_inputs(torch_device) - output = sd_pipe(**inputs) - image = output.images - - image_slice = image[0, -3:, -3:, -1] - - assert image.shape == (1, 32, 32, 3) - - expected_slice = np.array([0.4627, 0.3686, 0.3741, 0.5855, 0.6071, 0.4046, 0.1916, 0.3938, 0.4953]) - - assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2, ( - f"Expected: {expected_slice}, got: {image_slice.flatten()}" - ) - - @unittest.skip("xFormersAttnProcessor does not work with SD3 Joint Attention") - def test_xformers_attention_forwardGenerator_pass(self): - pass +# 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 torch +from transformers import ( + AutoConfig, + AutoTokenizer, + CLIPTextConfig, + CLIPTextModelWithProjection, + CLIPTokenizer, + T5EncoderModel, +) + +from diffusers import ( + AutoencoderKL, + FlowMatchEulerDiscreteScheduler, + SD3Transformer2DModel, + StableDiffusion3ControlNetInpaintingPipeline, +) +from diffusers.models import SD3ControlNetModel +from diffusers.utils.torch_utils import randn_tensor + +from ...testing_utils import assert_tensors_close, enable_full_determinism, torch_device +from ..testing_utils import BasePipelineTesterConfig, MemoryTesterMixin, PipelineTesterMixin + + +enable_full_determinism() + + +class StableDiffusion3ControlNetInpaintingPipelineTesterConfig(BasePipelineTesterConfig): + pipeline_class = StableDiffusion3ControlNetInpaintingPipeline + 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 = (3, 32, 32) + + def get_dummy_components(self): + torch.manual_seed(0) + transformer = SD3Transformer2DModel( + sample_size=32, + patch_size=1, + in_channels=8, + num_layers=4, + attention_head_dim=8, + num_attention_heads=4, + joint_attention_dim=32, + caption_projection_dim=32, + pooled_projection_dim=64, + out_channels=8, + ) + + torch.manual_seed(0) + controlnet = SD3ControlNetModel( + sample_size=32, + patch_size=1, + in_channels=8, + num_layers=1, + attention_head_dim=8, + num_attention_heads=4, + joint_attention_dim=32, + caption_projection_dim=32, + pooled_projection_dim=64, + out_channels=8, + extra_conditioning_channels=1, + ) + clip_text_encoder_config = CLIPTextConfig( + bos_token_id=0, + eos_token_id=2, + hidden_size=32, + intermediate_size=37, + layer_norm_eps=1e-05, + num_attention_heads=4, + num_hidden_layers=5, + pad_token_id=1, + vocab_size=1000, + hidden_act="gelu", + projection_dim=32, + ) + + torch.manual_seed(0) + text_encoder = CLIPTextModelWithProjection(clip_text_encoder_config) + + torch.manual_seed(0) + text_encoder_2 = CLIPTextModelWithProjection(clip_text_encoder_config) + + torch.manual_seed(0) + config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5") + # `eval()` because a directly constructed model stays in training mode, which leaves T5's + # dropout active and makes the pipeline outputs non-deterministic across calls. + text_encoder_3 = T5EncoderModel(config).eval() + + tokenizer = CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip") + tokenizer_2 = CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip") + tokenizer_3 = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5") + + torch.manual_seed(0) + vae = AutoencoderKL( + sample_size=32, + in_channels=3, + out_channels=3, + block_out_channels=(4,), + layers_per_block=1, + latent_channels=8, + norm_num_groups=1, + use_quant_conv=False, + use_post_quant_conv=False, + shift_factor=0.0609, + scaling_factor=1.5035, + ) + + scheduler = FlowMatchEulerDiscreteScheduler() + + return { + "scheduler": scheduler, + "text_encoder": text_encoder, + "text_encoder_2": text_encoder_2, + "text_encoder_3": text_encoder_3, + "tokenizer": tokenizer, + "tokenizer_2": tokenizer_2, + "tokenizer_3": tokenizer_3, + "transformer": transformer, + "vae": vae, + "controlnet": controlnet, + "image_encoder": None, + "feature_extractor": None, + } + + def get_dummy_inputs(self): + # The control image and mask are drawn from the same generator that is handed to the pipeline, so the + # pipeline sees an already-advanced generator state — keep the order to stay comparable with the expected + # slice below. + generator = self.get_generator(0) + control_image = randn_tensor( + (1, 3, 32, 32), + generator=generator, + device=torch.device(torch_device), + dtype=torch.float32, + ) + + control_mask = randn_tensor( + (1, 1, 32, 32), + generator=generator, + device=torch.device(torch_device), + dtype=torch.float32, + ) + + return { + "prompt": "A painting of a squirrel eating a burger", + "generator": generator, + "num_inference_steps": 2, + "guidance_scale": 7.0, + # 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", + "control_image": control_image, + "control_mask": control_mask, + "controlnet_conditioning_scale": 0.95, + } + + +class TestStableDiffusion3ControlNetInpaintingPipeline( + StableDiffusion3ControlNetInpaintingPipelineTesterConfig, PipelineTesterMixin +): + def test_controlnet_inpaint_sd3(self): + pipe = self.get_pipeline().to(torch_device, dtype=torch.float32) + + image = pipe(**self.get_dummy_inputs()).images + assert image.shape == (1, *self.output_shape) + + image_slice = image[0, -1, -3:, -3:] + # fmt: off + expected_slice = torch.tensor([0.4627, 0.3686, 0.3741, 0.5855, 0.6071, 0.4046, 0.1916, 0.3938, 0.4953]) + # fmt: on + + assert_tensors_close(image_slice.flatten().cpu(), expected_slice, atol=1e-2) + + +class TestStableDiffusion3ControlNetInpaintingPipelineMemory( + StableDiffusion3ControlNetInpaintingPipelineTesterConfig, MemoryTesterMixin +): + """Memory optimization tests (CPU offload, group offload, layerwise casting) for the SD3 ControlNet inpainting pipeline.""" diff --git a/tests/pipelines/controlnet_sd3/test_controlnet_sd3.py b/tests/pipelines/controlnet_sd3/test_controlnet_sd3.py index eca13af8340f..5fa6770dffdf 100644 --- a/tests/pipelines/controlnet_sd3/test_controlnet_sd3.py +++ b/tests/pipelines/controlnet_sd3/test_controlnet_sd3.py @@ -14,9 +14,9 @@ # limitations under the License. import gc -import unittest import numpy as np +import pytest import torch from transformers import ( AutoConfig, @@ -38,6 +38,7 @@ from diffusers.utils.torch_utils import randn_tensor from ...testing_utils import ( + assert_tensors_close, backend_empty_cache, enable_full_determinism, numpy_cosine_similarity_distance, @@ -45,15 +46,15 @@ slow, torch_device, ) -from ..test_pipelines_common import PipelineTesterMixin +from ..testing_utils import BasePipelineTesterConfig, MemoryTesterMixin, PipelineTesterMixin enable_full_determinism() -class StableDiffusion3ControlNetPipelineFastTests(unittest.TestCase, PipelineTesterMixin): +class StableDiffusion3ControlNetPipelineTesterConfig(BasePipelineTesterConfig): pipeline_class = StableDiffusion3ControlNetPipeline - params = frozenset( + required_input_params_in_call_signature = frozenset( [ "prompt", "height", @@ -64,9 +65,8 @@ class StableDiffusion3ControlNetPipelineFastTests(unittest.TestCase, PipelineTes "negative_prompt_embeds", ] ) - batch_params = frozenset(["prompt", "negative_prompt"]) - test_layerwise_casting = True - test_group_offloading = True + batch_input_params = frozenset(["prompt", "negative_prompt"]) + output_shape = (3, 32, 32) def get_dummy_components( self, num_controlnet_layers: int = 3, qk_norm: str | None = "rms_norm", use_dual_attention=False @@ -125,7 +125,9 @@ def get_dummy_components( torch.manual_seed(0) config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5") - text_encoder_3 = T5EncoderModel(config) + # `eval()` because a directly constructed model stays in training mode, which leaves T5's + # dropout active and makes the pipeline outputs non-deterministic across calls. + text_encoder_3 = T5EncoderModel(config).eval() tokenizer = CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip") tokenizer_2 = CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip") @@ -163,80 +165,68 @@ def get_dummy_components( "feature_extractor": None, } - 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) - + def get_dummy_inputs(self): + # The control image is drawn from the same generator that is handed to the pipeline, so the pipeline sees an + # already-advanced generator state — keep the order to stay comparable with the expected slices below. + generator = self.get_generator(0) control_image = randn_tensor( (1, 3, 32, 32), generator=generator, - device=torch.device(device), + device=torch.device(torch_device), dtype=torch.float32, ) - controlnet_conditioning_scale = 0.5 - - inputs = { + return { "prompt": "A painting of a squirrel eating a burger", "generator": generator, "num_inference_steps": 2, "guidance_scale": 5.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", "control_image": control_image, - "controlnet_conditioning_scale": controlnet_conditioning_scale, + "controlnet_conditioning_scale": 0.5, } - return inputs - - def run_pipe(self, components, use_sd35=False): - sd_pipe = StableDiffusion3ControlNetPipeline(**components) - sd_pipe = sd_pipe.to(torch_device, dtype=torch.float32) - sd_pipe.set_progress_bar_config(disable=None) - - inputs = self.get_dummy_inputs(torch_device) - output = sd_pipe(**inputs) - image = output.images - image_slice = image[0, -3:, -3:, -1] +class TestStableDiffusion3ControlNetPipeline(StableDiffusion3ControlNetPipelineTesterConfig, PipelineTesterMixin): + def _run_and_check_slice(self, components, expected_slice): + pipe = self.get_pipeline(**components).to(torch_device, dtype=torch.float32) - assert image.shape == (1, 32, 32, 3) + image = pipe(**self.get_dummy_inputs()).images + assert image.shape == (1, *self.output_shape) - if not use_sd35: - expected_slice = np.array([0.4121, 0.3775, 0.3734, 0.1509, 0.6324, 0.5503, 0.5425, 0.5614, 0.4061]) - else: - expected_slice = np.array([0.3793, 0.5179, 0.4389, 0.2820, 0.5148, 0.5565, 0.6282, 0.6891, 0.4197]) - - assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2, ( - f"Expected: {expected_slice}, got: {image_slice.flatten()}" - ) + image_slice = image[0, -1, -3:, -3:] + assert_tensors_close(image_slice.flatten().cpu(), expected_slice, atol=1e-2) def test_controlnet_sd3(self): - components = self.get_dummy_components() - self.run_pipe(components) + # fmt: off + expected_slice = torch.tensor([0.4121, 0.3775, 0.3734, 0.1509, 0.6324, 0.5503, 0.5425, 0.5614, 0.4061]) + # fmt: on + self._run_and_check_slice(self.get_dummy_components(), expected_slice) def test_controlnet_sd35(self): components = self.get_dummy_components(num_controlnet_layers=1, qk_norm="rms_norm", use_dual_attention=True) - self.run_pipe(components, use_sd35=True) + # fmt: off + expected_slice = torch.tensor([0.3793, 0.5179, 0.4389, 0.2820, 0.5148, 0.5565, 0.6282, 0.6891, 0.4197]) + # fmt: on + self._run_and_check_slice(components, expected_slice) - @unittest.skip("xFormersAttnProcessor does not work with SD3 Joint Attention") - def test_xformers_attention_forwardGenerator_pass(self): - pass + +class TestStableDiffusion3ControlNetPipelineMemory(StableDiffusion3ControlNetPipelineTesterConfig, MemoryTesterMixin): + """Memory optimization tests (CPU offload, group offload, layerwise casting) for the SD3 ControlNet pipeline.""" @slow @require_big_accelerator -class StableDiffusion3ControlNetPipelineSlowTests(unittest.TestCase): +class TestStableDiffusion3ControlNetPipelineSlow: pipeline_class = StableDiffusion3ControlNetPipeline - 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) diff --git a/tests/pipelines/cosmos/test_cosmos.py b/tests/pipelines/cosmos/test_cosmos.py index 85e9716b23fa..794a028add9a 100644 --- a/tests/pipelines/cosmos/test_cosmos.py +++ b/tests/pipelines/cosmos/test_cosmos.py @@ -12,22 +12,15 @@ # See the License for the specific language governing permissions and # limitations under the License. -import inspect -import json -import os -import tempfile -import unittest - -import numpy as np import torch from transformers import AutoConfig, AutoTokenizer, T5EncoderModel from diffusers import AutoencoderKLCosmos, CosmosTextToWorldPipeline, CosmosTransformer3DModel, EDMEulerScheduler -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 +from ..testing_utils import BasePipelineTesterConfig, MemoryTesterMixin, PipelineTesterMixin from .cosmos_guardrail import DummyCosmosSafetyChecker +from .testing_utils import CosmosSafetyCheckerTesterMixin enable_full_determinism() @@ -40,25 +33,17 @@ def from_pretrained(*args, **kwargs): return CosmosTextToWorldPipeline.from_pretrained(*args, **kwargs) -class CosmosTextToWorldPipelineFastTests(PipelineTesterMixin, unittest.TestCase): +class CosmosTextToWorldPipelineTesterConfig(BasePipelineTesterConfig): pipeline_class = CosmosTextToWorldPipelineWrapper - 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( - [ - "num_inference_steps", - "generator", - "latents", - "return_dict", - "callback_on_step_end", - "callback_on_step_end_tensor_inputs", - ] + 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 = (9, 3, 32, 32) + # Cosmos text-to-world is a video pipeline: it exposes `num_videos_per_prompt`, not `num_images_per_prompt`. + optional_input_params = frozenset( + ["num_inference_steps", "num_videos_per_prompt", "generator", "latents", "output_type", "return_dict"] ) - test_xformers_attention = False - test_layerwise_casting = True - test_group_offloading = True def get_dummy_components(self): torch.manual_seed(0) @@ -107,10 +92,12 @@ def get_dummy_components(self): final_sigmas_type="sigma_min", ) config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5") - text_encoder = T5EncoderModel(config) + # `eval()` because a directly constructed model stays in training mode, which leaves T5's + # dropout active and makes the pipeline outputs non-deterministic across calls. + text_encoder = T5EncoderModel(config).eval() tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5") - components = { + return { "transformer": transformer, "vae": vae, "scheduler": scheduler, @@ -119,41 +106,33 @@ def get_dummy_components(self): # We cannot run the Cosmos Guardrail for fast tests due to the large model size "safety_checker": DummyCosmosSafetyChecker(), } - 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": "dance monkey", "negative_prompt": "bad quality", - "generator": generator, + "generator": self.get_generator(0), "num_inference_steps": 2, "guidance_scale": 3.0, "height": 32, "width": 32, "num_frames": 9, "max_sequence_length": 16, + # Request torch outputs so tests compare torch tensors directly (see `BasePipelineTesterConfig`). "output_type": "pt", } - return inputs +class TestCosmosTextToWorldPipeline( + CosmosTextToWorldPipelineTesterConfig, CosmosSafetyCheckerTesterMixin, PipelineTesterMixin +): def test_inference(self): - device = "cpu" + # 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.to(device) - pipe.set_progress_bar_config(disable=None) - - inputs = self.get_dummy_inputs(device) - video = pipe(**inputs).frames + video = pipe(**self.get_dummy_inputs()).frames generated_video = video[0] - self.assertEqual(generated_video.shape, (9, 3, 32, 32)) + assert generated_video.shape == self.output_shape # fmt: off expected_slice = torch.tensor([0.0, 0.9686, 0.8549, 0.8078, 0.0, 0.8431, 1.0, 0.4863, 0.7098, 0.1098, 0.8157, 0.4235, 0.6353, 0.2549, 0.5137, 0.5333]) @@ -161,118 +140,16 @@ def test_inference(self): generated_slice = generated_video.flatten() generated_slice = torch.cat([generated_slice[:8], generated_slice[-8:]]) - self.assertTrue(torch.allclose(generated_slice, expected_slice, atol=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", - ) - - 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-2) + assert_tensors_close(generated_slice, expected_slice, atol=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() - for key in components: - if "text_encoder" in key and hasattr(components[key], "eval"): - components[key].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) - - 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-2): + 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() # 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] @@ -283,96 +160,14 @@ def test_vae_tiling(self, expected_diff_max: float = 0.2): 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).max() < expected_diff_max, ( + "VAE tiling should not affect the inference results" ) - def test_save_load_optional_components(self, expected_max_difference=1e-4): - self.pipeline_class._optional_components.remove("safety_checker") - super().test_save_load_optional_components(expected_max_difference=expected_max_difference) - self.pipeline_class._optional_components.append("safety_checker") - - def test_serialization_with_variants(self): - components = self.get_dummy_components() - pipe = self.pipeline_class(**components) - model_components = [ - component_name - for component_name, component in pipe.components.items() - if isinstance(component, torch.nn.Module) - ] - model_components.remove("safety_checker") - variant = "fp16" - - with tempfile.TemporaryDirectory() as tmpdir: - pipe.save_pretrained(tmpdir, variant=variant, safe_serialization=False) - - with open(f"{tmpdir}/model_index.json", "r") as f: - config = json.load(f) - - for subfolder in os.listdir(tmpdir): - if not os.path.isfile(subfolder) and subfolder in model_components: - folder_path = os.path.join(tmpdir, subfolder) - is_folder = os.path.isdir(folder_path) and subfolder in config - assert is_folder and any(p.split(".")[1].startswith(variant) for p in os.listdir(folder_path)) - def test_dtype_dict(self): - components = self.get_dummy_components() - if not components: - self.skipTest("No dummy components defined.") - - pipe = self.pipeline_class(**components) - - specified_key = next(iter(components.keys())) - - with tempfile.TemporaryDirectory(ignore_cleanup_errors=True) as tmpdirname: - pipe.save_pretrained(tmpdirname, safe_serialization=False) - dtype_dict = {specified_key: torch.bfloat16, "default": torch.float16} - loaded_pipe = self.pipeline_class.from_pretrained( - tmpdirname, safety_checker=DummyCosmosSafetyChecker(), dtype=dtype_dict - ) - - for name, component in loaded_pipe.components.items(): - if name == "safety_checker": - continue - if isinstance(component, torch.nn.Module) and hasattr(component, "dtype"): - expected_dtype = dtype_dict.get(name, dtype_dict.get("default", torch.float32)) - self.assertEqual( - component.dtype, - expected_dtype, - f"Component '{name}' has dtype {component.dtype} but expected {expected_dtype}", - ) - - def test_dtype_alias(self): - # `torch_dtype` is deprecated in favor of `dtype` in `from_pretrained`. - components = self.get_dummy_components() - pipe = self.pipeline_class(**components) - - with tempfile.TemporaryDirectory(ignore_cleanup_errors=True) as tmpdirname: - pipe.save_pretrained(tmpdirname, safe_serialization=False) - loaded_pipe = self.pipeline_class.from_pretrained( - tmpdirname, safety_checker=DummyCosmosSafetyChecker(), dtype=torch.float16 - ) - - for name, component in loaded_pipe.components.items(): - if name == "safety_checker": - continue - if isinstance(component, torch.nn.Module) and hasattr(component, "dtype"): - self.assertEqual( - component.dtype, - torch.float16, - f"Component '{name}' has dtype {component.dtype} but expected {torch.float16}", - ) - - @unittest.skip( - "The pipeline should not be runnable without a safety checker. The test creates a pipeline without passing in " - "a safety checker, which makes the pipeline default to the actual Cosmos Guardrail. The Cosmos Guardrail is " - "too large and slow to run on CI." - ) - def test_encode_prompt_works_in_isolation(self): - pass +class TestCosmosTextToWorldPipelineMemory(CosmosTextToWorldPipelineTesterConfig, MemoryTesterMixin): + """Memory optimization tests (CPU offload, group offload, layerwise casting) for the Cosmos text-to-world pipeline.""" diff --git a/tests/pipelines/cosmos/test_cosmos2_5_predict.py b/tests/pipelines/cosmos/test_cosmos2_5_predict.py index 32954d74683f..0d75e058362e 100644 --- a/tests/pipelines/cosmos/test_cosmos2_5_predict.py +++ b/tests/pipelines/cosmos/test_cosmos2_5_predict.py @@ -12,13 +12,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -import inspect -import json -import os -import tempfile -import unittest - -import numpy as np import torch from transformers import Qwen2_5_VLConfig, Qwen2_5_VLForConditionalGeneration, Qwen2Tokenizer @@ -29,10 +22,10 @@ UniPCMultistepScheduler, ) -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 enable_full_determinism +from ..testing_utils import BasePipelineTesterConfig, MemoryTesterMixin, PipelineTesterMixin from .cosmos_guardrail import DummyCosmosSafetyChecker +from .testing_utils import CosmosSafetyCheckerTesterMixin enable_full_determinism() @@ -51,25 +44,17 @@ def from_pretrained(*args, **kwargs): return Cosmos2_5_PredictBasePipeline.from_pretrained(*args, **kwargs) -class Cosmos2_5_PredictPipelineFastTests(PipelineTesterMixin, unittest.TestCase): +class Cosmos2_5_PredictPipelineTesterConfig(BasePipelineTesterConfig): pipeline_class = Cosmos2_5_PredictBaseWrapper - 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( - [ - "num_inference_steps", - "generator", - "latents", - "return_dict", - "callback_on_step_end", - "callback_on_step_end_tensor_inputs", - ] + 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 = (3, 3, 32, 32) + # Cosmos2.5 predict is a video pipeline: it exposes `num_videos_per_prompt`, not `num_images_per_prompt`. + optional_input_params = frozenset( + ["num_inference_steps", "num_videos_per_prompt", "generator", "latents", "output_type", "return_dict"] ) - test_xformers_attention = False - test_layerwise_casting = True - test_group_offloading = True def get_dummy_components(self): torch.manual_seed(0) @@ -132,7 +117,7 @@ def get_dummy_components(self): text_encoder = Qwen2_5_VLForConditionalGeneration(config) tokenizer = Qwen2Tokenizer.from_pretrained("hf-internal-testing/tiny-random-Qwen2VLForConditionalGeneration") - components = { + return { "transformer": transformer, "vae": vae, "scheduler": scheduler, @@ -140,218 +125,38 @@ def get_dummy_components(self): "tokenizer": tokenizer, "safety_checker": DummyCosmosSafetyChecker(), } - 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": "dance monkey", "negative_prompt": "bad quality", - "generator": generator, + "generator": self.get_generator(0), "num_inference_steps": 2, "guidance_scale": 3.0, "height": 32, "width": 32, "num_frames": 3, "max_sequence_length": 16, + # Request torch outputs so tests compare torch tensors directly (see `BasePipelineTesterConfig`). "output_type": "pt", } - return inputs - - def test_components_function(self): - init_components = self.get_dummy_components() - init_components = {k: v for k, v in init_components.items() if not isinstance(v, (str, int, float))} - pipe = self.pipeline_class(**init_components) - self.assertTrue(hasattr(pipe, "components")) - self.assertTrue(set(pipe.components.keys()) == set(init_components.keys())) +class TestCosmos2_5_PredictPipeline( + Cosmos2_5_PredictPipelineTesterConfig, CosmosSafetyCheckerTesterMixin, PipelineTesterMixin +): 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) + pipe = self.get_pipeline() - inputs = self.get_dummy_inputs(device) - video = pipe(**inputs).frames + video = pipe(**self.get_dummy_inputs()).frames generated_video = video[0] - self.assertEqual(generated_video.shape, (3, 3, 32, 32)) - self.assertTrue(torch.isfinite(generated_video).all()) - - 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): - for tensor_name in callback_kwargs.keys(): - 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 - for tensor_name in callback_kwargs.keys(): - assert tensor_name in pipe._callback_tensor_inputs - return callback_kwargs - - inputs = self.get_dummy_inputs(torch_device) - - inputs["callback_on_step_end"] = callback_inputs_subset - inputs["callback_on_step_end_tensor_inputs"] = ["latents"] - _ = pipe(**inputs)[0] - - inputs["callback_on_step_end"] = callback_inputs_all - inputs["callback_on_step_end_tensor_inputs"] = pipe._callback_tensor_inputs - _ = 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=2, expected_max_diff=1e-2) - def test_attention_slicing_forward_pass( - self, test_max_difference=True, test_mean_pixel_difference=True, expected_max_diff=1e-3 - ): - if not getattr(self, "test_attention_slicing", True): - return + assert generated_video.shape == self.output_shape + assert torch.isfinite(generated_video).all() - 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) + def test_inference_batch_single_identical(self, batch_size=2, expected_max_diff=1e-2): + super().test_inference_batch_single_identical(batch_size=batch_size, expected_max_diff=expected_max_diff) - 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_save_load_optional_components(self, expected_max_difference=1e-4): - self.pipeline_class._optional_components.remove("safety_checker") - super().test_save_load_optional_components(expected_max_difference=expected_max_difference) - self.pipeline_class._optional_components.append("safety_checker") - - def test_serialization_with_variants(self): - components = self.get_dummy_components() - pipe = self.pipeline_class(**components) - model_components = [ - component_name - for component_name, component in pipe.components.items() - if isinstance(component, torch.nn.Module) - ] - model_components.remove("safety_checker") - variant = "fp16" - - with tempfile.TemporaryDirectory() as tmpdir: - pipe.save_pretrained(tmpdir, variant=variant, safe_serialization=False) - - with open(f"{tmpdir}/model_index.json", "r") as f: - config = json.load(f) - - for subfolder in os.listdir(tmpdir): - if not os.path.isfile(subfolder) and subfolder in model_components: - folder_path = os.path.join(tmpdir, subfolder) - is_folder = os.path.isdir(folder_path) and subfolder in config - assert is_folder and any(p.split(".")[1].startswith(variant) for p in os.listdir(folder_path)) - - def test_dtype_dict(self): - components = self.get_dummy_components() - if not components: - self.skipTest("No dummy components defined.") - - pipe = self.pipeline_class(**components) - - specified_key = next(iter(components.keys())) - - with tempfile.TemporaryDirectory(ignore_cleanup_errors=True) as tmpdirname: - pipe.save_pretrained(tmpdirname, safe_serialization=False) - dtype_dict = {specified_key: torch.bfloat16, "default": torch.float16} - loaded_pipe = self.pipeline_class.from_pretrained( - tmpdirname, safety_checker=DummyCosmosSafetyChecker(), dtype=dtype_dict - ) - - for name, component in loaded_pipe.components.items(): - if name == "safety_checker": - continue - if isinstance(component, torch.nn.Module) and hasattr(component, "dtype"): - expected_dtype = dtype_dict.get(name, dtype_dict.get("default", torch.float32)) - self.assertEqual( - component.dtype, - expected_dtype, - f"Component '{name}' has dtype {component.dtype} but expected {expected_dtype}", - ) - - def test_dtype_alias(self): - # `torch_dtype` is deprecated in favor of `dtype` in `from_pretrained`. - components = self.get_dummy_components() - pipe = self.pipeline_class(**components) - - with tempfile.TemporaryDirectory(ignore_cleanup_errors=True) as tmpdirname: - pipe.save_pretrained(tmpdirname, safe_serialization=False) - loaded_pipe = self.pipeline_class.from_pretrained( - tmpdirname, safety_checker=DummyCosmosSafetyChecker(), dtype=torch.float16 - ) - - for name, component in loaded_pipe.components.items(): - if name == "safety_checker": - continue - if isinstance(component, torch.nn.Module) and hasattr(component, "dtype"): - self.assertEqual( - component.dtype, - torch.float16, - f"Component '{name}' has dtype {component.dtype} but expected {torch.float16}", - ) - - @unittest.skip( - "The pipeline should not be runnable without a safety checker. The test creates a pipeline without passing in " - "a safety checker, which makes the pipeline default to the actual Cosmos Guardrail. The Cosmos Guardrail is " - "too large and slow to run on CI." - ) - def test_encode_prompt_works_in_isolation(self): - pass +class TestCosmos2_5_PredictPipelineMemory(Cosmos2_5_PredictPipelineTesterConfig, MemoryTesterMixin): + """Memory optimization tests (CPU offload, group offload, layerwise casting) for the Cosmos2.5 predict pipeline.""" diff --git a/tests/pipelines/cosmos/test_cosmos2_5_transfer.py b/tests/pipelines/cosmos/test_cosmos2_5_transfer.py index f95caabb4243..a99f6d20bac6 100644 --- a/tests/pipelines/cosmos/test_cosmos2_5_transfer.py +++ b/tests/pipelines/cosmos/test_cosmos2_5_transfer.py @@ -12,13 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -import inspect -import json -import os -import tempfile -import unittest - -import numpy as np +import pytest import torch from transformers import Qwen2_5_VLConfig, Qwen2_5_VLForConditionalGeneration, Qwen2Tokenizer @@ -30,10 +24,10 @@ UniPCMultistepScheduler, ) -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 enable_full_determinism +from ..testing_utils import BasePipelineTesterConfig, MemoryTesterMixin, PipelineTesterMixin from .cosmos_guardrail import DummyCosmosSafetyChecker +from .testing_utils import CosmosSafetyCheckerTesterMixin enable_full_determinism() @@ -52,25 +46,17 @@ def from_pretrained(*args, **kwargs): return Cosmos2_5_TransferPipeline.from_pretrained(*args, **kwargs) -class Cosmos2_5_TransferPipelineFastTests(PipelineTesterMixin, unittest.TestCase): +class Cosmos2_5_TransferPipelineTesterConfig(BasePipelineTesterConfig): pipeline_class = Cosmos2_5_TransferWrapper - params = TEXT_TO_IMAGE_PARAMS - {"cross_attention_kwargs"} - batch_params = TEXT_TO_IMAGE_BATCH_PARAMS.union({"controls"}) - image_params = TEXT_TO_IMAGE_IMAGE_PARAMS - image_latents_params = TEXT_TO_IMAGE_IMAGE_PARAMS - required_optional_params = frozenset( - [ - "num_inference_steps", - "generator", - "latents", - "return_dict", - "callback_on_step_end", - "callback_on_step_end_tensor_inputs", - ] + 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", "controls"]) + output_shape = (3, 3, 32, 32) + # Cosmos2.5 transfer is a video pipeline: it exposes `num_videos_per_prompt`, not `num_images_per_prompt`. + optional_input_params = frozenset( + ["num_inference_steps", "num_videos_per_prompt", "generator", "latents", "output_type", "return_dict"] ) - test_xformers_attention = False - test_layerwise_casting = True - test_group_offloading = True def get_dummy_components(self): torch.manual_seed(0) @@ -158,7 +144,7 @@ def get_dummy_components(self): text_encoder = Qwen2_5_VLForConditionalGeneration(config) tokenizer = Qwen2Tokenizer.from_pretrained("hf-internal-testing/tiny-random-Qwen2VLForConditionalGeneration") - components = { + return { "transformer": transformer, "controlnet": controlnet, "vae": vae, @@ -167,21 +153,15 @@ def get_dummy_components(self): "tokenizer": tokenizer, "safety_checker": DummyCosmosSafetyChecker(), } - 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) - controls_generator = torch.Generator(device="cpu").manual_seed(seed) + def get_dummy_inputs(self): + controls_generator = torch.Generator(device="cpu").manual_seed(0) - inputs = { + return { "prompt": "dance monkey", "negative_prompt": "bad quality", "controls": [torch.randn(3, 32, 32, generator=controls_generator) for _ in range(5)], - "generator": generator, + "generator": self.get_generator(0), "num_inference_steps": 2, "guidance_scale": 3.0, "height": 32, @@ -189,273 +169,78 @@ def get_dummy_inputs(self, device, seed=0): "num_frames": 3, "num_frames_per_chunk": 16, "max_sequence_length": 16, + # Request torch outputs so tests compare torch tensors directly (see `BasePipelineTesterConfig`). "output_type": "pt", } - return inputs - - def test_components_function(self): - init_components = self.get_dummy_components() - init_components = {k: v for k, v in init_components.items() if not isinstance(v, (str, int, float))} - pipe = self.pipeline_class(**init_components) - self.assertTrue(hasattr(pipe, "components")) - self.assertTrue(set(pipe.components.keys()) == set(init_components.keys())) +class TestCosmos2_5_TransferPipeline( + Cosmos2_5_TransferPipelineTesterConfig, CosmosSafetyCheckerTesterMixin, PipelineTesterMixin +): 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) + pipe = self.get_pipeline() - inputs = self.get_dummy_inputs(device) - video = pipe(**inputs).frames + video = pipe(**self.get_dummy_inputs()).frames generated_video = video[0] - self.assertEqual(generated_video.shape, (3, 3, 32, 32)) - self.assertTrue(torch.isfinite(generated_video).all()) - def test_inference_autoregressive_multi_chunk(self): - device = "cpu" + assert generated_video.shape == self.output_shape + assert torch.isfinite(generated_video).all() - components = self.get_dummy_components() - pipe = self.pipeline_class(**components) - pipe.to(device) - pipe.set_progress_bar_config(disable=None) + def test_inference_autoregressive_multi_chunk(self): + pipe = self.get_pipeline() - inputs = self.get_dummy_inputs(device) + inputs = self.get_dummy_inputs() inputs["num_frames"] = 5 inputs["num_frames_per_chunk"] = 3 inputs["num_ar_conditional_frames"] = 1 video = pipe(**inputs).frames generated_video = video[0] - self.assertEqual(generated_video.shape, (5, 3, 32, 32)) - self.assertTrue(torch.isfinite(generated_video).all()) - def test_inference_autoregressive_multi_chunk_no_condition_frames(self): - device = "cpu" + assert generated_video.shape == (5, *self.output_shape[1:]) + assert torch.isfinite(generated_video).all() - components = self.get_dummy_components() - pipe = self.pipeline_class(**components) - pipe.to(device) - pipe.set_progress_bar_config(disable=None) + def test_inference_autoregressive_multi_chunk_no_condition_frames(self): + pipe = self.get_pipeline() - inputs = self.get_dummy_inputs(device) + inputs = self.get_dummy_inputs() inputs["num_frames"] = 5 inputs["num_frames_per_chunk"] = 3 inputs["num_ar_conditional_frames"] = 0 video = pipe(**inputs).frames generated_video = video[0] - self.assertEqual(generated_video.shape, (5, 3, 32, 32)) - self.assertTrue(torch.isfinite(generated_video).all()) - def test_num_frames_per_chunk_above_rope_raises(self): - device = "cpu" + assert generated_video.shape == (5, *self.output_shape[1:]) + assert torch.isfinite(generated_video).all() - components = self.get_dummy_components() - pipe = self.pipeline_class(**components) - pipe.to(device) - pipe.set_progress_bar_config(disable=None) + def test_num_frames_per_chunk_above_rope_raises(self): + pipe = self.get_pipeline() - inputs = self.get_dummy_inputs(device) + inputs = self.get_dummy_inputs() inputs["num_frames_per_chunk"] = 17 - with self.assertRaisesRegex(ValueError, "too large for RoPE setting"): + with pytest.raises(ValueError, match="too large for RoPE setting"): pipe(**inputs) def test_inference_with_controls(self): """Test inference with control inputs (ControlNet).""" - device = "cpu" + 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["controls"] = [torch.randn(3, 32, 32) for _ in range(5)] # list of 5 frames (C, H, W) inputs["controls_conditioning_scale"] = 1.0 inputs["num_frames"] = None video = pipe(**inputs).frames generated_video = video[0] - self.assertEqual(generated_video.shape, (5, 3, 32, 32)) - self.assertTrue(torch.isfinite(generated_video).all()) - - 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): - for tensor_name in callback_kwargs.keys(): - 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 - for tensor_name in callback_kwargs.keys(): - assert tensor_name in pipe._callback_tensor_inputs - return callback_kwargs - - inputs = self.get_dummy_inputs(torch_device) - - inputs["callback_on_step_end"] = callback_inputs_subset - inputs["callback_on_step_end_tensor_inputs"] = ["latents"] - _ = pipe(**inputs)[0] - - inputs["callback_on_step_end"] = callback_inputs_all - inputs["callback_on_step_end_tensor_inputs"] = pipe._callback_tensor_inputs - _ = 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=2, expected_max_diff=1e-2) - - def test_attention_slicing_forward_pass( - self, test_max_difference=True, test_mean_pixel_difference=True, expected_max_diff=1e-3 - ): - if not getattr(self, "test_attention_slicing", True): - 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_serialization_with_variants(self): - components = self.get_dummy_components() - pipe = self.pipeline_class(**components) - model_components = [ - component_name - for component_name, component in pipe.components.items() - if isinstance(component, torch.nn.Module) - ] - # Remove components that aren't saved as standard diffusers models - if "safety_checker" in model_components: - model_components.remove("safety_checker") - variant = "fp16" - - with tempfile.TemporaryDirectory() as tmpdir: - pipe.save_pretrained(tmpdir, variant=variant, safe_serialization=False) - - with open(f"{tmpdir}/model_index.json", "r") as f: - config = json.load(f) - - for subfolder in os.listdir(tmpdir): - if not os.path.isfile(subfolder) and subfolder in model_components: - folder_path = os.path.join(tmpdir, subfolder) - is_folder = os.path.isdir(folder_path) and subfolder in config - assert is_folder and any(p.split(".")[1].startswith(variant) for p in os.listdir(folder_path)) - - def test_dtype_dict(self): - components = self.get_dummy_components() - if not components: - self.skipTest("No dummy components defined.") - - pipe = self.pipeline_class(**components) - - specified_key = next(iter(components.keys())) - - with tempfile.TemporaryDirectory(ignore_cleanup_errors=True) as tmpdirname: - pipe.save_pretrained(tmpdirname, safe_serialization=False) - dtype_dict = {specified_key: torch.bfloat16, "default": torch.float16} - loaded_pipe = self.pipeline_class.from_pretrained( - tmpdirname, safety_checker=DummyCosmosSafetyChecker(), dtype=dtype_dict - ) - - for name, component in loaded_pipe.components.items(): - # Skip components that are not loaded from disk or have special handling - if name == "safety_checker": - continue - if isinstance(component, torch.nn.Module) and hasattr(component, "dtype"): - expected_dtype = dtype_dict.get(name, dtype_dict.get("default", torch.float32)) - self.assertEqual( - component.dtype, - expected_dtype, - f"Component '{name}' has dtype {component.dtype} but expected {expected_dtype}", - ) - - def test_dtype_alias(self): - # `torch_dtype` is deprecated in favor of `dtype` in `from_pretrained`. - components = self.get_dummy_components() - pipe = self.pipeline_class(**components) - - with tempfile.TemporaryDirectory(ignore_cleanup_errors=True) as tmpdirname: - pipe.save_pretrained(tmpdirname, safe_serialization=False) - loaded_pipe = self.pipeline_class.from_pretrained( - tmpdirname, safety_checker=DummyCosmosSafetyChecker(), dtype=torch.float16 - ) - - for name, component in loaded_pipe.components.items(): - # Skip components that are not loaded from disk or have special handling - if name == "safety_checker": - continue - if isinstance(component, torch.nn.Module) and hasattr(component, "dtype"): - self.assertEqual( - component.dtype, - torch.float16, - f"Component '{name}' has dtype {component.dtype} but expected {torch.float16}", - ) - - def test_save_load_optional_components(self, expected_max_difference=1e-4): - self.pipeline_class._optional_components.remove("safety_checker") - super().test_save_load_optional_components(expected_max_difference=expected_max_difference) - self.pipeline_class._optional_components.append("safety_checker") - - @unittest.skip( - "The pipeline should not be runnable without a safety checker. The test creates a pipeline without passing in " - "a safety checker, which makes the pipeline default to the actual Cosmos Guardrail. The Cosmos Guardrail is " - "too large and slow to run on CI." - ) - def test_encode_prompt_works_in_isolation(self): - pass + assert generated_video.shape == (5, *self.output_shape[1:]) + assert torch.isfinite(generated_video).all() + + def test_inference_batch_single_identical(self, batch_size=2, expected_max_diff=1e-2): + super().test_inference_batch_single_identical(batch_size=batch_size, expected_max_diff=expected_max_diff) + + +class TestCosmos2_5_TransferPipelineMemory(Cosmos2_5_TransferPipelineTesterConfig, MemoryTesterMixin): + """Memory optimization tests (CPU offload, group offload, layerwise casting) for the Cosmos2.5 transfer pipeline.""" diff --git a/tests/pipelines/cosmos/test_cosmos2_text2image.py b/tests/pipelines/cosmos/test_cosmos2_text2image.py index 232d5faf1d02..0abf9cb8677d 100644 --- a/tests/pipelines/cosmos/test_cosmos2_text2image.py +++ b/tests/pipelines/cosmos/test_cosmos2_text2image.py @@ -12,13 +12,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -import inspect -import json -import os -import tempfile -import unittest - -import numpy as np import torch from transformers import AutoConfig, AutoTokenizer, T5EncoderModel @@ -29,10 +22,10 @@ FlowMatchEulerDiscreteScheduler, ) -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 +from ..testing_utils import BasePipelineTesterConfig, MemoryTesterMixin, PipelineTesterMixin from .cosmos_guardrail import DummyCosmosSafetyChecker +from .testing_utils import CosmosSafetyCheckerTesterMixin enable_full_determinism() @@ -45,25 +38,13 @@ def from_pretrained(*args, **kwargs): return Cosmos2TextToImagePipeline.from_pretrained(*args, **kwargs) -class Cosmos2TextToImagePipelineFastTests(PipelineTesterMixin, unittest.TestCase): +class Cosmos2TextToImagePipelineTesterConfig(BasePipelineTesterConfig): pipeline_class = Cosmos2TextToImagePipelineWrapper - 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( - [ - "num_inference_steps", - "generator", - "latents", - "return_dict", - "callback_on_step_end", - "callback_on_step_end_tensor_inputs", - ] + required_input_params_in_call_signature = frozenset( + ["prompt", "height", "width", "guidance_scale", "negative_prompt", "prompt_embeds", "negative_prompt_embeds"] ) - test_xformers_attention = False - test_layerwise_casting = True - test_group_offloading = True + batch_input_params = frozenset(["prompt", "negative_prompt"]) + output_shape = (3, 32, 32) def get_dummy_components(self): torch.manual_seed(0) @@ -95,10 +76,12 @@ def get_dummy_components(self): torch.manual_seed(0) scheduler = FlowMatchEulerDiscreteScheduler(use_karras_sigmas=True) config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5") - text_encoder = T5EncoderModel(config) + # `eval()` because a directly constructed model stays in training mode, which leaves T5's + # dropout active and makes the pipeline outputs non-deterministic across calls. + text_encoder = T5EncoderModel(config).eval() tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5") - components = { + return { "transformer": transformer, "vae": vae, "scheduler": scheduler, @@ -107,40 +90,33 @@ def get_dummy_components(self): # We cannot run the Cosmos Guardrail for fast tests due to the large model size "safety_checker": DummyCosmosSafetyChecker(), } - 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": "dance monkey", "negative_prompt": "bad quality", - "generator": generator, + "generator": self.get_generator(0), "num_inference_steps": 2, "guidance_scale": 3.0, "height": 32, "width": 32, "max_sequence_length": 16, + # 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 TestCosmos2TextToImagePipeline( + Cosmos2TextToImagePipelineTesterConfig, CosmosSafetyCheckerTesterMixin, PipelineTesterMixin +): 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) + # Run on CPU: the expected slice below is CPU-specific. + pipe = self.get_pipeline() - inputs = self.get_dummy_inputs(device) - image = pipe(**inputs).images + image = pipe(**self.get_dummy_inputs()).images generated_image = image[0] - self.assertEqual(generated_image.shape, (3, 32, 32)) + assert generated_image.shape == self.output_shape # fmt: off expected_slice = torch.tensor([0.451, 0.451, 0.4471, 0.451, 0.451, 0.451, 0.451, 0.451, 0.4784, 0.4784, 0.4784, 0.4784, 0.4784, 0.4902, 0.4588, 0.5333]) @@ -148,115 +124,16 @@ def test_inference(self): generated_slice = generated_image.flatten() generated_slice = torch.cat([generated_slice[:8], generated_slice[-8:]]) - self.assertTrue(torch.allclose(generated_slice, expected_slice, atol=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", - ) - - 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-2) - - 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 + assert_tensors_close(generated_slice, expected_slice, atol=1e-3) - 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-2): + 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() # 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] @@ -267,96 +144,14 @@ def test_vae_tiling(self, expected_diff_max: float = 0.2): 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).max() < expected_diff_max, ( + "VAE tiling should not affect the inference results" ) - def test_save_load_optional_components(self, expected_max_difference=1e-4): - self.pipeline_class._optional_components.remove("safety_checker") - super().test_save_load_optional_components(expected_max_difference=expected_max_difference) - self.pipeline_class._optional_components.append("safety_checker") - - def test_serialization_with_variants(self): - components = self.get_dummy_components() - pipe = self.pipeline_class(**components) - model_components = [ - component_name - for component_name, component in pipe.components.items() - if isinstance(component, torch.nn.Module) - ] - model_components.remove("safety_checker") - variant = "fp16" - - with tempfile.TemporaryDirectory() as tmpdir: - pipe.save_pretrained(tmpdir, variant=variant, safe_serialization=False) - - with open(f"{tmpdir}/model_index.json", "r") as f: - config = json.load(f) - - for subfolder in os.listdir(tmpdir): - if not os.path.isfile(subfolder) and subfolder in model_components: - folder_path = os.path.join(tmpdir, subfolder) - is_folder = os.path.isdir(folder_path) and subfolder in config - assert is_folder and any(p.split(".")[1].startswith(variant) for p in os.listdir(folder_path)) - def test_dtype_dict(self): - components = self.get_dummy_components() - if not components: - self.skipTest("No dummy components defined.") - - pipe = self.pipeline_class(**components) - - specified_key = next(iter(components.keys())) - - with tempfile.TemporaryDirectory(ignore_cleanup_errors=True) as tmpdirname: - pipe.save_pretrained(tmpdirname, safe_serialization=False) - dtype_dict = {specified_key: torch.bfloat16, "default": torch.float16} - loaded_pipe = self.pipeline_class.from_pretrained( - tmpdirname, safety_checker=DummyCosmosSafetyChecker(), dtype=dtype_dict - ) - - for name, component in loaded_pipe.components.items(): - if name == "safety_checker": - continue - if isinstance(component, torch.nn.Module) and hasattr(component, "dtype"): - expected_dtype = dtype_dict.get(name, dtype_dict.get("default", torch.float32)) - self.assertEqual( - component.dtype, - expected_dtype, - f"Component '{name}' has dtype {component.dtype} but expected {expected_dtype}", - ) - - def test_dtype_alias(self): - # `torch_dtype` is deprecated in favor of `dtype` in `from_pretrained`. - components = self.get_dummy_components() - pipe = self.pipeline_class(**components) - - with tempfile.TemporaryDirectory(ignore_cleanup_errors=True) as tmpdirname: - pipe.save_pretrained(tmpdirname, safe_serialization=False) - loaded_pipe = self.pipeline_class.from_pretrained( - tmpdirname, safety_checker=DummyCosmosSafetyChecker(), dtype=torch.float16 - ) - - for name, component in loaded_pipe.components.items(): - if name == "safety_checker": - continue - if isinstance(component, torch.nn.Module) and hasattr(component, "dtype"): - self.assertEqual( - component.dtype, - torch.float16, - f"Component '{name}' has dtype {component.dtype} but expected {torch.float16}", - ) - - @unittest.skip( - "The pipeline should not be runnable without a safety checker. The test creates a pipeline without passing in " - "a safety checker, which makes the pipeline default to the actual Cosmos Guardrail. The Cosmos Guardrail is " - "too large and slow to run on CI." - ) - def test_encode_prompt_works_in_isolation(self): - pass +class TestCosmos2TextToImagePipelineMemory(Cosmos2TextToImagePipelineTesterConfig, MemoryTesterMixin): + """Memory optimization tests (CPU offload, group offload, layerwise casting) for the Cosmos2 text-to-image pipeline.""" diff --git a/tests/pipelines/cosmos/test_cosmos2_video2world.py b/tests/pipelines/cosmos/test_cosmos2_video2world.py index a951877db2b1..09a4c407272b 100644 --- a/tests/pipelines/cosmos/test_cosmos2_video2world.py +++ b/tests/pipelines/cosmos/test_cosmos2_video2world.py @@ -12,13 +12,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -import inspect -import json -import os -import tempfile -import unittest - -import numpy as np import PIL.Image import torch from transformers import AutoConfig, AutoTokenizer, T5EncoderModel @@ -30,10 +23,10 @@ FlowMatchEulerDiscreteScheduler, ) -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 +from ..testing_utils import BasePipelineTesterConfig, MemoryTesterMixin, PipelineTesterMixin from .cosmos_guardrail import DummyCosmosSafetyChecker +from .testing_utils import CosmosSafetyCheckerTesterMixin enable_full_determinism() @@ -46,25 +39,17 @@ def from_pretrained(*args, **kwargs): return Cosmos2VideoToWorldPipeline.from_pretrained(*args, **kwargs) -class Cosmos2VideoToWorldPipelineFastTests(PipelineTesterMixin, unittest.TestCase): +class Cosmos2VideoToWorldPipelineTesterConfig(BasePipelineTesterConfig): pipeline_class = Cosmos2VideoToWorldPipelineWrapper - params = TEXT_TO_IMAGE_PARAMS - {"cross_attention_kwargs"} - batch_params = TEXT_TO_IMAGE_BATCH_PARAMS.union({"image", "video"}) - image_params = TEXT_TO_IMAGE_IMAGE_PARAMS - image_latents_params = TEXT_TO_IMAGE_IMAGE_PARAMS - required_optional_params = frozenset( - [ - "num_inference_steps", - "generator", - "latents", - "return_dict", - "callback_on_step_end", - "callback_on_step_end_tensor_inputs", - ] + 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", "image", "video"]) + output_shape = (9, 3, 32, 32) + # Cosmos2 video-to-world is a video pipeline: it exposes `num_videos_per_prompt`, not `num_images_per_prompt`. + optional_input_params = frozenset( + ["num_inference_steps", "num_videos_per_prompt", "generator", "latents", "output_type", "return_dict"] ) - test_xformers_attention = False - test_layerwise_casting = True - test_group_offloading = True def get_dummy_components(self): torch.manual_seed(0) @@ -96,10 +81,12 @@ def get_dummy_components(self): torch.manual_seed(0) scheduler = FlowMatchEulerDiscreteScheduler(use_karras_sigmas=True) config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5") - text_encoder = T5EncoderModel(config) + # `eval()` because a directly constructed model stays in training mode, which leaves T5's + # dropout active and makes the pipeline outputs non-deterministic across calls. + text_encoder = T5EncoderModel(config).eval() tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5") - components = { + return { "transformer": transformer, "vae": vae, "scheduler": scheduler, @@ -108,46 +95,37 @@ def get_dummy_components(self): # We cannot run the Cosmos Guardrail for fast tests due to the large model size "safety_checker": DummyCosmosSafetyChecker(), } - 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): image_height = 32 image_width = 32 image = PIL.Image.new("RGB", (image_width, image_height)) - - inputs = { + return { "image": image, "prompt": "dance monkey", "negative_prompt": "bad quality", - "generator": generator, + "generator": self.get_generator(0), "num_inference_steps": 2, "guidance_scale": 3.0, "height": image_height, "width": image_width, "num_frames": 9, "max_sequence_length": 16, + # Request torch outputs so tests compare torch tensors directly (see `BasePipelineTesterConfig`). "output_type": "pt", } - return inputs +class TestCosmos2VideoToWorldPipeline( + Cosmos2VideoToWorldPipelineTesterConfig, CosmosSafetyCheckerTesterMixin, PipelineTesterMixin +): 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) + # Run on CPU: the expected slice below is CPU-specific. + pipe = self.get_pipeline() - inputs = self.get_dummy_inputs(device) - video = pipe(**inputs).frames + video = pipe(**self.get_dummy_inputs()).frames generated_video = video[0] - self.assertEqual(generated_video.shape, (9, 3, 32, 32)) + assert generated_video.shape == self.output_shape # fmt: off expected_slice = torch.tensor([0.451, 0.451, 0.4471, 0.451, 0.451, 0.451, 0.451, 0.451, 0.5098, 0.5137, 0.5176, 0.5098, 0.5255, 0.5412, 0.5098, 0.5059]) @@ -155,122 +133,16 @@ def test_inference(self): generated_slice = generated_video.flatten() generated_slice = torch.cat([generated_slice[:8], generated_slice[-8:]]) - self.assertTrue(torch.allclose(generated_slice, expected_slice, atol=1e-3)) - - def test_components_function(self): - init_components = self.get_dummy_components() - init_components = {k: v for k, v in init_components.items() if not isinstance(v, (str, int, float))} - pipe = self.pipeline_class(**init_components) - self.assertTrue(hasattr(pipe, "components")) - self.assertTrue(set(pipe.components.keys()) == set(init_components.keys())) - - 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 + assert_tensors_close(generated_slice, expected_slice, atol=1e-3) - def test_inference_batch_single_identical(self): - self._test_inference_batch_single_identical(batch_size=3, expected_max_diff=1e-2) - - 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-2): + 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() # 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] @@ -281,96 +153,14 @@ def test_vae_tiling(self, expected_diff_max: float = 0.2): 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).max() < expected_diff_max, ( + "VAE tiling should not affect the inference results" ) - def test_save_load_optional_components(self, expected_max_difference=1e-4): - self.pipeline_class._optional_components.remove("safety_checker") - super().test_save_load_optional_components(expected_max_difference=expected_max_difference) - self.pipeline_class._optional_components.append("safety_checker") - - def test_serialization_with_variants(self): - components = self.get_dummy_components() - pipe = self.pipeline_class(**components) - model_components = [ - component_name - for component_name, component in pipe.components.items() - if isinstance(component, torch.nn.Module) - ] - model_components.remove("safety_checker") - variant = "fp16" - - with tempfile.TemporaryDirectory() as tmpdir: - pipe.save_pretrained(tmpdir, variant=variant, safe_serialization=False) - - with open(f"{tmpdir}/model_index.json", "r") as f: - config = json.load(f) - for subfolder in os.listdir(tmpdir): - if not os.path.isfile(subfolder) and subfolder in model_components: - folder_path = os.path.join(tmpdir, subfolder) - is_folder = os.path.isdir(folder_path) and subfolder in config - assert is_folder and any(p.split(".")[1].startswith(variant) for p in os.listdir(folder_path)) - - def test_dtype_dict(self): - components = self.get_dummy_components() - if not components: - self.skipTest("No dummy components defined.") - - pipe = self.pipeline_class(**components) - - specified_key = next(iter(components.keys())) - - with tempfile.TemporaryDirectory(ignore_cleanup_errors=True) as tmpdirname: - pipe.save_pretrained(tmpdirname, safe_serialization=False) - dtype_dict = {specified_key: torch.bfloat16, "default": torch.float16} - loaded_pipe = self.pipeline_class.from_pretrained( - tmpdirname, safety_checker=DummyCosmosSafetyChecker(), dtype=dtype_dict - ) - - for name, component in loaded_pipe.components.items(): - if name == "safety_checker": - continue - if isinstance(component, torch.nn.Module) and hasattr(component, "dtype"): - expected_dtype = dtype_dict.get(name, dtype_dict.get("default", torch.float32)) - self.assertEqual( - component.dtype, - expected_dtype, - f"Component '{name}' has dtype {component.dtype} but expected {expected_dtype}", - ) - - def test_dtype_alias(self): - # `torch_dtype` is deprecated in favor of `dtype` in `from_pretrained`. - components = self.get_dummy_components() - pipe = self.pipeline_class(**components) - - with tempfile.TemporaryDirectory(ignore_cleanup_errors=True) as tmpdirname: - pipe.save_pretrained(tmpdirname, safe_serialization=False) - loaded_pipe = self.pipeline_class.from_pretrained( - tmpdirname, safety_checker=DummyCosmosSafetyChecker(), dtype=torch.float16 - ) - - for name, component in loaded_pipe.components.items(): - if name == "safety_checker": - continue - if isinstance(component, torch.nn.Module) and hasattr(component, "dtype"): - self.assertEqual( - component.dtype, - torch.float16, - f"Component '{name}' has dtype {component.dtype} but expected {torch.float16}", - ) - - @unittest.skip( - "The pipeline should not be runnable without a safety checker. The test creates a pipeline without passing in " - "a safety checker, which makes the pipeline default to the actual Cosmos Guardrail. The Cosmos Guardrail is " - "too large and slow to run on CI." - ) - def test_encode_prompt_works_in_isolation(self): - pass +class TestCosmos2VideoToWorldPipelineMemory(Cosmos2VideoToWorldPipelineTesterConfig, MemoryTesterMixin): + """Memory optimization tests (CPU offload, group offload, layerwise casting) for the Cosmos2 video-to-world pipeline.""" diff --git a/tests/pipelines/cosmos/test_cosmos3.py b/tests/pipelines/cosmos/test_cosmos3.py index dfc1dd0a1ea1..86bd371c83c5 100644 --- a/tests/pipelines/cosmos/test_cosmos3.py +++ b/tests/pipelines/cosmos/test_cosmos3.py @@ -12,10 +12,10 @@ # See the License for the specific language governing permissions and # limitations under the License. -import unittest from unittest import mock import numpy as np +import pytest import torch from PIL import Image from transformers import AutoTokenizer @@ -24,31 +24,22 @@ from diffusers.pipelines.cosmos.pipeline_cosmos3_omni import _preprocess_conditioning_image from ...testing_utils import enable_full_determinism, 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 Cosmos3OmniPipelineFastTests(PipelineTesterMixin, unittest.TestCase): +class Cosmos3OmniPipelineTesterConfig(BasePipelineTesterConfig): pipeline_class = Cosmos3OmniPipeline - params = TEXT_TO_IMAGE_PARAMS - {"cross_attention_kwargs", "negative_prompt_embeds", "prompt_embeds"} - batch_params = TEXT_TO_IMAGE_BATCH_PARAMS - required_optional_params = frozenset( - [ - "num_inference_steps", - "generator", - "latents", - "output_type", - "return_dict", - "callback_on_step_end", - "callback_on_step_end_tensor_inputs", - ] + required_input_params_in_call_signature = frozenset( + ["prompt", "height", "width", "guidance_scale", "negative_prompt"] ) - test_xformers_attention = False - test_layerwise_casting = False - test_group_offloading = True + batch_input_params = frozenset(["prompt", "negative_prompt"]) + output_shape = (3, 16, 16) + # Cosmos3 Omni generates one video per call, so it exposes neither `num_images_per_prompt` (the base default) + # nor the `num_videos_per_prompt` the other Cosmos pipelines take. + optional_input_params = frozenset(["num_inference_steps", "generator", "latents", "output_type", "return_dict"]) def get_dummy_components(self): torch.manual_seed(0) @@ -93,8 +84,7 @@ def get_dummy_components(self): "enable_safety_checker": False, } - def get_dummy_inputs(self, device, seed=0): - generator = torch.Generator(device="cpu").manual_seed(seed) + def get_dummy_inputs(self): return { "prompt": "a dog", "negative_prompt": "bad quality", @@ -103,25 +93,27 @@ def get_dummy_inputs(self, device, seed=0): "num_frames": 1, "num_inference_steps": 2, "guidance_scale": 1.0, - "generator": generator, - "output_type": "np", + "generator": self.get_generator(0), + # Request torch outputs so tests compare torch tensors directly (see `BasePipelineTesterConfig`). + "output_type": "pt", "use_system_prompt": False, "add_resolution_template": False, "add_duration_template": False, } + +class TestCosmos3OmniPipeline(Cosmos3OmniPipelineTesterConfig, PipelineTesterMixin): def test_inference(self): - pipeline = self.pipeline_class(**self.get_dummy_components()).to(torch_device) - pipeline.set_progress_bar_config(disable=None) + pipe = self.get_pipeline().to(torch_device) - video = pipeline(**self.get_dummy_inputs(torch_device)).video + video = pipe(**self.get_dummy_inputs()).video - self.assertEqual(video.shape, (1, 16, 16, 3)) + assert video[0].shape == self.output_shape def test_cosmos3_tokenize_prompt_uses_checkpoint_system_prompt_default(self): components = self.get_dummy_components() components["default_use_system_prompt"] = False - pipeline = self.pipeline_class(**components) + pipeline = self.get_pipeline(**components) with mock.patch.object( pipeline.text_tokenizer, @@ -148,27 +140,30 @@ def test_i2v_image_preprocessing_preserves_aspect_ratio(self): torch.testing.assert_close(actual, expected) def test_i2v_pipeline_uses_native_preprocessing(self): - pipeline = self.pipeline_class(**self.get_dummy_components()).to(torch_device) - pipeline.set_progress_bar_config(disable=None) + pipeline = self.get_pipeline().to(torch_device) image = np.zeros((16, 32, 3), dtype=np.uint8) image[:, :8] = [255, 0, 0] image[:, 8:24] = [0, 255, 0] image[:, 24:] = [0, 0, 255] center_crop = Image.fromarray(image[:, 8:24]) - inputs = self.get_dummy_inputs(torch_device) + inputs = self.get_dummy_inputs() inputs.update(image=Image.fromarray(image), num_frames=5, output_type="latent") wide_output = pipeline(**inputs).video - inputs.update(image=center_crop, generator=torch.Generator(device="cpu").manual_seed(0)) + inputs.update(image=center_crop, generator=self.get_generator(0)) crop_output = pipeline(**inputs).video torch.testing.assert_close(wide_output, crop_output) - @unittest.skip("Cosmos3 currently supports one prompt per pipeline call.") + @pytest.mark.skip("Cosmos3 currently supports one prompt per pipeline call.") def test_inference_batch_consistent(self): pass - @unittest.skip("Cosmos3 currently supports one prompt per pipeline call.") + @pytest.mark.skip("Cosmos3 currently supports one prompt per pipeline call.") def test_inference_batch_single_identical(self): pass + + +class TestCosmos3OmniPipelineMemory(Cosmos3OmniPipelineTesterConfig, MemoryTesterMixin): + """Memory optimization tests (CPU offload, group offload, layerwise casting) for the Cosmos3 Omni pipeline.""" diff --git a/tests/pipelines/cosmos/test_cosmos_video2world.py b/tests/pipelines/cosmos/test_cosmos_video2world.py index 7e4928ccf0df..600ddd1ed511 100644 --- a/tests/pipelines/cosmos/test_cosmos_video2world.py +++ b/tests/pipelines/cosmos/test_cosmos_video2world.py @@ -12,23 +12,16 @@ # See the License for the specific language governing permissions and # limitations under the License. -import inspect -import json -import os -import tempfile -import unittest - -import numpy as np import PIL.Image import torch from transformers import AutoConfig, AutoTokenizer, T5EncoderModel from diffusers import AutoencoderKLCosmos, CosmosTransformer3DModel, CosmosVideoToWorldPipeline, EDMEulerScheduler -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 +from ..testing_utils import BasePipelineTesterConfig, MemoryTesterMixin, PipelineTesterMixin from .cosmos_guardrail import DummyCosmosSafetyChecker +from .testing_utils import CosmosSafetyCheckerTesterMixin enable_full_determinism() @@ -41,25 +34,17 @@ def from_pretrained(*args, **kwargs): return CosmosVideoToWorldPipeline.from_pretrained(*args, **kwargs) -class CosmosVideoToWorldPipelineFastTests(PipelineTesterMixin, unittest.TestCase): +class CosmosVideoToWorldPipelineTesterConfig(BasePipelineTesterConfig): pipeline_class = CosmosVideoToWorldPipelineWrapper - params = TEXT_TO_IMAGE_PARAMS - {"cross_attention_kwargs"} - batch_params = TEXT_TO_IMAGE_BATCH_PARAMS.union({"image", "video"}) - image_params = TEXT_TO_IMAGE_IMAGE_PARAMS - image_latents_params = TEXT_TO_IMAGE_IMAGE_PARAMS - required_optional_params = frozenset( - [ - "num_inference_steps", - "generator", - "latents", - "return_dict", - "callback_on_step_end", - "callback_on_step_end_tensor_inputs", - ] + 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", "image", "video"]) + output_shape = (9, 3, 32, 32) + # Cosmos text-to-world is a video pipeline: it exposes `num_videos_per_prompt`, not `num_images_per_prompt`. + optional_input_params = frozenset( + ["num_inference_steps", "num_videos_per_prompt", "generator", "latents", "output_type", "return_dict"] ) - test_xformers_attention = False - test_layerwise_casting = True - test_group_offloading = True def get_dummy_components(self): torch.manual_seed(0) @@ -108,10 +93,12 @@ def get_dummy_components(self): final_sigmas_type="sigma_min", ) config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5") - text_encoder = T5EncoderModel(config) + # `eval()` because a directly constructed model stays in training mode, which leaves T5's + # dropout active and makes the pipeline outputs non-deterministic across calls. + text_encoder = T5EncoderModel(config).eval() tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5") - components = { + return { "transformer": transformer, "vae": vae, "scheduler": scheduler, @@ -120,46 +107,37 @@ def get_dummy_components(self): # We cannot run the Cosmos Guardrail for fast tests due to the large model size "safety_checker": DummyCosmosSafetyChecker(), } - 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): image_height = 32 image_width = 32 image = PIL.Image.new("RGB", (image_width, image_height)) - - inputs = { + return { "image": image, "prompt": "dance monkey", "negative_prompt": "bad quality", - "generator": generator, + "generator": self.get_generator(0), "num_inference_steps": 2, "guidance_scale": 3.0, "height": image_height, "width": image_width, "num_frames": 9, "max_sequence_length": 16, + # Request torch outputs so tests compare torch tensors directly (see `BasePipelineTesterConfig`). "output_type": "pt", } - return inputs +class TestCosmosVideoToWorldPipeline( + CosmosVideoToWorldPipelineTesterConfig, CosmosSafetyCheckerTesterMixin, PipelineTesterMixin +): 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) + # Run on CPU: the expected slice below is CPU-specific. + pipe = self.get_pipeline() - inputs = self.get_dummy_inputs(device) - video = pipe(**inputs).frames + video = pipe(**self.get_dummy_inputs()).frames generated_video = video[0] - self.assertEqual(generated_video.shape, (9, 3, 32, 32)) + assert generated_video.shape == self.output_shape # fmt: off expected_slice = torch.tensor([0.0, 0.8275, 0.7529, 0.7294, 0.0, 0.6, 1.0, 0.3804, 0.6667, 0.0863, 0.8784, 0.5922, 0.6627, 0.2784, 0.5725, 0.7765]) @@ -167,125 +145,16 @@ def test_inference(self): generated_slice = generated_video.flatten() generated_slice = torch.cat([generated_slice[:8], generated_slice[-8:]]) - self.assertTrue(torch.allclose(generated_slice, expected_slice, atol=1e-3)) - - def test_components_function(self): - init_components = self.get_dummy_components() - init_components = {k: v for k, v in init_components.items() if not isinstance(v, (str, int, float))} - pipe = self.pipeline_class(**init_components) - self.assertTrue(hasattr(pipe, "components")) - self.assertTrue(set(pipe.components.keys()) == set(init_components.keys())) - - 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 + assert_tensors_close(generated_slice, expected_slice, atol=1e-3) - def test_inference_batch_single_identical(self): - self._test_inference_batch_single_identical(batch_size=3, expected_max_diff=1e-2) - - 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() - for key in components: - if "text_encoder" in key and hasattr(components[key], "eval"): - components[key].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) - - 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-2): + 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() # 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] @@ -296,96 +165,14 @@ def test_vae_tiling(self, expected_diff_max: float = 0.2): 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).max() < expected_diff_max, ( + "VAE tiling should not affect the inference results" ) - def test_save_load_optional_components(self, expected_max_difference=1e-4): - self.pipeline_class._optional_components.remove("safety_checker") - super().test_save_load_optional_components(expected_max_difference=expected_max_difference) - self.pipeline_class._optional_components.append("safety_checker") - - def test_serialization_with_variants(self): - components = self.get_dummy_components() - pipe = self.pipeline_class(**components) - model_components = [ - component_name - for component_name, component in pipe.components.items() - if isinstance(component, torch.nn.Module) - ] - model_components.remove("safety_checker") - variant = "fp16" - - with tempfile.TemporaryDirectory() as tmpdir: - pipe.save_pretrained(tmpdir, variant=variant, safe_serialization=False) - - with open(f"{tmpdir}/model_index.json", "r") as f: - config = json.load(f) - for subfolder in os.listdir(tmpdir): - if not os.path.isfile(subfolder) and subfolder in model_components: - folder_path = os.path.join(tmpdir, subfolder) - is_folder = os.path.isdir(folder_path) and subfolder in config - assert is_folder and any(p.split(".")[1].startswith(variant) for p in os.listdir(folder_path)) - - def test_dtype_dict(self): - components = self.get_dummy_components() - if not components: - self.skipTest("No dummy components defined.") - - pipe = self.pipeline_class(**components) - - specified_key = next(iter(components.keys())) - - with tempfile.TemporaryDirectory(ignore_cleanup_errors=True) as tmpdirname: - pipe.save_pretrained(tmpdirname, safe_serialization=False) - dtype_dict = {specified_key: torch.bfloat16, "default": torch.float16} - loaded_pipe = self.pipeline_class.from_pretrained( - tmpdirname, safety_checker=DummyCosmosSafetyChecker(), dtype=dtype_dict - ) - - for name, component in loaded_pipe.components.items(): - if name == "safety_checker": - continue - if isinstance(component, torch.nn.Module) and hasattr(component, "dtype"): - expected_dtype = dtype_dict.get(name, dtype_dict.get("default", torch.float32)) - self.assertEqual( - component.dtype, - expected_dtype, - f"Component '{name}' has dtype {component.dtype} but expected {expected_dtype}", - ) - - def test_dtype_alias(self): - # `torch_dtype` is deprecated in favor of `dtype` in `from_pretrained`. - components = self.get_dummy_components() - pipe = self.pipeline_class(**components) - - with tempfile.TemporaryDirectory(ignore_cleanup_errors=True) as tmpdirname: - pipe.save_pretrained(tmpdirname, safe_serialization=False) - loaded_pipe = self.pipeline_class.from_pretrained( - tmpdirname, safety_checker=DummyCosmosSafetyChecker(), dtype=torch.float16 - ) - - for name, component in loaded_pipe.components.items(): - if name == "safety_checker": - continue - if isinstance(component, torch.nn.Module) and hasattr(component, "dtype"): - self.assertEqual( - component.dtype, - torch.float16, - f"Component '{name}' has dtype {component.dtype} but expected {torch.float16}", - ) - - @unittest.skip( - "The pipeline should not be runnable without a safety checker. The test creates a pipeline without passing in " - "a safety checker, which makes the pipeline default to the actual Cosmos Guardrail. The Cosmos Guardrail is " - "too large and slow to run on CI." - ) - def test_encode_prompt_works_in_isolation(self): - pass +class TestCosmosVideoToWorldPipelineMemory(CosmosVideoToWorldPipelineTesterConfig, MemoryTesterMixin): + """Memory optimization tests (CPU offload, group offload, layerwise casting) for the Cosmos video-to-world pipeline.""" diff --git a/tests/pipelines/cosmos/testing_utils.py b/tests/pipelines/cosmos/testing_utils.py new file mode 100644 index 000000000000..984019db610f --- /dev/null +++ b/tests/pipelines/cosmos/testing_utils.py @@ -0,0 +1,95 @@ +# 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. + +import json +import os + +import pytest +import torch + +from .cosmos_guardrail import DummyCosmosSafetyChecker + + +class CosmosSafetyCheckerTesterMixin: + """Overrides of the shared pipeline tests that the Cosmos `safety_checker` component gets in the way of. + + Every Cosmos pipeline takes a `safety_checker`, and the fast tests substitute `DummyCosmosSafetyChecker` for the + real Cosmos Guardrail, which is far too large to build on CI. The dummy is not serialized like the other model + components and a pipeline constructed without one falls back to the real guardrail, so the tests that save, + reload or enumerate components need the component handled explicitly. + + Compose it *before* `PipelineTesterMixin` so these overrides win, e.g. + `class TestCosmosXPipeline(CosmosXPipelineTesterConfig, CosmosSafetyCheckerTesterMixin, PipelineTesterMixin)`. + """ + + def test_save_load_optional_components(self, tmp_path, expected_max_difference=1e-4): + # `safety_checker` is listed as optional, but a pipeline built without one falls back to the real Cosmos + # Guardrail — so keep it out of the components the base test nulls. + self.pipeline_class._optional_components.remove("safety_checker") + try: + super().test_save_load_optional_components(tmp_path, expected_max_difference=expected_max_difference) + finally: + self.pipeline_class._optional_components.append("safety_checker") + + def test_serialization_with_variants(self, tmp_path): + # Same as the base test, except `safety_checker` is not serialized like the other model components. + pipe = self.get_pipeline() + model_components = [ + component_name + for component_name, component in pipe.components.items() + if isinstance(component, torch.nn.Module) + ] + model_components.remove("safety_checker") + variant = "fp16" + + pipe.save_pretrained(tmp_path, variant=variant, safe_serialization=False) + + with open(f"{tmp_path}/model_index.json", "r") as f: + config = json.load(f) + + for subfolder in os.listdir(tmp_path): + if not os.path.isfile(subfolder) and subfolder in model_components: + folder_path = os.path.join(tmp_path, subfolder) + is_folder = os.path.isdir(folder_path) and subfolder in config + assert is_folder and any(p.split(".")[1].startswith(variant) for p in os.listdir(folder_path)) + + def test_torch_dtype_dict(self, tmp_path): + # Same as the base test, except the safety checker has to be passed back in on load and is left out of the + # dtype check (the dummy tracks its dtype through a non-persistent buffer). + components = self.get_dummy_components() + pipe = self.get_pipeline(**components) + specified_key = next(iter(components.keys())) + + pipe.save_pretrained(str(tmp_path), safe_serialization=False) + torch_dtype_dict = {specified_key: torch.bfloat16, "default": torch.float16} + loaded_pipe = self.pipeline_class.from_pretrained( + str(tmp_path), safety_checker=DummyCosmosSafetyChecker(), dtype=torch_dtype_dict + ) + + for name, component in loaded_pipe.components.items(): + if name == "safety_checker": + continue + if isinstance(component, torch.nn.Module) and hasattr(component, "dtype"): + expected_dtype = torch_dtype_dict.get(name, torch_dtype_dict.get("default", torch.float32)) + assert component.dtype == expected_dtype, ( + f"Component '{name}' has dtype {component.dtype} but expected {expected_dtype}" + ) + + @pytest.mark.skip( + "The pipeline should not be runnable without a safety checker. The test creates a pipeline without passing in " + "a safety checker, which makes the pipeline default to the actual Cosmos Guardrail. The Cosmos Guardrail is " + "too large and slow to run on CI." + ) + def test_encode_prompt_works_in_isolation(self): + pass diff --git a/tests/pipelines/flux/testing_utils.py b/tests/pipelines/flux/testing_utils.py index 0168a23a78c2..f17c3916109f 100644 --- a/tests/pipelines/flux/testing_utils.py +++ b/tests/pipelines/flux/testing_utils.py @@ -26,7 +26,8 @@ @is_ip_adapter class FluxIPAdapterTesterMixin(BasePipelineOutputMixin): - """IP-Adapter tests shared by the Flux pipelines in this directory. + """IP-Adapter tests shared by the Flux pipelines in this directory and by the other pipelines built on the + Flux IP-Adapter API (Chroma, Flux ControlNet). Flux has its own IP-Adapter API (`FluxIPAdapterMixin`, image embeddings sized after the transformer's `pooled_projection_dim`), so it doesn't reuse the Stable Diffusion `IPAdapterTesterMixin`. Compose it with a diff --git a/tests/pipelines/testing_utils/common.py b/tests/pipelines/testing_utils/common.py index 23db523f1e1d..4ab410121d86 100644 --- a/tests/pipelines/testing_utils/common.py +++ b/tests/pipelines/testing_utils/common.py @@ -444,7 +444,7 @@ def test_save_load_float16(self, tmp_path, expected_max_diff=1e-2): output = pipe(**inputs)[0] pipe.save_pretrained(tmp_path) - pipe_loaded = self.pipeline_class.from_pretrained(tmp_path, torch_dtype=torch.float16) + pipe_loaded = self.pipeline_class.from_pretrained(tmp_path, dtype=torch.float16) pipe_loaded.to(torch_device) pipe_loaded.set_progress_bar_config(disable=None)