



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


snake_case__					:     Dict    		=   logging.get_logger(__name__)

snake_case__					:     List[Any]    		=   {
    """Helsinki-NLP/opus-mt-en-de""": """https://huggingface.co/Helsinki-NLP/opus-mt-en-de/resolve/main/config.json""",
    # See all Marian models at https://huggingface.co/models?filter=marian
}

class   _A   ( _lowercase	):


						'''simple docstring'''





						_snake_case			:     Tuple      							=      """marian"""
						_snake_case			:     int      							=      ["""past_key_values"""]
						_snake_case			:     Any      							=      {"""num_attention_heads""": """encoder_attention_heads""", """hidden_size""": """d_model"""}



						def __init__(	self	:							str							,      lowerCamelCase	:							Dict=58_101							,      lowerCamelCase	:							Any=None							,      lowerCamelCase	:							Optional[Any]=1_024							,      lowerCamelCase	:							Tuple=12							,      lowerCamelCase	:							Dict=4_096							,      lowerCamelCase	:							Optional[int]=16							,      lowerCamelCase	:							Union[str, Any]=12							,      lowerCamelCase	:							Optional[int]=4_096							,      lowerCamelCase	:							Any=16							,      lowerCamelCase	:							Any=0.0							,      lowerCamelCase	:							Optional[Any]=0.0							,      lowerCamelCase	:							List[Any]=True							,      lowerCamelCase	:							Optional[Any]=True							,      lowerCamelCase	:							Union[str, Any]="gelu"							,      lowerCamelCase	:							Any=1_024							,      lowerCamelCase	:							Dict=0.1							,      lowerCamelCase	:							List[Any]=0.0							,      lowerCamelCase	:							Optional[int]=0.0							,      lowerCamelCase	:							Dict=0.02							,      lowerCamelCase	:							Any=58_100							,      lowerCamelCase	:							Optional[int]=False							,      lowerCamelCase	:							Union[str, Any]=58_100							,      lowerCamelCase	:							Dict=0							,      lowerCamelCase	:							Union[str, Any]=0							,      lowerCamelCase	:							Union[str, Any]=True							,      **lowerCamelCase	:							Optional[Any]							,      ):






							'''simple docstring'''



							__lowercase		       =			vocab_size
							__lowercase		       =			decoder_vocab_size or vocab_size
							__lowercase		       =			max_position_embeddings
							__lowercase		       =			d_model
							__lowercase		       =			encoder_ffn_dim
							__lowercase		       =			encoder_layers
							__lowercase		       =			encoder_attention_heads
							__lowercase		       =			decoder_ffn_dim
							__lowercase		       =			decoder_layers
							__lowercase		       =			decoder_attention_heads
							__lowercase		       =			dropout
							__lowercase		       =			attention_dropout
							__lowercase		       =			activation_dropout
							__lowercase		       =			activation_function
							__lowercase		       =			init_std
							__lowercase		       =			encoder_layerdrop
							__lowercase		       =			decoder_layerdrop
							__lowercase		       =			use_cache
							__lowercase		       =			encoder_layers
							__lowercase		       =			scale_embedding  # scale factor will be sqrt(d_model) if True
							__lowercase		       =			share_encoder_decoder_embeddings
							super().__init__(
							    pad_token_id=lowerCamelCase							,      eos_token_id=lowerCamelCase							,      is_encoder_decoder=lowerCamelCase							,      decoder_start_token_id=lowerCamelCase							,      forced_eos_token_id=lowerCamelCase							,      **lowerCamelCase							,      )

class   _A   ( _lowercase	):


						'''simple docstring'''





						@property
						# Copied from transformers.models.bart.configuration_bart.BartOnnxConfig.inputs
						def     _snake_case	(	self	:							Union[str, Any]	):






							'''simple docstring'''



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

								if self.use_past:
									__lowercase		       =			{0: "batch"}
									__lowercase		       =			{0: "batch", 1: "past_decoder_sequence + sequence"}
								else:
									__lowercase		       =			{0: "batch", 1: "decoder_sequence"}
									__lowercase		       =			{0: "batch", 1: "decoder_sequence"}

								if self.use_past:
									self.fill_with_past_key_values_(lowerCamelCase							,      direction="inputs"	)
							elif self.task == "causal-lm":
								# TODO: figure this case out.
								__lowercase		       =			OrderedDict(
								    [
								        ("input_ids", {0: "batch", 1: "encoder_sequence"}),
								        ("attention_mask", {0: "batch", 1: "encoder_sequence"}),
								    ]	)
								if self.use_past:
									__lowercase       ,    __lowercase		       =			self.num_layers
									for i in range(lowerCamelCase	):
										__lowercase		       =			{0: "batch", 2: "past_sequence + sequence"}
										__lowercase		       =			{0: "batch", 2: "past_sequence + sequence"}
							else:
								__lowercase		       =			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
						# Copied from transformers.models.bart.configuration_bart.BartOnnxConfig.outputs
						def     _snake_case	(	self	:							Optional[int]	):






							'''simple docstring'''



							if self.task in ["default", "seq2seq-lm"]:
								__lowercase		       =			super().outputs
							else:
								__lowercase		       =			super(lowerCamelCase							,      self	).outputs
								if self.use_past:
									__lowercase       ,    __lowercase		       =			self.num_layers
									for i in range(lowerCamelCase	):
										__lowercase		       =			{0: "batch", 2: "past_sequence + sequence"}
										__lowercase		       =			{0: "batch", 2: "past_sequence + sequence"}
							return common_outputs



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






							'''simple docstring'''



							__lowercase		       =			self._generate_dummy_inputs_for_encoder_and_decoder(
							    lowerCamelCase							,      lowerCamelCase							,      lowerCamelCase							,      lowerCamelCase							,      lowerCamelCase	)

							# Generate decoder inputs
							__lowercase		       =			seq_length if not self.use_past else 1
							__lowercase		       =			self._generate_dummy_inputs_for_encoder_and_decoder(
							    lowerCamelCase							,      lowerCamelCase							,      lowerCamelCase							,      lowerCamelCase							,      lowerCamelCase	)
							__lowercase		       =			{f"""decoder_{name}""": tensor for name, tensor in decoder_inputs.items()}
							__lowercase		       =			dict(**lowerCamelCase							,      **lowerCamelCase	)

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

								__lowercase		       =			torch.cat(
								    [common_inputs["decoder_attention_mask"], torch.ones(lowerCamelCase							,      lowerCamelCase	)]							,      dim=1	)

								__lowercase		       =			[]
								# If the number of encoder and decoder layers are present in the model configuration, both are considered
								__lowercase       ,    __lowercase		       =			self.num_layers
								__lowercase		       =			min(lowerCamelCase							,      lowerCamelCase	)
								__lowercase		       =			max(lowerCamelCase							,      lowerCamelCase	) - min_num_layers
								__lowercase		       =			"encoder" if num_encoder_layers > num_decoder_layers else "decoder"

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



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






							'''simple docstring'''



							__lowercase		       =			self._generate_dummy_inputs_for_encoder_and_decoder(
							    lowerCamelCase							,      lowerCamelCase							,      lowerCamelCase							,      lowerCamelCase							,      lowerCamelCase	)

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

								__lowercase		       =			common_inputs["attention_mask"].dtype
								__lowercase		       =			torch.cat(
								    [common_inputs["attention_mask"], torch.ones(lowerCamelCase							,      lowerCamelCase							,      dtype=lowerCamelCase	)]							,      dim=1	)
								__lowercase		       =			[
								    (torch.zeros(lowerCamelCase	), torch.zeros(lowerCamelCase	)) for _ in range(lowerCamelCase	)
								]
							return common_inputs



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






							'''simple docstring'''



							__lowercase		       =			compute_effective_axis_dimension(
							    lowerCamelCase							,      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
							__lowercase		       =			tokenizer.num_special_tokens_to_add(lowerCamelCase	)
							__lowercase		       =			compute_effective_axis_dimension(
							    lowerCamelCase							,      fixed_dimension=OnnxConfig.default_fixed_sequence							,      num_token_to_add=lowerCamelCase	)

							# Generate dummy inputs according to compute batch and sequence
							__lowercase		       =			[" ".join([tokenizer.unk_token]	) * seq_length] * batch_size
							__lowercase		       =			dict(tokenizer(lowerCamelCase							,      return_tensors=lowerCamelCase	)	)
							return common_inputs



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






							'''simple docstring'''



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

							else:
								__lowercase		       =			self._generate_dummy_inputs_for_causal_lm(
								    lowerCamelCase							,      batch_size=lowerCamelCase							,      seq_length=lowerCamelCase							,      is_pair=lowerCamelCase							,      framework=lowerCamelCase	)

							return common_inputs



						def     _snake_case	(	self	:							str							,      lowerCamelCase	:							List[Any]							,      lowerCamelCase	:							str							,      lowerCamelCase	:							List[str]							,      lowerCamelCase	:							Tuple	):






							'''simple docstring'''



							if self.task in ["default", "seq2seq-lm"]:
								__lowercase		       =			super()._flatten_past_key_values_(lowerCamelCase							,      lowerCamelCase							,      lowerCamelCase							,      lowerCamelCase	)
							else:
								__lowercase		       =			super(lowerCamelCase							,      self	)._flatten_past_key_values_(
								    lowerCamelCase							,      lowerCamelCase							,      lowerCamelCase							,      lowerCamelCase	)



						@property
						def     _snake_case	(	self	:							Dict	):






							'''simple docstring'''



							return 1e-4








import tempfile

import numpy as np
import torch
from transformers import AutoTokenizer, TaEncoderModel

from diffusers import DDPMScheduler, UNetaDConditionModel
from diffusers.models.attention_processor import AttnAddedKVProcessor
from diffusers.pipelines.deepfloyd_if import IFWatermarker
from diffusers.utils.testing_utils import torch_device

from ..test_pipelines_common import to_np

class   _A   :


						'''simple docstring'''





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






							'''simple docstring'''



							torch.manual_seed(0	)
							__lowercase		       =			TaEncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5"	)

							torch.manual_seed(0	)
							__lowercase		       =			AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5"	)

							torch.manual_seed(0	)
							__lowercase		       =			UNetaDConditionModel(
							    sample_size=32							,      layers_per_block=1							,      block_out_channels=[32, 64]							,      down_block_types=[
							        "ResnetDownsampleBlock2D",
							        "SimpleCrossAttnDownBlock2D",
							    ]							,      mid_block_type="UNetMidBlock2DSimpleCrossAttn"							,      up_block_types=["SimpleCrossAttnUpBlock2D", "ResnetUpsampleBlock2D"]							,      in_channels=3							,      out_channels=6							,      cross_attention_dim=32							,      encoder_hid_dim=32							,      attention_head_dim=8							,      addition_embed_type="text"							,      addition_embed_type_num_heads=2							,      cross_attention_norm="group_norm"							,      resnet_time_scale_shift="scale_shift"							,      act_fn="gelu"							,      )
							unet.set_attn_processor(AttnAddedKVProcessor()	)  # For reproducibility tests

							torch.manual_seed(0	)
							__lowercase		       =			DDPMScheduler(
							    num_train_timesteps=1_000							,      beta_schedule="squaredcos_cap_v2"							,      beta_start=0.0001							,      beta_end=0.02							,      thresholding=lowerCamelCase							,      dynamic_thresholding_ratio=0.95							,      sample_max_value=1.0							,      prediction_type="epsilon"							,      variance_type="learned_range"							,      )

							torch.manual_seed(0	)
							__lowercase		       =			IFWatermarker()

							return {
							    "text_encoder": text_encoder,
							    "tokenizer": tokenizer,
							    "unet": unet,
							    "scheduler": scheduler,
							    "watermarker": watermarker,
							    "safety_checker": None,
							    "feature_extractor": None,
							}



						def     _snake_case	(	self	:							Tuple	):






							'''simple docstring'''



							torch.manual_seed(0	)
							__lowercase		       =			TaEncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5"	)

							torch.manual_seed(0	)
							__lowercase		       =			AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5"	)

							torch.manual_seed(0	)
							__lowercase		       =			UNetaDConditionModel(
							    sample_size=32							,      layers_per_block=[1, 2]							,      block_out_channels=[32, 64]							,      down_block_types=[
							        "ResnetDownsampleBlock2D",
							        "SimpleCrossAttnDownBlock2D",
							    ]							,      mid_block_type="UNetMidBlock2DSimpleCrossAttn"							,      up_block_types=["SimpleCrossAttnUpBlock2D", "ResnetUpsampleBlock2D"]							,      in_channels=6							,      out_channels=6							,      cross_attention_dim=32							,      encoder_hid_dim=32							,      attention_head_dim=8							,      addition_embed_type="text"							,      addition_embed_type_num_heads=2							,      cross_attention_norm="group_norm"							,      resnet_time_scale_shift="scale_shift"							,      act_fn="gelu"							,      class_embed_type="timestep"							,      mid_block_scale_factor=1.414							,      time_embedding_act_fn="gelu"							,      time_embedding_dim=32							,      )
							unet.set_attn_processor(AttnAddedKVProcessor()	)  # For reproducibility tests

							torch.manual_seed(0	)
							__lowercase		       =			DDPMScheduler(
							    num_train_timesteps=1_000							,      beta_schedule="squaredcos_cap_v2"							,      beta_start=0.0001							,      beta_end=0.02							,      thresholding=lowerCamelCase							,      dynamic_thresholding_ratio=0.95							,      sample_max_value=1.0							,      prediction_type="epsilon"							,      variance_type="learned_range"							,      )

							torch.manual_seed(0	)
							__lowercase		       =			DDPMScheduler(
							    num_train_timesteps=1_000							,      beta_schedule="squaredcos_cap_v2"							,      beta_start=0.0001							,      beta_end=0.02							,      )

							torch.manual_seed(0	)
							__lowercase		       =			IFWatermarker()

							return {
							    "text_encoder": text_encoder,
							    "tokenizer": tokenizer,
							    "unet": unet,
							    "scheduler": scheduler,
							    "image_noising_scheduler": image_noising_scheduler,
							    "watermarker": watermarker,
							    "safety_checker": None,
							    "feature_extractor": None,
							}



						def     _snake_case	(	self	:							str	):






							'''simple docstring'''



							__lowercase		       =			self.get_dummy_components()
							__lowercase		       =			self.pipeline_class(**lowerCamelCase	)
							pipe.to(lowerCamelCase	)
							pipe.set_progress_bar_config(disable=lowerCamelCase	)

							__lowercase		       =			self.get_dummy_inputs(lowerCamelCase	)

							__lowercase		       =			inputs["prompt"]
							__lowercase		       =			inputs["generator"]
							__lowercase		       =			inputs["num_inference_steps"]
							__lowercase		       =			inputs["output_type"]

							if "image" in inputs:
								__lowercase		       =			inputs["image"]
							else:
								__lowercase		       =			None

							if "mask_image" in inputs:
								__lowercase		       =			inputs["mask_image"]
							else:
								__lowercase		       =			None

							if "original_image" in inputs:
								__lowercase		       =			inputs["original_image"]
							else:
								__lowercase		       =			None

							__lowercase       ,    __lowercase		       =			pipe.encode_prompt(lowerCamelCase	)

							# inputs with prompt converted to embeddings
							__lowercase		       =			{
							    "prompt_embeds": prompt_embeds,
							    "negative_prompt_embeds": negative_prompt_embeds,
							    "generator": generator,
							    "num_inference_steps": num_inference_steps,
							    "output_type": output_type,
							}

							if image is not None:
								__lowercase		       =			image

							if mask_image is not None:
								__lowercase		       =			mask_image

							if original_image is not None:
								__lowercase		       =			original_image

							# set all optional components to None
							for optional_component in pipe._optional_components:
								setattr(lowerCamelCase							,      lowerCamelCase							,      lowerCamelCase	)

							__lowercase		       =			pipe(**lowerCamelCase	)[0]

							with tempfile.TemporaryDirectory() as tmpdir:
								pipe.save_pretrained(lowerCamelCase	)
								__lowercase		       =			self.pipeline_class.from_pretrained(lowerCamelCase	)
								pipe_loaded.to(lowerCamelCase	)
								pipe_loaded.set_progress_bar_config(disable=lowerCamelCase	)

							pipe_loaded.unet.set_attn_processor(AttnAddedKVProcessor()	)  # For reproducibility tests

							for optional_component in pipe._optional_components:
								self.assertTrue(
								    getattr(lowerCamelCase							,      lowerCamelCase	) is None							,      f"""`{optional_component}` did not stay set to None after loading."""							,      )

							__lowercase		       =			self.get_dummy_inputs(lowerCamelCase	)

							__lowercase		       =			inputs["generator"]
							__lowercase		       =			inputs["num_inference_steps"]
							__lowercase		       =			inputs["output_type"]

							# inputs with prompt converted to embeddings
							__lowercase		       =			{
							    "prompt_embeds": prompt_embeds,
							    "negative_prompt_embeds": negative_prompt_embeds,
							    "generator": generator,
							    "num_inference_steps": num_inference_steps,
							    "output_type": output_type,
							}

							if image is not None:
								__lowercase		       =			image

							if mask_image is not None:
								__lowercase		       =			mask_image

							if original_image is not None:
								__lowercase		       =			original_image

							__lowercase		       =			pipe_loaded(**lowerCamelCase	)[0]

							__lowercase		       =			np.abs(to_np(lowerCamelCase	) - to_np(lowerCamelCase	)	).max()
							self.assertLess(lowerCamelCase							,      1e-4	)



						def     _snake_case	(	self	:							Optional[Any]	):






							'''simple docstring'''



							__lowercase		       =			self.get_dummy_components()
							__lowercase		       =			self.pipeline_class(**lowerCamelCase	)
							pipe.to(lowerCamelCase	)
							pipe.set_progress_bar_config(disable=lowerCamelCase	)

							__lowercase		       =			self.get_dummy_inputs(lowerCamelCase	)
							__lowercase		       =			pipe(**lowerCamelCase	)[0]

							with tempfile.TemporaryDirectory() as tmpdir:
								pipe.save_pretrained(lowerCamelCase	)
								__lowercase		       =			self.pipeline_class.from_pretrained(lowerCamelCase	)
								pipe_loaded.to(lowerCamelCase	)
								pipe_loaded.set_progress_bar_config(disable=lowerCamelCase	)

							pipe_loaded.unet.set_attn_processor(AttnAddedKVProcessor()	)  # For reproducibility tests

							__lowercase		       =			self.get_dummy_inputs(lowerCamelCase	)
							__lowercase		       =			pipe_loaded(**lowerCamelCase	)[0]

							__lowercase		       =			np.abs(to_np(lowerCamelCase	) - to_np(lowerCamelCase	)	).max()
							self.assertLess(lowerCamelCase							,      1e-4	)



