






import tempfile
import unittest

import numpy as np
import torch

from diffusers import DiffusionPipeline
from diffusers.models.attention_processor import Attention, AttnAddedKVProcessor
class        __snake_case(					unittest.TestCase      ):



	def 						UpperCAmelCase    (	self   :   Union[str, Any] ,		UpperCAmelCase__   :   bool = False	):

				"""simple docstring"""






				lowercase__						:	List[str]         = 10

				if only_cross_attention:
							lowercase__						:	Any         = 12
				else:
							# when only cross attention is not set, the cross attention dim must be the same as the query dim
							lowercase__						:	Dict         = query_dim

				return {
				    "query_dim": query_dim,
				    "cross_attention_dim": cross_attention_dim,
				    "heads": 2,
				    "dim_head": 4,
				    "added_kv_proj_dim": 6,
				    "norm_num_groups": 1,
				    "only_cross_attention": only_cross_attention,
				    "processor": AttnAddedKVProcessor(),
				}



	def 						UpperCAmelCase    (	self   :   Union[str, Any] ,		UpperCAmelCase__   :   Optional[int] ,		UpperCAmelCase__   :   Optional[int]	):

				"""simple docstring"""






				lowercase__						:	List[Any]         = 2

				lowercase__						:	int         = torch.rand(UpperCAmelCase__ ,		UpperCAmelCase__ ,		3 ,		2	)
				lowercase__						:	Optional[int]         = torch.rand(UpperCAmelCase__ ,		4 ,		UpperCAmelCase__	)
				lowercase__						:	List[str]         = None

				return {
				    "hidden_states": hidden_states,
				    "encoder_hidden_states": encoder_hidden_states,
				    "attention_mask": attention_mask,
				}







	def 						UpperCAmelCase    (	self   :   Optional[Any]	):

				"""simple docstring"""






				torch.manual_seed(0	)

				lowercase__						:	Union[str, Any]         = self.get_constructor_arguments(only_cross_attention=UpperCAmelCase__	)
				lowercase__						:	Union[str, Any]         = Attention(**UpperCAmelCase__	)

				self.assertTrue(attn.to_k is not None	)
				self.assertTrue(attn.to_v is not None	)

				lowercase__						:	List[Any]         = self.get_forward_arguments(
				    query_dim=constructor_args["""query_dim"""] ,		added_kv_proj_dim=constructor_args["""added_kv_proj_dim"""]	)

				lowercase__						:	Any         = attn(**UpperCAmelCase__	)

				# only self attention

				torch.manual_seed(0	)

				lowercase__						:	List[str]         = self.get_constructor_arguments(only_cross_attention=UpperCAmelCase__	)
				lowercase__						:	int         = Attention(**UpperCAmelCase__	)

				self.assertTrue(attn.to_k is None	)
				self.assertTrue(attn.to_v is None	)

				lowercase__						:	Optional[Any]         = self.get_forward_arguments(
				    query_dim=constructor_args["""query_dim"""] ,		added_kv_proj_dim=constructor_args["""added_kv_proj_dim"""]	)

				lowercase__						:	int         = attn(**UpperCAmelCase__	)

				self.assertTrue((only_cross_attn_out != self_and_cross_attn_out).all()	)
class        __snake_case(					unittest.TestCase      ):



	def 						UpperCAmelCase    (	self   :   List[str]	):

				"""simple docstring"""






				lowercase__						:	Optional[Any]         = DiffusionPipeline.from_pretrained("""hf-internal-testing/tiny-stable-diffusion-pipe""" ,		safety_checker=UpperCAmelCase__	)

				lowercase__						:	Union[str, Any]         = pipe(
				    """foo""" ,		num_inference_steps=2 ,		generator=torch.Generator("""cpu"""	).manual_seed(0	) ,		output_type="""np""" ,		).images

				# the initial conversion succeeds
				lowercase__						:	Union[str, Any]         = DiffusionPipeline.from_pretrained(
				    """hf-internal-testing/tiny-stable-diffusion-pipe""" ,		device_map="""sequential""" ,		safety_checker=UpperCAmelCase__	)

				lowercase__						:	Dict         = pipe(
				    """foo""" ,		num_inference_steps=2 ,		generator=torch.Generator("""cpu"""	).manual_seed(0	) ,		output_type="""np""" ,		).images

				with tempfile.TemporaryDirectory() as tmpdir:
							# save the converted model
							pipe.save_pretrained(UpperCAmelCase__	)

							# can also load the converted weights
							lowercase__						:	str         = DiffusionPipeline.from_pretrained(UpperCAmelCase__ ,		device_map="""sequential""" ,		safety_checker=UpperCAmelCase__	)

				lowercase__						:	Optional[Any]         = pipe(
				    """foo""" ,		num_inference_steps=2 ,		generator=torch.Generator("""cpu"""	).manual_seed(0	) ,		output_type="""np""" ,		).images

				self.assertTrue(np.allclose(UpperCAmelCase__ ,		UpperCAmelCase__ ,		atol=1e-5	)	)
				self.assertTrue(np.allclose(UpperCAmelCase__ ,		UpperCAmelCase__ ,		atol=1e-5	)	)










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        __snake_case(					__UpperCamelCase      ):
 lowercase :       Dict		      =		'''ernie'''



 def __init__(	self   :   Tuple ,		UpperCAmelCase__   :   int=30_522 ,		UpperCAmelCase__   :   List[str]=768 ,		UpperCAmelCase__   :   Tuple=12 ,		UpperCAmelCase__   :   Optional[int]=12 ,		UpperCAmelCase__   :   Any=3_072 ,		UpperCAmelCase__   :   Tuple="gelu" ,		UpperCAmelCase__   :   Union[str, Any]=0.1 ,		UpperCAmelCase__   :   str=0.1 ,		UpperCAmelCase__   :   List[str]=512 ,		UpperCAmelCase__   :   str=2 ,		UpperCAmelCase__   :   Optional[int]=3 ,		UpperCAmelCase__   :   Dict=False ,		UpperCAmelCase__   :   Any=0.02 ,		UpperCAmelCase__   :   Any=1e-12 ,		UpperCAmelCase__   :   str=0 ,		UpperCAmelCase__   :   Any="absolute" ,		UpperCAmelCase__   :   List[str]=True ,		UpperCAmelCase__   :   Tuple=None ,		**UpperCAmelCase__   :   Dict ,		):

    """simple docstring"""






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

    lowercase__						:	List[str]         = vocab_size
    lowercase__						:	Optional[int]         = hidden_size
    lowercase__						:	Union[str, Any]         = num_hidden_layers
    lowercase__						:	Any         = num_attention_heads
    lowercase__						:	Optional[int]         = hidden_act
    lowercase__						:	str         = intermediate_size
    lowercase__						:	Tuple         = hidden_dropout_prob
    lowercase__						:	Tuple         = attention_probs_dropout_prob
    lowercase__						:	str         = max_position_embeddings
    lowercase__						:	List[Any]         = type_vocab_size
    lowercase__						:	Any         = task_type_vocab_size
    lowercase__						:	Optional[Any]         = use_task_id
    lowercase__						:	Optional[Any]         = initializer_range
    lowercase__						:	Union[str, Any]         = layer_norm_eps
    lowercase__						:	Dict         = position_embedding_type
    lowercase__						:	Union[str, Any]         = use_cache
    lowercase__						:	int         = classifier_dropout
class        __snake_case(					__UpperCamelCase      ):



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

    """simple docstring"""






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


