
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    _snake_case (						UpperCAmelCase_	,    UpperCAmelCase_	,    unittest.TestCase				):
 __lowerCAmelCase	:  Dict	   =      IFPipeline
 __lowerCAmelCase	:  Optional[Any]	   =      TEXT_TO_IMAGE_PARAMS - {'width', 'height', 'latents'}
 __lowerCAmelCase	:  Union[str, Any]	   =      TEXT_TO_IMAGE_BATCH_PARAMS
 __lowerCAmelCase	:  Optional[int]	   =      PipelineTesterMixin.required_optional_params - {'latents'}


 def   lowercase__  (       self):


   '''simple docstring'''




   return self._get_dummy_components()


 def   lowercase__  (       self       ,	SCREAMING_SNAKE_CASE_       ,	SCREAMING_SNAKE_CASE_=0):


   '''simple docstring'''




   if str(SCREAMING_SNAKE_CASE_).startswith("""mps"""):
     lowercase__	: Union[str, Any]     		=   torch.manual_seed(SCREAMING_SNAKE_CASE_)
   else:
     lowercase__	: int     		=   torch.Generator(device=SCREAMING_SNAKE_CASE_).manual_seed(SCREAMING_SNAKE_CASE_)

   lowercase__	: Optional[Any]     		=   {
       """prompt""": """A painting of a squirrel eating a burger""",
       """generator""": generator,
       """num_inference_steps""": 2,
       """output_type""": """numpy""",
   }

   return inputs


 def   lowercase__  (       self):


   '''simple docstring'''




   self._test_save_load_optional_components()


 @unittest.skipIf(torch_device != """cuda"""       ,	reason="""float16 requires CUDA""")
 def   lowercase__  (       self):


   '''simple docstring'''




   super().test_save_load_floataa(expected_max_diff=1E-1)


 def   lowercase__  (       self):


   '''simple docstring'''




   self._test_attention_slicing_forward_pass(expected_max_diff=1E-2)


 def   lowercase__  (       self):


   '''simple docstring'''




   self._test_save_load_local()


 def   lowercase__  (       self):


   '''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   lowercase__  (       self):


   '''simple docstring'''




   self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=1E-3)



@slow
@require_torch_gpu
class    _snake_case (						unittest.TestCase				):


 def   lowercase__  (       self):


   '''simple docstring'''




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


 def   lowercase__  (       self):


   '''simple docstring'''




   lowercase__	: Union[str, Any]     		=   IFPipeline.from_pretrained("""DeepFloyd/IF-I-XL-v1.0"""       ,	variant="""fp16"""       ,	torch_dtype=torch.floataa)

   lowercase__	: Union[str, Any]     		=   IFSuperResolutionPipeline.from_pretrained(
       """DeepFloyd/IF-II-L-v1.0"""       ,	variant="""fp16"""       ,	torch_dtype=torch.floataa       ,	text_encoder=SCREAMING_SNAKE_CASE_       ,	tokenizer=SCREAMING_SNAKE_CASE_)

   # pre compute text embeddings and remove T5 to save memory

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

   lowercase__		,     lowercase__	: Dict     		=   pipe_a.encode_prompt("""anime turtle"""       ,	device="""cuda""")

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

   lowercase__	: str     		=   None
   lowercase__	: Optional[Any]     		=   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(SCREAMING_SNAKE_CASE_       ,	SCREAMING_SNAKE_CASE_       ,	SCREAMING_SNAKE_CASE_       ,	SCREAMING_SNAKE_CASE_)

   pipe_a.remove_all_hooks()
   pipe_a.remove_all_hooks()

   # img2img

   lowercase__	: List[Any]     		=   IFImgaImgPipeline(**pipe_a.components)
   lowercase__	: str     		=   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(SCREAMING_SNAKE_CASE_       ,	SCREAMING_SNAKE_CASE_       ,	SCREAMING_SNAKE_CASE_       ,	SCREAMING_SNAKE_CASE_)

   pipe_a.remove_all_hooks()
   pipe_a.remove_all_hooks()

   # inpainting

   lowercase__	: str     		=   IFInpaintingPipeline(**pipe_a.components)
   lowercase__	: List[str]     		=   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(SCREAMING_SNAKE_CASE_       ,	SCREAMING_SNAKE_CASE_       ,	SCREAMING_SNAKE_CASE_       ,	SCREAMING_SNAKE_CASE_)


 def   lowercase__  (       self       ,	SCREAMING_SNAKE_CASE_       ,	SCREAMING_SNAKE_CASE_       ,	SCREAMING_SNAKE_CASE_       ,	SCREAMING_SNAKE_CASE_):


   '''simple docstring'''




   _start_torch_memory_measurement()

   lowercase__	: str     		=   torch.Generator(device="""cpu""").manual_seed(0)
   lowercase__	: List[str]     		=   pipe_a(
       prompt_embeds=SCREAMING_SNAKE_CASE_       ,	negative_prompt_embeds=SCREAMING_SNAKE_CASE_       ,	num_inference_steps=2       ,	generator=SCREAMING_SNAKE_CASE_       ,	output_type="""np"""       ,	)

   lowercase__	: List[str]     		=   output.images[0]

   assert image.shape == (64, 64, 3)

   lowercase__	: str     		=   torch.cuda.max_memory_allocated()
   assert mem_bytes < 13 * 10**9

   lowercase__	: Union[str, Any]     		=   load_numpy(
       """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/if/test_if.npy""")
   assert_mean_pixel_difference(SCREAMING_SNAKE_CASE_       ,	SCREAMING_SNAKE_CASE_)

   # pipeline 2

   _start_torch_memory_measurement()

   lowercase__	: Tuple     		=   torch.Generator(device="""cpu""").manual_seed(0)

   lowercase__	: Dict     		=   floats_tensor((1, 3, 64, 64)       ,	rng=random.Random(0)).to(SCREAMING_SNAKE_CASE_)

   lowercase__	: Optional[Any]     		=   pipe_a(
       prompt_embeds=SCREAMING_SNAKE_CASE_       ,	negative_prompt_embeds=SCREAMING_SNAKE_CASE_       ,	image=SCREAMING_SNAKE_CASE_       ,	generator=SCREAMING_SNAKE_CASE_       ,	num_inference_steps=2       ,	output_type="""np"""       ,	)

   lowercase__	: List[Any]     		=   output.images[0]

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

   lowercase__	: Union[str, Any]     		=   torch.cuda.max_memory_allocated()
   assert mem_bytes < 4 * 10**9

   lowercase__	: Optional[int]     		=   load_numpy(
       """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/if/test_if_superresolution_stage_II.npy""")
   assert_mean_pixel_difference(SCREAMING_SNAKE_CASE_       ,	SCREAMING_SNAKE_CASE_)


 def   lowercase__  (       self       ,	SCREAMING_SNAKE_CASE_       ,	SCREAMING_SNAKE_CASE_       ,	SCREAMING_SNAKE_CASE_       ,	SCREAMING_SNAKE_CASE_):


   '''simple docstring'''




   _start_torch_memory_measurement()

   lowercase__	: Optional[Any]     		=   floats_tensor((1, 3, 64, 64)       ,	rng=random.Random(0)).to(SCREAMING_SNAKE_CASE_)

   lowercase__	: int     		=   torch.Generator(device="""cpu""").manual_seed(0)

   lowercase__	: Optional[Any]     		=   pipe_a(
       prompt_embeds=SCREAMING_SNAKE_CASE_       ,	negative_prompt_embeds=SCREAMING_SNAKE_CASE_       ,	image=SCREAMING_SNAKE_CASE_       ,	num_inference_steps=2       ,	generator=SCREAMING_SNAKE_CASE_       ,	output_type="""np"""       ,	)

   lowercase__	: Tuple     		=   output.images[0]

   assert image.shape == (64, 64, 3)

   lowercase__	: str     		=   torch.cuda.max_memory_allocated()
   assert mem_bytes < 10 * 10**9

   lowercase__	: List[str]     		=   load_numpy(
       """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/if/test_if_img2img.npy""")
   assert_mean_pixel_difference(SCREAMING_SNAKE_CASE_       ,	SCREAMING_SNAKE_CASE_)

   # pipeline 2

   _start_torch_memory_measurement()

   lowercase__	: Union[str, Any]     		=   torch.Generator(device="""cpu""").manual_seed(0)

   lowercase__	: Dict     		=   floats_tensor((1, 3, 2_56, 2_56)       ,	rng=random.Random(0)).to(SCREAMING_SNAKE_CASE_)
   lowercase__	: Union[str, Any]     		=   floats_tensor((1, 3, 64, 64)       ,	rng=random.Random(0)).to(SCREAMING_SNAKE_CASE_)

   lowercase__	: str     		=   pipe_a(
       prompt_embeds=SCREAMING_SNAKE_CASE_       ,	negative_prompt_embeds=SCREAMING_SNAKE_CASE_       ,	image=SCREAMING_SNAKE_CASE_       ,	original_image=SCREAMING_SNAKE_CASE_       ,	generator=SCREAMING_SNAKE_CASE_       ,	num_inference_steps=2       ,	output_type="""np"""       ,	)

   lowercase__	: int     		=   output.images[0]

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

   lowercase__	: Optional[Any]     		=   torch.cuda.max_memory_allocated()
   assert mem_bytes < 4 * 10**9

   lowercase__	: Union[str, Any]     		=   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(SCREAMING_SNAKE_CASE_       ,	SCREAMING_SNAKE_CASE_)



 def   lowercase__  (       self       ,	SCREAMING_SNAKE_CASE_       ,	SCREAMING_SNAKE_CASE_       ,	SCREAMING_SNAKE_CASE_       ,	SCREAMING_SNAKE_CASE_):


   '''simple docstring'''




   _start_torch_memory_measurement()

   lowercase__	: List[str]     		=   floats_tensor((1, 3, 64, 64)       ,	rng=random.Random(0)).to(SCREAMING_SNAKE_CASE_)
   lowercase__	: Dict     		=   floats_tensor((1, 3, 64, 64)       ,	rng=random.Random(1)).to(SCREAMING_SNAKE_CASE_)

   lowercase__	: str     		=   torch.Generator(device="""cpu""").manual_seed(0)
   lowercase__	: List[Any]     		=   pipe_a(
       prompt_embeds=SCREAMING_SNAKE_CASE_       ,	negative_prompt_embeds=SCREAMING_SNAKE_CASE_       ,	image=SCREAMING_SNAKE_CASE_       ,	mask_image=SCREAMING_SNAKE_CASE_       ,	num_inference_steps=2       ,	generator=SCREAMING_SNAKE_CASE_       ,	output_type="""np"""       ,	)

   lowercase__	: Any     		=   output.images[0]

   assert image.shape == (64, 64, 3)

   lowercase__	: List[str]     		=   torch.cuda.max_memory_allocated()
   assert mem_bytes < 10 * 10**9

   lowercase__	: List[Any]     		=   load_numpy(
       """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/if/test_if_inpainting.npy""")
   assert_mean_pixel_difference(SCREAMING_SNAKE_CASE_       ,	SCREAMING_SNAKE_CASE_)

   # pipeline 2

   _start_torch_memory_measurement()

   lowercase__	: List[str]     		=   torch.Generator(device="""cpu""").manual_seed(0)

   lowercase__	: Union[str, Any]     		=   floats_tensor((1, 3, 64, 64)       ,	rng=random.Random(0)).to(SCREAMING_SNAKE_CASE_)
   lowercase__	: Dict     		=   floats_tensor((1, 3, 2_56, 2_56)       ,	rng=random.Random(0)).to(SCREAMING_SNAKE_CASE_)
   lowercase__	: Optional[int]     		=   floats_tensor((1, 3, 2_56, 2_56)       ,	rng=random.Random(1)).to(SCREAMING_SNAKE_CASE_)

   lowercase__	: Any     		=   pipe_a(
       prompt_embeds=SCREAMING_SNAKE_CASE_       ,	negative_prompt_embeds=SCREAMING_SNAKE_CASE_       ,	image=SCREAMING_SNAKE_CASE_       ,	mask_image=SCREAMING_SNAKE_CASE_       ,	original_image=SCREAMING_SNAKE_CASE_       ,	generator=SCREAMING_SNAKE_CASE_       ,	num_inference_steps=2       ,	output_type="""np"""       ,	)

   lowercase__	: Optional[Any]     		=   output.images[0]

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

   lowercase__	: str     		=   torch.cuda.max_memory_allocated()
   assert mem_bytes < 4 * 10**9

   lowercase__	: str     		=   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(SCREAMING_SNAKE_CASE_       ,	SCREAMING_SNAKE_CASE_)

def 		UpperCamelCase						(   )		->						List[str]:


  '''simple docstring'''






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








def 		UpperCamelCase						(   lowercase_ )		->						float:


  '''simple docstring'''






  if not nums:  # Makes sure that the list is not empty
    raise ValueError("""List is empty""" )

  lowercase__	: int     		=   sum(lowercase_ ) / len(lowercase_ )  # Calculate the average
  return sum(abs(x - average ) for x in nums ) / len(lowercase_ )


if __name__ == "__main__":
     import doctest

     doctest.testmod()






