


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







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

       '''simple docstring'''




       return self._get_dummy_components()







   def 				lowerCamelCase	(			self		,			UpperCamelCase__		,			UpperCamelCase__=0		)			->	Dict:

       '''simple docstring'''




       if str(UpperCamelCase__		).startswith("mps"		):
           snake_case					:   str         =					torch.manual_seed(UpperCamelCase__		)
       else:
           snake_case					:   Optional[int]         =					torch.Generator(device=UpperCamelCase__		).manual_seed(UpperCamelCase__		)

       snake_case					:   Any         =					{
           "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		)			->	Union[str, Any]:

       '''simple docstring'''




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







   def 				lowerCamelCase	(			self		)			->	Tuple:

       '''simple docstring'''




       self._test_attention_slicing_forward_pass(expected_max_diff=1e-2		)







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

       '''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		)			->	Union[str, Any]:

       '''simple docstring'''




       self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=1e-3		)







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







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

       '''simple docstring'''




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







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

       '''simple docstring'''




       snake_case					:   int         =					IFPipeline.from_pretrained("DeepFloyd/IF-I-XL-v1.0"		,			variant="fp16"		,			torch_dtype=torch.floataa		)

       snake_case					:   Optional[Any]         =					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"		)

       snake_case      ,snake_case					:   List[str]         =					pipe_a.encode_prompt("anime turtle"		,			device="cuda"		)

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

       snake_case					:   Union[str, Any]         =					None
       snake_case					:   str         =					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

       snake_case					:   Any         =					IFImgaImgPipeline(**pipe_a.components		)
       snake_case					:   Tuple         =					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

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







   def 				lowerCamelCase	(			self		,			UpperCamelCase__		,			UpperCamelCase__		,			UpperCamelCase__		,			UpperCamelCase__		)			->	Any:

       '''simple docstring'''




       _start_torch_memory_measurement()

       snake_case					:   int         =					torch.Generator(device="cpu"		).manual_seed(0		)
       snake_case					:   List[Any]         =					pipe_a(
           prompt_embeds=UpperCamelCase__		,			negative_prompt_embeds=UpperCamelCase__		,			num_inference_steps=2		,			generator=UpperCamelCase__		,			output_type="np"		,			)

       snake_case					:   str         =					output.images[0]

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

       snake_case					:   int         =					torch.cuda.max_memory_allocated()
       assert mem_bytes < 13 * 10**9

       snake_case					:   str         =					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()

       snake_case					:   List[Any]         =					torch.Generator(device="cpu"		).manual_seed(0		)

       snake_case					:   int         =					floats_tensor((1, 3, 64, 64)		,			rng=random.Random(0		)		).to(UpperCamelCase__		)

       snake_case					:   Union[str, Any]         =					pipe_a(
           prompt_embeds=UpperCamelCase__		,			negative_prompt_embeds=UpperCamelCase__		,			image=UpperCamelCase__		,			generator=UpperCamelCase__		,			num_inference_steps=2		,			output_type="np"		,			)

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

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

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

       snake_case					:   Optional[Any]         =					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 				lowerCamelCase	(			self		,			UpperCamelCase__		,			UpperCamelCase__		,			UpperCamelCase__		,			UpperCamelCase__		)			->	Tuple:

       '''simple docstring'''




       _start_torch_memory_measurement()

       snake_case					:   List[Any]         =					floats_tensor((1, 3, 64, 64)		,			rng=random.Random(0		)		).to(UpperCamelCase__		)

       snake_case					:   Dict         =					torch.Generator(device="cpu"		).manual_seed(0		)

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

       snake_case					:   str         =					output.images[0]

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

       snake_case					:   Optional[int]         =					torch.cuda.max_memory_allocated()
       assert mem_bytes < 10 * 10**9

       snake_case					:   List[Any]         =					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()

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

       snake_case					:   Dict         =					floats_tensor((1, 3, 256, 256)		,			rng=random.Random(0		)		).to(UpperCamelCase__		)
       snake_case					:   Optional[int]         =					floats_tensor((1, 3, 64, 64)		,			rng=random.Random(0		)		).to(UpperCamelCase__		)

       snake_case					:   Any         =					pipe_a(
           prompt_embeds=UpperCamelCase__		,			negative_prompt_embeds=UpperCamelCase__		,			image=UpperCamelCase__		,			original_image=UpperCamelCase__		,			generator=UpperCamelCase__		,			num_inference_steps=2		,			output_type="np"		,			)

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

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

       snake_case					:   int         =					torch.cuda.max_memory_allocated()
       assert mem_bytes < 4 * 10**9

       snake_case					:   Dict         =					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 				lowerCamelCase	(			self		,			UpperCamelCase__		,			UpperCamelCase__		,			UpperCamelCase__		,			UpperCamelCase__		)			->	List[str]:

       '''simple docstring'''




       _start_torch_memory_measurement()

       snake_case					:   List[Any]         =					floats_tensor((1, 3, 64, 64)		,			rng=random.Random(0		)		).to(UpperCamelCase__		)
       snake_case					:   Union[str, Any]         =					floats_tensor((1, 3, 64, 64)		,			rng=random.Random(1		)		).to(UpperCamelCase__		)

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

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

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

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

       snake_case					:   Tuple         =					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()

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

       snake_case					:   Union[str, Any]         =					floats_tensor((1, 3, 64, 64)		,			rng=random.Random(0		)		).to(UpperCamelCase__		)
       snake_case					:   Dict         =					floats_tensor((1, 3, 256, 256)		,			rng=random.Random(0		)		).to(UpperCamelCase__		)
       snake_case					:   List[str]         =					floats_tensor((1, 3, 256, 256)		,			rng=random.Random(1		)		).to(UpperCamelCase__		)

       snake_case					:   List[str]         =					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"		,			)

       snake_case					:   int         =					output.images[0]

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

       snake_case					:   int         =					torch.cuda.max_memory_allocated()
       assert mem_bytes < 4 * 10**9

       snake_case					:   int         =					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 						__lowerCAmelCase		(		)			->		str:


    """simple docstring"""





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









"""simple docstring"""



# this script reports modified .py files under the desired list of top-level sub-dirs passed as a list of arguments, e.g.:
#   python ./utils/get_modified_files.py utils src tests examples
#
# it uses git to find the forking point and which files were modified - i.e. files not under git won't be considered
# since the output of this script is fed into Makefile commands it doesn't print a newline after the results

import re
import subprocess
import sys


__snake_case         =					subprocess.check_output("""git merge-base main HEAD""".split()).decode("""utf-8""")
__snake_case         =					(
    subprocess.check_output(F'''git diff --diff-filter=d --name-only {fork_point_sha}'''.split()).decode("""utf-8""").split()
)

__snake_case         =					"""|""".join(sys.argv[1:])
__snake_case         =					re.compile(RF'''^({joined_dirs}).*?\.py$''')

__snake_case         =					[x for x in modified_files if regex.match(x)]
print(""" """.join(relevant_modified_files), end="""""")





