





"""simple docstring"""
import gc
import random
import unittest

import numpy as np
import torch

from diffusers import DDIMScheduler, KandinskyVaaPipeline, KandinskyVaaPriorPipeline, UNetaDConditionModel, VQModel
from diffusers.utils import floats_tensor, load_numpy, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..test_pipelines_common import PipelineTesterMixin, assert_mean_pixel_difference


enable_full_determinism()






class 					_SCREAMING_SNAKE_CASE(   A    ,	unittest.TestCase  ):
       SCREAMING_SNAKE_CASE_   :      int							=  KandinskyVaaPipeline
       SCREAMING_SNAKE_CASE_   :      int							=  [
           '''image_embeds''',
           '''negative_image_embeds''',
       ]
       SCREAMING_SNAKE_CASE_   :      List[str]							=  ['''image_embeds''', '''negative_image_embeds''']
       SCREAMING_SNAKE_CASE_   :      Dict							=  [
           '''generator''',
           '''height''',
           '''width''',
           '''latents''',
           '''guidance_scale''',
           '''num_inference_steps''',
           '''return_dict''',
           '''guidance_scale''',
           '''num_images_per_prompt''',
           '''output_type''',
           '''return_dict''',
       ]
       SCREAMING_SNAKE_CASE_   :      Any							=  False




       @property
       def     _UpperCamelCase     (						self							)			->      Any:





             """simple docstring"""






             return 32




       @property
       def     _UpperCamelCase     (						self							)			->      Optional[int]:





             """simple docstring"""






             return 32




       @property
       def     _UpperCamelCase     (						self							)			->      Dict:





             """simple docstring"""






             return self.time_input_dim




       @property
       def     _UpperCamelCase     (						self							)			->      str:





             """simple docstring"""






             return self.time_input_dim * 4




       @property
       def     _UpperCamelCase     (						self							)			->      List[Any]:





             """simple docstring"""






             return 1_00




       @property
       def     _UpperCamelCase     (						self							)			->      Union[str, Any]:





             """simple docstring"""






             torch.manual_seed(0							)

             __SCREAMING_SNAKE_CASE   :List[Any]						=  {
                 '''in_channels''': 4,
                 # Out channels is double in channels because predicts mean and variance
                 '''out_channels''': 8,
                 '''addition_embed_type''': '''image''',
                 '''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,
                 '''encoder_hid_dim''': self.text_embedder_hidden_size,
                 '''encoder_hid_dim_type''': '''image_proj''',
                 '''cross_attention_dim''': self.cross_attention_dim,
                 '''attention_head_dim''': 4,
                 '''resnet_time_scale_shift''': '''scale_shift''',
                 '''class_embed_type''': None,
             }

             __SCREAMING_SNAKE_CASE   :int						=  UNetaDConditionModel(**SCREAMING_SNAKE_CASE__							)
             return model




       @property
       def     _UpperCamelCase     (						self							)			->      Any:





             """simple docstring"""






             return {
                 "block_out_channels": [32, 64],
                 "down_block_types": ["DownEncoderBlock2D", "AttnDownEncoderBlock2D"],
                 "in_channels": 3,
                 "latent_channels": 4,
                 "layers_per_block": 1,
                 "norm_num_groups": 8,
                 "norm_type": "spatial",
                 "num_vq_embeddings": 12,
                 "out_channels": 3,
                 "up_block_types": [
                     "AttnUpDecoderBlock2D",
                     "UpDecoderBlock2D",
                 ],
                 "vq_embed_dim": 4,
             }




       @property
       def     _UpperCamelCase     (						self							)			->      List[Any]:





             """simple docstring"""






             torch.manual_seed(0							)
             __SCREAMING_SNAKE_CASE   :List[str]						=  VQModel(**self.dummy_movq_kwargs							)
             return model




       def     _UpperCamelCase     (						self							)			->      Union[str, Any]:





             """simple docstring"""






             __SCREAMING_SNAKE_CASE   :Optional[int]						=  self.dummy_unet
             __SCREAMING_SNAKE_CASE   :Optional[int]						=  self.dummy_movq

             __SCREAMING_SNAKE_CASE   :List[str]						=  DDIMScheduler(
                 num_train_timesteps=10_00		,beta_schedule='''linear'''		,beta_start=0.0_0_0_8_5		,beta_end=0.0_1_2		,clip_sample=SCREAMING_SNAKE_CASE__		,set_alpha_to_one=SCREAMING_SNAKE_CASE__		,steps_offset=1		,prediction_type='''epsilon'''		,thresholding=SCREAMING_SNAKE_CASE__		,)

             __SCREAMING_SNAKE_CASE   :Optional[Any]						=  {
                 '''unet''': unet,
                 '''scheduler''': scheduler,
                 '''movq''': movq,
             }
             return components




       def     _UpperCamelCase     (						self		,SCREAMING_SNAKE_CASE__		,SCREAMING_SNAKE_CASE__=0							)			->      List[Any]:





             """simple docstring"""






             __SCREAMING_SNAKE_CASE   :List[Any]						=  floats_tensor((1, self.text_embedder_hidden_size)		,rng=random.Random(SCREAMING_SNAKE_CASE__							)							).to(SCREAMING_SNAKE_CASE__							)
             __SCREAMING_SNAKE_CASE   :Dict						=  floats_tensor((1, self.text_embedder_hidden_size)		,rng=random.Random(seed + 1							)							).to(
                 SCREAMING_SNAKE_CASE__							)
             if str(SCREAMING_SNAKE_CASE__							).startswith('''mps'''							):
                   __SCREAMING_SNAKE_CASE   :List[Any]						=  torch.manual_seed(SCREAMING_SNAKE_CASE__							)
             else:
                   __SCREAMING_SNAKE_CASE   :Any						=  torch.Generator(device=SCREAMING_SNAKE_CASE__							).manual_seed(SCREAMING_SNAKE_CASE__							)
             __SCREAMING_SNAKE_CASE   :int						=  {
                 '''image_embeds''': image_embeds,
                 '''negative_image_embeds''': negative_image_embeds,
                 '''generator''': generator,
                 '''height''': 64,
                 '''width''': 64,
                 '''guidance_scale''': 4.0,
                 '''num_inference_steps''': 2,
                 '''output_type''': '''np''',
             }
             return inputs




       def     _UpperCamelCase     (						self							)			->      int:





             """simple docstring"""






             __SCREAMING_SNAKE_CASE   :Dict						=  '''cpu'''

             __SCREAMING_SNAKE_CASE   :Union[str, Any]						=  self.get_dummy_components()

             __SCREAMING_SNAKE_CASE   :Union[str, Any]						=  self.pipeline_class(**SCREAMING_SNAKE_CASE__							)
             __SCREAMING_SNAKE_CASE   :List[Any]						=  pipe.to(SCREAMING_SNAKE_CASE__							)

             pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE__							)

             __SCREAMING_SNAKE_CASE   :Optional[int]						=  pipe(**self.get_dummy_inputs(SCREAMING_SNAKE_CASE__							)							)
             __SCREAMING_SNAKE_CASE   :int						=  output.images

             __SCREAMING_SNAKE_CASE   :List[Any]						=  pipe(
                 **self.get_dummy_inputs(SCREAMING_SNAKE_CASE__							)		,return_dict=SCREAMING_SNAKE_CASE__		,)[0]

             __SCREAMING_SNAKE_CASE   :Optional[int]						=  image[0, -3:, -3:, -1]
             __SCREAMING_SNAKE_CASE   :Any						=  image_from_tuple[0, -3:, -3:, -1]

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

             __SCREAMING_SNAKE_CASE   :Optional[int]						=  np.array(
                 [0.6_2_3_7_9_7_6, 1.0, 0.3_6_4_4_1_3_3_2, 1.0, 0.7_0_6_3_9_6_3_4, 0.2_9_8_7_7_1_8_6, 0.8_5_6_5_2_1_2_5, 0.5_2_1_6_8_4_3, 0.5_4_4_5_4_0_4_6]							)

             assert (
                 np.abs(image_slice.flatten() - expected_slice							).max() < 1E-2
             ), f''' expected_slice {expected_slice}, but got {image_slice.flatten()}'''

             assert (
                 np.abs(image_from_tuple_slice.flatten() - expected_slice							).max() < 1E-2
             ), f''' expected_slice {expected_slice}, but got {image_from_tuple_slice.flatten()}'''






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




       def     _UpperCamelCase     (						self							)			->      Optional[int]:





             """simple docstring"""






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




       def     _UpperCamelCase     (						self							)			->      Any:





             """simple docstring"""






             __SCREAMING_SNAKE_CASE   :Tuple						=  load_numpy(
                 '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
                 '''/kandinskyv22/kandinskyv22_text2img_cat_fp16.npy'''							)

             __SCREAMING_SNAKE_CASE   :Optional[int]						=  KandinskyVaaPriorPipeline.from_pretrained(
                 '''kandinsky-community/kandinsky-2-2-prior'''		,torch_dtype=torch.floataa							)
             pipe_prior.to(SCREAMING_SNAKE_CASE__							)

             __SCREAMING_SNAKE_CASE   :Optional[int]						=  KandinskyVaaPipeline.from_pretrained(
                 '''kandinsky-community/kandinsky-2-2-decoder'''		,torch_dtype=torch.floataa							)
             __SCREAMING_SNAKE_CASE   :Dict						=  pipeline.to(SCREAMING_SNAKE_CASE__							)
             pipeline.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE__							)

             __SCREAMING_SNAKE_CASE   :List[Any]						=  '''red cat, 4k photo'''

             __SCREAMING_SNAKE_CASE   :Any						=  torch.Generator(device='''cuda'''							).manual_seed(0							)
             __SCREAMING_SNAKE_CASE				,__SCREAMING_SNAKE_CASE   :str						=  pipe_prior(
                 SCREAMING_SNAKE_CASE__		,generator=SCREAMING_SNAKE_CASE__		,num_inference_steps=5		,negative_prompt=''''''		,).to_tuple()

             __SCREAMING_SNAKE_CASE   :Optional[Any]						=  torch.Generator(device='''cuda'''							).manual_seed(0							)
             __SCREAMING_SNAKE_CASE   :List[Any]						=  pipeline(
                 image_embeds=SCREAMING_SNAKE_CASE__		,negative_image_embeds=SCREAMING_SNAKE_CASE__		,generator=SCREAMING_SNAKE_CASE__		,num_inference_steps=1_00		,output_type='''np'''		,)

             __SCREAMING_SNAKE_CASE   :Tuple						=  output.images[0]

             assert image.shape == (5_12, 5_12, 3)

             assert_mean_pixel_difference(SCREAMING_SNAKE_CASE__		,SCREAMING_SNAKE_CASE__							)






"""simple docstring"""
from ...configuration_utils import PretrainedConfig
from ...utils import logging


lowerCamelCase_								=		logging.get_logger(__name__)

lowerCamelCase_								=		{}






class 					_SCREAMING_SNAKE_CASE(   A  ):
       SCREAMING_SNAKE_CASE_   :      List[Any]							=  '''llama'''
       SCREAMING_SNAKE_CASE_   :      Optional[int]							=  ['''past_key_values''']




       def __init__(						self		,SCREAMING_SNAKE_CASE__=3_20_00		,SCREAMING_SNAKE_CASE__=40_96		,SCREAMING_SNAKE_CASE__=1_10_08		,SCREAMING_SNAKE_CASE__=32		,SCREAMING_SNAKE_CASE__=32		,SCREAMING_SNAKE_CASE__=None		,SCREAMING_SNAKE_CASE__="silu"		,SCREAMING_SNAKE_CASE__=20_48		,SCREAMING_SNAKE_CASE__=0.0_2		,SCREAMING_SNAKE_CASE__=1E-6		,SCREAMING_SNAKE_CASE__=True		,SCREAMING_SNAKE_CASE__=0		,SCREAMING_SNAKE_CASE__=1		,SCREAMING_SNAKE_CASE__=2		,SCREAMING_SNAKE_CASE__=1		,SCREAMING_SNAKE_CASE__=False		,SCREAMING_SNAKE_CASE__=None		,**SCREAMING_SNAKE_CASE__		,)			->      Union[str, Any]:





             """simple docstring"""






             __SCREAMING_SNAKE_CASE   :Optional[Any]						=  vocab_size
             __SCREAMING_SNAKE_CASE   :int						=  max_position_embeddings
             __SCREAMING_SNAKE_CASE   :List[str]						=  hidden_size
             __SCREAMING_SNAKE_CASE   :Tuple						=  intermediate_size
             __SCREAMING_SNAKE_CASE   :List[str]						=  num_hidden_layers
             __SCREAMING_SNAKE_CASE   :List[Any]						=  num_attention_heads

             # for backward compatibility
             if num_key_value_heads is None:
                   __SCREAMING_SNAKE_CASE   :Optional[int]						=  num_attention_heads

             __SCREAMING_SNAKE_CASE   :str						=  num_key_value_heads
             __SCREAMING_SNAKE_CASE   :Union[str, Any]						=  hidden_act
             __SCREAMING_SNAKE_CASE   :List[str]						=  initializer_range
             __SCREAMING_SNAKE_CASE   :Union[str, Any]						=  rms_norm_eps
             __SCREAMING_SNAKE_CASE   :Dict						=  pretraining_tp
             __SCREAMING_SNAKE_CASE   :Optional[Any]						=  use_cache
             __SCREAMING_SNAKE_CASE   :Optional[Any]						=  rope_scaling
             self._rope_scaling_validation()

             super().__init__(
                 pad_token_id=SCREAMING_SNAKE_CASE__		,bos_token_id=SCREAMING_SNAKE_CASE__		,eos_token_id=SCREAMING_SNAKE_CASE__		,tie_word_embeddings=SCREAMING_SNAKE_CASE__		,**SCREAMING_SNAKE_CASE__		,)




       def     _UpperCamelCase     (						self							)			->      List[Any]:





             """simple docstring"""






             if self.rope_scaling is None:
                   return

             if not isinstance(self.rope_scaling		,SCREAMING_SNAKE_CASE__							) or len(self.rope_scaling							) != 2:
                   raise ValueError(
                       '''`rope_scaling` must be a dictionary with with two fields, `name` and `factor`, '''
                       f'''got {self.rope_scaling}'''							)
             __SCREAMING_SNAKE_CASE   :Optional[Any]						=  self.rope_scaling.get('''type'''		,SCREAMING_SNAKE_CASE__							)
             __SCREAMING_SNAKE_CASE   :Optional[int]						=  self.rope_scaling.get('''factor'''		,SCREAMING_SNAKE_CASE__							)
             if rope_scaling_type is None or rope_scaling_type not in ["linear", "dynamic"]:
                   raise ValueError(
                       f'''`rope_scaling`\'s name field must be one of [\'linear\', \'dynamic\'], got {rope_scaling_type}'''							)
             if rope_scaling_factor is None or not isinstance(SCREAMING_SNAKE_CASE__		,SCREAMING_SNAKE_CASE__							) or rope_scaling_factor <= 1.0:
                   raise ValueError(f'''`rope_scaling`\'s factor field must be an float > 1, got {rope_scaling_factor}'''							)