diff --git a/src/diffusers/pipelines/visualcloze/pipeline_visualcloze_combined.py b/src/diffusers/pipelines/visualcloze/pipeline_visualcloze_combined.py index ec7a85f466..2d881e22a7 100644 --- a/src/diffusers/pipelines/visualcloze/pipeline_visualcloze_combined.py +++ b/src/diffusers/pipelines/visualcloze/pipeline_visualcloze_combined.py @@ -146,6 +146,9 @@ class VisualClozePipeline( transformer=transformer, scheduler=scheduler, ) + # `resolution` is not a module, so it has to be registered explicitly to survive a + # `save_pretrained` / `from_pretrained` round-trip. + self.register_to_config(resolution=resolution) self.generation_pipe = VisualClozeGenerationPipeline( vae=vae, @@ -376,6 +379,10 @@ class VisualClozePipeline( output_type=output_type if upsampling_strength == 0 else "pil", ) if upsampling_strength == 0: + # Offload all models. The inner pipelines free their own (empty) hooks, so the ones installed on this + # pipeline by `enable_model_cpu_offload` have to be freed here. + self.maybe_free_model_hooks() + if not return_dict: return (generation_output,) @@ -434,6 +441,9 @@ class VisualClozePipeline( else: output = image + # Offload all models + self.maybe_free_model_hooks() + if not return_dict: return (output,) diff --git a/src/diffusers/pipelines/visualcloze/pipeline_visualcloze_generation.py b/src/diffusers/pipelines/visualcloze/pipeline_visualcloze_generation.py index d3a6fda140..ed2cd519df 100644 --- a/src/diffusers/pipelines/visualcloze/pipeline_visualcloze_generation.py +++ b/src/diffusers/pipelines/visualcloze/pipeline_visualcloze_generation.py @@ -175,6 +175,9 @@ class VisualClozeGenerationPipeline( transformer=transformer, scheduler=scheduler, ) + # `resolution` is not a module, so it has to be registered explicitly to survive a + # `save_pretrained` / `from_pretrained` round-trip. + self.register_to_config(resolution=resolution) self.resolution = resolution self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1) if getattr(self, "vae", None) else 8 # Flux latents are turned into 2x2 patches and packed. This means the latent width and height has to be divisible @@ -715,8 +718,9 @@ class VisualClozeGenerationPipeline( Pre-generated pooled text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not provided, pooled text embeddings will be generated from `prompt` input argument. output_type (`str`, *optional*, defaults to `"pil"`): - The output format of the generate image. Choose between - [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`. + The output format of the generate image. Choose between `"pil"` + ([PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image`), `"np"` (`np.array`) or `"pt"` + (`torch.Tensor`). return_dict (`bool`, *optional*, defaults to `True`): Whether or not to return a [`~pipelines.flux.FluxPipelineOutput`] instead of a plain tuple. joint_attention_kwargs (`dict`, *optional*): @@ -907,11 +911,17 @@ class VisualClozeGenerationPipeline( if cur_target_position[i]: if output_type == "pil": cropped.append(cur_image.crop((start, 0, start + size[1], size[0]))) + elif output_type == "pt": + # `"pt"` images are `(channels, height, width)`, unlike the `(height, width, channels)` + # layout of `"np"`, so the spatial crop applies to the last two axes. + cropped.append(cur_image[:, 0 : size[0], start : start + size[1]]) else: cropped.append(cur_image[0 : size[0], start : start + size[1]]) start += size[1] image.append(cropped) - if output_type != "pil": + if output_type == "pt": + image = torch.stack([arr for sub_image in image for arr in sub_image], dim=0) + elif output_type != "pil": image = np.concatenate([arr[None] for sub_image in image for arr in sub_image], axis=0) # Offload all models diff --git a/tests/pipelines/visualcloze/test_pipeline_visualcloze_combined.py b/tests/pipelines/visualcloze/test_pipeline_visualcloze_combined.py index 1355f6dc06..7ef4249525 100644 --- a/tests/pipelines/visualcloze/test_pipeline_visualcloze_combined.py +++ b/tests/pipelines/visualcloze/test_pipeline_visualcloze_combined.py @@ -1,32 +1,20 @@ import random -import tempfile -import unittest import numpy as np +import pytest import torch from PIL import Image from transformers import AutoConfig, AutoTokenizer, CLIPTextConfig, CLIPTextModel, CLIPTokenizer, T5EncoderModel -import diffusers from diffusers import AutoencoderKL, FlowMatchEulerDiscreteScheduler, FluxTransformer2DModel, VisualClozePipeline -from diffusers.utils import logging -from ...testing_utils import ( - CaptureLogger, - enable_full_determinism, - floats_tensor, - require_accelerator, - torch_device, -) -from ..test_pipelines_common import PipelineTesterMixin, to_np +from ...testing_utils import floats_tensor, torch_device +from ..testing_utils import BasePipelineTesterConfig, MemoryTesterMixin, PipelineTesterMixin -enable_full_determinism() - - -class VisualClozePipelineFastTests(unittest.TestCase, PipelineTesterMixin): +class VisualClozePipelineTesterConfig(BasePipelineTesterConfig): pipeline_class = VisualClozePipeline - params = frozenset( + required_input_params_in_call_signature = frozenset( [ "task_prompt", "content_prompt", @@ -38,10 +26,8 @@ class VisualClozePipelineFastTests(unittest.TestCase, PipelineTesterMixin): "upsampling_strength", ] ) - batch_params = frozenset(["task_prompt", "content_prompt", "image"]) - test_xformers_attention = False - test_layerwise_casting = True - test_group_offloading = True + batch_input_params = frozenset(["task_prompt", "content_prompt", "image"]) + output_shape = (3, 32, 32) def get_dummy_components(self): torch.manual_seed(0) @@ -109,7 +95,7 @@ class VisualClozePipelineFastTests(unittest.TestCase, PipelineTesterMixin): "resolution": 32, } - def get_dummy_inputs(self, device, seed=0): + def get_dummy_inputs(self, seed=0): # Create example images to simulate the input format required by VisualCloze context_image = [ Image.fromarray(floats_tensor((32, 32, 3), rng=random.Random(seed), scale=255).numpy().astype(np.uint8)) @@ -128,44 +114,41 @@ class VisualClozePipelineFastTests(unittest.TestCase, PipelineTesterMixin): query_image, # Query image ] - if str(device).startswith("mps"): - generator = torch.manual_seed(seed) - else: - generator = torch.Generator(device="cpu").manual_seed(seed) - inputs = { "task_prompt": "Each row outlines a logical process, starting from [IMAGE1] gray-based depth map with detailed object contours, to achieve [IMAGE2] an image with flawless clarity.", "content_prompt": "A beautiful landscape with mountains and a lake", "image": image, - "generator": generator, + "generator": self.get_generator(seed), "num_inference_steps": 2, "guidance_scale": 5.0, "upsampling_height": 32, "upsampling_width": 32, "max_sequence_length": 77, - "output_type": "np", "upsampling_strength": 0.4, + # Request torch outputs so tests compare torch tensors directly (see `BasePipelineTesterConfig`). + # Note `"pt"` images are `(batch, channels, height, width)`, unlike `"np"` (`(batch, h, w, c)`). + "output_type": "pt", } return inputs - def test_visualcloze_different_prompts(self): - pipe = self.pipeline_class(**self.get_dummy_components()).to(torch_device) - inputs = self.get_dummy_inputs(torch_device) - output_same_prompt = pipe(**inputs).images[0] +class TestVisualClozePipeline(VisualClozePipelineTesterConfig, PipelineTesterMixin): + def test_visualcloze_different_task_prompts(self, expected_min_diff=1e-1): + pipe = self.get_pipeline().to(torch_device) - inputs = self.get_dummy_inputs(torch_device) - inputs["task_prompt"] = "A different task to perform." - output_different_prompts = pipe(**inputs).images[0] + inputs = self.get_dummy_inputs() + output_original = pipe(**inputs).images[0] - max_diff = np.abs(output_same_prompt - output_different_prompts).max() + inputs["task_prompt"] = "A different task description for image generation" + output_different_task = pipe(**inputs).images[0] - # Outputs should be different - assert max_diff > 1e-6 + # Different task prompts should produce different outputs + max_diff = (output_original - output_different_task).abs().max() + assert max_diff > expected_min_diff, "Outputs should be different for different task prompts." def test_visualcloze_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: @@ -174,15 +157,14 @@ class VisualClozePipelineFastTests(unittest.TestCase, PipelineTesterMixin): inputs.update({"upsampling_height": height, "upsampling_width": width}) image = pipe(**inputs).images[0] - output_height, output_width, _ = image.shape - assert (output_height, output_width) == (expected_height, expected_width) - - def test_inference_batch_single_identical(self): - self._test_inference_batch_single_identical(expected_max_diff=1e-3) + _, 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)}" + ) def test_upsampling_strength(self, expected_min_diff=1e-1): - 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() # Test different upsampling strengths inputs["upsampling_strength"] = 0.2 @@ -192,159 +174,19 @@ class VisualClozePipelineFastTests(unittest.TestCase, PipelineTesterMixin): output_full_upsampling = pipe(**inputs).images[0] # Different upsampling strengths should produce different outputs - max_diff = np.abs(output_no_upsampling - output_full_upsampling).max() - assert max_diff > expected_min_diff + max_diff = (output_no_upsampling - output_full_upsampling).abs().max() + assert max_diff > expected_min_diff, "Outputs should be different for different upsampling strengths." - def test_different_task_prompts(self, expected_min_diff=1e-1): - pipe = self.pipeline_class(**self.get_dummy_components()).to(torch_device) - 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) - output_original = pipe(**inputs).images[0] - - inputs["task_prompt"] = "A different task description for image generation" - output_different_task = pipe(**inputs).images[0] - - # Different task prompts should produce different outputs - max_diff = np.abs(output_original - output_different_task).max() - assert max_diff > expected_min_diff - - @unittest.skip( - "Test not applicable because the pipeline being tested is a wrapper pipeline. CFG tests should be done on the inner pipelines." + @pytest.mark.skip( + "The pipeline being tested is a wrapper around two inner pipelines, so `guidance_scale` is not tracked on it. " + "CFG tests should be done on the inner pipelines." ) def test_callback_cfg(self): pass - def test_save_load_local(self, expected_max_difference=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) - - inputs = self.get_dummy_inputs(torch_device) - output = pipe(**inputs)[0] - - logger = logging.get_logger("diffusers.pipelines.pipeline_utils") - logger.setLevel(diffusers.logging.INFO) - - with tempfile.TemporaryDirectory() as tmpdir: - pipe.save_pretrained(tmpdir, safe_serialization=False) - - with CaptureLogger(logger) as cap_logger: - # NOTE: Resolution must be set to 32 for loading otherwise will lead to OOM on CI hardware - # This attribute is not serialized in the config of the pipeline - pipe_loaded = self.pipeline_class.from_pretrained(tmpdir, resolution=32) - - for component in pipe_loaded.components.values(): - if hasattr(component, "set_default_attn_processor"): - component.set_default_attn_processor() - - for name in pipe_loaded.components.keys(): - if name not in pipe_loaded._optional_components: - assert name in str(cap_logger) - - pipe_loaded.to(torch_device) - pipe_loaded.set_progress_bar_config(disable=None) - - inputs = self.get_dummy_inputs(torch_device) - output_loaded = pipe_loaded(**inputs)[0] - - max_diff = np.abs(to_np(output) - to_np(output_loaded)).max() - self.assertLess(max_diff, expected_max_difference) - - def test_save_load_optional_components(self, expected_max_difference=1e-4): - if not hasattr(self.pipeline_class, "_optional_components"): - 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) - - # set all optional components to None - for optional_component in pipe._optional_components: - setattr(pipe, optional_component, None) - - generator_device = "cpu" - inputs = self.get_dummy_inputs(generator_device) - torch.manual_seed(0) - output = pipe(**inputs)[0] - - with tempfile.TemporaryDirectory() as tmpdir: - pipe.save_pretrained(tmpdir, safe_serialization=False) - # NOTE: Resolution must be set to 32 for loading otherwise will lead to OOM on CI hardware - # This attribute is not serialized in the config of the pipeline - pipe_loaded = self.pipeline_class.from_pretrained(tmpdir, resolution=32) - for component in pipe_loaded.components.values(): - if hasattr(component, "set_default_attn_processor"): - component.set_default_attn_processor() - pipe_loaded.to(torch_device) - pipe_loaded.set_progress_bar_config(disable=None) - - for optional_component in pipe._optional_components: - self.assertTrue( - getattr(pipe_loaded, optional_component) is None, - f"`{optional_component}` did not stay set to None after loading.", - ) - - inputs = self.get_dummy_inputs(generator_device) - torch.manual_seed(0) - output_loaded = pipe_loaded(**inputs)[0] - - max_diff = np.abs(to_np(output) - to_np(output_loaded)).max() - self.assertLess(max_diff, expected_max_difference) - - @unittest.skipIf(torch_device not in ["cuda", "xpu"], reason="float16 requires CUDA or XPU") - @require_accelerator - def test_save_load_float16(self, expected_max_diff=1e-2): - components = self.get_dummy_components() - for name, module in components.items(): - if hasattr(module, "half"): - components[name] = module.to(torch_device).half() - - pipe = self.pipeline_class(**components) - for component in pipe.components.values(): - if hasattr(component, "set_default_attn_processor"): - component.set_default_attn_processor() - pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) - - inputs = self.get_dummy_inputs(torch_device) - output = pipe(**inputs)[0] - - with tempfile.TemporaryDirectory() as tmpdir: - pipe.save_pretrained(tmpdir) - # NOTE: Resolution must be set to 32 for loading otherwise will lead to OOM on CI hardware - # This attribute is not serialized in the config of the pipeline - pipe_loaded = self.pipeline_class.from_pretrained(tmpdir, torch_dtype=torch.float16, resolution=32) - for component in pipe_loaded.components.values(): - if hasattr(component, "set_default_attn_processor"): - component.set_default_attn_processor() - pipe_loaded.to(torch_device) - pipe_loaded.set_progress_bar_config(disable=None) - - for name, component in pipe_loaded.components.items(): - if hasattr(component, "dtype"): - self.assertTrue( - component.dtype == torch.float16, - f"`{name}.dtype` switched from `float16` to {component.dtype} after loading.", - ) - - inputs = self.get_dummy_inputs(torch_device) - output_loaded = pipe_loaded(**inputs)[0] - max_diff = np.abs(to_np(output) - to_np(output_loaded)).max() - self.assertLess( - max_diff, expected_max_diff, "The output of the fp16 pipeline changed after saving and loading." - ) - - @unittest.skip("Test not supported.") - def test_pipeline_with_accelerator_device_map(self): - pass +class TestVisualClozePipelineMemory(VisualClozePipelineTesterConfig, MemoryTesterMixin): + """Memory optimization tests (CPU offload, group offload, layerwise casting) for the VisualCloze pipeline.""" diff --git a/tests/pipelines/visualcloze/test_pipeline_visualcloze_generation.py b/tests/pipelines/visualcloze/test_pipeline_visualcloze_generation.py index 81a73c8a75..7b0b5b5643 100644 --- a/tests/pipelines/visualcloze/test_pipeline_visualcloze_generation.py +++ b/tests/pipelines/visualcloze/test_pipeline_visualcloze_generation.py @@ -1,37 +1,25 @@ import random -import tempfile -import unittest import numpy as np +import pytest import torch from PIL import Image from transformers import AutoConfig, AutoTokenizer, CLIPTextConfig, CLIPTextModel, CLIPTokenizer, T5EncoderModel -import diffusers from diffusers import ( AutoencoderKL, FlowMatchEulerDiscreteScheduler, FluxTransformer2DModel, VisualClozeGenerationPipeline, ) -from diffusers.utils import logging -from ...testing_utils import ( - CaptureLogger, - enable_full_determinism, - floats_tensor, - require_accelerator, - torch_device, -) -from ..test_pipelines_common import PipelineTesterMixin, to_np +from ...testing_utils import floats_tensor, torch_device +from ..testing_utils import BasePipelineTesterConfig, MemoryTesterMixin, PipelineTesterMixin -enable_full_determinism() - - -class VisualClozeGenerationPipelineFastTests(unittest.TestCase, PipelineTesterMixin): +class VisualClozeGenerationPipelineTesterConfig(BasePipelineTesterConfig): pipeline_class = VisualClozeGenerationPipeline - params = frozenset( + required_input_params_in_call_signature = frozenset( [ "task_prompt", "content_prompt", @@ -40,10 +28,8 @@ class VisualClozeGenerationPipelineFastTests(unittest.TestCase, PipelineTesterMi "pooled_prompt_embeds", ] ) - batch_params = frozenset(["task_prompt", "content_prompt", "image"]) - test_xformers_attention = False - test_layerwise_casting = True - test_group_offloading = True + batch_input_params = frozenset(["task_prompt", "content_prompt", "image"]) + output_shape = (3, 32, 32) def get_dummy_components(self): torch.manual_seed(0) @@ -111,7 +97,7 @@ class VisualClozeGenerationPipelineFastTests(unittest.TestCase, PipelineTesterMi "resolution": 32, } - def get_dummy_inputs(self, device, seed=0): + def get_dummy_inputs(self, seed=0): # Create example images to simulate the input format required by VisualCloze context_image = [ Image.fromarray(floats_tensor((32, 32, 3), rng=random.Random(seed), scale=255).numpy().astype(np.uint8)) @@ -130,186 +116,46 @@ class VisualClozeGenerationPipelineFastTests(unittest.TestCase, PipelineTesterMi query_image, # Query image ] - if str(device).startswith("mps"): - generator = torch.manual_seed(seed) - else: - generator = torch.Generator(device="cpu").manual_seed(seed) - inputs = { "task_prompt": "Each row outlines a logical process, starting from [IMAGE1] gray-based depth map with detailed object contours, to achieve [IMAGE2] an image with flawless clarity.", "content_prompt": "A beautiful landscape with mountains and a lake", "image": image, - "generator": generator, + "generator": self.get_generator(seed), "num_inference_steps": 2, "guidance_scale": 5.0, "max_sequence_length": 77, - "output_type": "np", + # Request torch outputs so tests compare torch tensors directly (see `BasePipelineTesterConfig`). + # Note `"pt"` images are `(batch, channels, height, width)`, unlike `"np"` (`(batch, h, w, c)`). + "output_type": "pt", } return inputs - def test_visualcloze_different_prompts(self): - pipe = self.pipeline_class(**self.get_dummy_components()).to(torch_device) - inputs = self.get_dummy_inputs(torch_device) - output_same_prompt = pipe(**inputs).images[0] - - inputs = self.get_dummy_inputs(torch_device) - inputs["task_prompt"] = "A different task to perform." - output_different_prompts = pipe(**inputs).images[0] - - max_diff = np.abs(output_same_prompt - output_different_prompts).max() - - # Outputs should be different - assert max_diff > 1e-6 - - def test_inference_batch_single_identical(self): - self._test_inference_batch_single_identical(expected_max_diff=1e-3) - - def test_different_task_prompts(self, expected_min_diff=1e-1): - pipe = self.pipeline_class(**self.get_dummy_components()).to(torch_device) - inputs = self.get_dummy_inputs(torch_device) +class TestVisualClozeGenerationPipeline(VisualClozeGenerationPipelineTesterConfig, PipelineTesterMixin): + def test_visualcloze_different_task_prompts(self, expected_min_diff=1e-1): + pipe = self.get_pipeline().to(torch_device) + inputs = self.get_dummy_inputs() output_original = pipe(**inputs).images[0] inputs["task_prompt"] = "A different task description for image generation" output_different_task = pipe(**inputs).images[0] # Different task prompts should produce different outputs - max_diff = np.abs(output_original - output_different_task).max() - assert max_diff > expected_min_diff + max_diff = (output_original - output_different_task).abs().max() + assert max_diff > expected_min_diff, "Outputs should be different for different task prompts." - def test_save_load_local(self, expected_max_difference=5e-4): - 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() + 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) - pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) - - inputs = self.get_dummy_inputs(torch_device) - output = pipe(**inputs)[0] - - logger = logging.get_logger("diffusers.pipelines.pipeline_utils") - logger.setLevel(diffusers.logging.INFO) - - with tempfile.TemporaryDirectory() as tmpdir: - pipe.save_pretrained(tmpdir, safe_serialization=False) - - with CaptureLogger(logger) as cap_logger: - # NOTE: Resolution must be set to 32 for loading otherwise will lead to OOM on CI hardware - # This attribute is not serialized in the config of the pipeline - pipe_loaded = self.pipeline_class.from_pretrained(tmpdir, resolution=32) - - for component in pipe_loaded.components.values(): - if hasattr(component, "set_default_attn_processor"): - component.set_default_attn_processor() - - for name in pipe_loaded.components.keys(): - if name not in pipe_loaded._optional_components: - assert name in str(cap_logger) - - pipe_loaded.to(torch_device) - pipe_loaded.set_progress_bar_config(disable=None) - - inputs = self.get_dummy_inputs(torch_device) - output_loaded = pipe_loaded(**inputs)[0] - - max_diff = np.abs(to_np(output) - to_np(output_loaded)).max() - self.assertLess(max_diff, expected_max_difference) - - def test_save_load_optional_components(self, expected_max_difference=1e-4): - if not hasattr(self.pipeline_class, "_optional_components"): - 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) - - # set all optional components to None - for optional_component in pipe._optional_components: - setattr(pipe, optional_component, None) - - generator_device = "cpu" - inputs = self.get_dummy_inputs(generator_device) - torch.manual_seed(0) - output = pipe(**inputs)[0] - - with tempfile.TemporaryDirectory() as tmpdir: - pipe.save_pretrained(tmpdir, safe_serialization=False) - # NOTE: Resolution must be set to 32 for loading otherwise will lead to OOM on CI hardware - # This attribute is not serialized in the config of the pipeline - pipe_loaded = self.pipeline_class.from_pretrained(tmpdir, resolution=32) - for component in pipe_loaded.components.values(): - if hasattr(component, "set_default_attn_processor"): - component.set_default_attn_processor() - pipe_loaded.to(torch_device) - pipe_loaded.set_progress_bar_config(disable=None) - - for optional_component in pipe._optional_components: - self.assertTrue( - getattr(pipe_loaded, optional_component) is None, - f"`{optional_component}` did not stay set to None after loading.", - ) - - inputs = self.get_dummy_inputs(generator_device) - torch.manual_seed(0) - output_loaded = pipe_loaded(**inputs)[0] - - max_diff = np.abs(to_np(output) - to_np(output_loaded)).max() - self.assertLess(max_diff, expected_max_difference) - - @unittest.skipIf(torch_device not in ["cuda", "xpu"], reason="float16 requires CUDA or XPU") - @require_accelerator - def test_save_load_float16(self, expected_max_diff=1e-2): - components = self.get_dummy_components() - for name, module in components.items(): - if hasattr(module, "half"): - components[name] = module.to(torch_device).half() - - pipe = self.pipeline_class(**components) - for component in pipe.components.values(): - if hasattr(component, "set_default_attn_processor"): - component.set_default_attn_processor() - pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) - - inputs = self.get_dummy_inputs(torch_device) - output = pipe(**inputs)[0] - - with tempfile.TemporaryDirectory() as tmpdir: - pipe.save_pretrained(tmpdir) - # NOTE: Resolution must be set to 32 for loading otherwise will lead to OOM on CI hardware - # This attribute is not serialized in the config of the pipeline - pipe_loaded = self.pipeline_class.from_pretrained(tmpdir, torch_dtype=torch.float16, resolution=32) - for component in pipe_loaded.components.values(): - if hasattr(component, "set_default_attn_processor"): - component.set_default_attn_processor() - pipe_loaded.to(torch_device) - pipe_loaded.set_progress_bar_config(disable=None) - - for name, component in pipe_loaded.components.items(): - if hasattr(component, "dtype"): - self.assertTrue( - component.dtype == torch.float16, - f"`{name}.dtype` switched from `float16` to {component.dtype} after loading.", - ) - - inputs = self.get_dummy_inputs(torch_device) - output_loaded = pipe_loaded(**inputs)[0] - max_diff = np.abs(to_np(output) - to_np(output_loaded)).max() - self.assertLess( - max_diff, expected_max_diff, "The output of the fp16 pipeline changed after saving and loading." - ) - - @unittest.skip("Skipped due to missing layout_prompt. Needs further investigation.") - def test_encode_prompt_works_in_isolation(self, extra_required_param_value_dict=None, atol=0.0001, rtol=0.0001): + @pytest.mark.skip( + "`encode_prompt` requires a `layout_prompt` that the pipeline derives from the image grid and that is not a " + "`__call__` argument, and `__call__` cannot run without `task_prompt`, so the prompt encoder cannot be " + "swapped out the way this test does." + ) + def test_encode_prompt_works_in_isolation(self): pass - @unittest.skip("Needs to be revisited later.") - def test_pipeline_with_accelerator_device_map(self, expected_max_difference=0.0001): - pass + +class TestVisualClozeGenerationPipelineMemory(VisualClozeGenerationPipelineTesterConfig, MemoryTesterMixin): + """Memory optimization tests (CPU offload, group offload, layerwise casting) for the VisualCloze pipeline."""