



'''simple docstring'''





from __future__ import annotations

import unittest

from transformers import AutoTokenizer, MBartConfig, is_tf_available
from transformers.testing_utils import require_sentencepiece, require_tf, require_tokenizers, slow
from transformers.utils import cached_property

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
       import tensorflow as tf

       from transformers import TFAutoModelForSeqaSeqLM, TFMBartForConditionalGeneration, TFMBartModel







@require_tf
class     UpperCAmelCase_       :
    """simple docstring"""







    lowercase					=		MBartConfig
    lowercase					=		{}
    lowercase					=		"gelu"




    def __init__(    self					:     Dict ,     snake_case_					:     str ,     snake_case_					:     Tuple=13 ,     snake_case_					:     Optional[Any]=7 ,     snake_case_					:     List[Any]=True ,     snake_case_					:     List[str]=False ,     snake_case_					:     Any=99 ,     snake_case_					:     str=32 ,     snake_case_					:     int=2 ,     snake_case_					:     Optional[Any]=4 ,     snake_case_					:     Tuple=37 ,     snake_case_					:     Tuple=0.1 ,     snake_case_					:     Any=0.1 ,     snake_case_					:     Dict=20 ,     snake_case_					:     Optional[int]=2 ,     snake_case_					:     List[Any]=1 ,     snake_case_					:     Dict=0 ,     ):
          snake_case__  :    int    						=      parent
          snake_case__  :    Optional[Any]    						=      batch_size
          snake_case__  :    Optional[Any]    						=      seq_length
          snake_case__  :    Union[str, Any]    						=      is_training
          snake_case__  :    Optional[int]    						=      use_labels
          snake_case__  :    str    						=      vocab_size
          snake_case__  :    Optional[int]    						=      hidden_size
          snake_case__  :    str    						=      num_hidden_layers
          snake_case__  :    Optional[int]    						=      num_attention_heads
          snake_case__  :    int    						=      intermediate_size
          snake_case__  :    List[Any]    						=      hidden_dropout_prob
          snake_case__  :    Union[str, Any]    						=      attention_probs_dropout_prob
          snake_case__  :    Tuple    						=      max_position_embeddings
          snake_case__  :    List[Any]    						=      eos_token_id
          snake_case__  :    Dict    						=      pad_token_id
          snake_case__  :    Tuple    						=      bos_token_id




    def lowerCamelCase       (    self					:     Dict						):
          snake_case__  :    Union[str, Any]    						=      ids_tensor([self.batch_size, self.seq_length - 1] ,     self.vocab_size						)
          snake_case__  :    int    						=      tf.expand_dims(tf.constant([self.eos_token_id] * self.batch_size						) ,     1						)
          snake_case__  :    Any    						=      tf.concat([input_ids, eos_tensor] ,     axis=1						)

          snake_case__  :    Optional[int]    						=      ids_tensor([self.batch_size, self.seq_length] ,     self.vocab_size						)

          snake_case__  :    Union[str, Any]    						=      self.config_cls(
              vocab_size=self.vocab_size ,     d_model=self.hidden_size ,     encoder_layers=self.num_hidden_layers ,     decoder_layers=self.num_hidden_layers ,     encoder_attention_heads=self.num_attention_heads ,     decoder_attention_heads=self.num_attention_heads ,     encoder_ffn_dim=self.intermediate_size ,     decoder_ffn_dim=self.intermediate_size ,     dropout=self.hidden_dropout_prob ,     attention_dropout=self.attention_probs_dropout_prob ,     max_position_embeddings=self.max_position_embeddings ,     eos_token_ids=[2] ,     bos_token_id=self.bos_token_id ,     pad_token_id=self.pad_token_id ,     decoder_start_token_id=self.pad_token_id ,     **self.config_updates ,     )
          snake_case__  :    List[Any]    						=      prepare_mbart_inputs_dict(snake_case_ ,     snake_case_ ,     snake_case_						)
          return config, inputs_dict




    def lowerCamelCase       (    self					:     int ,     snake_case_					:     Dict ,     snake_case_					:     str						):
          snake_case__  :    Union[str, Any]    						=      TFMBartModel(config=snake_case_						).get_decoder()
          snake_case__  :    Any    						=      inputs_dict["""input_ids"""]

          snake_case__  :    Union[str, Any]    						=      input_ids[:1, :]
          snake_case__  :    Any    						=      inputs_dict["""attention_mask"""][:1, :]
          snake_case__  :    Optional[Any]    						=      inputs_dict["""head_mask"""]
          snake_case__  :    Optional[int]    						=      1

          # first forward pass
          snake_case__  :    Dict    						=      model(snake_case_ ,     attention_mask=snake_case_ ,     head_mask=snake_case_ ,     use_cache=snake_case_						)

          snake_case__						,				snake_case__  :    List[Any]    						=      outputs.to_tuple()
          snake_case__  :    Tuple    						=      past_key_values[1]





def 						__snake_case(				_lowerCAmelCase     ,   _lowerCAmelCase     ,   _lowerCAmelCase     ,   _lowerCAmelCase=None     ,   _lowerCAmelCase=None     ,   _lowerCAmelCase=None     ,   _lowerCAmelCase=None     ,   _lowerCAmelCase=None     ,   )		-> str:
      if attention_mask is None:
            snake_case__  :    Union[str, Any]    						=      tf.cast(tf.math.not_equal(_lowerCAmelCase     ,   config.pad_token_id		)     ,   tf.inta		)
      if decoder_attention_mask is None:
            snake_case__  :    Dict    						=      tf.concat(
                [
                    tf.ones(decoder_input_ids[:, :1].shape     ,   dtype=tf.inta		),
                    tf.cast(tf.math.not_equal(decoder_input_ids[:, 1:]     ,   config.pad_token_id		)     ,   tf.inta		),
                ]     ,   axis=-1     ,   )
      if head_mask is None:
            snake_case__  :    Dict    						=      tf.ones((config.encoder_layers, config.encoder_attention_heads)		)
      if decoder_head_mask is None:
            snake_case__  :    int    						=      tf.ones((config.decoder_layers, config.decoder_attention_heads)		)
      if cross_attn_head_mask is None:
            snake_case__  :    Union[str, Any]    						=      tf.ones((config.decoder_layers, config.decoder_attention_heads)		)
      return {
          "input_ids": input_ids,
          "decoder_input_ids": decoder_input_ids,
          "attention_mask": attention_mask,
          "decoder_attention_mask": decoder_attention_mask,
          "head_mask": head_mask,
          "decoder_head_mask": decoder_head_mask,
          "cross_attn_head_mask": cross_attn_head_mask,
      }







@require_tf
class     UpperCAmelCase_       (       _a		,		_a		,		unittest.TestCase	):
    """simple docstring"""







    lowercase					=		(TFMBartForConditionalGeneration, TFMBartModel) if is_tf_available() else ()
    lowercase					=		(TFMBartForConditionalGeneration,) if is_tf_available() else ()
    lowercase					=		(
        {
            "conversational": TFMBartForConditionalGeneration,
            "feature-extraction": TFMBartModel,
            "summarization": TFMBartForConditionalGeneration,
            "text2text-generation": TFMBartForConditionalGeneration,
            "translation": TFMBartForConditionalGeneration,
        }
        if is_tf_available()
        else {}
    )
    lowercase					=		True
    lowercase					=		False
    lowercase					=		False




    def lowerCamelCase       (    self					:     str ,     snake_case_					:     List[str] ,     snake_case_					:     Optional[Any] ,     snake_case_					:     Tuple ,     snake_case_					:     Any ,     snake_case_					:     List[str]						):
          if pipeline_test_casse_name != "FeatureExtractionPipelineTests":
                # Exception encountered when calling layer '...'
                return True

          return False




    def lowerCamelCase       (    self					:     str						):
          snake_case__  :    int    						=      TFMBartModelTester(self						)
          snake_case__  :    Optional[int]    						=      ConfigTester(self ,     config_class=snake_case_						)




    def lowerCamelCase       (    self					:     Optional[int]						):
          self.config_tester.run_common_tests()




    def lowerCamelCase       (    self					:     List[str]						):
          snake_case__  :    Optional[Any]    						=      self.model_tester.prepare_config_and_inputs_for_common()
          self.model_tester.check_decoder_model_past_large_inputs(*snake_case_						)







@require_sentencepiece
@require_tokenizers
@require_tf
class     UpperCAmelCase_       (       unittest.TestCase	):
    """simple docstring"""







    lowercase					=		[
        " UN Chief Says There Is No Military Solution in Syria",
    ]
    lowercase					=		[
        "Şeful ONU declară că nu există o soluţie militară în Siria",
    ]
    lowercase					=		"facebook/mbart-large-en-ro"




    @cached_property
    def lowerCamelCase       (    self					:     Optional[Any]						):
          return AutoTokenizer.from_pretrained(self.model_name						)




    @cached_property
    def lowerCamelCase       (    self					:     List[Any]						):
          snake_case__  :    Any    						=      TFAutoModelForSeqaSeqLM.from_pretrained(self.model_name						)
          return model




    def lowerCamelCase       (    self					:     Union[str, Any] ,     **snake_case_					:     int						):
          snake_case__  :    Tuple    						=      self.translate_src_text(**snake_case_						)
          self.assertListEqual(self.expected_text ,     snake_case_						)




    def lowerCamelCase       (    self					:     List[str] ,     **snake_case_					:     Any						):
          snake_case__  :    Tuple    						=      self.tokenizer(self.src_text ,     **snake_case_ ,     return_tensors="""tf"""						)
          snake_case__  :    List[str]    						=      self.model.generate(
              model_inputs.input_ids ,     attention_mask=model_inputs.attention_mask ,     num_beams=2						)
          snake_case__  :    str    						=      self.tokenizer.batch_decode(snake_case_ ,     skip_special_tokens=snake_case_						)
          return generated_words




    @slow
    def lowerCamelCase       (    self					:     int						):
          self._assert_generated_batch_equal_expected()







import random
import unittest

import numpy as np

from diffusers import (
    DPMSolverMultistepScheduler,
    EulerAncestralDiscreteScheduler,
    EulerDiscreteScheduler,
    LMSDiscreteScheduler,
    OnnxStableDiffusionImgaImgPipeline,
    PNDMScheduler,
)
from diffusers.utils import floats_tensor
from diffusers.utils.testing_utils import (
    is_onnx_available,
    load_image,
    nightly,
    require_onnxruntime,
    require_torch_gpu,
)

from ..test_pipelines_onnx_common import OnnxPipelineTesterMixin


if is_onnx_available():
    import onnxruntime as ort

class 		SCREAMING_SNAKE_CASE__				(       _UpperCAmelCase						,      unittest.TestCase  ):
      A_						:				Tuple							=    'hf-internal-testing/tiny-random-OnnxStableDiffusionPipeline'


      def    a (self       :      int				,			a__       :      List[Any]=0      ):





         """simple docstring"""



         __snake_case							       =   floats_tensor((1, 3, 128, 128)				,			rng=random.Random(a__      )      )
         __snake_case							       =   np.random.RandomState(a__      )
         __snake_case							       =   {
             '''prompt''': '''A painting of a squirrel eating a burger''',
             '''image''': image,
             '''generator''': generator,
             '''num_inference_steps''': 3,
             '''strength''': 0.7_5,
             '''guidance_scale''': 7.5,
             '''output_type''': '''numpy''',
         }
         return inputs


      def    a (self       :      List[Any]      ):





         """simple docstring"""



         __snake_case							       =   OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint				,			provider='''CPUExecutionProvider'''      )
         pipe.set_progress_bar_config(disable=a__      )

         __snake_case							       =   self.get_dummy_inputs()
         __snake_case							       =   pipe(**a__      ).images
         __snake_case							       =   image[0, -3:, -3:, -1].flatten()

         assert image.shape == (1, 128, 128, 3)
         __snake_case							       =   np.array([0.6_9_6_4_3, 0.5_8_4_8_4, 0.5_0_3_1_4, 0.5_8_7_6_0, 0.5_5_3_6_8, 0.5_9_6_4_3, 0.5_1_5_2_9, 0.4_1_2_1_7, 0.4_9_0_8_7]      )
         assert np.abs(image_slice - expected_slice      ).max() < 1E-1


      def    a (self       :      Dict      ):





         """simple docstring"""



         __snake_case							       =   OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint				,			provider='''CPUExecutionProvider'''      )
         __snake_case							       =   PNDMScheduler.from_config(pipe.scheduler.config				,			skip_prk_steps=a__      )
         pipe.set_progress_bar_config(disable=a__      )

         __snake_case							       =   self.get_dummy_inputs()
         __snake_case							       =   pipe(**a__      ).images
         __snake_case							       =   image[0, -3:, -3:, -1]

         assert image.shape == (1, 128, 128, 3)
         __snake_case							       =   np.array([0.6_1_7_3_7, 0.5_4_6_4_2, 0.5_3_1_8_3, 0.5_4_4_6_5, 0.5_2_7_4_2, 0.6_0_5_2_5, 0.4_9_9_6_9, 0.4_0_6_5_5, 0.4_8_1_5_4]      )

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


      def    a (self       :      List[str]      ):





         """simple docstring"""



         __snake_case							       =   OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint				,			provider='''CPUExecutionProvider'''      )
         __snake_case							       =   LMSDiscreteScheduler.from_config(pipe.scheduler.config      )
         pipe.set_progress_bar_config(disable=a__      )

         # warmup pass to apply optimizations
         __snake_case							       =   pipe(**self.get_dummy_inputs()      )

         __snake_case							       =   self.get_dummy_inputs()
         __snake_case							       =   pipe(**a__      ).images
         __snake_case							       =   image[0, -3:, -3:, -1]

         assert image.shape == (1, 128, 128, 3)
         __snake_case							       =   np.array([0.5_2_7_6_1, 0.5_9_9_7_7, 0.4_9_0_3_3, 0.4_9_6_1_9, 0.5_4_2_8_2, 0.5_0_3_1_1, 0.4_7_6_0_0, 0.4_0_9_1_8, 0.4_5_2_0_3]      )

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


      def    a (self       :      Any      ):





         """simple docstring"""



         __snake_case							       =   OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint				,			provider='''CPUExecutionProvider'''      )
         __snake_case							       =   EulerDiscreteScheduler.from_config(pipe.scheduler.config      )
         pipe.set_progress_bar_config(disable=a__      )

         __snake_case							       =   self.get_dummy_inputs()
         __snake_case							       =   pipe(**a__      ).images
         __snake_case							       =   image[0, -3:, -3:, -1]

         assert image.shape == (1, 128, 128, 3)
         __snake_case							       =   np.array([0.5_2_9_1_1, 0.6_0_0_0_4, 0.4_9_2_2_9, 0.4_9_8_0_5, 0.5_4_5_0_2, 0.5_0_6_8_0, 0.4_7_7_7_7, 0.4_1_0_2_8, 0.4_5_3_0_4]      )

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


      def    a (self       :      Dict      ):





         """simple docstring"""



         __snake_case							       =   OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint				,			provider='''CPUExecutionProvider'''      )
         __snake_case							       =   EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config      )
         pipe.set_progress_bar_config(disable=a__      )

         __snake_case							       =   self.get_dummy_inputs()
         __snake_case							       =   pipe(**a__      ).images
         __snake_case							       =   image[0, -3:, -3:, -1]

         assert image.shape == (1, 128, 128, 3)
         __snake_case							       =   np.array([0.5_2_9_1_1, 0.6_0_0_0_4, 0.4_9_2_2_9, 0.4_9_8_0_5, 0.5_4_5_0_2, 0.5_0_6_8_0, 0.4_7_7_7_7, 0.4_1_0_2_8, 0.4_5_3_0_4]      )

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


      def    a (self       :      List[str]      ):





         """simple docstring"""



         __snake_case							       =   OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint				,			provider='''CPUExecutionProvider'''      )
         __snake_case							       =   DPMSolverMultistepScheduler.from_config(pipe.scheduler.config      )
         pipe.set_progress_bar_config(disable=a__      )

         __snake_case							       =   self.get_dummy_inputs()
         __snake_case							       =   pipe(**a__      ).images
         __snake_case							       =   image[0, -3:, -3:, -1]

         assert image.shape == (1, 128, 128, 3)
         __snake_case							       =   np.array([0.6_5_3_3_1, 0.5_8_2_7_7, 0.4_8_2_0_4, 0.5_6_0_5_9, 0.5_3_6_6_5, 0.5_6_2_3_5, 0.5_0_9_6_9, 0.4_0_0_0_9, 0.4_6_5_5_2]      )

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






@nightly
@require_onnxruntime
@require_torch_gpu
class 		SCREAMING_SNAKE_CASE__				(       unittest.TestCase  ):


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





         """simple docstring"""



         return (
             "CUDAExecutionProvider",
             {
                 "gpu_mem_limit": "15000000000",  # 15GB
                 "arena_extend_strategy": "kSameAsRequested",
             },
         )


      @property
      def    a (self       :      Optional[Any]      ):





         """simple docstring"""



         __snake_case							       =   ort.SessionOptions()
         __snake_case							       =   False
         return options


      def    a (self       :      Optional[Any]      ):





         """simple docstring"""



         __snake_case							       =   load_image(
             '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
             '''/img2img/sketch-mountains-input.jpg'''      )
         __snake_case							       =   init_image.resize((768, 512)      )
         # using the PNDM scheduler by default
         __snake_case							       =   OnnxStableDiffusionImgaImgPipeline.from_pretrained(
             '''CompVis/stable-diffusion-v1-4'''				,			revision='''onnx'''				,			safety_checker=a__				,			feature_extractor=a__				,			provider=self.gpu_provider				,			sess_options=self.gpu_options				,			)
         pipe.set_progress_bar_config(disable=a__      )

         __snake_case							       =   '''A fantasy landscape, trending on artstation'''

         __snake_case							       =   np.random.RandomState(0      )
         __snake_case							       =   pipe(
             prompt=a__				,			image=a__				,			strength=0.7_5				,			guidance_scale=7.5				,			num_inference_steps=10				,			generator=a__				,			output_type='''np'''				,			)
         __snake_case							       =   output.images
         __snake_case							       =   images[0, 255:258, 383:386, -1]

         assert images.shape == (1, 512, 768, 3)
         __snake_case							       =   np.array([0.4_9_0_9, 0.5_0_5_9, 0.5_3_7_2, 0.4_6_2_3, 0.4_8_7_6, 0.5_0_4_9, 0.4_8_2_0, 0.4_9_5_6, 0.5_0_1_9]      )
         # TODO: lower the tolerance after finding the cause of onnxruntime reproducibility issues

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


      def    a (self       :      Dict      ):





         """simple docstring"""



         __snake_case							       =   load_image(
             '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
             '''/img2img/sketch-mountains-input.jpg'''      )
         __snake_case							       =   init_image.resize((768, 512)      )
         __snake_case							       =   LMSDiscreteScheduler.from_pretrained(
             '''runwayml/stable-diffusion-v1-5'''				,			subfolder='''scheduler'''				,			revision='''onnx'''      )
         __snake_case							       =   OnnxStableDiffusionImgaImgPipeline.from_pretrained(
             '''runwayml/stable-diffusion-v1-5'''				,			revision='''onnx'''				,			scheduler=a__				,			safety_checker=a__				,			feature_extractor=a__				,			provider=self.gpu_provider				,			sess_options=self.gpu_options				,			)
         pipe.set_progress_bar_config(disable=a__      )

         __snake_case							       =   '''A fantasy landscape, trending on artstation'''

         __snake_case							       =   np.random.RandomState(0      )
         __snake_case							       =   pipe(
             prompt=a__				,			image=a__				,			strength=0.7_5				,			guidance_scale=7.5				,			num_inference_steps=20				,			generator=a__				,			output_type='''np'''				,			)
         __snake_case							       =   output.images
         __snake_case							       =   images[0, 255:258, 383:386, -1]

         assert images.shape == (1, 512, 768, 3)
         __snake_case							       =   np.array([0.8_0_4_3, 0.9_2_6, 0.9_5_8_1, 0.8_1_1_9, 0.8_9_5_4, 0.9_1_3, 0.7_2_0_9, 0.7_4_6_3, 0.7_4_3_1]      )
         # TODO: lower the tolerance after finding the cause of onnxruntime reproducibility issues

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

