

'''simple docstring'''


import random

import torch
from huggingface_hub import HfApi

from diffusers import UNetaDModel


lowerCAmelCase__						:       List[Any]										=						HfApi()

lowerCAmelCase__						:       Any										=						{}
# fmt: off
lowerCAmelCase__						:       Union[str, Any]										=						torch.tensor([
    -0.7515, -1.6883, 0.2420, 0.0300, 0.6347, 1.3433, -1.1743, -3.7467,
    1.2342, -2.2485, 0.4636, 0.8076, -0.7991, 0.3969, 0.8498, 0.9189,
    -1.8887, -3.3522, 0.7639, 0.2040, 0.6271, -2.7148, -1.6316, 3.0839,
    0.3186, 0.2721, -0.9759, -1.2461, 2.6257, 1.3557
])
lowerCAmelCase__						:       int										=						torch.tensor([
    -2.3639, -2.5344, 0.0054, -0.6674, 1.5990, 1.0158, 0.3124, -2.1436,
    1.8795, -2.5429, -0.1566, -0.3973, 1.2490, 2.6447, 1.2283, -0.5208,
    -2.8154, -3.5119, 2.3838, 1.2033, 1.7201, -2.1256, -1.4576, 2.7948,
    2.4204, -0.9752, -1.2546, 0.8027, 3.2758, 3.1365
])
lowerCAmelCase__						:       str										=						torch.tensor([
    -0.6531, -0.6891, -0.3172, -0.5375, -0.9140, -0.5367, -0.1175, -0.7869,
    -0.3808, -0.4513, -0.2098, -0.0083, 0.3183, 0.5140, 0.2247, -0.1304,
    -0.1302, -0.2802, -0.2084, -0.2025, -0.4967, -0.4873, -0.0861, 0.6925,
    0.0250, 0.1290, -0.1543, 0.6316, 1.0460, 1.4943
])
lowerCAmelCase__						:       str										=						torch.tensor([
    0.0911, 0.1107, 0.0182, 0.0435, -0.0805, -0.0608, 0.0381, 0.2172,
    -0.0280, 0.1327, -0.0299, -0.0255, -0.0050, -0.1170, -0.1046, 0.0309,
    0.1367, 0.1728, -0.0533, -0.0748, -0.0534, 0.1624, 0.0384, -0.1805,
    -0.0707, 0.0642, 0.0220, -0.0134, -0.1333, -0.1505
])
lowerCAmelCase__						:       List[Any]										=						torch.tensor([
    0.1321, 0.1337, 0.0440, 0.0622, -0.0591, -0.0370, 0.0503, 0.2133,
    -0.0177, 0.1415, -0.0116, -0.0112, 0.0044, -0.0980, -0.0789, 0.0395,
    0.1502, 0.1785, -0.0488, -0.0514, -0.0404, 0.1539, 0.0454, -0.1559,
    -0.0665, 0.0659, 0.0383, -0.0005, -0.1266, -0.1386
])
lowerCAmelCase__						:       Union[str, Any]										=						torch.tensor([
    0.1154, 0.1218, 0.0307, 0.0526, -0.0711, -0.0541, 0.0366, 0.2078,
    -0.0267, 0.1317, -0.0226, -0.0193, -0.0014, -0.1055, -0.0902, 0.0330,
    0.1391, 0.1709, -0.0562, -0.0693, -0.0560, 0.1482, 0.0381, -0.1683,
    -0.0681, 0.0661, 0.0331, -0.0046, -0.1268, -0.1431
])
lowerCAmelCase__						:       int										=						torch.tensor([
    0.1192, 0.1240, 0.0414, 0.0606, -0.0557, -0.0412, 0.0430, 0.2042,
    -0.0200, 0.1385, -0.0115, -0.0132, 0.0017, -0.0965, -0.0802, 0.0398,
    0.1433, 0.1747, -0.0458, -0.0533, -0.0407, 0.1545, 0.0419, -0.1574,
    -0.0645, 0.0626, 0.0341, -0.0010, -0.1199, -0.1390
])
lowerCAmelCase__						:       Union[str, Any]										=						torch.tensor([
    0.1075, 0.1074, 0.0205, 0.0431, -0.0774, -0.0607, 0.0298, 0.2042,
    -0.0320, 0.1267, -0.0281, -0.0250, -0.0064, -0.1091, -0.0946, 0.0290,
    0.1328, 0.1650, -0.0580, -0.0738, -0.0586, 0.1440, 0.0337, -0.1746,
    -0.0712, 0.0605, 0.0250, -0.0099, -0.1316, -0.1473
])
lowerCAmelCase__						:       Optional[int]										=						torch.tensor([
    -1.4572, -2.0481, -0.0414, -0.6005, 1.4136, 0.5848, 0.4028, -2.7330,
    1.2212, -2.1228, 0.2155, 0.4039, 0.7662, 2.0535, 0.7477, -0.3243,
    -2.1758, -2.7648, 1.6947, 0.7026, 1.2338, -1.6078, -0.8682, 2.2810,
    1.8574, -0.5718, -0.5586, -0.0186, 2.3415, 2.1251])
lowerCAmelCase__						:       int										=						torch.tensor([
    -1.3690, -1.9720, -0.4090, -0.6966, 1.4660, 0.9938, -0.1385, -2.7324,
    0.7736, -1.8917, 0.2923, 0.4293, 0.1693, 1.4112, 1.1887, -0.3181,
    -2.2160, -2.6381, 1.3170, 0.8163, 0.9240, -1.6544, -0.6099, 2.5259,
    1.6430, -0.9090, -0.9392, -0.0126, 2.4268, 2.3266
])
lowerCAmelCase__						:       List[str]										=						torch.tensor([
    -1.3525, -1.9628, -0.3956, -0.6860, 1.4664, 1.0014, -0.1259, -2.7212,
    0.7772, -1.8811, 0.2996, 0.4388, 0.1704, 1.4029, 1.1701, -0.3027,
    -2.2053, -2.6287, 1.3350, 0.8131, 0.9274, -1.6292, -0.6098, 2.5131,
    1.6505, -0.8958, -0.9298, -0.0151, 2.4257, 2.3355
])
lowerCAmelCase__						:       Any										=						torch.tensor([
    -2.0585, -2.7897, -0.2850, -0.8940, 1.9052, 0.5702, 0.6345, -3.8959,
    1.5932, -3.2319, 0.1974, 0.0287, 1.7566, 2.6543, 0.8387, -0.5351,
    -3.2736, -4.3375, 2.9029, 1.6390, 1.4640, -2.1701, -1.9013, 2.9341,
    3.4981, -0.6255, -1.1644, -0.1591, 3.7097, 3.2066
])
lowerCAmelCase__						:       str										=						torch.tensor([
    -2.3139, -2.5594, -0.0197, -0.6785, 1.7001, 1.1606, 0.3075, -2.1740,
    1.8071, -2.5630, -0.0926, -0.3811, 1.2116, 2.6246, 1.2731, -0.5398,
    -2.8153, -3.6140, 2.3893, 1.3262, 1.6258, -2.1856, -1.3267, 2.8395,
    2.3779, -1.0623, -1.2468, 0.8959, 3.3367, 3.2243
])
lowerCAmelCase__						:       List[str]										=						torch.tensor([
    -2.0628, -2.7667, -0.2089, -0.8263, 2.0539, 0.5992, 0.6495, -3.8336,
    1.6025, -3.2817, 0.1721, -0.0633, 1.7516, 2.7039, 0.8100, -0.5908,
    -3.2113, -4.4343, 2.9257, 1.3632, 1.5562, -2.1489, -1.9894, 3.0560,
    3.3396, -0.7328, -1.0417, 0.0383, 3.7093, 3.2343
])
lowerCAmelCase__						:       str										=						torch.tensor([
    -1.4574, -2.0569, -0.0473, -0.6117, 1.4018, 0.5769, 0.4129, -2.7344,
    1.2241, -2.1397, 0.2000, 0.3937, 0.7616, 2.0453, 0.7324, -0.3391,
    -2.1746, -2.7744, 1.6963, 0.6921, 1.2187, -1.6172, -0.8877, 2.2439,
    1.8471, -0.5839, -0.5605, -0.0464, 2.3250, 2.1219
])
# fmt: on

lowerCAmelCase__						:       List[str]										=						api.list_models(filter="""diffusers""")
for mod in models:
			if "google" in mod.author or mod.modelId == "CompVis/ldm-celebahq-256":
						lowerCAmelCase__						:       str										=						"""/home/patrick/google_checkpoints/""" + mod.modelId.split("""/""")[-1]

						print(f"""Started running {mod.modelId}!!!""")

						if mod.modelId.startswith("""CompVis"""):
									lowerCAmelCase__						:       Tuple										=						UNetaDModel.from_pretrained(local_checkpoint, subfolder="""unet""")
						else:
									lowerCAmelCase__						:       str										=						UNetaDModel.from_pretrained(local_checkpoint)

						torch.manual_seed(0)
						random.seed(0)

						lowerCAmelCase__						:       Tuple										=						torch.randn(1, model.config.in_channels, model.config.sample_size, model.config.sample_size)
						lowerCAmelCase__						:       List[Any]										=						torch.tensor([10] * noise.shape[0])
						with torch.no_grad():
									lowerCAmelCase__						:       Union[str, Any]										=						model(noise, time_step).sample

						assert torch.allclose(
						    logits[0, 0, 0, :30], results["""_""".join("""_""".join(mod.modelId.split("""/""")).split("""-"""))], atol=1E-3
						)
						print(f"""{mod.modelId} has passed successfully!!!""")






'''simple docstring'''


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

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


lowerCAmelCase__						:       Optional[Any]										=						logging.get_logger(__name__)

lowerCAmelCase__						:       Optional[int]										=						{
    """facebook/bart-large""": """https://huggingface.co/facebook/bart-large/resolve/main/config.json""",
    # See all BART models at https://huggingface.co/models?filter=bart
}




class     a			(				SCREAMING_SNAKE_CASE       ):






						"""simple docstring"""




						__UpperCAmelCase			     =    """bart"""
						__UpperCAmelCase			     =    ["""past_key_values"""]
						__UpperCAmelCase			     =    {"""num_attention_heads""": """encoder_attention_heads""", """hidden_size""": """d_model"""}







						def __init__(							self					:      Optional[Any]			,			snake_case_					:      Tuple=5_0_2_6_5			,			snake_case_					:      Dict=1_0_2_4			,			snake_case_					:      int=1_2			,			snake_case_					:      int=4_0_9_6			,			snake_case_					:      str=1_6			,			snake_case_					:      List[Any]=1_2			,			snake_case_					:      List[Any]=4_0_9_6			,			snake_case_					:      Any=1_6			,			snake_case_					:      str=0.0			,			snake_case_					:      Optional[int]=0.0			,			snake_case_					:      List[Any]="gelu"			,			snake_case_					:      List[Any]=1_0_2_4			,			snake_case_					:      Union[str, Any]=0.1			,			snake_case_					:      int=0.0			,			snake_case_					:      Tuple=0.0			,			snake_case_					:      Optional[Any]=0.0_2			,			snake_case_					:      Dict=0.0			,			snake_case_					:      str=False			,			snake_case_					:      Optional[int]=True			,			snake_case_					:      Any=3			,			snake_case_					:      int=1			,			snake_case_					:      int=0			,			snake_case_					:      Optional[Any]=2			,			snake_case_					:      str=True			,			snake_case_					:      int=2			,			snake_case_					:      Union[str, Any]=2			,			**snake_case_					:      int			,			):



								'''simple docstring'''



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

								super().__init__(
								    num_labels=snake_case_			,			pad_token_id=snake_case_			,			bos_token_id=snake_case_			,			eos_token_id=snake_case_			,			is_encoder_decoder=snake_case_			,			decoder_start_token_id=snake_case_			,			forced_eos_token_id=snake_case_			,			**snake_case_			,			)

								# ensure backward compatibility for BART CNN models
								if self.forced_bos_token_id is None and kwargs.get('''force_bos_token_to_be_generated'''			,			snake_case_  ):
										snake_case__						:       Union[str, Any]       						=  self.bos_token_id
										warnings.warn(
										    F"""Please make sure the config includes `forced_bos_token_id={self.bos_token_id}` in future versions. """
										    '''The config can simply be saved and uploaded again to be fixed.'''  )





class     a			(				SCREAMING_SNAKE_CASE       ):






						"""simple docstring"""




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



								'''simple docstring'''



								if self.task in ["default", "seq2seq-lm"]:
										snake_case__						:       Optional[int]       						=  OrderedDict(
										    [
										        ('''input_ids''', {0: '''batch''', 1: '''encoder_sequence'''}),
										        ('''attention_mask''', {0: '''batch''', 1: '''encoder_sequence'''}),
										    ]  )

										if self.use_past:
												snake_case__						:       List[str]       						=  {0: '''batch'''}
												snake_case__						:       int       						=  {0: '''batch''', 1: '''past_decoder_sequence + sequence'''}
										else:
												snake_case__						:       List[Any]       						=  {0: '''batch''', 1: '''decoder_sequence'''}
												snake_case__						:       Optional[int]       						=  {0: '''batch''', 1: '''decoder_sequence'''}

										if self.use_past:
												self.fill_with_past_key_values_(snake_case_			,			direction='''inputs'''  )
								elif self.task == "causal-lm":
										# TODO: figure this case out.
										snake_case__						:       int       						=  OrderedDict(
										    [
										        ('''input_ids''', {0: '''batch''', 1: '''encoder_sequence'''}),
										        ('''attention_mask''', {0: '''batch''', 1: '''encoder_sequence'''}),
										    ]  )
										if self.use_past:
												snake_case__						,       snake_case__						:       Tuple       						=  self.num_layers
												for i in range(snake_case_  ):
														snake_case__						:       List[str]       						=  {0: '''batch''', 2: '''past_sequence + sequence'''}
														snake_case__						:       Tuple       						=  {0: '''batch''', 2: '''past_sequence + sequence'''}
								else:
										snake_case__						:       List[str]       						=  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
						def       __magic_name__							(							self					:      Optional[int]  ):



								'''simple docstring'''



								if self.task in ["default", "seq2seq-lm"]:
										snake_case__						:       str       						=  super().outputs
								else:
										snake_case__						:       List[Any]       						=  super(snake_case_			,			self  ).outputs
										if self.use_past:
												snake_case__						,       snake_case__						:       Dict       						=  self.num_layers
												for i in range(snake_case_  ):
														snake_case__						:       Any       						=  {0: '''batch''', 2: '''past_sequence + sequence'''}
														snake_case__						:       Optional[int]       						=  {0: '''batch''', 2: '''past_sequence + sequence'''}
								return common_outputs







						def       __magic_name__							(							self					:      Optional[Any]			,			snake_case_					:      PreTrainedTokenizer			,			snake_case_					:      int = -1			,			snake_case_					:      int = -1			,			snake_case_					:      bool = False			,			snake_case_					:      Optional[TensorType] = None			,			):



								'''simple docstring'''



								snake_case__						:       Union[str, Any]       						=  self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
								    snake_case_			,			snake_case_			,			snake_case_			,			snake_case_			,			snake_case_  )

								# Generate decoder inputs
								snake_case__						:       Dict       						=  seq_length if not self.use_past else 1
								snake_case__						:       Optional[Any]       						=  self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
								    snake_case_			,			snake_case_			,			snake_case_			,			snake_case_			,			snake_case_  )
								snake_case__						:       List[str]       						=  {F"""decoder_{name}""": tensor for name, tensor in decoder_inputs.items()}
								snake_case__						:       str       						=  dict(**snake_case_			,			**snake_case_  )

								if self.use_past:
										if not is_torch_available():
												raise ValueError('''Cannot generate dummy past_keys inputs without PyTorch installed.'''  )
										else:
												import torch
										snake_case__						,       snake_case__						:       str       						=  common_inputs['''input_ids'''].shape
										snake_case__						:       Dict       						=  common_inputs['''decoder_input_ids'''].shape[1]
										snake_case__						,       snake_case__						:       Optional[Any]       						=  self.num_attention_heads
										snake_case__						:       Tuple       						=  (
										    batch,
										    num_encoder_attention_heads,
										    encoder_seq_length,
										    self._config.hidden_size // num_encoder_attention_heads,
										)
										snake_case__						:       Optional[Any]       						=  decoder_seq_length + 3
										snake_case__						:       int       						=  (
										    batch,
										    num_decoder_attention_heads,
										    decoder_past_length,
										    self._config.hidden_size // num_decoder_attention_heads,
										)

										snake_case__						:       Optional[Any]       						=  torch.cat(
										    [common_inputs['''decoder_attention_mask'''], torch.ones(snake_case_			,			snake_case_  )]			,			dim=1  )

										snake_case__						:       Any       						=  []
										# If the number of encoder and decoder layers are present in the model configuration, both are considered
										snake_case__						,       snake_case__						:       str       						=  self.num_layers
										snake_case__						:       Any       						=  min(snake_case_			,			snake_case_  )
										snake_case__						:       Union[str, Any]       						=  max(snake_case_			,			snake_case_  ) - min_num_layers
										snake_case__						:       Optional[int]       						=  '''encoder''' if num_encoder_layers > num_decoder_layers else '''decoder'''

										for _ in range(snake_case_  ):
												common_inputs["past_key_values"].append(
												    (
												        torch.zeros(snake_case_  ),
												        torch.zeros(snake_case_  ),
												        torch.zeros(snake_case_  ),
												        torch.zeros(snake_case_  ),
												    )  )
										# TODO: test this.
										snake_case__						:       int       						=  encoder_shape if remaining_side_name == '''encoder''' else decoder_shape
										for _ in range(snake_case_			,			snake_case_  ):
												common_inputs["past_key_values"].append((torch.zeros(snake_case_  ), torch.zeros(snake_case_  ))  )
								return common_inputs







						def       __magic_name__							(							self					:      List[str]			,			snake_case_					:      PreTrainedTokenizer			,			snake_case_					:      int = -1			,			snake_case_					:      int = -1			,			snake_case_					:      bool = False			,			snake_case_					:      Optional[TensorType] = None			,			):



								'''simple docstring'''



								snake_case__						:       Any       						=  self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
								    snake_case_			,			snake_case_			,			snake_case_			,			snake_case_			,			snake_case_  )

								if self.use_past:
										if not is_torch_available():
												raise ValueError('''Cannot generate dummy past_keys inputs without PyTorch installed.'''  )
										else:
												import torch
										snake_case__						,       snake_case__						:       Dict       						=  common_inputs['''input_ids'''].shape
										# Not using the same length for past_key_values
										snake_case__						:       str       						=  seqlen + 2
										snake_case__						,       snake_case__						:       Union[str, Any]       						=  self.num_layers
										snake_case__						,       snake_case__						:       Optional[int]       						=  self.num_attention_heads
										snake_case__						:       str       						=  (
										    batch,
										    num_encoder_attention_heads,
										    past_key_values_length,
										    self._config.hidden_size // num_encoder_attention_heads,
										)

										snake_case__						:       List[Any]       						=  common_inputs['''attention_mask'''].dtype
										snake_case__						:       List[Any]       						=  torch.cat(
										    [common_inputs['''attention_mask'''], torch.ones(snake_case_			,			snake_case_			,			dtype=snake_case_  )]			,			dim=1  )
										snake_case__						:       Optional[Any]       						=  [
										    (torch.zeros(snake_case_  ), torch.zeros(snake_case_  )) for _ in range(snake_case_  )
										]
								return common_inputs







						def       __magic_name__							(							self					:      int			,			snake_case_					:      PreTrainedTokenizer			,			snake_case_					:      int = -1			,			snake_case_					:      int = -1			,			snake_case_					:      bool = False			,			snake_case_					:      Optional[TensorType] = None			,			):



								'''simple docstring'''



								snake_case__						:       str       						=  compute_effective_axis_dimension(
								    snake_case_			,			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
								snake_case__						:       List[Any]       						=  tokenizer.num_special_tokens_to_add(snake_case_  )
								snake_case__						:       List[str]       						=  compute_effective_axis_dimension(
								    snake_case_			,			fixed_dimension=OnnxConfig.default_fixed_sequence			,			num_token_to_add=snake_case_  )

								# Generate dummy inputs according to compute batch and sequence
								snake_case__						:       Any       						=  [''' '''.join([tokenizer.unk_token]  ) * seq_length] * batch_size
								snake_case__						:       Any       						=  dict(tokenizer(snake_case_			,			return_tensors=snake_case_  )  )
								return common_inputs







						def       __magic_name__							(							self					:      Optional[Any]			,			snake_case_					:      PreTrainedTokenizer			,			snake_case_					:      int = -1			,			snake_case_					:      int = -1			,			snake_case_					:      bool = False			,			snake_case_					:      Optional[TensorType] = None			,			):



								'''simple docstring'''



								if self.task in ["default", "seq2seq-lm"]:
										snake_case__						:       int       						=  self._generate_dummy_inputs_for_default_and_seqaseq_lm(
										    snake_case_			,			batch_size=snake_case_			,			seq_length=snake_case_			,			is_pair=snake_case_			,			framework=snake_case_  )

								elif self.task == "causal-lm":
										snake_case__						:       Optional[Any]       						=  self._generate_dummy_inputs_for_causal_lm(
										    snake_case_			,			batch_size=snake_case_			,			seq_length=snake_case_			,			is_pair=snake_case_			,			framework=snake_case_  )
								else:
										snake_case__						:       Optional[int]       						=  self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
										    snake_case_			,			batch_size=snake_case_			,			seq_length=snake_case_			,			is_pair=snake_case_			,			framework=snake_case_  )

								return common_inputs







						def       __magic_name__							(							self					:      Tuple			,			snake_case_					:      List[str]			,			snake_case_					:      Union[str, Any]			,			snake_case_					:      Any			,			snake_case_					:      Any  ):



								'''simple docstring'''



								if self.task in ["default", "seq2seq-lm"]:
										snake_case__						:       Optional[Any]       						=  super()._flatten_past_key_values_(snake_case_			,			snake_case_			,			snake_case_			,			snake_case_  )
								else:
										snake_case__						:       Union[str, Any]       						=  super(snake_case_			,			self  )._flatten_past_key_values_(
										    snake_case_			,			snake_case_			,			snake_case_			,			snake_case_  )



