Spaces:
No application file
No application file
| import tempfile | |
| import numpy as np | |
| import torch | |
| from transformers import AutoTokenizer, T5EncoderModel | |
| from diffusers import DDPMScheduler, UNet2DConditionModel | |
| from diffusers.models.attention_processor import AttnAddedKVProcessor | |
| from diffusers.pipelines.deepfloyd_if import IFWatermarker | |
| from diffusers.utils.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) | |
| text_encoder = T5EncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5") | |
| 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) | |
| text_encoder = T5EncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5") | |
| 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() | |
| 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) | |