


from dataclasses import dataclass
from typing import Optional

import numpy as np
import torch
import torch.nn as nn

from ..utils import BaseOutput, is_torch_version, randn_tensor
from .attention_processor import SpatialNorm
from .unet_ad_blocks import UNetMidBlockaD, get_down_block, get_up_block







@dataclass
class 							_UpperCamelCase (						A						):

   '''simple docstring'''







   a_				:							torch.FloatTensor







class 							_UpperCamelCase (						nn.Module						):

   '''simple docstring'''







   def __init__(  self  : int       ,  _lowerCamelCase  : str=3       ,  _lowerCamelCase  : Dict=3       ,  _lowerCamelCase  : Optional[int]=("DownEncoderBlock2D",)       ,  _lowerCamelCase  : List[str]=(6_4,)       ,  _lowerCamelCase  : List[str]=2       ,  _lowerCamelCase  : Optional[int]=3_2       ,  _lowerCamelCase  : int="silu"       ,  _lowerCamelCase  : str=True       ,  ):






         '''simple docstring'''



         super().__init__()
         __lowerCamelCase     :      List[str]       						=						layers_per_block

         __lowerCamelCase     :      Optional[int]       						=						torch.nn.Convad(
             _lowerCamelCase       ,  block_out_channels[0]       ,  kernel_size=3       ,  stride=1       ,  padding=1       ,  )

         __lowerCamelCase     :      Any       						=						None
         __lowerCamelCase     :      Optional[int]       						=						nn.ModuleList([]							)

         # down
         __lowerCamelCase     :      Dict       						=						block_out_channels[0]
         for i, down_block_type in enumerate(_lowerCamelCase							):
               __lowerCamelCase     :      Optional[Any]       						=						output_channel
               __lowerCamelCase     :      Dict       						=						block_out_channels[i]
               __lowerCamelCase     :      List[Any]       						=						i == len(_lowerCamelCase							) - 1

               __lowerCamelCase     :      Dict       						=						get_down_block(
                   _lowerCamelCase       ,  num_layers=self.layers_per_block       ,  in_channels=_lowerCamelCase       ,  out_channels=_lowerCamelCase       ,  add_downsample=not is_final_block       ,  resnet_eps=1E-6       ,  downsample_padding=0       ,  resnet_act_fn=_lowerCamelCase       ,  resnet_groups=_lowerCamelCase       ,  attention_head_dim=_lowerCamelCase       ,  temb_channels=_lowerCamelCase       ,  )
               self.down_blocks.append(_lowerCamelCase							)

         # mid
         __lowerCamelCase     :      str       						=						UNetMidBlockaD(
             in_channels=block_out_channels[-1]       ,  resnet_eps=1E-6       ,  resnet_act_fn=_lowerCamelCase       ,  output_scale_factor=1       ,  resnet_time_scale_shift="""default"""       ,  attention_head_dim=block_out_channels[-1]       ,  resnet_groups=_lowerCamelCase       ,  temb_channels=_lowerCamelCase       ,  )

         # out
         __lowerCamelCase     :      Any       						=						nn.GroupNorm(num_channels=block_out_channels[-1]       ,  num_groups=_lowerCamelCase       ,  eps=1E-6							)
         __lowerCamelCase     :      Optional[int]       						=						nn.SiLU()

         __lowerCamelCase     :      Dict       						=						2 * out_channels if double_z else out_channels
         __lowerCamelCase     :      Tuple       						=						nn.Convad(block_out_channels[-1]       ,  _lowerCamelCase       ,  3       ,  padding=1							)

         __lowerCamelCase     :      Any       						=						False


   def   _snake_case    (  self  : Optional[int]       ,  _lowerCamelCase  : Tuple							):






         '''simple docstring'''



         __lowerCamelCase     :      Tuple       						=						x
         __lowerCamelCase     :      int       						=						self.conv_in(_lowerCamelCase							)

         if self.training and self.gradient_checkpointing:

               def create_custom_forward(_lowerCamelCase  : Tuple							):
                     def custom_forward(*_lowerCamelCase  : Any							):
                           return module(*_lowerCamelCase							)

                     return custom_forward

               # down
               if is_torch_version(""">="""       ,  """1.11.0"""							):
                     for down_block in self.down_blocks:
                           __lowerCamelCase     :      int       						=						torch.utils.checkpoint.checkpoint(
                               create_custom_forward(_lowerCamelCase							)       ,  _lowerCamelCase       ,  use_reentrant=_lowerCamelCase							)
                     # middle
                     __lowerCamelCase     :      List[Any]       						=						torch.utils.checkpoint.checkpoint(
                         create_custom_forward(self.mid_block							)       ,  _lowerCamelCase       ,  use_reentrant=_lowerCamelCase							)
               else:
                     for down_block in self.down_blocks:
                           __lowerCamelCase     :      List[str]       						=						torch.utils.checkpoint.checkpoint(create_custom_forward(_lowerCamelCase							)       ,  _lowerCamelCase							)
                     # middle
                     __lowerCamelCase     :      str       						=						torch.utils.checkpoint.checkpoint(create_custom_forward(self.mid_block							)       ,  _lowerCamelCase							)

         else:
               # down
               for down_block in self.down_blocks:
                     __lowerCamelCase     :      List[Any]       						=						down_block(_lowerCamelCase							)

               # middle
               __lowerCamelCase     :      int       						=						self.mid_block(_lowerCamelCase							)

         # post-process
         __lowerCamelCase     :      str       						=						self.conv_norm_out(_lowerCamelCase							)
         __lowerCamelCase     :      int       						=						self.conv_act(_lowerCamelCase							)
         __lowerCamelCase     :      Dict       						=						self.conv_out(_lowerCamelCase							)

         return sample







class 							_UpperCamelCase (						nn.Module						):

   '''simple docstring'''







   def __init__(  self  : Optional[Any]       ,  _lowerCamelCase  : Optional[int]=3       ,  _lowerCamelCase  : List[str]=3       ,  _lowerCamelCase  : List[Any]=("UpDecoderBlock2D",)       ,  _lowerCamelCase  : str=(6_4,)       ,  _lowerCamelCase  : str=2       ,  _lowerCamelCase  : int=3_2       ,  _lowerCamelCase  : Union[str, Any]="silu"       ,  _lowerCamelCase  : Dict="group"       ,  ):






         '''simple docstring'''



         super().__init__()
         __lowerCamelCase     :      Optional[int]       						=						layers_per_block

         __lowerCamelCase     :      int       						=						nn.Convad(
             _lowerCamelCase       ,  block_out_channels[-1]       ,  kernel_size=3       ,  stride=1       ,  padding=1       ,  )

         __lowerCamelCase     :      Tuple       						=						None
         __lowerCamelCase     :      Optional[Any]       						=						nn.ModuleList([]							)

         __lowerCamelCase     :      Any       						=						in_channels if norm_type == """spatial""" else None

         # mid
         __lowerCamelCase     :      Any       						=						UNetMidBlockaD(
             in_channels=block_out_channels[-1]       ,  resnet_eps=1E-6       ,  resnet_act_fn=_lowerCamelCase       ,  output_scale_factor=1       ,  resnet_time_scale_shift="""default""" if norm_type == """group""" else norm_type       ,  attention_head_dim=block_out_channels[-1]       ,  resnet_groups=_lowerCamelCase       ,  temb_channels=_lowerCamelCase       ,  )

         # up
         __lowerCamelCase     :      List[Any]       						=						list(reversed(_lowerCamelCase							)							)
         __lowerCamelCase     :      Tuple       						=						reversed_block_out_channels[0]
         for i, up_block_type in enumerate(_lowerCamelCase							):
               __lowerCamelCase     :      Dict       						=						output_channel
               __lowerCamelCase     :      str       						=						reversed_block_out_channels[i]

               __lowerCamelCase     :      Tuple       						=						i == len(_lowerCamelCase							) - 1

               __lowerCamelCase     :      Tuple       						=						get_up_block(
                   _lowerCamelCase       ,  num_layers=self.layers_per_block + 1       ,  in_channels=_lowerCamelCase       ,  out_channels=_lowerCamelCase       ,  prev_output_channel=_lowerCamelCase       ,  add_upsample=not is_final_block       ,  resnet_eps=1E-6       ,  resnet_act_fn=_lowerCamelCase       ,  resnet_groups=_lowerCamelCase       ,  attention_head_dim=_lowerCamelCase       ,  temb_channels=_lowerCamelCase       ,  resnet_time_scale_shift=_lowerCamelCase       ,  )
               self.up_blocks.append(_lowerCamelCase							)
               __lowerCamelCase     :      Tuple       						=						output_channel

         # out
         if norm_type == "spatial":
               __lowerCamelCase     :      Union[str, Any]       						=						SpatialNorm(block_out_channels[0]       ,  _lowerCamelCase							)
         else:
               __lowerCamelCase     :      int       						=						nn.GroupNorm(num_channels=block_out_channels[0]       ,  num_groups=_lowerCamelCase       ,  eps=1E-6							)
         __lowerCamelCase     :      int       						=						nn.SiLU()
         __lowerCamelCase     :      Dict       						=						nn.Convad(block_out_channels[0]       ,  _lowerCamelCase       ,  3       ,  padding=1							)

         __lowerCamelCase     :      int       						=						False


   def   _snake_case    (  self  : List[str]       ,  _lowerCamelCase  : Union[str, Any]       ,  _lowerCamelCase  : Tuple=None							):






         '''simple docstring'''



         __lowerCamelCase     :      List[Any]       						=						z
         __lowerCamelCase     :      Union[str, Any]       						=						self.conv_in(_lowerCamelCase							)

         __lowerCamelCase     :      Dict       						=						next(iter(self.up_blocks.parameters()							)							).dtype
         if self.training and self.gradient_checkpointing:

               def create_custom_forward(_lowerCamelCase  : int							):
                     def custom_forward(*_lowerCamelCase  : Any							):
                           return module(*_lowerCamelCase							)

                     return custom_forward

               if is_torch_version(""">="""       ,  """1.11.0"""							):
                     # middle
                     __lowerCamelCase     :      Dict       						=						torch.utils.checkpoint.checkpoint(
                         create_custom_forward(self.mid_block							)       ,  _lowerCamelCase       ,  _lowerCamelCase       ,  use_reentrant=_lowerCamelCase							)
                     __lowerCamelCase     :      List[str]       						=						sample.to(_lowerCamelCase							)

                     # up
                     for up_block in self.up_blocks:
                           __lowerCamelCase     :      List[str]       						=						torch.utils.checkpoint.checkpoint(
                               create_custom_forward(_lowerCamelCase							)       ,  _lowerCamelCase       ,  _lowerCamelCase       ,  use_reentrant=_lowerCamelCase							)
               else:
                     # middle
                     __lowerCamelCase     :      int       						=						torch.utils.checkpoint.checkpoint(
                         create_custom_forward(self.mid_block							)       ,  _lowerCamelCase       ,  _lowerCamelCase							)
                     __lowerCamelCase     :      Optional[Any]       						=						sample.to(_lowerCamelCase							)

                     # up
                     for up_block in self.up_blocks:
                           __lowerCamelCase     :      Optional[int]       						=						torch.utils.checkpoint.checkpoint(create_custom_forward(_lowerCamelCase							)       ,  _lowerCamelCase       ,  _lowerCamelCase							)
         else:
               # middle
               __lowerCamelCase     :      Dict       						=						self.mid_block(_lowerCamelCase       ,  _lowerCamelCase							)
               __lowerCamelCase     :      Optional[int]       						=						sample.to(_lowerCamelCase							)

               # up
               for up_block in self.up_blocks:
                     __lowerCamelCase     :      Any       						=						up_block(_lowerCamelCase       ,  _lowerCamelCase							)

        # post-process
         if latent_embeds is None:
               __lowerCamelCase     :      Any       						=						self.conv_norm_out(_lowerCamelCase							)
         else:
               __lowerCamelCase     :      Any       						=						self.conv_norm_out(_lowerCamelCase       ,  _lowerCamelCase							)
         __lowerCamelCase     :      List[str]       						=						self.conv_act(_lowerCamelCase							)
         __lowerCamelCase     :      str       						=						self.conv_out(_lowerCamelCase							)

         return sample







class 							_UpperCamelCase (						nn.Module						):

   '''simple docstring'''







   def __init__(  self  : List[str]       ,  _lowerCamelCase  : int       ,  _lowerCamelCase  : Union[str, Any]       ,  _lowerCamelCase  : Optional[Any]       ,  _lowerCamelCase  : Optional[int]=None       ,  _lowerCamelCase  : int="random"       ,  _lowerCamelCase  : List[Any]=False       ,  _lowerCamelCase  : Union[str, Any]=True							):






         '''simple docstring'''



         super().__init__()
         __lowerCamelCase     :      str       						=						n_e
         __lowerCamelCase     :      Optional[Any]       						=						vq_embed_dim
         __lowerCamelCase     :      Tuple       						=						beta
         __lowerCamelCase     :      int       						=						legacy

         __lowerCamelCase     :      Tuple       						=						nn.Embedding(self.n_e       ,  self.vq_embed_dim							)
         self.embedding.weight.data.uniform_(-1.0 / self.n_e       ,  1.0 / self.n_e							)

         __lowerCamelCase     :      Optional[Any]       						=						remap
         if self.remap is not None:
               self.register_buffer("""used"""       ,  torch.tensor(np.load(self.remap							)							)							)
               __lowerCamelCase     :      Union[str, Any]       						=						self.used.shape[0]
               __lowerCamelCase     :      Tuple       						=						unknown_index  # "random" or "extra" or integer
               if self.unknown_index == "extra":
                     __lowerCamelCase     :      int       						=						self.re_embed
                     __lowerCamelCase     :      Tuple       						=						self.re_embed + 1
               print(
                   F"""Remapping {self.n_e} indices to {self.re_embed} indices. """
                   F"""Using {self.unknown_index} for unknown indices."""							)
         else:
               __lowerCamelCase     :      Dict       						=						n_e

         __lowerCamelCase     :      Union[str, Any]       						=						sane_index_shape

   def   _snake_case    (  self  : Tuple       ,  _lowerCamelCase  : int							):






         '''simple docstring'''



         __lowerCamelCase     :      Optional[int]       						=						inds.shape
         assert len(_lowerCamelCase							) > 1
         __lowerCamelCase     :      Optional[Any]       						=						inds.reshape(ishape[0]       ,  -1							)
         __lowerCamelCase     :      List[Any]       						=						self.used.to(_lowerCamelCase							)
         __lowerCamelCase     :      Union[str, Any]       						=						(inds[:, :, None] == used[None, None, ...]).long()
         __lowerCamelCase     :      Tuple       						=						match.argmax(-1							)
         __lowerCamelCase     :      Union[str, Any]       						=						match.sum(2							) < 1
         if self.unknown_index == "random":
               __lowerCamelCase     :      List[Any]       						=						torch.randint(0       ,  self.re_embed       ,  size=new[unknown].shape							).to(device=new.device							)
         else:
               __lowerCamelCase     :      Any       						=						self.unknown_index
         return new.reshape(_lowerCamelCase							)

   def   _snake_case    (  self  : List[Any]       ,  _lowerCamelCase  : Union[str, Any]							):






         '''simple docstring'''



         __lowerCamelCase     :      str       						=						inds.shape
         assert len(_lowerCamelCase							) > 1
         __lowerCamelCase     :      str       						=						inds.reshape(ishape[0]       ,  -1							)
         __lowerCamelCase     :      List[str]       						=						self.used.to(_lowerCamelCase							)
         if self.re_embed > self.used.shape[0]:  # extra token
               __lowerCamelCase     :      Dict       						=						0  # simply set to zero
         __lowerCamelCase     :      Union[str, Any]       						=						torch.gather(used[None, :][inds.shape[0] * [0], :]       ,  1       ,  _lowerCamelCase							)
         return back.reshape(_lowerCamelCase							)

   def   _snake_case    (  self  : Any       ,  _lowerCamelCase  : Optional[Any]							):






         '''simple docstring'''



         __lowerCamelCase     :      Any       						=						z.permute(0       ,  2       ,  3       ,  1							).contiguous()
         __lowerCamelCase     :      Union[str, Any]       						=						z.view(-1       ,  self.vq_embed_dim							)

         # distances from z to embeddings e_j (z - e)^2 = z^2 + e^2 - 2 e * z
         __lowerCamelCase     :      Union[str, Any]       						=						torch.argmin(torch.cdist(_lowerCamelCase       ,  self.embedding.weight							)       ,  dim=1							)

         __lowerCamelCase     :      List[str]       						=						self.embedding(_lowerCamelCase							).view(z.shape							)
         __lowerCamelCase     :      Optional[Any]       						=						None
         __lowerCamelCase     :      Any       						=						None

         # compute loss for embedding
         if not self.legacy:
               __lowerCamelCase     :      List[str]       						=						self.beta * torch.mean((z_q.detach() - z) ** 2							) + torch.mean((z_q - z.detach()) ** 2							)
         else:
               __lowerCamelCase     :      Tuple       						=						torch.mean((z_q.detach() - z) ** 2							) + self.beta * torch.mean((z_q - z.detach()) ** 2							)

         # preserve gradients
         __lowerCamelCase     :      str       						=						z + (z_q - z).detach()

         # reshape back to match original input shape
         __lowerCamelCase     :      Dict       						=						z_q.permute(0       ,  3       ,  1       ,  2							).contiguous()

         if self.remap is not None:
               __lowerCamelCase     :      Optional[Any]       						=						min_encoding_indices.reshape(z.shape[0]       ,  -1							)  # add batch axis
               __lowerCamelCase     :      List[str]       						=						self.remap_to_used(_lowerCamelCase							)
               __lowerCamelCase     :      int       						=						min_encoding_indices.reshape(-1       ,  1							)  # flatten

         if self.sane_index_shape:
               __lowerCamelCase     :      Dict       						=						min_encoding_indices.reshape(z_q.shape[0]       ,  z_q.shape[2]       ,  z_q.shape[3]							)

         return z_q, loss, (perplexity, min_encodings, min_encoding_indices)


   def   _snake_case    (  self  : List[Any]       ,  _lowerCamelCase  : Optional[int]       ,  _lowerCamelCase  : Optional[Any]							):






         '''simple docstring'''



         if self.remap is not None:
               __lowerCamelCase     :      List[str]       						=						indices.reshape(shape[0]       ,  -1							)  # add batch axis
               __lowerCamelCase     :      str       						=						self.unmap_to_all(_lowerCamelCase							)
               __lowerCamelCase     :      List[str]       						=						indices.reshape(-1							)  # flatten again

         # get quantized latent vectors
         __lowerCamelCase     :      int       						=						self.embedding(_lowerCamelCase							)

         if shape is not None:
               __lowerCamelCase     :      Dict       						=						z_q.view(_lowerCamelCase							)
               # reshape back to match original input shape
               __lowerCamelCase     :      str       						=						z_q.permute(0       ,  3       ,  1       ,  2							).contiguous()

         return z_q







class 							_UpperCamelCase (						A						):

   '''simple docstring'''







   def __init__(  self  : Optional[Any]       ,  _lowerCamelCase  : str       ,  _lowerCamelCase  : Optional[int]=False							):






         '''simple docstring'''



         __lowerCamelCase     :      Tuple       						=						parameters
         __lowerCamelCase     :      Union[str, Any]       						=						torch.chunk(_lowerCamelCase       ,  2       ,  dim=1							)
         __lowerCamelCase     :      Dict       						=						torch.clamp(self.logvar       ,  -30.0       ,  20.0							)
         __lowerCamelCase     :      Dict       						=						deterministic
         __lowerCamelCase     :      List[Any]       						=						torch.exp(0.5 * self.logvar							)
         __lowerCamelCase     :      Optional[int]       						=						torch.exp(self.logvar							)
         if self.deterministic:
               __lowerCamelCase     :      Optional[Any]       						=						torch.zeros_like(
                   self.mean       ,  device=self.parameters.device       ,  dtype=self.parameters.dtype							)

   def   _snake_case    (  self  : List[Any]       ,  _lowerCamelCase  : Optional[torch.Generator] = None							):






         '''simple docstring'''



         __lowerCamelCase     :      str       						=						randn_tensor(
             self.mean.shape       ,  generator=_lowerCamelCase       ,  device=self.parameters.device       ,  dtype=self.parameters.dtype							)
         __lowerCamelCase     :      Union[str, Any]       						=						self.mean + self.std * sample
         return x

   def   _snake_case    (  self  : List[str]       ,  _lowerCamelCase  : List[str]=None							):






         '''simple docstring'''



         if self.deterministic:
               return torch.Tensor([0.0]							)
         else:
               if other is None:
                     return 0.5 * torch.sum(torch.pow(self.mean       ,  2							) + self.var - 1.0 - self.logvar       ,  dim=[1, 2, 3]							)
               else:
                     return 0.5 * torch.sum(
                         torch.pow(self.mean - other.mean       ,  2							) / other.var
                         + self.var / other.var
                         - 1.0
                         - self.logvar
                         + other.logvar       ,  dim=[1, 2, 3]       ,  )

   def   _snake_case    (  self  : str       ,  _lowerCamelCase  : Optional[int]       ,  _lowerCamelCase  : Optional[Any]=[1, 2, 3]							):






         '''simple docstring'''



         if self.deterministic:
               return torch.Tensor([0.0]							)
         __lowerCamelCase     :      Any       						=						np.log(2.0 * np.pi							)
         return 0.5 * torch.sum(logtwopi + self.logvar + torch.pow(sample - self.mean       ,  2							) / self.var       ,  dim=_lowerCamelCase							)


   def   _snake_case    (  self  : Optional[Any]							):






         '''simple docstring'''



         return self.mean







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

   '''simple docstring'''







   a_				:							Any    = "ssube/stable-diffusion-x4-upscaler-onnx"

   def   _snake_case    (  self  : Any       ,  _lowerCamelCase  : List[str]=0							):






         '''simple docstring'''



         __lowerCamelCase     :      Union[str, Any]       						=						floats_tensor((1, 3, 1_2_8, 1_2_8)       ,  rng=random.Random(_lowerCamelCase							)							)
         __lowerCamelCase     :      int       						=						torch.manual_seed(_lowerCamelCase							)
         __lowerCamelCase     :      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   _snake_case    (  self  : Optional[Any]							):






         '''simple docstring'''



         __lowerCamelCase     :      Union[str, Any]       						=						OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint       ,  provider="""CPUExecutionProvider"""							)
         pipe.set_progress_bar_config(disable=_lowerCamelCase							)

         __lowerCamelCase     :      Tuple       						=						self.get_dummy_inputs()
         __lowerCamelCase     :      Dict       						=						pipe(**_lowerCamelCase							).images
         __lowerCamelCase     :      Dict       						=						image[0, -3:, -3:, -1].flatten()

         # started as 128, should now be 512
         assert image.shape == (1, 5_1_2, 5_1_2, 3)
         __lowerCamelCase     :      str       						=						np.array(
             [0.6_974_782, 0.68_902_093, 0.70_135_885, 0.7_583_618, 0.7_804_545, 0.7_854_912, 0.78_667_426, 0.78_743_863, 0.78_070_223]							)
         assert np.abs(image_slice - expected_slice							).max() < 1E-1

   def   _snake_case    (  self  : Dict							):






         '''simple docstring'''



         __lowerCamelCase     :      Optional[Any]       						=						OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint       ,  provider="""CPUExecutionProvider"""							)
         __lowerCamelCase     :      str       						=						PNDMScheduler.from_config(pipe.scheduler.config       ,  skip_prk_steps=_lowerCamelCase							)
         pipe.set_progress_bar_config(disable=_lowerCamelCase							)

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

         assert image.shape == (1, 5_1_2, 5_1_2, 3)
         __lowerCamelCase     :      int       						=						np.array(
             [0.6_898_892, 0.59_240_556, 0.52_499_527, 0.58_866_215, 0.52_258_235, 0.52_572_715, 0.62_414_473, 0.6_174_387, 0.6_214_964]							)
         assert np.abs(image_slice.flatten() - expected_slice							).max() < 1E-1

   def   _snake_case    (  self  : Tuple							):






         '''simple docstring'''



         __lowerCamelCase     :      Optional[int]       						=						OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint       ,  provider="""CPUExecutionProvider"""							)
         __lowerCamelCase     :      int       						=						DPMSolverMultistepScheduler.from_config(pipe.scheduler.config							)
         pipe.set_progress_bar_config(disable=_lowerCamelCase							)

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

         assert image.shape == (1, 5_1_2, 5_1_2, 3)
         __lowerCamelCase     :      str       						=						np.array(
             [0.7_659_278, 0.76_437_664, 0.75_579_107, 0.7_691_116, 0.77_666_986, 0.7_727_672, 0.7_758_664, 0.7_812_226, 0.76_942_515]							)

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

   def   _snake_case    (  self  : str							):






         '''simple docstring'''



         __lowerCamelCase     :      Optional[int]       						=						OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint       ,  provider="""CPUExecutionProvider"""							)
         __lowerCamelCase     :      Optional[Any]       						=						EulerDiscreteScheduler.from_config(pipe.scheduler.config							)
         pipe.set_progress_bar_config(disable=_lowerCamelCase							)

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

         assert image.shape == (1, 5_1_2, 5_1_2, 3)
         __lowerCamelCase     :      List[str]       						=						np.array(
             [0.6_974_782, 0.68_902_093, 0.70_135_885, 0.7_583_618, 0.7_804_545, 0.7_854_912, 0.78_667_426, 0.78_743_863, 0.78_070_223]							)
         assert np.abs(image_slice.flatten() - expected_slice							).max() < 1E-1


   def   _snake_case    (  self  : Any							):






         '''simple docstring'''



         __lowerCamelCase     :      List[str]       						=						OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint       ,  provider="""CPUExecutionProvider"""							)
         __lowerCamelCase     :      Union[str, Any]       						=						EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config							)
         pipe.set_progress_bar_config(disable=_lowerCamelCase							)

         __lowerCamelCase     :      Optional[int]       						=						self.get_dummy_inputs()
         __lowerCamelCase     :      int       						=						pipe(**_lowerCamelCase							).images
         __lowerCamelCase     :      int       						=						image[0, -3:, -3:, -1]

         assert image.shape == (1, 5_1_2, 5_1_2, 3)
         __lowerCamelCase     :      Optional[int]       						=						np.array(
             [0.77_424_496, 0.773_601, 0.7_645_288, 0.7_769_598, 0.7_772_739, 0.7_738_688, 0.78_187_233, 0.77_879_584, 0.767_043]							)

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



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

   '''simple docstring'''







   @property
   def   _snake_case    (  self  : str							):






         '''simple docstring'''



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

   @property
   def   _snake_case    (  self  : Tuple							):






         '''simple docstring'''



         __lowerCamelCase     :      Any       						=						ort.SessionOptions()
         __lowerCamelCase     :      str       						=						False
         return options

   def   _snake_case    (  self  : Any							):






         '''simple docstring'''



         __lowerCamelCase     :      str       						=						load_image(
             """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"""
             """/img2img/sketch-mountains-input.jpg"""							)
         __lowerCamelCase     :      Optional[Any]       						=						init_image.resize((1_2_8, 1_2_8)							)
         # using the PNDM scheduler by default
         __lowerCamelCase     :      List[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=_lowerCamelCase							)

         __lowerCamelCase     :      List[Any]       						=						"""A fantasy landscape, trending on artstation"""

         __lowerCamelCase     :      str       						=						torch.manual_seed(0							)
         __lowerCamelCase     :      Any       						=						pipe(
             prompt=_lowerCamelCase       ,  image=_lowerCamelCase       ,  guidance_scale=7.5       ,  num_inference_steps=1_0       ,  generator=_lowerCamelCase       ,  output_type="""np"""       ,  )
         __lowerCamelCase     :      List[str]       						=						output.images
         __lowerCamelCase     :      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)
         __lowerCamelCase     :      int       						=						np.array([0.4_883, 0.4_947, 0.4_980, 0.4_975, 0.4_982, 0.4_980, 0.5_000, 0.5_006, 0.4_972]							)
         # 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  : Any							):






         '''simple docstring'''



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

         __lowerCamelCase     :      str       						=						"""A fantasy landscape, trending on artstation"""

         __lowerCamelCase     :      Tuple       						=						torch.manual_seed(0							)
         __lowerCamelCase     :      Dict       						=						pipe(
             prompt=_lowerCamelCase       ,  image=_lowerCamelCase       ,  guidance_scale=7.5       ,  num_inference_steps=2_0       ,  generator=_lowerCamelCase       ,  output_type="""np"""       ,  )
         __lowerCamelCase     :      int       						=						output.images
         __lowerCamelCase     :      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)
         __lowerCamelCase     :      List[str]       						=						np.array(
             [0.50_173_753, 0.50_223_356, 0.502_039, 0.50_233_036, 0.5_023_725, 0.5_022_601, 0.5_018_758, 0.50_234_085, 0.50_241_566]							)
         # TODO: lower the tolerance after finding the cause of onnxruntime reproducibility issues

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



