


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        __snake_case		(  UpperCamelCase__		,unittest.TestCase		):
							_a       =				KandinskyPriorPipeline
							_a       =				["""prompt"""]
							_a       =				["""prompt""", """negative_prompt"""]
							_a       =				[
							    """num_images_per_prompt""",
							    """generator""",
							    """num_inference_steps""",
							    """latents""",
							    """negative_prompt""",
							    """guidance_scale""",
							    """output_type""",
							    """return_dict""",
							]
							_a       =				False




							@property
							def      UpperCAmelCase__    ( self  : Union[str, Any]):
									return 3_2




							@property
							def      UpperCAmelCase__    ( self  : int):
									return 3_2




							@property
							def      UpperCAmelCase__    ( self  : Optional[int]):
									return self.time_input_dim




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




							@property
							def      UpperCAmelCase__    ( self  : int):
									return 1_0_0




							@property
							def      UpperCAmelCase__    ( self  : Dict):
									lowerCAmelCase_    :  Union[str, Any]             =				CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip''')
									return tokenizer




							@property
							def      UpperCAmelCase__    ( self  : str):
									torch.manual_seed(0)
									lowerCAmelCase_    :  str             =				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      UpperCAmelCase__    ( self  : Any):
									torch.manual_seed(0)

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

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




							@property
							def      UpperCAmelCase__    ( self  : List[Any]):
									torch.manual_seed(0)
									lowerCAmelCase_    :  List[Any]             =				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       ,       )

									lowerCAmelCase_    :  List[str]             =				CLIPVisionModelWithProjection(lowerCAmelCase__)
									return model




							@property
							def      UpperCAmelCase__    ( self  : Tuple):
									lowerCAmelCase_    :  Optional[int]             =				CLIPImageProcessor(
									    crop_size=2_2_4       ,       do_center_crop=lowerCAmelCase__       ,       do_normalize=lowerCAmelCase__       ,       do_resize=lowerCAmelCase__       ,       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=2_2_4       ,       )

									return image_processor




							def      UpperCAmelCase__    ( self  : Tuple):
									lowerCAmelCase_    :  Any             =				self.dummy_prior
									lowerCAmelCase_    :  List[Any]             =				self.dummy_image_encoder
									lowerCAmelCase_    :  List[Any]             =				self.dummy_text_encoder
									lowerCAmelCase_    :  Optional[int]             =				self.dummy_tokenizer
									lowerCAmelCase_    :  Union[str, Any]             =				self.dummy_image_processor

									lowerCAmelCase_    :  List[str]             =				UnCLIPScheduler(
									    variance_type='''fixed_small_log'''       ,       prediction_type='''sample'''       ,       num_train_timesteps=1_0_0_0       ,       clip_sample=lowerCAmelCase__       ,       clip_sample_range=10.0       ,       )

									lowerCAmelCase_    :  Optional[Any]             =				{
									    "prior": prior,
									    "image_encoder": image_encoder,
									    "text_encoder": text_encoder,
									    "tokenizer": tokenizer,
									    "scheduler": scheduler,
									    "image_processor": image_processor,
									}

									return components




							def      UpperCAmelCase__    ( self  : Optional[int]       ,       A_  : Optional[int]       ,       A_  : Any=0):




									if str(lowerCAmelCase__).startswith('''mps'''):
											lowerCAmelCase_    :  List[str]             =				torch.manual_seed(lowerCAmelCase__)
									else:
											lowerCAmelCase_    :  str             =				torch.Generator(device=lowerCAmelCase__).manual_seed(lowerCAmelCase__)
									lowerCAmelCase_    :  Tuple             =				{
									    "prompt": "horse",
									    "generator": generator,
									    "guidance_scale": 4.0,
									    "num_inference_steps": 2,
									    "output_type": "np",
									}
									return inputs




							def      UpperCAmelCase__    ( self  : Dict):
									lowerCAmelCase_    :  int             =				"cpu"

									lowerCAmelCase_    :  Dict             =				self.get_dummy_components()

									lowerCAmelCase_    :  Optional[Any]             =				self.pipeline_class(**lowerCAmelCase__)
									lowerCAmelCase_    :  str             =				pipe.to(lowerCAmelCase__)

									pipe.set_progress_bar_config(disable=lowerCAmelCase__)

									lowerCAmelCase_    :  Dict             =				pipe(**self.get_dummy_inputs(lowerCAmelCase__))
									lowerCAmelCase_    :  str             =				output.image_embeds

									lowerCAmelCase_    :  List[Any]             =				pipe(
									    **self.get_dummy_inputs(lowerCAmelCase__)       ,       return_dict=lowerCAmelCase__       ,       )[0]

									lowerCAmelCase_    :  int             =				image[0, -1_0:]
									lowerCAmelCase_    :  Tuple             =				image_from_tuple[0, -1_0:]

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

									lowerCAmelCase_    :  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      UpperCAmelCase__    ( self  : Optional[int]):
									lowerCAmelCase_    :  Optional[Any]             =				torch_device == "cpu"
									lowerCAmelCase_    :  Any             =				True
									lowerCAmelCase_    :  str             =				False

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







							@skip_mps
							def      UpperCAmelCase__    ( self  : Union[str, Any]):
									lowerCAmelCase_    :  Union[str, Any]             =				torch_device == "cpu"
									lowerCAmelCase_    :  Optional[Any]             =				False

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










from collections import OrderedDict
from typing import Mapping

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig
from ...utils import logging


A__						:						int          =      logging.get_logger(__name__)

A__						:						Dict          =      {
    '''mnaylor/mega-base-wikitext''': '''https://huggingface.co/mnaylor/mega-base-wikitext/resolve/main/config.json''',
}



class        __snake_case		(  UpperCamelCase_		):
							_a       =				'''mega'''




							def __init__( self  : Any       ,       A_  : Dict=3_0_5_2_2       ,       A_  : int=1_2_8       ,       A_  : str=4       ,       A_  : int=2_5_6       ,       A_  : int=1_6       ,       A_  : str=True       ,       A_  : List[str]=6_4       ,       A_  : Optional[Any]=False       ,       A_  : str=-1       ,       A_  : Optional[int]=None       ,       A_  : Tuple=True       ,       A_  : Tuple="scalenorm"       ,       A_  : Optional[Any]=True       ,       A_  : int="silu"       ,       A_  : Optional[int]="softmax"       ,       A_  : Optional[int]=0.1       ,       A_  : Union[str, Any]=0.1       ,       A_  : Any=0.1       ,       A_  : Dict=False       ,       A_  : int=True       ,       A_  : Tuple=2_5_6       ,       A_  : int=True       ,       A_  : List[Any]=0.1       ,       A_  : Optional[Any]=2_0_4_8       ,       A_  : Dict=False       ,       A_  : Dict=2       ,       A_  : Tuple=0.02       ,       A_  : List[Any]=0.2       ,       A_  : str=0.02       ,       A_  : List[str]=1.0       ,       A_  : Union[str, Any]=1       ,       A_  : List[Any]=0       ,       A_  : int=2       ,       A_  : Tuple="rotary"       ,       A_  : Any=None       ,       A_  : Tuple=True       ,       A_  : Any=True       ,       **A_  : Optional[int]       ,       ):
									super().__init__(pad_token_id=A_       ,       bos_token_id=A_       ,       eos_token_id=A_       ,       **A_)

									lowerCAmelCase_    :  Tuple             =				vocab_size
									lowerCAmelCase_    :  str             =				hidden_size
									lowerCAmelCase_    :  Any             =				num_hidden_layers
									lowerCAmelCase_    :  Optional[int]             =				activation
									lowerCAmelCase_    :  Any             =				attention_activation
									lowerCAmelCase_    :  List[Any]             =				intermediate_size
									lowerCAmelCase_    :  Any             =				ema_projection_size
									lowerCAmelCase_    :  Optional[int]             =				bidirectional
									lowerCAmelCase_    :  Tuple             =				shared_representation_size
									lowerCAmelCase_    :  Any             =				use_chunking
									lowerCAmelCase_    :  List[str]             =				chunk_size
									lowerCAmelCase_    :  int             =				truncation
									lowerCAmelCase_    :  Tuple             =				normalize_before_mega
									lowerCAmelCase_    :  Tuple             =				normalization_type
									lowerCAmelCase_    :  List[str]             =				norm_affine
									lowerCAmelCase_    :  List[str]             =				dropout_prob
									lowerCAmelCase_    :  List[str]             =				hidden_dropout_prob
									lowerCAmelCase_    :  Tuple             =				attention_probs_dropout_prob
									lowerCAmelCase_    :  str             =				use_feature_dropout
									lowerCAmelCase_    :  Union[str, Any]             =				use_normalized_ffn
									lowerCAmelCase_    :  Union[str, Any]             =				nffn_hidden_size
									lowerCAmelCase_    :  int             =				normalize_before_ffn
									lowerCAmelCase_    :  Tuple             =				nffn_activation_dropout_prob
									lowerCAmelCase_    :  List[str]             =				max_positions
									lowerCAmelCase_    :  Union[str, Any]             =				add_token_type_embeddings
									lowerCAmelCase_    :  str             =				type_vocab_size
									lowerCAmelCase_    :  int             =				initializer_range
									lowerCAmelCase_    :  str             =				ema_delta_alpha_range
									lowerCAmelCase_    :  Tuple             =				ema_beta_range
									lowerCAmelCase_    :  Any             =				ema_gamma_omega_range
									lowerCAmelCase_    :  Any             =				relative_positional_bias
									lowerCAmelCase_    :  Optional[int]             =				use_cache
									lowerCAmelCase_    :  int             =				classifier_dropout
									lowerCAmelCase_    :  List[Any]             =				add_lm_hidden_dense_layer
									lowerCAmelCase_    :  Optional[Any]             =				1  # not used but required by Hugging Face





class        __snake_case		(  UpperCamelCase_		):




							@property
							def      UpperCAmelCase__    ( self  : Any):
									if self.task == "multiple-choice":
											lowerCAmelCase_    :  Dict             =				{0: '''batch''', 1: '''choice''', 2: '''sequence'''}
									else:
											lowerCAmelCase_    :  int             =				{0: '''batch''', 1: '''sequence'''}
									return OrderedDict(
									    [
									        ('''input_ids''', dynamic_axis),
									        ('''attention_mask''', dynamic_axis),
									    ])





