


import inspect
import unittest

import numpy as np

from tests.test_modeling_common import floats_tensor
from transformers import MaskaFormerConfig, is_torch_available, is_vision_available
from transformers.testing_utils import require_torch, require_torch_multi_gpu, require_vision, slow, torch_device
from transformers.utils import cached_property

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
   import torch

   from transformers import MaskaFormerForUniversalSegmentation, MaskaFormerModel

   if is_vision_available():
      from transformers import MaskaFormerImageProcessor

if is_vision_available():
   from PIL import Image
class  SCREAMING_SNAKE_CASE__  :
      def __init__(       self				:       Union[str, Any]     ,							SCREAMING_SNAKE_CASE__				:       List[Any]     ,							SCREAMING_SNAKE_CASE__				:       Dict=2     ,							SCREAMING_SNAKE_CASE__				:       Dict=True     ,							SCREAMING_SNAKE_CASE__				:       Optional[Any]=False     ,							SCREAMING_SNAKE_CASE__				:       int=1_0     ,							SCREAMING_SNAKE_CASE__				:       List[Any]=3     ,							SCREAMING_SNAKE_CASE__				:       Optional[Any]=3_2 * 8     ,							SCREAMING_SNAKE_CASE__				:       Dict=3_2 * 8     ,							SCREAMING_SNAKE_CASE__				:       List[str]=4     ,							SCREAMING_SNAKE_CASE__				:       Optional[Any]=6_4     ,							)		->			str:
             a_					:		Optional[Any]							=  parent
             a_					:		List[Any]							=  batch_size
             a_					:		List[Any]							=  is_training
             a_					:		str							=  use_auxiliary_loss
             a_					:		str							=  num_queries
             a_					:		str							=  num_channels
             a_					:		Union[str, Any]							=  min_size
             a_					:		Union[str, Any]							=  max_size
             a_					:		int							=  num_labels
             a_					:		List[Any]							=  hidden_dim
             a_					:		int							=  hidden_dim
      def        SCREAMING_SNAKE_CASE							(       self				:       Optional[Any]     )		->			int:
             a_					:		Any							=  floats_tensor([self.batch_size, self.num_channels, self.min_size, self.max_size]     ).to(
                 SCREAMING_SNAKE_CASE__     )

             a_					:		int							=  torch.ones([self.batch_size, self.min_size, self.max_size]     ,							device=SCREAMING_SNAKE_CASE__     )

             a_					:		Any							=  (
                 torch.rand([self.batch_size, self.num_labels, self.min_size, self.max_size]     ,							device=SCREAMING_SNAKE_CASE__     ) > 0.5
             ).float()
             a_					:		int							=  (torch.rand((self.batch_size, self.num_labels)     ,							device=SCREAMING_SNAKE_CASE__     ) > 0.5).long()

             a_					:		Any							=  self.get_config()
             return config, pixel_values, pixel_mask, mask_labels, class_labels
      def        SCREAMING_SNAKE_CASE							(       self				:       Any     )		->			Union[str, Any]:
             a_					:		Any							=  MaskaFormerConfig(
                 hidden_size=self.hidden_dim     ,							)
             a_					:		Tuple							=  self.num_queries
             a_					:		Tuple							=  self.num_labels

             a_					:		Optional[int]							=  [1, 1, 1, 1]
             a_					:		Union[str, Any]							=  self.num_channels

             a_					:		str							=  6_4
             a_					:		List[Any]							=  1_2_8
             a_					:		Tuple							=  self.hidden_dim
             a_					:		Dict							=  self.hidden_dim
             a_					:		Optional[Any]							=  self.hidden_dim
             return config
      def        SCREAMING_SNAKE_CASE							(       self				:       Any     )		->			str:
             a_							,		a_							,		a_							,		a_							,		a_					:		List[str]							=  self.prepare_config_and_inputs()
             a_					:		Dict							=  {'pixel_values': pixel_values, 'pixel_mask': pixel_mask}
             return config, inputs_dict
      def        SCREAMING_SNAKE_CASE							(       self				:       Optional[Any]     ,							SCREAMING_SNAKE_CASE__				:       Optional[int]     ,							SCREAMING_SNAKE_CASE__				:       int     )		->			Optional[Any]:
             a_					:		List[str]							=  output.encoder_hidden_states
             a_					:		Any							=  output.pixel_decoder_hidden_states
             a_					:		List[Any]							=  output.transformer_decoder_hidden_states

             self.parent.assertTrue(len(SCREAMING_SNAKE_CASE__     )     ,							len(config.backbone_config.depths     )     )
             self.parent.assertTrue(len(SCREAMING_SNAKE_CASE__     )     ,							len(config.backbone_config.depths     )     )
             self.parent.assertTrue(len(SCREAMING_SNAKE_CASE__     )     ,							config.decoder_layers     )
      def        SCREAMING_SNAKE_CASE							(       self				:       Optional[Any]     ,							SCREAMING_SNAKE_CASE__				:       Dict     ,							SCREAMING_SNAKE_CASE__				:       List[Any]     ,							SCREAMING_SNAKE_CASE__				:       int     ,							SCREAMING_SNAKE_CASE__				:       List[str]=False     )		->			Optional[int]:
             with torch.no_grad():
                    a_					:		int							=  MaskaFormerModel(config=SCREAMING_SNAKE_CASE__     )
                    model.to(SCREAMING_SNAKE_CASE__     )
                    model.eval()

                    a_					:		Dict							=  model(pixel_values=SCREAMING_SNAKE_CASE__     ,							pixel_mask=SCREAMING_SNAKE_CASE__     )
                    a_					:		Union[str, Any]							=  model(SCREAMING_SNAKE_CASE__     ,							output_hidden_states=SCREAMING_SNAKE_CASE__     )

             self.parent.assertEqual(
                 output.transformer_decoder_last_hidden_state.shape     ,							(self.batch_size, self.num_queries, self.hidden_dim)     ,							)
             # let's ensure the other two hidden state exists
             self.parent.assertTrue(output.pixel_decoder_last_hidden_state is not None     )
             self.parent.assertTrue(output.encoder_last_hidden_state is not None     )

             if output_hidden_states:
                    self.check_output_hidden_state(SCREAMING_SNAKE_CASE__     ,							SCREAMING_SNAKE_CASE__     )

      def        SCREAMING_SNAKE_CASE							(       self				:       Dict     ,							SCREAMING_SNAKE_CASE__				:       str     ,							SCREAMING_SNAKE_CASE__				:       int     ,							SCREAMING_SNAKE_CASE__				:       Optional[Any]     ,							SCREAMING_SNAKE_CASE__				:       Optional[Any]     ,							SCREAMING_SNAKE_CASE__				:       Dict     )		->			str:
             a_					:		Dict							=  MaskaFormerForUniversalSegmentation(config=SCREAMING_SNAKE_CASE__     )
             model.to(SCREAMING_SNAKE_CASE__     )
             model.eval()

             def comm_check_on_output(SCREAMING_SNAKE_CASE__				:       List[Any]     ):
                    # let's still check that all the required stuff is there
                    self.parent.assertTrue(result.transformer_decoder_last_hidden_state is not None     )
                    self.parent.assertTrue(result.pixel_decoder_last_hidden_state is not None     )
                    self.parent.assertTrue(result.encoder_last_hidden_state is not None     )
                    # okay, now we need to check the logits shape
                    # due to the encoder compression, masks have a //4 spatial size
                    self.parent.assertEqual(
                        result.masks_queries_logits.shape     ,							(self.batch_size, self.num_queries, self.min_size // 4, self.max_size // 4)     ,							)
                    # + 1 for null class
                    self.parent.assertEqual(
                        result.class_queries_logits.shape     ,							(self.batch_size, self.num_queries, self.num_labels + 1)     )

             with torch.no_grad():
                    a_					:		str							=  model(pixel_values=SCREAMING_SNAKE_CASE__     ,							pixel_mask=SCREAMING_SNAKE_CASE__     )
                    a_					:		str							=  model(SCREAMING_SNAKE_CASE__     )

                    comm_check_on_output(SCREAMING_SNAKE_CASE__     )

                    a_					:		Tuple							=  model(
                        pixel_values=SCREAMING_SNAKE_CASE__     ,							pixel_mask=SCREAMING_SNAKE_CASE__     ,							mask_labels=SCREAMING_SNAKE_CASE__     ,							class_labels=SCREAMING_SNAKE_CASE__     )

             comm_check_on_output(SCREAMING_SNAKE_CASE__     )

             self.parent.assertTrue(result.loss is not None     )
             self.parent.assertEqual(result.loss.shape     ,							torch.Size([1]     )     )
@require_torch
class  SCREAMING_SNAKE_CASE__  (				lowercase__						,    lowercase__						,    unittest.TestCase			):
      snake_case__				:     Optional[int]        =						(MaskaFormerModel, MaskaFormerForUniversalSegmentation) if is_torch_available() else ()
      snake_case__				:     List[Any]        =						{'''feature-extraction''': MaskaFormerModel} if is_torch_available() else {}

      snake_case__				:     Dict        =						False
      snake_case__				:     Optional[int]        =						False
      snake_case__				:     Tuple        =						False
      snake_case__				:     Union[str, Any]        =						False
      def        SCREAMING_SNAKE_CASE							(       self				:       str     )		->			str:
             a_					:		Optional[Any]							=  MaskaFormerModelTester(self     )
             a_					:		Union[str, Any]							=  ConfigTester(self     ,							config_class=SCREAMING_SNAKE_CASE__     ,							has_text_modality=SCREAMING_SNAKE_CASE__     )
      def        SCREAMING_SNAKE_CASE							(       self				:       Dict     )		->			List[str]:
             self.config_tester.run_common_tests()
      def        SCREAMING_SNAKE_CASE							(       self				:       Dict     )		->			List[Any]:
             a_							,		a_					:		List[Any]							=  self.model_tester.prepare_config_and_inputs_for_common()
             self.model_tester.create_and_check_maskaformer_model(SCREAMING_SNAKE_CASE__     ,							**SCREAMING_SNAKE_CASE__     ,							output_hidden_states=SCREAMING_SNAKE_CASE__     )
      def        SCREAMING_SNAKE_CASE							(       self				:       Dict     )		->			Tuple:
             a_					:		List[Any]							=  self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_maskaformer_instance_segmentation_head_model(*SCREAMING_SNAKE_CASE__     )
      @unittest.skip(reason='Mask2Former does not use inputs_embeds'     )
      def        SCREAMING_SNAKE_CASE							(       self				:       Any     )		->			Tuple:
             pass
      @unittest.skip(reason='Mask2Former does not have a get_input_embeddings method'     )
      def        SCREAMING_SNAKE_CASE							(       self				:       Tuple     )		->			Dict:
             pass
      @unittest.skip(reason='Mask2Former is not a generative model'     )
      def        SCREAMING_SNAKE_CASE							(       self				:       Any     )		->			int:
             pass
      @unittest.skip(reason='Mask2Former does not use token embeddings'     )
      def        SCREAMING_SNAKE_CASE							(       self				:       Union[str, Any]     )		->			Dict:
             pass
      @require_torch_multi_gpu
      @unittest.skip(
          reason='Mask2Former has some layers using `add_module` which doesn\'t work well with `nn.DataParallel`'     )
      def        SCREAMING_SNAKE_CASE							(       self				:       int     )		->			Tuple:
             pass
      @unittest.skip('Will be fixed soon by reducing the size of the model used for common tests.'     )
      def        SCREAMING_SNAKE_CASE							(       self				:       Union[str, Any]     )		->			List[str]:
             pass
      def        SCREAMING_SNAKE_CASE							(       self				:       List[str]     )		->			List[str]:
             a_							,		a_					:		Tuple							=  self.model_tester.prepare_config_and_inputs_for_common()

             for model_class in self.all_model_classes:
                    a_					:		Dict							=  model_class(SCREAMING_SNAKE_CASE__     )
                    a_					:		str							=  inspect.signature(model.forward     )
                    # signature.parameters is an OrderedDict => so arg_names order is deterministic
                    a_					:		Optional[Any]							=  [*signature.parameters.keys()]

                    a_					:		List[str]							=  ['pixel_values']
                    self.assertListEqual(arg_names[:1]     ,							SCREAMING_SNAKE_CASE__     )
      @slow
      def        SCREAMING_SNAKE_CASE							(       self				:       int     )		->			str:
             for model_name in ["facebook/mask2former-swin-small-coco-instance"]:
                    a_					:		Any							=  MaskaFormerModel.from_pretrained(SCREAMING_SNAKE_CASE__     )
                    self.assertIsNotNone(SCREAMING_SNAKE_CASE__     )
      def        SCREAMING_SNAKE_CASE							(       self				:       List[Any]     )		->			str:
             a_					:		List[str]							=  (self.model_tester.min_size,) * 2
             a_					:		Dict							=  {
                 'pixel_values': torch.randn((2, 3, *size)     ,							device=SCREAMING_SNAKE_CASE__     ),
                 'mask_labels': torch.randn((2, 1_0, *size)     ,							device=SCREAMING_SNAKE_CASE__     ),
                 'class_labels': torch.zeros(2     ,							1_0     ,							device=SCREAMING_SNAKE_CASE__     ).long(),
             }
             a_					:		Optional[int]							=  self.model_tester.get_config()

             a_					:		int							=  MaskaFormerForUniversalSegmentation(SCREAMING_SNAKE_CASE__     ).to(SCREAMING_SNAKE_CASE__     )
             a_					:		Tuple							=  model(**SCREAMING_SNAKE_CASE__     )
             self.assertTrue(outputs.loss is not None     )
      def        SCREAMING_SNAKE_CASE							(       self				:       Optional[Any]     )		->			Optional[Any]:
             a_							,		a_					:		int							=  self.model_tester.prepare_config_and_inputs_for_common()
             self.model_tester.create_and_check_maskaformer_model(SCREAMING_SNAKE_CASE__     ,							**SCREAMING_SNAKE_CASE__     ,							output_hidden_states=SCREAMING_SNAKE_CASE__     )
      def        SCREAMING_SNAKE_CASE							(       self				:       Any     )		->			Tuple:
             a_							,		a_					:		Optional[Any]							=  self.model_tester.prepare_config_and_inputs_for_common()

             for model_class in self.all_model_classes:
                    a_					:		Any							=  model_class(SCREAMING_SNAKE_CASE__     ).to(SCREAMING_SNAKE_CASE__     )
                    a_					:		str							=  model(**SCREAMING_SNAKE_CASE__     ,							output_attentions=SCREAMING_SNAKE_CASE__     )
                    self.assertTrue(outputs.attentions is not None     )
      def        SCREAMING_SNAKE_CASE							(       self				:       Optional[int]     )		->			Dict:
             if not self.model_tester.is_training:
                    return

             a_					:		Optional[int]							=  self.all_model_classes[1]
             a_							,		a_							,		a_							,		a_							,		a_					:		int							=  self.model_tester.prepare_config_and_inputs()

             a_					:		int							=  model_class(SCREAMING_SNAKE_CASE__     )
             model.to(SCREAMING_SNAKE_CASE__     )
             model.train()

             a_					:		List[str]							=  model(SCREAMING_SNAKE_CASE__     ,							mask_labels=SCREAMING_SNAKE_CASE__     ,							class_labels=SCREAMING_SNAKE_CASE__     ).loss
             loss.backward()

      def        SCREAMING_SNAKE_CASE							(       self				:       List[str]     )		->			Optional[int]:
             a_					:		Dict							=  self.all_model_classes[1]
             a_							,		a_							,		a_							,		a_							,		a_					:		Tuple							=  self.model_tester.prepare_config_and_inputs()
             a_					:		Dict							=  True
             a_					:		Optional[int]							=  True

             a_					:		List[Any]							=  model_class(SCREAMING_SNAKE_CASE__     ).to(SCREAMING_SNAKE_CASE__     )
             model.train()

             a_					:		Union[str, Any]							=  model(SCREAMING_SNAKE_CASE__     ,							mask_labels=SCREAMING_SNAKE_CASE__     ,							class_labels=SCREAMING_SNAKE_CASE__     )

             a_					:		Union[str, Any]							=  outputs.encoder_hidden_states[0]
             encoder_hidden_states.retain_grad()

             a_					:		Dict							=  outputs.pixel_decoder_hidden_states[0]
             pixel_decoder_hidden_states.retain_grad()

             a_					:		Union[str, Any]							=  outputs.transformer_decoder_hidden_states[0]
             transformer_decoder_hidden_states.retain_grad()

             a_					:		List[str]							=  outputs.attentions[0]
             attentions.retain_grad()

             outputs.loss.backward(retain_graph=SCREAMING_SNAKE_CASE__     )

             self.assertIsNotNone(encoder_hidden_states.grad     )
             self.assertIsNotNone(pixel_decoder_hidden_states.grad     )
             self.assertIsNotNone(transformer_decoder_hidden_states.grad     )
             self.assertIsNotNone(attentions.grad     )
UpperCAmelCase_      : Dict  					=       1e-4







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

       """simple docstring"""

       a_					:		Any							=  Image.open('./tests/fixtures/tests_samples/COCO/000000039769.png'				)
       return image
@require_vision
@slow
class  SCREAMING_SNAKE_CASE__  (				unittest.TestCase			):
      @cached_property
      def        SCREAMING_SNAKE_CASE							(       self				:       Tuple     )		->			Tuple:
             return "facebook/mask2former-swin-small-coco-instance"
      @cached_property
      def        SCREAMING_SNAKE_CASE							(       self				:       Tuple     )		->			Tuple:
             return MaskaFormerImageProcessor.from_pretrained(self.model_checkpoints     ) if is_vision_available() else None
      def        SCREAMING_SNAKE_CASE							(       self				:       Any     )		->			List[Any]:
             a_					:		Tuple							=  MaskaFormerModel.from_pretrained(self.model_checkpoints     ).to(SCREAMING_SNAKE_CASE__     )
             a_					:		Any							=  self.default_image_processor
             a_					:		Union[str, Any]							=  prepare_img()
             a_					:		List[Any]							=  image_processor(SCREAMING_SNAKE_CASE__     ,							return_tensors='pt'     ).to(SCREAMING_SNAKE_CASE__     )
             a_					:		List[str]							=  inputs['pixel_values'].shape
             # check size is divisible by 32
             self.assertTrue((inputs_shape[-1] % 3_2) == 0 and (inputs_shape[-2] % 3_2) == 0     )
             # check size
             self.assertEqual(SCREAMING_SNAKE_CASE__     ,							(1, 3, 3_8_4, 3_8_4)     )

             with torch.no_grad():
                    a_					:		Union[str, Any]							=  model(**SCREAMING_SNAKE_CASE__     )

             a_					:		Optional[Any]							=  torch.tensor(
                 [[-0.2790, -1.0717, -1.1668], [-0.5128, -0.3128, -0.4987], [-0.5832, 0.1971, -0.0197]]     ).to(SCREAMING_SNAKE_CASE__     )
             self.assertTrue(
                 torch.allclose(
                     outputs.encoder_last_hidden_state[0, 0, :3, :3]     ,							SCREAMING_SNAKE_CASE__     ,							atol=SCREAMING_SNAKE_CASE__     )     )

             a_					:		str							=  torch.tensor(
                 [[0.8973, 1.1847, 1.1776], [1.1934, 1.5040, 1.5128], [1.1153, 1.4486, 1.4951]]     ).to(SCREAMING_SNAKE_CASE__     )
             self.assertTrue(
                 torch.allclose(
                     outputs.pixel_decoder_last_hidden_state[0, 0, :3, :3]     ,							SCREAMING_SNAKE_CASE__     ,							atol=SCREAMING_SNAKE_CASE__     )     )

             a_					:		str							=  torch.tensor(
                 [[2.1152, 1.7000, -0.8603], [1.5808, 1.8004, -0.9353], [1.6043, 1.7495, -0.5999]]     ).to(SCREAMING_SNAKE_CASE__     )
             self.assertTrue(
                 torch.allclose(
                     outputs.transformer_decoder_last_hidden_state[0, :3, :3]     ,							SCREAMING_SNAKE_CASE__     ,							atol=SCREAMING_SNAKE_CASE__     )     )
      def        SCREAMING_SNAKE_CASE							(       self				:       Any     )		->			Dict:
             a_					:		str							=  MaskaFormerForUniversalSegmentation.from_pretrained(self.model_checkpoints     ).to(SCREAMING_SNAKE_CASE__     ).eval()
             a_					:		List[Any]							=  self.default_image_processor
             a_					:		Tuple							=  prepare_img()
             a_					:		Optional[int]							=  image_processor(SCREAMING_SNAKE_CASE__     ,							return_tensors='pt'     ).to(SCREAMING_SNAKE_CASE__     )
             a_					:		Optional[Any]							=  inputs['pixel_values'].shape
             # check size is divisible by 32
             self.assertTrue((inputs_shape[-1] % 3_2) == 0 and (inputs_shape[-2] % 3_2) == 0     )
             # check size
             self.assertEqual(SCREAMING_SNAKE_CASE__     ,							(1, 3, 3_8_4, 3_8_4)     )

             with torch.no_grad():
                    a_					:		Union[str, Any]							=  model(**SCREAMING_SNAKE_CASE__     )
             # masks_queries_logits
             a_					:		Dict							=  outputs.masks_queries_logits
             self.assertEqual(
                 masks_queries_logits.shape     ,							(1, model.config.num_queries, inputs_shape[-2] // 4, inputs_shape[-1] // 4)     )
             a_					:		List[str]							=  [
                 [-8.7839, -9.0056, -8.8121],
                 [-7.4104, -7.0313, -6.5401],
                 [-6.6105, -6.3427, -6.4675],
             ]
             a_					:		Optional[Any]							=  torch.tensor(SCREAMING_SNAKE_CASE__     ).to(SCREAMING_SNAKE_CASE__     )
             self.assertTrue(torch.allclose(masks_queries_logits[0, 0, :3, :3]     ,							SCREAMING_SNAKE_CASE__     ,							atol=SCREAMING_SNAKE_CASE__     )     )
             # class_queries_logits
             a_					:		Tuple							=  outputs.class_queries_logits
             self.assertEqual(class_queries_logits.shape     ,							(1, model.config.num_queries, model.config.num_labels + 1)     )
             a_					:		Optional[int]							=  torch.tensor(
                 [
                     [1.8324, -8.0835, -4.1922],
                     [0.8450, -9.0050, -3.6053],
                     [0.3045, -7.7293, -3.0275],
                 ]     ).to(SCREAMING_SNAKE_CASE__     )
             self.assertTrue(torch.allclose(outputs.class_queries_logits[0, :3, :3]     ,							SCREAMING_SNAKE_CASE__     ,							atol=SCREAMING_SNAKE_CASE__     )     )

      def        SCREAMING_SNAKE_CASE							(       self				:       Optional[Any]     )		->			Any:
             a_					:		Tuple							=  MaskaFormerForUniversalSegmentation.from_pretrained(self.model_checkpoints     ).to(SCREAMING_SNAKE_CASE__     ).eval()
             a_					:		List[str]							=  self.default_image_processor

             a_					:		List[str]							=  image_processor(
                 [np.zeros((3, 8_0_0, 1_3_3_3)     ), np.zeros((3, 8_0_0, 1_3_3_3)     )]     ,							segmentation_maps=[np.zeros((3_8_4, 3_8_4)     ).astype(np.floataa     ), np.zeros((3_8_4, 3_8_4)     ).astype(np.floataa     )]     ,							return_tensors='pt'     ,							)

             a_					:		Any							=  inputs['pixel_values'].to(SCREAMING_SNAKE_CASE__     )
             a_					:		List[Any]							=  [el.to(SCREAMING_SNAKE_CASE__     ) for el in inputs['mask_labels']]
             a_					:		str							=  [el.to(SCREAMING_SNAKE_CASE__     ) for el in inputs['class_labels']]

             with torch.no_grad():
                    a_					:		Tuple							=  model(**SCREAMING_SNAKE_CASE__     )

             self.assertTrue(outputs.loss is not None     )






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  SCREAMING_SNAKE_CASE__  (				lowercase__						,    unittest.TestCase			):
      # TODO: is there an appropriate internal test set?
      snake_case__				:     Any        =						'''ssube/stable-diffusion-x4-upscaler-onnx'''
      def        SCREAMING_SNAKE_CASE							(       self				:       List[str]     ,							SCREAMING_SNAKE_CASE__				:       int=0     )		->			Tuple:
             a_					:		Union[str, Any]							=  floats_tensor((1, 3, 1_2_8, 1_2_8)     ,							rng=random.Random(SCREAMING_SNAKE_CASE__     )     )
             a_					:		Dict							=  torch.manual_seed(SCREAMING_SNAKE_CASE__     )
             a_					:		Union[str, 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        SCREAMING_SNAKE_CASE							(       self				:       List[str]     )		->			Dict:
             a_					:		Dict							=  OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint     ,							provider='CPUExecutionProvider'     )
             pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE__     )

             a_					:		Optional[int]							=  self.get_dummy_inputs()
             a_					:		int							=  pipe(**SCREAMING_SNAKE_CASE__     ).images
             a_					:		Tuple							=  image[0, -3:, -3:, -1].flatten()

             # started as 128, should now be 512
             assert image.shape == (1, 5_1_2, 5_1_2, 3)
             a_					:		List[Any]							=  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        SCREAMING_SNAKE_CASE							(       self				:       Dict     )		->			Dict:
             a_					:		Tuple							=  OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint     ,							provider='CPUExecutionProvider'     )
             a_					:		int							=  PNDMScheduler.from_config(pipe.scheduler.config     ,							skip_prk_steps=SCREAMING_SNAKE_CASE__     )
             pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE__     )

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

             assert image.shape == (1, 5_1_2, 5_1_2, 3)
             a_					:		str							=  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        SCREAMING_SNAKE_CASE							(       self				:       Any     )		->			Dict:
             a_					:		Optional[Any]							=  OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint     ,							provider='CPUExecutionProvider'     )
             a_					:		Optional[Any]							=  DPMSolverMultistepScheduler.from_config(pipe.scheduler.config     )
             pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE__     )

             a_					:		Optional[int]							=  self.get_dummy_inputs()
             a_					:		Dict							=  pipe(**SCREAMING_SNAKE_CASE__     ).images
             a_					:		str							=  image[0, -3:, -3:, -1]

             assert image.shape == (1, 5_1_2, 5_1_2, 3)
             a_					:		Optional[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        SCREAMING_SNAKE_CASE							(       self				:       List[Any]     )		->			int:
             a_					:		List[str]							=  OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint     ,							provider='CPUExecutionProvider'     )
             a_					:		int							=  EulerDiscreteScheduler.from_config(pipe.scheduler.config     )
             pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE__     )

             a_					:		Optional[Any]							=  self.get_dummy_inputs()
             a_					:		Dict							=  pipe(**SCREAMING_SNAKE_CASE__     ).images
             a_					:		List[str]							=  image[0, -3:, -3:, -1]

             assert image.shape == (1, 5_1_2, 5_1_2, 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.flatten() - expected_slice     ).max() < 1E-1

      def        SCREAMING_SNAKE_CASE							(       self				:       str     )		->			Optional[Any]:
             a_					:		Tuple							=  OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint     ,							provider='CPUExecutionProvider'     )
             a_					:		Optional[int]							=  EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config     )
             pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE__     )

             a_					:		Union[str, Any]							=  self.get_dummy_inputs()
             a_					:		List[str]							=  pipe(**SCREAMING_SNAKE_CASE__     ).images
             a_					:		List[str]							=  image[0, -3:, -3:, -1]

             assert image.shape == (1, 5_1_2, 5_1_2, 3)
             a_					:		Union[str, Any]							=  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  SCREAMING_SNAKE_CASE__  (				unittest.TestCase			):
      @property
      def        SCREAMING_SNAKE_CASE							(       self				:       Tuple     )		->			Union[str, Any]:
             return (
                 "CUDAExecutionProvider",
                 {
                     "gpu_mem_limit": "15000000000",  # 15GB
                     "arena_extend_strategy": "kSameAsRequested",
                 },
             )
      @property
      def        SCREAMING_SNAKE_CASE							(       self				:       Optional[Any]     )		->			List[str]:
             a_					:		List[str]							=  ort.SessionOptions()
             a_					:		int							=  False
             return options
      def        SCREAMING_SNAKE_CASE							(       self				:       List[str]     )		->			Tuple:
             a_					:		str							=  load_image(
                 'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
                 '/img2img/sketch-mountains-input.jpg'     )
             a_					:		int							=  init_image.resize((1_2_8, 1_2_8)     )
             # 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=SCREAMING_SNAKE_CASE__     )

             a_					:		Tuple							=  'A fantasy landscape, trending on artstation'

             a_					:		str							=  torch.manual_seed(0     )
             a_					:		List[str]							=  pipe(
                 prompt=SCREAMING_SNAKE_CASE__     ,							image=SCREAMING_SNAKE_CASE__     ,							guidance_scale=7.5     ,							num_inference_steps=1_0     ,							generator=SCREAMING_SNAKE_CASE__     ,							output_type='np'     ,							)
             a_					:		Dict							=  output.images
             a_					:		Any							=  images[0, 2_5_5:2_5_8, 3_8_3:3_8_6, -1]

             assert images.shape == (1, 5_1_2, 5_1_2, 3)
             a_					:		str							=  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        SCREAMING_SNAKE_CASE							(       self				:       Optional[int]     )		->			List[str]:
             a_					:		Dict							=  load_image(
                 'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
                 '/img2img/sketch-mountains-input.jpg'     )
             a_					:		List[str]							=  init_image.resize((1_2_8, 1_2_8)     )
             a_					:		Dict							=  LMSDiscreteScheduler.from_pretrained(
                 'ssube/stable-diffusion-x4-upscaler-onnx'     ,							subfolder='scheduler'     )
             a_					:		Union[str, Any]							=  OnnxStableDiffusionUpscalePipeline.from_pretrained(
                 'ssube/stable-diffusion-x4-upscaler-onnx'     ,							scheduler=SCREAMING_SNAKE_CASE__     ,							provider=self.gpu_provider     ,							sess_options=self.gpu_options     ,							)
             pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE__     )

             a_					:		Any							=  'A fantasy landscape, trending on artstation'

             a_					:		Tuple							=  torch.manual_seed(0     )
             a_					:		Optional[Any]							=  pipe(
                 prompt=SCREAMING_SNAKE_CASE__     ,							image=SCREAMING_SNAKE_CASE__     ,							guidance_scale=7.5     ,							num_inference_steps=2_0     ,							generator=SCREAMING_SNAKE_CASE__     ,							output_type='np'     ,							)
             a_					:		str							=  output.images
             a_					:		List[Any]							=  images[0, 2_5_5:2_5_8, 3_8_3:3_8_6, -1]

             assert images.shape == (1, 5_1_2, 5_1_2, 3)
             a_					:		Tuple							=  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


