import gc
import tempfile
import unittest

import numpy as np
import torch

from diffusers import VersatileDiffusionDualGuidedPipeline
from diffusers.utils.testing_utils import load_image, nightly, require_torch_gpu, torch_device


_SCREAMING_SNAKE_CASE   		=			False
@nightly
@require_torch_gpu
class SCREAMING_SNAKE_CASE_		(			unittest.TestCase		):







   def 	lowerCamelCase_  (    self	:	str							):





    """simple docstring"""


    super().tearDown()
    gc.collect()
    torch.cuda.empty_cache()







   def 	lowerCamelCase_  (    self	:	Optional[int]							):





    """simple docstring"""


    UpperCamelCase      =			VersatileDiffusionDualGuidedPipeline.from_pretrained("""shi-labs/versatile-diffusion"""							)
    # remove text_unet
    pipe.remove_unused_weights()
    pipe.to(lowerCamelCase_							)
    pipe.set_progress_bar_config(disable=lowerCamelCase_							)

    UpperCamelCase      =			load_image(
        """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/versatile_diffusion/benz.jpg"""							)

    UpperCamelCase      =			torch.manual_seed(0							)
    UpperCamelCase      =			pipe(
        prompt="""first prompt"""      ,      image=lowerCamelCase_      ,      text_to_image_strength=0.7_5      ,      generator=lowerCamelCase_      ,      guidance_scale=7.5      ,      num_inference_steps=2      ,      output_type="""numpy"""      ,      ).images

    with tempfile.TemporaryDirectory() as tmpdirname:
     pipe.save_pretrained(lowerCamelCase_							)
     UpperCamelCase      =			VersatileDiffusionDualGuidedPipeline.from_pretrained(lowerCamelCase_							)

    pipe.to(lowerCamelCase_							)
    pipe.set_progress_bar_config(disable=lowerCamelCase_							)

    UpperCamelCase      =			generator.manual_seed(0							)
    UpperCamelCase      =			pipe(
        prompt="""first prompt"""      ,      image=lowerCamelCase_      ,      text_to_image_strength=0.7_5      ,      generator=lowerCamelCase_      ,      guidance_scale=7.5      ,      num_inference_steps=2      ,      output_type="""numpy"""      ,      ).images

    assert np.abs(image - new_image							).sum() < 1E-5, "Models don't have the same forward pass"







   def 	lowerCamelCase_  (    self	:	int							):





    """simple docstring"""


    UpperCamelCase      =			VersatileDiffusionDualGuidedPipeline.from_pretrained("""shi-labs/versatile-diffusion"""							)
    pipe.remove_unused_weights()
    pipe.to(lowerCamelCase_							)
    pipe.set_progress_bar_config(disable=lowerCamelCase_							)

    UpperCamelCase      =			"""cyberpunk 2077"""
    UpperCamelCase      =			load_image(
        """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/versatile_diffusion/benz.jpg"""							)
    UpperCamelCase      =			torch.manual_seed(0							)
    UpperCamelCase      =			pipe(
        prompt=lowerCamelCase_      ,      image=lowerCamelCase_      ,      text_to_image_strength=0.7_5      ,      generator=lowerCamelCase_      ,      guidance_scale=7.5      ,      num_inference_steps=50      ,      output_type="""numpy"""      ,      ).images

    UpperCamelCase      =			image[0, 253:256, 253:256, -1]

    assert image.shape == (1, 512, 512, 3)
    UpperCamelCase      =			np.array([0.0_7_8_7, 0.0_8_4_9, 0.0_8_2_6, 0.0_8_1_2, 0.0_8_0_7, 0.0_7_9_5, 0.0_8_1_8, 0.0_7_9_8, 0.0_7_7_9]							)

    assert np.abs(image_slice.flatten() - expected_slice							).max() < 1E-2

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


_SCREAMING_SNAKE_CASE   		=			logging.get_logger(__name__)

_SCREAMING_SNAKE_CASE   		=			{
    """facebook/mbart-large-cc25""": """https://huggingface.co/facebook/mbart-large-cc25/resolve/main/config.json""",
    # See all MBART models at https://huggingface.co/models?filter=mbart
}
class SCREAMING_SNAKE_CASE_		(			__lowerCAmelCase		):
   __lowerCAmelCase											=							"""mbart"""
   __lowerCAmelCase											=							["""past_key_values"""]
   __lowerCAmelCase											=							{"""num_attention_heads""": """encoder_attention_heads""", """hidden_size""": """d_model"""}







   def __init__(    self	:	Optional[int]      ,      lowerCamelCase_	:	Dict=5_0265      ,      lowerCamelCase_	:	Tuple=1024      ,      lowerCamelCase_	:	int=12      ,      lowerCamelCase_	:	Tuple=4096      ,      lowerCamelCase_	:	int=16      ,      lowerCamelCase_	:	Dict=12      ,      lowerCamelCase_	:	Dict=4096      ,      lowerCamelCase_	:	Optional[int]=16      ,      lowerCamelCase_	:	int=0.0      ,      lowerCamelCase_	:	Tuple=0.0      ,      lowerCamelCase_	:	int=True      ,      lowerCamelCase_	:	List[str]=True      ,      lowerCamelCase_	:	List[str]="gelu"      ,      lowerCamelCase_	:	str=1024      ,      lowerCamelCase_	:	int=0.1      ,      lowerCamelCase_	:	List[Any]=0.0      ,      lowerCamelCase_	:	Dict=0.0      ,      lowerCamelCase_	:	Union[str, Any]=0.0_2      ,      lowerCamelCase_	:	Tuple=0.0      ,      lowerCamelCase_	:	int=False      ,      lowerCamelCase_	:	List[Any]=1      ,      lowerCamelCase_	:	Union[str, Any]=0      ,      lowerCamelCase_	:	Optional[Any]=2      ,      lowerCamelCase_	:	Any=2      ,      **lowerCamelCase_	:	int      ,      ):





    """simple docstring"""


    UpperCamelCase      =			vocab_size
    UpperCamelCase      =			max_position_embeddings
    UpperCamelCase      =			d_model
    UpperCamelCase      =			encoder_ffn_dim
    UpperCamelCase      =			encoder_layers
    UpperCamelCase      =			encoder_attention_heads
    UpperCamelCase      =			decoder_ffn_dim
    UpperCamelCase      =			decoder_layers
    UpperCamelCase      =			decoder_attention_heads
    UpperCamelCase      =			dropout
    UpperCamelCase      =			attention_dropout
    UpperCamelCase      =			activation_dropout
    UpperCamelCase      =			activation_function
    UpperCamelCase      =			init_std
    UpperCamelCase      =			encoder_layerdrop
    UpperCamelCase      =			decoder_layerdrop
    UpperCamelCase      =			classifier_dropout
    UpperCamelCase      =			use_cache
    UpperCamelCase      =			encoder_layers
    UpperCamelCase      =			scale_embedding  # scale factor will be sqrt(d_model) if True
    super().__init__(
        pad_token_id=lowerCamelCase_      ,      bos_token_id=lowerCamelCase_      ,      eos_token_id=lowerCamelCase_      ,      is_encoder_decoder=lowerCamelCase_      ,      forced_eos_token_id=lowerCamelCase_      ,      **lowerCamelCase_      ,      )





class SCREAMING_SNAKE_CASE_		(			__lowerCAmelCase		):







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





    """simple docstring"""


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

     if self.use_past:
      UpperCamelCase      =			{0: """batch"""}
      UpperCamelCase      =			{0: """batch""", 1: """past_decoder_sequence + sequence"""}
     else:
      UpperCamelCase      =			{0: """batch""", 1: """decoder_sequence"""}
      UpperCamelCase      =			{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.
     UpperCamelCase      =			OrderedDict(
         [
             ("""input_ids""", {0: """batch""", 1: """encoder_sequence"""}),
             ("""attention_mask""", {0: """batch""", 1: """encoder_sequence"""}),
         ]							)
     if self.use_past:
      UpperCamelCase     ,    UpperCamelCase      =			self.num_layers
      for i in range(lowerCamelCase_							):
       UpperCamelCase      =			{0: """batch""", 2: """past_sequence + sequence"""}
       UpperCamelCase      =			{0: """batch""", 2: """past_sequence + sequence"""}
    else:
     UpperCamelCase      =			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	:	str							):





    """simple docstring"""


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







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





    """simple docstring"""


    UpperCamelCase      =			self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
        lowerCamelCase_      ,      lowerCamelCase_      ,      lowerCamelCase_      ,      lowerCamelCase_      ,      lowerCamelCase_							)

    # Generate decoder inputs
    UpperCamelCase      =			seq_length if not self.use_past else 1
    UpperCamelCase      =			self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
        lowerCamelCase_      ,      lowerCamelCase_      ,      lowerCamelCase_      ,      lowerCamelCase_      ,      lowerCamelCase_							)
    UpperCamelCase      =			{f"""decoder_{name}""": tensor for name, tensor in decoder_inputs.items()}
    UpperCamelCase      =			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
     UpperCamelCase     ,    UpperCamelCase      =			common_inputs["""input_ids"""].shape
     UpperCamelCase      =			common_inputs["""decoder_input_ids"""].shape[1]
     UpperCamelCase     ,    UpperCamelCase      =			self.num_attention_heads
     UpperCamelCase      =			(
         batch,
         num_encoder_attention_heads,
         encoder_seq_length,
         self._config.hidden_size // num_encoder_attention_heads,
     )
     UpperCamelCase      =			decoder_seq_length + 3
     UpperCamelCase      =			(
         batch,
         num_decoder_attention_heads,
         decoder_past_length,
         self._config.hidden_size // num_decoder_attention_heads,
     )

     UpperCamelCase      =			torch.cat(
         [common_inputs["""decoder_attention_mask"""], torch.ones(lowerCamelCase_      ,      lowerCamelCase_							)]      ,      dim=1							)

     UpperCamelCase      =			[]
     # If the number of encoder and decoder layers are present in the model configuration, both are considered
     UpperCamelCase     ,    UpperCamelCase      =			self.num_layers
     UpperCamelCase      =			min(lowerCamelCase_      ,      lowerCamelCase_							)
     UpperCamelCase      =			max(lowerCamelCase_      ,      lowerCamelCase_							) - min_num_layers
     UpperCamelCase      =			"""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.
     UpperCamelCase      =			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 	lowerCamelCase_  (    self	:	Dict      ,      lowerCamelCase_	:	PreTrainedTokenizer      ,      lowerCamelCase_	:	int = -1      ,      lowerCamelCase_	:	int = -1      ,      lowerCamelCase_	:	bool = False      ,      lowerCamelCase_	:	Optional[TensorType] = None      ,      ):





    """simple docstring"""


    UpperCamelCase      =			self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
        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
     UpperCamelCase     ,    UpperCamelCase      =			common_inputs["""input_ids"""].shape
     # Not using the same length for past_key_values
     UpperCamelCase      =			seqlen + 2
     UpperCamelCase     ,    UpperCamelCase      =			self.num_layers
     UpperCamelCase     ,    UpperCamelCase      =			self.num_attention_heads
     UpperCamelCase      =			(
         batch,
         num_encoder_attention_heads,
         past_key_values_length,
         self._config.hidden_size // num_encoder_attention_heads,
     )

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







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





    """simple docstring"""


    UpperCamelCase      =			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
    UpperCamelCase      =			tokenizer.num_special_tokens_to_add(lowerCamelCase_							)
    UpperCamelCase      =			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
    UpperCamelCase      =			[""" """.join([tokenizer.unk_token]							) * seq_length] * batch_size
    UpperCamelCase      =			dict(tokenizer(lowerCamelCase_      ,      return_tensors=lowerCamelCase_							)							)
    return common_inputs







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





    """simple docstring"""


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

    elif self.task == "causal-lm":
     UpperCamelCase      =			self._generate_dummy_inputs_for_causal_lm(
         lowerCamelCase_      ,      batch_size=lowerCamelCase_      ,      seq_length=lowerCamelCase_      ,      is_pair=lowerCamelCase_      ,      framework=lowerCamelCase_							)
    else:
     UpperCamelCase      =			self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
         lowerCamelCase_      ,      batch_size=lowerCamelCase_      ,      seq_length=lowerCamelCase_      ,      is_pair=lowerCamelCase_      ,      framework=lowerCamelCase_							)

    return common_inputs







   def 	lowerCamelCase_  (    self	:	Optional[Any]      ,      lowerCamelCase_	:	int      ,      lowerCamelCase_	:	Union[str, Any]      ,      lowerCamelCase_	:	Optional[Any]      ,      lowerCamelCase_	:	List[str]							):





    """simple docstring"""


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