


"""simple docstring"""

import gc
import random
import unittest

import numpy as np
import torch
from PIL import Image

from diffusers import (
    DDIMScheduler,
    KandinskyVaaInpaintPipeline,
    KandinskyVaaPriorPipeline,
    UNetaDConditionModel,
    VQModel,
)
from diffusers.utils import floats_tensor, load_image, load_numpy, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..test_pipelines_common import PipelineTesterMixin, assert_mean_pixel_difference


enable_full_determinism()





class  lowerCamelCase__    (		_a      , unittest.TestCase     ):
      a			:  List[str]           =		KandinskyVaaInpaintPipeline
      a			:  List[str]           =		["""image_embeds""", """negative_image_embeds""", """image""", """mask_image"""]
      a			:  str           =		[
          """image_embeds""",
          """negative_image_embeds""",
          """image""",
          """mask_image""",
      ]
      a			:  List[str]           =		[
          """generator""",
          """height""",
          """width""",
          """latents""",
          """guidance_scale""",
          """num_inference_steps""",
          """return_dict""",
          """guidance_scale""",
          """num_images_per_prompt""",
          """output_type""",
          """return_dict""",
      ]
      a			:  Tuple           =		False
      @property
      def SCREAMING_SNAKE_CASE_						(		self						:						Tuple				):






       '''simple docstring'''


       return 3_2
      @property
      def SCREAMING_SNAKE_CASE_						(		self						:						int				):






       '''simple docstring'''


       return 3_2
      @property
      def SCREAMING_SNAKE_CASE_						(		self						:						str				):






       '''simple docstring'''


       return self.time_input_dim
      @property
      def SCREAMING_SNAKE_CASE_						(		self						:						str				):






       '''simple docstring'''


       return self.time_input_dim * 4
      @property
      def SCREAMING_SNAKE_CASE_						(		self						:						Union[str, Any]				):






       '''simple docstring'''


       return 1_0_0
      @property
      def SCREAMING_SNAKE_CASE_						(		self						:						int				):






       '''simple docstring'''


       torch.manual_seed(0				)

       __lowercase						=						{
           """in_channels""": 9,
           # Out channels is double in channels because predicts mean and variance
           """out_channels""": 8,
           """addition_embed_type""": """image""",
           """down_block_types""": ("""ResnetDownsampleBlock2D""", """SimpleCrossAttnDownBlock2D"""),
           """up_block_types""": ("""SimpleCrossAttnUpBlock2D""", """ResnetUpsampleBlock2D"""),
           """mid_block_type""": """UNetMidBlock2DSimpleCrossAttn""",
           """block_out_channels""": (self.block_out_channels_a, self.block_out_channels_a * 2),
           """layers_per_block""": 1,
           """encoder_hid_dim""": self.text_embedder_hidden_size,
           """encoder_hid_dim_type""": """image_proj""",
           """cross_attention_dim""": self.cross_attention_dim,
           """attention_head_dim""": 4,
           """resnet_time_scale_shift""": """scale_shift""",
           """class_embed_type""": None,
       }

       __lowercase						=						UNetaDConditionModel(**A_				)
       return model
      @property
      def SCREAMING_SNAKE_CASE_						(		self						:						Optional[Any]				):






       '''simple docstring'''


       return {
           "block_out_channels": [3_2, 6_4],
           "down_block_types": ["DownEncoderBlock2D", "AttnDownEncoderBlock2D"],
           "in_channels": 3,
           "latent_channels": 4,
           "layers_per_block": 1,
           "norm_num_groups": 8,
           "norm_type": "spatial",
           "num_vq_embeddings": 1_2,
           "out_channels": 3,
           "up_block_types": [
               "AttnUpDecoderBlock2D",
               "UpDecoderBlock2D",
           ],
           "vq_embed_dim": 4,
       }
      @property
      def SCREAMING_SNAKE_CASE_						(		self						:						Union[str, Any]				):






       '''simple docstring'''


       torch.manual_seed(0				)
       __lowercase						=						VQModel(**self.dummy_movq_kwargs				)
       return model
      def SCREAMING_SNAKE_CASE_						(		self						:						List[Any]				):






       '''simple docstring'''


       __lowercase						=						self.dummy_unet
       __lowercase						=						self.dummy_movq

       __lowercase						=						DDIMScheduler(
           num_train_timesteps=1_0_0_0   ,							beta_schedule="""linear"""   ,							beta_start=0.0_00_85   ,							beta_end=0.0_12   ,							clip_sample=A_   ,							set_alpha_to_one=A_   ,							steps_offset=1   ,							prediction_type="""epsilon"""   ,							thresholding=A_   ,							)

       __lowercase						=						{
           """unet""": unet,
           """scheduler""": scheduler,
           """movq""": movq,
       }

       return components
      def SCREAMING_SNAKE_CASE_						(		self						:						List[Any]   ,							A_						:						List[Any]   ,							A_						:						List[Any]=0				):






       '''simple docstring'''


       __lowercase						=						floats_tensor((1, self.text_embedder_hidden_size)   ,							rng=random.Random(A_				)				).to(A_				)
       __lowercase						=						floats_tensor((1, self.text_embedder_hidden_size)   ,							rng=random.Random(seed + 1				)				).to(
           A_				)
       # create init_image
       __lowercase						=						floats_tensor((1, 3, 6_4, 6_4)   ,							rng=random.Random(A_				)				).to(A_				)
       __lowercase						=						image.cpu().permute(0   ,							2   ,							3   ,							1				)[0]
       __lowercase						=						Image.fromarray(np.uinta(A_				)				).convert("""RGB"""				).resize((2_5_6, 2_5_6)				)
       # create mask
       __lowercase						=						np.ones((6_4, 6_4)   ,							dtype=np.floataa				)
       __lowercase						=						0

       if str(A_				).startswith("""mps"""				):
        __lowercase						=						torch.manual_seed(A_				)
       else:
        __lowercase						=						torch.Generator(device=A_				).manual_seed(A_				)
       __lowercase						=						{
           """image""": init_image,
           """mask_image""": mask,
           """image_embeds""": image_embeds,
           """negative_image_embeds""": negative_image_embeds,
           """generator""": generator,
           """height""": 6_4,
           """width""": 6_4,
           """num_inference_steps""": 2,
           """guidance_scale""": 4.0,
           """output_type""": """np""",
       }
       return inputs
      def SCREAMING_SNAKE_CASE_						(		self						:						List[str]				):






       '''simple docstring'''


       __lowercase						=						"""cpu"""

       __lowercase						=						self.get_dummy_components()

       __lowercase						=						self.pipeline_class(**A_				)
       __lowercase						=						pipe.to(A_				)

       pipe.set_progress_bar_config(disable=A_				)

       __lowercase						=						pipe(**self.get_dummy_inputs(A_				)				)
       __lowercase						=						output.images

       __lowercase						=						pipe(
           **self.get_dummy_inputs(A_				)   ,							return_dict=A_   ,							)[0]

       __lowercase						=						image[0, -3:, -3:, -1]
       __lowercase						=						image_from_tuple[0, -3:, -3:, -1]

       print(F'''image.shape {image.shape}'''				)

       assert image.shape == (1, 6_4, 6_4, 3)

       __lowercase						=						np.array(
           [0.50_77_59_03, 0.49_52_71_95, 0.48_82_45_43, 0.50_19_22_37, 0.48_64_49_06, 0.49_37_38_14, 0.4_78_05_98, 0.47_23_48_27, 0.48_32_78_48]				)

       assert (
           np.abs(image_slice.flatten() - expected_slice				).max() < 1e-2
       ), F''' expected_slice {expected_slice}, but got {image_slice.flatten()}'''
       assert (
           np.abs(image_from_tuple_slice.flatten() - expected_slice				).max() < 1e-2
       ), F''' expected_slice {expected_slice}, but got {image_from_tuple_slice.flatten()}'''






      def SCREAMING_SNAKE_CASE_						(		self						:						List[Any]				):






       '''simple docstring'''


       super().test_inference_batch_single_identical(expected_max_diff=3e-3				)






@slow
@require_torch_gpu
class  lowerCamelCase__    (		unittest.TestCase     ):
      def SCREAMING_SNAKE_CASE_						(		self						:						str				):






       '''simple docstring'''


       super().tearDown()
       gc.collect()
       torch.cuda.empty_cache()






      def SCREAMING_SNAKE_CASE_						(		self						:						Optional[int]				):






       '''simple docstring'''


       __lowercase						=						load_numpy(
           """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"""
           """/kandinskyv22/kandinskyv22_inpaint_cat_with_hat_fp16.npy"""				)

       __lowercase						=						load_image(
           """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main""" """/kandinsky/cat.png"""				)
       __lowercase						=						np.ones((7_6_8, 7_6_8)   ,							dtype=np.floataa				)
       __lowercase						=						0

       __lowercase						=						"""a hat"""

       __lowercase						=						KandinskyVaaPriorPipeline.from_pretrained(
           """kandinsky-community/kandinsky-2-2-prior"""   ,							torch_dtype=torch.floataa				)
       pipe_prior.to(A_				)

       __lowercase						=						KandinskyVaaInpaintPipeline.from_pretrained(
           """kandinsky-community/kandinsky-2-2-decoder-inpaint"""   ,							torch_dtype=torch.floataa				)
       __lowercase						=						pipeline.to(A_				)
       pipeline.set_progress_bar_config(disable=A_				)

       __lowercase						=						torch.Generator(device="""cpu"""				).manual_seed(0				)
       __lowercase				,		__lowercase						=						pipe_prior(
           A_   ,							generator=A_   ,							num_inference_steps=5   ,							negative_prompt=""""""   ,							).to_tuple()

       __lowercase						=						pipeline(
           image=A_   ,							mask_image=A_   ,							image_embeds=A_   ,							negative_image_embeds=A_   ,							generator=A_   ,							num_inference_steps=1_0_0   ,							height=7_6_8   ,							width=7_6_8   ,							output_type="""np"""   ,							)

       __lowercase						=						output.images[0]

       assert image.shape == (7_6_8, 7_6_8, 3)

       assert_mean_pixel_difference(A_   ,							A_				)




"""simple docstring"""

import gc
import random
import unittest

import torch

from diffusers import (
    IFImgaImgPipeline,
    IFImgaImgSuperResolutionPipeline,
    IFInpaintingPipeline,
    IFInpaintingSuperResolutionPipeline,
    IFPipeline,
    IFSuperResolutionPipeline,
)
from diffusers.models.attention_processor import AttnAddedKVProcessor
from diffusers.utils.import_utils import is_xformers_available
from diffusers.utils.testing_utils import floats_tensor, load_numpy, require_torch_gpu, skip_mps, slow, torch_device

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import PipelineTesterMixin, assert_mean_pixel_difference
from . import IFPipelineTesterMixin





@skip_mps
class  lowerCamelCase__    (		_a      , _a      , unittest.TestCase     ):
      a			:  Any           =		IFPipeline
      a			:  str           =		TEXT_TO_IMAGE_PARAMS - {"""width""", """height""", """latents"""}
      a			:  Dict           =		TEXT_TO_IMAGE_BATCH_PARAMS
      a			:  int           =		PipelineTesterMixin.required_optional_params - {"""latents"""}
      def SCREAMING_SNAKE_CASE_						(		self						:						List[Any]				):






       '''simple docstring'''


       return self._get_dummy_components()
      def SCREAMING_SNAKE_CASE_						(		self						:						Any   ,							A_						:						int   ,							A_						:						Dict=0				):






       '''simple docstring'''


       if str(A_				).startswith("""mps"""				):
        __lowercase						=						torch.manual_seed(A_				)
       else:
        __lowercase						=						torch.Generator(device=A_				).manual_seed(A_				)

       __lowercase						=						{
           """prompt""": """A painting of a squirrel eating a burger""",
           """generator""": generator,
           """num_inference_steps""": 2,
           """output_type""": """numpy""",
       }

       return inputs
      def SCREAMING_SNAKE_CASE_						(		self						:						Union[str, Any]				):






       '''simple docstring'''


       self._test_save_load_optional_components()
      @unittest.skipIf(torch_device != """cuda"""   ,							reason="""float16 requires CUDA"""				)
      def SCREAMING_SNAKE_CASE_						(		self						:						List[str]				):






       '''simple docstring'''


       super().test_save_load_floataa(expected_max_diff=1e-1				)
      def SCREAMING_SNAKE_CASE_						(		self						:						List[Any]				):






       '''simple docstring'''


       self._test_attention_slicing_forward_pass(expected_max_diff=1e-2				)
      def SCREAMING_SNAKE_CASE_						(		self						:						List[str]				):






       '''simple docstring'''


       self._test_save_load_local()
      def SCREAMING_SNAKE_CASE_						(		self						:						Union[str, Any]				):






       '''simple docstring'''


       self._test_inference_batch_single_identical(
           expected_max_diff=1e-2   ,							)






      @unittest.skipIf(
          torch_device != """cuda""" or not is_xformers_available()   ,							reason="""XFormers attention is only available with CUDA and `xformers` installed"""   ,							)
      def SCREAMING_SNAKE_CASE_						(		self						:						Tuple				):






       '''simple docstring'''


       self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=1e-3				)






@slow
@require_torch_gpu
class  lowerCamelCase__    (		unittest.TestCase     ):
      def SCREAMING_SNAKE_CASE_						(		self						:						str				):






       '''simple docstring'''


       super().tearDown()
       gc.collect()
       torch.cuda.empty_cache()
      def SCREAMING_SNAKE_CASE_						(		self						:						Union[str, Any]				):






       '''simple docstring'''


       __lowercase						=						IFPipeline.from_pretrained("""DeepFloyd/IF-I-XL-v1.0"""   ,							variant="""fp16"""   ,							torch_dtype=torch.floataa				)

       __lowercase						=						IFSuperResolutionPipeline.from_pretrained(
           """DeepFloyd/IF-II-L-v1.0"""   ,							variant="""fp16"""   ,							torch_dtype=torch.floataa   ,							text_encoder=A_   ,							tokenizer=A_				)

       # pre compute text embeddings and remove T5 to save memory

       pipe_a.text_encoder.to("""cuda"""				)

       __lowercase				,		__lowercase						=						pipe_a.encode_prompt("""anime turtle"""   ,							device="""cuda"""				)

       del pipe_a.tokenizer
       del pipe_a.text_encoder
       gc.collect()

       __lowercase						=						None
       __lowercase						=						None

       pipe_a.enable_model_cpu_offload()
       pipe_a.enable_model_cpu_offload()

       pipe_a.unet.set_attn_processor(AttnAddedKVProcessor()				)
       pipe_a.unet.set_attn_processor(AttnAddedKVProcessor()				)

       self._test_if(A_   ,							A_   ,							A_   ,							A_				)

       pipe_a.remove_all_hooks()
       pipe_a.remove_all_hooks()

       # img2img

       __lowercase						=						IFImgaImgPipeline(**pipe_a.components				)
       __lowercase						=						IFImgaImgSuperResolutionPipeline(**pipe_a.components				)

       pipe_a.enable_model_cpu_offload()
       pipe_a.enable_model_cpu_offload()

       pipe_a.unet.set_attn_processor(AttnAddedKVProcessor()				)
       pipe_a.unet.set_attn_processor(AttnAddedKVProcessor()				)

       self._test_if_imgaimg(A_   ,							A_   ,							A_   ,							A_				)

       pipe_a.remove_all_hooks()
       pipe_a.remove_all_hooks()

       # inpainting

       __lowercase						=						IFInpaintingPipeline(**pipe_a.components				)
       __lowercase						=						IFInpaintingSuperResolutionPipeline(**pipe_a.components				)

       pipe_a.enable_model_cpu_offload()
       pipe_a.enable_model_cpu_offload()

       pipe_a.unet.set_attn_processor(AttnAddedKVProcessor()				)
       pipe_a.unet.set_attn_processor(AttnAddedKVProcessor()				)

       self._test_if_inpainting(A_   ,							A_   ,							A_   ,							A_				)
      def SCREAMING_SNAKE_CASE_						(		self						:						Optional[int]   ,							A_						:						Any   ,							A_						:						int   ,							A_						:						str   ,							A_						:						Dict				):






       '''simple docstring'''


       _start_torch_memory_measurement()

       __lowercase						=						torch.Generator(device="""cpu"""				).manual_seed(0				)
       __lowercase						=						pipe_a(
           prompt_embeds=A_   ,							negative_prompt_embeds=A_   ,							num_inference_steps=2   ,							generator=A_   ,							output_type="""np"""   ,							)

       __lowercase						=						output.images[0]

       assert image.shape == (6_4, 6_4, 3)

       __lowercase						=						torch.cuda.max_memory_allocated()
       assert mem_bytes < 1_3 * 1_0**9

       __lowercase						=						load_numpy(
           """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/if/test_if.npy"""				)
       assert_mean_pixel_difference(A_   ,							A_				)

       # pipeline 2

       _start_torch_memory_measurement()

       __lowercase						=						torch.Generator(device="""cpu"""				).manual_seed(0				)

       __lowercase						=						floats_tensor((1, 3, 6_4, 6_4)   ,							rng=random.Random(0				)				).to(A_				)

       __lowercase						=						pipe_a(
           prompt_embeds=A_   ,							negative_prompt_embeds=A_   ,							image=A_   ,							generator=A_   ,							num_inference_steps=2   ,							output_type="""np"""   ,							)

       __lowercase						=						output.images[0]

       assert image.shape == (2_5_6, 2_5_6, 3)

       __lowercase						=						torch.cuda.max_memory_allocated()
       assert mem_bytes < 4 * 1_0**9

       __lowercase						=						load_numpy(
           """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/if/test_if_superresolution_stage_II.npy"""				)
       assert_mean_pixel_difference(A_   ,							A_				)
      def SCREAMING_SNAKE_CASE_						(		self						:						Optional[Any]   ,							A_						:						Tuple   ,							A_						:						List[Any]   ,							A_						:						List[Any]   ,							A_						:						Any				):






       '''simple docstring'''


       _start_torch_memory_measurement()

       __lowercase						=						floats_tensor((1, 3, 6_4, 6_4)   ,							rng=random.Random(0				)				).to(A_				)

       __lowercase						=						torch.Generator(device="""cpu"""				).manual_seed(0				)

       __lowercase						=						pipe_a(
           prompt_embeds=A_   ,							negative_prompt_embeds=A_   ,							image=A_   ,							num_inference_steps=2   ,							generator=A_   ,							output_type="""np"""   ,							)

       __lowercase						=						output.images[0]

       assert image.shape == (6_4, 6_4, 3)

       __lowercase						=						torch.cuda.max_memory_allocated()
       assert mem_bytes < 1_0 * 1_0**9

       __lowercase						=						load_numpy(
           """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/if/test_if_img2img.npy"""				)
       assert_mean_pixel_difference(A_   ,							A_				)

       # pipeline 2

       _start_torch_memory_measurement()

       __lowercase						=						torch.Generator(device="""cpu"""				).manual_seed(0				)

       __lowercase						=						floats_tensor((1, 3, 2_5_6, 2_5_6)   ,							rng=random.Random(0				)				).to(A_				)
       __lowercase						=						floats_tensor((1, 3, 6_4, 6_4)   ,							rng=random.Random(0				)				).to(A_				)

       __lowercase						=						pipe_a(
           prompt_embeds=A_   ,							negative_prompt_embeds=A_   ,							image=A_   ,							original_image=A_   ,							generator=A_   ,							num_inference_steps=2   ,							output_type="""np"""   ,							)

       __lowercase						=						output.images[0]

       assert image.shape == (2_5_6, 2_5_6, 3)

       __lowercase						=						torch.cuda.max_memory_allocated()
       assert mem_bytes < 4 * 1_0**9

       __lowercase						=						load_numpy(
           """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/if/test_if_img2img_superresolution_stage_II.npy"""				)
       assert_mean_pixel_difference(A_   ,							A_				)






      def SCREAMING_SNAKE_CASE_						(		self						:						List[str]   ,							A_						:						List[Any]   ,							A_						:						str   ,							A_						:						List[Any]   ,							A_						:						List[Any]				):






       '''simple docstring'''


       _start_torch_memory_measurement()

       __lowercase						=						floats_tensor((1, 3, 6_4, 6_4)   ,							rng=random.Random(0				)				).to(A_				)
       __lowercase						=						floats_tensor((1, 3, 6_4, 6_4)   ,							rng=random.Random(1				)				).to(A_				)

       __lowercase						=						torch.Generator(device="""cpu"""				).manual_seed(0				)
       __lowercase						=						pipe_a(
           prompt_embeds=A_   ,							negative_prompt_embeds=A_   ,							image=A_   ,							mask_image=A_   ,							num_inference_steps=2   ,							generator=A_   ,							output_type="""np"""   ,							)

       __lowercase						=						output.images[0]

       assert image.shape == (6_4, 6_4, 3)

       __lowercase						=						torch.cuda.max_memory_allocated()
       assert mem_bytes < 1_0 * 1_0**9

       __lowercase						=						load_numpy(
           """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/if/test_if_inpainting.npy"""				)
       assert_mean_pixel_difference(A_   ,							A_				)

       # pipeline 2

       _start_torch_memory_measurement()

       __lowercase						=						torch.Generator(device="""cpu"""				).manual_seed(0				)

       __lowercase						=						floats_tensor((1, 3, 6_4, 6_4)   ,							rng=random.Random(0				)				).to(A_				)
       __lowercase						=						floats_tensor((1, 3, 2_5_6, 2_5_6)   ,							rng=random.Random(0				)				).to(A_				)
       __lowercase						=						floats_tensor((1, 3, 2_5_6, 2_5_6)   ,							rng=random.Random(1				)				).to(A_				)

       __lowercase						=						pipe_a(
           prompt_embeds=A_   ,							negative_prompt_embeds=A_   ,							image=A_   ,							mask_image=A_   ,							original_image=A_   ,							generator=A_   ,							num_inference_steps=2   ,							output_type="""np"""   ,							)

       __lowercase						=						output.images[0]

       assert image.shape == (2_5_6, 2_5_6, 3)

       __lowercase						=						torch.cuda.max_memory_allocated()
       assert mem_bytes < 4 * 1_0**9

       __lowercase						=						load_numpy(
           """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/if/test_if_inpainting_superresolution_stage_II.npy"""				)
       assert_mean_pixel_difference(A_   ,							A_				)






def lowerCAmelCase_   (					):






 """simple docstring"""







 torch.cuda.empty_cache()
 torch.cuda.reset_max_memory_allocated()
 torch.cuda.reset_peak_memory_stats()
