
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 			__A( SCREAMING_SNAKE_CASE__  ,   SCREAMING_SNAKE_CASE__  ,   unittest.TestCase			):
				snake_case_						   =    IFPipeline
				snake_case_						   =    TEXT_TO_IMAGE_PARAMS - {'''width''', '''height''', '''latents'''}
				snake_case_						   =    TEXT_TO_IMAGE_BATCH_PARAMS
				snake_case_						   =    PipelineTesterMixin.required_optional_params - {'''latents'''}




				def 				SCREAMING_SNAKE_CASE_		(     self				)    ->      Optional[int]:


								'''simple docstring'''




								return self._get_dummy_components()




				def 				SCREAMING_SNAKE_CASE_		(     self		,     _snake_case		,     _snake_case=0				)    ->      str:


								'''simple docstring'''




								if str(A__				).startswith('''mps'''				):
												__a										=    torch.manual_seed(A__				)
								else:
												__a										=    torch.Generator(device=A__				).manual_seed(A__				)

								__a										=    {
								    '''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				)    ->      Dict:


								'''simple docstring'''




								self._test_save_load_optional_components()




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


								'''simple docstring'''




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




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


								'''simple docstring'''




								self._test_attention_slicing_forward_pass(expected_max_diff=1E-2				)




				def 				SCREAMING_SNAKE_CASE_		(     self				)    ->      str:


								'''simple docstring'''




								self._test_save_load_local()




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


								'''simple docstring'''




								self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=1E-3				)







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




				def 				SCREAMING_SNAKE_CASE_		(     self				)    ->      Tuple:


								'''simple docstring'''




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




				def 				SCREAMING_SNAKE_CASE_		(     self				)    ->      Any:


								'''simple docstring'''




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

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

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

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

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

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

								__a										=    IFInpaintingPipeline(**pipe_a.components				)
								__a										=    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		,     _snake_case		,     _snake_case		,     _snake_case		,     _snake_case				)    ->      str:


								'''simple docstring'''




								_start_torch_memory_measurement()

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

								__a										=    output.images[0]

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

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

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

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

								__a										=    floats_tensor((1, 3, 64, 64)		,     rng=random.Random(0				)				).to(A__				)

								__a										=    pipe_a(
								    prompt_embeds=A__		,     negative_prompt_embeds=A__		,     image=A__		,     generator=A__		,     num_inference_steps=2		,     output_type='''np'''		,     )

								__a										=    output.images[0]

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

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

								__a										=    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		,     _snake_case		,     _snake_case		,     _snake_case		,     _snake_case				)    ->      List[Any]:


								'''simple docstring'''




								_start_torch_memory_measurement()

								__a										=    floats_tensor((1, 3, 64, 64)		,     rng=random.Random(0				)				).to(A__				)

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

								__a										=    pipe_a(
								    prompt_embeds=A__		,     negative_prompt_embeds=A__		,     image=A__		,     num_inference_steps=2		,     generator=A__		,     output_type='''np'''		,     )

								__a										=    output.images[0]

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

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

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

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

								__a										=    floats_tensor((1, 3, 256, 256)		,     rng=random.Random(0				)				).to(A__				)
								__a										=    floats_tensor((1, 3, 64, 64)		,     rng=random.Random(0				)				).to(A__				)

								__a										=    pipe_a(
								    prompt_embeds=A__		,     negative_prompt_embeds=A__		,     image=A__		,     original_image=A__		,     generator=A__		,     num_inference_steps=2		,     output_type='''np'''		,     )

								__a										=    output.images[0]

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

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

								__a										=    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		,     _snake_case		,     _snake_case		,     _snake_case		,     _snake_case				)    ->      Tuple:


								'''simple docstring'''




								_start_torch_memory_measurement()

								__a										=    floats_tensor((1, 3, 64, 64)		,     rng=random.Random(0				)				).to(A__				)
								__a										=    floats_tensor((1, 3, 64, 64)		,     rng=random.Random(1				)				).to(A__				)

								__a										=    torch.Generator(device='''cpu'''				).manual_seed(0				)
								__a										=    pipe_a(
								    prompt_embeds=A__		,     negative_prompt_embeds=A__		,     image=A__		,     mask_image=A__		,     num_inference_steps=2		,     generator=A__		,     output_type='''np'''		,     )

								__a										=    output.images[0]

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

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

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

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

								__a										=    floats_tensor((1, 3, 64, 64)		,     rng=random.Random(0				)				).to(A__				)
								__a										=    floats_tensor((1, 3, 256, 256)		,     rng=random.Random(0				)				).to(A__				)
								__a										=    floats_tensor((1, 3, 256, 256)		,     rng=random.Random(1				)				).to(A__				)

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

								__a										=    output.images[0]

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

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

								__a										=    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			(			)       ->  str:
				torch.cuda.empty_cache()
				torch.cuda.reset_max_memory_allocated()
				torch.cuda.reset_peak_memory_stats()

import random
import unittest

import torch

from diffusers import IFInpaintingPipeline
from diffusers.utils import floats_tensor
from diffusers.utils.import_utils import is_xformers_available
from diffusers.utils.testing_utils import skip_mps, torch_device

from ..pipeline_params import (
    TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS,
    TEXT_GUIDED_IMAGE_INPAINTING_PARAMS,
)
from ..test_pipelines_common import PipelineTesterMixin
from . import IFPipelineTesterMixin







@skip_mps
class 				lowercase		(					SCREAMING_SNAKE_CASE__       ,  SCREAMING_SNAKE_CASE__       ,  unittest.TestCase		):
							lowercase_       :       List[Any]       					=IFInpaintingPipeline
							lowercase_       :       Optional[int]       					=TEXT_GUIDED_IMAGE_INPAINTING_PARAMS - {'''width''', '''height'''}
							lowercase_       :       Any       					=TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS
							lowercase_       :       str       					=PipelineTesterMixin.required_optional_params - {'''latents'''}




							def  A__      (				self):
												return self._get_dummy_components()




							def  A__      (				self   ,A__   ,A__=0):
												if str(A__).startswith('''mps'''):
																	lowercase							=						torch.manual_seed(A__)
												else:
																	lowercase							=						torch.Generator(device=A__).manual_seed(A__)

												lowercase							=						floats_tensor((1, 3, 3_2, 3_2)   ,rng=random.Random(A__)).to(A__)
												lowercase							=						floats_tensor((1, 3, 3_2, 3_2)   ,rng=random.Random(A__)).to(A__)

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

												return inputs




							@unittest.skipIf(
							    torch_device != '''cuda''' or not is_xformers_available()   ,reason='''XFormers attention is only available with CUDA and `xformers` installed'''   ,)
							def  A__      (				self):
												self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=1E-3)




							def  A__      (				self):
												self._test_save_load_optional_components()




							@unittest.skipIf(torch_device != '''cuda'''   ,reason='''float16 requires CUDA''')
							def  A__      (				self):
												# Due to non-determinism in save load of the hf-internal-testing/tiny-random-t5 text encoder
												super().test_save_load_floataa(expected_max_diff=1E-1)




							def  A__      (				self):
												self._test_attention_slicing_forward_pass(expected_max_diff=1E-2)




							def  A__      (				self):
												self._test_save_load_local()





							def  A__      (				self):
												self._test_inference_batch_single_identical(
												    expected_max_diff=1E-2   ,)




