





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 							__SCREAMING_SNAKE_CASE       (  _a ,		_a ,		unittest.TestCase							):
							snake_case  :   Dict      =							IFPipeline
							snake_case  :   List[str]      =							TEXT_TO_IMAGE_PARAMS - {"""width""", """height""", """latents"""}
							snake_case  :   str      =							TEXT_TO_IMAGE_BATCH_PARAMS
							snake_case  :   List[str]      =							PipelineTesterMixin.required_optional_params - {"""latents"""}






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






							def 					_lowerCamelCase		(    self	,			__lowerCAmelCase	,			__lowerCAmelCase=0	):
														if str(__lowerCAmelCase	).startswith("""mps"""	):
																					UpperCamelCase__					  =      torch.manual_seed(__lowerCAmelCase	)
														else:
																					UpperCamelCase__					  =      torch.Generator(device=__lowerCAmelCase	).manual_seed(__lowerCAmelCase	)

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

														return inputs






							def 					_lowerCamelCase		(    self	):
														self._test_save_load_optional_components()






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






							def 					_lowerCamelCase		(    self	):
														self._test_save_load_local()






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






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






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






							def 					_lowerCamelCase		(    self	):
														# if

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

														UpperCamelCase__					  =      IFSuperResolutionPipeline.from_pretrained(
														    """DeepFloyd/IF-II-L-v1.0"""	,			variant="""fp16"""	,			torch_dtype=torch.floataa	,			text_encoder=__lowerCAmelCase	,			tokenizer=__lowerCAmelCase	)

														# pre compute text embeddings and remove T5 to save memory

														pipe_a.text_encoder.to("""cuda"""	)

														UpperCamelCase__  ,  UpperCamelCase__					  =      pipe_a.encode_prompt("""anime turtle"""	,			device="""cuda"""	)

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

														UpperCamelCase__					  =      None
														UpperCamelCase__					  =      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(__lowerCAmelCase	,			__lowerCAmelCase	,			__lowerCAmelCase	,			__lowerCAmelCase	)

														pipe_a.remove_all_hooks()
														pipe_a.remove_all_hooks()

														# img2img

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

														pipe_a.remove_all_hooks()
														pipe_a.remove_all_hooks()

														# inpainting

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






							def 					_lowerCamelCase		(    self	,			__lowerCAmelCase	,			__lowerCAmelCase	,			__lowerCAmelCase	,			__lowerCAmelCase	):
														# pipeline 1

														_start_torch_memory_measurement()

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

														UpperCamelCase__					  =      output.images[0]

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

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

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

														# pipeline 2

														_start_torch_memory_measurement()

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

														UpperCamelCase__					  =      floats_tensor((1, 3, 64, 64)	,			rng=random.Random(0	)	).to(__lowerCAmelCase	)

														UpperCamelCase__					  =      pipe_a(
														    prompt_embeds=__lowerCAmelCase	,			negative_prompt_embeds=__lowerCAmelCase	,			image=__lowerCAmelCase	,			generator=__lowerCAmelCase	,			num_inference_steps=2	,			output_type="""np"""	,			)

														UpperCamelCase__					  =      output.images[0]

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

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

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






							def 					_lowerCamelCase		(    self	,			__lowerCAmelCase	,			__lowerCAmelCase	,			__lowerCAmelCase	,			__lowerCAmelCase	):
														# pipeline 1

														_start_torch_memory_measurement()

														UpperCamelCase__					  =      floats_tensor((1, 3, 64, 64)	,			rng=random.Random(0	)	).to(__lowerCAmelCase	)

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

														UpperCamelCase__					  =      pipe_a(
														    prompt_embeds=__lowerCAmelCase	,			negative_prompt_embeds=__lowerCAmelCase	,			image=__lowerCAmelCase	,			num_inference_steps=2	,			generator=__lowerCAmelCase	,			output_type="""np"""	,			)

														UpperCamelCase__					  =      output.images[0]

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

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

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

														# pipeline 2

														_start_torch_memory_measurement()

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

														UpperCamelCase__					  =      floats_tensor((1, 3, 256, 256)	,			rng=random.Random(0	)	).to(__lowerCAmelCase	)
														UpperCamelCase__					  =      floats_tensor((1, 3, 64, 64)	,			rng=random.Random(0	)	).to(__lowerCAmelCase	)

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

														UpperCamelCase__					  =      output.images[0]

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

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

														UpperCamelCase__					  =      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(__lowerCAmelCase	,			__lowerCAmelCase	)







							def 					_lowerCamelCase		(    self	,			__lowerCAmelCase	,			__lowerCAmelCase	,			__lowerCAmelCase	,			__lowerCAmelCase	):
														# pipeline 1

														_start_torch_memory_measurement()

														UpperCamelCase__					  =      floats_tensor((1, 3, 64, 64)	,			rng=random.Random(0	)	).to(__lowerCAmelCase	)
														UpperCamelCase__					  =      floats_tensor((1, 3, 64, 64)	,			rng=random.Random(1	)	).to(__lowerCAmelCase	)

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

														UpperCamelCase__					  =      output.images[0]

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

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

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

														# pipeline 2

														_start_torch_memory_measurement()

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

														UpperCamelCase__					  =      floats_tensor((1, 3, 64, 64)	,			rng=random.Random(0	)	).to(__lowerCAmelCase	)
														UpperCamelCase__					  =      floats_tensor((1, 3, 256, 256)	,			rng=random.Random(0	)	).to(__lowerCAmelCase	)
														UpperCamelCase__					  =      floats_tensor((1, 3, 256, 256)	,			rng=random.Random(1	)	).to(__lowerCAmelCase	)

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

														UpperCamelCase__					  =      output.images[0]

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

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

														UpperCamelCase__					  =      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(__lowerCAmelCase	,			__lowerCAmelCase	)






def      _UpperCamelCase  ():






							"""simple docstring"""





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











from manim import *

class 							__SCREAMING_SNAKE_CASE       (  _a							):






							def 					_lowerCamelCase		(    self	):
														UpperCamelCase__					  =      Rectangle(height=0.5	,			width=0.5	)
														UpperCamelCase__					  =      Rectangle(height=0.25	,			width=0.25	)
														UpperCamelCase__					  =      Rectangle(height=0.46	,			width=0.46	).set_stroke(width=0	)

														UpperCamelCase__					  =      [mem.copy() for i in range(6	)]
														UpperCamelCase__					  =      [mem.copy() for i in range(6	)]
														UpperCamelCase__					  =      VGroup(*__lowerCAmelCase	).arrange(__lowerCAmelCase	,			buff=0	)
														UpperCamelCase__					  =      VGroup(*__lowerCAmelCase	).arrange(__lowerCAmelCase	,			buff=0	)
														UpperCamelCase__					  =      VGroup(__lowerCAmelCase	,			__lowerCAmelCase	).arrange(__lowerCAmelCase	,			buff=0	)
														UpperCamelCase__					  =      Text("""CPU"""	,			font_size=24	)
														UpperCamelCase__					  =      Group(__lowerCAmelCase	,			__lowerCAmelCase	).arrange(__lowerCAmelCase	,			buff=0.5	,			aligned_edge=__lowerCAmelCase	)
														cpu.move_to([-2.5, -0.5, 0]	)
														self.add(__lowerCAmelCase	)

														UpperCamelCase__					  =      [mem.copy() for i in range(4	)]
														UpperCamelCase__					  =      VGroup(*__lowerCAmelCase	).arrange(__lowerCAmelCase	,			buff=0	)
														UpperCamelCase__					  =      Text("""GPU"""	,			font_size=24	)
														UpperCamelCase__					  =      Group(__lowerCAmelCase	,			__lowerCAmelCase	).arrange(__lowerCAmelCase	,			buff=0.5	,			aligned_edge=__lowerCAmelCase	)
														gpu.move_to([-1, -1, 0]	)
														self.add(__lowerCAmelCase	)

														UpperCamelCase__					  =      [mem.copy() for i in range(6	)]
														UpperCamelCase__					  =      VGroup(*__lowerCAmelCase	).arrange(__lowerCAmelCase	,			buff=0	)

														UpperCamelCase__					  =      Text("""Model"""	,			font_size=24	)
														UpperCamelCase__					  =      Group(__lowerCAmelCase	,			__lowerCAmelCase	).arrange(__lowerCAmelCase	,			buff=0.5	,			aligned_edge=__lowerCAmelCase	)
														model.move_to([3, -1.0, 0]	)
														self.add(__lowerCAmelCase	)

														UpperCamelCase__					  =      []
														UpperCamelCase__					  =      []
														UpperCamelCase__					  =      []

														for i, rect in enumerate(__lowerCAmelCase	):
																					rect.set_stroke(__lowerCAmelCase	)

																					UpperCamelCase__					  =      Rectangle(height=0.46 / 4	,			width=0.46 / 3	).set_stroke(width=0.0	).set_fill(__lowerCAmelCase	,			opacity=0.7	)

																					if i == 0:
																												cpu_target.next_to(cpu_left_col_base[0].get_corner(DOWN + LEFT	)	,			buff=0.02	,			direction=__lowerCAmelCase	)
																												cpu_target.set_x(cpu_target.get_x() + 0.1	)
																					elif i == 3:
																												cpu_target.next_to(model_cpu_arr[0]	,			direction=__lowerCAmelCase	,			buff=0.0	)
																					else:
																												cpu_target.next_to(model_cpu_arr[i - 1]	,			direction=__lowerCAmelCase	,			buff=0.0	)
																					self.add(__lowerCAmelCase	)
																					model_cpu_arr.append(__lowerCAmelCase	)

														self.add(*__lowerCAmelCase	,			*__lowerCAmelCase	,			*__lowerCAmelCase	)

														UpperCamelCase__					  =      [mem.copy() for i in range(6	)]
														UpperCamelCase__					  =      VGroup(*__lowerCAmelCase	).arrange(__lowerCAmelCase	,			buff=0	)

														UpperCamelCase__					  =      Text("""Loaded Checkpoint"""	,			font_size=24	)
														UpperCamelCase__					  =      Group(__lowerCAmelCase	,			__lowerCAmelCase	).arrange(__lowerCAmelCase	,			buff=0.5	,			aligned_edge=__lowerCAmelCase	)
														checkpoint.move_to([3, 0.5, 0]	)

														self.add(__lowerCAmelCase	)

														UpperCamelCase__					  =      []
														UpperCamelCase__					  =      []

														for i, rect in enumerate(__lowerCAmelCase	):
																					UpperCamelCase__					  =      fill.copy().set_fill(__lowerCAmelCase	,			opacity=0.7	)
																					target.move_to(__lowerCAmelCase	)
																					ckpt_arr.append(__lowerCAmelCase	)

																					UpperCamelCase__					  =      target.copy()
																					if i < 5:
																												cpu_target.move_to(cpu_left_col_base[i + 1]	)
																					else:
																												cpu_target.move_to(cpu_right_col_base[i - 5]	)
																					ckpt_cpu_arr.append(__lowerCAmelCase	)
														self.add(*__lowerCAmelCase	,			*__lowerCAmelCase	)

														UpperCamelCase__					  =      Square(side_length=2.2	)
														key.move_to([-5, 2, 0]	)

														UpperCamelCase__					  =      MarkupText(
														    f"""<b>Key:</b>\n\n<span fgcolor='{YELLOW}'>●</span> Empty Model"""	,			font_size=18	,			)

														key_text.move_to([-5, 2.4, 0]	)

														self.add(__lowerCAmelCase	,			__lowerCAmelCase	)

														UpperCamelCase__					  =      MarkupText(
														    f"""<span fgcolor='{BLUE}'>●</span> Checkpoint"""	,			font_size=18	,			)

														blue_text.next_to(__lowerCAmelCase	,			DOWN * 2.4	,			aligned_edge=key_text.get_left()	)
														self.add(__lowerCAmelCase	)

														UpperCamelCase__					  =      MarkupText(
														    f"""Based on the passed in configuration, weights are stored in\na variety of np.memmaps on disk or to a particular device."""	,			font_size=24	,			)
														step_a.move_to([2, 2, 0]	)

														UpperCamelCase__					  =      [meta_mem.copy() for i in range(6	)]
														UpperCamelCase__					  =      [meta_mem.copy() for i in range(6	)]
														UpperCamelCase__					  =      VGroup(*__lowerCAmelCase	).arrange(__lowerCAmelCase	,			buff=0	)
														UpperCamelCase__					  =      VGroup(*__lowerCAmelCase	).arrange(__lowerCAmelCase	,			buff=0	)
														UpperCamelCase__					  =      VGroup(__lowerCAmelCase	,			__lowerCAmelCase	).arrange(__lowerCAmelCase	,			buff=0	)
														UpperCamelCase__					  =      Text("""Disk"""	,			font_size=24	)
														UpperCamelCase__					  =      Group(__lowerCAmelCase	,			__lowerCAmelCase	).arrange(__lowerCAmelCase	,			buff=0.5	,			aligned_edge=__lowerCAmelCase	)
														disk.move_to([-4.0, -1.25, 0]	)
														self.play(Write(__lowerCAmelCase	,			run_time=3	)	,			Write(__lowerCAmelCase	,			run_time=1	)	,			Create(__lowerCAmelCase	,			run_time=1	)	)

														UpperCamelCase__					  =      []
														for i, rect in enumerate(__lowerCAmelCase	):
																					UpperCamelCase__					  =      rect.copy()
																					target.generate_target()
																					target.target.move_to(disk_left_col_base[i]	).scale(0.5	)
																					animations.append(MoveToTarget(__lowerCAmelCase	,			run_time=1.5	)	)
														self.play(*__lowerCAmelCase	)

														self.play(FadeOut(__lowerCAmelCase	)	)

														UpperCamelCase__					  =      MarkupText(f"""Then, the checkpoint is removed from memory\nthrough garbage collection."""	,			font_size=24	)
														step_a.move_to([2, 2, 0]	)

														self.play(Write(__lowerCAmelCase	,			run_time=3	)	)

														self.play(
														    FadeOut(__lowerCAmelCase	,			__lowerCAmelCase	,			*__lowerCAmelCase	,			*__lowerCAmelCase	)	,			)

														self.wait()




