
import gc
import random
import unittest

import numpy as np
import torch
from PIL import Image
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    EulerAncestralDiscreteScheduler,
    LMSDiscreteScheduler,
    PNDMScheduler,
    StableDiffusionInstructPixaPixPipeline,
    UNetaDConditionModel,
)
from diffusers.image_processor import VaeImageProcessor
from diffusers.utils import floats_tensor, load_image, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..pipeline_params import (
    IMAGE_TO_IMAGE_IMAGE_PARAMS,
    TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS,
    TEXT_GUIDED_IMAGE_VARIATION_PARAMS,
)
from ..test_pipelines_common import PipelineKarrasSchedulerTesterMixin, PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()



class 			lowerCAmelCase_		(a__				,      a__				,      a__				,      unittest.TestCase						):


				"""simple docstring"""



				__UpperCamelCase   :		Union[str, Any]       =     StableDiffusionInstructPixaPixPipeline
				__UpperCamelCase   :		Optional[Any]       =     TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {'''height''', '''width''', '''cross_attention_kwargs'''}
				__UpperCamelCase   :		Dict       =     TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS
				__UpperCamelCase   :		Tuple       =     IMAGE_TO_IMAGE_IMAGE_PARAMS
				__UpperCamelCase   :		Dict       =     IMAGE_TO_IMAGE_IMAGE_PARAMS



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



											"""simple docstring"""




											torch.manual_seed(0					)
											SCREAMING_SNAKE_CASE__		: Any     =     UNetaDConditionModel(
											    block_out_channels=(32, 64)     , layers_per_block=2     , sample_size=32     , in_channels=8     , out_channels=4     , down_block_types=("""DownBlock2D""", """CrossAttnDownBlock2D""")     , up_block_types=("""CrossAttnUpBlock2D""", """UpBlock2D""")     , cross_attention_dim=32     , )
											SCREAMING_SNAKE_CASE__		: Union[str, Any]     =     PNDMScheduler(skip_prk_steps=SCREAMING_SNAKE_CASE__					)
											torch.manual_seed(0					)
											SCREAMING_SNAKE_CASE__		: List[Any]     =     AutoencoderKL(
											    block_out_channels=[32, 64]     , in_channels=3     , out_channels=3     , down_block_types=["""DownEncoderBlock2D""", """DownEncoderBlock2D"""]     , up_block_types=["""UpDecoderBlock2D""", """UpDecoderBlock2D"""]     , latent_channels=4     , )
											torch.manual_seed(0					)
											SCREAMING_SNAKE_CASE__		: str     =     CLIPTextConfig(
											    bos_token_id=0     , eos_token_id=2     , hidden_size=32     , intermediate_size=37     , layer_norm_eps=1E-05     , num_attention_heads=4     , num_hidden_layers=5     , pad_token_id=1     , vocab_size=10_00     , )
											SCREAMING_SNAKE_CASE__		: List[Any]     =     CLIPTextModel(SCREAMING_SNAKE_CASE__					)
											SCREAMING_SNAKE_CASE__		: List[Any]     =     CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""					)

											SCREAMING_SNAKE_CASE__		: Any     =     {
											    """unet""": unet,
											    """scheduler""": scheduler,
											    """vae""": vae,
											    """text_encoder""": text_encoder,
											    """tokenizer""": tokenizer,
											    """safety_checker""": None,
											    """feature_extractor""": None,
											}
											return components



				def 			__magic_name__					(self     , SCREAMING_SNAKE_CASE__     , SCREAMING_SNAKE_CASE__=0					)      ->	Union[str, Any]:



											"""simple docstring"""




											SCREAMING_SNAKE_CASE__		: Union[str, Any]     =     floats_tensor((1, 3, 32, 32)     , rng=random.Random(SCREAMING_SNAKE_CASE__					)					).to(SCREAMING_SNAKE_CASE__					)
											SCREAMING_SNAKE_CASE__		: Optional[Any]     =     image.cpu().permute(0     , 2     , 3     , 1					)[0]
											SCREAMING_SNAKE_CASE__		: Union[str, Any]     =     Image.fromarray(np.uinta(SCREAMING_SNAKE_CASE__					)					).convert("""RGB"""					)
											if str(SCREAMING_SNAKE_CASE__					).startswith("""mps"""					):
																		SCREAMING_SNAKE_CASE__		: Optional[Any]     =     torch.manual_seed(SCREAMING_SNAKE_CASE__					)
											else:
																		SCREAMING_SNAKE_CASE__		: Any     =     torch.Generator(device=SCREAMING_SNAKE_CASE__					).manual_seed(SCREAMING_SNAKE_CASE__					)
											SCREAMING_SNAKE_CASE__		: Optional[int]     =     {
											    """prompt""": """A painting of a squirrel eating a burger""",
											    """image""": image,
											    """generator""": generator,
											    """num_inference_steps""": 2,
											    """guidance_scale""": 6.0,
											    """image_guidance_scale""": 1,
											    """output_type""": """numpy""",
											}
											return inputs



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



											"""simple docstring"""




											SCREAMING_SNAKE_CASE__		: Any     =     """cpu"""  # ensure determinism for the device-dependent torch.Generator
											SCREAMING_SNAKE_CASE__		: List[str]     =     self.get_dummy_components()
											SCREAMING_SNAKE_CASE__		: int     =     StableDiffusionInstructPixaPixPipeline(**SCREAMING_SNAKE_CASE__					)
											SCREAMING_SNAKE_CASE__		: int     =     sd_pipe.to(SCREAMING_SNAKE_CASE__					)
											sd_pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE__					)

											SCREAMING_SNAKE_CASE__		: Union[str, Any]     =     self.get_dummy_inputs(SCREAMING_SNAKE_CASE__					)
											SCREAMING_SNAKE_CASE__		: Optional[Any]     =     sd_pipe(**SCREAMING_SNAKE_CASE__					).images
											SCREAMING_SNAKE_CASE__		: int     =     image[0, -3:, -3:, -1]
											assert image.shape == (1, 32, 32, 3)
											SCREAMING_SNAKE_CASE__		: Dict     =     np.array([0.7526, 0.3750, 0.4547, 0.6117, 0.5866, 0.5016, 0.4327, 0.5642, 0.4815]					)

											assert np.abs(image_slice.flatten() - expected_slice					).max() < 1E-3



				def 			__magic_name__					(self					)      ->	Any:



											"""simple docstring"""




											SCREAMING_SNAKE_CASE__		: Union[str, Any]     =     """cpu"""  # ensure determinism for the device-dependent torch.Generator
											SCREAMING_SNAKE_CASE__		: int     =     self.get_dummy_components()
											SCREAMING_SNAKE_CASE__		: Union[str, Any]     =     StableDiffusionInstructPixaPixPipeline(**SCREAMING_SNAKE_CASE__					)
											SCREAMING_SNAKE_CASE__		: List[str]     =     sd_pipe.to(SCREAMING_SNAKE_CASE__					)
											sd_pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE__					)

											SCREAMING_SNAKE_CASE__		: Union[str, Any]     =     self.get_dummy_inputs(SCREAMING_SNAKE_CASE__					)
											SCREAMING_SNAKE_CASE__		: Optional[int]     =     """french fries"""
											SCREAMING_SNAKE_CASE__		: int     =     sd_pipe(**SCREAMING_SNAKE_CASE__     , negative_prompt=SCREAMING_SNAKE_CASE__					)
											SCREAMING_SNAKE_CASE__		: Tuple     =     output.images
											SCREAMING_SNAKE_CASE__		: Optional[Any]     =     image[0, -3:, -3:, -1]

											assert image.shape == (1, 32, 32, 3)
											SCREAMING_SNAKE_CASE__		: Tuple     =     np.array([0.7511, 0.3642, 0.4553, 0.6236, 0.5797, 0.5013, 0.4343, 0.5611, 0.4831]					)

											assert np.abs(image_slice.flatten() - expected_slice					).max() < 1E-3



				def 			__magic_name__					(self					)      ->	str:



											"""simple docstring"""




											SCREAMING_SNAKE_CASE__		: Optional[int]     =     """cpu"""  # ensure determinism for the device-dependent torch.Generator
											SCREAMING_SNAKE_CASE__		: Optional[Any]     =     self.get_dummy_components()
											SCREAMING_SNAKE_CASE__		: Any     =     StableDiffusionInstructPixaPixPipeline(**SCREAMING_SNAKE_CASE__					)
											SCREAMING_SNAKE_CASE__		: int     =     sd_pipe.to(SCREAMING_SNAKE_CASE__					)
											sd_pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE__					)

											SCREAMING_SNAKE_CASE__		: str     =     self.get_dummy_inputs(SCREAMING_SNAKE_CASE__					)
											SCREAMING_SNAKE_CASE__		: str     =     [inputs["""prompt"""]] * 2

											SCREAMING_SNAKE_CASE__		: List[str]     =     np.array(inputs["""image"""]					).astype(np.floataa					) / 255.0
											SCREAMING_SNAKE_CASE__		: Union[str, Any]     =     torch.from_numpy(SCREAMING_SNAKE_CASE__					).unsqueeze(0					).to(SCREAMING_SNAKE_CASE__					)
											SCREAMING_SNAKE_CASE__		: List[str]     =     image / 2 + 0.5
											SCREAMING_SNAKE_CASE__		: Optional[int]     =     image.permute(0     , 3     , 1     , 2					)
											SCREAMING_SNAKE_CASE__		: Optional[int]     =     image.repeat(2     , 1     , 1     , 1					)

											SCREAMING_SNAKE_CASE__		: List[Any]     =     sd_pipe(**SCREAMING_SNAKE_CASE__					).images
											SCREAMING_SNAKE_CASE__		: Tuple     =     image[-1, -3:, -3:, -1]

											assert image.shape == (2, 32, 32, 3)
											SCREAMING_SNAKE_CASE__		: str     =     np.array([0.5812, 0.5748, 0.5222, 0.5908, 0.5695, 0.7174, 0.6804, 0.5523, 0.5579]					)

											assert np.abs(image_slice.flatten() - expected_slice					).max() < 1E-3



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



											"""simple docstring"""




											SCREAMING_SNAKE_CASE__		: List[str]     =     """cpu"""  # ensure determinism for the device-dependent torch.Generator
											SCREAMING_SNAKE_CASE__		: Optional[int]     =     self.get_dummy_components()
											SCREAMING_SNAKE_CASE__		: str     =     EulerAncestralDiscreteScheduler(
											    beta_start=0.00085     , beta_end=0.012     , beta_schedule="""scaled_linear"""					)
											SCREAMING_SNAKE_CASE__		: Any     =     StableDiffusionInstructPixaPixPipeline(**SCREAMING_SNAKE_CASE__					)
											SCREAMING_SNAKE_CASE__		: Union[str, Any]     =     sd_pipe.to(SCREAMING_SNAKE_CASE__					)
											sd_pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE__					)

											SCREAMING_SNAKE_CASE__		: str     =     self.get_dummy_inputs(SCREAMING_SNAKE_CASE__					)
											SCREAMING_SNAKE_CASE__		: Optional[Any]     =     sd_pipe(**SCREAMING_SNAKE_CASE__					).images
											SCREAMING_SNAKE_CASE__		: Any     =     image[0, -3:, -3:, -1]

											SCREAMING_SNAKE_CASE__		: Optional[Any]     =     [round(SCREAMING_SNAKE_CASE__     , 4					) for x in image_slice.flatten().tolist()]
											print(""",""".join([str(SCREAMING_SNAKE_CASE__					) for x in slice]					)					)

											assert image.shape == (1, 32, 32, 3)
											SCREAMING_SNAKE_CASE__		: List[str]     =     np.array([0.7417, 0.3842, 0.4732, 0.5776, 0.5891, 0.5139, 0.4052, 0.5673, 0.4986]					)

											assert np.abs(image_slice.flatten() - expected_slice					).max() < 1E-3



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



											"""simple docstring"""




											super().test_inference_batch_single_identical(expected_max_diff=3E-3					)



				def 			__magic_name__					(self					)      ->	str:



											"""simple docstring"""




											SCREAMING_SNAKE_CASE__		: Optional[int]     =     self.get_dummy_components()
											SCREAMING_SNAKE_CASE__		: Tuple     =     StableDiffusionInstructPixaPixPipeline(**SCREAMING_SNAKE_CASE__					)
											SCREAMING_SNAKE_CASE__		: int     =     VaeImageProcessor(do_resize=SCREAMING_SNAKE_CASE__     , do_normalize=SCREAMING_SNAKE_CASE__					)
											SCREAMING_SNAKE_CASE__		: List[str]     =     pipe.to(SCREAMING_SNAKE_CASE__					)
											pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE__					)

											SCREAMING_SNAKE_CASE__		: str     =     pipe(**self.get_dummy_inputs_by_type(SCREAMING_SNAKE_CASE__     , input_image_type="""pt"""					)					)[0]

											SCREAMING_SNAKE_CASE__		: int     =     components["""vae"""]
											SCREAMING_SNAKE_CASE__		: List[Any]     =     self.get_dummy_inputs_by_type(SCREAMING_SNAKE_CASE__     , input_image_type="""pt"""					)

											for image_param in self.image_latents_params:
																		if image_param in inputs.keys():
																									SCREAMING_SNAKE_CASE__		: Optional[Any]     =     vae.encode(inputs[image_param]					).latent_dist.mode()

											SCREAMING_SNAKE_CASE__		: Any     =     pipe(**SCREAMING_SNAKE_CASE__					)[0]

											SCREAMING_SNAKE_CASE__		: Any     =     np.abs(out - out_latents_inputs					).max()
											self.assertLess(SCREAMING_SNAKE_CASE__     , 1E-4     , """passing latents as image input generate different result from passing image"""					)






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


				"""simple docstring"""



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



											"""simple docstring"""




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



				def 			__magic_name__					(self     , SCREAMING_SNAKE_CASE__=0					)      ->	Optional[int]:



											"""simple docstring"""




											SCREAMING_SNAKE_CASE__		: Optional[Any]     =     torch.manual_seed(SCREAMING_SNAKE_CASE__					)
											SCREAMING_SNAKE_CASE__		: Optional[Any]     =     load_image(
											    """https://huggingface.co/datasets/diffusers/test-arrays/resolve/main/stable_diffusion_pix2pix/example.jpg"""					)
											SCREAMING_SNAKE_CASE__		: str     =     {
											    """prompt""": """turn him into a cyborg""",
											    """image""": image,
											    """generator""": generator,
											    """num_inference_steps""": 3,
											    """guidance_scale""": 7.5,
											    """image_guidance_scale""": 1.0,
											    """output_type""": """numpy""",
											}
											return inputs



				def 			__magic_name__					(self					)      ->	str:



											"""simple docstring"""




											SCREAMING_SNAKE_CASE__		: List[str]     =     StableDiffusionInstructPixaPixPipeline.from_pretrained(
											    """timbrooks/instruct-pix2pix"""     , safety_checker=SCREAMING_SNAKE_CASE__					)
											pipe.to(SCREAMING_SNAKE_CASE__					)
											pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE__					)
											pipe.enable_attention_slicing()

											SCREAMING_SNAKE_CASE__		: Tuple     =     self.get_inputs()
											SCREAMING_SNAKE_CASE__		: Optional[Any]     =     pipe(**SCREAMING_SNAKE_CASE__					).images
											SCREAMING_SNAKE_CASE__		: Tuple     =     image[0, -3:, -3:, -1].flatten()

											assert image.shape == (1, 5_12, 5_12, 3)
											SCREAMING_SNAKE_CASE__		: Optional[int]     =     np.array([0.5902, 0.6015, 0.6027, 0.5983, 0.6092, 0.6061, 0.5765, 0.5785, 0.5555]					)

											assert np.abs(expected_slice - image_slice					).max() < 1E-3



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



											"""simple docstring"""




											SCREAMING_SNAKE_CASE__		: Optional[Any]     =     StableDiffusionInstructPixaPixPipeline.from_pretrained(
											    """timbrooks/instruct-pix2pix"""     , safety_checker=SCREAMING_SNAKE_CASE__					)
											SCREAMING_SNAKE_CASE__		: int     =     LMSDiscreteScheduler.from_config(pipe.scheduler.config					)
											pipe.to(SCREAMING_SNAKE_CASE__					)
											pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE__					)
											pipe.enable_attention_slicing()

											SCREAMING_SNAKE_CASE__		: List[str]     =     self.get_inputs()
											SCREAMING_SNAKE_CASE__		: str     =     pipe(**SCREAMING_SNAKE_CASE__					).images
											SCREAMING_SNAKE_CASE__		: Any     =     image[0, -3:, -3:, -1].flatten()

											assert image.shape == (1, 5_12, 5_12, 3)
											SCREAMING_SNAKE_CASE__		: Tuple     =     np.array([0.6578, 0.6817, 0.6972, 0.6761, 0.6856, 0.6916, 0.6428, 0.6516, 0.6301]					)

											assert np.abs(expected_slice - image_slice					).max() < 1E-3



				def 			__magic_name__					(self					)      ->	str:



											"""simple docstring"""




											SCREAMING_SNAKE_CASE__		: str     =     StableDiffusionInstructPixaPixPipeline.from_pretrained(
											    """timbrooks/instruct-pix2pix"""     , safety_checker=SCREAMING_SNAKE_CASE__					)
											SCREAMING_SNAKE_CASE__		: Any     =     DDIMScheduler.from_config(pipe.scheduler.config					)
											pipe.to(SCREAMING_SNAKE_CASE__					)
											pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE__					)
											pipe.enable_attention_slicing()

											SCREAMING_SNAKE_CASE__		: Tuple     =     self.get_inputs()
											SCREAMING_SNAKE_CASE__		: Optional[int]     =     pipe(**SCREAMING_SNAKE_CASE__					).images
											SCREAMING_SNAKE_CASE__		: Tuple     =     image[0, -3:, -3:, -1].flatten()

											assert image.shape == (1, 5_12, 5_12, 3)
											SCREAMING_SNAKE_CASE__		: Optional[Any]     =     np.array([0.3828, 0.3834, 0.3818, 0.3792, 0.3865, 0.3752, 0.3792, 0.3847, 0.3753]					)

											assert np.abs(expected_slice - image_slice					).max() < 1E-3



				def 			__magic_name__					(self					)      ->	str:



											"""simple docstring"""




											SCREAMING_SNAKE_CASE__		: Any     =     0

											def callback_fn(SCREAMING_SNAKE_CASE__     , SCREAMING_SNAKE_CASE__     , SCREAMING_SNAKE_CASE__					) -> None:
																		SCREAMING_SNAKE_CASE__		: Optional[int]     =     True
																		nonlocal number_of_steps
																		number_of_steps += 1
																		if step == 1:
																									SCREAMING_SNAKE_CASE__		: Dict     =     latents.detach().cpu().numpy()
																									assert latents.shape == (1, 4, 64, 64)
																									SCREAMING_SNAKE_CASE__		: Optional[int]     =     latents[0, -3:, -3:, -1]
																									SCREAMING_SNAKE_CASE__		: int     =     np.array([-0.2463, -0.4644, -0.9756, 1.5176, 1.4414, 0.7866, 0.9897, 0.8521, 0.7983]					)

																									assert np.abs(latents_slice.flatten() - expected_slice					).max() < 5E-2
																		elif step == 2:
																									SCREAMING_SNAKE_CASE__		: int     =     latents.detach().cpu().numpy()
																									assert latents.shape == (1, 4, 64, 64)
																									SCREAMING_SNAKE_CASE__		: Optional[Any]     =     latents[0, -3:, -3:, -1]
																									SCREAMING_SNAKE_CASE__		: List[Any]     =     np.array([-0.2644, -0.4626, -0.9653, 1.5176, 1.4551, 0.7686, 0.9805, 0.8452, 0.8115]					)

																									assert np.abs(latents_slice.flatten() - expected_slice					).max() < 5E-2

											SCREAMING_SNAKE_CASE__		: Optional[int]     =     False

											SCREAMING_SNAKE_CASE__		: Optional[Any]     =     StableDiffusionInstructPixaPixPipeline.from_pretrained(
											    """timbrooks/instruct-pix2pix"""     , safety_checker=SCREAMING_SNAKE_CASE__     , torch_dtype=torch.floataa					)
											SCREAMING_SNAKE_CASE__		: List[str]     =     pipe.to(SCREAMING_SNAKE_CASE__					)
											pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE__					)
											pipe.enable_attention_slicing()

											SCREAMING_SNAKE_CASE__		: List[Any]     =     self.get_inputs()
											pipe(**SCREAMING_SNAKE_CASE__     , callback=SCREAMING_SNAKE_CASE__     , callback_steps=1					)
											assert callback_fn.has_been_called
											assert number_of_steps == 3



				def 			__magic_name__					(self					)      ->	Tuple:



											"""simple docstring"""




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

											SCREAMING_SNAKE_CASE__		: str     =     StableDiffusionInstructPixaPixPipeline.from_pretrained(
											    """timbrooks/instruct-pix2pix"""     , safety_checker=SCREAMING_SNAKE_CASE__     , torch_dtype=torch.floataa					)
											SCREAMING_SNAKE_CASE__		: str     =     pipe.to(SCREAMING_SNAKE_CASE__					)
											pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE__					)
											pipe.enable_attention_slicing(1					)
											pipe.enable_sequential_cpu_offload()

											SCREAMING_SNAKE_CASE__		: Any     =     self.get_inputs()
											SCREAMING_SNAKE_CASE__		: List[Any]     =     pipe(**SCREAMING_SNAKE_CASE__					)

											SCREAMING_SNAKE_CASE__		: str     =     torch.cuda.max_memory_allocated()
											# make sure that less than 2.2 GB is allocated
											assert mem_bytes < 2.2 * 10**9



				def 			__magic_name__					(self					)      ->	List[Any]:



											"""simple docstring"""




											SCREAMING_SNAKE_CASE__		: Tuple     =     self.get_inputs()
											# resize to resolution that is divisible by 8 but not 16 or 32
											SCREAMING_SNAKE_CASE__		: List[str]     =     inputs["""image"""].resize((5_04, 5_04)					)

											SCREAMING_SNAKE_CASE__		: Any     =     """timbrooks/instruct-pix2pix"""
											SCREAMING_SNAKE_CASE__		: Dict     =     StableDiffusionInstructPixaPixPipeline.from_pretrained(
											    SCREAMING_SNAKE_CASE__     , safety_checker=SCREAMING_SNAKE_CASE__     , )
											pipe.to(SCREAMING_SNAKE_CASE__					)
											pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE__					)
											pipe.enable_attention_slicing()

											SCREAMING_SNAKE_CASE__		: int     =     pipe(**SCREAMING_SNAKE_CASE__					)
											SCREAMING_SNAKE_CASE__		: Optional[int]     =     output.images[0]

											SCREAMING_SNAKE_CASE__		: Optional[Any]     =     image[2_55:2_58, 3_83:3_86, -1]

											assert image.shape == (5_04, 5_04, 3)
											SCREAMING_SNAKE_CASE__		: Optional[int]     =     np.array([0.2726, 0.2529, 0.2664, 0.2655, 0.2641, 0.2642, 0.2591, 0.2649, 0.2590]					)

											assert np.abs(image_slice.flatten() - expected_slice					).max() < 5E-3






"""simple docstring"""




from manim import *



class 			lowerCAmelCase_		(a__						):


				"""simple docstring"""



				def 			__magic_name__					(self					)      ->	Any:



											"""simple docstring"""




											SCREAMING_SNAKE_CASE__		: Dict     =     Rectangle(height=0.5     , width=0.5					)
											SCREAMING_SNAKE_CASE__		: Dict     =     Rectangle(height=0.46     , width=0.46					).set_stroke(width=0					)

											SCREAMING_SNAKE_CASE__		: Optional[Any]     =     Rectangle(height=0.25     , width=0.25					)

											SCREAMING_SNAKE_CASE__		: Any     =     [mem.copy() for i in range(6					)]
											SCREAMING_SNAKE_CASE__		: int     =     [mem.copy() for i in range(6					)]
											SCREAMING_SNAKE_CASE__		: Any     =     VGroup(*SCREAMING_SNAKE_CASE__					).arrange(SCREAMING_SNAKE_CASE__     , buff=0					)
											SCREAMING_SNAKE_CASE__		: Any     =     VGroup(*SCREAMING_SNAKE_CASE__					).arrange(SCREAMING_SNAKE_CASE__     , buff=0					)
											SCREAMING_SNAKE_CASE__		: Optional[int]     =     VGroup(SCREAMING_SNAKE_CASE__     , SCREAMING_SNAKE_CASE__					).arrange(SCREAMING_SNAKE_CASE__     , buff=0					)
											SCREAMING_SNAKE_CASE__		: List[str]     =     Text("""CPU"""     , font_size=24					)
											SCREAMING_SNAKE_CASE__		: Union[str, Any]     =     Group(SCREAMING_SNAKE_CASE__     , SCREAMING_SNAKE_CASE__					).arrange(SCREAMING_SNAKE_CASE__     , buff=0.5     , aligned_edge=SCREAMING_SNAKE_CASE__					)
											cpu.move_to([-2.5, -0.5, 0]					)
											self.add(SCREAMING_SNAKE_CASE__					)

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

											SCREAMING_SNAKE_CASE__		: List[Any]     =     [mem.copy() for i in range(6					)]
											SCREAMING_SNAKE_CASE__		: Optional[Any]     =     VGroup(*SCREAMING_SNAKE_CASE__					).arrange(SCREAMING_SNAKE_CASE__     , buff=0					)

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

											SCREAMING_SNAKE_CASE__		: Optional[Any]     =     []
											SCREAMING_SNAKE_CASE__		: Any     =     []

											for i, rect in enumerate(SCREAMING_SNAKE_CASE__					):
																		SCREAMING_SNAKE_CASE__		: List[str]     =     fill.copy().set_fill(SCREAMING_SNAKE_CASE__     , opacity=0.8					)
																		target.move_to(SCREAMING_SNAKE_CASE__					)
																		model_arr.append(SCREAMING_SNAKE_CASE__					)

																		SCREAMING_SNAKE_CASE__		: List[str]     =     Rectangle(height=0.46     , width=0.46					).set_stroke(width=0.0					).set_fill(SCREAMING_SNAKE_CASE__     , opacity=0.8					)
																		cpu_target.move_to(cpu_left_col_base[i]					)
																		model_cpu_arr.append(SCREAMING_SNAKE_CASE__					)

											self.add(*SCREAMING_SNAKE_CASE__     , *SCREAMING_SNAKE_CASE__					)

											SCREAMING_SNAKE_CASE__		: Tuple     =     [meta_mem.copy() for i in range(6					)]
											SCREAMING_SNAKE_CASE__		: List[str]     =     [meta_mem.copy() for i in range(6					)]
											SCREAMING_SNAKE_CASE__		: Any     =     VGroup(*SCREAMING_SNAKE_CASE__					).arrange(SCREAMING_SNAKE_CASE__     , buff=0					)
											SCREAMING_SNAKE_CASE__		: int     =     VGroup(*SCREAMING_SNAKE_CASE__					).arrange(SCREAMING_SNAKE_CASE__     , buff=0					)
											SCREAMING_SNAKE_CASE__		: Union[str, Any]     =     VGroup(SCREAMING_SNAKE_CASE__     , SCREAMING_SNAKE_CASE__					).arrange(SCREAMING_SNAKE_CASE__     , buff=0					)
											SCREAMING_SNAKE_CASE__		: Union[str, Any]     =     Text("""Disk"""     , font_size=24					)
											SCREAMING_SNAKE_CASE__		: Any     =     Group(SCREAMING_SNAKE_CASE__     , SCREAMING_SNAKE_CASE__					).arrange(SCREAMING_SNAKE_CASE__     , buff=0.5     , aligned_edge=SCREAMING_SNAKE_CASE__					)
											disk.move_to([-4, -1.25, 0]					)
											self.add(SCREAMING_SNAKE_CASE__     , SCREAMING_SNAKE_CASE__					)

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

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

											SCREAMING_SNAKE_CASE__		: Any     =     MarkupText(
											    F'''<span fgcolor=\'{BLUE}\'>●</span> Checkpoint'''     , font_size=18     , )

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

											SCREAMING_SNAKE_CASE__		: Union[str, Any]     =     MarkupText(
											    F'''Now watch as an input is passed through the model\nand how the memory is utilized and handled.'''     , font_size=24     , )
											step_a.move_to([2, 2, 0]					)

											self.play(Write(SCREAMING_SNAKE_CASE__					)					)

											SCREAMING_SNAKE_CASE__		: Dict     =     Square(0.3					)
											input.set_fill(SCREAMING_SNAKE_CASE__     , opacity=1.0					)
											input.set_stroke(width=0.0					)
											input.next_to(model_base[0]     , SCREAMING_SNAKE_CASE__     , buff=0.5					)

											self.play(Write(SCREAMING_SNAKE_CASE__					)					)

											input.generate_target()
											input.target.next_to(model_arr[0]     , direction=SCREAMING_SNAKE_CASE__     , buff=0.02					)
											self.play(MoveToTarget(SCREAMING_SNAKE_CASE__					)					)

											self.play(FadeOut(SCREAMING_SNAKE_CASE__					)					)

											SCREAMING_SNAKE_CASE__		: Optional[Any]     =     Arrow(start=SCREAMING_SNAKE_CASE__     , end=SCREAMING_SNAKE_CASE__     , color=SCREAMING_SNAKE_CASE__     , buff=0.5					)
											a.next_to(model_arr[0].get_left()     , SCREAMING_SNAKE_CASE__     , buff=0.2					)

											model_cpu_arr[0].generate_target()
											model_cpu_arr[0].target.move_to(gpu_rect[0]					)

											SCREAMING_SNAKE_CASE__		: Tuple     =     MarkupText(
											    F'''As the input reaches a layer, the hook triggers\nand weights are moved from the CPU\nto the GPU and back.'''     , font_size=24     , )
											step_a.move_to([2, 2, 0]					)

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

											SCREAMING_SNAKE_CASE__		: Tuple     =     {"""run_time""": 1, """fade_in""": True, """fade_out""": True, """buff""": 0.02}

											self.play(
											    Write(SCREAMING_SNAKE_CASE__					)     , Circumscribe(model_arr[0]     , color=SCREAMING_SNAKE_CASE__     , **SCREAMING_SNAKE_CASE__					)     , Circumscribe(model_cpu_arr[0]     , color=SCREAMING_SNAKE_CASE__     , **SCREAMING_SNAKE_CASE__					)     , Circumscribe(gpu_rect[0]     , color=SCREAMING_SNAKE_CASE__     , **SCREAMING_SNAKE_CASE__					)     , )
											self.play(MoveToTarget(model_cpu_arr[0]					)					)

											SCREAMING_SNAKE_CASE__		: Optional[Any]     =     a.copy()
											for i in range(6					):
																		a_c.next_to(model_arr[i].get_right() + 0.02     , SCREAMING_SNAKE_CASE__     , buff=0.2					)

																		input.generate_target()
																		input.target.move_to(model_arr[i].get_right() + 0.02					)

																		SCREAMING_SNAKE_CASE__		: Optional[Any]     =     AnimationGroup(
																		    FadeOut(SCREAMING_SNAKE_CASE__     , run_time=0.5					)     , MoveToTarget(SCREAMING_SNAKE_CASE__     , run_time=0.5					)     , FadeIn(SCREAMING_SNAKE_CASE__     , run_time=0.5					)     , lag_ratio=0.2					)

																		self.play(SCREAMING_SNAKE_CASE__					)

																		model_cpu_arr[i].generate_target()
																		model_cpu_arr[i].target.move_to(cpu_left_col_base[i]					)

																		if i < 5:
																									model_cpu_arr[i + 1].generate_target()
																									model_cpu_arr[i + 1].target.move_to(gpu_rect[0]					)
																									if i >= 1:
																																SCREAMING_SNAKE_CASE__		: Union[str, Any]     =     0.7

																									self.play(
																									    Circumscribe(model_arr[i]     , **SCREAMING_SNAKE_CASE__					)     , Circumscribe(cpu_left_col_base[i]     , **SCREAMING_SNAKE_CASE__					)     , Circumscribe(cpu_left_col_base[i + 1]     , color=SCREAMING_SNAKE_CASE__     , **SCREAMING_SNAKE_CASE__					)     , Circumscribe(gpu_rect[0]     , color=SCREAMING_SNAKE_CASE__     , **SCREAMING_SNAKE_CASE__					)     , Circumscribe(model_arr[i + 1]     , color=SCREAMING_SNAKE_CASE__     , **SCREAMING_SNAKE_CASE__					)     , )
																									if i < 1:
																																self.play(
																																    MoveToTarget(model_cpu_arr[i]					)     , MoveToTarget(model_cpu_arr[i + 1]					)     , )
																									else:
																																self.play(
																																    MoveToTarget(model_cpu_arr[i]     , run_time=0.7					)     , MoveToTarget(model_cpu_arr[i + 1]     , run_time=0.7					)     , )
																		else:
																									model_cpu_arr[i].generate_target()
																									model_cpu_arr[i].target.move_to(cpu_left_col_base[-1]					)
																									input.generate_target()
																									input.target.next_to(model_arr[-1].get_right()     , RIGHT + 0.02     , buff=0.2					)

																									self.play(
																									    Circumscribe(model_arr[-1]     , color=SCREAMING_SNAKE_CASE__     , **SCREAMING_SNAKE_CASE__					)     , Circumscribe(cpu_left_col_base[-1]     , color=SCREAMING_SNAKE_CASE__     , **SCREAMING_SNAKE_CASE__					)     , Circumscribe(gpu_rect[0]     , color=SCREAMING_SNAKE_CASE__     , **SCREAMING_SNAKE_CASE__					)     , )

																									self.play(MoveToTarget(model_cpu_arr[i]					)					)

																		SCREAMING_SNAKE_CASE__		: Union[str, Any]     =     a_c
																		SCREAMING_SNAKE_CASE__		: Optional[Any]     =     a_c.copy()

											input.generate_target()
											input.target.next_to(model_base[-1]     , RIGHT + 0.02     , buff=0.5					)
											self.play(
											    FadeOut(SCREAMING_SNAKE_CASE__					)     , FadeOut(SCREAMING_SNAKE_CASE__     , run_time=0.5					)     , )

											SCREAMING_SNAKE_CASE__		: Dict     =     MarkupText(F'''Inference on a model too large for GPU memory\nis successfully completed.'''     , font_size=24					)
											step_a.move_to([2, 2, 0]					)

											self.play(Write(SCREAMING_SNAKE_CASE__     , run_time=3					)     , MoveToTarget(SCREAMING_SNAKE_CASE__					)					)

											self.wait()




