code: 
  
  
  
  
  
  '''simple docstring'''
  
  
  
  
  
  import random
  import unittest
  
  import numpy as np
  import torch
  from PIL import Image
  from torch import nn
  from transformers import (
      CLIPImageProcessor,
      CLIPTextConfig,
      CLIPTextModelWithProjection,
      CLIPTokenizer,
      CLIPVisionConfig,
      CLIPVisionModelWithProjection,
  )
  
  from diffusers import KandinskyVaaPriorEmbaEmbPipeline, PriorTransformer, UnCLIPScheduler
  from diffusers.utils import floats_tensor, torch_device
  from diffusers.utils.testing_utils import enable_full_determinism, skip_mps
  
  from ..test_pipelines_common import PipelineTesterMixin
  
  
  enable_full_determinism()
  
  
  
  
  
  
  
  class        __lowerCAmelCase (					UpperCamelCase__		,       unittest.TestCase):
        _lowercase       :      int        =    KandinskyVaaPriorEmbaEmbPipeline
        _lowercase       :      Union[str, Any]        =    ["""prompt""", """image"""]
        _lowercase       :      Any        =    ["""prompt""", """image"""]
        _lowercase       :      Tuple        =    [
            """num_images_per_prompt""",
            """strength""",
            """generator""",
            """num_inference_steps""",
            """latents""",
            """negative_prompt""",
            """guidance_scale""",
            """output_type""",
            """return_dict""",
        ]
        _lowercase       :      Tuple        =    False
  
  
  
  
  
  
  
        @property
        def      _lowercase						(					self							)     ->		Union[str, Any]:
  
  
  
          '''simple docstring'''
  
  
  
  
          return 3_2
  
  
  
  
  
  
  
        @property
        def      _lowercase						(					self							)     ->		int:
  
  
  
          '''simple docstring'''
  
  
  
  
          return 3_2
  
  
  
  
  
  
  
        @property
        def      _lowercase						(					self							)     ->		List[Any]:
  
  
  
          '''simple docstring'''
  
  
  
  
          return self.time_input_dim
  
  
  
  
  
  
  
        @property
        def      _lowercase						(					self							)     ->		Any:
  
  
  
          '''simple docstring'''
  
  
  
  
          return self.time_input_dim * 4
  
  
  
  
  
  
  
        @property
        def      _lowercase						(					self							)     ->		str:
  
  
  
          '''simple docstring'''
  
  
  
  
          return 1_0_0
  
  
  
  
  
  
  
        @property
        def      _lowercase						(					self							)     ->		List[str]:
  
  
  
          '''simple docstring'''
  
  
  
  
          a__			:      Optional[Any]            =CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"							)
          return tokenizer
  
  
  
  
  
  
  
        @property
        def      _lowercase						(					self							)     ->		Tuple:
  
  
  
          '''simple docstring'''
  
  
  
  
          torch.manual_seed(0							)
          a__			:      Optional[int]            =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=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	,   )
          return CLIPTextModelWithProjection(lowerCAmelCase__							)
  
  
  
  
  
  
  
        @property
        def      _lowercase						(					self							)     ->		str:
  
  
  
          '''simple docstring'''
  
  
  
  
          torch.manual_seed(0							)
  
          a__			:      Optional[int]            ={
              "num_attention_heads": 2,
              "attention_head_dim": 1_2,
              "embedding_dim": self.text_embedder_hidden_size,
              "num_layers": 1,
          }
  
          a__			:      Dict            =PriorTransformer(**lowerCAmelCase__							)
          # 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__			:      List[str]            =nn.Parameter(torch.ones(model.clip_std.shape							)							)
          return model
  
  
  
  
  
  
  
        @property
        def      _lowercase						(					self							)     ->		str:
  
  
  
          '''simple docstring'''
  
  
  
  
          torch.manual_seed(0							)
          a__			:      Union[str, Any]            =CLIPVisionConfig(
              hidden_size=self.text_embedder_hidden_size	,   image_size=2_2_4	,   projection_dim=self.text_embedder_hidden_size	,   intermediate_size=3_7	,   num_attention_heads=4	,   num_channels=3	,   num_hidden_layers=5	,   patch_size=1_4	,   )
  
          a__			:      List[str]            =CLIPVisionModelWithProjection(lowerCAmelCase__							)
          return model
  
  
  
  
  
  
  
        @property
        def      _lowercase						(					self							)     ->		Optional[int]:
  
  
  
          '''simple docstring'''
  
  
  
  
          a__			:      Optional[Any]            =CLIPImageProcessor(
              crop_size=2_2_4	,   do_center_crop=lowerCAmelCase__	,   do_normalize=lowerCAmelCase__	,   do_resize=lowerCAmelCase__	,   image_mean=[0.48_14_54_66, 0.4_57_82_75, 0.40_82_10_73]	,   image_std=[0.26_86_29_54, 0.26_13_02_58, 0.27_57_77_11]	,   resample=3	,   size=2_2_4	,   )
  
          return image_processor
  
  
  
  
  
  
  
        def      _lowercase						(					self							)     ->		Dict:
  
  
  
          '''simple docstring'''
  
  
  
  
          a__			:      Optional[Any]            =self.dummy_prior
          a__			:      Union[str, Any]            =self.dummy_image_encoder
          a__			:      Optional[int]            =self.dummy_text_encoder
          a__			:      Union[str, Any]            =self.dummy_tokenizer
          a__			:      Optional[int]            =self.dummy_image_processor
  
          a__			:      Optional[int]            =UnCLIPScheduler(
              variance_type="fixed_small_log"	,   prediction_type="sample"	,   num_train_timesteps=1_0_0_0	,   clip_sample=lowerCAmelCase__	,   clip_sample_range=10.0	,   )
  
          a__			:      List[Any]            ={
              "prior": prior,
              "image_encoder": image_encoder,
              "text_encoder": text_encoder,
              "tokenizer": tokenizer,
              "scheduler": scheduler,
              "image_processor": image_processor,
          }
  
          return components
  
  
  
  
  
  
  
        def      _lowercase						(					self	,   lowerCAmelCase__	,   lowerCAmelCase__=0							)     ->		Tuple:
  
  
  
          '''simple docstring'''
  
  
  
  
          if str(lowerCAmelCase__							).startswith("mps"							):
            a__			:      List[Any]            =torch.manual_seed(lowerCAmelCase__							)
          else:
            a__			:      Union[str, Any]            =torch.Generator(device=lowerCAmelCase__							).manual_seed(lowerCAmelCase__							)
  
          a__			:      Optional[int]            =floats_tensor((1, 3, 6_4, 6_4)	,   rng=random.Random(lowerCAmelCase__							)							).to(lowerCAmelCase__							)
          a__			:      Tuple            =image.cpu().permute(0	,   2	,   3	,   1							)[0]
          a__			:      List[str]            =Image.fromarray(np.uinta(lowerCAmelCase__							)							).convert("RGB"							).resize((2_5_6, 2_5_6)							)
  
          a__			:      List[Any]            ={
              "prompt": "horse",
              "image": init_image,
              "strength": 0.5,
              "generator": generator,
              "guidance_scale": 4.0,
              "num_inference_steps": 2,
              "output_type": "np",
          }
          return inputs
  
  
  
  
  
  
  
        def      _lowercase						(					self							)     ->		List[Any]:
  
  
  
          '''simple docstring'''
  
  
  
  
          a__			:      Optional[int]            ="cpu"
  
          a__			:      Optional[int]            =self.get_dummy_components()
  
          a__			:      str            =self.pipeline_class(**lowerCAmelCase__							)
          a__			:      List[str]            =pipe.to(lowerCAmelCase__							)
  
          pipe.set_progress_bar_config(disable=lowerCAmelCase__							)
  
          a__			:      Tuple            =pipe(**self.get_dummy_inputs(lowerCAmelCase__							)							)
          a__			:      int            =output.image_embeds
  
          a__			:      Optional[Any]            =pipe(
              **self.get_dummy_inputs(lowerCAmelCase__							)	,   return_dict=lowerCAmelCase__	,   )[0]
  
          a__			:      str            =image[0, -1_0:]
          a__			:      Tuple            =image_from_tuple[0, -1_0:]
  
          assert image.shape == (1, 3_2)
  
          a__			:      Tuple            =np.array(
              [
                  0.1_07_12_84,
                  1.3_33_02_71,
                  0.61_26_02_23,
                  -0.6_69_10_65,
                  -0.3_84_68_52,
                  -1.0_30_36_61,
                  0.22_71_61_11,
                  0.03_34_89_01,
                  0.30_04_06_75,
                  -0.24_80_50_29,
              ]							)
  
          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      _lowercase						(					self							)     ->		Optional[int]:
  
  
  
          '''simple docstring'''
  
  
  
  
          a__			:      Optional[Any]            =torch_device == "cpu"
          a__			:      List[Any]            =True
          a__			:      int            =False
  
          self._test_inference_batch_single_identical(
              test_max_difference=lowerCAmelCase__	,   relax_max_difference=lowerCAmelCase__	,   test_mean_pixel_difference=lowerCAmelCase__	,   )
  
  
  
  
  
  
  
        @skip_mps
        def      _lowercase						(					self							)     ->		Union[str, Any]:
  
  
  
          '''simple docstring'''
  
  
  
  
          a__			:      Any            =torch_device == "cpu"
          a__			:      Dict            =False
  
          self._test_attention_slicing_forward_pass(
              test_max_difference=lowerCAmelCase__	,   test_mean_pixel_difference=lowerCAmelCase__	,   )
  
  
  
  
  
  
  
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):
        _lowercase       :      str        =    """hf-internal-testing/tiny-stable-diffusion-pipe"""
        _lowercase       :      Optional[int]        =    1
        _lowercase       :      List[str]        =    77
        _lowercase       :      int        =    32
        _lowercase       :      Dict        =    4
        _lowercase       :      str        =    64
        _lowercase       :      Optional[Any]        =    torch.manual_seed(0)
  
  
  
  
  
  
  
        def      _lowercase						(					self	,   lowerCAmelCase__=0.99_99							)     ->		Tuple:
  
  
  
          '''simple docstring'''
  
  
  
  
          a__			:      int            =UNetaDConditionModel.from_pretrained(self.model_id	,   subfolder="unet"							)
          a__			:      Any            =unet.to(lowerCAmelCase__							)
          a__			:      int            =EMAModel(unet.parameters()	,   decay=lowerCAmelCase__	,   model_cls=lowerCAmelCase__	,   model_config=unet.config							)
          return unet, ema_unet
  
  
  
  
  
  
  
        def      _lowercase						(					self							)     ->		Dict:
  
  
  
          '''simple docstring'''
  
  
  
  
          a__			:      str            =torch.randn(
              self.batch_size	,   self.num_in_channels	,   self.latent_height	,   self.latent_width	,   generator=self.generator							).to(lowerCAmelCase__							)
          a__			:      Any            =torch.randint(0	,   1_0_0_0	,   size=(self.batch_size,)	,   generator=self.generator							).to(lowerCAmelCase__							)
          a__			:      Tuple            =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      _lowercase						(					self	,   lowerCAmelCase__							)     ->		Any:
  
  
  
          '''simple docstring'''
  
  
  
  
          a__			:      Any            ={}
          for k, param in unet.state_dict().items():
            a__			:      Tuple            =torch.randn_like(lowerCAmelCase__							) + (param * torch.randn_like(lowerCAmelCase__							))
            updated_state_dict.update({k: updated_param}							)
          unet.load_state_dict(lowerCAmelCase__							)
          return unet
  
  
  
  
  
  
  
        def      _lowercase						(					self							)     ->		Tuple:
  
  
  
          '''simple docstring'''
  
  
  
  
          a__      ,							a__			:      List[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      _lowercase						(					self							)     ->		List[str]:
  
  
  
          '''simple docstring'''
  
  
  
  
          a__      ,							a__			:      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()							)
          a__			:      str            =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      _lowercase						(					self							)     ->		List[Any]:
  
  
  
          '''simple docstring'''
  
  
  
  
          a__      ,							a__			:      Optional[Any]            =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.
          a__			:      List[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.
          a__			:      Optional[int]            =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      _lowercase						(					self							)     ->		Union[str, Any]:
  
  
  
          '''simple docstring'''
  
  
  
  
          a__      ,							a__			:      Dict            =self.get_models()
  
          # First backprop + EMA
          a__			:      Tuple            =self.simulate_backprop(lowerCAmelCase__							)
          ema_unet.step(unet_step_one.parameters()							)
          a__			:      Dict            =ema_unet.shadow_params
  
          # Second backprop + EMA
          a__			:      str            =self.simulate_backprop(lowerCAmelCase__							)
          ema_unet.step(unet_step_two.parameters()							)
          a__			:      List[str]            =ema_unet.shadow_params
  
          for step_one, step_two in zip(lowerCAmelCase__	,   lowerCAmelCase__							):
            assert ~torch.allclose(lowerCAmelCase__	,   lowerCAmelCase__							)
  
  
  
  
  
  
  
        def      _lowercase						(					self							)     ->		List[str]:
  
  
  
          '''simple docstring'''
  
  
  
  
          a__      ,							a__			:      int            =self.get_models(decay=0.0							)
          a__			:      Optional[int]            =self.simulate_backprop(lowerCAmelCase__							)
          ema_unet.step(unet_step_one.parameters()							)
          a__			:      List[Any]            =ema_unet.shadow_params
  
          a__			:      Tuple            =self.simulate_backprop(lowerCAmelCase__							)
          ema_unet.step(unet_step_two.parameters()							)
          a__			:      List[Any]            =ema_unet.shadow_params
  
          for step_one, step_two in zip(lowerCAmelCase__	,   lowerCAmelCase__							):
            assert torch.allclose(lowerCAmelCase__	,   lowerCAmelCase__							)
  
  
  
  
  
  
  
        @skip_mps
        def      _lowercase						(					self							)     ->		Optional[Any]:
  
  
  
          '''simple docstring'''
  
  
  
  
          a__      ,							a__			:      Union[str, Any]            =self.get_models()
          a__      ,							a__      ,							a__			:      Tuple            =self.get_dummy_inputs()
  
          with tempfile.TemporaryDirectory() as tmpdir:
            ema_unet.save_pretrained(lowerCAmelCase__							)
            a__			:      List[str]            =UNetaDConditionModel.from_pretrained(lowerCAmelCase__	,   model_cls=lowerCAmelCase__							)
            a__			:      Tuple            =loaded_unet.to(unet.device							)
  
          # Since no EMA step has been performed the outputs should match.
          a__			:      Optional[Any]            =unet(lowerCAmelCase__	,   lowerCAmelCase__	,   lowerCAmelCase__							).sample
          a__			:      Optional[int]            =loaded_unet(lowerCAmelCase__	,   lowerCAmelCase__	,   lowerCAmelCase__							).sample
  
          assert torch.allclose(lowerCAmelCase__	,   lowerCAmelCase__	,   atol=1E-4							)
  
  
  
  
  
  
  