





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


   def 	SCREAMING_SNAKE_CASE_							(   self      :    Dict      ,					a      :    int=0	)->       Any:





     """simple docstring"""







     lowercase__           =  floats_tensor((1, 3, 128, 128)      ,					rng=random.Random(__A	)	)
     lowercase__           =  torch.manual_seed(__A	)
     lowercase__           =  {
         '''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      :    Dict	)->       Dict:





     """simple docstring"""







     lowercase__           =  OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint      ,					provider='CPUExecutionProvider'	)
     pipe.set_progress_bar_config(disable=__A	)

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

     # started as 128, should now be 512
     assert image.shape == (1, 512, 512, 3)
     lowercase__           =  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      :    List[Any]	)->       List[str]:





     """simple docstring"""







     lowercase__           =  OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint      ,					provider='CPUExecutionProvider'	)
     lowercase__           =  PNDMScheduler.from_config(pipe.scheduler.config      ,					skip_prk_steps=__A	)
     pipe.set_progress_bar_config(disable=__A	)

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

     assert image.shape == (1, 512, 512, 3)
     lowercase__           =  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      :    Optional[int]	)->       int:





     """simple docstring"""







     lowercase__           =  OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint      ,					provider='CPUExecutionProvider'	)
     lowercase__           =  DPMSolverMultistepScheduler.from_config(pipe.scheduler.config	)
     pipe.set_progress_bar_config(disable=__A	)

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

     assert image.shape == (1, 512, 512, 3)
     lowercase__           =  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      :    Any	)->       str:





     """simple docstring"""







     lowercase__           =  OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint      ,					provider='CPUExecutionProvider'	)
     lowercase__           =  EulerDiscreteScheduler.from_config(pipe.scheduler.config	)
     pipe.set_progress_bar_config(disable=__A	)

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

     assert image.shape == (1, 512, 512, 3)
     lowercase__           =  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      :    Dict	)->       Tuple:





     """simple docstring"""







     lowercase__           =  OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint      ,					provider='CPUExecutionProvider'	)
     lowercase__           =  EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config	)
     pipe.set_progress_bar_config(disable=__A	)

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

     assert image.shape == (1, 512, 512, 3)
     lowercase__           =  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      :    str	)->       str:





     """simple docstring"""







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


   @property
   def 	SCREAMING_SNAKE_CASE_							(   self      :    List[str]	)->       List[str]:





     """simple docstring"""







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


   def 	SCREAMING_SNAKE_CASE_							(   self      :    Any	)->       List[Any]:





     """simple docstring"""







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

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

     lowercase__           =  torch.manual_seed(0	)
     lowercase__           =  pipe(
         prompt=__A      ,					image=__A      ,					guidance_scale=7.5      ,					num_inference_steps=10      ,					generator=__A      ,					output_type='np'      ,					)
     lowercase__           =  output.images
     lowercase__           =  images[0, 255:258, 383:386, -1]

     assert images.shape == (1, 512, 512, 3)
     lowercase__           =  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      :    Union[str, Any]	)->       List[str]:





     """simple docstring"""







     lowercase__           =  load_image(
         'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
         '/img2img/sketch-mountains-input.jpg'	)
     lowercase__           =  init_image.resize((128, 128)	)
     lowercase__           =  LMSDiscreteScheduler.from_pretrained(
         'ssube/stable-diffusion-x4-upscaler-onnx'      ,					subfolder='scheduler'	)
     lowercase__           =  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	)

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

     lowercase__           =  torch.manual_seed(0	)
     lowercase__           =  pipe(
         prompt=__A      ,					image=__A      ,					guidance_scale=7.5      ,					num_inference_steps=20      ,					generator=__A      ,					output_type='np'      ,					)
     lowercase__           =  output.images
     lowercase__           =  images[0, 255:258, 383:386, -1]

     assert images.shape == (1, 512, 512, 3)
     lowercase__           =  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









from ...configuration_utils import PretrainedConfig
from ...utils import logging


lowercase_          =		logging.get_logger(__name__)

lowercase_          =		{
    """transfo-xl-wt103""": """https://huggingface.co/transfo-xl-wt103/resolve/main/config.json""",
}
class     SCREAMING_SNAKE_CASE       (UpperCAmelCase				):
   _UpperCamelCase   :		Optional[Any]			 =	'transfo-xl'
   _UpperCamelCase   :		Any			 =	['mems']
   _UpperCamelCase   :		Any			 =	{
       'n_token': 'vocab_size',
       'hidden_size': 'd_model',
       'num_attention_heads': 'n_head',
       'num_hidden_layers': 'n_layer',
   }


   def __init__(   self      :    Optional[Any]      ,					a      :    Optional[int]=267_735      ,					a      :    str=[20_000, 40_000, 200_000]      ,					a      :    str=1_024      ,					a      :    str=1_024      ,					a      :    int=16      ,					a      :    Optional[int]=64      ,					a      :    Optional[int]=4_096      ,					a      :    int=4      ,					a      :    Tuple=False      ,					a      :    Any=18      ,					a      :    Tuple=1_600      ,					a      :    Union[str, Any]=1_000      ,					a      :    str=True      ,					a      :    Dict=True      ,					a      :    Any=0      ,					a      :    List[Any]=-1      ,					a      :    List[Any]=True      ,					a      :    Tuple=0.1      ,					a      :    List[Any]=0.0      ,					a      :    Optional[Any]=True      ,					a      :    int="normal"      ,					a      :    Optional[Any]=0.01      ,					a      :    str=0.01      ,					a      :    List[Any]=0.02      ,					a      :    List[Any]=1E-5      ,					a      :    Optional[Any]=0      ,					**a      :    Optional[int]      ,					)->       Optional[int]:





     """simple docstring"""







     lowercase__           =  vocab_size
     lowercase__           =  []
     self.cutoffs.extend(a	)
     if proj_share_all_but_first:
       lowercase__           =  [False] + [True] * len(self.cutoffs	)
     else:
       lowercase__           =  [False] + [False] * len(self.cutoffs	)
     lowercase__           =  d_model
     lowercase__           =  d_embed
     lowercase__           =  d_head
     lowercase__           =  d_inner
     lowercase__           =  div_val
     lowercase__           =  pre_lnorm
     lowercase__           =  n_layer
     lowercase__           =  n_head
     lowercase__           =  mem_len
     lowercase__           =  same_length
     lowercase__           =  attn_type
     lowercase__           =  clamp_len
     lowercase__           =  sample_softmax
     lowercase__           =  adaptive
     lowercase__           =  dropout
     lowercase__           =  dropatt
     lowercase__           =  untie_r
     lowercase__           =  init
     lowercase__           =  init_range
     lowercase__           =  proj_init_std
     lowercase__           =  init_std
     lowercase__           =  layer_norm_epsilon
     super().__init__(eos_token_id=a      ,					**a	)


   @property
   def 	SCREAMING_SNAKE_CASE_							(   self      :    Optional[Any]	)->       Union[str, Any]:





     """simple docstring"""







     logger.info(f"""The model {self.model_type} is one of the few models that has no sequence length limit."""	)
     return -1


   @max_position_embeddings.setter
   def 	SCREAMING_SNAKE_CASE_							(   self      :    Any      ,					a      :    Optional[int]	)->       Optional[int]:





     """simple docstring"""







     raise NotImplementedError(
         f"""The model {self.model_type} is one of the few models that has no sequence length limit."""	)


