
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			(   UpperCamelCase_						,				UpperCamelCase_						,				unittest.TestCase			):





     """simple docstring"""







     SCREAMING_SNAKE_CASE_		 =	IFPipeline
     SCREAMING_SNAKE_CASE_		 =	TEXT_TO_IMAGE_PARAMS - {'''width''', '''height''', '''latents'''}
     SCREAMING_SNAKE_CASE_		 =	TEXT_TO_IMAGE_BATCH_PARAMS
     SCREAMING_SNAKE_CASE_		 =	PipelineTesterMixin.required_optional_params - {'''latents'''}






     def  __SCREAMING_SNAKE_CASE      (				self	)	-> List[str]:




           """simple docstring"""







           return self._get_dummy_components()






     def  __SCREAMING_SNAKE_CASE      (				self    ,    __SCREAMING_SNAKE_CASE    ,    __SCREAMING_SNAKE_CASE=0	)	-> Any:




           """simple docstring"""







           if str(snake_case__	).startswith('''mps'''	):
                 UpperCamelCase__     : Optional[int]											=					torch.manual_seed(snake_case__	)
           else:
                 UpperCamelCase__     : Optional[Any]											=					torch.Generator(device=snake_case__	).manual_seed(snake_case__	)

           UpperCamelCase__     : Dict											=					{
               '''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	)	-> List[Any]:




           """simple docstring"""







           self._test_save_load_optional_components()






     @unittest.skipIf(torch_device != '''cuda'''    ,    reason='''float16 requires CUDA'''	)
     def  __SCREAMING_SNAKE_CASE      (				self	)	-> str:




           """simple docstring"""







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






     def  __SCREAMING_SNAKE_CASE      (				self	)	-> Dict:




           """simple docstring"""







           self._test_attention_slicing_forward_pass(expected_max_diff=1e-2	)






     def  __SCREAMING_SNAKE_CASE      (				self	)	-> Tuple:




           """simple docstring"""







           self._test_save_load_local()






     def  __SCREAMING_SNAKE_CASE      (				self	)	-> Optional[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	)	-> Optional[Any]:




           """simple docstring"""







           self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=1e-3	)




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





     """simple docstring"""







     def  __SCREAMING_SNAKE_CASE      (				self	)	-> Union[str, Any]:




           """simple docstring"""







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






     def  __SCREAMING_SNAKE_CASE      (				self	)	-> int:




           """simple docstring"""







           UpperCamelCase__     : Optional[int]											=					IFPipeline.from_pretrained('''DeepFloyd/IF-I-XL-v1.0'''    ,    variant='''fp16'''    ,    torch_dtype=torch.floataa	)

           UpperCamelCase__     : List[str]											=					IFSuperResolutionPipeline.from_pretrained(
               '''DeepFloyd/IF-II-L-v1.0'''    ,    variant='''fp16'''    ,    torch_dtype=torch.floataa    ,    text_encoder=snake_case__    ,    tokenizer=snake_case__	)

           # pre compute text embeddings and remove T5 to save memory

           pipe_a.text_encoder.to('''cuda'''	)

           UpperCamelCase__					,UpperCamelCase__     : str											=					pipe_a.encode_prompt('''anime turtle'''    ,    device='''cuda'''	)

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

           UpperCamelCase__     : int											=					None
           UpperCamelCase__     : Union[str, 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(snake_case__    ,    snake_case__    ,    snake_case__    ,    snake_case__	)

           pipe_a.remove_all_hooks()
           pipe_a.remove_all_hooks()

           # img2img

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

           pipe_a.remove_all_hooks()
           pipe_a.remove_all_hooks()

           # inpainting

           UpperCamelCase__     : Any											=					IFInpaintingPipeline(**pipe_a.components	)
           UpperCamelCase__     : Optional[int]											=					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(snake_case__    ,    snake_case__    ,    snake_case__    ,    snake_case__	)






     def  __SCREAMING_SNAKE_CASE      (				self    ,    __SCREAMING_SNAKE_CASE    ,    __SCREAMING_SNAKE_CASE    ,    __SCREAMING_SNAKE_CASE    ,    __SCREAMING_SNAKE_CASE	)	-> Dict:




           """simple docstring"""







           _start_torch_memory_measurement()

           UpperCamelCase__     : Optional[int]											=					torch.Generator(device='''cpu'''	).manual_seed(0	)
           UpperCamelCase__     : Union[str, Any]											=					pipe_a(
               prompt_embeds=snake_case__    ,    negative_prompt_embeds=snake_case__    ,    num_inference_steps=2    ,    generator=snake_case__    ,    output_type='''np'''    ,    )

           UpperCamelCase__     : Tuple											=					output.images[0]

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

           UpperCamelCase__     : Tuple											=					torch.cuda.max_memory_allocated()
           assert mem_bytes < 1_3 * 1_0**9

           UpperCamelCase__     : Dict											=					load_numpy(
               '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/if/test_if.npy'''	)
           assert_mean_pixel_difference(snake_case__    ,    snake_case__	)

           # pipeline 2

           _start_torch_memory_measurement()

           UpperCamelCase__     : str											=					torch.Generator(device='''cpu'''	).manual_seed(0	)

           UpperCamelCase__     : Dict											=					floats_tensor((1, 3, 6_4, 6_4)    ,    rng=random.Random(0	)	).to(snake_case__	)

           UpperCamelCase__     : Optional[Any]											=					pipe_a(
               prompt_embeds=snake_case__    ,    negative_prompt_embeds=snake_case__    ,    image=snake_case__    ,    generator=snake_case__    ,    num_inference_steps=2    ,    output_type='''np'''    ,    )

           UpperCamelCase__     : Union[str, Any]											=					output.images[0]

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

           UpperCamelCase__     : Dict											=					torch.cuda.max_memory_allocated()
           assert mem_bytes < 4 * 1_0**9

           UpperCamelCase__     : str											=					load_numpy(
               '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/if/test_if_superresolution_stage_II.npy'''	)
           assert_mean_pixel_difference(snake_case__    ,    snake_case__	)






     def  __SCREAMING_SNAKE_CASE      (				self    ,    __SCREAMING_SNAKE_CASE    ,    __SCREAMING_SNAKE_CASE    ,    __SCREAMING_SNAKE_CASE    ,    __SCREAMING_SNAKE_CASE	)	-> Any:




           """simple docstring"""







           _start_torch_memory_measurement()

           UpperCamelCase__     : Optional[int]											=					floats_tensor((1, 3, 6_4, 6_4)    ,    rng=random.Random(0	)	).to(snake_case__	)

           UpperCamelCase__     : Tuple											=					torch.Generator(device='''cpu'''	).manual_seed(0	)

           UpperCamelCase__     : List[str]											=					pipe_a(
               prompt_embeds=snake_case__    ,    negative_prompt_embeds=snake_case__    ,    image=snake_case__    ,    num_inference_steps=2    ,    generator=snake_case__    ,    output_type='''np'''    ,    )

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

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

           UpperCamelCase__     : List[str]											=					torch.cuda.max_memory_allocated()
           assert mem_bytes < 1_0 * 1_0**9

           UpperCamelCase__     : List[Any]											=					load_numpy(
               '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/if/test_if_img2img.npy'''	)
           assert_mean_pixel_difference(snake_case__    ,    snake_case__	)

           # pipeline 2

           _start_torch_memory_measurement()

           UpperCamelCase__     : str											=					torch.Generator(device='''cpu'''	).manual_seed(0	)

           UpperCamelCase__     : Optional[Any]											=					floats_tensor((1, 3, 2_5_6, 2_5_6)    ,    rng=random.Random(0	)	).to(snake_case__	)
           UpperCamelCase__     : int											=					floats_tensor((1, 3, 6_4, 6_4)    ,    rng=random.Random(0	)	).to(snake_case__	)

           UpperCamelCase__     : Optional[int]											=					pipe_a(
               prompt_embeds=snake_case__    ,    negative_prompt_embeds=snake_case__    ,    image=snake_case__    ,    original_image=snake_case__    ,    generator=snake_case__    ,    num_inference_steps=2    ,    output_type='''np'''    ,    )

           UpperCamelCase__     : Any											=					output.images[0]

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

           UpperCamelCase__     : int											=					torch.cuda.max_memory_allocated()
           assert mem_bytes < 4 * 1_0**9

           UpperCamelCase__     : Optional[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(snake_case__    ,    snake_case__	)






     def  __SCREAMING_SNAKE_CASE      (				self    ,    __SCREAMING_SNAKE_CASE    ,    __SCREAMING_SNAKE_CASE    ,    __SCREAMING_SNAKE_CASE    ,    __SCREAMING_SNAKE_CASE	)	-> str:




           """simple docstring"""







           _start_torch_memory_measurement()

           UpperCamelCase__     : int											=					floats_tensor((1, 3, 6_4, 6_4)    ,    rng=random.Random(0	)	).to(snake_case__	)
           UpperCamelCase__     : Any											=					floats_tensor((1, 3, 6_4, 6_4)    ,    rng=random.Random(1	)	).to(snake_case__	)

           UpperCamelCase__     : Tuple											=					torch.Generator(device='''cpu'''	).manual_seed(0	)
           UpperCamelCase__     : Optional[int]											=					pipe_a(
               prompt_embeds=snake_case__    ,    negative_prompt_embeds=snake_case__    ,    image=snake_case__    ,    mask_image=snake_case__    ,    num_inference_steps=2    ,    generator=snake_case__    ,    output_type='''np'''    ,    )

           UpperCamelCase__     : Union[str, Any]											=					output.images[0]

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

           UpperCamelCase__     : List[str]											=					torch.cuda.max_memory_allocated()
           assert mem_bytes < 1_0 * 1_0**9

           UpperCamelCase__     : Dict											=					load_numpy(
               '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/if/test_if_inpainting.npy'''	)
           assert_mean_pixel_difference(snake_case__    ,    snake_case__	)

           # pipeline 2

           _start_torch_memory_measurement()

           UpperCamelCase__     : Tuple											=					torch.Generator(device='''cpu'''	).manual_seed(0	)

           UpperCamelCase__     : Union[str, Any]											=					floats_tensor((1, 3, 6_4, 6_4)    ,    rng=random.Random(0	)	).to(snake_case__	)
           UpperCamelCase__     : List[str]											=					floats_tensor((1, 3, 2_5_6, 2_5_6)    ,    rng=random.Random(0	)	).to(snake_case__	)
           UpperCamelCase__     : Tuple											=					floats_tensor((1, 3, 2_5_6, 2_5_6)    ,    rng=random.Random(1	)	).to(snake_case__	)

           UpperCamelCase__     : List[Any]											=					pipe_a(
               prompt_embeds=snake_case__    ,    negative_prompt_embeds=snake_case__    ,    image=snake_case__    ,    mask_image=snake_case__    ,    original_image=snake_case__    ,    generator=snake_case__    ,    num_inference_steps=2    ,    output_type='''np'''    ,    )

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

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

           UpperCamelCase__     : int											=					torch.cuda.max_memory_allocated()
           assert mem_bytes < 4 * 1_0**9

           UpperCamelCase__     : Dict											=					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(snake_case__    ,    snake_case__	)


def 				SCREAMING_SNAKE_CASE_   (      ):
      torch.cuda.empty_cache()
      torch.cuda.reset_max_memory_allocated()
      torch.cuda.reset_peak_memory_stats()





import warnings

from ...processing_utils import ProcessorMixin
from ...tokenization_utils_base import BatchEncoding




class   UpperCamelCase			(  SCREAMING_SNAKE_CASE ):
    __UpperCamelCase    =["image_processor", "tokenizer"]
    __UpperCamelCase    ="CLIPImageProcessor"
    __UpperCamelCase    =("CLIPTokenizer", "CLIPTokenizerFast")

    def __init__(      self			:   int				,     snake_case__			:   List[Any]=None				,     snake_case__			:   Any=None				,     **snake_case__			:   Any     ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE		    =      None
     if "feature_extractor" in kwargs:
      warnings.warn(
          'The `feature_extractor` argument is deprecated and will be removed in v5, use `image_processor`'
          ' instead.'				,     snake_case__				,     )
      SCREAMING_SNAKE_CASE		    =      kwargs.pop('feature_extractor'     )

     SCREAMING_SNAKE_CASE		    =      image_processor if image_processor is not None else feature_extractor
     if image_processor is None:
      raise ValueError('You need to specify an `image_processor`.'     )
     if tokenizer is None:
      raise ValueError('You need to specify a `tokenizer`.'     )

     super().__init__(snake_case__				,     snake_case__     )

    def __call__(      self			:   List[Any]				,     snake_case__			:   Tuple=None				,     snake_case__			:   Optional[int]=None				,     snake_case__			:   List[str]=None				,     **snake_case__			:   Dict     ):

     """simple docstring"""






     if text is None and images is None:
      raise ValueError('You have to specify either text or images. Both cannot be none.'     )

     if text is not None:
      SCREAMING_SNAKE_CASE		    =      self.tokenizer(snake_case__				,     return_tensors=snake_case__				,     **snake_case__     )

     if images is not None:
      SCREAMING_SNAKE_CASE		    =      self.image_processor(snake_case__				,     return_tensors=snake_case__				,     **snake_case__     )

     if text is not None and images is not None:
      SCREAMING_SNAKE_CASE		    =      image_features.pixel_values
      return encoding
     elif text is not None:
      return encoding
     else:
      return BatchEncoding(data=dict(**snake_case__     )				,     tensor_type=snake_case__     )

    def UpperCamelCase      (      self			:   Union[str, Any]				,     *snake_case__			:   Optional[Any]				,     **snake_case__			:   Optional[int]     ):

     """simple docstring"""






     return self.tokenizer.batch_decode(*snake_case__				,     **snake_case__     )

    def UpperCamelCase      (      self			:   Dict				,     *snake_case__			:   Optional[Any]				,     **snake_case__			:   Dict     ):

     """simple docstring"""






     return self.tokenizer.decode(*snake_case__				,     **snake_case__     )

    @property
    def UpperCamelCase      (      self			:   Any     ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE		    =      self.tokenizer.model_input_names
     SCREAMING_SNAKE_CASE		    =      self.image_processor.model_input_names
     return list(dict.fromkeys(tokenizer_input_names + image_processor_input_names     )     )

    @property
    def UpperCamelCase      (      self			:   int     ):

     """simple docstring"""






     warnings.warn(
         '`feature_extractor_class` is deprecated and will be removed in v5. Use `image_processor_class` instead.'				,     snake_case__				,     )
     return self.image_processor_class




    @property
    def UpperCamelCase      (      self			:   Tuple     ):

     """simple docstring"""






     warnings.warn(
         '`feature_extractor` is deprecated and will be removed in v5. Use `image_processor` instead.'				,     snake_case__				,     )
     return self.image_processor

