






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        _lowerCAmelCase( UpperCAmelCase_							,			unittest.TestCase      ):





       """simple docstring"""







       a       :	str								=KandinskyPriorPipeline
       a       :	Dict								=['''prompt''']
       a       :	Tuple								=['''prompt''', '''negative_prompt''']
       a       :	Union[str, Any]								=[
           '''num_images_per_prompt''',
           '''generator''',
           '''num_inference_steps''',
           '''latents''',
           '''negative_prompt''',
           '''guidance_scale''',
           '''output_type''',
           '''return_dict''',
       ]
       a       :	int								=False





       @property
       def     _a					(      self							):
         return 3_2





       @property
       def     _a					(      self							):
         return 3_2





       @property
       def     _a					(      self							):
         return self.time_input_dim





       @property
       def     _a					(      self							):
         return self.time_input_dim * 4





       @property
       def     _a					(      self							):
         return 1_0_0





       @property
       def     _a					(      self							):
         UpperCamelCase_:      int					   =       CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'							)
         return tokenizer





       @property
       def     _a					(      self							):
         torch.manual_seed(0							)
         UpperCamelCase_:      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     _a					(      self							):
         torch.manual_seed(0							)

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

         UpperCamelCase_:      int					   =       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_:      List[Any]					   =       nn.Parameter(torch.ones(model.clip_std.shape							)							)
         return model





       @property
       def     _a					(      self							):
         torch.manual_seed(0							)
         UpperCamelCase_:      List[str]					   =       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_:      str					   =       CLIPVisionModelWithProjection(_lowerCamelCase							)
         return model





       @property
       def     _a					(      self							):
         UpperCamelCase_:      Any					   =       CLIPImageProcessor(
             crop_size=2_2_4						,    do_center_crop=_lowerCamelCase						,    do_normalize=_lowerCamelCase						,    do_resize=_lowerCamelCase						,    image_mean=[0.4_8_1_4_5_4_6_6, 0.4_5_7_8_2_7_5, 0.4_0_8_2_1_0_7_3]						,    image_std=[0.2_6_8_6_2_9_5_4, 0.2_6_1_3_0_2_5_8, 0.2_7_5_7_7_7_1_1]						,    resample=3						,    size=2_2_4						,    )

         return image_processor





       def     _a					(      self							):
         UpperCamelCase_:      Dict					   =       self.dummy_prior
         UpperCamelCase_:      Union[str, Any]					   =       self.dummy_image_encoder
         UpperCamelCase_:      Optional[int]					   =       self.dummy_text_encoder
         UpperCamelCase_:      List[Any]					   =       self.dummy_tokenizer
         UpperCamelCase_:      Union[str, Any]					   =       self.dummy_image_processor

         UpperCamelCase_:      Optional[int]					   =       UnCLIPScheduler(
             variance_type='fixed_small_log'						,    prediction_type='sample'						,    num_train_timesteps=1_0_0_0						,    clip_sample=_lowerCamelCase						,    clip_sample_range=1_0.0						,    )

         UpperCamelCase_:      Optional[Any]					   =       {
             'prior': prior,
             'image_encoder': image_encoder,
             'text_encoder': text_encoder,
             'tokenizer': tokenizer,
             'scheduler': scheduler,
             'image_processor': image_processor,
         }

         return components





       def     _a					(      self						,    _lowerCamelCase						,    _lowerCamelCase=0							):
         if str(_lowerCamelCase							).startswith('mps'							):
           UpperCamelCase_:      Dict					   =       torch.manual_seed(_lowerCamelCase							)
         else:
           UpperCamelCase_:      Optional[int]					   =       torch.Generator(device=_lowerCamelCase							).manual_seed(_lowerCamelCase							)
         UpperCamelCase_:      int					   =       {
             'prompt': 'horse',
             'generator': generator,
             'guidance_scale': 4.0,
             'num_inference_steps': 2,
             'output_type': 'np',
         }
         return inputs





       def     _a					(      self							):
         UpperCamelCase_:      List[Any]					   =       'cpu'

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

         UpperCamelCase_:      str					   =       self.pipeline_class(**_lowerCamelCase							)
         UpperCamelCase_:      str					   =       pipe.to(_lowerCamelCase							)

         pipe.set_progress_bar_config(disable=_lowerCamelCase							)

         UpperCamelCase_:      Optional[Any]					   =       pipe(**self.get_dummy_inputs(_lowerCamelCase							)							)
         UpperCamelCase_:      Any					   =       output.image_embeds

         UpperCamelCase_:      Tuple					   =       pipe(
             **self.get_dummy_inputs(_lowerCamelCase							)						,    return_dict=_lowerCamelCase						,    )[0]

         UpperCamelCase_:      Tuple					   =       image[0, -1_0:]
         UpperCamelCase_:      Any					   =       image_from_tuple[0, -1_0:]

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

         UpperCamelCase_:      Optional[Any]					   =       np.array(
             [-0.0_5_3_2, 1.7_1_2_0, 0.3_6_5_6, -1.0_8_5_2, -0.8_9_4_6, -1.1_7_5_6, 0.4_3_4_8, 0.2_4_8_2, 0.5_1_4_6, -0.1_1_5_6]							)

         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     _a					(      self							):
         UpperCamelCase_:      Any					   =       torch_device == 'cpu'
         UpperCamelCase_:      Tuple					   =       True
         UpperCamelCase_:      int					   =       False

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





       @skip_mps
       def     _a					(      self							):
         UpperCamelCase_:      List[Any]					   =       torch_device == 'cpu'
         UpperCamelCase_:      str					   =       False

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







# Lint as: python3

import hashlib
from typing import Union

def 						snake_case      (UpperCAmelCase__						) ->       bytes:
  if isinstance(UpperCAmelCase__			,					UpperCAmelCase__						):
    # Data already in bytes, returns as it as
    return hash_data
  elif isinstance(UpperCAmelCase__			,					UpperCAmelCase__						):
    # We keep the data as it as for it ot be later encoded to UTF-8
    # However replace `\\` with `/` for Windows compatibility
    UpperCamelCase_:      List[str]					   =       hash_data.replace('\\'			,					'/'						)
  elif isinstance(UpperCAmelCase__			,					UpperCAmelCase__						):
    UpperCamelCase_:      Union[str, Any]					   =       str(UpperCAmelCase__						)
  else:
    # If data is not of the required type, raise error
    raise InvalidKeyError(UpperCAmelCase__						)

  return hash_data.encode('utf-8'						)




class        _lowerCAmelCase( UpperCAmelCase_      ):





       """simple docstring"""







       def __init__(      self						,    _lowerCamelCase							):
         UpperCamelCase_:      int					   =       '\nFAILURE TO GENERATE DATASET: Invalid key type detected'
         UpperCamelCase_:      Dict					   =       f'''\nFound Key {hash_data} of type {type(_lowerCamelCase							)}'''
         UpperCamelCase_:      Any					   =       '\nKeys should be either str, int or bytes type'
         super().__init__(f'''{self.prefix}{self.err_msg}{self.suffix}'''							)







class        _lowerCAmelCase( UpperCAmelCase_      ):





       """simple docstring"""







       def __init__(      self						,    _lowerCamelCase						,    _lowerCamelCase						,    _lowerCamelCase=""							):
         UpperCamelCase_:      int					   =       key
         UpperCamelCase_:      Dict					   =       duplicate_key_indices
         UpperCamelCase_:      int					   =       fix_msg
         UpperCamelCase_:      List[Any]					   =       'Found multiple examples generated with the same key'
         if len(_lowerCamelCase							) <= 2_0:
           UpperCamelCase_:      Tuple					   =       f'''\nThe examples at index {', '.join(_lowerCamelCase							)} have the key {key}'''
         else:
           UpperCamelCase_:      Optional[int]					   =       f'''\nThe examples at index {', '.join(duplicate_key_indices[:2_0]							)}... ({len(_lowerCamelCase							) - 2_0} more) have the key {key}'''
         UpperCamelCase_:      List[Any]					   =       '\n' + fix_msg if fix_msg else ''
         super().__init__(f'''{self.prefix}{self.err_msg}{self.suffix}'''							)







class        _lowerCAmelCase:





       """simple docstring"""







       def __init__(      self						,    _lowerCamelCase							):
         UpperCamelCase_:      int					   =       hashlib.mda(_as_bytes(_lowerCamelCase							)							)





       def     _a					(      self						,    _lowerCamelCase							):
         UpperCamelCase_:      Optional[int]					   =       self._split_mda.copy()
         UpperCamelCase_:      str					   =       _as_bytes(_lowerCamelCase							)
         mda.update(_lowerCamelCase							)
         # Convert to integer with hexadecimal conversion
         return int(mda.hexdigest()						,    1_6							)