"""simple docstring"""
import random
import unittest

import numpy as np
import torch

from diffusers import (
    DPMSolverMultistepScheduler,
    EulerAncestralDiscreteScheduler,
    EulerDiscreteScheduler,
    LMSDiscreteScheduler,
    OnnxStableDiffusionUpscalePipeline,
    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        __snake_case	(					__lowerCAmelCase	,				unittest.TestCase    ):
      # TODO: is there an appropriate internal test set?
      a__				=		"""ssube/stable-diffusion-x4-upscaler-onnx"""



      def   lowerCamelCase_       (						self   ,      lowercase=0)     ->      List[Any]:





        '''simple docstring'''





        a__:					List[Any]      							=     floats_tensor((1, 3, 1_28, 1_28)   ,      rng=random.Random(lowercase))
        a__:					int      							=     torch.manual_seed(lowercase)
        a__:					List[Any]      							=     {
            'prompt': 'A painting of a squirrel eating a burger',
            'image': image,
            'generator': generator,
            'num_inference_steps': 3,
            'guidance_scale': 7.5,
            'output_type': 'numpy',
        }
        return inputs



      def   lowerCamelCase_       (						self)     ->      int:





        '''simple docstring'''





        a__:					Optional[Any]      							=     OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint   ,      provider='CPUExecutionProvider')
        pipe.set_progress_bar_config(disable=lowercase)

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

        # started as 128, should now be 512
        assert image.shape == (1, 5_12, 5_12, 3)
        a__:					int      							=     np.array(
            [0.6974782, 0.68902093, 0.70135885, 0.7583618, 0.7804545, 0.7854912, 0.78667426, 0.78743863, 0.78070223])
        assert np.abs(image_slice - expected_slice).max() < 1e-1



      def   lowerCamelCase_       (						self)     ->      Any:





        '''simple docstring'''





        a__:					int      							=     OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint   ,      provider='CPUExecutionProvider')
        a__:					Optional[int]      							=     PNDMScheduler.from_config(pipe.scheduler.config   ,      skip_prk_steps=lowercase)
        pipe.set_progress_bar_config(disable=lowercase)

        a__:					List[str]      							=     self.get_dummy_inputs()
        a__:					List[Any]      							=     pipe(**lowercase).images
        a__:					int      							=     image[0, -3:, -3:, -1]

        assert image.shape == (1, 5_12, 5_12, 3)
        a__:					Optional[Any]      							=     np.array(
            [0.6898892, 0.59240556, 0.52499527, 0.58866215, 0.52258235, 0.52572715, 0.62414473, 0.6174387, 0.6214964])
        assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-1



      def   lowerCamelCase_       (						self)     ->      Optional[int]:





        '''simple docstring'''





        a__:					Optional[Any]      							=     OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint   ,      provider='CPUExecutionProvider')
        a__:					int      							=     DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
        pipe.set_progress_bar_config(disable=lowercase)

        a__:					List[str]      							=     self.get_dummy_inputs()
        a__:					Dict      							=     pipe(**lowercase).images
        a__:					List[str]      							=     image[0, -3:, -3:, -1]

        assert image.shape == (1, 5_12, 5_12, 3)
        a__:					List[Any]      							=     np.array(
            [0.7659278, 0.76437664, 0.75579107, 0.7691116, 0.77666986, 0.7727672, 0.7758664, 0.7812226, 0.76942515])

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



      def   lowerCamelCase_       (						self)     ->      int:





        '''simple docstring'''





        a__:					Dict      							=     OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint   ,      provider='CPUExecutionProvider')
        a__:					int      							=     EulerDiscreteScheduler.from_config(pipe.scheduler.config)
        pipe.set_progress_bar_config(disable=lowercase)

        a__:					Union[str, Any]      							=     self.get_dummy_inputs()
        a__:					Tuple      							=     pipe(**lowercase).images
        a__:					Tuple      							=     image[0, -3:, -3:, -1]

        assert image.shape == (1, 5_12, 5_12, 3)
        a__:					Tuple      							=     np.array(
            [0.6974782, 0.68902093, 0.70135885, 0.7583618, 0.7804545, 0.7854912, 0.78667426, 0.78743863, 0.78070223])
        assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-1



      def   lowerCamelCase_       (						self)     ->      Optional[Any]:





        '''simple docstring'''





        a__:					List[str]      							=     OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint   ,      provider='CPUExecutionProvider')
        a__:					Optional[Any]      							=     EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
        pipe.set_progress_bar_config(disable=lowercase)

        a__:					Optional[Any]      							=     self.get_dummy_inputs()
        a__:					List[Any]      							=     pipe(**lowercase).images
        a__:					List[Any]      							=     image[0, -3:, -3:, -1]

        assert image.shape == (1, 5_12, 5_12, 3)
        a__:					Tuple      							=     np.array(
            [0.77424496, 0.773601, 0.7645288, 0.7769598, 0.7772739, 0.7738688, 0.78187233, 0.77879584, 0.767043])

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







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



      @property
      def   lowerCamelCase_       (						self)     ->      int:





        '''simple docstring'''





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



      @property
      def   lowerCamelCase_       (						self)     ->      Tuple:





        '''simple docstring'''





        a__:					Any      							=     ort.SessionOptions()
        a__:					List[str]      							=     False
        return options



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





        '''simple docstring'''





        a__:					List[str]      							=     load_image(
            'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
            '/img2img/sketch-mountains-input.jpg')
        a__:					Tuple      							=     init_image.resize((1_28, 1_28))
        # using the PNDM scheduler by default
        a__:					Dict      							=     OnnxStableDiffusionUpscalePipeline.from_pretrained(
            'ssube/stable-diffusion-x4-upscaler-onnx'   ,      provider=self.gpu_provider   ,      sess_options=self.gpu_options   ,      )
        pipe.set_progress_bar_config(disable=lowercase)

        a__:					int      							=     'A fantasy landscape, trending on artstation'

        a__:					int      							=     torch.manual_seed(0)
        a__:					Union[str, Any]      							=     pipe(
            prompt=lowercase   ,      image=lowercase   ,      guidance_scale=7.5   ,      num_inference_steps=10   ,      generator=lowercase   ,      output_type='np'   ,      )
        a__:					List[str]      							=     output.images
        a__:					Tuple      							=     images[0, 2_55:2_58, 3_83:3_86, -1]

        assert images.shape == (1, 5_12, 5_12, 3)
        a__:					Any      							=     np.array([0.4883, 0.4947, 0.4980, 0.4975, 0.4982, 0.4980, 0.5000, 0.5006, 0.4972])
        # TODO: lower the tolerance after finding the cause of onnxruntime reproducibility issues

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



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





        '''simple docstring'''





        a__:					List[str]      							=     load_image(
            'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
            '/img2img/sketch-mountains-input.jpg')
        a__:					Union[str, Any]      							=     init_image.resize((1_28, 1_28))
        a__:					Any      							=     LMSDiscreteScheduler.from_pretrained(
            'ssube/stable-diffusion-x4-upscaler-onnx'   ,      subfolder='scheduler')
        a__:					int      							=     OnnxStableDiffusionUpscalePipeline.from_pretrained(
            'ssube/stable-diffusion-x4-upscaler-onnx'   ,      scheduler=lowercase   ,      provider=self.gpu_provider   ,      sess_options=self.gpu_options   ,      )
        pipe.set_progress_bar_config(disable=lowercase)

        a__:					Optional[Any]      							=     'A fantasy landscape, trending on artstation'

        a__:					Dict      							=     torch.manual_seed(0)
        a__:					Union[str, Any]      							=     pipe(
            prompt=lowercase   ,      image=lowercase   ,      guidance_scale=7.5   ,      num_inference_steps=20   ,      generator=lowercase   ,      output_type='np'   ,      )
        a__:					List[str]      							=     output.images
        a__:					int      							=     images[0, 2_55:2_58, 3_83:3_86, -1]

        assert images.shape == (1, 5_12, 5_12, 3)
        a__:					int      							=     np.array(
            [0.50173753, 0.50223356, 0.502039, 0.50233036, 0.5023725, 0.5022601, 0.5018758, 0.50234085, 0.50241566])
        # TODO: lower the tolerance after finding the cause of onnxruntime reproducibility issues

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



"""simple docstring"""
import unittest

from transformers import AutoConfig, AutoTokenizer, BertConfig, TensorType, is_flax_available
from transformers.testing_utils import DUMMY_UNKNOWN_IDENTIFIER, require_flax, slow


if is_flax_available():
 import jax

 from transformers.models.auto.modeling_flax_auto import FlaxAutoModel
 from transformers.models.bert.modeling_flax_bert import FlaxBertModel
 from transformers.models.roberta.modeling_flax_roberta import FlaxRobertaModel




@require_flax
class        __snake_case	(					unittest.TestCase    ):



      @slow
      def   lowerCamelCase_       (						self)     ->      Tuple:





        '''simple docstring'''





        for model_name in ["bert-base-cased", "bert-large-uncased"]:
          with self.subTest(lowercase):
            a__:					Optional[Any]      							=     AutoConfig.from_pretrained(lowercase)
            self.assertIsNotNone(lowercase)
            self.assertIsInstance(lowercase   ,      lowercase)

            a__:					Any      							=     FlaxAutoModel.from_pretrained(lowercase)
            self.assertIsNotNone(lowercase)
            self.assertIsInstance(lowercase   ,      lowercase)



      @slow
      def   lowerCamelCase_       (						self)     ->      Tuple:





        '''simple docstring'''





        for model_name in ["roberta-base", "roberta-large"]:
          with self.subTest(lowercase):
            a__:					str      							=     AutoConfig.from_pretrained(lowercase)
            self.assertIsNotNone(lowercase)
            self.assertIsInstance(lowercase   ,      lowercase)

            a__:					Any      							=     FlaxAutoModel.from_pretrained(lowercase)
            self.assertIsNotNone(lowercase)
            self.assertIsInstance(lowercase   ,      lowercase)



      @slow
      def   lowerCamelCase_       (						self)     ->      str:





        '''simple docstring'''





        for model_name in ["bert-base-cased", "bert-large-uncased"]:
          a__:					Dict      							=     AutoTokenizer.from_pretrained(lowercase)
          a__:					Union[str, Any]      							=     FlaxBertModel.from_pretrained(lowercase)
          a__:					List[str]      							=     tokenizer('Do you support jax jitted function?'   ,      return_tensors=TensorType.JAX)

          @jax.jit
          def eval(**lowercase):
            return model(**lowercase)

          eval(**lowercase).block_until_ready()



      @slow
      def   lowerCamelCase_       (						self)     ->      str:





        '''simple docstring'''





        for model_name in ["roberta-base", "roberta-large"]:
          a__:					Optional[Any]      							=     AutoTokenizer.from_pretrained(lowercase)
          a__:					Any      							=     FlaxRobertaModel.from_pretrained(lowercase)
          a__:					int      							=     tokenizer('Do you support jax jitted function?'   ,      return_tensors=TensorType.JAX)

          @jax.jit
          def eval(**lowercase):
            return model(**lowercase)

          eval(**lowercase).block_until_ready()



      def   lowerCamelCase_       (						self)     ->      Any:





        '''simple docstring'''





        with self.assertRaisesRegex(
            lowercase   ,      'bert-base is not a local folder and is not a valid model identifier'):
          a__:					str      							=     FlaxAutoModel.from_pretrained('bert-base')



      def   lowerCamelCase_       (						self)     ->      Optional[int]:





        '''simple docstring'''





        with self.assertRaisesRegex(
            lowercase   ,      r'aaaaaa is not a valid git identifier \(branch name, tag name or commit id\)'):
          a__:					List[str]      							=     FlaxAutoModel.from_pretrained(lowercase   ,      revision='aaaaaa')



      def   lowerCamelCase_       (						self)     ->      Tuple:





        '''simple docstring'''





        with self.assertRaisesRegex(
            lowercase   ,      'hf-internal-testing/config-no-model does not appear to have a file named flax_model.msgpack'   ,      ):
          a__:					List[str]      							=     FlaxAutoModel.from_pretrained('hf-internal-testing/config-no-model')



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





        '''simple docstring'''





        with self.assertRaisesRegex(lowercase   ,      'Use `from_pt=True` to load this model'):
          a__:					List[str]      							=     FlaxAutoModel.from_pretrained('hf-internal-testing/tiny-bert-pt-only')


