code: 
  
  
  """simple docstring"""
  
  
  
  import unittest
  
  import numpy as np
  import torch
  from torch import nn
  from transformers import (
      CLIPImageProcessor,
      CLIPTextConfig,
      CLIPTextModelWithProjection,
      CLIPTokenizer,
      CLIPVisionConfig,
      CLIPVisionModelWithProjection,
  )
  
  from diffusers import KandinskyPriorPipeline, PriorTransformer, UnCLIPScheduler
  from diffusers.utils import torch_device
  from diffusers.utils.testing_utils import enable_full_determinism, skip_mps
  
  from ..test_pipelines_common import PipelineTesterMixin
  
  
  enable_full_determinism()
  
  
  
  
  class 			UpperCAmelCase_						(		A_,     unittest.TestCase							):
      lowercase__		      =   KandinskyPriorPipeline
      lowercase__		      =   ['''prompt''']
      lowercase__		      =   ['''prompt''', '''negative_prompt''']
      lowercase__		      =   [
          '''num_images_per_prompt''',
          '''generator''',
          '''num_inference_steps''',
          '''latents''',
          '''negative_prompt''',
          '''guidance_scale''',
          '''output_type''',
          '''return_dict''',
      ]
      lowercase__		      =   False
  
  
  
  
      @property
      def 							__magic_name__							(					self		: Dict   )   ->							Optional[int]:
  
  
  
  
       '''simple docstring'''
  
       return 32
  
  
  
  
      @property
      def 							__magic_name__							(					self		: Optional[int]   )   ->							Tuple:
  
  
  
  
       '''simple docstring'''
  
       return 32
  
  
  
  
      @property
      def 							__magic_name__							(					self		: Union[str, Any]   )   ->							Any:
  
  
  
  
       '''simple docstring'''
  
       return self.time_input_dim
  
  
  
  
      @property
      def 							__magic_name__							(					self		: int   )   ->							Optional[Any]:
  
  
  
  
       '''simple docstring'''
  
       return self.time_input_dim * 4
  
  
  
  
      @property
      def 							__magic_name__							(					self		: Any   )   ->							Tuple:
  
  
  
  
       '''simple docstring'''
  
       return 100
  
  
  
  
      @property
      def 							__magic_name__							(					self		: List[Any]   )   ->							Union[str, Any]:
  
  
  
  
       '''simple docstring'''
  
       A__ =       CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"   )
       return tokenizer
  
  
  
  
      @property
      def 							__magic_name__							(					self		: Optional[int]   )   ->							Union[str, Any]:
  
  
  
  
       '''simple docstring'''
  
       torch.manual_seed(0   )
       A__ =       CLIPTextConfig(
           bos_token_id=0			,				eos_token_id=2			,				hidden_size=self.text_embedder_hidden_size			,				projection_dim=self.text_embedder_hidden_size			,				intermediate_size=37			,				layer_norm_eps=1e-05			,				num_attention_heads=4			,				num_hidden_layers=5			,				pad_token_id=1			,				vocab_size=1_000			,				)
       return CLIPTextModelWithProjection(snake_case_   )
  
  
  
  
      @property
      def 							__magic_name__							(					self		: List[str]   )   ->							str:
  
  
  
  
       '''simple docstring'''
  
       torch.manual_seed(0   )
  
       A__ =       {
           "num_attention_heads": 2,
           "attention_head_dim": 12,
           "embedding_dim": self.text_embedder_hidden_size,
           "num_layers": 1,
       }
  
       A__ =       PriorTransformer(**snake_case_   )
       # clip_std and clip_mean is initialized to be 0 so PriorTransformer.post_process_latents will always return 0 - set clip_std to be 1 so it won't return 0
       A__ =       nn.Parameter(torch.ones(model.clip_std.shape   )   )
       return model
  
  
  
  
      @property
      def 							__magic_name__							(					self		: List[str]   )   ->							List[Any]:
  
  
  
  
       '''simple docstring'''
  
       torch.manual_seed(0   )
       A__ =       CLIPVisionConfig(
           hidden_size=self.text_embedder_hidden_size			,				image_size=224			,				projection_dim=self.text_embedder_hidden_size			,				intermediate_size=37			,				num_attention_heads=4			,				num_channels=3			,				num_hidden_layers=5			,				patch_size=14			,				)
  
       A__ =       CLIPVisionModelWithProjection(snake_case_   )
       return model
  
  
  
  
      @property
      def 							__magic_name__							(					self		: Optional[int]   )   ->							Dict:
  
  
  
  
       '''simple docstring'''
  
       A__ =       CLIPImageProcessor(
           crop_size=224			,				do_center_crop=snake_case_			,				do_normalize=snake_case_			,				do_resize=snake_case_			,				image_mean=[0.4814_5466, 0.457_8275, 0.4082_1073]			,				image_std=[0.2686_2954, 0.2613_0258, 0.2757_7711]			,				resample=3			,				size=224			,				)
  
       return image_processor
  
  
  
  
      def 							__magic_name__							(					self		: List[Any]   )   ->							Dict:
  
  
  
  
       '''simple docstring'''
  
       A__ =       self.dummy_prior
       A__ =       self.dummy_image_encoder
       A__ =       self.dummy_text_encoder
       A__ =       self.dummy_tokenizer
       A__ =       self.dummy_image_processor
  
       A__ =       UnCLIPScheduler(
           variance_type="fixed_small_log"			,				prediction_type="sample"			,				num_train_timesteps=1_000			,				clip_sample=snake_case_			,				clip_sample_range=10.0			,				)
  
       A__ =       {
           "prior": prior,
           "image_encoder": image_encoder,
           "text_encoder": text_encoder,
           "tokenizer": tokenizer,
           "scheduler": scheduler,
           "image_processor": image_processor,
       }
  
       return components
  
  
  
  
      def 							__magic_name__							(					self		: List[str]			,				snake_case_		: Any			,				snake_case_		: Optional[Any]=0   )   ->							int:
  
  
  
  
       '''simple docstring'''
  
       if str(snake_case_   ).startswith("mps"   ):
        A__ =       torch.manual_seed(snake_case_   )
       else:
        A__ =       torch.Generator(device=snake_case_   ).manual_seed(snake_case_   )
       A__ =       {
           "prompt": "horse",
           "generator": generator,
           "guidance_scale": 4.0,
           "num_inference_steps": 2,
           "output_type": "np",
       }
       return inputs
  
  
  
  
      def 							__magic_name__							(					self		: Any   )   ->							str:
  
  
  
  
       '''simple docstring'''
  
       A__ =       "cpu"
  
       A__ =       self.get_dummy_components()
  
       A__ =       self.pipeline_class(**snake_case_   )
       A__ =       pipe.to(snake_case_   )
  
       pipe.set_progress_bar_config(disable=snake_case_   )
  
       A__ =       pipe(**self.get_dummy_inputs(snake_case_   )   )
       A__ =       output.image_embeds
  
       A__ =       pipe(
           **self.get_dummy_inputs(snake_case_   )			,				return_dict=snake_case_			,				)[0]
  
       A__ =       image[0, -10:]
       A__ =       image_from_tuple[0, -10:]
  
       assert image.shape == (1, 32)
  
       A__ =       np.array(
           [-0.0532, 1.7120, 0.3656, -1.0852, -0.8946, -1.1756, 0.4348, 0.2482, 0.5146, -0.1156]   )
  
       assert np.abs(image_slice.flatten() - expected_slice   ).max() < 1e-2
       assert np.abs(image_from_tuple_slice.flatten() - expected_slice   ).max() < 1e-2
  
  
  
  
      @skip_mps
      def 							__magic_name__							(					self		: List[Any]   )   ->							Dict:
  
  
  
  
       '''simple docstring'''
  
       A__ =       torch_device == "cpu"
       A__ =       True
       A__ =       False
  
       self._test_inference_batch_single_identical(
           test_max_difference=snake_case_			,				relax_max_difference=snake_case_			,				test_mean_pixel_difference=snake_case_			,				)
  
  
  
  
      @skip_mps
      def 							__magic_name__							(					self		: List[Any]   )   ->							Any:
  
  
  
  
       '''simple docstring'''
  
       A__ =       torch_device == "cpu"
       A__ =       False
  
       self._test_attention_slicing_forward_pass(
           test_max_difference=snake_case_			,				test_mean_pixel_difference=snake_case_			,				)
  
  
  
  
  
  
style_context: 
  
  
  
  
  """simple docstring"""
  
  
  
  from manim import *
  
  
  
  
  class 			UpperCAmelCase_						(		A_							):
  
  
  
  
      def 							__magic_name__							(					self		: Any   )   ->							Optional[int]:
  
  
  
  
       '''simple docstring'''
  
       A__ =       Rectangle(height=0.5			,				width=0.5   )
       A__ =       Rectangle(height=0.46			,				width=0.46   ).set_stroke(width=0   )
       A__ =       Rectangle(height=0.25			,				width=0.25   )
  
       A__ =       [mem.copy() for i in range(6   )]
       A__ =       [mem.copy() for i in range(6   )]
       A__ =       VGroup(*snake_case_   ).arrange(snake_case_			,				buff=0   )
       A__ =       VGroup(*snake_case_   ).arrange(snake_case_			,				buff=0   )
       A__ =       VGroup(snake_case_			,				snake_case_   ).arrange(snake_case_			,				buff=0   )
       A__ =       Text("CPU"			,				font_size=24   )
       A__ =       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_   )
  
       A__ =       [mem.copy() for i in range(4   )]
       A__ =       VGroup(*snake_case_   ).arrange(snake_case_			,				buff=0   )
       A__ =       Text("GPU"			,				font_size=24   )
       A__ =       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_   )
  
       A__ =       [mem.copy() for i in range(6   )]
       A__ =       VGroup(*snake_case_   ).arrange(snake_case_			,				buff=0   )
  
       A__ =       Text("Model"			,				font_size=24   )
       A__ =       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_   )
  
       A__ =       []
       A__ =       []
  
       for i, rect in enumerate(snake_case_   ):
        rect.set_stroke(snake_case_   )
  
        A__ =       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_   )
  
       A__ =       [meta_mem.copy() for i in range(6   )]
       A__ =       [meta_mem.copy() for i in range(6   )]
       A__ =       VGroup(*snake_case_   ).arrange(snake_case_			,				buff=0   )
       A__ =       VGroup(*snake_case_   ).arrange(snake_case_			,				buff=0   )
       A__ =       VGroup(snake_case_			,				snake_case_   ).arrange(snake_case_			,				buff=0   )
       A__ =       Text("Disk"			,				font_size=24   )
       A__ =       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_   )
  
       A__ =       []
  
       for i in range(6   ):
        A__ =       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_   )
  
       A__ =       Square(side_length=2.2   )
       key.move_to([-5, 2, 0]   )
  
       A__ =       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(snake_case_			,				snake_case_   )
  
       A__ =       MarkupText(
           F"""<span fgcolor='{BLUE}'>●</span> Checkpoint"""			,				font_size=18			,				)
  
       blue_text.next_to(snake_case_			,				DOWN * 2.4			,				aligned_edge=key_text.get_left()   )
       self.add(snake_case_   )
  
       A__ =       MarkupText(F"""The offloaded weights are all sent to the CPU."""			,				font_size=24   )
       step_a.move_to([2, 2, 0]   )
  
       self.play(Write(snake_case_			,				run_time=3   )   )
  
       for i in range(6   ):
        A__ =       cpu_disk_arr[i]
        A__ =       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_   )   )
  
       A__ =       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=24			,				)
       step_a.move_to([2, 2, 0]   )
  
       self.play(Write(snake_case_			,				run_time=3   )   )
  
       A__ =       []
       A__ =       []
       for i in range(6   ):
        A__ =       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()
  
  
  
  
  
  