

"""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 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       _UpperCAmelCase     (						__a      ,       unittest.TestCase):
			__a	:    List[Any]			   =							KandinskyPriorPipeline
			__a	:    List[Any]			   =							["""prompt"""]
			__a	:    List[str]			   =							["""prompt""", """negative_prompt"""]
			__a	:    Tuple			   =							[
			    """num_images_per_prompt""",
			    """generator""",
			    """num_inference_steps""",
			    """latents""",
			    """negative_prompt""",
			    """guidance_scale""",
			    """output_type""",
			    """return_dict""",
			]
			__a	:    str			   =							False
			@property
			def 							__snake_case	(						self       )      ->      Any:






				'''simple docstring'''







				return 32
			@property
			def 							__snake_case	(						self       )      ->      Union[str, Any]:






				'''simple docstring'''







				return 32
			@property
			def 							__snake_case	(						self       )      ->      Dict:






				'''simple docstring'''







				return self.time_input_dim
			@property
			def 							__snake_case	(						self       )      ->      Optional[Any]:






				'''simple docstring'''







				return self.time_input_dim * 4
			@property
			def 							__snake_case	(						self       )      ->      List[Any]:






				'''simple docstring'''







				return 1_00
			@property
			def 							__snake_case	(						self       )      ->      Optional[int]:






				'''simple docstring'''







				_UpperCAmelCase      :							List[str]       =   CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""       )
				return tokenizer
			@property
			def 							__snake_case	(						self       )      ->      int:






				'''simple docstring'''







				torch.manual_seed(0       )
				_UpperCAmelCase      :							List[Any]       =   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=10_00 ,						)
				return CLIPTextModelWithProjection(_A       )
			@property
			def 							__snake_case	(						self       )      ->      List[Any]:






				'''simple docstring'''







				torch.manual_seed(0       )

				_UpperCAmelCase      :							Optional[int]       =   {
				    """num_attention_heads""": 2,
				    """attention_head_dim""": 12,
				    """embedding_dim""": self.text_embedder_hidden_size,
				    """num_layers""": 1,
				}

				_UpperCAmelCase      :							List[Any]       =   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
				_UpperCAmelCase      :							List[Any]       =   nn.Parameter(torch.ones(model.clip_std.shape       )       )
				return model
			@property
			def 							__snake_case	(						self       )      ->      Dict:






				'''simple docstring'''







				torch.manual_seed(0       )
				_UpperCAmelCase      :							Optional[Any]       =   CLIPVisionConfig(
				    hidden_size=self.text_embedder_hidden_size ,						image_size=2_24 ,						projection_dim=self.text_embedder_hidden_size ,						intermediate_size=37 ,						num_attention_heads=4 ,						num_channels=3 ,						num_hidden_layers=5 ,						patch_size=14 ,						)

				_UpperCAmelCase      :							Tuple       =   CLIPVisionModelWithProjection(_A       )
				return model
			@property
			def 							__snake_case	(						self       )      ->      List[str]:






				'''simple docstring'''







				_UpperCAmelCase      :							str       =   CLIPImageProcessor(
				    crop_size=2_24 ,						do_center_crop=_A ,						do_normalize=_A ,						do_resize=_A ,						image_mean=[0.48145466, 0.4578275, 0.40821073] ,						image_std=[0.26862954, 0.26130258, 0.27577711] ,						resample=3 ,						size=2_24 ,						)

				return image_processor
			def 							__snake_case	(						self       )      ->      Dict:






				'''simple docstring'''







				_UpperCAmelCase      :							Optional[int]       =   self.dummy_prior
				_UpperCAmelCase      :							int       =   self.dummy_image_encoder
				_UpperCAmelCase      :							int       =   self.dummy_text_encoder
				_UpperCAmelCase      :							str       =   self.dummy_tokenizer
				_UpperCAmelCase      :							Any       =   self.dummy_image_processor

				_UpperCAmelCase      :							Any       =   UnCLIPScheduler(
				    variance_type="""fixed_small_log""" ,						prediction_type="""sample""" ,						num_train_timesteps=10_00 ,						clip_sample=_A ,						clip_sample_range=10.0 ,						)

				_UpperCAmelCase      :							Any       =   {
				    """prior""": prior,
				    """image_encoder""": image_encoder,
				    """text_encoder""": text_encoder,
				    """tokenizer""": tokenizer,
				    """scheduler""": scheduler,
				    """image_processor""": image_processor,
				}

				return components
			def 							__snake_case	(						self ,						_A ,						_A=0       )      ->      Optional[Any]:






				'''simple docstring'''







				if str(_A       ).startswith("""mps"""       ):
					_UpperCAmelCase      :							Any       =   torch.manual_seed(_A       )
				else:
					_UpperCAmelCase      :							Optional[int]       =   torch.Generator(device=_A       ).manual_seed(_A       )
				_UpperCAmelCase      :							Union[str, Any]       =   {
				    """prompt""": """horse""",
				    """generator""": generator,
				    """guidance_scale""": 4.0,
				    """num_inference_steps""": 2,
				    """output_type""": """np""",
				}
				return inputs
			def 							__snake_case	(						self       )      ->      Union[str, Any]:






				'''simple docstring'''







				_UpperCAmelCase      :							List[str]       =   """cpu"""

				_UpperCAmelCase      :							int       =   self.get_dummy_components()

				_UpperCAmelCase      :							Optional[Any]       =   self.pipeline_class(**_A       )
				_UpperCAmelCase      :							List[Any]       =   pipe.to(_A       )

				pipe.set_progress_bar_config(disable=_A       )

				_UpperCAmelCase      :							str       =   pipe(**self.get_dummy_inputs(_A       )       )
				_UpperCAmelCase      :							Any       =   output.image_embeds

				_UpperCAmelCase      :							Optional[int]       =   pipe(
				    **self.get_dummy_inputs(_A       ) ,						return_dict=_A ,						)[0]

				_UpperCAmelCase      :							int       =   image[0, -10:]
				_UpperCAmelCase      :							Optional[Any]       =   image_from_tuple[0, -10:]

				assert image.shape == (1, 32)

				_UpperCAmelCase      :							Optional[Any]       =   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 							__snake_case	(						self       )      ->      Tuple:






				'''simple docstring'''







				_UpperCAmelCase      :							str       =   torch_device == """cpu"""
				_UpperCAmelCase      :							Optional[Any]       =   True
				_UpperCAmelCase      :							List[Any]       =   False

				self._test_inference_batch_single_identical(
				    test_max_difference=_A ,						relax_max_difference=_A ,						test_mean_pixel_difference=_A ,						)
			@skip_mps
			def 							__snake_case	(						self       )      ->      Any:






				'''simple docstring'''







				_UpperCAmelCase      :							Any       =   torch_device == """cpu"""
				_UpperCAmelCase      :							List[str]       =   False

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








"""simple docstring"""






from collections import OrderedDict
from typing import Any, Mapping, Optional

from ... import PreTrainedTokenizer
from ...configuration_utils import PretrainedConfig
from ...file_utils import TensorType, is_torch_available
from ...onnx import OnnxConfig, OnnxConfigWithPast, OnnxSeqaSeqConfigWithPast
from ...onnx.utils import compute_effective_axis_dimension
from ...utils import logging


lowerCamelCase__			:					Dict							  =					logging.get_logger(__name__)

lowerCamelCase__			:					Optional[Any]							  =					{
    '''facebook/blenderbot-3B''': '''https://huggingface.co/facebook/blenderbot-3B/resolve/main/config.json''',
    # See all Blenderbot models at https://huggingface.co/models?filter=blenderbot
}






class       _UpperCAmelCase     (						__a):
			__a	:    Optional[int]			   =							"""blenderbot"""
			__a	:    Tuple			   =							["""past_key_values"""]
			__a	:    int			   =							{"""num_attention_heads""": """encoder_attention_heads""", """hidden_size""": """d_model"""}
			def __init__(						self ,						_A=80_08 ,						_A=1_28 ,						_A=2 ,						_A=1_02_40 ,						_A=32 ,						_A=24 ,						_A=1_02_40 ,						_A=32 ,						_A=0.0 ,						_A=0.0 ,						_A=True ,						_A=True ,						_A="gelu" ,						_A=25_60 ,						_A=0.1 ,						_A=0.0 ,						_A=0.0 ,						_A=0.02 ,						_A=1 ,						_A=False ,						_A=0 ,						_A=1 ,						_A=2 ,						_A=3 ,						_A=2 ,						**_A ,						)      ->      Union[str, Any]:






				'''simple docstring'''







				_UpperCAmelCase      :							Any       =   vocab_size
				_UpperCAmelCase      :							Dict       =   max_position_embeddings
				_UpperCAmelCase      :							Optional[int]       =   d_model
				_UpperCAmelCase      :							Optional[Any]       =   encoder_ffn_dim
				_UpperCAmelCase      :							Optional[Any]       =   encoder_layers
				_UpperCAmelCase      :							List[Any]       =   encoder_attention_heads
				_UpperCAmelCase      :							str       =   decoder_ffn_dim
				_UpperCAmelCase      :							int       =   decoder_layers
				_UpperCAmelCase      :							List[str]       =   decoder_attention_heads
				_UpperCAmelCase      :							List[Any]       =   dropout
				_UpperCAmelCase      :							List[str]       =   attention_dropout
				_UpperCAmelCase      :							int       =   activation_dropout
				_UpperCAmelCase      :							Union[str, Any]       =   activation_function
				_UpperCAmelCase      :							List[str]       =   init_std
				_UpperCAmelCase      :							Dict       =   encoder_layerdrop
				_UpperCAmelCase      :							Union[str, Any]       =   decoder_layerdrop
				_UpperCAmelCase      :							str       =   use_cache
				_UpperCAmelCase      :							Optional[int]       =   encoder_layers
				_UpperCAmelCase      :							Union[str, Any]       =   scale_embedding  # scale factor will be sqrt(d_model) if True

				super().__init__(
				    pad_token_id=_A ,						bos_token_id=_A ,						eos_token_id=_A ,						is_encoder_decoder=_A ,						decoder_start_token_id=_A ,						encoder_no_repeat_ngram_size=_A ,						forced_eos_token_id=_A ,						**_A ,						)






class       _UpperCAmelCase     (						__a):
			@property
			def 							__snake_case	(						self       )      ->      Mapping[str, Mapping[int, str]]:






				'''simple docstring'''







				if self.task in ["default", "seq2seq-lm"]:
					_UpperCAmelCase      :							Tuple       =   OrderedDict(
					    [
					        ("""input_ids""", {0: """batch""", 1: """encoder_sequence"""}),
					        ("""attention_mask""", {0: """batch""", 1: """encoder_sequence"""}),
					    ]       )
					if self.use_past:
						_UpperCAmelCase      :							List[str]       =   {0: """batch"""}
						_UpperCAmelCase      :							List[str]       =   {0: """batch""", 1: """past_decoder_sequence + sequence"""}
					else:
						_UpperCAmelCase      :							Optional[int]       =   {0: """batch""", 1: """decoder_sequence"""}
						_UpperCAmelCase      :							int       =   {0: """batch""", 1: """decoder_sequence"""}
					if self.use_past:
						self.fill_with_past_key_values_(_A ,						direction="""inputs"""       )
				elif self.task == "causal-lm":
					_UpperCAmelCase      :							List[str]       =   OrderedDict(
					    [
					        ("""input_ids""", {0: """batch""", 1: """encoder_sequence"""}),
					        ("""attention_mask""", {0: """batch""", 1: """encoder_sequence"""}),
					    ]       )
					if self.use_past:
						_UpperCAmelCase  ,				_UpperCAmelCase      :							Optional[Any]       =   self.num_layers
						for i in range(_A       ):
							_UpperCAmelCase      :							Any       =   {0: """batch""", 2: """past_sequence + sequence"""}
							_UpperCAmelCase      :							Dict       =   {0: """batch""", 2: """past_sequence + sequence"""}
				else:
					_UpperCAmelCase      :							List[Any]       =   OrderedDict(
					    [
					        ("""input_ids""", {0: """batch""", 1: """encoder_sequence"""}),
					        ("""attention_mask""", {0: """batch""", 1: """encoder_sequence"""}),
					        ("""decoder_input_ids""", {0: """batch""", 1: """decoder_sequence"""}),
					        ("""decoder_attention_mask""", {0: """batch""", 1: """decoder_sequence"""}),
					    ]       )

				return common_inputs
			@property
			# Copied from transformers.models.bart.configuration_bart.BartOnnxConfig.outputs
			def 							__snake_case	(						self       )      ->      Mapping[str, Mapping[int, str]]:






				'''simple docstring'''







				if self.task in ["default", "seq2seq-lm"]:
					_UpperCAmelCase      :							Optional[Any]       =   super().outputs
				else:
					_UpperCAmelCase      :							Optional[int]       =   super(_A ,						self       ).outputs
					if self.use_past:
						_UpperCAmelCase  ,				_UpperCAmelCase      :							int       =   self.num_layers
						for i in range(_A       ):
							_UpperCAmelCase      :							Tuple       =   {0: """batch""", 2: """past_sequence + sequence"""}
							_UpperCAmelCase      :							Optional[Any]       =   {0: """batch""", 2: """past_sequence + sequence"""}
				return common_outputs
			def 							__snake_case	(						self ,						_A ,						_A = -1 ,						_A = -1 ,						_A = False ,						_A = None ,						)      ->      Mapping[str, Any]:






				'''simple docstring'''







				_UpperCAmelCase      :							Dict       =   self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
				    _A ,						_A ,						_A ,						_A ,						_A       )
				# Generate decoder inputs
				_UpperCAmelCase      :							Any       =   seq_length if not self.use_past else 1
				_UpperCAmelCase      :							List[str]       =   self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
				    _A ,						_A ,						_A ,						_A ,						_A       )
				_UpperCAmelCase      :							str       =   {f'''decoder_{name}''': tensor for name, tensor in decoder_inputs.items()}
				_UpperCAmelCase      :							Dict       =   dict(**_A ,						**_A       )

				if self.use_past:
					if not is_torch_available():
						raise ValueError("""Cannot generate dummy past_keys inputs without PyTorch installed."""       )
					else:
						import torch
					_UpperCAmelCase  ,				_UpperCAmelCase      :							Tuple       =   common_inputs["""input_ids"""].shape
					_UpperCAmelCase      :							List[Any]       =   common_inputs["""decoder_input_ids"""].shape[1]
					_UpperCAmelCase  ,				_UpperCAmelCase      :							List[str]       =   self.num_attention_heads
					_UpperCAmelCase      :							List[str]       =   (
					    batch,
					    num_encoder_attention_heads,
					    encoder_seq_length,
					    self._config.hidden_size // num_encoder_attention_heads,
					)
					_UpperCAmelCase      :							List[str]       =   decoder_seq_length
					_UpperCAmelCase      :							Tuple       =   (
					    batch,
					    num_decoder_attention_heads,
					    decoder_past_length,
					    self._config.hidden_size // num_decoder_attention_heads,
					)
					_UpperCAmelCase      :							Tuple       =   torch.cat(
					    [common_inputs["""decoder_attention_mask"""], torch.ones(_A ,						_A       )] ,						dim=1       )
					_UpperCAmelCase      :							Tuple       =   []
					_UpperCAmelCase  ,				_UpperCAmelCase      :							int       =   self.num_layers

					for _ in range(_A       ):
						common_inputs["past_key_values"].append(
						    (
						        torch.zeros(_A       ),
						        torch.zeros(_A       ),
						        torch.zeros(_A       ),
						        torch.zeros(_A       ),
						    )       )
				return common_inputs
			def 							__snake_case	(						self ,						_A ,						_A = -1 ,						_A = -1 ,						_A = False ,						_A = None ,						)      ->      Mapping[str, Any]:






				'''simple docstring'''







				_UpperCAmelCase      :							Union[str, Any]       =   self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
				    _A ,						_A ,						_A ,						_A ,						_A       )

				if self.use_past:
					if not is_torch_available():
						raise ValueError("""Cannot generate dummy past_keys inputs without PyTorch installed."""       )
					else:
						import torch
					_UpperCAmelCase  ,				_UpperCAmelCase      :							Any       =   common_inputs["""input_ids"""].shape
					_UpperCAmelCase      :							Union[str, Any]       =   seqlen
					_UpperCAmelCase  ,				_UpperCAmelCase      :							Optional[Any]       =   self.num_layers
					_UpperCAmelCase  ,				_UpperCAmelCase      :							Tuple       =   self.num_attention_heads
					_UpperCAmelCase      :							Optional[Any]       =   (
					    batch,
					    num_encoder_attention_heads,
					    past_key_values_length,
					    self._config.hidden_size // num_encoder_attention_heads,
					)
					_UpperCAmelCase      :							str       =   common_inputs["""attention_mask"""].dtype
					_UpperCAmelCase      :							Tuple       =   torch.cat(
					    [common_inputs["""attention_mask"""], torch.ones(_A ,						_A ,						dtype=_A       )] ,						dim=1       )
					_UpperCAmelCase      :							str       =   [
					    (torch.zeros(_A       ), torch.zeros(_A       )) for _ in range(_A       )
					]
				return common_inputs
			def 							__snake_case	(						self ,						_A ,						_A = -1 ,						_A = -1 ,						_A = False ,						_A = None ,						)      ->      Mapping[str, Any]:






				'''simple docstring'''







				_UpperCAmelCase      :							str       =   compute_effective_axis_dimension(
				    _A ,						fixed_dimension=OnnxConfig.default_fixed_batch ,						num_token_to_add=0       )

				# If dynamic axis (-1) we forward with a fixed dimension of 8 tokens to avoid optimizations made by ONNX
				_UpperCAmelCase      :							List[str]       =   tokenizer.num_special_tokens_to_add(_A       )
				_UpperCAmelCase      :							Tuple       =   compute_effective_axis_dimension(
				    _A ,						fixed_dimension=OnnxConfig.default_fixed_sequence ,						num_token_to_add=_A       )

				# Generate dummy inputs according to compute batch and sequence
				_UpperCAmelCase      :							str       =   [""" """.join([tokenizer.unk_token]       ) * seq_length] * batch_size
				_UpperCAmelCase      :							str       =   dict(tokenizer(_A ,						return_tensors=_A       )       )
				return common_inputs
			def 							__snake_case	(						self ,						_A ,						_A = -1 ,						_A = -1 ,						_A = False ,						_A = None ,						)      ->      Mapping[str, Any]:






				'''simple docstring'''







				if self.task in ["default", "seq2seq-lm"]:
					_UpperCAmelCase      :							Any       =   self._generate_dummy_inputs_for_default_and_seqaseq_lm(
					    _A ,						batch_size=_A ,						seq_length=_A ,						is_pair=_A ,						framework=_A       )

				elif self.task == "causal-lm":
					_UpperCAmelCase      :							int       =   self._generate_dummy_inputs_for_causal_lm(
					    _A ,						batch_size=_A ,						seq_length=_A ,						is_pair=_A ,						framework=_A       )
				else:
					_UpperCAmelCase      :							Tuple       =   self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
					    _A ,						batch_size=_A ,						seq_length=_A ,						is_pair=_A ,						framework=_A       )

				return common_inputs
			def 							__snake_case	(						self ,						_A ,						_A ,						_A ,						_A       )      ->      Optional[int]:






				'''simple docstring'''







				if self.task in ["default", "seq2seq-lm"]:
					_UpperCAmelCase      :							Union[str, Any]       =   super()._flatten_past_key_values_(_A ,						_A ,						_A ,						_A       )
				else:
					_UpperCAmelCase      :							Tuple       =   super(_A ,						self       )._flatten_past_key_values_(
					    _A ,						_A ,						_A ,						_A       )
			def 							__snake_case	(						self ,						_A ,						_A       )      ->      Optional[Any]:






				'''simple docstring'''







				if direction not in ["inputs", "outputs"]:
					raise ValueError(f'''direction must either be "inputs" or "outputs", but {direction} was given'''       )

				_UpperCAmelCase      :							Any       =   """past_key_values""" if direction == """inputs""" else """present"""
				_UpperCAmelCase  ,				_UpperCAmelCase      :							Any       =   self.num_layers

				_UpperCAmelCase      :							List[str]       =   """past_encoder_sequence"""
				_UpperCAmelCase      :							int       =   """past_decoder_sequence""" if direction == """inputs""" else """past_decoder_sequence + sequence"""

				for i in range(_A       ):
					_UpperCAmelCase      :							Union[str, Any]       =   {0: """batch""", 2: decoder_sequence}
					_UpperCAmelCase      :							Optional[int]       =   {0: """batch""", 2: decoder_sequence}
					_UpperCAmelCase      :							Tuple       =   {0: """batch""", 2: encoder_sequence}
					_UpperCAmelCase      :							Tuple       =   {0: """batch""", 2: encoder_sequence}





