


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   A__							(						UpperCamelCase__      ,   unittest.TestCase     ):
						UpperCAmelCase									=  KandinskyPriorPipeline
						UpperCAmelCase									=  ["prompt"]
						UpperCAmelCase									=  ["prompt", "negative_prompt"]
						UpperCAmelCase									=  [
						    "num_images_per_prompt",
						    "generator",
						    "num_inference_steps",
						    "latents",
						    "negative_prompt",
						    "guidance_scale",
						    "output_type",
						    "return_dict",
						]
						UpperCAmelCase									=  False




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


										"""simple docstring"""





										return 32




						@property
						def 							__UpperCamelCase (				self    :					Optional[Any]     )							->   Optional[Any]:


										"""simple docstring"""





										return 32




						@property
						def 							__UpperCamelCase (				self    :					str     )							->   int:


										"""simple docstring"""





										return self.time_input_dim




						@property
						def 							__UpperCamelCase (				self    :					List[str]     )							->   Tuple:


										"""simple docstring"""





										return self.time_input_dim * 4




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


										"""simple docstring"""





										return 100




						@property
						def 							__UpperCamelCase (				self    :					Tuple     )							->   str:


										"""simple docstring"""





										_SCREAMING_SNAKE_CASE	     =CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''     )
										return tokenizer




						@property
						def 							__UpperCamelCase (				self    :					Dict     )							->   List[str]:


										"""simple docstring"""





										torch.manual_seed(0     )
										_SCREAMING_SNAKE_CASE	     =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=1000 ,			)
										return CLIPTextModelWithProjection(_a     )




						@property
						def 							__UpperCamelCase (				self    :					str     )							->   int:


										"""simple docstring"""





										torch.manual_seed(0     )

										_SCREAMING_SNAKE_CASE	     ={
										    '''num_attention_heads''': 2,
										    '''attention_head_dim''': 12,
										    '''embedding_dim''': self.text_embedder_hidden_size,
										    '''num_layers''': 1,
										}

										_SCREAMING_SNAKE_CASE	     =PriorTransformer(**_a     )
										# 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
										_SCREAMING_SNAKE_CASE	     =nn.Parameter(torch.ones(model.clip_std.shape     )     )
										return model




						@property
						def 							__UpperCamelCase (				self    :					Tuple     )							->   int:


										"""simple docstring"""





										torch.manual_seed(0     )
										_SCREAMING_SNAKE_CASE	     =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 ,			)

										_SCREAMING_SNAKE_CASE	     =CLIPVisionModelWithProjection(_a     )
										return model




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


										"""simple docstring"""





										_SCREAMING_SNAKE_CASE	     =CLIPImageProcessor(
										    crop_size=224 ,			do_center_crop=_a ,			do_normalize=_a ,			do_resize=_a ,			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=224 ,			)

										return image_processor




						def 							__UpperCamelCase (				self    :					Optional[Any]     )							->   Optional[Any]:


										"""simple docstring"""





										_SCREAMING_SNAKE_CASE	     =self.dummy_prior
										_SCREAMING_SNAKE_CASE	     =self.dummy_image_encoder
										_SCREAMING_SNAKE_CASE	     =self.dummy_text_encoder
										_SCREAMING_SNAKE_CASE	     =self.dummy_tokenizer
										_SCREAMING_SNAKE_CASE	     =self.dummy_image_processor

										_SCREAMING_SNAKE_CASE	     =UnCLIPScheduler(
										    variance_type='''fixed_small_log''' ,			prediction_type='''sample''' ,			num_train_timesteps=1000 ,			clip_sample=_a ,			clip_sample_range=10.0 ,			)

										_SCREAMING_SNAKE_CASE	     ={
										    '''prior''': prior,
										    '''image_encoder''': image_encoder,
										    '''text_encoder''': text_encoder,
										    '''tokenizer''': tokenizer,
										    '''scheduler''': scheduler,
										    '''image_processor''': image_processor,
										}

										return components




						def 							__UpperCamelCase (				self    :					str ,			_a    :					Dict ,			_a    :					List[str]=0     )							->   Dict:


										"""simple docstring"""





										if str(_a     ).startswith('''mps'''     ):
														_SCREAMING_SNAKE_CASE	     =torch.manual_seed(_a     )
										else:
														_SCREAMING_SNAKE_CASE	     =torch.Generator(device=_a     ).manual_seed(_a     )
										_SCREAMING_SNAKE_CASE	     ={
										    '''prompt''': '''horse''',
										    '''generator''': generator,
										    '''guidance_scale''': 4.0,
										    '''num_inference_steps''': 2,
										    '''output_type''': '''np''',
										}
										return inputs




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


										"""simple docstring"""





										_SCREAMING_SNAKE_CASE	     ='''cpu'''

										_SCREAMING_SNAKE_CASE	     =self.get_dummy_components()

										_SCREAMING_SNAKE_CASE	     =self.pipeline_class(**_a     )
										_SCREAMING_SNAKE_CASE	     =pipe.to(_a     )

										pipe.set_progress_bar_config(disable=_a     )

										_SCREAMING_SNAKE_CASE	     =pipe(**self.get_dummy_inputs(_a     )     )
										_SCREAMING_SNAKE_CASE	     =output.image_embeds

										_SCREAMING_SNAKE_CASE	     =pipe(
										    **self.get_dummy_inputs(_a     ) ,			return_dict=_a ,			)[0]

										_SCREAMING_SNAKE_CASE	     =image[0, -10:]
										_SCREAMING_SNAKE_CASE	     =image_from_tuple[0, -10:]

										assert image.shape == (1, 32)

										_SCREAMING_SNAKE_CASE	     =np.array(
										    [-0.05_32, 1.71_20, 0.36_56, -1.08_52, -0.89_46, -1.17_56, 0.43_48, 0.24_82, 0.51_46, -0.11_56]     )

										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 							__UpperCamelCase (				self    :					List[Any]     )							->   Optional[int]:


										"""simple docstring"""





										_SCREAMING_SNAKE_CASE	     =torch_device == '''cpu'''
										_SCREAMING_SNAKE_CASE	     =True
										_SCREAMING_SNAKE_CASE	     =False

										self._test_inference_batch_single_identical(
										    test_max_difference=_a ,			relax_max_difference=_a ,			test_mean_pixel_difference=_a ,			)




						@skip_mps
						def 							__UpperCamelCase (				self    :					Tuple     )							->   Optional[Any]:


										"""simple docstring"""





										_SCREAMING_SNAKE_CASE	     =torch_device == '''cpu'''
										_SCREAMING_SNAKE_CASE	     =False

										self._test_attention_slicing_forward_pass(
										    test_max_difference=_a ,			test_mean_pixel_difference=_a ,			)


import copy

from ...configuration_utils import PretrainedConfig
from ...utils import logging
from ..auto.configuration_auto import AutoConfig


UpperCamelCase__     :					Optional[Any]								=			logging.get_logger(__name__)

UpperCamelCase__     :					List[Any]								=			{
    "facebook/musicgen-small": "https://huggingface.co/facebook/musicgen-small/resolve/main/config.json",
    # See all Musicgen models at https://huggingface.co/models?filter=musicgen
}



class 					__snake_case     (			lowerCAmelCase__				):
  __lowerCAmelCase   :		Any         =     'musicgen_decoder'
  __lowerCAmelCase   :		Optional[Any]         =     ['past_key_values']

  def __init__(    self  ,						_A=2048  ,						_A=2048  ,						_A=24  ,						_A=4096  ,						_A=16  ,						_A=0.0  ,						_A=True  ,						_A="gelu"  ,						_A=1024  ,						_A=0.1  ,						_A=0.0  ,						_A=0.0  ,						_A=0.0_2  ,						_A=False  ,						_A=4  ,						_A=2048  ,						_A=2048  ,						_A=None  ,						_A=False  ,						**_A  ,						):
      SCREAMING_SNAKE_CASE_													=  vocab_size
      SCREAMING_SNAKE_CASE_													=  max_position_embeddings
      SCREAMING_SNAKE_CASE_													=  hidden_size
      SCREAMING_SNAKE_CASE_													=  ffn_dim
      SCREAMING_SNAKE_CASE_													=  num_hidden_layers
      SCREAMING_SNAKE_CASE_													=  num_attention_heads
      SCREAMING_SNAKE_CASE_													=  dropout
      SCREAMING_SNAKE_CASE_													=  attention_dropout
      SCREAMING_SNAKE_CASE_													=  activation_dropout
      SCREAMING_SNAKE_CASE_													=  activation_function
      SCREAMING_SNAKE_CASE_													=  initializer_factor
      SCREAMING_SNAKE_CASE_													=  layerdrop
      SCREAMING_SNAKE_CASE_													=  use_cache
      SCREAMING_SNAKE_CASE_													=  scale_embedding  # scale factor will be sqrt(d_model) if True
      SCREAMING_SNAKE_CASE_													=  num_codebooks
      super().__init__(
          pad_token_id=_A  ,						bos_token_id=_A  ,						eos_token_id=_A  ,						tie_word_embeddings=_A  ,						**_A  ,						)






class 					__snake_case     (			lowerCAmelCase__				):
  __lowerCAmelCase   :		int         =     'musicgen'
  __lowerCAmelCase   :		Optional[Any]         =     True

  def __init__(    self  ,						**_A):
      super().__init__(**_A)
      if "text_encoder" not in kwargs or "audio_encoder" not in kwargs or "decoder" not in kwargs:
          raise ValueError('Config has to be initialized with text_encoder, audio_encoder and decoder config')

      SCREAMING_SNAKE_CASE_													=  kwargs.pop('text_encoder')
      SCREAMING_SNAKE_CASE_													=  text_encoder_config.pop('model_type')

      SCREAMING_SNAKE_CASE_													=  kwargs.pop('audio_encoder')
      SCREAMING_SNAKE_CASE_													=  audio_encoder_config.pop('model_type')

      SCREAMING_SNAKE_CASE_													=  kwargs.pop('decoder')

      SCREAMING_SNAKE_CASE_													=  AutoConfig.for_model(_A  ,						**_A)
      SCREAMING_SNAKE_CASE_													=  AutoConfig.for_model(_A  ,						**_A)
      SCREAMING_SNAKE_CASE_													=  MusicgenDecoderConfig(**_A)
      SCREAMING_SNAKE_CASE_													=  True

  @classmethod
  def   lowerCAmelCase__   (    cls  ,						_A  ,						_A  ,						_A  ,						**_A  ,						):
      return cls(
          text_encoder=text_encoder_config.to_dict()  ,						audio_encoder=audio_encoder_config.to_dict()  ,						decoder=decoder_config.to_dict()  ,						**_A  ,						)


  def   lowerCAmelCase__   (    self):
      SCREAMING_SNAKE_CASE_													=  copy.deepcopy(self.__dict__)
      SCREAMING_SNAKE_CASE_													=  self.text_encoder.to_dict()
      SCREAMING_SNAKE_CASE_													=  self.audio_encoder.to_dict()
      SCREAMING_SNAKE_CASE_													=  self.decoder.to_dict()
      SCREAMING_SNAKE_CASE_													=  self.__class__.model_type
      return output


