'''simple docstring'''


import gc
import random
import tempfile
import unittest

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

from diffusers import (
    AutoencoderKL,
    DDIMInverseScheduler,
    DDIMScheduler,
    DPMSolverMultistepInverseScheduler,
    DPMSolverMultistepScheduler,
    StableDiffusionDiffEditPipeline,
    UNetaDConditionModel,
)
from diffusers.utils import load_image, slow
from diffusers.utils.testing_utils import enable_full_determinism, floats_tensor, require_torch_gpu, torch_device

from ..pipeline_params import TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS, TEXT_GUIDED_IMAGE_INPAINTING_PARAMS
from ..test_pipelines_common import PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()






class 					__SCREAMING_SNAKE_CASE						(	_UpperCAmelCase						,     _UpperCAmelCase						,     unittest.TestCase	):







  '''simple docstring'''






  lowerCamelCase_		:List[str]						   =		StableDiffusionDiffEditPipeline
  lowerCamelCase_		:List[Any]						   =		TEXT_GUIDED_IMAGE_INPAINTING_PARAMS - {'height', 'width', 'image'} | {'image_latents'}
  lowerCamelCase_		:str						   =		TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS - {'image'} | {'image_latents'}
  lowerCamelCase_		:List[Any]						   =		frozenset(
      []	)  # TO-DO: update image_params once pipeline is refactored with VaeImageProcessor.preprocess
  lowerCamelCase_		:Union[str, Any]						   =		frozenset([]	)


  def _UpperCamelCase					(			self						):




       '''simple docstring'''

       torch.manual_seed(0						)
       UpperCAmelCase_							:       List[str]       =				UNetaDConditionModel(
           block_out_channels=(3_2, 6_4)				,		layers_per_block=2				,		sample_size=3_2				,		in_channels=4				,		out_channels=4				,		down_block_types=('DownBlock2D', 'CrossAttnDownBlock2D')				,		up_block_types=('CrossAttnUpBlock2D', 'UpBlock2D')				,		cross_attention_dim=3_2				,		attention_head_dim=(2, 4)				,		use_linear_projection=SCREAMING_SNAKE_CASE_				,		)
       UpperCAmelCase_							:       Union[str, Any]       =				DDIMScheduler(
           beta_start=0.0_00_85				,		beta_end=0.0_12				,		beta_schedule='scaled_linear'				,		clip_sample=SCREAMING_SNAKE_CASE_				,		set_alpha_to_one=SCREAMING_SNAKE_CASE_				,		)
       UpperCAmelCase_							:       Dict       =				DDIMInverseScheduler(
           beta_start=0.0_00_85				,		beta_end=0.0_12				,		beta_schedule='scaled_linear'				,		clip_sample=SCREAMING_SNAKE_CASE_				,		set_alpha_to_zero=SCREAMING_SNAKE_CASE_				,		)
       torch.manual_seed(0						)
       UpperCAmelCase_							:       Any       =				AutoencoderKL(
           block_out_channels=[3_2, 6_4]				,		in_channels=3				,		out_channels=3				,		down_block_types=['DownEncoderBlock2D', 'DownEncoderBlock2D']				,		up_block_types=['UpDecoderBlock2D', 'UpDecoderBlock2D']				,		latent_channels=4				,		sample_size=1_2_8				,		)
       torch.manual_seed(0						)
       UpperCAmelCase_							:       Optional[int]       =				CLIPTextConfig(
           bos_token_id=0				,		eos_token_id=2				,		hidden_size=3_2				,		intermediate_size=3_7				,		layer_norm_eps=1E-05				,		num_attention_heads=4				,		num_hidden_layers=5				,		pad_token_id=1				,		vocab_size=1_0_0_0				,		hidden_act='gelu'				,		projection_dim=5_1_2				,		)
       UpperCAmelCase_							:       List[str]       =				CLIPTextModel(SCREAMING_SNAKE_CASE_						)
       UpperCAmelCase_							:       Union[str, Any]       =				CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'						)

       UpperCAmelCase_							:       Any       =				{
           """unet""": unet,
           """scheduler""": scheduler,
           """inverse_scheduler""": inverse_scheduler,
           """vae""": vae,
           """text_encoder""": text_encoder,
           """tokenizer""": tokenizer,
           """safety_checker""": None,
           """feature_extractor""": None,
       }

       return components


  def _UpperCamelCase					(			self				,		snake_case_				,		snake_case_=0						):




       '''simple docstring'''

       UpperCAmelCase_							:       str       =				floats_tensor((1, 1_6, 1_6)				,		rng=random.Random(SCREAMING_SNAKE_CASE_						)						).to(SCREAMING_SNAKE_CASE_						)
       UpperCAmelCase_							:       Any       =				floats_tensor((1, 2, 4, 1_6, 1_6)				,		rng=random.Random(SCREAMING_SNAKE_CASE_						)						).to(SCREAMING_SNAKE_CASE_						)
       if str(SCREAMING_SNAKE_CASE_						).startswith('mps'						):
            UpperCAmelCase_							:       List[Any]       =				torch.manual_seed(SCREAMING_SNAKE_CASE_						)
       else:
            UpperCAmelCase_							:       Dict       =				torch.Generator(device=SCREAMING_SNAKE_CASE_						).manual_seed(SCREAMING_SNAKE_CASE_						)
       UpperCAmelCase_							:       int       =				{
           """prompt""": """a dog and a newt""",
           """mask_image""": mask,
           """image_latents""": latents,
           """generator""": generator,
           """num_inference_steps""": 2,
           """inpaint_strength""": 1.0,
           """guidance_scale""": 6.0,
           """output_type""": """numpy""",
       }

       return inputs


  def _UpperCamelCase					(			self				,		snake_case_				,		snake_case_=0						):




       '''simple docstring'''

       UpperCAmelCase_							:       Optional[int]       =				floats_tensor((1, 3, 3_2, 3_2)				,		rng=random.Random(SCREAMING_SNAKE_CASE_						)						).to(SCREAMING_SNAKE_CASE_						)
       UpperCAmelCase_							:       Dict       =				image.cpu().permute(0				,		2				,		3				,		1						)[0]
       UpperCAmelCase_							:       Optional[int]       =				Image.fromarray(np.uinta(SCREAMING_SNAKE_CASE_						)						).convert('RGB'						)
       if str(SCREAMING_SNAKE_CASE_						).startswith('mps'						):
            UpperCAmelCase_							:       Optional[Any]       =				torch.manual_seed(SCREAMING_SNAKE_CASE_						)
       else:
            UpperCAmelCase_							:       str       =				torch.Generator(device=SCREAMING_SNAKE_CASE_						).manual_seed(SCREAMING_SNAKE_CASE_						)
       UpperCAmelCase_							:       List[Any]       =				{
           """image""": image,
           """source_prompt""": """a cat and a frog""",
           """target_prompt""": """a dog and a newt""",
           """generator""": generator,
           """num_inference_steps""": 2,
           """num_maps_per_mask""": 2,
           """mask_encode_strength""": 1.0,
           """guidance_scale""": 6.0,
           """output_type""": """numpy""",
       }

       return inputs


  def _UpperCamelCase					(			self				,		snake_case_				,		snake_case_=0						):




       '''simple docstring'''

       UpperCAmelCase_							:       Any       =				floats_tensor((1, 3, 3_2, 3_2)				,		rng=random.Random(SCREAMING_SNAKE_CASE_						)						).to(SCREAMING_SNAKE_CASE_						)
       UpperCAmelCase_							:       List[Any]       =				image.cpu().permute(0				,		2				,		3				,		1						)[0]
       UpperCAmelCase_							:       List[Any]       =				Image.fromarray(np.uinta(SCREAMING_SNAKE_CASE_						)						).convert('RGB'						)
       if str(SCREAMING_SNAKE_CASE_						).startswith('mps'						):
            UpperCAmelCase_							:       Any       =				torch.manual_seed(SCREAMING_SNAKE_CASE_						)
       else:
            UpperCAmelCase_							:       Optional[Any]       =				torch.Generator(device=SCREAMING_SNAKE_CASE_						).manual_seed(SCREAMING_SNAKE_CASE_						)
       UpperCAmelCase_							:       Tuple       =				{
           """image""": image,
           """prompt""": """a cat and a frog""",
           """generator""": generator,
           """num_inference_steps""": 2,
           """inpaint_strength""": 1.0,
           """guidance_scale""": 6.0,
           """decode_latents""": True,
           """output_type""": """numpy""",
       }
       return inputs


  def _UpperCamelCase					(			self						):




       '''simple docstring'''

       if not hasattr(self.pipeline_class				,		'_optional_components'						):
            return

       UpperCAmelCase_							:       Optional[Any]       =				self.get_dummy_components()
       UpperCAmelCase_							:       str       =				self.pipeline_class(**SCREAMING_SNAKE_CASE_						)
       pipe.to(SCREAMING_SNAKE_CASE_						)
       pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE_						)

       # set all optional components to None and update pipeline config accordingly
       for optional_component in pipe._optional_components:
            setattr(SCREAMING_SNAKE_CASE_				,		SCREAMING_SNAKE_CASE_				,		SCREAMING_SNAKE_CASE_						)
       pipe.register_modules(**{optional_component: None for optional_component in pipe._optional_components}						)

       UpperCAmelCase_							:       Union[str, Any]       =				self.get_dummy_inputs(SCREAMING_SNAKE_CASE_						)
       UpperCAmelCase_							:       Dict       =				pipe(**SCREAMING_SNAKE_CASE_						)[0]

       with tempfile.TemporaryDirectory() as tmpdir:
            pipe.save_pretrained(SCREAMING_SNAKE_CASE_						)
            UpperCAmelCase_							:       str       =				self.pipeline_class.from_pretrained(SCREAMING_SNAKE_CASE_						)
            pipe_loaded.to(SCREAMING_SNAKE_CASE_						)
            pipe_loaded.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE_						)

       for optional_component in pipe._optional_components:
            self.assertTrue(
                getattr(SCREAMING_SNAKE_CASE_				,		SCREAMING_SNAKE_CASE_						) is None				,		F'''`{optional_component}` did not stay set to None after loading.'''				,		)

       UpperCAmelCase_							:       Optional[int]       =				self.get_dummy_inputs(SCREAMING_SNAKE_CASE_						)
       UpperCAmelCase_							:       Tuple       =				pipe_loaded(**SCREAMING_SNAKE_CASE_						)[0]

       UpperCAmelCase_							:       Dict       =				np.abs(output - output_loaded						).max()
       self.assertLess(SCREAMING_SNAKE_CASE_				,		1E-4						)


  def _UpperCamelCase					(			self						):




       '''simple docstring'''

       UpperCAmelCase_							:       Dict       =				"""cpu"""

       UpperCAmelCase_							:       Dict       =				self.get_dummy_components()
       UpperCAmelCase_							:       Dict       =				self.pipeline_class(**SCREAMING_SNAKE_CASE_						)
       pipe.to(SCREAMING_SNAKE_CASE_						)
       pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE_						)

       UpperCAmelCase_							:       Union[str, Any]       =				self.get_dummy_mask_inputs(SCREAMING_SNAKE_CASE_						)
       UpperCAmelCase_							:       int       =				pipe.generate_mask(**SCREAMING_SNAKE_CASE_						)
       UpperCAmelCase_							:       Optional[int]       =				mask[0, -3:, -3:]

       self.assertEqual(mask.shape				,		(1, 1_6, 1_6)						)
       UpperCAmelCase_							:       List[Any]       =				np.array([0] * 9						)
       UpperCAmelCase_							:       List[Any]       =				np.abs(mask_slice.flatten() - expected_slice						).max()
       self.assertLessEqual(SCREAMING_SNAKE_CASE_				,		1E-3						)
       self.assertEqual(mask[0, -3, -4]				,		0						)


  def _UpperCamelCase					(			self						):




       '''simple docstring'''

       UpperCAmelCase_							:       Dict       =				"""cpu"""

       UpperCAmelCase_							:       Tuple       =				self.get_dummy_components()
       UpperCAmelCase_							:       Any       =				self.pipeline_class(**SCREAMING_SNAKE_CASE_						)
       pipe.to(SCREAMING_SNAKE_CASE_						)
       pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE_						)

       UpperCAmelCase_							:       Tuple       =				self.get_dummy_inversion_inputs(SCREAMING_SNAKE_CASE_						)
       UpperCAmelCase_							:       List[str]       =				pipe.invert(**SCREAMING_SNAKE_CASE_						).images
       UpperCAmelCase_							:       str       =				image[0, -1, -3:, -3:]

       self.assertEqual(image.shape				,		(2, 3_2, 3_2, 3)						)
       UpperCAmelCase_							:       Union[str, Any]       =				np.array(
           [0.51_50, 0.51_34, 0.50_43, 0.53_76, 0.46_94, 0.5_10_50, 0.50_15, 0.44_07, 0.47_99]				,		)
       UpperCAmelCase_							:       Dict       =				np.abs(image_slice.flatten() - expected_slice						).max()
       self.assertLessEqual(SCREAMING_SNAKE_CASE_				,		1E-3						)


  def _UpperCamelCase					(			self						):




       '''simple docstring'''

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






  def _UpperCamelCase					(			self						):




       '''simple docstring'''

       UpperCAmelCase_							:       str       =				"""cpu"""

       UpperCAmelCase_							:       int       =				self.get_dummy_components()

       UpperCAmelCase_							:       str       =				{"""beta_start""": 0.0_00_85, """beta_end""": 0.0_12, """beta_schedule""": """scaled_linear"""}
       UpperCAmelCase_							:       List[Any]       =				DPMSolverMultistepScheduler(**SCREAMING_SNAKE_CASE_						)
       UpperCAmelCase_							:       List[Any]       =				DPMSolverMultistepInverseScheduler(**SCREAMING_SNAKE_CASE_						)

       UpperCAmelCase_							:       Optional[int]       =				self.pipeline_class(**SCREAMING_SNAKE_CASE_						)
       pipe.to(SCREAMING_SNAKE_CASE_						)
       pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE_						)

       UpperCAmelCase_							:       Optional[int]       =				self.get_dummy_inversion_inputs(SCREAMING_SNAKE_CASE_						)
       UpperCAmelCase_							:       Optional[int]       =				pipe.invert(**SCREAMING_SNAKE_CASE_						).images
       UpperCAmelCase_							:       Optional[Any]       =				image[0, -1, -3:, -3:]

       self.assertEqual(image.shape				,		(2, 3_2, 3_2, 3)						)
       UpperCAmelCase_							:       Tuple       =				np.array(
           [0.51_50, 0.51_34, 0.50_43, 0.53_76, 0.46_94, 0.5_10_50, 0.50_15, 0.44_07, 0.47_99]				,		)
       UpperCAmelCase_							:       Optional[Any]       =				np.abs(image_slice.flatten() - expected_slice						).max()
       self.assertLessEqual(SCREAMING_SNAKE_CASE_				,		1E-3						)







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







  '''simple docstring'''






  def _UpperCamelCase					(			self						):




       '''simple docstring'''

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


  @classmethod
  def _UpperCamelCase					(			cls						):




       '''simple docstring'''

       UpperCAmelCase_							:       Any       =				load_image(
           'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/diffedit/fruit.png'						)

       UpperCAmelCase_							:       List[str]       =				raw_image.convert('RGB'						).resize((7_6_8, 7_6_8)						)

       UpperCAmelCase_							:       Optional[int]       =				raw_image


  def _UpperCamelCase					(			self						):




       '''simple docstring'''

       UpperCAmelCase_							:       List[Any]       =				torch.manual_seed(0						)

       UpperCAmelCase_							:       List[Any]       =				StableDiffusionDiffEditPipeline.from_pretrained(
           'stabilityai/stable-diffusion-2-1'				,		safety_checker=SCREAMING_SNAKE_CASE_				,		torch_dtype=torch.floataa						)
       UpperCAmelCase_							:       List[Any]       =				DDIMScheduler.from_config(pipe.scheduler.config						)
       UpperCAmelCase_							:       List[Any]       =				DDIMInverseScheduler.from_config(pipe.scheduler.config						)
       pipe.enable_model_cpu_offload()
       pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE_						)

       UpperCAmelCase_							:       int       =				"""a bowl of fruit"""
       UpperCAmelCase_							:       int       =				"""a bowl of pears"""

       UpperCAmelCase_							:       Tuple       =				pipe.generate_mask(
           image=self.raw_image				,		source_prompt=SCREAMING_SNAKE_CASE_				,		target_prompt=SCREAMING_SNAKE_CASE_				,		generator=SCREAMING_SNAKE_CASE_				,		)

       UpperCAmelCase_							:       Any       =				pipe.invert(
           prompt=SCREAMING_SNAKE_CASE_				,		image=self.raw_image				,		inpaint_strength=0.7				,		generator=SCREAMING_SNAKE_CASE_						).latents

       UpperCAmelCase_							:       Optional[int]       =				pipe(
           prompt=SCREAMING_SNAKE_CASE_				,		mask_image=SCREAMING_SNAKE_CASE_				,		image_latents=SCREAMING_SNAKE_CASE_				,		generator=SCREAMING_SNAKE_CASE_				,		negative_prompt=SCREAMING_SNAKE_CASE_				,		inpaint_strength=0.7				,		output_type='numpy'				,		).images[0]

       UpperCAmelCase_							:       Dict       =				(
           np.array(
               load_image(
                   'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
                   '/diffedit/pears.png'						).resize((7_6_8, 7_6_8)						)						)
           / 2_5_5
       )
       assert np.abs((expected_image - image).max()						) < 5E-1






  def _UpperCamelCase					(			self						):




       '''simple docstring'''

       UpperCAmelCase_							:       Union[str, Any]       =				torch.manual_seed(0						)

       UpperCAmelCase_							:       int       =				StableDiffusionDiffEditPipeline.from_pretrained(
           'stabilityai/stable-diffusion-2-1'				,		safety_checker=SCREAMING_SNAKE_CASE_				,		torch_dtype=torch.floataa						)
       UpperCAmelCase_							:       Optional[int]       =				DPMSolverMultistepScheduler.from_config(pipe.scheduler.config						)
       UpperCAmelCase_							:       Dict       =				DPMSolverMultistepInverseScheduler.from_config(pipe.scheduler.config						)
       pipe.enable_model_cpu_offload()
       pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE_						)

       UpperCAmelCase_							:       Tuple       =				"""a bowl of fruit"""
       UpperCAmelCase_							:       List[Any]       =				"""a bowl of pears"""

       UpperCAmelCase_							:       List[str]       =				pipe.generate_mask(
           image=self.raw_image				,		source_prompt=SCREAMING_SNAKE_CASE_				,		target_prompt=SCREAMING_SNAKE_CASE_				,		generator=SCREAMING_SNAKE_CASE_				,		)

       UpperCAmelCase_							:       Dict       =				pipe.invert(
           prompt=SCREAMING_SNAKE_CASE_				,		image=self.raw_image				,		inpaint_strength=0.7				,		generator=SCREAMING_SNAKE_CASE_				,		num_inference_steps=2_5				,		).latents

       UpperCAmelCase_							:       Any       =				pipe(
           prompt=SCREAMING_SNAKE_CASE_				,		mask_image=SCREAMING_SNAKE_CASE_				,		image_latents=SCREAMING_SNAKE_CASE_				,		generator=SCREAMING_SNAKE_CASE_				,		negative_prompt=SCREAMING_SNAKE_CASE_				,		inpaint_strength=0.7				,		num_inference_steps=2_5				,		output_type='numpy'				,		).images[0]

       UpperCAmelCase_							:       Any       =				(
           np.array(
               load_image(
                   'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
                   '/diffedit/pears.png'						).resize((7_6_8, 7_6_8)						)						)
           / 2_5_5
       )
       assert np.abs((expected_image - image).max()						) < 5E-1







'''simple docstring'''


from manim import *






class 					__SCREAMING_SNAKE_CASE						(	lowerCamelCase_	):







  '''simple docstring'''






  def _UpperCamelCase					(			self						):




       '''simple docstring'''

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

       UpperCAmelCase_							:       List[str]       =				Rectangle(height=0.25				,		width=0.25						)

       UpperCAmelCase_							:       Any       =				[mem.copy() for i in range(6						)]
       UpperCAmelCase_							:       Union[str, Any]       =				[mem.copy() for i in range(6						)]
       UpperCAmelCase_							:       Optional[int]       =				VGroup(*snake_case_						).arrange(snake_case_				,		buff=0						)
       UpperCAmelCase_							:       Any       =				VGroup(*snake_case_						).arrange(snake_case_				,		buff=0						)
       UpperCAmelCase_							:       str       =				VGroup(snake_case_				,		snake_case_						).arrange(snake_case_				,		buff=0						)
       UpperCAmelCase_							:       Any       =				Text('CPU'				,		font_size=2_4						)
       UpperCAmelCase_							:       Tuple       =				Group(snake_case_				,		snake_case_						).arrange(snake_case_				,		buff=0.5				,		aligned_edge=snake_case_						)
       cpu.move_to([-2.5, -0.5, 0]						)
       self.add(snake_case_						)

       UpperCAmelCase_							:       str       =				[mem.copy() for i in range(4						)]
       UpperCAmelCase_							:       Dict       =				VGroup(*snake_case_						).arrange(snake_case_				,		buff=0						)
       UpperCAmelCase_							:       List[str]       =				Text('GPU'				,		font_size=2_4						)
       UpperCAmelCase_							:       Optional[int]       =				Group(snake_case_				,		snake_case_						).arrange(snake_case_				,		buff=0.5				,		aligned_edge=snake_case_						)
       gpu.move_to([-1, -1, 0]						)
       self.add(snake_case_						)

       UpperCAmelCase_							:       str       =				[mem.copy() for i in range(6						)]
       UpperCAmelCase_							:       Tuple       =				VGroup(*snake_case_						).arrange(snake_case_				,		buff=0						)

       UpperCAmelCase_							:       str       =				Text('Model'				,		font_size=2_4						)
       UpperCAmelCase_							:       Optional[int]       =				Group(snake_case_				,		snake_case_						).arrange(snake_case_				,		buff=0.5				,		aligned_edge=snake_case_						)
       model.move_to([3, -1.0, 0]						)
       self.add(snake_case_						)

       UpperCAmelCase_							:       str       =				[]
       UpperCAmelCase_							:       Optional[Any]       =				[]

       for i, rect in enumerate(snake_case_						):
            UpperCAmelCase_							:       str       =				fill.copy().set_fill(snake_case_				,		opacity=0.8						)
            target.move_to(snake_case_						)
            model_arr.append(snake_case_						)

            UpperCAmelCase_							:       int       =				Rectangle(height=0.46				,		width=0.46						).set_stroke(width=0.0						).set_fill(snake_case_				,		opacity=0.8						)
            cpu_target.move_to(cpu_left_col_base[i]						)
            model_cpu_arr.append(snake_case_						)

       self.add(*snake_case_				,		*snake_case_						)

       UpperCAmelCase_							:       List[Any]       =				[meta_mem.copy() for i in range(6						)]
       UpperCAmelCase_							:       List[str]       =				[meta_mem.copy() for i in range(6						)]
       UpperCAmelCase_							:       Tuple       =				VGroup(*snake_case_						).arrange(snake_case_				,		buff=0						)
       UpperCAmelCase_							:       Dict       =				VGroup(*snake_case_						).arrange(snake_case_				,		buff=0						)
       UpperCAmelCase_							:       Optional[Any]       =				VGroup(snake_case_				,		snake_case_						).arrange(snake_case_				,		buff=0						)
       UpperCAmelCase_							:       Tuple       =				Text('Disk'				,		font_size=2_4						)
       UpperCAmelCase_							:       Union[str, Any]       =				Group(snake_case_				,		snake_case_						).arrange(snake_case_				,		buff=0.5				,		aligned_edge=snake_case_						)
       disk.move_to([-4, -1.25, 0]						)
       self.add(snake_case_				,		snake_case_						)

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

       UpperCAmelCase_							:       Any       =				MarkupText(
           F'''<b>Key:</b>\n\n<span fgcolor=\'{YELLOW}\'>●</span> Empty Model'''				,		font_size=1_8				,		)

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

       self.add(snake_case_				,		snake_case_						)

       UpperCAmelCase_							:       Dict       =				MarkupText(
           F'''<span fgcolor=\'{BLUE}\'>●</span> Checkpoint'''				,		font_size=1_8				,		)

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

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

       self.play(Write(snake_case_						)						)

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

       self.play(Write(snake_case_						)						)

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

       self.play(FadeOut(snake_case_						)						)

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

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

       UpperCAmelCase_							:       List[str]       =				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=2_4				,		)
       step_a.move_to([2, 2, 0]						)

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

       UpperCAmelCase_							:       List[Any]       =				{'run_time': 1, 'fade_in': True, 'fade_out': True, 'buff': 0.02}

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

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

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

            UpperCAmelCase_							:       Tuple       =				AnimationGroup(
                FadeOut(snake_case_				,		run_time=0.5						)				,		MoveToTarget(snake_case_				,		run_time=0.5						)				,		FadeIn(snake_case_				,		run_time=0.5						)				,		lag_ratio=0.2						)

            self.play(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:
                      UpperCAmelCase_							:       Any       =				0.7

                 self.play(
                     Circumscribe(model_arr[i]				,		**snake_case_						)				,		Circumscribe(cpu_left_col_base[i]				,		**snake_case_						)				,		Circumscribe(cpu_left_col_base[i + 1]				,		color=snake_case_				,		**snake_case_						)				,		Circumscribe(gpu_rect[0]				,		color=snake_case_				,		**snake_case_						)				,		Circumscribe(model_arr[i + 1]				,		color=snake_case_				,		**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=snake_case_				,		**snake_case_						)				,		Circumscribe(cpu_left_col_base[-1]				,		color=snake_case_				,		**snake_case_						)				,		Circumscribe(gpu_rect[0]				,		color=snake_case_				,		**snake_case_						)				,		)

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

            UpperCAmelCase_							:       Any       =				a_c
            UpperCAmelCase_							:       int       =				a_c.copy()

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

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

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

       self.wait()






