






'''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 	SCREAMING_SNAKE_CASE					(					SCREAMING_SNAKE_CASE_	,   unittest.TestCase       ):
			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 			SCREAMING_SNAKE_CASE      ( self      :			Optional[Any]					)			->			Tuple:
							return 32






			@property
			def 			SCREAMING_SNAKE_CASE      ( self      :			Optional[int]					)			->			Tuple:
							return 32






			@property
			def 			SCREAMING_SNAKE_CASE      ( self      :			Optional[int]					)			->			List[str]:
							return self.time_input_dim






			@property
			def 			SCREAMING_SNAKE_CASE      ( self      :			Any					)			->			Dict:
							return self.time_input_dim * 4






			@property
			def 			SCREAMING_SNAKE_CASE      ( self      :			Dict					)			->			List[Any]:
							return 100






			@property
			def 			SCREAMING_SNAKE_CASE      ( self      :			Dict					)			->			Dict:
							a_		:  Tuple											=      CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''					)
							return tokenizer






			@property
			def 			SCREAMING_SNAKE_CASE      ( self      :			Optional[int]					)			->			List[str]:
							torch.manual_seed(0					)
							a_		:  Dict											=      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(__SCREAMING_SNAKE_CASE					)






			@property
			def 			SCREAMING_SNAKE_CASE      ( self      :			List[Any]					)			->			Tuple:
							torch.manual_seed(0					)

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

							a_		:  Dict											=      PriorTransformer(**__SCREAMING_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
							a_		:  Any											=      nn.Parameter(torch.ones(model.clip_std.shape					)					)
							return model






			@property
			def 			SCREAMING_SNAKE_CASE      ( self      :			Dict					)			->			Union[str, Any]:
							torch.manual_seed(0					)
							a_		:  Any											=      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  ,						)

							a_		:  Optional[int]											=      CLIPVisionModelWithProjection(__SCREAMING_SNAKE_CASE					)
							return model






			@property
			def 			SCREAMING_SNAKE_CASE      ( self      :			List[Any]					)			->			str:
							a_		:  Optional[int]											=      CLIPImageProcessor(
							    crop_size=224  ,						do_center_crop=__SCREAMING_SNAKE_CASE  ,						do_normalize=__SCREAMING_SNAKE_CASE  ,						do_resize=__SCREAMING_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 			SCREAMING_SNAKE_CASE      ( self      :			Union[str, Any]					)			->			Any:
							a_		:  Dict											=      self.dummy_prior
							a_		:  int											=      self.dummy_image_encoder
							a_		:  List[str]											=      self.dummy_text_encoder
							a_		:  Tuple											=      self.dummy_tokenizer
							a_		:  int											=      self.dummy_image_processor

							a_		:  Optional[int]											=      UnCLIPScheduler(
							    variance_type='''fixed_small_log'''  ,						prediction_type='''sample'''  ,						num_train_timesteps=1000  ,						clip_sample=__SCREAMING_SNAKE_CASE  ,						clip_sample_range=10.0  ,						)

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

							return components






			def 			SCREAMING_SNAKE_CASE      ( self      :			Tuple  ,						__SCREAMING_SNAKE_CASE      :			Dict  ,						__SCREAMING_SNAKE_CASE      :			Union[str, Any]=0					)			->			Tuple:
							if str(__SCREAMING_SNAKE_CASE					).startswith('''mps'''					):
											a_		:  Dict											=      torch.manual_seed(__SCREAMING_SNAKE_CASE					)
							else:
											a_		:  str											=      torch.Generator(device=__SCREAMING_SNAKE_CASE					).manual_seed(__SCREAMING_SNAKE_CASE					)
							a_		:  int											=      {
							    '''prompt''': '''horse''',
							    '''generator''': generator,
							    '''guidance_scale''': 4.0,
							    '''num_inference_steps''': 2,
							    '''output_type''': '''np''',
							}
							return inputs






			def 			SCREAMING_SNAKE_CASE      ( self      :			Tuple					)			->			Tuple:
							a_		:  str											=      '''cpu'''

							a_		:  str											=      self.get_dummy_components()

							a_		:  Dict											=      self.pipeline_class(**__SCREAMING_SNAKE_CASE					)
							a_		:  Dict											=      pipe.to(__SCREAMING_SNAKE_CASE					)

							pipe.set_progress_bar_config(disable=__SCREAMING_SNAKE_CASE					)

							a_		:  Optional[int]											=      pipe(**self.get_dummy_inputs(__SCREAMING_SNAKE_CASE					)					)
							a_		:  Optional[Any]											=      output.image_embeds

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

							a_		:  Any											=      image[0, -10:]
							a_		:  List[str]											=      image_from_tuple[0, -10:]

							assert image.shape == (1, 32)

							a_		:  List[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 			SCREAMING_SNAKE_CASE      ( self      :			Optional[Any]					)			->			List[Any]:
							a_		:  Union[str, Any]											=      torch_device == '''cpu'''
							a_		:  Optional[Any]											=      True
							a_		:  Dict											=      False

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







			@skip_mps
			def 			SCREAMING_SNAKE_CASE      ( self      :			Tuple					)			->			str:
							a_		:  str											=      torch_device == '''cpu'''
							a_		:  List[Any]											=      False

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












'''simple docstring'''







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

from packaging import version


if TYPE_CHECKING:
		from ... import PreTrainedTokenizer, TensorType

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfigWithPast, PatchingSpec
from ...utils import is_torch_available, logging


__lowerCAmelCase       		=      logging.get_logger(__name__)

__lowerCAmelCase       		=      {
    'bigscience/bloom': 'https://huggingface.co/bigscience/bloom/resolve/main/config.json',
    'bigscience/bloom-560m': 'https://huggingface.co/bigscience/bloom-560m/blob/main/config.json',
    'bigscience/bloom-1b1': 'https://huggingface.co/bigscience/bloom-1b1/blob/main/config.json',
    'bigscience/bloom-1b7': 'https://huggingface.co/bigscience/bloom-1b7/blob/main/config.json',
    'bigscience/bloom-3b': 'https://huggingface.co/bigscience/bloom-3b/blob/main/config.json',
    'bigscience/bloom-7b1': 'https://huggingface.co/bigscience/bloom-7b1/blob/main/config.json',
}





class 	SCREAMING_SNAKE_CASE					(					SCREAMING_SNAKE_CASE_       ):
			snake_case__  =				"bloom"
			snake_case__  =				["past_key_values"]
			snake_case__  =				{
			    "num_hidden_layers": "n_layer",
			    "num_attention_heads": "n_head",
			}






			def __init__( self      :			Union[str, Any]  ,						__SCREAMING_SNAKE_CASE      :			int=25_0880  ,						__SCREAMING_SNAKE_CASE      :			Dict=64  ,						__SCREAMING_SNAKE_CASE      :			Tuple=2  ,						__SCREAMING_SNAKE_CASE      :			int=8  ,						__SCREAMING_SNAKE_CASE      :			Any=1e-5  ,						__SCREAMING_SNAKE_CASE      :			Optional[Any]=0.02  ,						__SCREAMING_SNAKE_CASE      :			Union[str, Any]=True  ,						__SCREAMING_SNAKE_CASE      :			int=1  ,						__SCREAMING_SNAKE_CASE      :			Any=2  ,						__SCREAMING_SNAKE_CASE      :			Optional[Any]=False  ,						__SCREAMING_SNAKE_CASE      :			Optional[Any]=0.0  ,						__SCREAMING_SNAKE_CASE      :			str=0.0  ,						__SCREAMING_SNAKE_CASE      :			List[Any]=1  ,						__SCREAMING_SNAKE_CASE      :			List[str]=False  ,						**__SCREAMING_SNAKE_CASE      :			str  ,						)			->			Any:
							a_		:  Optional[int]											=      vocab_size
							# Backward compatibility with n_embed kwarg
							a_		:  Any											=      kwargs.pop('''n_embed'''  ,						__SCREAMING_SNAKE_CASE					)
							a_		:  Optional[int]											=      hidden_size if n_embed is None else n_embed
							a_		:  int											=      n_layer
							a_		:  str											=      n_head
							a_		:  Optional[int]											=      layer_norm_epsilon
							a_		:  Dict											=      initializer_range
							a_		:  List[str]											=      use_cache
							a_		:  Dict											=      pretraining_tp
							a_		:  Optional[Any]											=      apply_residual_connection_post_layernorm
							a_		:  Optional[Any]											=      hidden_dropout
							a_		:  List[str]											=      attention_dropout

							a_		:  Dict											=      bos_token_id
							a_		:  Optional[int]											=      eos_token_id
							a_		:  Any											=      slow_but_exact

							super().__init__(bos_token_id=__SCREAMING_SNAKE_CASE  ,						eos_token_id=__SCREAMING_SNAKE_CASE  ,						**__SCREAMING_SNAKE_CASE					)





class 	SCREAMING_SNAKE_CASE					(					SCREAMING_SNAKE_CASE_       ):
			snake_case__  =				version.parse("1.12"       )






			def __init__( self      :			Tuple  ,						__SCREAMING_SNAKE_CASE      :			PretrainedConfig  ,						__SCREAMING_SNAKE_CASE      :			str = "default"  ,						__SCREAMING_SNAKE_CASE      :			List[PatchingSpec] = None  ,						__SCREAMING_SNAKE_CASE      :			bool = False  ,						)			->			Optional[Any]:
							super().__init__(__SCREAMING_SNAKE_CASE  ,						task=__SCREAMING_SNAKE_CASE  ,						patching_specs=__SCREAMING_SNAKE_CASE  ,						use_past=__SCREAMING_SNAKE_CASE					)
							if not getattr(self._config  ,						'''pad_token_id'''  ,						__SCREAMING_SNAKE_CASE					):
											# TODO: how to do that better?
											a_		:  Tuple											=      0






			@property
			def 			SCREAMING_SNAKE_CASE      ( self      :			Tuple					)			->			Mapping[str, Mapping[int, str]]:
							a_		:  Optional[Any]											=      OrderedDict({'''input_ids''': {0: '''batch''', 1: '''sequence'''}}					)
							if self.use_past:
											# BLOOM stores values on dynamic axis 2. For more details see: https://github.com/huggingface/transformers/pull/18344
											self.fill_with_past_key_values_(__SCREAMING_SNAKE_CASE  ,						direction='''inputs'''  ,						inverted_values_shape=__SCREAMING_SNAKE_CASE					)
											a_		:  Union[str, Any]											=      {0: '''batch''', 1: '''past_sequence + sequence'''}
							else:
											a_		:  Union[str, Any]											=      {0: '''batch''', 1: '''sequence'''}

							return common_inputs






			@property
			def 			SCREAMING_SNAKE_CASE      ( self      :			Any					)			->			int:
							return self._config.n_layer






			@property
			def 			SCREAMING_SNAKE_CASE      ( self      :			int					)			->			int:
							return self._config.n_head






			@property
			def 			SCREAMING_SNAKE_CASE      ( self      :			int					)			->			float:
							return 1e-3






			def 			SCREAMING_SNAKE_CASE      ( self      :			Dict  ,						__SCREAMING_SNAKE_CASE      :			"PreTrainedTokenizer"  ,						__SCREAMING_SNAKE_CASE      :			int = -1  ,						__SCREAMING_SNAKE_CASE      :			int = -1  ,						__SCREAMING_SNAKE_CASE      :			bool = False  ,						__SCREAMING_SNAKE_CASE      :			Optional["TensorType"] = None  ,						)			->			Mapping[str, Any]:
							a_		:  Dict											=      super(__SCREAMING_SNAKE_CASE  ,						self					).generate_dummy_inputs(
							    __SCREAMING_SNAKE_CASE  ,						batch_size=__SCREAMING_SNAKE_CASE  ,						seq_length=__SCREAMING_SNAKE_CASE  ,						is_pair=__SCREAMING_SNAKE_CASE  ,						framework=__SCREAMING_SNAKE_CASE					)

							# We need to order the input in the way they appears in the forward()
							a_		:  Union[str, Any]											=      OrderedDict({'''input_ids''': common_inputs['''input_ids''']}					)

							# Need to add the past_keys
							if self.use_past:
											if not is_torch_available():
															raise ValueError('''Cannot generate dummy past_keys inputs without PyTorch installed.'''					)
											else:
															import torch

															a_				,			a_		:  Any											=      common_inputs['''input_ids'''].shape
															# Not using the same length for past_key_values
															a_		:  str											=      seqlen + 2
															a_		:  Any											=      self._config.hidden_size // self.num_attention_heads
															a_		:  Optional[int]											=      (
															    batch * self.num_attention_heads,
															    head_dim,
															    past_key_values_length,
															)
															a_		:  Any											=      (
															    batch * self.num_attention_heads,
															    past_key_values_length,
															    head_dim,
															)
															a_		:  List[str]											=      [
															    (torch.zeros(__SCREAMING_SNAKE_CASE					), torch.zeros(__SCREAMING_SNAKE_CASE					)) for _ in range(self.num_layers					)
															]

							a_		:  Union[str, Any]											=      common_inputs['''attention_mask''']
							if self.use_past:
											a_		:  Optional[int]											=      ordered_inputs['''attention_mask'''].dtype
											a_		:  List[Any]											=      torch.cat(
											    [ordered_inputs['''attention_mask'''], torch.ones(__SCREAMING_SNAKE_CASE  ,						__SCREAMING_SNAKE_CASE  ,						dtype=__SCREAMING_SNAKE_CASE					)]  ,						dim=1					)

							return ordered_inputs







			@property
			def 			SCREAMING_SNAKE_CASE      ( self      :			Optional[int]					)			->			int:
							return 13




