code: 
  
  
  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 						A__      (					unittest.TestCase			):
  
  
  
  
  
  
  
      """simple docstring"""
  
      __magic_name__    						=	'hf-internal-testing/tiny-stable-diffusion-pipe'
      __magic_name__    						=	1
      __magic_name__    						=	77
      __magic_name__    						=	32
      __magic_name__    						=	4
      __magic_name__    						=	__magic_name__    						=	64
      __magic_name__    						=	torch.manual_seed(0			)
      def 	a_    ( self    ,      __snake_case=0.9999					):
           snake_case					=       UNetaDConditionModel.from_pretrained(self.model_id    ,      subfolder='''unet'''					)
           snake_case					=       unet.to(__UpperCamelCase					)
           snake_case					=       EMAModel(unet.parameters()    ,      decay=__UpperCamelCase    ,      model_cls=__UpperCamelCase    ,      model_config=unet.config					)
           return unet, ema_unet
      def 	a_    ( self					):
           snake_case					=       torch.randn(
               self.batch_size    ,      self.num_in_channels    ,      self.latent_height    ,      self.latent_width    ,      generator=self.generator					).to(__UpperCamelCase					)
           snake_case					=       torch.randint(0    ,      1_0_0_0    ,      size=(self.batch_size,)    ,      generator=self.generator					).to(__UpperCamelCase					)
           snake_case					=       torch.randn(
               self.batch_size    ,      self.prompt_length    ,      self.text_encoder_hidden_dim    ,      generator=self.generator					).to(__UpperCamelCase					)
           return noisy_latents, timesteps, encoder_hidden_states
      def 	a_    ( self    ,      __snake_case					):
           snake_case					=       {}
           for k, param in unet.state_dict().items():
                snake_case					=       torch.randn_like(__UpperCamelCase					) + (param * torch.randn_like(__UpperCamelCase					))
                updated_state_dict.update({k: updated_param}					)
           unet.load_state_dict(__UpperCamelCase					)
           return unet
      def 	a_    ( self					):
           snake_case				,						snake_case					=       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 	a_    ( self					):
           snake_case				,						snake_case					=       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()					)
           snake_case					=       list(unet.parameters()					)
           for s_param, param in zip(ema_unet.shadow_params    ,      __UpperCamelCase					):
                assert torch.allclose(__UpperCamelCase    ,      __UpperCamelCase					)
  
           # 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    ,      __UpperCamelCase					):
                assert torch.allclose(__UpperCamelCase    ,      __UpperCamelCase					)
      def 	a_    ( self					):
           snake_case				,						snake_case					=       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.
           snake_case					=       self.simulate_backprop(__UpperCamelCase					)
  
           # 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.
           snake_case					=       list(unet_pseudo_updated_step_one.parameters()					)
           for s_param, param in zip(ema_unet.shadow_params    ,      __UpperCamelCase					):
                assert ~torch.allclose(__UpperCamelCase    ,      __UpperCamelCase					)
  
           # 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    ,      __UpperCamelCase					):
                assert ~torch.allclose(__UpperCamelCase    ,      __UpperCamelCase					)
      def 	a_    ( self					):
           # If we call EMA step after a backpropagation consecutively for two times,
           # the shadow params from those two steps should be different.
           snake_case				,						snake_case					=       self.get_models()
  
           # First backprop + EMA
           snake_case					=       self.simulate_backprop(__UpperCamelCase					)
           ema_unet.step(unet_step_one.parameters()					)
           snake_case					=       ema_unet.shadow_params
  
           # Second backprop + EMA
           snake_case					=       self.simulate_backprop(__UpperCamelCase					)
           ema_unet.step(unet_step_two.parameters()					)
           snake_case					=       ema_unet.shadow_params
  
           for step_one, step_two in zip(__UpperCamelCase    ,      __UpperCamelCase					):
                assert ~torch.allclose(__UpperCamelCase    ,      __UpperCamelCase					)
      def 	a_    ( self					):
           # 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.
           snake_case				,						snake_case					=       self.get_models(decay=0.0					)
           snake_case					=       self.simulate_backprop(__UpperCamelCase					)
           ema_unet.step(unet_step_one.parameters()					)
           snake_case					=       ema_unet.shadow_params
  
           snake_case					=       self.simulate_backprop(__UpperCamelCase					)
           ema_unet.step(unet_step_two.parameters()					)
           snake_case					=       ema_unet.shadow_params
  
           for step_one, step_two in zip(__UpperCamelCase    ,      __UpperCamelCase					):
                assert torch.allclose(__UpperCamelCase    ,      __UpperCamelCase					)
  
  
  
  
      @skip_mps
      def 	a_    ( self					):
           snake_case				,						snake_case					=       self.get_models()
           snake_case				,						snake_case				,						snake_case					=       self.get_dummy_inputs()
  
           with tempfile.TemporaryDirectory() as tmpdir:
                ema_unet.save_pretrained(__UpperCamelCase					)
                snake_case					=       UNetaDConditionModel.from_pretrained(__UpperCamelCase    ,      model_cls=__UpperCamelCase					)
                snake_case					=       loaded_unet.to(unet.device					)
  
           # Since no EMA step has been performed the outputs should match.
           snake_case					=       unet(__UpperCamelCase    ,      __UpperCamelCase    ,      __UpperCamelCase					).sample
           snake_case					=       loaded_unet(__UpperCamelCase    ,      __UpperCamelCase    ,      __UpperCamelCase					).sample
  
           assert torch.allclose(__UpperCamelCase    ,      __UpperCamelCase    ,      atol=1E-4					)
  
  
  
  
  
  
  
style_context: 
  
  
  from manim import *
  
  class 						A__      (					snake_case__			):
  
  
  
  
  
  
  
      """simple docstring"""
  
      def 	a_    ( self					):
           snake_case					=       Rectangle(height=0.5    ,      width=0.5					)
           snake_case					=       Rectangle(height=0.46    ,      width=0.46					).set_stroke(width=0					)
           snake_case					=       Rectangle(height=0.25    ,      width=0.25					)
  
           snake_case					=       [mem.copy() for i in range(6					)]
           snake_case					=       [mem.copy() for i in range(6					)]
           snake_case					=       VGroup(*__snake_case					).arrange(__snake_case    ,      buff=0					)
           snake_case					=       VGroup(*__snake_case					).arrange(__snake_case    ,      buff=0					)
           snake_case					=       VGroup(__snake_case    ,      __snake_case					).arrange(__snake_case    ,      buff=0					)
           snake_case					=       Text('''CPU'''    ,      font_size=2_4					)
           snake_case					=       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					)
  
           snake_case					=       [mem.copy() for i in range(4					)]
           snake_case					=       VGroup(*__snake_case					).arrange(__snake_case    ,      buff=0					)
           snake_case					=       Text('''GPU'''    ,      font_size=2_4					)
           snake_case					=       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					)
  
           snake_case					=       [mem.copy() for i in range(6					)]
           snake_case					=       VGroup(*__snake_case					).arrange(__snake_case    ,      buff=0					)
  
           snake_case					=       Text('''Model'''    ,      font_size=2_4					)
           snake_case					=       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					)
  
           snake_case					=       []
           snake_case					=       []
  
           for i, rect in enumerate(__snake_case					):
                rect.set_stroke(__snake_case					)
  
                snake_case					=       Rectangle(height=0.46 / 4    ,      width=0.46 / 3					).set_stroke(width=0.0					).set_fill(__snake_case    ,      opacity=0.7					)
  
                if i == 0:
                     cpu_target.next_to(cpu_left_col_base[0].get_corner(DOWN + LEFT					)    ,      buff=0.02    ,      direction=__snake_case					)
                     cpu_target.set_x(cpu_target.get_x() + 0.1					)
                elif i == 3:
                     cpu_target.next_to(model_cpu_arr[0]    ,      direction=__snake_case    ,      buff=0.0					)
                else:
                     cpu_target.next_to(model_cpu_arr[i - 1]    ,      direction=__snake_case    ,      buff=0.0					)
                self.add(__snake_case					)
                model_cpu_arr.append(__snake_case					)
  
           self.add(*__snake_case    ,      *__snake_case					)
  
           snake_case					=       [meta_mem.copy() for i in range(6					)]
           snake_case					=       [meta_mem.copy() for i in range(6					)]
           snake_case					=       VGroup(*__snake_case					).arrange(__snake_case    ,      buff=0					)
           snake_case					=       VGroup(*__snake_case					).arrange(__snake_case    ,      buff=0					)
           snake_case					=       VGroup(__snake_case    ,      __snake_case					).arrange(__snake_case    ,      buff=0					)
           snake_case					=       Text('''Disk'''    ,      font_size=2_4					)
           snake_case					=       Group(__snake_case    ,      __snake_case					).arrange(__snake_case    ,      buff=0.5    ,      aligned_edge=__snake_case					)
           disk.move_to([-4.0, -1.25, 0]					)
           self.add(__snake_case    ,      __snake_case					)
  
           snake_case					=       []
  
           for i in range(6					):
                snake_case					=       fill.copy().set_fill(__snake_case    ,      opacity=0.8					)
                target.move_to(disk_left_col_base[i]					).scale(0.5					)
                cpu_disk_arr.append(__snake_case					)
  
           self.add(*__snake_case					)
  
           snake_case					=       Square(side_length=2.2					)
           key.move_to([-5, 2, 0]					)
  
           snake_case					=       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					)
  
           snake_case					=       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					)
  
           snake_case					=       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(__snake_case    ,      run_time=3					)					)
  
           for i in range(6					):
                snake_case					=       cpu_disk_arr[i]
                snake_case					=       rect.copy().set_fill(__snake_case    ,      opacity=0.8					).scale(2.0					)
                cpa.generate_target()
                cpa.target.move_to(model_base[i]					)
  
                if i == 0:
                     rect.set_fill(__snake_case    ,      opacity=0.8					)
                     rect.generate_target()
                     rect.target.move_to(cpu_left_col_base[0]					).scale(2.0					)
  
                     self.remove(
                         *__snake_case    ,      *__snake_case    ,      )
  
                else:
                     rect.generate_target()
                     rect.target.move_to(cpu_left_col_base[i]					).scale(2.0					)
                self.play(MoveToTarget(__snake_case					)    ,      MoveToTarget(__snake_case					)    ,      model_base[i].animate.set_stroke(__snake_case					)					)
           self.play(FadeOut(__snake_case					)					)
  
           snake_case					=       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(__snake_case    ,      run_time=3					)					)
  
           snake_case					=       []
           snake_case					=       []
           for i in range(6					):
                snake_case					=       Arrow(start=__snake_case    ,      end=__snake_case    ,      color=__snake_case    ,      buff=0.5					)
                a.next_to(model_base[i].get_left()    ,      __snake_case    ,      buff=0.2					)
                arrows.append(__snake_case					)
                animations.append(Write(__snake_case					)					)
           self.play(*__snake_case					)
           self.wait()
  
  
  
  
  