





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__	,   )













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							)






