


import gc
import inspect
import unittest

import torch
from parameterized import parameterized

from diffusers import PriorTransformer
from diffusers.utils import floats_tensor, slow, torch_all_close, torch_device
from diffusers.utils.testing_utils import enable_full_determinism

from .test_modeling_common import ModelTesterMixin


enable_full_determinism()






class     UpperCAmelCase__       (					A_   ,  unittest.TestCase		):



    """simple docstring"""






    UpperCAmelCase__						:		Tuple      =  PriorTransformer
    UpperCAmelCase__						:		List[str]      =  "hidden_states"






    @property
    def _a					(			self    )  ->       int:
         __UpperCamelCase    							=4
         __UpperCamelCase    							=8
         __UpperCamelCase    							=7

         __UpperCamelCase    							=floats_tensor((batch_size, embedding_dim)    ).to(A_    )

         __UpperCamelCase    							=floats_tensor((batch_size, embedding_dim)    ).to(A_    )
         __UpperCamelCase    							=floats_tensor((batch_size, num_embeddings, embedding_dim)    ).to(A_    )

         return {
             "hidden_states": hidden_states,
             "timestep": 2,
             "proj_embedding": proj_embedding,
             "encoder_hidden_states": encoder_hidden_states,
         }






    def _a					(			self					,   A_=0    )  ->       Dict:
         torch.manual_seed(A_    )
         __UpperCamelCase    							=4
         __UpperCamelCase    							=8
         __UpperCamelCase    							=7

         __UpperCamelCase    							=torch.randn((batch_size, embedding_dim)    ).to(A_    )

         __UpperCamelCase    							=torch.randn((batch_size, embedding_dim)    ).to(A_    )
         __UpperCamelCase    							=torch.randn((batch_size, num_embeddings, embedding_dim)    ).to(A_    )

         return {
             "hidden_states": hidden_states,
             "timestep": 2,
             "proj_embedding": proj_embedding,
             "encoder_hidden_states": encoder_hidden_states,
         }






    @property
    def _a					(			self    )  ->       Tuple:
         return (4, 8)






    @property
    def _a					(			self    )  ->       List[Any]:
         return (4, 8)






    def _a					(			self    )  ->       str:
         __UpperCamelCase    							={
             'num_attention_heads': 2,
             'attention_head_dim': 4,
             'num_layers': 2,
             'embedding_dim': 8,
             'num_embeddings': 7,
             'additional_embeddings': 4,
         }
         __UpperCamelCase    							=self.dummy_input
         return init_dict, inputs_dict






    def _a					(			self    )  ->       Union[str, Any]:
         __UpperCamelCase							,      __UpperCamelCase    							=PriorTransformer.from_pretrained(
             'hf-internal-testing/prior-dummy'					,   output_loading_info=A_    )
         self.assertIsNotNone(A_    )
         self.assertEqual(len(loading_info['missing_keys']    )					,   0    )

         model.to(A_    )
         __UpperCamelCase    							=model(**self.dummy_input    )[0]

         assert hidden_states is not None, "Make sure output is not None"






    def _a					(			self    )  ->       Any:
         __UpperCamelCase							,      __UpperCamelCase    							=self.prepare_init_args_and_inputs_for_common()

         __UpperCamelCase    							=self.model_class(**A_    )
         __UpperCamelCase    							=inspect.signature(model.forward    )
         # signature.parameters is an OrderedDict => so arg_names order is deterministic
         __UpperCamelCase    							=[*signature.parameters.keys()]

         __UpperCamelCase    							=['hidden_states', 'timestep']
         self.assertListEqual(arg_names[:2]					,   A_    )






    def _a					(			self    )  ->       Union[str, Any]:
         __UpperCamelCase    							=PriorTransformer.from_pretrained('hf-internal-testing/prior-dummy'    )
         __UpperCamelCase    							=model.to(A_    )

         if hasattr(A_					,   'set_default_attn_processor'    ):
              model.set_default_attn_processor()

         __UpperCamelCase    							=self.get_dummy_seed_input()

         with torch.no_grad():
              __UpperCamelCase    							=model(**A_    )[0]

         __UpperCamelCase    							=output[0, :5].flatten().cpu()
         print(A_    )

         # Since the VAE Gaussian prior's generator is seeded on the appropriate device,
         # the expected output slices are not the same for CPU and GPU.
         __UpperCamelCase    							=torch.tensor([-1.3436, -0.2870, 0.7538, 0.4368, -0.0239]    )
         self.assertTrue(torch_all_close(A_					,   A_					,   rtol=1E-2    )    )






@slow
class     UpperCAmelCase__       (					unittest.TestCase		):



    """simple docstring"""






    def _a					(			self					,   A_=1					,   A_=768					,   A_=77					,   A_=0    )  ->       Union[str, Any]:
         torch.manual_seed(A_    )
         __UpperCamelCase    							=batch_size
         __UpperCamelCase    							=embedding_dim
         __UpperCamelCase    							=num_embeddings

         __UpperCamelCase    							=torch.randn((batch_size, embedding_dim)    ).to(A_    )

         __UpperCamelCase    							=torch.randn((batch_size, embedding_dim)    ).to(A_    )
         __UpperCamelCase    							=torch.randn((batch_size, num_embeddings, embedding_dim)    ).to(A_    )

         return {
             "hidden_states": hidden_states,
             "timestep": 2,
             "proj_embedding": proj_embedding,
             "encoder_hidden_states": encoder_hidden_states,
         }






    def _a					(			self    )  ->       Any:
         # clean up the VRAM after each test
         super().tearDown()
         gc.collect()
         torch.cuda.empty_cache()






    @parameterized.expand(
        [
            # fmt: off
            [13, [-0.5861,  0.1283, -0.0931,  0.0882,  0.4476,  0.1329, -0.0498,  0.0640]],
            [37, [-0.4913,  0.0110, -0.0483,  0.0541,  0.4954, -0.0170,  0.0354,  0.1651]],
            # fmt: on
        ]    )
    def _a					(			self					,   A_					,   A_    )  ->       int:
         __UpperCamelCase    							=PriorTransformer.from_pretrained('kandinsky-community/kandinsky-2-1-prior'					,   subfolder='prior'    )
         model.to(A_    )
         __UpperCamelCase    							=self.get_dummy_seed_input(seed=A_    )

         with torch.no_grad():
              __UpperCamelCase    							=model(**A_    )[0]

         assert list(sample.shape    ) == [1, 768]

         __UpperCamelCase    							=sample[0, :8].flatten().cpu()
         print(A_    )
         __UpperCamelCase    							=torch.tensor(A_    )

         assert torch_all_close(A_					,   A_					,   atol=1E-3    )










"""simple docstring"""







from collections import UserDict
from typing import List, Union

from ..utils import (
    add_end_docstrings,
    is_tf_available,
    is_torch_available,
    is_vision_available,
    logging,
    requires_backends,
)
from .base import PIPELINE_INIT_ARGS, Pipeline


if is_vision_available():
							from PIL import Image

							from ..image_utils import load_image

if is_torch_available():
							from ..models.auto.modeling_auto import MODEL_FOR_ZERO_SHOT_IMAGE_CLASSIFICATION_MAPPING

if is_tf_available():
							from ..models.auto.modeling_tf_auto import TF_MODEL_FOR_ZERO_SHOT_IMAGE_CLASSIFICATION_MAPPING
							from ..tf_utils import stable_softmax

lowerCAmelCase_   						=   logging.get_logger(__name__)
@add_end_docstrings(A_    )
class        __A       (	A_    ):







		'''simple docstring'''





		def __init__(	self						:  List[str]			,**_snake_case						:  Dict						)					->		List[Any]:






				"""simple docstring"""



				super().__init__(**_snake_case						)

				requires_backends(self			,'''vision'''						)
				self.check_model_type(
				    TF_MODEL_FOR_ZERO_SHOT_IMAGE_CLASSIFICATION_MAPPING
				    if self.framework == '''tf'''
				    else MODEL_FOR_ZERO_SHOT_IMAGE_CLASSIFICATION_MAPPING						)







		def __call__(	self						:  Optional[int]			,_snake_case						:  Union[str, List[str], "Image", List["Image"]]			,**_snake_case						:  int						)					->		Optional[Any]:






				"""simple docstring"""



				return super().__call__(_snake_case			,**_snake_case						)







		def 					UpperCAmelCase    (	self						:  Dict			,**_snake_case						:  Optional[int]						)					->		List[Any]:






				"""simple docstring"""



				lowercase__   :    List[str]				=			{}
				if "candidate_labels" in kwargs:
						lowercase__   :    Any				=			kwargs['''candidate_labels''']
				if "hypothesis_template" in kwargs:
						lowercase__   :    Optional[Any]				=			kwargs['''hypothesis_template''']

				return preprocess_params, {}, {}







		def 					UpperCAmelCase    (	self						:  Optional[Any]			,_snake_case						:  Optional[int]			,_snake_case						:  Dict=None			,_snake_case						:  Union[str, Any]="This is a photo of {}."						)					->		List[str]:






				"""simple docstring"""



				lowercase__   :    List[Any]				=			load_image(_snake_case						)
				lowercase__   :    int				=			self.image_processor(images=[image]			,return_tensors=self.framework						)
				lowercase__   :    str				=			candidate_labels
				lowercase__   :    Dict				=			[hypothesis_template.format(_snake_case						) for x in candidate_labels]
				lowercase__   :    Any				=			self.tokenizer(_snake_case			,return_tensors=self.framework			,padding=_snake_case						)
				lowercase__   :    Optional[int]				=			[text_inputs]
				return inputs







		def 					UpperCAmelCase    (	self						:  Optional[Any]			,_snake_case						:  Optional[int]						)					->		List[Any]:






				"""simple docstring"""



				lowercase__   :    Optional[int]				=			model_inputs.pop('''candidate_labels'''						)
				lowercase__   :    Union[str, Any]				=			model_inputs.pop('''text_inputs'''						)
				if isinstance(text_inputs[0]			,_snake_case						):
						lowercase__   :    List[str]				=			text_inputs[0]
				else:
						# Batching case.
						lowercase__   :    int				=			text_inputs[0][0]

				lowercase__   :    Tuple				=			self.model(**_snake_case			,**_snake_case						)

				lowercase__   :    Union[str, Any]				=			{
				    '''candidate_labels''': candidate_labels,
				    '''logits''': outputs.logits_per_image,
				}
				return model_outputs







		def 					UpperCAmelCase    (	self						:  Any			,_snake_case						:  Tuple						)					->		Any:






				"""simple docstring"""



				lowercase__   :    Dict				=			model_outputs.pop('''candidate_labels'''						)
				lowercase__   :    Optional[Any]				=			model_outputs['''logits'''][0]
				if self.framework == "pt":
						lowercase__   :    Optional[int]				=			logits.softmax(dim=-1						).squeeze(-1						)
						lowercase__   :    Tuple				=			probs.tolist()
						if not isinstance(_snake_case			,_snake_case						):
								lowercase__   :    Any				=			[scores]
				elif self.framework == "tf":
						lowercase__   :    List[str]				=			stable_softmax(_snake_case			,axis=-1						)
						lowercase__   :    Optional[Any]				=			probs.numpy().tolist()
				else:
						raise ValueError(f"""Unsupported framework: {self.framework}"""						)

				lowercase__   :    Union[str, Any]				=			[
				    {'''score''': score, '''label''': candidate_label}
				    for score, candidate_label in sorted(zip(_snake_case			,_snake_case						)			,key=lambda _snake_case						: -x[0]						)
				]
				return result
