




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


       def    lowerCAmelCase   ( self	)   ->	Optional[Any]:



            '''simple docstring'''




            return self._get_dummy_components()


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



            '''simple docstring'''




            if str(__SCREAMING_SNAKE_CASE	).startswith('''mps'''	):
                 __snake_case            =	torch.manual_seed(__SCREAMING_SNAKE_CASE	)
            else:
                 __snake_case            =	torch.Generator(device=__SCREAMING_SNAKE_CASE	).manual_seed(__SCREAMING_SNAKE_CASE	)

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

            return inputs


       def    lowerCAmelCase   ( self	)   ->	int:



            '''simple docstring'''




            self._test_save_load_optional_components()


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



            '''simple docstring'''




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


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



            '''simple docstring'''




            self._test_attention_slicing_forward_pass(expected_max_diff=1E-2	)


       def    lowerCAmelCase   ( self	)   ->	int:



            '''simple docstring'''




            self._test_save_load_local()


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



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



            '''simple docstring'''




            self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=1E-3	)





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


       def    lowerCAmelCase   ( self	)   ->	Tuple:



            '''simple docstring'''




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


       def    lowerCAmelCase   ( self	)   ->	Dict:



            '''simple docstring'''




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

            __snake_case            =	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'''	)

            __snake_case       ,     __snake_case            =	pipe_a.encode_prompt('''anime turtle'''     ,	device='''cuda'''	)

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

            __snake_case            =	None
            __snake_case            =	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

            __snake_case            =	IFImgaImgPipeline(**pipe_a.components	)
            __snake_case            =	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

            __snake_case            =	IFInpaintingPipeline(**pipe_a.components	)
            __snake_case            =	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    lowerCAmelCase   ( self     ,	__SCREAMING_SNAKE_CASE     ,	__SCREAMING_SNAKE_CASE     ,	__SCREAMING_SNAKE_CASE     ,	__SCREAMING_SNAKE_CASE	)   ->	int:



            '''simple docstring'''




            _start_torch_memory_measurement()

            __snake_case            =	torch.Generator(device='''cpu'''	).manual_seed(0	)
            __snake_case            =	pipe_a(
                prompt_embeds=__SCREAMING_SNAKE_CASE     ,	negative_prompt_embeds=__SCREAMING_SNAKE_CASE     ,	num_inference_steps=2     ,	generator=__SCREAMING_SNAKE_CASE     ,	output_type='''np'''     ,	)

            __snake_case            =	output.images[0]

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

            __snake_case            =	torch.cuda.max_memory_allocated()
            assert mem_bytes < 13 * 10**9

            __snake_case            =	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()

            __snake_case            =	torch.Generator(device='''cpu'''	).manual_seed(0	)

            __snake_case            =	floats_tensor((1, 3, 64, 64)     ,	rng=random.Random(0	)	).to(__SCREAMING_SNAKE_CASE	)

            __snake_case            =	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'''     ,	)

            __snake_case            =	output.images[0]

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

            __snake_case            =	torch.cuda.max_memory_allocated()
            assert mem_bytes < 4 * 10**9

            __snake_case            =	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    lowerCAmelCase   ( self     ,	__SCREAMING_SNAKE_CASE     ,	__SCREAMING_SNAKE_CASE     ,	__SCREAMING_SNAKE_CASE     ,	__SCREAMING_SNAKE_CASE	)   ->	Union[str, Any]:



            '''simple docstring'''




            _start_torch_memory_measurement()

            __snake_case            =	floats_tensor((1, 3, 64, 64)     ,	rng=random.Random(0	)	).to(__SCREAMING_SNAKE_CASE	)

            __snake_case            =	torch.Generator(device='''cpu'''	).manual_seed(0	)

            __snake_case            =	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'''     ,	)

            __snake_case            =	output.images[0]

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

            __snake_case            =	torch.cuda.max_memory_allocated()
            assert mem_bytes < 10 * 10**9

            __snake_case            =	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()

            __snake_case            =	torch.Generator(device='''cpu'''	).manual_seed(0	)

            __snake_case            =	floats_tensor((1, 3, 256, 256)     ,	rng=random.Random(0	)	).to(__SCREAMING_SNAKE_CASE	)
            __snake_case            =	floats_tensor((1, 3, 64, 64)     ,	rng=random.Random(0	)	).to(__SCREAMING_SNAKE_CASE	)

            __snake_case            =	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'''     ,	)

            __snake_case            =	output.images[0]

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

            __snake_case            =	torch.cuda.max_memory_allocated()
            assert mem_bytes < 4 * 10**9

            __snake_case            =	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    lowerCAmelCase   ( self     ,	__SCREAMING_SNAKE_CASE     ,	__SCREAMING_SNAKE_CASE     ,	__SCREAMING_SNAKE_CASE     ,	__SCREAMING_SNAKE_CASE	)   ->	Optional[Any]:



            '''simple docstring'''




            _start_torch_memory_measurement()

            __snake_case            =	floats_tensor((1, 3, 64, 64)     ,	rng=random.Random(0	)	).to(__SCREAMING_SNAKE_CASE	)
            __snake_case            =	floats_tensor((1, 3, 64, 64)     ,	rng=random.Random(1	)	).to(__SCREAMING_SNAKE_CASE	)

            __snake_case            =	torch.Generator(device='''cpu'''	).manual_seed(0	)
            __snake_case            =	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'''     ,	)

            __snake_case            =	output.images[0]

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

            __snake_case            =	torch.cuda.max_memory_allocated()
            assert mem_bytes < 10 * 10**9

            __snake_case            =	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()

            __snake_case            =	torch.Generator(device='''cpu'''	).manual_seed(0	)

            __snake_case            =	floats_tensor((1, 3, 64, 64)     ,	rng=random.Random(0	)	).to(__SCREAMING_SNAKE_CASE	)
            __snake_case            =	floats_tensor((1, 3, 256, 256)     ,	rng=random.Random(0	)	).to(__SCREAMING_SNAKE_CASE	)
            __snake_case            =	floats_tensor((1, 3, 256, 256)     ,	rng=random.Random(1	)	).to(__SCREAMING_SNAKE_CASE	)

            __snake_case            =	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'''     ,	)

            __snake_case            =	output.images[0]

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

            __snake_case            =	torch.cuda.max_memory_allocated()
            assert mem_bytes < 4 * 10**9

            __snake_case            =	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()










'''simple docstring'''






import warnings
from functools import wraps
from typing import Callable





def lowerCamelCase__     (		__lowerCamelCase			:				Callable				):







   '''simple docstring'''




   @wraps(__lowerCamelCase				)
   def _inner_fn(*__lowerCamelCase			:				int			,       **__lowerCamelCase			:				Optional[Any]				):
      warnings.warn(
          (f"'{fn.__name__}' is experimental and might be subject to breaking changes in the future.")			,       __lowerCamelCase			,       )
      return fn(*__lowerCamelCase			,       **__lowerCamelCase				)

   return _inner_fn

