code: from manim import *
  
  
  
  
  class    __lowerCAmelCase  (       a_						):
  
  
  
   '''simple docstring'''
  
   def 					__snake_case					(			self       :		Tuple  ) ->							Tuple:
          _UpperCamelCase     :       Tuple  						=  Rectangle(height=0.5       ,			width=0.5  )
          _UpperCamelCase     :       Dict  						=  Rectangle(height=0.4_6       ,			width=0.4_6  ).set_stroke(width=0  )
          _UpperCamelCase     :       str  						=  Rectangle(height=0.2_5       ,			width=0.2_5  )
  
          _UpperCamelCase     :       Union[str, Any]  						=  [mem.copy() for i in range(6  )]
          _UpperCamelCase     :       int  						=  [mem.copy() for i in range(6  )]
          _UpperCamelCase     :       List[str]  						=  VGroup(*__lowerCamelCase  ).arrange(__lowerCamelCase       ,			buff=0  )
          _UpperCamelCase     :       Union[str, Any]  						=  VGroup(*__lowerCamelCase  ).arrange(__lowerCamelCase       ,			buff=0  )
          _UpperCamelCase     :       Union[str, Any]  						=  VGroup(__lowerCamelCase       ,			__lowerCamelCase  ).arrange(__lowerCamelCase       ,			buff=0  )
          _UpperCamelCase     :       List[Any]  						=  Text('CPU'       ,			font_size=2_4  )
          _UpperCamelCase     :       List[str]  						=  Group(__lowerCamelCase       ,			__lowerCamelCase  ).arrange(__lowerCamelCase       ,			buff=0.5       ,			aligned_edge=__lowerCamelCase  )
          cpu.move_to([-2.5, -0.5, 0]  )
          self.add(__lowerCamelCase  )
  
          _UpperCamelCase     :       Optional[int]  						=  [mem.copy() for i in range(4  )]
          _UpperCamelCase     :       Dict  						=  VGroup(*__lowerCamelCase  ).arrange(__lowerCamelCase       ,			buff=0  )
          _UpperCamelCase     :       Optional[int]  						=  Text('GPU'       ,			font_size=2_4  )
          _UpperCamelCase     :       Optional[Any]  						=  Group(__lowerCamelCase       ,			__lowerCamelCase  ).arrange(__lowerCamelCase       ,			buff=0.5       ,			aligned_edge=__lowerCamelCase  )
          gpu.move_to([-1, -1, 0]  )
          self.add(__lowerCamelCase  )
  
          _UpperCamelCase     :       int  						=  [mem.copy() for i in range(6  )]
          _UpperCamelCase     :       Any  						=  VGroup(*__lowerCamelCase  ).arrange(__lowerCamelCase       ,			buff=0  )
  
          _UpperCamelCase     :       int  						=  Text('Model'       ,			font_size=2_4  )
          _UpperCamelCase     :       str  						=  Group(__lowerCamelCase       ,			__lowerCamelCase  ).arrange(__lowerCamelCase       ,			buff=0.5       ,			aligned_edge=__lowerCamelCase  )
          model.move_to([3, -1.0, 0]  )
          self.add(__lowerCamelCase  )
  
          _UpperCamelCase     :       Tuple  						=  []
          _UpperCamelCase     :       Any  						=  []
  
          for i, rect in enumerate(__lowerCamelCase  ):
                 rect.set_stroke(__lowerCamelCase  )
  
                 _UpperCamelCase     :       Any  						=  Rectangle(height=0.4_6 / 4       ,			width=0.4_6 / 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.0_2       ,			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  )
  
          _UpperCamelCase     :       List[Any]  						=  [meta_mem.copy() for i in range(6  )]
          _UpperCamelCase     :       int  						=  [meta_mem.copy() for i in range(6  )]
          _UpperCamelCase     :       int  						=  VGroup(*__lowerCamelCase  ).arrange(__lowerCamelCase       ,			buff=0  )
          _UpperCamelCase     :       Optional[int]  						=  VGroup(*__lowerCamelCase  ).arrange(__lowerCamelCase       ,			buff=0  )
          _UpperCamelCase     :       Optional[int]  						=  VGroup(__lowerCamelCase       ,			__lowerCamelCase  ).arrange(__lowerCamelCase       ,			buff=0  )
          _UpperCamelCase     :       Dict  						=  Text('Disk'       ,			font_size=2_4  )
          _UpperCamelCase     :       Optional[Any]  						=  Group(__lowerCamelCase       ,			__lowerCamelCase  ).arrange(__lowerCamelCase       ,			buff=0.5       ,			aligned_edge=__lowerCamelCase  )
          disk.move_to([-4.0, -1.2_5, 0]  )
          self.add(__lowerCamelCase       ,			__lowerCamelCase  )
  
          _UpperCamelCase     :       Union[str, Any]  						=  []
  
          for i in range(6  ):
                 _UpperCamelCase     :       Optional[int]  						=  fill.copy().set_fill(__lowerCamelCase       ,			opacity=0.8  )
                 target.move_to(disk_left_col_base[i]  ).scale(0.5  )
                 cpu_disk_arr.append(__lowerCamelCase  )
  
          self.add(*__lowerCamelCase  )
  
          _UpperCamelCase     :       Dict  						=  Square(side_length=2.2  )
          key.move_to([-5, 2, 0]  )
  
          _UpperCamelCase     :       Optional[int]  						=  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(__lowerCamelCase       ,			__lowerCamelCase  )
  
          _UpperCamelCase     :       Tuple  						=  MarkupText(
              f"<span fgcolor='{BLUE}'>●</span> Checkpoint"       ,			font_size=1_8       ,			)
  
          blue_text.next_to(__lowerCamelCase       ,			DOWN * 2.4       ,			aligned_edge=key_text.get_left()  )
          self.add(__lowerCamelCase  )
  
          _UpperCamelCase     :       int  						=  MarkupText(f"The offloaded weights are all sent to the CPU."       ,			font_size=2_4  )
          step_a.move_to([2, 2, 0]  )
  
          self.play(Write(__lowerCamelCase       ,			run_time=3  )  )
  
          for i in range(6  ):
                 _UpperCamelCase     :       str  						=  cpu_disk_arr[i]
                 _UpperCamelCase     :       Dict  						=  rect.copy().set_fill(__lowerCamelCase       ,			opacity=0.8  ).scale(2.0  )
                 cpa.generate_target()
                 cpa.target.move_to(model_base[i]  )
  
                 if i == 0:
                        rect.set_fill(__lowerCamelCase       ,			opacity=0.8  )
                        rect.generate_target()
                        rect.target.move_to(cpu_left_col_base[0]  ).scale(2.0  )
  
                        self.remove(
                            *__lowerCamelCase       ,			*__lowerCamelCase       ,			)
  
                 else:
                        rect.generate_target()
                        rect.target.move_to(cpu_left_col_base[i]  ).scale(2.0  )
                 self.play(MoveToTarget(__lowerCamelCase  )       ,			MoveToTarget(__lowerCamelCase  )       ,			model_base[i].animate.set_stroke(__lowerCamelCase  )  )
          self.play(FadeOut(__lowerCamelCase  )  )
  
          _UpperCamelCase     :       Any  						=  MarkupText(
              f"Finally, hooks are added to each weight in the model\nto transfer the weights from CPU to GPU\n\t\tand back when needed."       ,			font_size=2_4       ,			)
          step_a.move_to([2, 2, 0]  )
  
          self.play(Write(__lowerCamelCase       ,			run_time=3  )  )
  
          _UpperCamelCase     :       Any  						=  []
          _UpperCamelCase     :       str  						=  []
          for i in range(6  ):
                 _UpperCamelCase     :       Optional[Any]  						=  Arrow(start=__lowerCamelCase       ,			end=__lowerCamelCase       ,			color=__lowerCamelCase       ,			buff=0.5  )
                 a.next_to(model_base[i].get_left()       ,			__lowerCamelCase       ,			buff=0.2  )
                 arrows.append(__lowerCamelCase  )
                 animations.append(Write(__lowerCamelCase  )  )
          self.play(*__lowerCamelCase  )
          self.wait()
  
  
  
  
style_context: import tempfile
  import unittest
  
  import torch
  
  from diffusers import UNetaDConditionModel
  from diffusers.training_utils import EMAModel
  from diffusers.utils.testing_utils import enable_full_determinism, skip_mps, torch_device
  
  
  enable_full_determinism()
  
  
  
  
  class    __lowerCAmelCase  (       unittest.TestCase						):
  
  
  
   '''simple docstring'''
   _lowerCamelCase						:	str				    = 'hf-internal-testing/tiny-stable-diffusion-pipe'
   _lowerCamelCase						:	int				    = 1
   _lowerCamelCase						:	Any				    = 77
   _lowerCamelCase						:	List[Any]				    = 32
   _lowerCamelCase						:	List[str]				    = 4
   _lowerCamelCase						:	Dict				    = 64
   _lowerCamelCase						:	List[str]				    = torch.manual_seed(0						)
  
   def 					__snake_case					(			self       :		Any       ,			lowerCamelCase       :		Union[str, Any]=0.9_9_9_9  ) ->							Dict:
          _UpperCamelCase     :       int  						=  UNetaDConditionModel.from_pretrained(self.model_id       ,			subfolder='unet'  )
          _UpperCamelCase     :       Any  						=  unet.to(lowerCamelCase  )
          _UpperCamelCase     :       Union[str, Any]  						=  EMAModel(unet.parameters()       ,			decay=lowerCamelCase       ,			model_cls=lowerCamelCase       ,			model_config=unet.config  )
          return unet, ema_unet
  
   def 					__snake_case					(			self       :		Tuple  ) ->							Union[str, Any]:
          _UpperCamelCase     :       Tuple  						=  torch.randn(
              self.batch_size       ,			self.num_in_channels       ,			self.latent_height       ,			self.latent_width       ,			generator=self.generator  ).to(lowerCamelCase  )
          _UpperCamelCase     :       Union[str, Any]  						=  torch.randint(0       ,			1_0_0_0       ,			size=(self.batch_size,)       ,			generator=self.generator  ).to(lowerCamelCase  )
          _UpperCamelCase     :       int  						=  torch.randn(
              self.batch_size       ,			self.prompt_length       ,			self.text_encoder_hidden_dim       ,			generator=self.generator  ).to(lowerCamelCase  )
          return noisy_latents, timesteps, encoder_hidden_states
  
   def 					__snake_case					(			self       :		str       ,			lowerCamelCase       :		List[Any]  ) ->							List[Any]:
          _UpperCamelCase     :       int  						=  {}
          for k, param in unet.state_dict().items():
                 _UpperCamelCase     :       Optional[Any]  						=  torch.randn_like(lowerCamelCase  ) + (param * torch.randn_like(lowerCamelCase  ))
                 updated_state_dict.update({k: updated_param}  )
          unet.load_state_dict(lowerCamelCase  )
          return unet
  
   def 					__snake_case					(			self       :		List[Any]  ) ->							Tuple:
          _UpperCamelCase						,      _UpperCamelCase     :       Union[str, Any]  						=  self.get_models()
          # Take the first (hypothetical) EMA step.
          ema_unet.step(unet.parameters()  )
          assert ema_unet.optimization_step == 1
  
          # Take two more.
          for _ in range(2  ):
                 ema_unet.step(unet.parameters()  )
          assert ema_unet.optimization_step == 3
  
   def 					__snake_case					(			self       :		int  ) ->							Any:
          _UpperCamelCase						,      _UpperCamelCase     :       Any  						=  self.get_models()
          # Since the `unet` is not being updated (i.e., backprop'd)
          # there won't be any difference between the `params` of `unet`
          # and `ema_unet` even if we call `ema_unet.step(unet.parameters())`.
          ema_unet.step(unet.parameters()  )
          _UpperCamelCase     :       List[Any]  						=  list(unet.parameters()  )
          for s_param, param in zip(ema_unet.shadow_params       ,			lowerCamelCase  ):
                 assert torch.allclose(lowerCamelCase       ,			lowerCamelCase  )
  
          # The above holds true even if we call `ema.step()` multiple times since
          # `unet` params are still not being updated.
          for _ in range(4  ):
                 ema_unet.step(unet.parameters()  )
          for s_param, param in zip(ema_unet.shadow_params       ,			lowerCamelCase  ):
                 assert torch.allclose(lowerCamelCase       ,			lowerCamelCase  )
  
   def 					__snake_case					(			self       :		Dict  ) ->							Dict:
          _UpperCamelCase						,      _UpperCamelCase     :       Dict  						=  self.get_models()
          # Here we simulate the parameter updates for `unet`. Since there might
          # be some parameters which are initialized to zero we take extra care to
          # initialize their values to something non-zero before the multiplication.
          _UpperCamelCase     :       str  						=  self.simulate_backprop(lowerCamelCase  )
  
          # Take the EMA step.
          ema_unet.step(unet_pseudo_updated_step_one.parameters()  )
  
          # Now the EMA'd parameters won't be equal to the original model parameters.
          _UpperCamelCase     :       Optional[Any]  						=  list(unet_pseudo_updated_step_one.parameters()  )
          for s_param, param in zip(ema_unet.shadow_params       ,			lowerCamelCase  ):
                 assert ~torch.allclose(lowerCamelCase       ,			lowerCamelCase  )
  
          # Ensure this is the case when we take multiple EMA steps.
          for _ in range(4  ):
                 ema_unet.step(unet.parameters()  )
          for s_param, param in zip(ema_unet.shadow_params       ,			lowerCamelCase  ):
                 assert ~torch.allclose(lowerCamelCase       ,			lowerCamelCase  )
  
   def 					__snake_case					(			self       :		Union[str, Any]  ) ->							Tuple:
          # If we call EMA step after a backpropagation consecutively for two times,
          # the shadow params from those two steps should be different.
          _UpperCamelCase						,      _UpperCamelCase     :       Dict  						=  self.get_models()
  
          # First backprop + EMA
          _UpperCamelCase     :       Optional[Any]  						=  self.simulate_backprop(lowerCamelCase  )
          ema_unet.step(unet_step_one.parameters()  )
          _UpperCamelCase     :       Optional[Any]  						=  ema_unet.shadow_params
  
          # Second backprop + EMA
          _UpperCamelCase     :       Optional[int]  						=  self.simulate_backprop(lowerCamelCase  )
          ema_unet.step(unet_step_two.parameters()  )
          _UpperCamelCase     :       List[str]  						=  ema_unet.shadow_params
  
          for step_one, step_two in zip(lowerCamelCase       ,			lowerCamelCase  ):
                 assert ~torch.allclose(lowerCamelCase       ,			lowerCamelCase  )
  
   def 					__snake_case					(			self       :		str  ) ->							int:
          # If there's no decay even if there are backprops, EMA steps
          # won't take any effect i.e., the shadow params would remain the
          # same.
          _UpperCamelCase						,      _UpperCamelCase     :       List[str]  						=  self.get_models(decay=0.0  )
          _UpperCamelCase     :       Optional[Any]  						=  self.simulate_backprop(lowerCamelCase  )
          ema_unet.step(unet_step_one.parameters()  )
          _UpperCamelCase     :       Dict  						=  ema_unet.shadow_params
  
          _UpperCamelCase     :       List[str]  						=  self.simulate_backprop(lowerCamelCase  )
          ema_unet.step(unet_step_two.parameters()  )
          _UpperCamelCase     :       Union[str, Any]  						=  ema_unet.shadow_params
  
          for step_one, step_two in zip(lowerCamelCase       ,			lowerCamelCase  ):
                 assert torch.allclose(lowerCamelCase       ,			lowerCamelCase  )
  
   @skip_mps
   def 					__snake_case					(			self       :		Dict  ) ->							str:
          _UpperCamelCase						,      _UpperCamelCase     :       Optional[Any]  						=  self.get_models()
          _UpperCamelCase						,      _UpperCamelCase						,      _UpperCamelCase     :       Any  						=  self.get_dummy_inputs()
  
          with tempfile.TemporaryDirectory() as tmpdir:
                 ema_unet.save_pretrained(lowerCamelCase  )
                 _UpperCamelCase     :       List[str]  						=  UNetaDConditionModel.from_pretrained(lowerCamelCase       ,			model_cls=lowerCamelCase  )
                 _UpperCamelCase     :       Dict  						=  loaded_unet.to(unet.device  )
  
          # Since no EMA step has been performed the outputs should match.
          _UpperCamelCase     :       Optional[Any]  						=  unet(lowerCamelCase       ,			lowerCamelCase       ,			lowerCamelCase  ).sample
          _UpperCamelCase     :       List[str]  						=  loaded_unet(lowerCamelCase       ,			lowerCamelCase       ,			lowerCamelCase  ).sample
  
          assert torch.allclose(lowerCamelCase       ,			lowerCamelCase       ,			atol=1e-4  )
  
  
  
  