



'''simple docstring'''





import unittest

from transformers import BigBirdConfig, is_flax_available
from transformers.testing_utils import require_flax, slow

from ...test_modeling_flax_common import FlaxModelTesterMixin, ids_tensor, random_attention_mask


if is_flax_available():
  import jax

  from transformers.models.big_bird.modeling_flax_big_bird import (
      FlaxBigBirdForCausalLM,
      FlaxBigBirdForMaskedLM,
      FlaxBigBirdForMultipleChoice,
      FlaxBigBirdForPreTraining,
      FlaxBigBirdForQuestionAnswering,
      FlaxBigBirdForSequenceClassification,
      FlaxBigBirdForTokenClassification,
      FlaxBigBirdModel,
  )






class __A    (						unittest.TestCase		):



 '''simple docstring'''







 def __init__(self				,     A				,     A=2				,     A=56				,     A=True				,     A=True				,     A=True				,     A=True				,     A=99				,     A=32				,     A=2				,     A=2				,     A=7				,     A="gelu_new"				,     A=0.1				,     A=0.1				,     A=512				,     A=16				,     A=2				,     A=0.02				,     A=4				,     A="block_sparse"				,     A=True				,     A=False				,     A=2				,     A=3				,     )					->  List[Any]:


      """simple docstring"""



      _a											=						parent
      _a											=						batch_size
      _a											=						seq_length
      _a											=						is_training
      _a											=						use_attention_mask
      _a											=						use_token_type_ids
      _a											=						use_labels
      _a											=						vocab_size
      _a											=						hidden_size
      _a											=						num_hidden_layers
      _a											=						num_attention_heads
      _a											=						intermediate_size
      _a											=						hidden_act
      _a											=						hidden_dropout_prob
      _a											=						attention_probs_dropout_prob
      _a											=						max_position_embeddings
      _a											=						type_vocab_size
      _a											=						type_sequence_label_size
      _a											=						initializer_range
      _a											=						num_choices

      _a											=						rescale_embeddings
      _a											=						attention_type
      _a											=						use_bias
      _a											=						block_size
      _a											=						num_random_blocks



 def       a__						(self )					->  List[str]:


      """simple docstring"""



      _a											=						ids_tensor([self.batch_size, self.seq_length]				,     self.vocab_size )

      _a											=						None
      if self.use_attention_mask:
           _a											=						random_attention_mask([self.batch_size, self.seq_length] )

      _a											=						None
      if self.use_token_type_ids:
           _a											=						ids_tensor([self.batch_size, self.seq_length]				,     self.type_vocab_size )

      _a											=						BigBirdConfig(
          vocab_size=self.vocab_size				,     hidden_size=self.hidden_size				,     num_hidden_layers=self.num_hidden_layers				,     num_attention_heads=self.num_attention_heads				,     intermediate_size=self.intermediate_size				,     hidden_act=self.hidden_act				,     hidden_dropout_prob=self.hidden_dropout_prob				,     attention_probs_dropout_prob=self.attention_probs_dropout_prob				,     max_position_embeddings=self.max_position_embeddings				,     type_vocab_size=self.type_vocab_size				,     is_decoder=A				,     initializer_range=self.initializer_range				,     attention_type=self.attention_type				,     block_size=self.block_size				,     num_random_blocks=self.num_random_blocks				,     use_bias=self.use_bias				,     rescale_embeddings=self.rescale_embeddings				,     )

      return config, input_ids, token_type_ids, attention_mask



 def       a__						(self )					->  Tuple:


      """simple docstring"""



      _a											=						self.prepare_config_and_inputs()
      _a						,					_a						,					_a						,					_a											=						config_and_inputs
      _a											=						{
          '''input_ids''': input_ids,
          '''token_type_ids''': token_type_ids,
          '''attention_mask''': attention_mask,
      }
      return config, inputs_dict






@require_flax
class __A    (						A     ,     unittest.TestCase		):



 '''simple docstring'''







 __lowerCamelCase					: str    = (
     (
         FlaxBigBirdForCausalLM,
         FlaxBigBirdModel,
         FlaxBigBirdForPreTraining,
         FlaxBigBirdForMaskedLM,
         FlaxBigBirdForMultipleChoice,
         FlaxBigBirdForQuestionAnswering,
         FlaxBigBirdForSequenceClassification,
         FlaxBigBirdForTokenClassification,
     )
     if is_flax_available()
     else ()
 )

 __lowerCamelCase					: Optional[Any]    = False
 __lowerCamelCase					: Tuple    = False



 def       a__						(self )					->  Union[str, Any]:


      """simple docstring"""



      _a											=						FlaxBigBirdModelTester(self )



 @slow
 # copied from `test_modeling_flax_common` because it takes much longer than other models
 def       a__						(self )					->  Tuple:


      """simple docstring"""



      super().test_from_pretrained_save_pretrained()



 @slow
 # copied from `test_modeling_flax_common` because it takes much longer than other models
 def       a__						(self )					->  Optional[int]:


      """simple docstring"""



      super().test_from_pretrained_with_no_automatic_init()



 @slow
 # copied from `test_modeling_flax_common` because it takes much longer than other models
 def       a__						(self )					->  Tuple:


      """simple docstring"""



      super().test_no_automatic_init()



 @slow
 # copied from `test_modeling_flax_common` because it takes much longer than other models
 def       a__						(self )					->  Dict:


      """simple docstring"""



      super().test_hidden_states_output()



 @slow
 def       a__						(self )					->  int:


      """simple docstring"""



      for model_class_name in self.all_model_classes:
           _a											=						model_class_name.from_pretrained('''google/bigbird-roberta-base''' )
           self.assertIsNotNone(A )



 def       a__						(self )					->  int:


      """simple docstring"""



      if self.test_attn_probs:
           super().test_attention_outputs()



 @slow
 # copied from `test_modeling_flax_common` because it takes much longer than other models
 def       a__						(self )					->  Optional[int]:


      """simple docstring"""



      _a						,					_a											=						self.model_tester.prepare_config_and_inputs_for_common()

      for model_class in self.all_model_classes:
           with self.subTest(model_class.__name__ ):
                _a											=						self._prepare_for_class(A				,     A )
                _a											=						model_class(A )

                @jax.jit
                def model_jitted(A				,     A=None				,     **A ):
                     return model(input_ids=A				,     attention_mask=A				,     **A )

                with self.subTest('''JIT Enabled''' ):
                     _a											=						model_jitted(**A ).to_tuple()

                with self.subTest('''JIT Disabled''' ):
                     with jax.disable_jit():
                          _a											=						model_jitted(**A ).to_tuple()

                self.assertEqual(len(A )				,     len(A ) )
                for jitted_output, output in zip(A				,     A ):
                     self.assertEqual(jitted_output.shape				,     output.shape )



 def       a__						(self				,     A				,     A				,     A				,     A=1E-5				,     A="outputs"				,     A=None )					->  Any:


      """simple docstring"""



      if name.startswith('''outputs.attentions''' ):
           return
      else:
           super().check_pt_flax_outputs(A				,     A				,     A				,     A				,     A				,     A )








'''simple docstring'''





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 __A    (						A     ,     unittest.TestCase		):



 '''simple docstring'''







 __lowerCamelCase					: Dict    = 'hf-internal-testing/tiny-random-OnnxStableDiffusionPipeline'



 def       a__						(self				,     A=0 )					->  str:


      """simple docstring"""



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



 def       a__						(self )					->  str:


      """simple docstring"""



      _a											=						OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint				,     provider='''CPUExecutionProvider''' )
      pipe.set_progress_bar_config(disable=A )

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

      assert image.shape == (1, 128, 128, 3)
      _a											=						np.array([0.69643, 0.58484, 0.50314, 0.58760, 0.55368, 0.59643, 0.51529, 0.41217, 0.49087] )
      assert np.abs(image_slice - expected_slice ).max() < 1E-1



 def       a__						(self )					->  Tuple:


      """simple docstring"""



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

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

      assert image.shape == (1, 128, 128, 3)
      _a											=						np.array([0.61737, 0.54642, 0.53183, 0.54465, 0.52742, 0.60525, 0.49969, 0.40655, 0.48154] )

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



 def       a__						(self )					->  int:


      """simple docstring"""



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

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

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

      assert image.shape == (1, 128, 128, 3)
      _a											=						np.array([0.52761, 0.59977, 0.49033, 0.49619, 0.54282, 0.50311, 0.47600, 0.40918, 0.45203] )

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



 def       a__						(self )					->  Dict:


      """simple docstring"""



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

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

      assert image.shape == (1, 128, 128, 3)
      _a											=						np.array([0.52911, 0.60004, 0.49229, 0.49805, 0.54502, 0.50680, 0.47777, 0.41028, 0.45304] )

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



 def       a__						(self )					->  Tuple:


      """simple docstring"""



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

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

      assert image.shape == (1, 128, 128, 3)
      _a											=						np.array([0.52911, 0.60004, 0.49229, 0.49805, 0.54502, 0.50680, 0.47777, 0.41028, 0.45304] )

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



 def       a__						(self )					->  List[str]:


      """simple docstring"""



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

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

      assert image.shape == (1, 128, 128, 3)
      _a											=						np.array([0.65331, 0.58277, 0.48204, 0.56059, 0.53665, 0.56235, 0.50969, 0.40009, 0.46552] )

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






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



 '''simple docstring'''







 @property
 def       a__						(self )					->  Tuple:


      """simple docstring"""



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



 @property
 def       a__						(self )					->  Any:


      """simple docstring"""



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



 def       a__						(self )					->  str:


      """simple docstring"""



      _a											=						load_image(
          '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
          '''/img2img/sketch-mountains-input.jpg''' )
      _a											=						init_image.resize((768, 512) )
      # using the PNDM scheduler by default
      _a											=						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 )

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

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

      assert images.shape == (1, 512, 768, 3)
      _a											=						np.array([0.4909, 0.5059, 0.5372, 0.4623, 0.4876, 0.5049, 0.4820, 0.4956, 0.5019] )
      # 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 )					->  Optional[int]:


      """simple docstring"""



      _a											=						load_image(
          '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
          '''/img2img/sketch-mountains-input.jpg''' )
      _a											=						init_image.resize((768, 512) )
      _a											=						LMSDiscreteScheduler.from_pretrained(
          '''runwayml/stable-diffusion-v1-5'''				,     subfolder='''scheduler'''				,     revision='''onnx''' )
      _a											=						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 )

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

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

      assert images.shape == (1, 512, 768, 3)
      _a											=						np.array([0.8043, 0.926, 0.9581, 0.8119, 0.8954, 0.913, 0.7209, 0.7463, 0.7431] )
      # TODO: lower the tolerance after finding the cause of onnxruntime reproducibility issues

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



