

import gc
import random
import unittest

import numpy as np
import torch
from transformers import XLMRobertaTokenizerFast

from diffusers import DDIMScheduler, KandinskyPipeline, KandinskyPriorPipeline, UNetaDConditionModel, VQModel
from diffusers.pipelines.kandinsky.text_encoder import MCLIPConfig, MultilingualCLIP
from diffusers.utils import floats_tensor, load_numpy, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..test_pipelines_common import PipelineTesterMixin, assert_mean_pixel_difference


enable_full_determinism()

class 					__UpperCamelCase					(       SCREAMING_SNAKE_CASE	,    unittest.TestCase      ):
	_lowercase	: List[str]												=	KandinskyPipeline
	_lowercase	: Optional[int]												=	[
	    "prompt",
	    "image_embeds",
	    "negative_image_embeds",
	]
	_lowercase	: Tuple												=	["prompt", "negative_prompt", "image_embeds", "negative_image_embeds"]
	_lowercase	: Union[str, Any]												=	[
	    "generator",
	    "height",
	    "width",
	    "latents",
	    "guidance_scale",
	    "negative_prompt",
	    "num_inference_steps",
	    "return_dict",
	    "guidance_scale",
	    "num_images_per_prompt",
	    "output_type",
	    "return_dict",
	]
	_lowercase	: Any												=	False







	@property
	def 						_SCREAMING_SNAKE_CASE				( self:			Optional[Any]   ):

						'''simple docstring'''





						return 32







	@property
	def 						_SCREAMING_SNAKE_CASE				( self:			Tuple   ):

						'''simple docstring'''





						return 32







	@property
	def 						_SCREAMING_SNAKE_CASE				( self:			Union[str, Any]   ):

						'''simple docstring'''





						return self.time_input_dim







	@property
	def 						_SCREAMING_SNAKE_CASE				( self:			Union[str, Any]   ):

						'''simple docstring'''





						return self.time_input_dim * 4







	@property
	def 						_SCREAMING_SNAKE_CASE				( self:			Optional[int]   ):

						'''simple docstring'''





						return 1_00







	@property
	def 						_SCREAMING_SNAKE_CASE				( self:			Optional[int]   ):

						'''simple docstring'''





						__magic_name__				=     XLMRobertaTokenizerFast.from_pretrained('YiYiXu/tiny-random-mclip-base'   )
						return tokenizer







	@property
	def 						_SCREAMING_SNAKE_CASE				( self:			Optional[Any]   ):

						'''simple docstring'''





						torch.manual_seed(0   )
						__magic_name__				=     MCLIPConfig(
						    numDims=self.cross_attention_dim	,			transformerDimensions=self.text_embedder_hidden_size	,			hidden_size=self.text_embedder_hidden_size	,			intermediate_size=37	,			num_attention_heads=4	,			num_hidden_layers=5	,			vocab_size=10_05	,			)

						__magic_name__				=     MultilingualCLIP(__UpperCamelCase   )
						__magic_name__				=     text_encoder.eval()

						return text_encoder







	@property
	def 						_SCREAMING_SNAKE_CASE				( self:			str   ):

						'''simple docstring'''





						torch.manual_seed(0   )

						__magic_name__				=     {
						    'in_channels': 4,
						    # Out channels is double in channels because predicts mean and variance
						    'out_channels': 8,
						    'addition_embed_type': 'text_image',
						    'down_block_types': ('ResnetDownsampleBlock2D', 'SimpleCrossAttnDownBlock2D'),
						    'up_block_types': ('SimpleCrossAttnUpBlock2D', 'ResnetUpsampleBlock2D'),
						    'mid_block_type': 'UNetMidBlock2DSimpleCrossAttn',
						    'block_out_channels': (self.block_out_channels_a, self.block_out_channels_a * 2),
						    'layers_per_block': 1,
						    'encoder_hid_dim': self.text_embedder_hidden_size,
						    'encoder_hid_dim_type': 'text_image_proj',
						    'cross_attention_dim': self.cross_attention_dim,
						    'attention_head_dim': 4,
						    'resnet_time_scale_shift': 'scale_shift',
						    'class_embed_type': None,
						}

						__magic_name__				=     UNetaDConditionModel(**__UpperCamelCase   )
						return model







	@property
	def 						_SCREAMING_SNAKE_CASE				( self:			List[Any]   ):

						'''simple docstring'''





						return {
						    "block_out_channels": [32, 64],
						    "down_block_types": ["DownEncoderBlock2D", "AttnDownEncoderBlock2D"],
						    "in_channels": 3,
						    "latent_channels": 4,
						    "layers_per_block": 1,
						    "norm_num_groups": 8,
						    "norm_type": "spatial",
						    "num_vq_embeddings": 12,
						    "out_channels": 3,
						    "up_block_types": [
						        "AttnUpDecoderBlock2D",
						        "UpDecoderBlock2D",
						    ],
						    "vq_embed_dim": 4,
						}







	@property
	def 						_SCREAMING_SNAKE_CASE				( self:			List[str]   ):

						'''simple docstring'''





						torch.manual_seed(0   )
						__magic_name__				=     VQModel(**self.dummy_movq_kwargs   )
						return model







	def 						_SCREAMING_SNAKE_CASE				( self:			List[Any]   ):

						'''simple docstring'''





						__magic_name__				=     self.dummy_text_encoder
						__magic_name__				=     self.dummy_tokenizer
						__magic_name__				=     self.dummy_unet
						__magic_name__				=     self.dummy_movq

						__magic_name__				=     DDIMScheduler(
						    num_train_timesteps=10_00	,			beta_schedule='linear'	,			beta_start=0.00085	,			beta_end=0.012	,			clip_sample=__UpperCamelCase	,			set_alpha_to_one=__UpperCamelCase	,			steps_offset=1	,			prediction_type='epsilon'	,			thresholding=__UpperCamelCase	,			)

						__magic_name__				=     {
						    'text_encoder': text_encoder,
						    'tokenizer': tokenizer,
						    'unet': unet,
						    'scheduler': scheduler,
						    'movq': movq,
						}
						return components







	def 						_SCREAMING_SNAKE_CASE				( self:			int	,			__UpperCamelCase:			Optional[Any]	,			__UpperCamelCase:			Any=0   ):

						'''simple docstring'''





						__magic_name__				=     floats_tensor((1, self.cross_attention_dim)	,			rng=random.Random(__UpperCamelCase   )   ).to(__UpperCamelCase   )
						__magic_name__				=     floats_tensor((1, self.cross_attention_dim)	,			rng=random.Random(seed + 1   )   ).to(__UpperCamelCase   )
						if str(__UpperCamelCase   ).startswith('mps'   ):
											__magic_name__				=     torch.manual_seed(__UpperCamelCase   )
						else:
											__magic_name__				=     torch.Generator(device=__UpperCamelCase   ).manual_seed(__UpperCamelCase   )
						__magic_name__				=     {
						    'prompt': 'horse',
						    'image_embeds': image_embeds,
						    'negative_image_embeds': negative_image_embeds,
						    'generator': generator,
						    'height': 64,
						    'width': 64,
						    'guidance_scale': 4.0,
						    'num_inference_steps': 2,
						    'output_type': 'np',
						}
						return inputs







	def 						_SCREAMING_SNAKE_CASE				( self:			Union[str, Any]   ):

						'''simple docstring'''





						__magic_name__				=     'cpu'

						__magic_name__				=     self.get_dummy_components()

						__magic_name__				=     self.pipeline_class(**__UpperCamelCase   )
						__magic_name__				=     pipe.to(__UpperCamelCase   )

						pipe.set_progress_bar_config(disable=__UpperCamelCase   )

						__magic_name__				=     pipe(**self.get_dummy_inputs(__UpperCamelCase   )   )
						__magic_name__				=     output.images

						__magic_name__				=     pipe(
						    **self.get_dummy_inputs(__UpperCamelCase   )	,			return_dict=__UpperCamelCase	,			)[0]

						__magic_name__				=     image[0, -3:, -3:, -1]
						__magic_name__				=     image_from_tuple[0, -3:, -3:, -1]

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

						__magic_name__				=     np.array(
						    [0.328663, 1.0, 0.23216873, 1.0, 0.92717564, 0.4639046, 0.96894777, 0.31713378, 0.6293953]   )

						assert (
						    np.abs(image_slice.flatten() - expected_slice   ).max() < 1E-2
						), F' expected_slice {expected_slice}, but got {image_slice.flatten()}'
						assert (
						    np.abs(image_from_tuple_slice.flatten() - expected_slice   ).max() < 1E-2
						), F' expected_slice {expected_slice}, but got {image_from_tuple_slice.flatten()}'




@slow
@require_torch_gpu
class 					__UpperCamelCase					(       unittest.TestCase      ):







	def 						_SCREAMING_SNAKE_CASE				( self:			List[str]   ):

						'''simple docstring'''





						super().tearDown()
						gc.collect()
						torch.cuda.empty_cache()







	def 						_SCREAMING_SNAKE_CASE				( self:			Dict   ):

						'''simple docstring'''





						__magic_name__				=     load_numpy(
						    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
						    '/kandinsky/kandinsky_text2img_cat_fp16.npy'   )

						__magic_name__				=     KandinskyPriorPipeline.from_pretrained(
						    'kandinsky-community/kandinsky-2-1-prior'	,			torch_dtype=torch.floataa   )
						pipe_prior.to(__UpperCamelCase   )

						__magic_name__				=     KandinskyPipeline.from_pretrained('kandinsky-community/kandinsky-2-1'	,			torch_dtype=torch.floataa   )
						__magic_name__				=     pipeline.to(__UpperCamelCase   )
						pipeline.set_progress_bar_config(disable=__UpperCamelCase   )

						__magic_name__				=     'red cat, 4k photo'

						__magic_name__				=     torch.Generator(device='cuda'   ).manual_seed(0   )
						__magic_name__,						__magic_name__				=     pipe_prior(
						    __UpperCamelCase	,			generator=__UpperCamelCase	,			num_inference_steps=5	,			negative_prompt=''	,			).to_tuple()

						__magic_name__				=     torch.Generator(device='cuda'   ).manual_seed(0   )
						__magic_name__				=     pipeline(
						    __UpperCamelCase	,			image_embeds=__UpperCamelCase	,			negative_image_embeds=__UpperCamelCase	,			generator=__UpperCamelCase	,			num_inference_steps=1_00	,			output_type='np'	,			)

						__magic_name__				=     output.images[0]

						assert image.shape == (5_12, 5_12, 3)

						assert_mean_pixel_difference(__UpperCamelCase	,			__UpperCamelCase   )







from collections import OrderedDict
from typing import Mapping

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


A__							=				logging.get_logger(__name__)

A__							=				{
    "nghuyong/ernie-1.0-base-zh": "https://huggingface.co/nghuyong/ernie-1.0-base-zh/resolve/main/config.json",
    "nghuyong/ernie-2.0-base-en": "https://huggingface.co/nghuyong/ernie-2.0-base-en/resolve/main/config.json",
    "nghuyong/ernie-2.0-large-en": "https://huggingface.co/nghuyong/ernie-2.0-large-en/resolve/main/config.json",
    "nghuyong/ernie-3.0-base-zh": "https://huggingface.co/nghuyong/ernie-3.0-base-zh/resolve/main/config.json",
    "nghuyong/ernie-3.0-medium-zh": "https://huggingface.co/nghuyong/ernie-3.0-medium-zh/resolve/main/config.json",
    "nghuyong/ernie-3.0-mini-zh": "https://huggingface.co/nghuyong/ernie-3.0-mini-zh/resolve/main/config.json",
    "nghuyong/ernie-3.0-micro-zh": "https://huggingface.co/nghuyong/ernie-3.0-micro-zh/resolve/main/config.json",
    "nghuyong/ernie-3.0-nano-zh": "https://huggingface.co/nghuyong/ernie-3.0-nano-zh/resolve/main/config.json",
    "nghuyong/ernie-gram-zh": "https://huggingface.co/nghuyong/ernie-gram-zh/resolve/main/config.json",
    "nghuyong/ernie-health-zh": "https://huggingface.co/nghuyong/ernie-health-zh/resolve/main/config.json",
}

class 					__UpperCamelCase					(       SCREAMING_SNAKE_CASE      ):
	_lowercase	: Dict												=	"ernie"







	def __init__( self:			List[Any]	,			__UpperCamelCase:			Dict=3_05_22	,			__UpperCamelCase:			str=7_68	,			__UpperCamelCase:			Tuple=12	,			__UpperCamelCase:			int=12	,			__UpperCamelCase:			Optional[Any]=30_72	,			__UpperCamelCase:			str="gelu"	,			__UpperCamelCase:			Any=0.1	,			__UpperCamelCase:			int=0.1	,			__UpperCamelCase:			int=5_12	,			__UpperCamelCase:			Any=2	,			__UpperCamelCase:			Tuple=3	,			__UpperCamelCase:			Any=False	,			__UpperCamelCase:			Optional[int]=0.02	,			__UpperCamelCase:			List[Any]=1E-12	,			__UpperCamelCase:			Union[str, Any]=0	,			__UpperCamelCase:			Tuple="absolute"	,			__UpperCamelCase:			Union[str, Any]=True	,			__UpperCamelCase:			Optional[Any]=None	,			**__UpperCamelCase:			List[Any]	,			):

						'''simple docstring'''





						super().__init__(pad_token_id=__UpperCamelCase	,			**__UpperCamelCase   )

						__magic_name__				=     vocab_size
						__magic_name__				=     hidden_size
						__magic_name__				=     num_hidden_layers
						__magic_name__				=     num_attention_heads
						__magic_name__				=     hidden_act
						__magic_name__				=     intermediate_size
						__magic_name__				=     hidden_dropout_prob
						__magic_name__				=     attention_probs_dropout_prob
						__magic_name__				=     max_position_embeddings
						__magic_name__				=     type_vocab_size
						__magic_name__				=     task_type_vocab_size
						__magic_name__				=     use_task_id
						__magic_name__				=     initializer_range
						__magic_name__				=     layer_norm_eps
						__magic_name__				=     position_embedding_type
						__magic_name__				=     use_cache
						__magic_name__				=     classifier_dropout




class 					__UpperCamelCase					(       SCREAMING_SNAKE_CASE      ):







	@property
	def 						_SCREAMING_SNAKE_CASE				( self:			Tuple   ):

						'''simple docstring'''





						if self.task == "multiple-choice":
											__magic_name__				=     {0: 'batch', 1: 'choice', 2: 'sequence'}
						else:
											__magic_name__				=     {0: 'batch', 1: 'sequence'}
						return OrderedDict(
						    [
						        ('input_ids', dynamic_axis),
						        ('attention_mask', dynamic_axis),
						        ('token_type_ids', dynamic_axis),
						        ('task_type_ids', dynamic_axis),
						    ]   )




