



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





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





						def        _snake_case    (self							,    __magic_name__							,    __magic_name__=0				):


									if str(_a				).startswith("""mps"""				):
												lowerCamelCase__ :      List[Any]								=		torch.manual_seed(_a				)
									else:
												lowerCamelCase__ :      List[Any]								=		torch.Generator(device=_a				).manual_seed(_a				)

									lowerCamelCase__ :      Union[str, Any]								=		{
									    """prompt""": """A painting of a squirrel eating a burger""",
									    """generator""": generator,
									    """num_inference_steps""": 2,
									    """output_type""": """numpy""",
									}

									return inputs





						def        _snake_case    (self				):
									self._test_save_load_optional_components()





						@unittest.skipIf(torch_device != """cuda"""							,    reason="""float16 requires CUDA"""				)
						def        _snake_case    (self				):
									super().test_save_load_floataa(expected_max_diff=1E-1				)





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





						def        _snake_case    (self				):
									self._test_save_load_local()





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







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





						def        _snake_case    (self				):
									super().tearDown()
									gc.collect()
									torch.cuda.empty_cache()





						def        _snake_case    (self				):
									lowerCamelCase__ :      List[str]								=		IFPipeline.from_pretrained("""DeepFloyd/IF-I-XL-v1.0"""							,    variant="""fp16"""							,    torch_dtype=torch.floataa				)

									lowerCamelCase__ :      str								=		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"""				)

									lowerCamelCase__ :      Dict								=		pipe_a.encode_prompt("""anime turtle"""							,    device="""cuda"""				)

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

									lowerCamelCase__ :      List[str]								=		None
									lowerCamelCase__ :      List[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(_a							,    _a							,    _a							,    _a				)

									pipe_a.remove_all_hooks()
									pipe_a.remove_all_hooks()

									# img2img

									lowerCamelCase__ :      List[str]								=		IFImgaImgPipeline(**pipe_a.components				)
									lowerCamelCase__ :      int								=		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

									lowerCamelCase__ :      Dict								=		IFInpaintingPipeline(**pipe_a.components				)
									lowerCamelCase__ :      Union[str, Any]								=		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        _snake_case    (self							,    __magic_name__							,    __magic_name__							,    __magic_name__							,    __magic_name__				):
									_start_torch_memory_measurement()

									lowerCamelCase__ :      Tuple								=		torch.Generator(device="""cpu"""				).manual_seed(0				)
									lowerCamelCase__ :      Optional[int]								=		pipe_a(
									    prompt_embeds=_a							,    negative_prompt_embeds=_a							,    num_inference_steps=2							,    generator=_a							,    output_type="""np"""							,    )

									lowerCamelCase__ :      Optional[int]								=		output.images[0]

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

									lowerCamelCase__ :      List[str]								=		torch.cuda.max_memory_allocated()
									assert mem_bytes < 13 * 10**9

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

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

									lowerCamelCase__ :      Tuple								=		floats_tensor((1, 3, 64, 64)							,    rng=random.Random(0				)				).to(_a				)

									lowerCamelCase__ :      Optional[int]								=		pipe_a(
									    prompt_embeds=_a							,    negative_prompt_embeds=_a							,    image=_a							,    generator=_a							,    num_inference_steps=2							,    output_type="""np"""							,    )

									lowerCamelCase__ :      List[str]								=		output.images[0]

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

									lowerCamelCase__ :      List[str]								=		torch.cuda.max_memory_allocated()
									assert mem_bytes < 4 * 10**9

									lowerCamelCase__ :      Dict								=		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        _snake_case    (self							,    __magic_name__							,    __magic_name__							,    __magic_name__							,    __magic_name__				):
									_start_torch_memory_measurement()

									lowerCamelCase__ :      Dict								=		floats_tensor((1, 3, 64, 64)							,    rng=random.Random(0				)				).to(_a				)

									lowerCamelCase__ :      Union[str, Any]								=		torch.Generator(device="""cpu"""				).manual_seed(0				)

									lowerCamelCase__ :      Dict								=		pipe_a(
									    prompt_embeds=_a							,    negative_prompt_embeds=_a							,    image=_a							,    num_inference_steps=2							,    generator=_a							,    output_type="""np"""							,    )

									lowerCamelCase__ :      List[str]								=		output.images[0]

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

									lowerCamelCase__ :      Tuple								=		torch.cuda.max_memory_allocated()
									assert mem_bytes < 10 * 10**9

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

									lowerCamelCase__ :      Optional[int]								=		torch.Generator(device="""cpu"""				).manual_seed(0				)

									lowerCamelCase__ :      List[str]								=		floats_tensor((1, 3, 256, 256)							,    rng=random.Random(0				)				).to(_a				)
									lowerCamelCase__ :      str								=		floats_tensor((1, 3, 64, 64)							,    rng=random.Random(0				)				).to(_a				)

									lowerCamelCase__ :      List[str]								=		pipe_a(
									    prompt_embeds=_a							,    negative_prompt_embeds=_a							,    image=_a							,    original_image=_a							,    generator=_a							,    num_inference_steps=2							,    output_type="""np"""							,    )

									lowerCamelCase__ :      Any								=		output.images[0]

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

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

									lowerCamelCase__ :      int								=		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        _snake_case    (self							,    __magic_name__							,    __magic_name__							,    __magic_name__							,    __magic_name__				):
									_start_torch_memory_measurement()

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

									lowerCamelCase__ :      Any								=		torch.Generator(device="""cpu"""				).manual_seed(0				)
									lowerCamelCase__ :      Dict								=		pipe_a(
									    prompt_embeds=_a							,    negative_prompt_embeds=_a							,    image=_a							,    mask_image=_a							,    num_inference_steps=2							,    generator=_a							,    output_type="""np"""							,    )

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

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

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

									lowerCamelCase__ :      Optional[Any]								=		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()

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

									lowerCamelCase__ :      int								=		floats_tensor((1, 3, 64, 64)							,    rng=random.Random(0				)				).to(_a				)
									lowerCamelCase__ :      Union[str, Any]								=		floats_tensor((1, 3, 256, 256)							,    rng=random.Random(0				)				).to(_a				)
									lowerCamelCase__ :      Union[str, Any]								=		floats_tensor((1, 3, 256, 256)							,    rng=random.Random(1				)				).to(_a				)

									lowerCamelCase__ :      Optional[int]								=		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"""							,    )

									lowerCamelCase__ :      Any								=		output.images[0]

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

									lowerCamelCase__ :      List[str]								=		torch.cuda.max_memory_allocated()
									assert mem_bytes < 4 * 10**9

									lowerCamelCase__ :      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(_a							,    _a				)







def      _A				()  ->Optional[int]:



			'''simple docstring'''







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








import unittest

from .lib import (
    Matrix,
    Vector,
    axpy,
    square_zero_matrix,
    unit_basis_vector,
    zero_vector,
)
class 				__A					( unittest.TestCase      ):





			def        _snake_case    (self				):
						lowerCamelCase__ :      List[str]								=		Vector([1, 2, 3]				)
						self.assertEqual(x.component(0				)							,    1				)
						self.assertEqual(x.component(2				)							,    3				)
						lowerCamelCase__ :      str								=		Vector()





			def        _snake_case    (self				):
						lowerCamelCase__ :      Dict								=		Vector([0, 0, 0, 0, 0, 1]				)
						self.assertEqual(str(__magic_name__				)							,    """(0,0,0,0,0,1)"""				)





			def        _snake_case    (self				):
						lowerCamelCase__ :      Tuple								=		Vector([1, 2, 3, 4]				)
						self.assertEqual(len(__magic_name__				)							,    4				)





			def        _snake_case    (self				):
						lowerCamelCase__ :      Dict								=		Vector([1, 2]				)
						lowerCamelCase__ :      str								=		Vector([1, 2, 3, 4, 5]				)
						lowerCamelCase__ :      int								=		Vector([0, 0, 0, 0, 0, 0, 0, 0, 0, 0]				)
						lowerCamelCase__ :      int								=		Vector([1, -1, 1, -1, 2, -3, 4, -5]				)
						self.assertAlmostEqual(x.euclidean_length()							,    2.2_36							,    3				)
						self.assertAlmostEqual(y.euclidean_length()							,    7.4_16							,    3				)
						self.assertEqual(z.euclidean_length()							,    0				)
						self.assertAlmostEqual(w.euclidean_length()							,    7.6_16							,    3				)





			def        _snake_case    (self				):
						lowerCamelCase__ :      str								=		Vector([1, 2, 3]				)
						lowerCamelCase__ :      Union[str, Any]								=		Vector([1, 1, 1]				)
						self.assertEqual((x + y).component(0				)							,    2				)
						self.assertEqual((x + y).component(1				)							,    3				)
						self.assertEqual((x + y).component(2				)							,    4				)





			def        _snake_case    (self				):
						lowerCamelCase__ :      Optional[int]								=		Vector([1, 2, 3]				)
						lowerCamelCase__ :      Any								=		Vector([1, 1, 1]				)
						self.assertEqual((x - y).component(0				)							,    0				)
						self.assertEqual((x - y).component(1				)							,    1				)
						self.assertEqual((x - y).component(2				)							,    2				)





			def        _snake_case    (self				):
						lowerCamelCase__ :      Dict								=		Vector([1, 2, 3]				)
						lowerCamelCase__ :      int								=		Vector([2, -1, 4]				)  # for test of dot product
						lowerCamelCase__ :      List[Any]								=		Vector([1, -2, -1]				)
						self.assertEqual(str(x * 3.0				)							,    """(3.0,6.0,9.0)"""				)
						self.assertEqual((a * b)							,    0				)





			def        _snake_case    (self				):
						self.assertEqual(str(zero_vector(10				)				).count("""0"""				)							,    10				)





			def        _snake_case    (self				):
						self.assertEqual(str(unit_basis_vector(3							,    1				)				)							,    """(0,1,0)"""				)





			def        _snake_case    (self				):
						lowerCamelCase__ :      Any								=		Vector([1, 2, 3]				)
						lowerCamelCase__ :      str								=		Vector([1, 0, 1]				)
						self.assertEqual(str(axpy(2							,    __magic_name__							,    __magic_name__				)				)							,    """(3,4,7)"""				)





			def        _snake_case    (self				):
						lowerCamelCase__ :      Optional[Any]								=		Vector([1, 0, 0, 0, 0, 0]				)
						lowerCamelCase__ :      int								=		x.copy()
						self.assertEqual(str(__magic_name__				)							,    str(__magic_name__				)				)





			def        _snake_case    (self				):
						lowerCamelCase__ :      int								=		Vector([1, 0, 0]				)
						x.change_component(0							,    0				)
						x.change_component(1							,    1				)
						self.assertEqual(str(__magic_name__				)							,    """(0,1,0)"""				)





			def        _snake_case    (self				):
						lowerCamelCase__ :      Optional[Any]								=		Matrix([[1, 2, 3], [2, 4, 5], [6, 7, 8]]							,    3							,    3				)
						self.assertEqual("""|1,2,3|\n|2,4,5|\n|6,7,8|\n"""							,    str(__magic_name__				)				)





			def        _snake_case    (self				):
						lowerCamelCase__ :      int								=		Matrix([[1, 2, 3], [2, 4, 5], [6, 7, 8]]							,    3							,    3				)
						lowerCamelCase__ :      int								=		[[-3, -14, -10], [-5, -10, -5], [-2, -1, 0]]
						for x in range(a.height()				):
									for y in range(a.width()				):
												self.assertEqual(minors[x][y]							,    a.minor(__magic_name__							,    __magic_name__				)				)





			def        _snake_case    (self				):
						lowerCamelCase__ :      Any								=		Matrix([[1, 2, 3], [2, 4, 5], [6, 7, 8]]							,    3							,    3				)
						lowerCamelCase__ :      Optional[Any]								=		[[-3, 14, -10], [5, -10, 5], [-2, 1, 0]]
						for x in range(a.height()				):
									for y in range(a.width()				):
												self.assertEqual(cofactors[x][y]							,    a.cofactor(__magic_name__							,    __magic_name__				)				)





			def        _snake_case    (self				):
						lowerCamelCase__ :      Optional[int]								=		Matrix([[1, 2, 3], [2, 4, 5], [6, 7, 8]]							,    3							,    3				)
						self.assertEqual(-5							,    a.determinant()				)





			def        _snake_case    (self				):
						lowerCamelCase__ :      List[str]								=		Matrix([[1, 2, 3], [4, 5, 6], [7, 8, 9]]							,    3							,    3				)
						lowerCamelCase__ :      Optional[int]								=		Vector([1, 2, 3]				)
						self.assertEqual("""(14,32,50)"""							,    str(a * x				)				)
						self.assertEqual("""|2,4,6|\n|8,10,12|\n|14,16,18|\n"""							,    str(a * 2				)				)





			def        _snake_case    (self				):
						lowerCamelCase__ :      str								=		Matrix([[1, 2, 3], [2, 4, 5], [6, 7, 8]]							,    3							,    3				)
						a.change_component(0							,    2							,    5				)
						self.assertEqual("""|1,2,5|\n|2,4,5|\n|6,7,8|\n"""							,    str(__magic_name__				)				)





			def        _snake_case    (self				):
						lowerCamelCase__ :      List[Any]								=		Matrix([[1, 2, 3], [2, 4, 5], [6, 7, 8]]							,    3							,    3				)
						self.assertEqual(7							,    a.component(2							,    1				)							,    0.01				)





			def        _snake_case    (self				):
						lowerCamelCase__ :      Tuple								=		Matrix([[1, 2, 3], [2, 4, 5], [6, 7, 8]]							,    3							,    3				)
						lowerCamelCase__ :      int								=		Matrix([[1, 2, 7], [2, 4, 5], [6, 7, 10]]							,    3							,    3				)
						self.assertEqual("""|2,4,10|\n|4,8,10|\n|12,14,18|\n"""							,    str(a + b				)				)





			def        _snake_case    (self				):
						lowerCamelCase__ :      str								=		Matrix([[1, 2, 3], [2, 4, 5], [6, 7, 8]]							,    3							,    3				)
						lowerCamelCase__ :      Optional[Any]								=		Matrix([[1, 2, 7], [2, 4, 5], [6, 7, 10]]							,    3							,    3				)
						self.assertEqual("""|0,0,-4|\n|0,0,0|\n|0,0,-2|\n"""							,    str(a - b				)				)





			def        _snake_case    (self				):
						self.assertEqual(
						    """|0,0,0,0,0|\n|0,0,0,0,0|\n|0,0,0,0,0|\n|0,0,0,0,0|\n|0,0,0,0,0|\n"""							,    str(square_zero_matrix(5				)				)							,    )


if __name__ == "__main__":
				unittest.main()



