code: 
  
  
  
  
  """simple docstring"""
  
  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
  
  
  __A 		=			logging.get_logger(__name__)
  
  __A 		=			{
      "google/bigbird-pegasus-large-arxiv": (
          "https://huggingface.co/google/bigbird-pegasus-large-arxiv/resolve/main/config.json"
      ),
      "google/bigbird-pegasus-large-pubmed": (
          "https://huggingface.co/google/bigbird-pegasus-large-pubmed/resolve/main/config.json"
      ),
      "google/bigbird-pegasus-large-bigpatent": (
          "https://huggingface.co/google/bigbird-pegasus-large-bigpatent/resolve/main/config.json"
      ),
      # See all BigBirdPegasus models at https://huggingface.co/models?filter=bigbird_pegasus
  }
  class 			lowerCamelCase__    (  lowerCamelCase__			):
  			a__ :      Optional[int]								=     'bigbird_pegasus'
  			a__ :      str								=     ['past_key_values']
  			a__ :      Dict								=     {
  			    'num_attention_heads': 'encoder_attention_heads',
  			    'hidden_size': 'd_model',
  			    'attention_probs_dropout_prob': 'attention_dropout',
  			}
  
  
  			def __init__(     self				, SCREAMING_SNAKE_CASE=96_103				, SCREAMING_SNAKE_CASE=4_096				, SCREAMING_SNAKE_CASE=16				, SCREAMING_SNAKE_CASE=4_096				, SCREAMING_SNAKE_CASE=16				, SCREAMING_SNAKE_CASE=16				, SCREAMING_SNAKE_CASE=4_096				, SCREAMING_SNAKE_CASE=16				, SCREAMING_SNAKE_CASE=0.0				, SCREAMING_SNAKE_CASE=0.0				, SCREAMING_SNAKE_CASE=True				, SCREAMING_SNAKE_CASE=True				, SCREAMING_SNAKE_CASE="gelu_new"				, SCREAMING_SNAKE_CASE=1_024				, SCREAMING_SNAKE_CASE=0.1				, SCREAMING_SNAKE_CASE=0.0				, SCREAMING_SNAKE_CASE=0.0				, SCREAMING_SNAKE_CASE=0.02				, SCREAMING_SNAKE_CASE=2				, SCREAMING_SNAKE_CASE=0.0				, SCREAMING_SNAKE_CASE=True				, SCREAMING_SNAKE_CASE=0				, SCREAMING_SNAKE_CASE=2				, SCREAMING_SNAKE_CASE=1				, SCREAMING_SNAKE_CASE="block_sparse"				, SCREAMING_SNAKE_CASE=64				, SCREAMING_SNAKE_CASE=3				, SCREAMING_SNAKE_CASE=False				, **SCREAMING_SNAKE_CASE				, ):
  
  
  
  					"""simple docstring"""
  
  
  
  
  
  
  
  					snake_case		:						List[str]								=				vocab_size
  					snake_case		:						int								=				max_position_embeddings
  					snake_case		:						List[Any]								=				d_model
  					snake_case		:						Optional[int]								=				encoder_ffn_dim
  					snake_case		:						List[str]								=				encoder_layers
  					snake_case		:						Union[str, Any]								=				encoder_attention_heads
  					snake_case		:						Union[str, Any]								=				decoder_ffn_dim
  					snake_case		:						Any								=				decoder_layers
  					snake_case		:						Optional[int]								=				decoder_attention_heads
  					snake_case		:						str								=				dropout
  					snake_case		:						Optional[int]								=				attention_dropout
  					snake_case		:						Union[str, Any]								=				activation_dropout
  					snake_case		:						int								=				activation_function
  					snake_case		:						str								=				init_std
  					snake_case		:						List[Any]								=				encoder_layerdrop
  					snake_case		:						List[str]								=				decoder_layerdrop
  					snake_case		:						int								=				classifier_dropout
  					snake_case		:						Optional[int]								=				use_cache
  					snake_case		:						Optional[Any]								=				encoder_layers
  					snake_case		:						Any								=				scale_embedding  # scale factor will be sqrt(d_model) if True
  
  					# extra config
  					snake_case		:						List[str]								=				attention_type
  					snake_case		:						Optional[int]								=				block_size
  					snake_case		:						Union[str, Any]								=				num_random_blocks
  					snake_case		:						str								=				use_bias
  
  					super().__init__(
  					    pad_token_id=lowercase_				, bos_token_id=lowercase_				, eos_token_id=lowercase_				, is_encoder_decoder=lowercase_				, decoder_start_token_id=lowercase_				, **lowercase_				, )
  
  
  
  
  
  
  
  class 			lowerCamelCase__    (  lowerCamelCase__			):
  
  
  			@property
  			def 							lowerCamelCase_				(     self							):
  
  
  
  					"""simple docstring"""
  
  
  
  
  
  
  
  					if self.task in ["default", "seq2seq-lm"]:
  							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		:						str								=				{0: "batch"}
  									snake_case		:						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_(lowercase_				, 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(lowercase_							):
  											snake_case		:						Union[str, Any]								=				{0: "batch", 2: "past_sequence + sequence"}
  											snake_case		:						Optional[Any]								=				{0: "batch", 2: "past_sequence + sequence"}
  					else:
  							snake_case		:						Dict								=				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 							lowerCamelCase_				(     self							):
  
  
  
  					"""simple docstring"""
  
  
  
  
  
  
  
  					if self.task in ["default", "seq2seq-lm"]:
  							snake_case		:						Optional[int]								=				super().outputs
  					else:
  							snake_case		:						Optional[Any]								=				super(lowercase_				, self							).outputs
  							if self.use_past:
  									snake_case			,				snake_case		:						List[str]								=				self.num_layers
  									for i in range(lowercase_							):
  											snake_case		:						List[str]								=				{0: "batch", 2: "past_sequence + sequence"}
  											snake_case		:						Union[str, Any]								=				{0: "batch", 2: "past_sequence + sequence"}
  					return common_outputs
  
  
  			def 							lowerCamelCase_				(     self				, SCREAMING_SNAKE_CASE				, SCREAMING_SNAKE_CASE = -1				, SCREAMING_SNAKE_CASE = -1				, SCREAMING_SNAKE_CASE = False				, SCREAMING_SNAKE_CASE = None				, ):
  
  
  
  					"""simple docstring"""
  
  
  
  
  
  
  
  					snake_case		:						Dict								=				self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
  					    lowercase_				, lowercase_				, lowercase_				, lowercase_				, lowercase_							)
  
  					# Generate decoder inputs
  					snake_case		:						Tuple								=				seq_length if not self.use_past else 1
  					snake_case		:						Union[str, Any]								=				self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
  					    lowercase_				, lowercase_				, lowercase_				, lowercase_				, lowercase_							)
  					snake_case		:						Any								=				{F'''decoder_{name}''': tensor for name, tensor in decoder_inputs.items()}
  					snake_case		:						Any								=				dict(**lowercase_				, **lowercase_							)
  
  					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
  							snake_case		:						Any								=				common_inputs["decoder_input_ids"].shape[1]
  							snake_case			,				snake_case		:						Dict								=				self.num_attention_heads
  							snake_case		:						Dict								=				(
  							    batch,
  							    num_encoder_attention_heads,
  							    encoder_seq_length,
  							    self._config.hidden_size // num_encoder_attention_heads,
  							)
  							snake_case		:						str								=				decoder_seq_length + 3
  							snake_case		:						str								=				(
  							    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(lowercase_				, lowercase_							)]				, dim=1							)
  
  							snake_case		:						Optional[int]								=				[]
  							# If the number of encoder and decoder layers are present in the model configuration, both are considered
  							snake_case			,				snake_case		:						int								=				self.num_layers
  							snake_case		:						List[Any]								=				min(lowercase_				, lowercase_							)
  							snake_case		:						List[str]								=				max(lowercase_				, lowercase_							) - min_num_layers
  							snake_case		:						Dict								=				"encoder" if num_encoder_layers > num_decoder_layers else "decoder"
  
  							for _ in range(lowercase_							):
  									common_inputs["past_key_values"].append(
  									    (
  									        torch.zeros(lowercase_							),
  									        torch.zeros(lowercase_							),
  									        torch.zeros(lowercase_							),
  									        torch.zeros(lowercase_							),
  									    )							)
  							# TODO: test this.
  							snake_case		:						Optional[int]								=				encoder_shape if remaining_side_name == "encoder" else decoder_shape
  							for _ in range(lowercase_				, lowercase_							):
  									common_inputs["past_key_values"].append((torch.zeros(lowercase_							), torch.zeros(lowercase_							))							)
  					return common_inputs
  
  
  			def 							lowerCamelCase_				(     self				, SCREAMING_SNAKE_CASE				, SCREAMING_SNAKE_CASE = -1				, SCREAMING_SNAKE_CASE = -1				, SCREAMING_SNAKE_CASE = False				, SCREAMING_SNAKE_CASE = None				, ):
  
  
  
  					"""simple docstring"""
  
  
  
  
  
  
  
  					snake_case		:						List[Any]								=				self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
  					    lowercase_				, lowercase_				, lowercase_				, lowercase_				, lowercase_							)
  
  					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		:						Optional[Any]								=				seqlen + 2
  							snake_case			,				snake_case		:						Optional[Any]								=				self.num_layers
  							snake_case			,				snake_case		:						Dict								=				self.num_attention_heads
  							snake_case		:						List[Any]								=				(
  							    batch,
  							    num_encoder_attention_heads,
  							    past_key_values_length,
  							    self._config.hidden_size // num_encoder_attention_heads,
  							)
  
  							snake_case		:						Optional[int]								=				common_inputs["attention_mask"].dtype
  							snake_case		:						Union[str, Any]								=				torch.cat(
  							    [common_inputs["attention_mask"], torch.ones(lowercase_				, lowercase_				, dtype=lowercase_							)]				, dim=1							)
  							snake_case		:						Tuple								=				[
  							    (torch.zeros(lowercase_							), torch.zeros(lowercase_							)) for _ in range(lowercase_							)
  							]
  					return common_inputs
  
  
  			def 							lowerCamelCase_				(     self				, SCREAMING_SNAKE_CASE				, SCREAMING_SNAKE_CASE = -1				, SCREAMING_SNAKE_CASE = -1				, SCREAMING_SNAKE_CASE = False				, SCREAMING_SNAKE_CASE = None				, ):
  
  
  
  					"""simple docstring"""
  
  
  
  
  
  
  
  					snake_case		:						Union[str, Any]								=				compute_effective_axis_dimension(
  					    lowercase_				, 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		:						str								=				tokenizer.num_special_tokens_to_add(lowercase_							)
  					snake_case		:						Dict								=				compute_effective_axis_dimension(
  					    lowercase_				, fixed_dimension=OnnxConfig.default_fixed_sequence				, num_token_to_add=lowercase_							)
  
  					# Generate dummy inputs according to compute batch and sequence
  					snake_case		:						Tuple								=				[" ".join([tokenizer.unk_token]							) * seq_length] * batch_size
  					snake_case		:						Any								=				dict(tokenizer(lowercase_				, return_tensors=lowercase_							)							)
  					return common_inputs
  
  
  			def 							lowerCamelCase_				(     self				, SCREAMING_SNAKE_CASE				, SCREAMING_SNAKE_CASE = -1				, SCREAMING_SNAKE_CASE = -1				, SCREAMING_SNAKE_CASE = False				, SCREAMING_SNAKE_CASE = None				, ):
  
  
  
  					"""simple docstring"""
  
  
  
  
  
  
  
  					if self.task in ["default", "seq2seq-lm"]:
  							snake_case		:						int								=				self._generate_dummy_inputs_for_default_and_seqaseq_lm(
  							    lowercase_				, batch_size=lowercase_				, seq_length=lowercase_				, is_pair=lowercase_				, framework=lowercase_							)
  
  					elif self.task == "causal-lm":
  							snake_case		:						str								=				self._generate_dummy_inputs_for_causal_lm(
  							    lowercase_				, batch_size=lowercase_				, seq_length=lowercase_				, is_pair=lowercase_				, framework=lowercase_							)
  					else:
  							snake_case		:						Optional[Any]								=				self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
  							    lowercase_				, batch_size=lowercase_				, seq_length=lowercase_				, is_pair=lowercase_				, framework=lowercase_							)
  
  					return common_inputs
  
  
  
  			def 							lowerCamelCase_				(     self				, SCREAMING_SNAKE_CASE				, SCREAMING_SNAKE_CASE				, SCREAMING_SNAKE_CASE				, SCREAMING_SNAKE_CASE							):
  
  
  
  					"""simple docstring"""
  
  
  
  
  
  
  
  					if self.task in ["default", "seq2seq-lm"]:
  							snake_case		:						Optional[int]								=				super()._flatten_past_key_values_(lowercase_				, lowercase_				, lowercase_				, lowercase_							)
  					else:
  							snake_case		:						List[Any]								=				super(lowercase_				, self							)._flatten_past_key_values_(
  							    lowercase_				, lowercase_				, lowercase_				, lowercase_							)
  
style_context: 
  
  
  
  
  """simple docstring"""
  
  import gc
  import unittest
  
  import numpy as np
  import torch
  
  from diffusers import RePaintPipeline, RePaintScheduler, UNetaDModel
  from diffusers.utils.testing_utils import (
      enable_full_determinism,
      load_image,
      load_numpy,
      nightly,
      require_torch_gpu,
      skip_mps,
      torch_device,
  )
  
  from ..pipeline_params import IMAGE_INPAINTING_BATCH_PARAMS, IMAGE_INPAINTING_PARAMS
  from ..test_pipelines_common import PipelineTesterMixin
  
  
  enable_full_determinism()
  class 			lowerCamelCase__    (  lowerCamelCase_				,  unittest.TestCase			):
  			a__ :      Optional[Any]								=     RePaintPipeline
  			a__ :      List[Any]								=     IMAGE_INPAINTING_PARAMS - {"""width""", """height""", """guidance_scale"""}
  			a__ :      List[Any]								=     PipelineTesterMixin.required_optional_params - {
  			    """latents""",
  			    """num_images_per_prompt""",
  			    """callback""",
  			    """callback_steps""",
  			}
  			a__ :      List[Any]								=     IMAGE_INPAINTING_BATCH_PARAMS
  			a__ :      List[str]								=     False
  
  
  			def 							lowerCamelCase_				(     self							):
  
  
  
  					"""simple docstring"""
  
  
  
  
  
  
  
  					torch.manual_seed(0							)
  					torch.manual_seed(0							)
  					snake_case		:						Union[str, Any]								=				UNetaDModel(
  					    block_out_channels=(32, 64)				, layers_per_block=2				, sample_size=32				, in_channels=3				, out_channels=3				, down_block_types=("DownBlock2D", "AttnDownBlock2D")				, up_block_types=("AttnUpBlock2D", "UpBlock2D")				, )
  					snake_case		:						Union[str, Any]								=				RePaintScheduler()
  					snake_case		:						Optional[Any]								=				{"unet": unet, "scheduler": scheduler}
  					return components
  
  
  			def 							lowerCamelCase_				(     self				, SCREAMING_SNAKE_CASE				, SCREAMING_SNAKE_CASE=0							):
  
  
  
  					"""simple docstring"""
  
  
  
  
  
  
  
  					if str(SCREAMING_SNAKE_CASE							).startswith("mps"							):
  							snake_case		:						str								=				torch.manual_seed(SCREAMING_SNAKE_CASE							)
  					else:
  							snake_case		:						List[Any]								=				torch.Generator(device=SCREAMING_SNAKE_CASE							).manual_seed(SCREAMING_SNAKE_CASE							)
  					snake_case		:						Any								=				np.random.RandomState(SCREAMING_SNAKE_CASE							).standard_normal((1, 3, 32, 32)							)
  					snake_case		:						Any								=				torch.from_numpy(SCREAMING_SNAKE_CASE							).to(device=SCREAMING_SNAKE_CASE				, dtype=torch.floataa							)
  					snake_case		:						List[str]								=				(image > 0).to(device=SCREAMING_SNAKE_CASE				, dtype=torch.floataa							)
  					snake_case		:						int								=				{
  					    "image": image,
  					    "mask_image": mask,
  					    "generator": generator,
  					    "num_inference_steps": 5,
  					    "eta": 0.0,
  					    "jump_length": 2,
  					    "jump_n_sample": 2,
  					    "output_type": "numpy",
  					}
  					return inputs
  
  
  			def 							lowerCamelCase_				(     self							):
  
  
  
  					"""simple docstring"""
  
  
  
  
  
  
  
  					snake_case		:						List[Any]								=				"cpu"  # ensure determinism for the device-dependent torch.Generator
  					snake_case		:						Any								=				self.get_dummy_components()
  					snake_case		:						Optional[int]								=				RePaintPipeline(**SCREAMING_SNAKE_CASE							)
  					snake_case		:						Optional[Any]								=				sd_pipe.to(SCREAMING_SNAKE_CASE							)
  					sd_pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE							)
  
  					snake_case		:						Any								=				self.get_dummy_inputs(SCREAMING_SNAKE_CASE							)
  					snake_case		:						Dict								=				sd_pipe(**SCREAMING_SNAKE_CASE							).images
  					snake_case		:						Optional[int]								=				image[0, -3:, -3:, -1]
  
  					assert image.shape == (1, 32, 32, 3)
  					snake_case		:						List[str]								=				np.array([1.00_00, 0.54_26, 0.54_97, 0.22_00, 1.00_00, 1.00_00, 0.56_23, 1.00_00, 0.62_74]							)
  
  					assert np.abs(image_slice.flatten() - expected_slice							).max() < 1E-3
  
  
  			@skip_mps
  			def 							lowerCamelCase_				(     self							):
  
  
  
  					"""simple docstring"""
  
  
  
  
  
  
  
  					return super().test_save_load_local()
  
  
  			@unittest.skip("non-deterministic pipeline"							)
  			def 							lowerCamelCase_				(     self							):
  
  
  
  					"""simple docstring"""
  
  
  
  
  
  
  
  					return super().test_inference_batch_single_identical()
  
  
  			@skip_mps
  			def 							lowerCamelCase_				(     self							):
  
  
  
  					"""simple docstring"""
  
  
  
  
  
  
  
  					return super().test_dict_tuple_outputs_equivalent()
  
  
  			@skip_mps
  			def 							lowerCamelCase_				(     self							):
  
  
  
  					"""simple docstring"""
  
  
  
  
  
  
  
  					return super().test_save_load_optional_components()
  
  
  
  			@skip_mps
  			def 							lowerCamelCase_				(     self							):
  
  
  
  					"""simple docstring"""
  
  
  
  
  
  
  
  					return super().test_attention_slicing_forward_pass()
  
  
  
  
  
  
  
  @nightly
  @require_torch_gpu
  class 			lowerCamelCase__    (  unittest.TestCase			):
  
  
  			def 							lowerCamelCase_				(     self							):
  
  
  
  					"""simple docstring"""
  
  
  
  
  
  
  
  					super().tearDown()
  					gc.collect()
  					torch.cuda.empty_cache()
  
  
  
  			def 							lowerCamelCase_				(     self							):
  
  
  
  					"""simple docstring"""
  
  
  
  
  
  
  
  					snake_case		:						Optional[int]								=				load_image(
  					    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/"
  					    "repaint/celeba_hq_256.png"							)
  					snake_case		:						Any								=				load_image(
  					    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/repaint/mask_256.png"							)
  					snake_case		:						List[Any]								=				load_numpy(
  					    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/"
  					    "repaint/celeba_hq_256_result.npy"							)
  
  					snake_case		:						str								=				"google/ddpm-ema-celebahq-256"
  					snake_case		:						Optional[Any]								=				UNetaDModel.from_pretrained(SCREAMING_SNAKE_CASE							)
  					snake_case		:						Optional[int]								=				RePaintScheduler.from_pretrained(SCREAMING_SNAKE_CASE							)
  
  					snake_case		:						str								=				RePaintPipeline(unet=SCREAMING_SNAKE_CASE				, scheduler=SCREAMING_SNAKE_CASE							).to(SCREAMING_SNAKE_CASE							)
  					repaint.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE							)
  					repaint.enable_attention_slicing()
  
  					snake_case		:						Union[str, Any]								=				torch.manual_seed(0							)
  					snake_case		:						int								=				repaint(
  					    SCREAMING_SNAKE_CASE				, SCREAMING_SNAKE_CASE				, num_inference_steps=250				, eta=0.0				, jump_length=10				, jump_n_sample=10				, generator=SCREAMING_SNAKE_CASE				, output_type="np"				, )
  					snake_case		:						List[Any]								=				output.images[0]
  
  					assert image.shape == (256, 256, 3)
  					assert np.abs(expected_image - image							).mean() < 1E-2
  
  