




import gc
import random
import tempfile
import unittest

import numpy as np
import torch
from PIL import Image
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    DDIMInverseScheduler,
    DDIMScheduler,
    DPMSolverMultistepInverseScheduler,
    DPMSolverMultistepScheduler,
    StableDiffusionDiffEditPipeline,
    UNetaDConditionModel,
)
from diffusers.utils import load_image, slow
from diffusers.utils.testing_utils import enable_full_determinism, floats_tensor, require_torch_gpu, torch_device

from ..pipeline_params import TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS, TEXT_GUIDED_IMAGE_INPAINTING_PARAMS
from ..test_pipelines_common import PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()




class 						__lowercase		(UpperCamelCase__					,		UpperCamelCase__					,		unittest.TestCase							):




      """simple docstring"""





      _snake_case												=						StableDiffusionDiffEditPipeline
      _snake_case												=						TEXT_GUIDED_IMAGE_INPAINTING_PARAMS - {"""height""", """width""", """image"""} | {"""image_latents"""}
      _snake_case												=						TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS - {"""image"""} | {"""image_latents"""}
      _snake_case												=						frozenset(
          []							)  # TO-DO: update image_params once pipeline is refactored with VaeImageProcessor.preprocess
      _snake_case												=						frozenset([]							)

      def      UpperCAmelCase				(				self     )		->    Tuple:
         torch.manual_seed(0     )
         snake_case				:    List[Any]           =	UNetaDConditionModel(
             block_out_channels=(3_2, 6_4)    ,				layers_per_block=2    ,				sample_size=3_2    ,				in_channels=4    ,				out_channels=4    ,				down_block_types=("""DownBlock2D""", """CrossAttnDownBlock2D""")    ,				up_block_types=("""CrossAttnUpBlock2D""", """UpBlock2D""")    ,				cross_attention_dim=3_2    ,				attention_head_dim=(2, 4)    ,				use_linear_projection=__A    ,				)
         snake_case				:    Any           =	DDIMScheduler(
             beta_start=0.0_00_85    ,				beta_end=0.0_12    ,				beta_schedule="""scaled_linear"""    ,				clip_sample=__A    ,				set_alpha_to_one=__A    ,				)
         snake_case				:    Optional[int]           =	DDIMInverseScheduler(
             beta_start=0.0_00_85    ,				beta_end=0.0_12    ,				beta_schedule="""scaled_linear"""    ,				clip_sample=__A    ,				set_alpha_to_zero=__A    ,				)
         torch.manual_seed(0     )
         snake_case				:    Any           =	AutoencoderKL(
             block_out_channels=[3_2, 6_4]    ,				in_channels=3    ,				out_channels=3    ,				down_block_types=["""DownEncoderBlock2D""", """DownEncoderBlock2D"""]    ,				up_block_types=["""UpDecoderBlock2D""", """UpDecoderBlock2D"""]    ,				latent_channels=4    ,				sample_size=1_2_8    ,				)
         torch.manual_seed(0     )
         snake_case				:    Optional[int]           =	CLIPTextConfig(
             bos_token_id=0    ,				eos_token_id=2    ,				hidden_size=3_2    ,				intermediate_size=3_7    ,				layer_norm_eps=1e-0_5    ,				num_attention_heads=4    ,				num_hidden_layers=5    ,				pad_token_id=1    ,				vocab_size=1_0_0_0    ,				hidden_act="""gelu"""    ,				projection_dim=5_1_2    ,				)
         snake_case				:    List[Any]           =	CLIPTextModel(__A     )
         snake_case				:    int           =	CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""     )

         snake_case				:    Union[str, Any]           =	{
             """unet""": unet,
             """scheduler""": scheduler,
             """inverse_scheduler""": inverse_scheduler,
             """vae""": vae,
             """text_encoder""": text_encoder,
             """tokenizer""": tokenizer,
             """safety_checker""": None,
             """feature_extractor""": None,
         }

         return components

      def      UpperCAmelCase				(				self    ,				A    ,				A=0     )		->    Any:
         snake_case				:    List[str]           =	floats_tensor((1, 1_6, 1_6)    ,				rng=random.Random(__A     )     ).to(__A     )
         snake_case				:    List[Any]           =	floats_tensor((1, 2, 4, 1_6, 1_6)    ,				rng=random.Random(__A     )     ).to(__A     )
         if str(__A     ).startswith("""mps"""     ):
            snake_case				:    Optional[Any]           =	torch.manual_seed(__A     )
         else:
            snake_case				:    Optional[int]           =	torch.Generator(device=__A     ).manual_seed(__A     )
         snake_case				:    int           =	{
             """prompt""": """a dog and a newt""",
             """mask_image""": mask,
             """image_latents""": latents,
             """generator""": generator,
             """num_inference_steps""": 2,
             """inpaint_strength""": 1.0,
             """guidance_scale""": 6.0,
             """output_type""": """numpy""",
         }

         return inputs

      def      UpperCAmelCase				(				self    ,				A    ,				A=0     )		->    Union[str, Any]:
         snake_case				:    Tuple           =	floats_tensor((1, 3, 3_2, 3_2)    ,				rng=random.Random(__A     )     ).to(__A     )
         snake_case				:    List[Any]           =	image.cpu().permute(0    ,				2    ,				3    ,				1     )[0]
         snake_case				:    Tuple           =	Image.fromarray(np.uinta(__A     )     ).convert("""RGB"""     )
         if str(__A     ).startswith("""mps"""     ):
            snake_case				:    Tuple           =	torch.manual_seed(__A     )
         else:
            snake_case				:    Tuple           =	torch.Generator(device=__A     ).manual_seed(__A     )
         snake_case				:    List[Any]           =	{
             """image""": image,
             """source_prompt""": """a cat and a frog""",
             """target_prompt""": """a dog and a newt""",
             """generator""": generator,
             """num_inference_steps""": 2,
             """num_maps_per_mask""": 2,
             """mask_encode_strength""": 1.0,
             """guidance_scale""": 6.0,
             """output_type""": """numpy""",
         }

         return inputs

      def      UpperCAmelCase				(				self    ,				A    ,				A=0     )		->    int:
         snake_case				:    Tuple           =	floats_tensor((1, 3, 3_2, 3_2)    ,				rng=random.Random(__A     )     ).to(__A     )
         snake_case				:    Any           =	image.cpu().permute(0    ,				2    ,				3    ,				1     )[0]
         snake_case				:    Optional[Any]           =	Image.fromarray(np.uinta(__A     )     ).convert("""RGB"""     )
         if str(__A     ).startswith("""mps"""     ):
            snake_case				:    int           =	torch.manual_seed(__A     )
         else:
            snake_case				:    List[str]           =	torch.Generator(device=__A     ).manual_seed(__A     )
         snake_case				:    Union[str, Any]           =	{
             """image""": image,
             """prompt""": """a cat and a frog""",
             """generator""": generator,
             """num_inference_steps""": 2,
             """inpaint_strength""": 1.0,
             """guidance_scale""": 6.0,
             """decode_latents""": True,
             """output_type""": """numpy""",
         }
         return inputs

      def      UpperCAmelCase				(				self     )		->    Optional[int]:
         if not hasattr(self.pipeline_class    ,				"""_optional_components"""     ):
            return

         snake_case				:    str           =	self.get_dummy_components()
         snake_case				:    Optional[int]           =	self.pipeline_class(**__A     )
         pipe.to(__A     )
         pipe.set_progress_bar_config(disable=__A     )

         # set all optional components to None and update pipeline config accordingly
         for optional_component in pipe._optional_components:
            setattr(__A    ,				__A    ,				__A     )
         pipe.register_modules(**{optional_component: None for optional_component in pipe._optional_components}     )

         snake_case				:    Union[str, Any]           =	self.get_dummy_inputs(__A     )
         snake_case				:    Any           =	pipe(**__A     )[0]

         with tempfile.TemporaryDirectory() as tmpdir:
            pipe.save_pretrained(__A     )
            snake_case				:    Optional[Any]           =	self.pipeline_class.from_pretrained(__A     )
            pipe_loaded.to(__A     )
            pipe_loaded.set_progress_bar_config(disable=__A     )

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

         snake_case				:    Union[str, Any]           =	self.get_dummy_inputs(__A     )
         snake_case				:    Tuple           =	pipe_loaded(**__A     )[0]

         snake_case				:    List[Any]           =	np.abs(output - output_loaded     ).max()
         self.assertLess(__A    ,				1e-4     )

      def      UpperCAmelCase				(				self     )		->    Union[str, Any]:
         snake_case				:    Any           =	"""cpu"""

         snake_case				:    str           =	self.get_dummy_components()
         snake_case				:    str           =	self.pipeline_class(**__A     )
         pipe.to(__A     )
         pipe.set_progress_bar_config(disable=__A     )

         snake_case				:    Any           =	self.get_dummy_mask_inputs(__A     )
         snake_case				:    List[Any]           =	pipe.generate_mask(**__A     )
         snake_case				:    int           =	mask[0, -3:, -3:]

         self.assertEqual(mask.shape    ,				(1, 1_6, 1_6)     )
         snake_case				:    List[str]           =	np.array([0] * 9     )
         snake_case				:    Union[str, Any]           =	np.abs(mask_slice.flatten() - expected_slice     ).max()
         self.assertLessEqual(__A    ,				1e-3     )
         self.assertEqual(mask[0, -3, -4]    ,				0     )

      def      UpperCAmelCase				(				self     )		->    Optional[int]:
         snake_case				:    Optional[Any]           =	"""cpu"""

         snake_case				:    int           =	self.get_dummy_components()
         snake_case				:    Tuple           =	self.pipeline_class(**__A     )
         pipe.to(__A     )
         pipe.set_progress_bar_config(disable=__A     )

         snake_case				:    Optional[int]           =	self.get_dummy_inversion_inputs(__A     )
         snake_case				:    str           =	pipe.invert(**__A     ).images
         snake_case				:    Any           =	image[0, -1, -3:, -3:]

         self.assertEqual(image.shape    ,				(2, 3_2, 3_2, 3)     )
         snake_case				:    Optional[int]           =	np.array(
             [0.51_50, 0.51_34, 0.50_43, 0.53_76, 0.46_94, 0.5_10_50, 0.50_15, 0.44_07, 0.47_99]    ,				)
         snake_case				:    Optional[Any]           =	np.abs(image_slice.flatten() - expected_slice     ).max()
         self.assertLessEqual(__A    ,				1e-3     )

      def      UpperCAmelCase				(				self     )		->    Any:
         super().test_inference_batch_single_identical(expected_max_diff=5e-3     )

      def      UpperCAmelCase				(				self     )		->    Union[str, Any]:
         snake_case				:    str           =	"""cpu"""

         snake_case				:    Tuple           =	self.get_dummy_components()

         snake_case				:    List[Any]           =	{"""beta_start""": 0.0_00_85, """beta_end""": 0.0_12, """beta_schedule""": """scaled_linear"""}
         snake_case				:    List[Any]           =	DPMSolverMultistepScheduler(**__A     )
         snake_case				:    List[str]           =	DPMSolverMultistepInverseScheduler(**__A     )

         snake_case				:    Union[str, Any]           =	self.pipeline_class(**__A     )
         pipe.to(__A     )
         pipe.set_progress_bar_config(disable=__A     )

         snake_case				:    Optional[int]           =	self.get_dummy_inversion_inputs(__A     )
         snake_case				:    Optional[int]           =	pipe.invert(**__A     ).images
         snake_case				:    Optional[Any]           =	image[0, -1, -3:, -3:]

         self.assertEqual(image.shape    ,				(2, 3_2, 3_2, 3)     )
         snake_case				:    Optional[Any]           =	np.array(
             [0.51_50, 0.51_34, 0.50_43, 0.53_76, 0.46_94, 0.5_10_50, 0.50_15, 0.44_07, 0.47_99]    ,				)
         snake_case				:    Dict           =	np.abs(image_slice.flatten() - expected_slice     ).max()
         self.assertLessEqual(__A    ,				1e-3     )




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




      """simple docstring"""





      def      UpperCAmelCase				(				self     )		->    Tuple:
         super().tearDown()
         gc.collect()
         torch.cuda.empty_cache()

      @classmethod
      def      UpperCAmelCase				(				cls     )		->    Union[str, Any]:
         snake_case				:    Tuple           =	load_image(
             """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/diffedit/fruit.png"""     )

         snake_case				:    Any           =	raw_image.convert("""RGB"""     ).resize((7_6_8, 7_6_8)     )

         snake_case				:    List[Any]           =	raw_image

      def      UpperCAmelCase				(				self     )		->    str:
         snake_case				:    Any           =	torch.manual_seed(0     )

         snake_case				:    str           =	StableDiffusionDiffEditPipeline.from_pretrained(
             """stabilityai/stable-diffusion-2-1"""    ,				safety_checker=__A    ,				torch_dtype=torch.floataa     )
         snake_case				:    Union[str, Any]           =	DDIMScheduler.from_config(pipe.scheduler.config     )
         snake_case				:    List[Any]           =	DDIMInverseScheduler.from_config(pipe.scheduler.config     )
         pipe.enable_model_cpu_offload()
         pipe.set_progress_bar_config(disable=__A     )

         snake_case				:    List[Any]           =	"""a bowl of fruit"""
         snake_case				:    int           =	"""a bowl of pears"""

         snake_case				:    Any           =	pipe.generate_mask(
             image=self.raw_image    ,				source_prompt=__A    ,				target_prompt=__A    ,				generator=__A    ,				)

         snake_case				:    str           =	pipe.invert(
             prompt=__A    ,				image=self.raw_image    ,				inpaint_strength=0.7    ,				generator=__A     ).latents

         snake_case				:    Optional[int]           =	pipe(
             prompt=__A    ,				mask_image=__A    ,				image_latents=__A    ,				generator=__A    ,				negative_prompt=__A    ,				inpaint_strength=0.7    ,				output_type="""numpy"""    ,				).images[0]

         snake_case				:    List[Any]           =	(
             np.array(
                 load_image(
                     """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"""
                     """/diffedit/pears.png"""     ).resize((7_6_8, 7_6_8)     )     )
             / 2_5_5
         )
         assert np.abs((expected_image - image).max()     ) < 5e-1

      def      UpperCAmelCase				(				self     )		->    Optional[int]:
         snake_case				:    Optional[Any]           =	torch.manual_seed(0     )

         snake_case				:    Tuple           =	StableDiffusionDiffEditPipeline.from_pretrained(
             """stabilityai/stable-diffusion-2-1"""    ,				safety_checker=__A    ,				torch_dtype=torch.floataa     )
         snake_case				:    Dict           =	DPMSolverMultistepScheduler.from_config(pipe.scheduler.config     )
         snake_case				:    Optional[int]           =	DPMSolverMultistepInverseScheduler.from_config(pipe.scheduler.config     )
         pipe.enable_model_cpu_offload()
         pipe.set_progress_bar_config(disable=__A     )

         snake_case				:    Union[str, Any]           =	"""a bowl of fruit"""
         snake_case				:    Dict           =	"""a bowl of pears"""

         snake_case				:    str           =	pipe.generate_mask(
             image=self.raw_image    ,				source_prompt=__A    ,				target_prompt=__A    ,				generator=__A    ,				)

         snake_case				:    Tuple           =	pipe.invert(
             prompt=__A    ,				image=self.raw_image    ,				inpaint_strength=0.7    ,				generator=__A    ,				num_inference_steps=2_5    ,				).latents

         snake_case				:    Tuple           =	pipe(
             prompt=__A    ,				mask_image=__A    ,				image_latents=__A    ,				generator=__A    ,				negative_prompt=__A    ,				inpaint_strength=0.7    ,				num_inference_steps=2_5    ,				output_type="""numpy"""    ,				).images[0]

         snake_case				:    Union[str, Any]           =	(
             np.array(
                 load_image(
                     """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"""
                     """/diffedit/pears.png"""     ).resize((7_6_8, 7_6_8)     )     )
             / 2_5_5
         )
         assert np.abs((expected_image - image).max()     ) < 5e-1





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 						__lowercase		(UpperCamelCase__							):




      """simple docstring"""





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

      def __init__(				self    ,				A=5_0_2_6_5    ,				A=1_0_2_4    ,				A=1_2    ,				A=4_0_9_6    ,				A=1_6    ,				A=1_2    ,				A=4_0_9_6    ,				A=1_6    ,				A=0.0    ,				A=0.0    ,				A="gelu"    ,				A=1_0_2_4    ,				A=0.1    ,				A=0.0    ,				A=0.0    ,				A=0.02    ,				A=0.0    ,				A=False    ,				A=True    ,				A=3    ,				A=1    ,				A=0    ,				A=2    ,				A=True    ,				A=2    ,				A=2    ,				**A    ,				)		->    Any:
         snake_case				:    Optional[int]           =	vocab_size
         snake_case				:    Union[str, Any]           =	max_position_embeddings
         snake_case				:    List[str]           =	d_model
         snake_case				:    List[Any]           =	encoder_ffn_dim
         snake_case				:    Optional[Any]           =	encoder_layers
         snake_case				:    Union[str, Any]           =	encoder_attention_heads
         snake_case				:    str           =	decoder_ffn_dim
         snake_case				:    Union[str, Any]           =	decoder_layers
         snake_case				:    Any           =	decoder_attention_heads
         snake_case				:    Union[str, Any]           =	dropout
         snake_case				:    List[str]           =	attention_dropout
         snake_case				:    List[Any]           =	activation_dropout
         snake_case				:    Optional[int]           =	activation_function
         snake_case				:    Union[str, Any]           =	init_std
         snake_case				:    List[str]           =	encoder_layerdrop
         snake_case				:    int           =	decoder_layerdrop
         snake_case				:    str           =	classifier_dropout
         snake_case				:    List[str]           =	use_cache
         snake_case				:    Tuple           =	encoder_layers
         snake_case				:    int           =	scale_embedding  # scale factor will be sqrt(d_model) if True

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

         # ensure backward compatibility for BART CNN models
         if self.forced_bos_token_id is None and kwargs.get("""force_bos_token_to_be_generated"""    ,				A     ):
            snake_case				:    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 						__lowercase		(UpperCamelCase__							):




      """simple docstring"""





      @property
      def      UpperCAmelCase				(				self     )		->    Mapping[str, Mapping[int, str]]:
         if self.task in ["default", "seq2seq-lm"]:
            snake_case				:    Optional[Any]           =	OrderedDict(
                [
                    ("""input_ids""", {0: """batch""", 1: """encoder_sequence"""}),
                    ("""attention_mask""", {0: """batch""", 1: """encoder_sequence"""}),
                ]     )

            if self.use_past:
               snake_case				:    Tuple           =	{0: """batch"""}
               snake_case				:    List[Any]           =	{0: """batch""", 1: """past_decoder_sequence + sequence"""}
            else:
               snake_case				:    Union[str, Any]           =	{0: """batch""", 1: """decoder_sequence"""}
               snake_case				:    Any           =	{0: """batch""", 1: """decoder_sequence"""}

            if self.use_past:
               self.fill_with_past_key_values_(A    ,				direction="""inputs"""     )
         elif self.task == "causal-lm":
            # TODO: figure this case out.
            snake_case				:    Union[str, Any]           =	OrderedDict(
                [
                    ("""input_ids""", {0: """batch""", 1: """encoder_sequence"""}),
                    ("""attention_mask""", {0: """batch""", 1: """encoder_sequence"""}),
                ]     )
            if self.use_past:
               snake_case     , snake_case				:    List[Any]           =	self.num_layers
               for i in range(A     ):
                  snake_case				:    List[Any]           =	{0: """batch""", 2: """past_sequence + sequence"""}
                  snake_case				:    Optional[int]           =	{0: """batch""", 2: """past_sequence + sequence"""}
         else:
            snake_case				:    Union[str, Any]           =	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      UpperCAmelCase				(				self     )		->    Mapping[str, Mapping[int, str]]:
         if self.task in ["default", "seq2seq-lm"]:
            snake_case				:    Any           =	super().outputs
         else:
            snake_case				:    Any           =	super(A    ,				self     ).outputs
            if self.use_past:
               snake_case     , snake_case				:    Any           =	self.num_layers
               for i in range(A     ):
                  snake_case				:    Any           =	{0: """batch""", 2: """past_sequence + sequence"""}
                  snake_case				:    Union[str, Any]           =	{0: """batch""", 2: """past_sequence + sequence"""}
         return common_outputs

      def      UpperCAmelCase				(				self    ,				A    ,				A = -1    ,				A = -1    ,				A = False    ,				A = None    ,				)		->    Mapping[str, Any]:
         snake_case				:    Optional[Any]           =	self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
             A    ,				A    ,				A    ,				A    ,				A     )

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

         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				:    Optional[int]           =	common_inputs["""input_ids"""].shape
            snake_case				:    Any           =	common_inputs["""decoder_input_ids"""].shape[1]
            snake_case     , snake_case				:    Optional[Any]           =	self.num_attention_heads
            snake_case				:    Optional[int]           =	(
                batch,
                num_encoder_attention_heads,
                encoder_seq_length,
                self._config.hidden_size // num_encoder_attention_heads,
            )
            snake_case				:    Any           =	decoder_seq_length + 3
            snake_case				:    List[Any]           =	(
                batch,
                num_decoder_attention_heads,
                decoder_past_length,
                self._config.hidden_size // num_decoder_attention_heads,
            )

            snake_case				:    str           =	torch.cat(
                [common_inputs["""decoder_attention_mask"""], torch.ones(A    ,				A     )]    ,				dim=1     )

            snake_case				:    str           =	[]
            # If the number of encoder and decoder layers are present in the model configuration, both are considered
            snake_case     , snake_case				:    Any           =	self.num_layers
            snake_case				:    List[str]           =	min(A    ,				A     )
            snake_case				:    Dict           =	max(A    ,				A     ) - min_num_layers
            snake_case				:    List[str]           =	"""encoder""" if num_encoder_layers > num_decoder_layers else """decoder"""

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

      def      UpperCAmelCase				(				self    ,				A    ,				A = -1    ,				A = -1    ,				A = False    ,				A = None    ,				)		->    Mapping[str, Any]:
         snake_case				:    str           =	self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
             A    ,				A    ,				A    ,				A    ,				A     )

         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
            # Not using the same length for past_key_values
            snake_case				:    Optional[int]           =	seqlen + 2
            snake_case     , snake_case				:    Tuple           =	self.num_layers
            snake_case     , snake_case				:    Optional[Any]           =	self.num_attention_heads
            snake_case				:    Union[str, Any]           =	(
                batch,
                num_encoder_attention_heads,
                past_key_values_length,
                self._config.hidden_size // num_encoder_attention_heads,
            )

            snake_case				:    Optional[Any]           =	common_inputs["""attention_mask"""].dtype
            snake_case				:    int           =	torch.cat(
                [common_inputs["""attention_mask"""], torch.ones(A    ,				A    ,				dtype=A     )]    ,				dim=1     )
            snake_case				:    Union[str, Any]           =	[
                (torch.zeros(A     ), torch.zeros(A     )) for _ in range(A     )
            ]
         return common_inputs

      def      UpperCAmelCase				(				self    ,				A    ,				A = -1    ,				A = -1    ,				A = False    ,				A = None    ,				)		->    Mapping[str, Any]:
         # Copied from OnnxConfig.generate_dummy_inputs
         # Did not use super(OnnxConfigWithPast, self).generate_dummy_inputs for code clarity.
         # If dynamic axis (-1) we forward with a fixed dimension of 2 samples to avoid optimizations made by ONNX
         snake_case				:    int           =	compute_effective_axis_dimension(
             A    ,				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				:    int           =	tokenizer.num_special_tokens_to_add(A     )
         snake_case				:    Tuple           =	compute_effective_axis_dimension(
             A    ,				fixed_dimension=OnnxConfig.default_fixed_sequence    ,				num_token_to_add=A     )

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

      def      UpperCAmelCase				(				self    ,				A    ,				A = -1    ,				A = -1    ,				A = False    ,				A = None    ,				)		->    Mapping[str, Any]:
         if self.task in ["default", "seq2seq-lm"]:
            snake_case				:    Optional[Any]           =	self._generate_dummy_inputs_for_default_and_seqaseq_lm(
                A    ,				batch_size=A    ,				seq_length=A    ,				is_pair=A    ,				framework=A     )

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

         return common_inputs

      def      UpperCAmelCase				(				self    ,				A    ,				A    ,				A    ,				A     )		->    Union[str, Any]:
         if self.task in ["default", "seq2seq-lm"]:
            snake_case				:    Optional[Any]           =	super()._flatten_past_key_values_(A    ,				A    ,				A    ,				A     )
         else:
            snake_case				:    Union[str, Any]           =	super(A    ,				self     )._flatten_past_key_values_(
                A    ,				A    ,				A    ,				A     )





