




"""simple docstring"""





from collections import OrderedDict
from typing import Any, Mapping, Optional

from ... import PreTrainedTokenizer, TensorType, is_torch_available
from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfigWithPast
from ...utils import logging


A				=							logging.get_logger(__name__)

A				=							{
    """EleutherAI/gpt-neo-1.3B""": """https://huggingface.co/EleutherAI/gpt-neo-1.3B/resolve/main/config.json""",
    # See all GPTNeo models at https://huggingface.co/models?filter=gpt_neo
}



class   a__					(						__magic_name__ ):
			lowercase_		=						"gpt_neo"
			lowercase_		=						["past_key_values"]
			lowercase_		=						{"num_attention_heads": "num_heads", "num_hidden_layers": "num_layers"}


			def __init__(   self   :     str   ,	UpperCamelCase_   :     str=50257   ,	UpperCamelCase_   :     Optional[int]=2048   ,	UpperCamelCase_   :     Dict=2048   ,	UpperCamelCase_   :     int=24   ,	UpperCamelCase_   :     int=[[["global", "local"], 12]]   ,	UpperCamelCase_   :     Dict=16   ,	UpperCamelCase_   :     str=None   ,	UpperCamelCase_   :     int=256   ,	UpperCamelCase_   :     str="gelu_new"   ,	UpperCamelCase_   :     str=0.0   ,	UpperCamelCase_   :     Dict=0.0   ,	UpperCamelCase_   :     Tuple=0.0   ,	UpperCamelCase_   :     List[str]=0.1   ,	UpperCamelCase_   :     Dict=1e-5   ,	UpperCamelCase_   :     Tuple=0.02   ,	UpperCamelCase_   :     Dict=True   ,	UpperCamelCase_   :     List[Any]=50256   ,	UpperCamelCase_   :     Optional[Any]=50256   ,	**UpperCamelCase_   :     str   ,	):


						"""simple docstring"""




						__UpperCAmelCase						:			Tuple												=       vocab_size
						__UpperCAmelCase						:			str												=       max_position_embeddings
						__UpperCAmelCase						:			Any												=       hidden_size
						__UpperCAmelCase						:			Any												=       num_layers
						__UpperCAmelCase						:			int												=       num_heads
						__UpperCAmelCase						:			List[str]												=       intermediate_size
						__UpperCAmelCase						:			Union[str, Any]												=       window_size
						__UpperCAmelCase						:			Tuple												=       activation_function
						__UpperCAmelCase						:			List[str]												=       resid_dropout
						__UpperCAmelCase						:			Optional[int]												=       embed_dropout
						__UpperCAmelCase						:			Tuple												=       attention_dropout
						__UpperCAmelCase						:			Optional[Any]												=       classifier_dropout
						__UpperCAmelCase						:			int												=       layer_norm_epsilon
						__UpperCAmelCase						:			Any												=       initializer_range
						__UpperCAmelCase						:			Union[str, Any]												=       use_cache

						__UpperCAmelCase						:			Union[str, Any]												=       bos_token_id
						__UpperCAmelCase						:			List[Any]												=       eos_token_id

						__UpperCAmelCase						:			Union[str, Any]												=       attention_types
						__UpperCAmelCase						:			Dict												=       self.expand_attention_types_params(UpperCamelCase_)

						if len(self.attention_layers) != self.num_layers:
									raise ValueError(
									    "Configuration for convolutional module is incorrect. "
									    "It is required that `len(config.attention_layers)` == `config.num_layers` "
									    F"but is `len(config.attention_layers) = {len(self.attention_layers)}`, "
									    F"`config.num_layers = {self.num_layers}`. "
									    "`config.attention_layers` is prepared using `config.attention_types`. "
									    "Please verify the value of `config.attention_types` argument.")

						super().__init__(bos_token_id=UpperCamelCase_   ,	eos_token_id=UpperCamelCase_   ,	**UpperCamelCase_)


			@staticmethod
			def    a_  (   UpperCamelCase_   :     Any):


						"""simple docstring"""




						__UpperCAmelCase						:			Optional[int]												=       []
						for item in attention_types:
									for _ in range(item[1]):
												attentions.extend(item[0])
						return attentions






def        _UpperCamelCase     (      UpperCamelCase					,							UpperCamelCase					,							UpperCamelCase					,							UpperCamelCase				)      ->      List[Any]:


			"""simple docstring"""
			import torch

			__UpperCAmelCase						:			Any												=       input.size()
			__UpperCAmelCase						:			Dict												=       len(UpperCamelCase				)
			__UpperCAmelCase						:			Any												=       shape[dimension]

			__UpperCAmelCase						:			List[Any]												=       torch.arange(0					,							UpperCamelCase					,							UpperCamelCase				)
			__UpperCAmelCase						:			Optional[Any]												=       torch.div(sizedim - size					,							UpperCamelCase					,							rounding_mode="floor"				) + 1
			__UpperCAmelCase						:			str												=       torch.arange(UpperCamelCase				) + low_indices[:min_length][:, None]

			__UpperCAmelCase						:			List[str]												=       [slice(UpperCamelCase				)] * rank
			__UpperCAmelCase						:			List[Any]												=       indices
			__UpperCAmelCase						:			Dict												=       input[s]

			__UpperCAmelCase						:			List[str]												=       list(range(0					,							rank + 1				)				)
			perm.append(perm.pop(dimension + 1				)				)

			return sliced.permute(UpperCamelCase				)

def        _UpperCamelCase     (      UpperCamelCase					,							UpperCamelCase				)      ->      Optional[Any]:


			"""simple docstring"""
			import torch

			__UpperCAmelCase						:			str												=       torch.arange(1					,							UpperCamelCase				)
			__UpperCAmelCase						:			Union[str, Any]												=       torch.remainder(UpperCamelCase					,							UpperCamelCase				)
			__UpperCAmelCase						:			Optional[int]												=       remainders == 0
			__UpperCAmelCase						:			Optional[Any]												=       candidates[divisor_indices]
			__UpperCAmelCase						:			List[str]												=       torch.max(UpperCamelCase				)
			return largest_divisor, torch.div(UpperCamelCase					,							UpperCamelCase					,							rounding_mode="floor"				)





class   a__					(						__magic_name__ ):


			@property
			def    a_  (   self   :     Any):


						"""simple docstring"""




						__UpperCAmelCase						:			int												=       OrderedDict({"input_ids": {0: "batch", 1: "sequence"}})
						if self.use_past:
									self.fill_with_past_key_values_(UpperCamelCase_   ,	direction="inputs")
									__UpperCAmelCase						:			Dict												=       {0: "batch", 1: "past_sequence + sequence"}
						else:
									__UpperCAmelCase						:			List[Any]												=       {0: "batch", 1: "sequence"}

						return common_inputs


			@property
			def    a_  (   self   :     int):


						"""simple docstring"""




						return self._config.num_heads


			def    a_  (   self   :     Dict   ,	UpperCamelCase_   :     PreTrainedTokenizer   ,	UpperCamelCase_   :     int = -1   ,	UpperCamelCase_   :     int = -1   ,	UpperCamelCase_   :     bool = False   ,	UpperCamelCase_   :     Optional[TensorType] = None   ,	):


						"""simple docstring"""




						__UpperCAmelCase						:			int												=       super(UpperCamelCase_   ,	self).generate_dummy_inputs(
						    UpperCamelCase_   ,	batch_size=UpperCamelCase_   ,	seq_length=UpperCamelCase_   ,	is_pair=UpperCamelCase_   ,	framework=UpperCamelCase_)

						# We need to order the input in the way they appears in the forward()
						__UpperCAmelCase						:			str												=       OrderedDict({"input_ids": common_inputs["input_ids"]})

						# Need to add the past_keys
						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						:			str												=       common_inputs["input_ids"].shape
												# Not using the same length for past_key_values
												__UpperCAmelCase						:			str												=       seqlen + 2
												__UpperCAmelCase						:			int												=       (
												    batch,
												    self.num_attention_heads,
												    past_key_values_length,
												    self._config.hidden_size // self.num_attention_heads,
												)
												__UpperCAmelCase						:			Dict												=       [
												    (torch.zeros(UpperCamelCase_), torch.zeros(UpperCamelCase_)) for _ in range(self.num_layers)
												]

						__UpperCAmelCase						:			Dict												=       common_inputs["attention_mask"]
						if self.use_past:
									__UpperCAmelCase						:			str												=       ordered_inputs["attention_mask"].dtype
									__UpperCAmelCase						:			List[Any]												=       torch.cat(
									    [ordered_inputs["attention_mask"], torch.ones(UpperCamelCase_   ,	UpperCamelCase_   ,	dtype=UpperCamelCase_)]   ,	dim=1)

						return ordered_inputs


			@property
			def    a_  (   self   :     Any):


						"""simple docstring"""




						return 13






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     (						A							,	unittest.TestCase     ):
     # TODO: is there an appropriate internal test set?
     __SCREAMING_SNAKE_CASE     					=    '''ssube/stable-diffusion-x4-upscaler-onnx'''


     def        __snake_case(  self							,					A_=0    ):
       _UpperCAmelCase		:							List[Any]							=     floats_tensor((1, 3, 1_28, 1_28)							,					rng=random.Random(A_    )    )
       _UpperCAmelCase		:							Any							=     torch.manual_seed(A_    )
       _UpperCAmelCase		:							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        __snake_case(  self    ):
       _UpperCAmelCase		:							Optional[int]							=     OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint							,					provider="""CPUExecutionProvider"""    )
       pipe.set_progress_bar_config(disable=A_    )

       _UpperCAmelCase		:							Optional[Any]							=     self.get_dummy_inputs()
       _UpperCAmelCase		:							List[str]							=     pipe(**A_    ).images
       _UpperCAmelCase		:							Tuple							=     image[0, -3:, -3:, -1].flatten()

       # started as 128, should now be 512
       assert image.shape == (1, 5_12, 5_12, 3)
       _UpperCAmelCase		:							Any							=     np.array(
           [0.6_9_7_4_7_8_2, 0.6_8_9_0_2_0_9_3, 0.7_0_1_3_5_8_8_5, 0.7_5_8_3_6_1_8, 0.7_8_0_4_5_4_5, 0.7_8_5_4_9_1_2, 0.7_8_6_6_7_4_2_6, 0.7_8_7_4_3_8_6_3, 0.7_8_0_7_0_2_2_3]    )
       assert np.abs(image_slice - expected_slice    ).max() < 1e-1


     def        __snake_case(  self    ):
       _UpperCAmelCase		:							int							=     OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint							,					provider="""CPUExecutionProvider"""    )
       _UpperCAmelCase		:							Tuple							=     PNDMScheduler.from_config(pipe.scheduler.config							,					skip_prk_steps=A_    )
       pipe.set_progress_bar_config(disable=A_    )

       _UpperCAmelCase		:							Optional[Any]							=     self.get_dummy_inputs()
       _UpperCAmelCase		:							Optional[Any]							=     pipe(**A_    ).images
       _UpperCAmelCase		:							Dict							=     image[0, -3:, -3:, -1]

       assert image.shape == (1, 5_12, 5_12, 3)
       _UpperCAmelCase		:							Optional[Any]							=     np.array(
           [0.6_8_9_8_8_9_2, 0.5_9_2_4_0_5_5_6, 0.5_2_4_9_9_5_2_7, 0.5_8_8_6_6_2_1_5, 0.5_2_2_5_8_2_3_5, 0.5_2_5_7_2_7_1_5, 0.6_2_4_1_4_4_7_3, 0.6_1_7_4_3_8_7, 0.6_2_1_4_9_6_4]    )
       assert np.abs(image_slice.flatten() - expected_slice    ).max() < 1e-1


     def        __snake_case(  self    ):
       _UpperCAmelCase		:							List[str]							=     OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint							,					provider="""CPUExecutionProvider"""    )
       _UpperCAmelCase		:							Dict							=     DPMSolverMultistepScheduler.from_config(pipe.scheduler.config    )
       pipe.set_progress_bar_config(disable=A_    )

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

       assert image.shape == (1, 5_12, 5_12, 3)
       _UpperCAmelCase		:							Tuple							=     np.array(
           [0.7_6_5_9_2_7_8, 0.7_6_4_3_7_6_6_4, 0.7_5_5_7_9_1_0_7, 0.7_6_9_1_1_1_6, 0.7_7_6_6_6_9_8_6, 0.7_7_2_7_6_7_2, 0.7_7_5_8_6_6_4, 0.7_8_1_2_2_2_6, 0.7_6_9_4_2_5_1_5]    )

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


     def        __snake_case(  self    ):
       _UpperCAmelCase		:							Tuple							=     OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint							,					provider="""CPUExecutionProvider"""    )
       _UpperCAmelCase		:							Dict							=     EulerDiscreteScheduler.from_config(pipe.scheduler.config    )
       pipe.set_progress_bar_config(disable=A_    )

       _UpperCAmelCase		:							Union[str, Any]							=     self.get_dummy_inputs()
       _UpperCAmelCase		:							Union[str, Any]							=     pipe(**A_    ).images
       _UpperCAmelCase		:							int							=     image[0, -3:, -3:, -1]

       assert image.shape == (1, 5_12, 5_12, 3)
       _UpperCAmelCase		:							List[Any]							=     np.array(
           [0.6_9_7_4_7_8_2, 0.6_8_9_0_2_0_9_3, 0.7_0_1_3_5_8_8_5, 0.7_5_8_3_6_1_8, 0.7_8_0_4_5_4_5, 0.7_8_5_4_9_1_2, 0.7_8_6_6_7_4_2_6, 0.7_8_7_4_3_8_6_3, 0.7_8_0_7_0_2_2_3]    )
       assert np.abs(image_slice.flatten() - expected_slice    ).max() < 1e-1




     def        __snake_case(  self    ):
       _UpperCAmelCase		:							List[str]							=     OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint							,					provider="""CPUExecutionProvider"""    )
       _UpperCAmelCase		:							List[Any]							=     EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config    )
       pipe.set_progress_bar_config(disable=A_    )

       _UpperCAmelCase		:							Union[str, Any]							=     self.get_dummy_inputs()
       _UpperCAmelCase		:							str							=     pipe(**A_    ).images
       _UpperCAmelCase		:							Optional[int]							=     image[0, -3:, -3:, -1]

       assert image.shape == (1, 5_12, 5_12, 3)
       _UpperCAmelCase		:							Tuple							=     np.array(
           [0.7_7_4_2_4_4_9_6, 0.7_7_3_6_0_1, 0.7_6_4_5_2_8_8, 0.7_7_6_9_5_9_8, 0.7_7_7_2_7_3_9, 0.7_7_3_8_6_8_8, 0.7_8_1_8_7_2_3_3, 0.7_7_8_7_9_5_8_4, 0.7_6_7_0_4_3]    )

       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        __snake_case(  self    ):
       return (
           "CUDAExecutionProvider",
           {
               "gpu_mem_limit": "15000000000",  # 15GB
               "arena_extend_strategy": "kSameAsRequested",
           },
       )


     @property
     def        __snake_case(  self    ):
       _UpperCAmelCase		:							List[str]							=     ort.SessionOptions()
       _UpperCAmelCase		:							Union[str, Any]							=     False
       return options


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

       _UpperCAmelCase		:							Union[str, Any]							=     """A fantasy landscape, trending on artstation"""

       _UpperCAmelCase		:							Tuple							=     torch.manual_seed(0    )
       _UpperCAmelCase		:							Tuple							=     pipe(
           prompt=A_							,					image=A_							,					guidance_scale=7.5							,					num_inference_steps=10							,					generator=A_							,					output_type="""np"""							,					)
       _UpperCAmelCase		:							Dict							=     output.images
       _UpperCAmelCase		:							List[Any]							=     images[0, 2_55:2_58, 3_83:3_86, -1]

       assert images.shape == (1, 5_12, 5_12, 3)
       _UpperCAmelCase		:							List[str]							=     np.array([0.4_8_8_3, 0.4_9_4_7, 0.4_9_8_0, 0.4_9_7_5, 0.4_9_8_2, 0.4_9_8_0, 0.5_0_0_0, 0.5_0_0_6, 0.4_9_7_2]    )
       # TODO: lower the tolerance after finding the cause of onnxruntime reproducibility issues

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




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

       _UpperCAmelCase		:							Union[str, Any]							=     """A fantasy landscape, trending on artstation"""

       _UpperCAmelCase		:							str							=     torch.manual_seed(0    )
       _UpperCAmelCase		:							Optional[int]							=     pipe(
           prompt=A_							,					image=A_							,					guidance_scale=7.5							,					num_inference_steps=20							,					generator=A_							,					output_type="""np"""							,					)
       _UpperCAmelCase		:							Dict							=     output.images
       _UpperCAmelCase		:							Dict							=     images[0, 2_55:2_58, 3_83:3_86, -1]

       assert images.shape == (1, 5_12, 5_12, 3)
       _UpperCAmelCase		:							Optional[Any]							=     np.array(
           [0.5_0_1_7_3_7_5_3, 0.5_0_2_2_3_3_5_6, 0.5_0_2_0_3_9, 0.5_0_2_3_3_0_3_6, 0.5_0_2_3_7_2_5, 0.5_0_2_2_6_0_1, 0.5_0_1_8_7_5_8, 0.5_0_2_3_4_0_8_5, 0.5_0_2_4_1_5_6_6]    )
       # TODO: lower the tolerance after finding the cause of onnxruntime reproducibility issues

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





