

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'''







       snake_case    							=	KandinskyPriorPipeline
       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   lowerCamelCase__						(	self					:  List[Any]     )   -> List[Any]:







        '''simple docstring'''

        return 32

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







        '''simple docstring'''

        return 32

       @property
       def   lowerCamelCase__						(	self					:  Any     )   -> Tuple:







        '''simple docstring'''

        return self.time_input_dim

       @property
       def   lowerCamelCase__						(	self					:  Dict     )   -> Dict:







        '''simple docstring'''

        return self.time_input_dim * 4

       @property
       def   lowerCamelCase__						(	self					:  str     )   -> Dict:







        '''simple docstring'''

        return 100

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







        '''simple docstring'''

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

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







        '''simple docstring'''

        torch.manual_seed(0     )
        lowerCamelCase      						=    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(__snake_case     )

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







        '''simple docstring'''

        torch.manual_seed(0     )

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

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

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







        '''simple docstring'''

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

        lowerCamelCase      						=    CLIPVisionModelWithProjection(__snake_case     )
        return model

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







        '''simple docstring'''

        lowerCamelCase      						=    CLIPImageProcessor(
            crop_size=224						,	do_center_crop=__snake_case						,	do_normalize=__snake_case						,	do_resize=__snake_case						,	image_mean=[0.4814_5466, 0.457_8275, 0.4082_1073]						,	image_std=[0.2686_2954, 0.2613_0258, 0.2757_7711]						,	resample=3						,	size=224						,	)

        return image_processor

       def   lowerCamelCase__						(	self					:  List[Any]     )   -> Optional[int]:







        '''simple docstring'''

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

        lowerCamelCase      						=    UnCLIPScheduler(
            variance_type='fixed_small_log'						,	prediction_type='sample'						,	num_train_timesteps=1000						,	clip_sample=__snake_case						,	clip_sample_range=10.0						,	)

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

        return components

       def   lowerCamelCase__						(	self					:  Any						,	__snake_case					:  int						,	__snake_case					:  List[str]=0     )   -> Any:







        '''simple docstring'''

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

       def   lowerCamelCase__						(	self					:  str     )   -> Optional[Any]:







        '''simple docstring'''

        lowerCamelCase      						=    'cpu'

        lowerCamelCase      						=    self.get_dummy_components()

        lowerCamelCase      						=    self.pipeline_class(**__snake_case     )
        lowerCamelCase      						=    pipe.to(__snake_case     )

        pipe.set_progress_bar_config(disable=__snake_case     )

        lowerCamelCase      						=    pipe(**self.get_dummy_inputs(__snake_case     )     )
        lowerCamelCase      						=    output.image_embeds

        lowerCamelCase      						=    pipe(
            **self.get_dummy_inputs(__snake_case     )						,	return_dict=__snake_case						,	)[0]

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

        assert image.shape == (1, 32)

        lowerCamelCase      						=    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   lowerCamelCase__						(	self					:  List[str]     )   -> Any:







        '''simple docstring'''

        lowerCamelCase      						=    torch_device == 'cpu'
        lowerCamelCase      						=    True
        lowerCamelCase      						=    False

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





       @skip_mps
       def   lowerCamelCase__						(	self					:  List[Any]     )   -> int:







        '''simple docstring'''

        lowerCamelCase      						=    torch_device == 'cpu'
        lowerCamelCase      						=    False

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







from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_flax_available,
    is_sentencepiece_available,
    is_tf_available,
    is_tokenizers_available,
    is_torch_available,
)


_lowerCAmelCase :	List[Any]     				=    {
    'configuration_switch_transformers': [
        'SWITCH_TRANSFORMERS_PRETRAINED_CONFIG_ARCHIVE_MAP',
        'SwitchTransformersConfig',
        'SwitchTransformersOnnxConfig',
    ]
}

try:
     if not is_torch_available():
          raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
     pass
else:
     _lowerCAmelCase :	Tuple     				=    [
         'SWITCH_TRANSFORMERS_PRETRAINED_MODEL_ARCHIVE_LIST',
         'SwitchTransformersEncoderModel',
         'SwitchTransformersForConditionalGeneration',
         'SwitchTransformersModel',
         'SwitchTransformersPreTrainedModel',
         'SwitchTransformersTop1Router',
         'SwitchTransformersSparseMLP',
     ]


if TYPE_CHECKING:
     from .configuration_switch_transformers import (
         SWITCH_TRANSFORMERS_PRETRAINED_CONFIG_ARCHIVE_MAP,
         SwitchTransformersConfig,
         SwitchTransformersOnnxConfig,
     )

     try:
          if not is_torch_available():
               raise OptionalDependencyNotAvailable()
     except OptionalDependencyNotAvailable:
          pass
     else:
          from .modeling_switch_transformers import (
              SWITCH_TRANSFORMERS_PRETRAINED_MODEL_ARCHIVE_LIST,
              SwitchTransformersEncoderModel,
              SwitchTransformersForConditionalGeneration,
              SwitchTransformersModel,
              SwitchTransformersPreTrainedModel,
              SwitchTransformersSparseMLP,
              SwitchTransformersTopaRouter,
          )


else:
     import sys

     _lowerCAmelCase :	List[Any]     				=    _LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)




