
'''simple docstring'''


import unittest

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

from diffusers import KandinskyVaaPriorPipeline, 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 						lowercase__	(  snake_case_,					unittest.TestCase   ):






   '''simple docstring'''







   _snake_case		   =							KandinskyVaaPriorPipeline
   _snake_case		   =							['''prompt''']
   _snake_case		   =							['''prompt''', '''negative_prompt''']
   _snake_case		   =							[
       '''num_images_per_prompt''',
       '''generator''',
       '''num_inference_steps''',
       '''latents''',
       '''negative_prompt''',
       '''guidance_scale''',
       '''output_type''',
       '''return_dict''',
   ]
   _snake_case		   =							False







   @property
   def       UpperCAmelCase				(			self   ):

      '''simple docstring'''







      return 3_2







   @property
   def       UpperCAmelCase				(			self   ):

      '''simple docstring'''







      return 3_2







   @property
   def       UpperCAmelCase				(			self   ):

      '''simple docstring'''







      return self.time_input_dim







   @property
   def       UpperCAmelCase				(			self   ):

      '''simple docstring'''







      return self.time_input_dim * 4







   @property
   def       UpperCAmelCase				(			self   ):

      '''simple docstring'''







      return 1_0_0







   @property
   def       UpperCAmelCase				(			self   ):

      '''simple docstring'''







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







   @property
   def       UpperCAmelCase				(			self   ):

      '''simple docstring'''







      torch.manual_seed(0   )
      UpperCamelCase       =  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-0_5 ,   num_attention_heads=4 ,   num_hidden_layers=5 ,   pad_token_id=1 ,   vocab_size=1_0_0_0 ,   )
      return CLIPTextModelWithProjection(lowerCamelCase__   )







   @property
   def       UpperCAmelCase				(			self   ):

      '''simple docstring'''







      torch.manual_seed(0   )

      UpperCamelCase       =  {
          '''num_attention_heads''': 2,
          '''attention_head_dim''': 1_2,
          '''embedding_dim''': self.text_embedder_hidden_size,
          '''num_layers''': 1,
      }

      UpperCamelCase       =  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
      UpperCamelCase       =  nn.Parameter(torch.ones(model.clip_std.shape   )   )
      return model







   @property
   def       UpperCAmelCase				(			self   ):

      '''simple docstring'''







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

      UpperCamelCase       =  CLIPVisionModelWithProjection(lowerCamelCase__   )
      return model







   @property
   def       UpperCAmelCase				(			self   ):

      '''simple docstring'''







      UpperCamelCase       =  CLIPImageProcessor(
          crop_size=2_2_4 ,   do_center_crop=lowerCamelCase__ ,   do_normalize=lowerCamelCase__ ,   do_resize=lowerCamelCase__ ,   image_mean=[0.48145466, 0.4578275, 0.40821073] ,   image_std=[0.26862954, 0.26130258, 0.27577711] ,   resample=3 ,   size=2_2_4 ,   )

      return image_processor







   def       UpperCAmelCase				(			self   ):

      '''simple docstring'''







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

      UpperCamelCase       =  UnCLIPScheduler(
          variance_type='''fixed_small_log''' ,   prediction_type='''sample''' ,   num_train_timesteps=1_0_0_0 ,   clip_sample=lowerCamelCase__ ,   clip_sample_range=10.0 ,   )

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

      return components







   def       UpperCAmelCase				(			self ,   lowerCamelCase__ ,   lowerCamelCase__=0   ):

      '''simple docstring'''







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







   def       UpperCAmelCase				(			self   ):

      '''simple docstring'''







      UpperCamelCase       =  '''cpu'''

      UpperCamelCase       =  self.get_dummy_components()

      UpperCamelCase       =  self.pipeline_class(**lowerCamelCase__   )
      UpperCamelCase       =  pipe.to(lowerCamelCase__   )

      pipe.set_progress_bar_config(disable=lowerCamelCase__   )

      UpperCamelCase       =  pipe(**self.get_dummy_inputs(lowerCamelCase__   )   )
      UpperCamelCase       =  output.image_embeds

      UpperCamelCase       =  pipe(
          **self.get_dummy_inputs(lowerCamelCase__   ) ,   return_dict=lowerCamelCase__ ,   )[0]

      UpperCamelCase       =  image[0, -1_0:]
      UpperCamelCase       =  image_from_tuple[0, -1_0:]

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

      UpperCamelCase       =  np.array(
          [-0.0532, 1.7120, 0.3656, -1.0852, -0.8946, -1.1756, 0.4348, 0.2482, 0.5146, -0.1156]   )

      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   ):

      '''simple docstring'''







      UpperCamelCase       =  torch_device == '''cpu'''
      UpperCamelCase       =  True
      UpperCamelCase       =  False

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







   @skip_mps
   def       UpperCAmelCase				(			self   ):

      '''simple docstring'''







      UpperCamelCase       =  torch_device == '''cpu'''
      UpperCamelCase       =  False

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







'''simple docstring'''


from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_tf_available, is_torch_available


snake_case_ :  List[Any]     =					{
    'configuration_data2vec_audio': ['DATA2VEC_AUDIO_PRETRAINED_CONFIG_ARCHIVE_MAP', 'Data2VecAudioConfig'],
    'configuration_data2vec_text': [
        'DATA2VEC_TEXT_PRETRAINED_CONFIG_ARCHIVE_MAP',
        'Data2VecTextConfig',
        'Data2VecTextOnnxConfig',
    ],
    'configuration_data2vec_vision': [
        'DATA2VEC_VISION_PRETRAINED_CONFIG_ARCHIVE_MAP',
        'Data2VecVisionConfig',
        'Data2VecVisionOnnxConfig',
    ],
}

try:
   if not is_torch_available():
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   pass
else:
   snake_case_ :  Union[str, Any]     =					[
       'DATA2VEC_AUDIO_PRETRAINED_MODEL_ARCHIVE_LIST',
       'Data2VecAudioForAudioFrameClassification',
       'Data2VecAudioForCTC',
       'Data2VecAudioForSequenceClassification',
       'Data2VecAudioForXVector',
       'Data2VecAudioModel',
       'Data2VecAudioPreTrainedModel',
   ]
   snake_case_ :  str     =					[
       'DATA2VEC_TEXT_PRETRAINED_MODEL_ARCHIVE_LIST',
       'Data2VecTextForCausalLM',
       'Data2VecTextForMaskedLM',
       'Data2VecTextForMultipleChoice',
       'Data2VecTextForQuestionAnswering',
       'Data2VecTextForSequenceClassification',
       'Data2VecTextForTokenClassification',
       'Data2VecTextModel',
       'Data2VecTextPreTrainedModel',
   ]
   snake_case_ :  Optional[int]     =					[
       'DATA2VEC_VISION_PRETRAINED_MODEL_ARCHIVE_LIST',
       'Data2VecVisionForImageClassification',
       'Data2VecVisionForMaskedImageModeling',
       'Data2VecVisionForSemanticSegmentation',
       'Data2VecVisionModel',
       'Data2VecVisionPreTrainedModel',
   ]

if is_tf_available():
   snake_case_ :  List[str]     =					[
       'TFData2VecVisionForImageClassification',
       'TFData2VecVisionForSemanticSegmentation',
       'TFData2VecVisionModel',
       'TFData2VecVisionPreTrainedModel',
   ]

if TYPE_CHECKING:
   from .configuration_dataavec_audio import DATA2VEC_AUDIO_PRETRAINED_CONFIG_ARCHIVE_MAP, DataaVecAudioConfig
   from .configuration_dataavec_text import (
       DATA2VEC_TEXT_PRETRAINED_CONFIG_ARCHIVE_MAP,
       DataaVecTextConfig,
       DataaVecTextOnnxConfig,
   )
   from .configuration_dataavec_vision import (
       DATA2VEC_VISION_PRETRAINED_CONFIG_ARCHIVE_MAP,
       DataaVecVisionConfig,
       DataaVecVisionOnnxConfig,
   )

   try:
      if not is_torch_available():
         raise OptionalDependencyNotAvailable()
   except OptionalDependencyNotAvailable:
      pass
   else:
      from .modeling_dataavec_audio import (
          DATA2VEC_AUDIO_PRETRAINED_MODEL_ARCHIVE_LIST,
          DataaVecAudioForAudioFrameClassification,
          DataaVecAudioForCTC,
          DataaVecAudioForSequenceClassification,
          DataaVecAudioForXVector,
          DataaVecAudioModel,
          DataaVecAudioPreTrainedModel,
      )
      from .modeling_dataavec_text import (
          DATA2VEC_TEXT_PRETRAINED_MODEL_ARCHIVE_LIST,
          DataaVecTextForCausalLM,
          DataaVecTextForMaskedLM,
          DataaVecTextForMultipleChoice,
          DataaVecTextForQuestionAnswering,
          DataaVecTextForSequenceClassification,
          DataaVecTextForTokenClassification,
          DataaVecTextModel,
          DataaVecTextPreTrainedModel,
      )
      from .modeling_dataavec_vision import (
          DATA2VEC_VISION_PRETRAINED_MODEL_ARCHIVE_LIST,
          DataaVecVisionForImageClassification,
          DataaVecVisionForMaskedImageModeling,
          DataaVecVisionForSemanticSegmentation,
          DataaVecVisionModel,
          DataaVecVisionPreTrainedModel,
      )
   if is_tf_available():
      from .modeling_tf_dataavec_vision import (
          TFDataaVecVisionForImageClassification,
          TFDataaVecVisionForSemanticSegmentation,
          TFDataaVecVisionModel,
          TFDataaVecVisionPreTrainedModel,
      )

else:
   import sys

   snake_case_ :  int     =					_LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)





