
'''simple docstring'''





import gc
import random
import unittest

import numpy as np
import torch
from transformers import (
    CLIPImageProcessor,
    CLIPTextConfig,
    CLIPTextModelWithProjection,
    CLIPTokenizer,
    CLIPVisionConfig,
    CLIPVisionModelWithProjection,
)

from diffusers import (
    DiffusionPipeline,
    UnCLIPImageVariationPipeline,
    UnCLIPScheduler,
    UNetaDConditionModel,
    UNetaDModel,
)
from diffusers.pipelines.unclip.text_proj import UnCLIPTextProjModel
from diffusers.utils import floats_tensor, load_numpy, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, load_image, require_torch_gpu, skip_mps

from ..pipeline_params import IMAGE_VARIATION_BATCH_PARAMS, IMAGE_VARIATION_PARAMS
from ..test_pipelines_common import PipelineTesterMixin, assert_mean_pixel_difference


enable_full_determinism()






class 							UpperCAmelCase_							(    A					,				unittest.TestCase				):

 '''simple docstring'''
 lowercase_					:    int				=							UnCLIPImageVariationPipeline
 lowercase_					:    Optional[int]				=							IMAGE_VARIATION_PARAMS - {"height", "width", "guidance_scale"}
 lowercase_					:    List[str]				=							IMAGE_VARIATION_BATCH_PARAMS

 lowercase_					:    Optional[Any]				=							[
     "generator",
     "return_dict",
     "decoder_num_inference_steps",
     "super_res_num_inference_steps",
 ]
 lowercase_					:    Union[str, Any]				=							False







 @property
 def 							UpperCamelCase  (					self		:  Dict ):



   '''simple docstring'''




   return 32







 @property
 def 							UpperCamelCase  (					self		:  List[str] ):



   '''simple docstring'''




   return 32







 @property
 def 							UpperCamelCase  (					self		:  Any ):



   '''simple docstring'''




   return self.time_input_dim







 @property
 def 							UpperCamelCase  (					self		:  Union[str, Any] ):



   '''simple docstring'''




   return self.time_input_dim * 4







 @property
 def 							UpperCamelCase  (					self		:  int ):



   '''simple docstring'''




   return 1_00







 @property
 def 							UpperCamelCase  (					self		:  Optional[int] ):



   '''simple docstring'''




   UpperCAmelCase__      :					List[Any]				       =							CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip" )
   return tokenizer







 @property
 def 							UpperCamelCase  (					self		:  Tuple ):



   '''simple docstring'''




   torch.manual_seed(0 )
   UpperCAmelCase__      :					List[str]				       =							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=10_00			,				)
   return CLIPTextModelWithProjection(snake_case__ )







 @property
 def 							UpperCamelCase  (					self		:  Any ):



   '''simple docstring'''




   torch.manual_seed(0 )
   UpperCAmelCase__      :					Optional[int]				       =							CLIPVisionConfig(
       hidden_size=self.text_embedder_hidden_size			,				projection_dim=self.text_embedder_hidden_size			,				num_hidden_layers=5			,				num_attention_heads=4			,				image_size=32			,				intermediate_size=37			,				patch_size=1			,				)
   return CLIPVisionModelWithProjection(snake_case__ )







 @property
 def 							UpperCamelCase  (					self		:  List[Any] ):



   '''simple docstring'''




   torch.manual_seed(0 )

   UpperCAmelCase__      :					Optional[int]				       =							{
       "clip_embeddings_dim": self.text_embedder_hidden_size,
       "time_embed_dim": self.time_embed_dim,
       "cross_attention_dim": self.cross_attention_dim,
   }

   UpperCAmelCase__      :					Union[str, Any]				       =							UnCLIPTextProjModel(**snake_case__ )
   return model







 @property
 def 							UpperCamelCase  (					self		:  Optional[int] ):



   '''simple docstring'''




   torch.manual_seed(0 )

   UpperCAmelCase__      :					str				       =							{
       "sample_size": 32,
       # RGB in channels
       "in_channels": 3,
       # Out channels is double in channels because predicts mean and variance
       "out_channels": 6,
       "down_block_types": ("ResnetDownsampleBlock2D", "SimpleCrossAttnDownBlock2D"),
       "up_block_types": ("SimpleCrossAttnUpBlock2D", "ResnetUpsampleBlock2D"),
       "mid_block_type": "UNetMidBlock2DSimpleCrossAttn",
       "block_out_channels": (self.block_out_channels_a, self.block_out_channels_a * 2),
       "layers_per_block": 1,
       "cross_attention_dim": self.cross_attention_dim,
       "attention_head_dim": 4,
       "resnet_time_scale_shift": "scale_shift",
       "class_embed_type": "identity",
   }

   UpperCAmelCase__      :					Any				       =							UNetaDConditionModel(**snake_case__ )
   return model







 @property
 def 							UpperCamelCase  (					self		:  List[Any] ):



   '''simple docstring'''




   return {
       "sample_size": 64,
       "layers_per_block": 1,
       "down_block_types": ("ResnetDownsampleBlock2D", "ResnetDownsampleBlock2D"),
       "up_block_types": ("ResnetUpsampleBlock2D", "ResnetUpsampleBlock2D"),
       "block_out_channels": (self.block_out_channels_a, self.block_out_channels_a * 2),
       "in_channels": 6,
       "out_channels": 3,
   }







 @property
 def 							UpperCamelCase  (					self		:  Optional[Any] ):



   '''simple docstring'''




   torch.manual_seed(0 )

   UpperCAmelCase__      :					str				       =							UNetaDModel(**self.dummy_super_res_kwargs )
   return model







 @property
 def 							UpperCamelCase  (					self		:  Dict ):



   '''simple docstring'''




   torch.manual_seed(1 )

   UpperCAmelCase__      :					List[Any]				       =							UNetaDModel(**self.dummy_super_res_kwargs )
   return model







 def 							UpperCamelCase  (					self		:  List[str] ):



   '''simple docstring'''




   UpperCAmelCase__      :					Union[str, Any]				       =							self.dummy_decoder
   UpperCAmelCase__      :					Optional[Any]				       =							self.dummy_text_proj
   UpperCAmelCase__      :					Dict				       =							self.dummy_text_encoder
   UpperCAmelCase__      :					List[Any]				       =							self.dummy_tokenizer
   UpperCAmelCase__      :					Any				       =							self.dummy_super_res_first
   UpperCAmelCase__      :					Optional[int]				       =							self.dummy_super_res_last

   UpperCAmelCase__      :					int				       =							UnCLIPScheduler(
       variance_type="learned_range"			,				prediction_type="epsilon"			,				num_train_timesteps=10_00			,				)

   UpperCAmelCase__      :					List[str]				       =							UnCLIPScheduler(
       variance_type="fixed_small_log"			,				prediction_type="epsilon"			,				num_train_timesteps=10_00			,				)

   UpperCAmelCase__      :					Union[str, Any]				       =							CLIPImageProcessor(crop_size=32			,				size=32 )

   UpperCAmelCase__      :					str				       =							self.dummy_image_encoder

   return {
       "decoder": decoder,
       "text_encoder": text_encoder,
       "tokenizer": tokenizer,
       "text_proj": text_proj,
       "feature_extractor": feature_extractor,
       "image_encoder": image_encoder,
       "super_res_first": super_res_first,
       "super_res_last": super_res_last,
       "decoder_scheduler": decoder_scheduler,
       "super_res_scheduler": super_res_scheduler,
   }







 def 							UpperCamelCase  (					self		:  List[Any]			,				snake_case__		:  Union[str, Any]			,				snake_case__		:  Union[str, Any]=0			,				snake_case__		:  Tuple=True ):



   '''simple docstring'''




   UpperCAmelCase__      :					Any				       =							floats_tensor((1, 3, 32, 32)			,				rng=random.Random(snake_case__ ) ).to(snake_case__ )
   if str(snake_case__ ).startswith("mps" ):
     UpperCAmelCase__      :					Tuple				       =							torch.manual_seed(snake_case__ )
   else:
     UpperCAmelCase__      :					Dict				       =							torch.Generator(device=snake_case__ ).manual_seed(snake_case__ )

   if pil_image:
     UpperCAmelCase__      :					Optional[Any]				       =							input_image * 0.5 + 0.5
     UpperCAmelCase__      :					Optional[Any]				       =							input_image.clamp(0			,				1 )
     UpperCAmelCase__      :					List[str]				       =							input_image.cpu().permute(0			,				2			,				3			,				1 ).float().numpy()
     UpperCAmelCase__      :					Optional[int]				       =							DiffusionPipeline.numpy_to_pil(snake_case__ )[0]

   return {
       "image": input_image,
       "generator": generator,
       "decoder_num_inference_steps": 2,
       "super_res_num_inference_steps": 2,
       "output_type": "np",
   }







 def 							UpperCamelCase  (					self		:  Optional[Any] ):



   '''simple docstring'''




   UpperCAmelCase__      :					str				       =							"cpu"

   UpperCAmelCase__      :					int				       =							self.get_dummy_components()

   UpperCAmelCase__      :					List[str]				       =							self.pipeline_class(**snake_case__ )
   UpperCAmelCase__      :					Optional[int]				       =							pipe.to(snake_case__ )

   pipe.set_progress_bar_config(disable=snake_case__ )

   UpperCAmelCase__      :					Dict				       =							self.get_dummy_inputs(snake_case__			,				pil_image=snake_case__ )

   UpperCAmelCase__      :					Dict				       =							pipe(**snake_case__ )
   UpperCAmelCase__      :					Tuple				       =							output.images

   UpperCAmelCase__      :					Tuple				       =							self.get_dummy_inputs(snake_case__			,				pil_image=snake_case__ )

   UpperCAmelCase__      :					Optional[Any]				       =							pipe(
       **snake_case__			,				return_dict=snake_case__			,				)[0]

   UpperCAmelCase__      :					Optional[Any]				       =							image[0, -3:, -3:, -1]
   UpperCAmelCase__      :					List[Any]				       =							image_from_tuple[0, -3:, -3:, -1]

   assert image.shape == (1, 64, 64, 3)

   UpperCAmelCase__      :					Optional[Any]				       =							np.array(
       [
           0.9997,
           0.0002,
           0.9997,
           0.9997,
           0.9969,
           0.0023,
           0.9997,
           0.9969,
           0.9970,
       ] )

   assert np.abs(image_slice.flatten() - expected_slice ).max() < 1e-2
   assert np.abs(image_from_tuple_slice.flatten() - expected_slice ).max() < 1e-2







 def 							UpperCamelCase  (					self		:  int ):



   '''simple docstring'''




   UpperCAmelCase__      :					Dict				       =							"cpu"

   UpperCAmelCase__      :					Any				       =							self.get_dummy_components()

   UpperCAmelCase__      :					List[Any]				       =							self.pipeline_class(**snake_case__ )
   UpperCAmelCase__      :					str				       =							pipe.to(snake_case__ )

   pipe.set_progress_bar_config(disable=snake_case__ )

   UpperCAmelCase__      :					List[str]				       =							self.get_dummy_inputs(snake_case__			,				pil_image=snake_case__ )

   UpperCAmelCase__      :					Union[str, Any]				       =							pipe(**snake_case__ )
   UpperCAmelCase__      :					str				       =							output.images

   UpperCAmelCase__      :					Tuple				       =							self.get_dummy_inputs(snake_case__			,				pil_image=snake_case__ )

   UpperCAmelCase__      :					Dict				       =							pipe(
       **snake_case__			,				return_dict=snake_case__			,				)[0]

   UpperCAmelCase__      :					List[str]				       =							image[0, -3:, -3:, -1]
   UpperCAmelCase__      :					Any				       =							image_from_tuple[0, -3:, -3:, -1]

   assert image.shape == (1, 64, 64, 3)

   UpperCAmelCase__      :					List[str]				       =							np.array([0.9997, 0.0003, 0.9997, 0.9997, 0.9970, 0.0024, 0.9997, 0.9971, 0.9971] )

   assert np.abs(image_slice.flatten() - expected_slice ).max() < 1e-2
   assert np.abs(image_from_tuple_slice.flatten() - expected_slice ).max() < 1e-2







 def 							UpperCamelCase  (					self		:  Optional[int] ):



   '''simple docstring'''




   UpperCAmelCase__      :					List[str]				       =							"cpu"

   UpperCAmelCase__      :					List[Any]				       =							self.get_dummy_components()

   UpperCAmelCase__      :					Tuple				       =							self.pipeline_class(**snake_case__ )
   UpperCAmelCase__      :					Tuple				       =							pipe.to(snake_case__ )

   pipe.set_progress_bar_config(disable=snake_case__ )

   UpperCAmelCase__      :					str				       =							self.get_dummy_inputs(snake_case__			,				pil_image=snake_case__ )
   UpperCAmelCase__      :					Any				       =							[
       pipeline_inputs["image"],
       pipeline_inputs["image"],
   ]

   UpperCAmelCase__      :					List[Any]				       =							pipe(**snake_case__ )
   UpperCAmelCase__      :					int				       =							output.images

   UpperCAmelCase__      :					Optional[int]				       =							self.get_dummy_inputs(snake_case__			,				pil_image=snake_case__ )
   UpperCAmelCase__      :					List[str]				       =							[
       tuple_pipeline_inputs["image"],
       tuple_pipeline_inputs["image"],
   ]

   UpperCAmelCase__      :					Optional[int]				       =							pipe(
       **snake_case__			,				return_dict=snake_case__			,				)[0]

   UpperCAmelCase__      :					List[str]				       =							image[0, -3:, -3:, -1]
   UpperCAmelCase__      :					Any				       =							image_from_tuple[0, -3:, -3:, -1]

   assert image.shape == (2, 64, 64, 3)

   UpperCAmelCase__      :					Any				       =							np.array(
       [
           0.9997,
           0.9989,
           0.0008,
           0.0021,
           0.9960,
           0.0018,
           0.0014,
           0.0002,
           0.9933,
       ] )

   assert np.abs(image_slice.flatten() - expected_slice ).max() < 1e-2
   assert np.abs(image_from_tuple_slice.flatten() - expected_slice ).max() < 1e-2







 def 							UpperCamelCase  (					self		:  Dict ):



   '''simple docstring'''




   UpperCAmelCase__      :					Optional[int]				       =							torch.device("cpu" )






   class 							UpperCAmelCase_							:

    '''simple docstring'''
    lowercase_					:    str				=							1



   UpperCAmelCase__      :					Tuple				       =							self.get_dummy_components()

   UpperCAmelCase__      :					Dict				       =							self.pipeline_class(**snake_case__ )
   UpperCAmelCase__      :					Union[str, Any]				       =							pipe.to(snake_case__ )

   pipe.set_progress_bar_config(disable=snake_case__ )

   UpperCAmelCase__      :					int				       =							torch.Generator(device=snake_case__ ).manual_seed(0 )
   UpperCAmelCase__      :					List[Any]				       =							pipe.decoder.dtype
   UpperCAmelCase__      :					List[Any]				       =							1

   UpperCAmelCase__      :					Optional[int]				       =							(
       batch_size,
       pipe.decoder.config.in_channels,
       pipe.decoder.config.sample_size,
       pipe.decoder.config.sample_size,
   )
   UpperCAmelCase__      :					List[str]				       =							pipe.prepare_latents(
       snake_case__			,				dtype=snake_case__			,				device=snake_case__			,				generator=snake_case__			,				latents=snake_case__			,				scheduler=DummyScheduler() )

   UpperCAmelCase__      :					Dict				       =							(
       batch_size,
       pipe.super_res_first.config.in_channels // 2,
       pipe.super_res_first.config.sample_size,
       pipe.super_res_first.config.sample_size,
   )
   UpperCAmelCase__      :					Union[str, Any]				       =							pipe.prepare_latents(
       snake_case__			,				dtype=snake_case__			,				device=snake_case__			,				generator=snake_case__			,				latents=snake_case__			,				scheduler=DummyScheduler() )

   UpperCAmelCase__      :					Tuple				       =							self.get_dummy_inputs(snake_case__			,				pil_image=snake_case__ )

   UpperCAmelCase__      :					List[Any]				       =							pipe(
       **snake_case__			,				decoder_latents=snake_case__			,				super_res_latents=snake_case__ ).images

   UpperCAmelCase__      :					Optional[int]				       =							self.get_dummy_inputs(snake_case__			,				pil_image=snake_case__ )
   # Don't pass image, instead pass embedding
   UpperCAmelCase__      :					str				       =							pipeline_inputs.pop("image" )
   UpperCAmelCase__      :					Any				       =							pipe.image_encoder(snake_case__ ).image_embeds

   UpperCAmelCase__      :					Any				       =							pipe(
       **snake_case__			,				decoder_latents=snake_case__			,				super_res_latents=snake_case__			,				image_embeddings=snake_case__			,				).images

   # make sure passing text embeddings manually is identical
   assert np.abs(img_out_a - img_out_a ).max() < 1e-4







 @skip_mps
 def 							UpperCamelCase  (					self		:  List[str] ):



   '''simple docstring'''




   UpperCAmelCase__      :					Tuple				       =							torch_device == "cpu"

   # Check is relaxed because there is not a torch 2.0 sliced attention added kv processor
   UpperCAmelCase__      :					Dict				       =							1e-2

   self._test_attention_slicing_forward_pass(
       test_max_difference=snake_case__			,				expected_max_diff=snake_case__ )







 @skip_mps
 def 							UpperCamelCase  (					self		:  Union[str, Any] ):



   '''simple docstring'''




   UpperCAmelCase__      :					Tuple				       =							torch_device == "cpu"
   UpperCAmelCase__      :					Tuple				       =							True
   UpperCAmelCase__      :					Dict				       =							[
       "decoder_num_inference_steps",
       "super_res_num_inference_steps",
   ]

   self._test_inference_batch_single_identical(
       test_max_difference=snake_case__			,				relax_max_difference=snake_case__			,				additional_params_copy_to_batched_inputs=snake_case__			,				)







 def 							UpperCamelCase  (					self		:  Dict ):



   '''simple docstring'''




   UpperCAmelCase__      :					str				       =							[
       "decoder_num_inference_steps",
       "super_res_num_inference_steps",
   ]

   if torch_device == "mps":
     # TODO: MPS errors with larger batch sizes
     UpperCAmelCase__      :					str				       =							[2, 3]
     self._test_inference_batch_consistent(
         batch_sizes=snake_case__			,				additional_params_copy_to_batched_inputs=snake_case__			,				)
   else:
     self._test_inference_batch_consistent(
         additional_params_copy_to_batched_inputs=snake_case__ )







 @skip_mps
 def 							UpperCamelCase  (					self		:  List[str] ):



   '''simple docstring'''




   return super().test_dict_tuple_outputs_equivalent()







 @skip_mps
 def 							UpperCamelCase  (					self		:  List[Any] ):



   '''simple docstring'''




   return super().test_save_load_local()







 @skip_mps
 def 							UpperCamelCase  (					self		:  int ):



   '''simple docstring'''




   return super().test_save_load_optional_components()






@slow
@require_torch_gpu
class 							UpperCAmelCase_							(    unittest.TestCase				):

 '''simple docstring'''







 def 							UpperCamelCase  (					self		:  Union[str, Any] ):



   '''simple docstring'''




   super().tearDown()
   gc.collect()
   torch.cuda.empty_cache()







 def 							UpperCamelCase  (					self		:  Tuple ):



   '''simple docstring'''




   UpperCAmelCase__      :					Tuple				       =							load_image(
       "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/unclip/cat.png" )
   UpperCAmelCase__      :					Optional[Any]				       =							load_numpy(
       "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
       "/unclip/karlo_v1_alpha_cat_variation_fp16.npy" )

   UpperCAmelCase__      :					List[str]				       =							UnCLIPImageVariationPipeline.from_pretrained(
       "kakaobrain/karlo-v1-alpha-image-variations"			,				torch_dtype=torch.floataa )
   UpperCAmelCase__      :					Optional[Any]				       =							pipeline.to(snake_case__ )
   pipeline.set_progress_bar_config(disable=snake_case__ )

   UpperCAmelCase__      :					Any				       =							torch.Generator(device="cpu" ).manual_seed(0 )
   UpperCAmelCase__      :					int				       =							pipeline(
       snake_case__			,				generator=snake_case__			,				output_type="np"			,				)

   UpperCAmelCase__      :					Optional[int]				       =							output.images[0]

   assert image.shape == (2_56, 2_56, 3)

   assert_mean_pixel_difference(snake_case__			,				snake_case__			,				15 )





'''simple docstring'''





import gc
import inspect
import unittest

import torch
from parameterized import parameterized

from diffusers import PriorTransformer
from diffusers.utils import floats_tensor, slow, torch_all_close, torch_device
from diffusers.utils.testing_utils import enable_full_determinism

from .test_modeling_common import ModelTesterMixin


enable_full_determinism()






class 							UpperCAmelCase_							(    A					,				unittest.TestCase				):

 '''simple docstring'''
 lowercase_					:    Tuple				=							PriorTransformer
 lowercase_					:    Dict				=							"hidden_states"







 @property
 def 							UpperCamelCase  (					self		:  List[Any] ):



   '''simple docstring'''




   UpperCAmelCase__      :					List[Any]				       =							4
   UpperCAmelCase__      :					List[Any]				       =							8
   UpperCAmelCase__      :					Union[str, Any]				       =							7

   UpperCAmelCase__      :					Tuple				       =							floats_tensor((batch_size, embedding_dim) ).to(snake_case__ )

   UpperCAmelCase__      :					List[Any]				       =							floats_tensor((batch_size, embedding_dim) ).to(snake_case__ )
   UpperCAmelCase__      :					Union[str, Any]				       =							floats_tensor((batch_size, num_embeddings, embedding_dim) ).to(snake_case__ )

   return {
       "hidden_states": hidden_states,
       "timestep": 2,
       "proj_embedding": proj_embedding,
       "encoder_hidden_states": encoder_hidden_states,
   }







 def 							UpperCamelCase  (					self		:  int			,				snake_case__		:  List[str]=0 ):



   '''simple docstring'''




   torch.manual_seed(snake_case__ )
   UpperCAmelCase__      :					str				       =							4
   UpperCAmelCase__      :					int				       =							8
   UpperCAmelCase__      :					Optional[int]				       =							7

   UpperCAmelCase__      :					Any				       =							torch.randn((batch_size, embedding_dim) ).to(snake_case__ )

   UpperCAmelCase__      :					Union[str, Any]				       =							torch.randn((batch_size, embedding_dim) ).to(snake_case__ )
   UpperCAmelCase__      :					List[str]				       =							torch.randn((batch_size, num_embeddings, embedding_dim) ).to(snake_case__ )

   return {
       "hidden_states": hidden_states,
       "timestep": 2,
       "proj_embedding": proj_embedding,
       "encoder_hidden_states": encoder_hidden_states,
   }







 @property
 def 							UpperCamelCase  (					self		:  Tuple ):



   '''simple docstring'''




   return (4, 8)







 @property
 def 							UpperCamelCase  (					self		:  Optional[int] ):



   '''simple docstring'''




   return (4, 8)







 def 							UpperCamelCase  (					self		:  int ):



   '''simple docstring'''




   UpperCAmelCase__      :					Optional[Any]				       =							{
       "num_attention_heads": 2,
       "attention_head_dim": 4,
       "num_layers": 2,
       "embedding_dim": 8,
       "num_embeddings": 7,
       "additional_embeddings": 4,
   }
   UpperCAmelCase__      :					Optional[Any]				       =							self.dummy_input
   return init_dict, inputs_dict







 def 							UpperCamelCase  (					self		:  Union[str, Any] ):



   '''simple docstring'''




   UpperCAmelCase__       ,			UpperCAmelCase__      :					List[str]				       =							PriorTransformer.from_pretrained(
       "hf-internal-testing/prior-dummy"			,				output_loading_info=snake_case__ )
   self.assertIsNotNone(snake_case__ )
   self.assertEqual(len(loading_info["missing_keys"] )			,				0 )

   model.to(snake_case__ )
   UpperCAmelCase__      :					List[Any]				       =							model(**self.dummy_input )[0]

   assert hidden_states is not None, "Make sure output is not None"







 def 							UpperCamelCase  (					self		:  Any ):



   '''simple docstring'''




   UpperCAmelCase__       ,			UpperCAmelCase__      :					int				       =							self.prepare_init_args_and_inputs_for_common()

   UpperCAmelCase__      :					Tuple				       =							self.model_class(**snake_case__ )
   UpperCAmelCase__      :					Optional[int]				       =							inspect.signature(model.forward )
   # signature.parameters is an OrderedDict => so arg_names order is deterministic
   UpperCAmelCase__      :					int				       =							[*signature.parameters.keys()]

   UpperCAmelCase__      :					Optional[Any]				       =							["hidden_states", "timestep"]
   self.assertListEqual(arg_names[:2]			,				snake_case__ )







 def 							UpperCamelCase  (					self		:  Tuple ):



   '''simple docstring'''




   UpperCAmelCase__      :					List[Any]				       =							PriorTransformer.from_pretrained("hf-internal-testing/prior-dummy" )
   UpperCAmelCase__      :					List[Any]				       =							model.to(snake_case__ )

   if hasattr(snake_case__			,				"set_default_attn_processor" ):
     model.set_default_attn_processor()

   UpperCAmelCase__      :					int				       =							self.get_dummy_seed_input()

   with torch.no_grad():
     UpperCAmelCase__      :					Tuple				       =							model(**snake_case__ )[0]

   UpperCAmelCase__      :					Any				       =							output[0, :5].flatten().cpu()
   print(snake_case__ )

   # Since the VAE Gaussian prior's generator is seeded on the appropriate device,
   # the expected output slices are not the same for CPU and GPU.
   UpperCAmelCase__      :					Optional[int]				       =							torch.tensor([-1.3436, -0.2870, 0.7538, 0.4368, -0.0239] )
   self.assertTrue(torch_all_close(snake_case__			,				snake_case__			,				rtol=1e-2 ) )






@slow
class 							UpperCAmelCase_							(    unittest.TestCase				):

 '''simple docstring'''







 def 							UpperCamelCase  (					self		:  Union[str, Any]			,				snake_case__		:  Optional[int]=1			,				snake_case__		:  Any=7_68			,				snake_case__		:  List[Any]=77			,				snake_case__		:  Tuple=0 ):



   '''simple docstring'''




   torch.manual_seed(snake_case__ )
   UpperCAmelCase__      :					Optional[int]				       =							batch_size
   UpperCAmelCase__      :					Any				       =							embedding_dim
   UpperCAmelCase__      :					Any				       =							num_embeddings

   UpperCAmelCase__      :					str				       =							torch.randn((batch_size, embedding_dim) ).to(snake_case__ )

   UpperCAmelCase__      :					Optional[Any]				       =							torch.randn((batch_size, embedding_dim) ).to(snake_case__ )
   UpperCAmelCase__      :					Any				       =							torch.randn((batch_size, num_embeddings, embedding_dim) ).to(snake_case__ )

   return {
       "hidden_states": hidden_states,
       "timestep": 2,
       "proj_embedding": proj_embedding,
       "encoder_hidden_states": encoder_hidden_states,
   }







 def 							UpperCamelCase  (					self		:  str ):



   '''simple docstring'''




   super().tearDown()
   gc.collect()
   torch.cuda.empty_cache()







 @parameterized.expand(
     [
         # fmt: off
         [13, [-0.5861,  0.1283, -0.0931,  0.0882,  0.4476,  0.1329, -0.0498,  0.0640]],
         [37, [-0.4913,  0.0110, -0.0483,  0.0541,  0.4954, -0.0170,  0.0354,  0.1651]],
         # fmt: on
     ] )
 def 							UpperCamelCase  (					self		:  Optional[Any]			,				snake_case__		:  List[str]			,				snake_case__		:  Optional[int] ):



   '''simple docstring'''




   UpperCAmelCase__      :					int				       =							PriorTransformer.from_pretrained("kandinsky-community/kandinsky-2-1-prior"			,				subfolder="prior" )
   model.to(snake_case__ )
   UpperCAmelCase__      :					List[Any]				       =							self.get_dummy_seed_input(seed=snake_case__ )

   with torch.no_grad():
     UpperCAmelCase__      :					Any				       =							model(**snake_case__ )[0]

   assert list(sample.shape ) == [1, 7_68]

   UpperCAmelCase__      :					List[str]				       =							sample[0, :8].flatten().cpu()
   print(snake_case__ )
   UpperCAmelCase__      :					Tuple				       =							torch.tensor(snake_case__ )

   assert torch_all_close(snake_case__			,				snake_case__			,				atol=1e-3 )



