| | import gc |
| | import random |
| | import tempfile |
| | import unittest |
| |
|
| | import numpy as np |
| | import torch |
| | from transformers import ( |
| | CLIPImageProcessor, |
| | CLIPVisionConfig, |
| | CLIPVisionModelWithProjection, |
| | ) |
| |
|
| | import diffusers |
| | from diffusers import ( |
| | AutoencoderKLTemporalDecoder, |
| | EulerDiscreteScheduler, |
| | StableVideoDiffusionPipeline, |
| | UNetSpatioTemporalConditionModel, |
| | ) |
| | from diffusers.utils import is_accelerate_available, is_accelerate_version, load_image, logging |
| | from diffusers.utils.import_utils import is_xformers_available |
| | from diffusers.utils.testing_utils import ( |
| | CaptureLogger, |
| | enable_full_determinism, |
| | floats_tensor, |
| | numpy_cosine_similarity_distance, |
| | require_torch_gpu, |
| | slow, |
| | torch_device, |
| | ) |
| |
|
| | from ..test_pipelines_common import PipelineTesterMixin |
| |
|
| |
|
| | enable_full_determinism() |
| |
|
| |
|
| | def to_np(tensor): |
| | if isinstance(tensor, torch.Tensor): |
| | tensor = tensor.detach().cpu().numpy() |
| |
|
| | return tensor |
| |
|
| |
|
| | class StableVideoDiffusionPipelineFastTests(PipelineTesterMixin, unittest.TestCase): |
| | pipeline_class = StableVideoDiffusionPipeline |
| | params = frozenset(["image"]) |
| | batch_params = frozenset(["image", "generator"]) |
| | required_optional_params = frozenset( |
| | [ |
| | "num_inference_steps", |
| | "generator", |
| | "latents", |
| | "return_dict", |
| | ] |
| | ) |
| |
|
| | def get_dummy_components(self): |
| | torch.manual_seed(0) |
| | unet = UNetSpatioTemporalConditionModel( |
| | block_out_channels=(32, 64), |
| | layers_per_block=2, |
| | sample_size=32, |
| | in_channels=8, |
| | out_channels=4, |
| | down_block_types=( |
| | "CrossAttnDownBlockSpatioTemporal", |
| | "DownBlockSpatioTemporal", |
| | ), |
| | up_block_types=("UpBlockSpatioTemporal", "CrossAttnUpBlockSpatioTemporal"), |
| | cross_attention_dim=32, |
| | num_attention_heads=8, |
| | projection_class_embeddings_input_dim=96, |
| | addition_time_embed_dim=32, |
| | ) |
| | scheduler = EulerDiscreteScheduler( |
| | beta_start=0.00085, |
| | beta_end=0.012, |
| | beta_schedule="scaled_linear", |
| | interpolation_type="linear", |
| | num_train_timesteps=1000, |
| | prediction_type="v_prediction", |
| | sigma_max=700.0, |
| | sigma_min=0.002, |
| | steps_offset=1, |
| | timestep_spacing="leading", |
| | timestep_type="continuous", |
| | trained_betas=None, |
| | use_karras_sigmas=True, |
| | ) |
| |
|
| | torch.manual_seed(0) |
| | vae = AutoencoderKLTemporalDecoder( |
| | block_out_channels=[32, 64], |
| | in_channels=3, |
| | out_channels=3, |
| | down_block_types=["DownEncoderBlock2D", "DownEncoderBlock2D"], |
| | latent_channels=4, |
| | ) |
| |
|
| | torch.manual_seed(0) |
| | config = CLIPVisionConfig( |
| | hidden_size=32, |
| | projection_dim=32, |
| | num_hidden_layers=5, |
| | num_attention_heads=4, |
| | image_size=32, |
| | intermediate_size=37, |
| | patch_size=1, |
| | ) |
| | image_encoder = CLIPVisionModelWithProjection(config) |
| |
|
| | torch.manual_seed(0) |
| | feature_extractor = CLIPImageProcessor(crop_size=32, size=32) |
| | components = { |
| | "unet": unet, |
| | "image_encoder": image_encoder, |
| | "scheduler": scheduler, |
| | "vae": vae, |
| | "feature_extractor": feature_extractor, |
| | } |
| | 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) |
| |
|
| | image = floats_tensor((1, 3, 32, 32), rng=random.Random(0)).to(device) |
| | inputs = { |
| | "generator": generator, |
| | "image": image, |
| | "num_inference_steps": 2, |
| | "output_type": "pt", |
| | "min_guidance_scale": 1.0, |
| | "max_guidance_scale": 2.5, |
| | "num_frames": 2, |
| | "height": 32, |
| | "width": 32, |
| | } |
| | return inputs |
| |
|
| | @unittest.skip("Deprecated functionality") |
| | def test_attention_slicing_forward_pass(self): |
| | pass |
| |
|
| | @unittest.skip("Batched inference works and outputs look correct, but the test is failing") |
| | def test_inference_batch_single_identical( |
| | self, |
| | batch_size=2, |
| | expected_max_diff=1e-4, |
| | ): |
| | components = self.get_dummy_components() |
| | pipe = self.pipeline_class(**components) |
| | for components in pipe.components.values(): |
| | if hasattr(components, "set_default_attn_processor"): |
| | components.set_default_attn_processor() |
| | pipe.to(torch_device) |
| |
|
| | pipe.set_progress_bar_config(disable=None) |
| | inputs = self.get_dummy_inputs(torch_device) |
| |
|
| | |
| | inputs["generator"] = torch.Generator("cpu").manual_seed(0) |
| |
|
| | logger = logging.get_logger(pipe.__module__) |
| | logger.setLevel(level=diffusers.logging.FATAL) |
| |
|
| | |
| | batched_inputs = {} |
| | batched_inputs.update(inputs) |
| |
|
| | batched_inputs["generator"] = [torch.Generator("cpu").manual_seed(0) for i in range(batch_size)] |
| | batched_inputs["image"] = torch.cat([inputs["image"]] * batch_size, dim=0) |
| |
|
| | output = pipe(**inputs).frames |
| | output_batch = pipe(**batched_inputs).frames |
| |
|
| | assert len(output_batch) == batch_size |
| |
|
| | max_diff = np.abs(to_np(output_batch[0]) - to_np(output[0])).max() |
| | assert max_diff < expected_max_diff |
| |
|
| | @unittest.skip("Test is similar to test_inference_batch_single_identical") |
| | def test_inference_batch_consistent(self): |
| | pass |
| |
|
| | def test_np_output_type(self): |
| | 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) |
| | inputs["output_type"] = "np" |
| | output = pipe(**inputs).frames |
| | self.assertTrue(isinstance(output, np.ndarray)) |
| | self.assertEqual(len(output.shape), 5) |
| |
|
| | def test_dict_tuple_outputs_equivalent(self, expected_max_difference=1e-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() |
| |
|
| | pipe.to(torch_device) |
| | pipe.set_progress_bar_config(disable=None) |
| |
|
| | generator_device = "cpu" |
| | output = pipe(**self.get_dummy_inputs(generator_device)).frames[0] |
| | output_tuple = pipe(**self.get_dummy_inputs(generator_device), return_dict=False)[0] |
| |
|
| | max_diff = np.abs(to_np(output) - to_np(output_tuple)).max() |
| | self.assertLess(max_diff, expected_max_difference) |
| |
|
| | @unittest.skip("Test is currently failing") |
| | def test_float16_inference(self, expected_max_diff=5e-2): |
| | 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) |
| |
|
| | components = self.get_dummy_components() |
| | pipe_fp16 = self.pipeline_class(**components) |
| | for component in pipe_fp16.components.values(): |
| | if hasattr(component, "set_default_attn_processor"): |
| | component.set_default_attn_processor() |
| |
|
| | pipe_fp16.to(torch_device, torch.float16) |
| | pipe_fp16.set_progress_bar_config(disable=None) |
| |
|
| | inputs = self.get_dummy_inputs(torch_device) |
| | output = pipe(**inputs).frames[0] |
| |
|
| | fp16_inputs = self.get_dummy_inputs(torch_device) |
| | output_fp16 = pipe_fp16(**fp16_inputs).frames[0] |
| |
|
| | max_diff = np.abs(to_np(output) - to_np(output_fp16)).max() |
| | self.assertLess(max_diff, expected_max_diff, "The outputs of the fp16 and fp32 pipelines are too different.") |
| |
|
| | @unittest.skipIf(torch_device != "cuda", reason="float16 requires CUDA") |
| | 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).frames[0] |
| |
|
| | with tempfile.TemporaryDirectory() as tmpdir: |
| | pipe.save_pretrained(tmpdir) |
| | pipe_loaded = self.pipeline_class.from_pretrained(tmpdir, torch_dtype=torch.float16) |
| | for component in pipe_loaded.components.values(): |
| | if hasattr(component, "set_default_attn_processor"): |
| | component.set_default_attn_processor() |
| | pipe_loaded.to(torch_device) |
| | pipe_loaded.set_progress_bar_config(disable=None) |
| |
|
| | for name, component in pipe_loaded.components.items(): |
| | if hasattr(component, "dtype"): |
| | self.assertTrue( |
| | component.dtype == torch.float16, |
| | f"`{name}.dtype` switched from `float16` to {component.dtype} after loading.", |
| | ) |
| |
|
| | inputs = self.get_dummy_inputs(torch_device) |
| | output_loaded = pipe_loaded(**inputs).frames[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." |
| | ) |
| |
|
| | 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) |
| |
|
| | |
| | for optional_component in pipe._optional_components: |
| | setattr(pipe, optional_component, None) |
| |
|
| | generator_device = "cpu" |
| | inputs = self.get_dummy_inputs(generator_device) |
| | output = pipe(**inputs).frames[0] |
| |
|
| | with tempfile.TemporaryDirectory() as tmpdir: |
| | pipe.save_pretrained(tmpdir, safe_serialization=False) |
| | pipe_loaded = self.pipeline_class.from_pretrained(tmpdir) |
| | 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) |
| | output_loaded = pipe_loaded(**inputs).frames[0] |
| |
|
| | max_diff = np.abs(to_np(output) - to_np(output_loaded)).max() |
| | self.assertLess(max_diff, expected_max_difference) |
| |
|
| | def test_save_load_local(self, expected_max_difference=9e-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() |
| |
|
| | pipe.to(torch_device) |
| | pipe.set_progress_bar_config(disable=None) |
| |
|
| | inputs = self.get_dummy_inputs(torch_device) |
| | output = pipe(**inputs).frames[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: |
| | pipe_loaded = self.pipeline_class.from_pretrained(tmpdir) |
| |
|
| | 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).frames[0] |
| |
|
| | max_diff = np.abs(to_np(output) - to_np(output_loaded)).max() |
| | self.assertLess(max_diff, expected_max_difference) |
| |
|
| | @unittest.skipIf(torch_device != "cuda", reason="CUDA and CPU are required to switch devices") |
| | def test_to_device(self): |
| | components = self.get_dummy_components() |
| | pipe = self.pipeline_class(**components) |
| | pipe.set_progress_bar_config(disable=None) |
| |
|
| | pipe.to("cpu") |
| | model_devices = [ |
| | component.device.type for component in pipe.components.values() if hasattr(component, "device") |
| | ] |
| | self.assertTrue(all(device == "cpu" for device in model_devices)) |
| |
|
| | output_cpu = pipe(**self.get_dummy_inputs("cpu")).frames[0] |
| | self.assertTrue(np.isnan(output_cpu).sum() == 0) |
| |
|
| | pipe.to("cuda") |
| | model_devices = [ |
| | component.device.type for component in pipe.components.values() if hasattr(component, "device") |
| | ] |
| | self.assertTrue(all(device == "cuda" for device in model_devices)) |
| |
|
| | output_cuda = pipe(**self.get_dummy_inputs("cuda")).frames[0] |
| | self.assertTrue(np.isnan(to_np(output_cuda)).sum() == 0) |
| |
|
| | def test_to_dtype(self): |
| | components = self.get_dummy_components() |
| | pipe = self.pipeline_class(**components) |
| | pipe.set_progress_bar_config(disable=None) |
| |
|
| | model_dtypes = [component.dtype for component in pipe.components.values() if hasattr(component, "dtype")] |
| | self.assertTrue(all(dtype == torch.float32 for dtype in model_dtypes)) |
| |
|
| | pipe.to(dtype=torch.float16) |
| | model_dtypes = [component.dtype for component in pipe.components.values() if hasattr(component, "dtype")] |
| | self.assertTrue(all(dtype == torch.float16 for dtype in model_dtypes)) |
| |
|
| | @unittest.skipIf( |
| | torch_device != "cuda" or not is_accelerate_available() or is_accelerate_version("<", "0.14.0"), |
| | reason="CPU offload is only available with CUDA and `accelerate v0.14.0` or higher", |
| | ) |
| | def test_sequential_cpu_offload_forward_pass(self, expected_max_diff=1e-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() |
| | pipe.to(torch_device) |
| | pipe.set_progress_bar_config(disable=None) |
| |
|
| | generator_device = "cpu" |
| | inputs = self.get_dummy_inputs(generator_device) |
| | output_without_offload = pipe(**inputs).frames[0] |
| |
|
| | pipe.enable_sequential_cpu_offload() |
| |
|
| | inputs = self.get_dummy_inputs(generator_device) |
| | output_with_offload = pipe(**inputs).frames[0] |
| |
|
| | max_diff = np.abs(to_np(output_with_offload) - to_np(output_without_offload)).max() |
| | self.assertLess(max_diff, expected_max_diff, "CPU offloading should not affect the inference results") |
| |
|
| | @unittest.skipIf( |
| | torch_device != "cuda" or not is_accelerate_available() or is_accelerate_version("<", "0.17.0"), |
| | reason="CPU offload is only available with CUDA and `accelerate v0.17.0` or higher", |
| | ) |
| | def test_model_cpu_offload_forward_pass(self, expected_max_diff=2e-4): |
| | generator_device = "cpu" |
| | 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 = pipe.to(torch_device) |
| | pipe.set_progress_bar_config(disable=None) |
| |
|
| | inputs = self.get_dummy_inputs(generator_device) |
| | output_without_offload = pipe(**inputs).frames[0] |
| |
|
| | pipe.enable_model_cpu_offload() |
| | inputs = self.get_dummy_inputs(generator_device) |
| | output_with_offload = pipe(**inputs).frames[0] |
| |
|
| | max_diff = np.abs(to_np(output_with_offload) - to_np(output_without_offload)).max() |
| | self.assertLess(max_diff, expected_max_diff, "CPU offloading should not affect the inference results") |
| | offloaded_modules = [ |
| | v |
| | for k, v in pipe.components.items() |
| | if isinstance(v, torch.nn.Module) and k not in pipe._exclude_from_cpu_offload |
| | ] |
| | ( |
| | self.assertTrue(all(v.device.type == "cpu" for v in offloaded_modules)), |
| | f"Not offloaded: {[v for v in offloaded_modules if v.device.type != 'cpu']}", |
| | ) |
| |
|
| | @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): |
| | expected_max_diff = 9e-4 |
| |
|
| | if not self.test_xformers_attention: |
| | 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) |
| |
|
| | inputs = self.get_dummy_inputs(torch_device) |
| | output_without_offload = pipe(**inputs).frames[0] |
| | output_without_offload = ( |
| | output_without_offload.cpu() if torch.is_tensor(output_without_offload) else output_without_offload |
| | ) |
| |
|
| | pipe.enable_xformers_memory_efficient_attention() |
| | inputs = self.get_dummy_inputs(torch_device) |
| | output_with_offload = pipe(**inputs).frames[0] |
| | output_with_offload = ( |
| | output_with_offload.cpu() if torch.is_tensor(output_with_offload) else output_without_offload |
| | ) |
| |
|
| | max_diff = np.abs(to_np(output_with_offload) - to_np(output_without_offload)).max() |
| | self.assertLess(max_diff, expected_max_diff, "XFormers attention should not affect the inference results") |
| |
|
| | def test_disable_cfg(self): |
| | 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) |
| | inputs["max_guidance_scale"] = 1.0 |
| | output = pipe(**inputs).frames |
| | self.assertEqual(len(output.shape), 5) |
| |
|
| |
|
| | @slow |
| | @require_torch_gpu |
| | class StableVideoDiffusionPipelineSlowTests(unittest.TestCase): |
| | def tearDown(self): |
| | |
| | super().tearDown() |
| | gc.collect() |
| | torch.cuda.empty_cache() |
| |
|
| | def test_sd_video(self): |
| | pipe = StableVideoDiffusionPipeline.from_pretrained( |
| | "stabilityai/stable-video-diffusion-img2vid", |
| | variant="fp16", |
| | torch_dtype=torch.float16, |
| | ) |
| | pipe = pipe.to(torch_device) |
| | pipe.enable_model_cpu_offload() |
| | pipe.set_progress_bar_config(disable=None) |
| | image = load_image( |
| | "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/pix2pix/cat_6.png?download=true" |
| | ) |
| |
|
| | generator = torch.Generator("cpu").manual_seed(0) |
| | num_frames = 3 |
| |
|
| | output = pipe( |
| | image=image, |
| | num_frames=num_frames, |
| | generator=generator, |
| | num_inference_steps=3, |
| | output_type="np", |
| | ) |
| |
|
| | image = output.frames[0] |
| | assert image.shape == (num_frames, 576, 1024, 3) |
| |
|
| | image_slice = image[0, -3:, -3:, -1] |
| | expected_slice = np.array([0.8592, 0.8645, 0.8499, 0.8722, 0.8769, 0.8421, 0.8557, 0.8528, 0.8285]) |
| | assert numpy_cosine_similarity_distance(image_slice.flatten(), expected_slice.flatten()) < 1e-3 |
| |
|
