From 08ef86ded1af6e81e213af58df711ba81a791228 Mon Sep 17 00:00:00 2001 From: Akshan Krithick Date: Wed, 26 Aug 2026 13:48:00 -0700 Subject: [PATCH] remove dead deepfloyd if shared test mixin --- tests/pipelines/deepfloyd_if/__init__.py | 277 ----------------------- 1 file changed, 277 deletions(-) diff --git a/tests/pipelines/deepfloyd_if/__init__.py b/tests/pipelines/deepfloyd_if/__init__.py index d8f72bd2ca5b..e69de29bb2d1 100644 --- a/tests/pipelines/deepfloyd_if/__init__.py +++ b/tests/pipelines/deepfloyd_if/__init__.py @@ -1,277 +0,0 @@ -import tempfile - -import numpy as np -import torch -from transformers import AutoConfig, AutoTokenizer, T5EncoderModel - -from diffusers import DDPMScheduler, UNet2DConditionModel -from diffusers.models.attention_processor import AttnAddedKVProcessor -from diffusers.pipelines.deepfloyd_if import IFWatermarker - -from ...testing_utils import torch_device -from ..test_pipelines_common import to_np - - -# WARN: the hf-internal-testing/tiny-random-t5 text encoder has some non-determinism in the `save_load` tests. - - -class IFPipelineTesterMixin: - def _get_dummy_components(self): - torch.manual_seed(0) - config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5") - text_encoder = T5EncoderModel(config) - - torch.manual_seed(0) - tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5") - - torch.manual_seed(0) - unet = UNet2DConditionModel( - sample_size=32, - layers_per_block=1, - block_out_channels=[32, 64], - down_block_types=[ - "ResnetDownsampleBlock2D", - "SimpleCrossAttnDownBlock2D", - ], - mid_block_type="UNetMidBlock2DSimpleCrossAttn", - up_block_types=["SimpleCrossAttnUpBlock2D", "ResnetUpsampleBlock2D"], - in_channels=3, - out_channels=6, - cross_attention_dim=32, - encoder_hid_dim=32, - attention_head_dim=8, - addition_embed_type="text", - addition_embed_type_num_heads=2, - cross_attention_norm="group_norm", - resnet_time_scale_shift="scale_shift", - act_fn="gelu", - ) - unet.set_attn_processor(AttnAddedKVProcessor()) # For reproducibility tests - - torch.manual_seed(0) - scheduler = DDPMScheduler( - num_train_timesteps=1000, - beta_schedule="squaredcos_cap_v2", - beta_start=0.0001, - beta_end=0.02, - thresholding=True, - dynamic_thresholding_ratio=0.95, - sample_max_value=1.0, - prediction_type="epsilon", - variance_type="learned_range", - ) - - torch.manual_seed(0) - watermarker = IFWatermarker() - - return { - "text_encoder": text_encoder, - "tokenizer": tokenizer, - "unet": unet, - "scheduler": scheduler, - "watermarker": watermarker, - "safety_checker": None, - "feature_extractor": None, - } - - def _get_superresolution_dummy_components(self): - torch.manual_seed(0) - config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5") - text_encoder = T5EncoderModel(config) - - torch.manual_seed(0) - tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5") - - torch.manual_seed(0) - unet = UNet2DConditionModel( - sample_size=32, - layers_per_block=[1, 2], - block_out_channels=[32, 64], - down_block_types=[ - "ResnetDownsampleBlock2D", - "SimpleCrossAttnDownBlock2D", - ], - mid_block_type="UNetMidBlock2DSimpleCrossAttn", - up_block_types=["SimpleCrossAttnUpBlock2D", "ResnetUpsampleBlock2D"], - in_channels=6, - out_channels=6, - cross_attention_dim=32, - encoder_hid_dim=32, - attention_head_dim=8, - addition_embed_type="text", - addition_embed_type_num_heads=2, - cross_attention_norm="group_norm", - resnet_time_scale_shift="scale_shift", - act_fn="gelu", - class_embed_type="timestep", - mid_block_scale_factor=1.414, - time_embedding_act_fn="gelu", - time_embedding_dim=32, - ) - unet.set_attn_processor(AttnAddedKVProcessor()) # For reproducibility tests - - torch.manual_seed(0) - scheduler = DDPMScheduler( - num_train_timesteps=1000, - beta_schedule="squaredcos_cap_v2", - beta_start=0.0001, - beta_end=0.02, - thresholding=True, - dynamic_thresholding_ratio=0.95, - sample_max_value=1.0, - prediction_type="epsilon", - variance_type="learned_range", - ) - - torch.manual_seed(0) - image_noising_scheduler = DDPMScheduler( - num_train_timesteps=1000, - beta_schedule="squaredcos_cap_v2", - beta_start=0.0001, - beta_end=0.02, - ) - - torch.manual_seed(0) - watermarker = IFWatermarker() - - return { - "text_encoder": text_encoder, - "tokenizer": tokenizer, - "unet": unet, - "scheduler": scheduler, - "image_noising_scheduler": image_noising_scheduler, - "watermarker": watermarker, - "safety_checker": None, - "feature_extractor": None, - } - - # this test is modified from the base class because if pipelines set the text encoder - # as optional with the intention that the user is allowed to encode the prompt once - # and then pass the embeddings directly to the pipeline. The base class test uses - # the unmodified arguments from `self.get_dummy_inputs` which will pass the unencoded - # prompt to the pipeline when the text encoder is set to None, throwing an error. - # So we make the test reflect the intended usage of setting the text encoder to None. - def _test_save_load_optional_components(self): - components = self.get_dummy_components() - pipe = self.pipeline_class(**components) - pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) - - inputs = self.get_dummy_inputs(torch_device) - - prompt = inputs["prompt"] - generator = inputs["generator"] - num_inference_steps = inputs["num_inference_steps"] - output_type = inputs["output_type"] - - if "image" in inputs: - image = inputs["image"] - else: - image = None - - if "mask_image" in inputs: - mask_image = inputs["mask_image"] - else: - mask_image = None - - if "original_image" in inputs: - original_image = inputs["original_image"] - else: - original_image = None - - prompt_embeds, negative_prompt_embeds = pipe.encode_prompt(prompt) - - # inputs with prompt converted to embeddings - inputs = { - "prompt_embeds": prompt_embeds, - "negative_prompt_embeds": negative_prompt_embeds, - "generator": generator, - "num_inference_steps": num_inference_steps, - "output_type": output_type, - } - - if image is not None: - inputs["image"] = image - - if mask_image is not None: - inputs["mask_image"] = mask_image - - if original_image is not None: - inputs["original_image"] = original_image - - # set all optional components to None - for optional_component in pipe._optional_components: - setattr(pipe, optional_component, None) - - output = pipe(**inputs)[0] - - with tempfile.TemporaryDirectory() as tmpdir: - pipe.save_pretrained(tmpdir) - pipe_loaded = self.pipeline_class.from_pretrained(tmpdir) - pipe_loaded.to(torch_device) - pipe_loaded.set_progress_bar_config(disable=None) - - pipe_loaded.unet.set_attn_processor(AttnAddedKVProcessor()) # For reproducibility tests - - 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(torch_device) - - generator = inputs["generator"] - num_inference_steps = inputs["num_inference_steps"] - output_type = inputs["output_type"] - - # inputs with prompt converted to embeddings - inputs = { - "prompt_embeds": prompt_embeds, - "negative_prompt_embeds": negative_prompt_embeds, - "generator": generator, - "num_inference_steps": num_inference_steps, - "output_type": output_type, - } - - if image is not None: - inputs["image"] = image - - if mask_image is not None: - inputs["mask_image"] = mask_image - - if original_image is not None: - inputs["original_image"] = original_image - - output_loaded = pipe_loaded(**inputs)[0] - - max_diff = np.abs(to_np(output) - to_np(output_loaded)).max() - self.assertLess(max_diff, 1e-4) - - # Modified from `PipelineTesterMixin` to set the attn processor as it's not serialized. - # This should be handled in the base test and then this method can be removed. - def _test_save_load_local(self): - 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) - 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) - pipe_loaded = self.pipeline_class.from_pretrained(tmpdir) - pipe_loaded.to(torch_device) - pipe_loaded.set_progress_bar_config(disable=None) - - pipe_loaded.unet.set_attn_processor(AttnAddedKVProcessor()) # For reproducibility tests - - 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, 1e-4)