
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_ (			_a					,			_a					,			unittest.TestCase    ):
			a	:			Any					   =							IFPipeline
			a	:			str					   =							TEXT_TO_IMAGE_PARAMS - {'''width''', '''height''', '''latents'''}
			a	:			Dict					   =							TEXT_TO_IMAGE_BATCH_PARAMS
			a	:			Tuple					   =							PipelineTesterMixin.required_optional_params - {'''latents'''}






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






			def    UpperCamelCase_				(      self				,			__UpperCamelCase				,			__UpperCamelCase=0  ):
									if str(__UpperCamelCase  ).startswith("""mps"""  ):
															_lowercase					       =					torch.manual_seed(__UpperCamelCase  )
									else:
															_lowercase					       =					torch.Generator(device=__UpperCamelCase  ).manual_seed(__UpperCamelCase  )

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

									return inputs






			def    UpperCamelCase_				(      self  ):
									self._test_save_load_optional_components()






			@unittest.skipIf(torch_device != """cuda"""				,			reason="""float16 requires CUDA"""  )
			def    UpperCamelCase_				(      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    UpperCamelCase_				(      self  ):
									self._test_attention_slicing_forward_pass(expected_max_diff=1E-2  )






			def    UpperCamelCase_				(      self  ):
									self._test_save_load_local()






			def    UpperCamelCase_				(      self  ):
									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    UpperCamelCase_				(      self  ):
									self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=1E-3  )







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






			def    UpperCamelCase_				(      self  ):
									# clean up the VRAM after each test
									super().tearDown()
									gc.collect()
									torch.cuda.empty_cache()






			def    UpperCamelCase_				(      self  ):
									# if

									_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=__UpperCamelCase				,			tokenizer=__UpperCamelCase  )

									# 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(__UpperCamelCase				,			__UpperCamelCase				,			__UpperCamelCase				,			__UpperCamelCase  )

									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(__UpperCamelCase				,			__UpperCamelCase				,			__UpperCamelCase				,			__UpperCamelCase  )

									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(__UpperCamelCase				,			__UpperCamelCase				,			__UpperCamelCase				,			__UpperCamelCase  )






			def    UpperCamelCase_				(      self				,			__UpperCamelCase				,			__UpperCamelCase				,			__UpperCamelCase				,			__UpperCamelCase  ):
									# pipeline 1

									_start_torch_memory_measurement()

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

									_lowercase					       =					output.images[0]

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

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

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

									# pipeline 2

									_start_torch_memory_measurement()

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

									_lowercase					       =					floats_tensor((1, 3, 64, 64)				,			rng=random.Random(0  )  ).to(__UpperCamelCase  )

									_lowercase					       =					pipe_a(
									    prompt_embeds=__UpperCamelCase				,			negative_prompt_embeds=__UpperCamelCase				,			image=__UpperCamelCase				,			generator=__UpperCamelCase				,			num_inference_steps=2				,			output_type="""np"""				,			)

									_lowercase					       =					output.images[0]

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

									_lowercase					       =					torch.cuda.max_memory_allocated()
									assert mem_bytes < 4 * 10**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(__UpperCamelCase				,			__UpperCamelCase  )






			def    UpperCamelCase_				(      self				,			__UpperCamelCase				,			__UpperCamelCase				,			__UpperCamelCase				,			__UpperCamelCase  ):
									# pipeline 1

									_start_torch_memory_measurement()

									_lowercase					       =					floats_tensor((1, 3, 64, 64)				,			rng=random.Random(0  )  ).to(__UpperCamelCase  )

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

									_lowercase					       =					pipe_a(
									    prompt_embeds=__UpperCamelCase				,			negative_prompt_embeds=__UpperCamelCase				,			image=__UpperCamelCase				,			num_inference_steps=2				,			generator=__UpperCamelCase				,			output_type="""np"""				,			)

									_lowercase					       =					output.images[0]

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

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

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

									# pipeline 2

									_start_torch_memory_measurement()

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

									_lowercase					       =					floats_tensor((1, 3, 256, 256)				,			rng=random.Random(0  )  ).to(__UpperCamelCase  )
									_lowercase					       =					floats_tensor((1, 3, 64, 64)				,			rng=random.Random(0  )  ).to(__UpperCamelCase  )

									_lowercase					       =					pipe_a(
									    prompt_embeds=__UpperCamelCase				,			negative_prompt_embeds=__UpperCamelCase				,			image=__UpperCamelCase				,			original_image=__UpperCamelCase				,			generator=__UpperCamelCase				,			num_inference_steps=2				,			output_type="""np"""				,			)

									_lowercase					       =					output.images[0]

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

									_lowercase					       =					torch.cuda.max_memory_allocated()
									assert mem_bytes < 4 * 10**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(__UpperCamelCase				,			__UpperCamelCase  )






			def    UpperCamelCase_				(      self				,			__UpperCamelCase				,			__UpperCamelCase				,			__UpperCamelCase				,			__UpperCamelCase  ):
									# pipeline 1

									_start_torch_memory_measurement()

									_lowercase					       =					floats_tensor((1, 3, 64, 64)				,			rng=random.Random(0  )  ).to(__UpperCamelCase  )
									_lowercase					       =					floats_tensor((1, 3, 64, 64)				,			rng=random.Random(1  )  ).to(__UpperCamelCase  )

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

									_lowercase					       =					output.images[0]

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

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

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

									# pipeline 2

									_start_torch_memory_measurement()

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

									_lowercase					       =					floats_tensor((1, 3, 64, 64)				,			rng=random.Random(0  )  ).to(__UpperCamelCase  )
									_lowercase					       =					floats_tensor((1, 3, 256, 256)				,			rng=random.Random(0  )  ).to(__UpperCamelCase  )
									_lowercase					       =					floats_tensor((1, 3, 256, 256)				,			rng=random.Random(1  )  ).to(__UpperCamelCase  )

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

									_lowercase					       =					output.images[0]

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

									_lowercase					       =					torch.cuda.max_memory_allocated()
									assert mem_bytes < 4 * 10**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(__UpperCamelCase				,			__UpperCamelCase  )







def 	UpperCamelCase__  (     )  ->     List[Any]:
						torch.cuda.empty_cache()
						torch.cuda.reset_max_memory_allocated()
						torch.cuda.reset_peak_memory_stats()

from datasets.utils.patching import _PatchedModuleObj, patch_submodule

from . import _test_patching




def 	UpperCamelCase__  (     )  ->     List[str]:
						import os as original_os
						from os import path as original_path
						from os import rename as original_rename
						from os.path import dirname as original_dirname
						from os.path import join as original_join

						assert _test_patching.os is original_os
						assert _test_patching.path is original_path
						assert _test_patching.join is original_join

						assert _test_patching.renamed_os is original_os
						assert _test_patching.renamed_path is original_path
						assert _test_patching.renamed_join is original_join

						_lowercase					       =					"""__test_patch_submodule_mock__"""
						with patch_submodule(_test_patching		,   """os.path.join"""		,   SCREAMING_SNAKE_CASE_       ):
												# Every way to access os.path.join must be patched, and the rest must stay untouched

												# check os.path.join
												assert isinstance(_test_patching.os		,   _PatchedModuleObj       )
												assert isinstance(_test_patching.os.path		,   _PatchedModuleObj       )
												assert _test_patching.os.path.join is mock

												# check path.join
												assert isinstance(_test_patching.path		,   _PatchedModuleObj       )
												assert _test_patching.path.join is mock

												# check join
												assert _test_patching.join is mock

												# check that the other attributes are untouched
												assert _test_patching.os.rename is original_rename
												assert _test_patching.path.dirname is original_dirname
												assert _test_patching.os.path.dirname is original_dirname

												# Even renamed modules or objects must be patched

												# check renamed_os.path.join
												assert isinstance(_test_patching.renamed_os		,   _PatchedModuleObj       )
												assert isinstance(_test_patching.renamed_os.path		,   _PatchedModuleObj       )
												assert _test_patching.renamed_os.path.join is mock

												# check renamed_path.join
												assert isinstance(_test_patching.renamed_path		,   _PatchedModuleObj       )
												assert _test_patching.renamed_path.join is mock

												# check renamed_join
												assert _test_patching.renamed_join is mock

												# check that the other attributes are untouched
												assert _test_patching.renamed_os.rename is original_rename
												assert _test_patching.renamed_path.dirname is original_dirname
												assert _test_patching.renamed_os.path.dirname is original_dirname

						# check that everthing is back to normal when the patch is over

						assert _test_patching.os is original_os
						assert _test_patching.path is original_path
						assert _test_patching.join is original_join

						assert _test_patching.renamed_os is original_os
						assert _test_patching.renamed_path is original_path
						assert _test_patching.renamed_join is original_join




def 	UpperCamelCase__  (     )  ->     int:
						assert _test_patching.open is open

						_lowercase					       =					"""__test_patch_submodule_builtin_mock__"""
						# _test_patching has "open" in its globals
						assert _test_patching.open is open
						with patch_submodule(_test_patching		,   """open"""		,   SCREAMING_SNAKE_CASE_       ):
												assert _test_patching.open is mock

						# check that everthing is back to normal when the patch is over

						assert _test_patching.open is open




def 	UpperCamelCase__  (     )  ->     List[Any]:
						# pandas.read_csv is not present in _test_patching
						_lowercase					       =					"""__test_patch_submodule_missing_mock__"""
						with patch_submodule(_test_patching		,   """pandas.read_csv"""		,   SCREAMING_SNAKE_CASE_       ):
												pass




def 	UpperCamelCase__  (     )  ->     Tuple:
						# builtin should always be mocked even if they're not in the globals
						# in case they're loaded at one point
						_lowercase					       =					"""__test_patch_submodule_missing_builtin_mock__"""
						# _test_patching doesn't have "len" in its globals
						assert getattr(_test_patching		,   """len"""		,   SCREAMING_SNAKE_CASE_       ) is None
						with patch_submodule(_test_patching		,   """len"""		,   SCREAMING_SNAKE_CASE_       ):
												assert _test_patching.len is mock
						assert _test_patching.len is len




def 	UpperCamelCase__  (     )  ->     List[str]:
						_lowercase					       =					"""__test_patch_submodule_start_and_stop_mock__"""
						_lowercase					       =					patch_submodule(_test_patching		,   """open"""		,   SCREAMING_SNAKE_CASE_       )
						assert _test_patching.open is open
						patch.start()
						assert _test_patching.open is mock
						patch.stop()
						assert _test_patching.open is open




def 	UpperCamelCase__  (     )  ->     Union[str, Any]:
						from os import rename as original_rename
						from os.path import dirname as original_dirname
						from os.path import join as original_join

						_lowercase					       =					"""__test_patch_submodule_successive_join__"""
						_lowercase					       =					"""__test_patch_submodule_successive_dirname__"""
						_lowercase					       =					"""__test_patch_submodule_successive_rename__"""
						assert _test_patching.os.path.join is original_join
						assert _test_patching.os.path.dirname is original_dirname
						assert _test_patching.os.rename is original_rename

						with patch_submodule(_test_patching		,   """os.path.join"""		,   SCREAMING_SNAKE_CASE_       ):
												with patch_submodule(_test_patching		,   """os.rename"""		,   SCREAMING_SNAKE_CASE_       ):
																		with patch_submodule(_test_patching		,   """os.path.dirname"""		,   SCREAMING_SNAKE_CASE_       ):
																								assert _test_patching.os.path.join is mock_join
																								assert _test_patching.os.path.dirname is mock_dirname
																								assert _test_patching.os.rename is mock_rename

    # try another order
						with patch_submodule(_test_patching		,   """os.rename"""		,   SCREAMING_SNAKE_CASE_       ):
												with patch_submodule(_test_patching		,   """os.path.join"""		,   SCREAMING_SNAKE_CASE_       ):
																		with patch_submodule(_test_patching		,   """os.path.dirname"""		,   SCREAMING_SNAKE_CASE_       ):
																								assert _test_patching.os.path.join is mock_join
																								assert _test_patching.os.path.dirname is mock_dirname
																								assert _test_patching.os.rename is mock_rename

						assert _test_patching.os.path.join is original_join
						assert _test_patching.os.path.dirname is original_dirname
						assert _test_patching.os.rename is original_rename




def 	UpperCamelCase__  (     )  ->     List[Any]:
						_lowercase					       =					"""__test_patch_submodule_doesnt_exist_mock__"""
						with patch_submodule(_test_patching		,   """__module_that_doesn_exist__.__attribute_that_doesn_exist__"""		,   SCREAMING_SNAKE_CASE_       ):
												pass
						with patch_submodule(_test_patching		,   """os.__attribute_that_doesn_exist__"""		,   SCREAMING_SNAKE_CASE_       ):
												pass