






"""simple docstring"""






import tensorflow as tf

from ...tf_utils import shape_list




class 							lowerCAmelCase__				(			tf.keras.layers.Layer ):






      '''simple docstring'''

      def __init__(						self     :  Optional[int]						,	lowercase_     :  Tuple						,	lowercase_     :  int						,	lowercase_     :  List[Any]						,	lowercase_     :  int						,	lowercase_     :  Optional[int]=1						,	lowercase_     :  List[Any]=False						,	**lowercase_     :  List[Any]):




          '''simple docstring'''


          super().__init__(**lowercase_)

          SCREAMING_SNAKE_CASE_      :  List[str]    						=		vocab_size
          SCREAMING_SNAKE_CASE_      :  Any    						=		d_embed
          SCREAMING_SNAKE_CASE_      :  int    						=		d_proj

          SCREAMING_SNAKE_CASE_      :  List[Any]    						=		cutoffs + [vocab_size]
          SCREAMING_SNAKE_CASE_      :  Dict    						=		[0] + self.cutoffs
          SCREAMING_SNAKE_CASE_      :  str    						=		div_val

          SCREAMING_SNAKE_CASE_      :  List[Any]    						=		self.cutoffs[0]
          SCREAMING_SNAKE_CASE_      :  List[Any]    						=		len(self.cutoffs) - 1
          SCREAMING_SNAKE_CASE_      :  Dict    						=		self.shortlist_size + self.n_clusters
          SCREAMING_SNAKE_CASE_      :  Optional[int]    						=		keep_order

          SCREAMING_SNAKE_CASE_      :  List[Any]    						=		[]
          SCREAMING_SNAKE_CASE_      :  Optional[int]    						=		[]
      def _SCREAMING_SNAKE_CASE (						self     :  Optional[int]						,	lowercase_     :  List[str]):




          '''simple docstring'''


          if self.n_clusters > 0:
              SCREAMING_SNAKE_CASE_      :  Any    						=		self.add_weight(
                  shape=(self.n_clusters, self.d_embed)						,	initializer='''zeros'''						,	trainable=lowercase_						,	name='''cluster_weight''')
              SCREAMING_SNAKE_CASE_      :  Union[str, Any]    						=		self.add_weight(
                  shape=(self.n_clusters,)						,	initializer='''zeros'''						,	trainable=lowercase_						,	name='''cluster_bias''')

          if self.div_val == 1:
              for i in range(len(self.cutoffs)):
                  if self.d_proj != self.d_embed:
                      SCREAMING_SNAKE_CASE_      :  Union[str, Any]    						=		self.add_weight(
                          shape=(self.d_embed, self.d_proj)						,	initializer='''zeros'''						,	trainable=lowercase_						,	name=F'out_projs_._{i}'						,	)
                      self.out_projs.append(lowercase_)
                  else:
                      self.out_projs.append(lowercase_)
                  SCREAMING_SNAKE_CASE_      :  Tuple    						=		self.add_weight(
                      shape=(self.vocab_size, self.d_embed)						,	initializer='''zeros'''						,	trainable=lowercase_						,	name=F'out_layers_._{i}_._weight'						,	)
                  SCREAMING_SNAKE_CASE_      :  List[Any]    						=		self.add_weight(
                      shape=(self.vocab_size,)						,	initializer='''zeros'''						,	trainable=lowercase_						,	name=F'out_layers_._{i}_._bias'						,	)
                  self.out_layers.append((weight, bias))
          else:
              for i in range(len(self.cutoffs)):
                  SCREAMING_SNAKE_CASE_						,			SCREAMING_SNAKE_CASE_      :  int    						=		self.cutoff_ends[i], self.cutoff_ends[i + 1]
                  SCREAMING_SNAKE_CASE_      :  Union[str, Any]    						=		self.d_embed // (self.div_val**i)

                  SCREAMING_SNAKE_CASE_      :  List[Any]    						=		self.add_weight(
                      shape=(d_emb_i, self.d_proj)						,	initializer='''zeros'''						,	trainable=lowercase_						,	name=F'out_projs_._{i}')
                  self.out_projs.append(lowercase_)
                  SCREAMING_SNAKE_CASE_      :  Optional[int]    						=		self.add_weight(
                      shape=(r_idx - l_idx, d_emb_i)						,	initializer='''zeros'''						,	trainable=lowercase_						,	name=F'out_layers_._{i}_._weight'						,	)
                  SCREAMING_SNAKE_CASE_      :  Union[str, Any]    						=		self.add_weight(
                      shape=(r_idx - l_idx,)						,	initializer='''zeros'''						,	trainable=lowercase_						,	name=F'out_layers_._{i}_._bias'						,	)
                  self.out_layers.append((weight, bias))
          super().build(lowercase_)
      @staticmethod
      def _SCREAMING_SNAKE_CASE (						lowercase_     :  Tuple						,	lowercase_     :  Dict						,	lowercase_     :  Optional[int]						,	lowercase_     :  List[str]=None):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_      :  Union[str, Any]    						=		x
          if proj is not None:
              SCREAMING_SNAKE_CASE_      :  Optional[int]    						=		tf.einsum('''ibd,ed->ibe'''						,	lowercase_						,	lowercase_)
          return tf.einsum('''ibd,nd->ibn'''						,	lowercase_						,	lowercase_) + b
      @staticmethod
      def _SCREAMING_SNAKE_CASE (						lowercase_     :  Optional[Any]						,	lowercase_     :  Optional[Any]):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_      :  List[Any]    						=		shape_list(lowercase_)
          SCREAMING_SNAKE_CASE_      :  int    						=		tf.range(lp_size[0]						,	dtype=target.dtype)
          SCREAMING_SNAKE_CASE_      :  Dict    						=		tf.stack([r, target]						,	1)
          return tf.gather_nd(lowercase_						,	lowercase_)



      def _SCREAMING_SNAKE_CASE (						self     :  Tuple						,	lowercase_     :  Optional[Any]						,	lowercase_     :  int						,	lowercase_     :  Dict=True						,	lowercase_     :  str=False):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_      :  Dict    						=		0
          if self.n_clusters == 0:
              SCREAMING_SNAKE_CASE_      :  Optional[Any]    						=		self._logit(lowercase_						,	self.out_layers[0][0]						,	self.out_layers[0][1]						,	self.out_projs[0])
              if target is not None:
                  SCREAMING_SNAKE_CASE_      :  Tuple    						=		tf.nn.sparse_softmax_cross_entropy_with_logits(labels=lowercase_						,	logits=lowercase_)
              SCREAMING_SNAKE_CASE_      :  Dict    						=		tf.nn.log_softmax(lowercase_						,	axis=-1)
          else:
              SCREAMING_SNAKE_CASE_      :  Tuple    						=		shape_list(lowercase_)
              SCREAMING_SNAKE_CASE_      :  Union[str, Any]    						=		[]
              SCREAMING_SNAKE_CASE_      :  str    						=		tf.zeros(hidden_sizes[:2])
              for i in range(len(self.cutoffs)):
                  SCREAMING_SNAKE_CASE_						,			SCREAMING_SNAKE_CASE_      :  Optional[int]    						=		self.cutoff_ends[i], self.cutoff_ends[i + 1]
                  if target is not None:
                      SCREAMING_SNAKE_CASE_      :  Union[str, Any]    						=		(target >= l_idx) & (target < r_idx)
                      SCREAMING_SNAKE_CASE_      :  Any    						=		tf.where(lowercase_)
                      SCREAMING_SNAKE_CASE_      :  Optional[Any]    						=		tf.boolean_mask(lowercase_						,	lowercase_) - l_idx

                  if self.div_val == 1:
                      SCREAMING_SNAKE_CASE_      :  Optional[Any]    						=		self.out_layers[0][0][l_idx:r_idx]
                      SCREAMING_SNAKE_CASE_      :  Any    						=		self.out_layers[0][1][l_idx:r_idx]
                  else:
                      SCREAMING_SNAKE_CASE_      :  str    						=		self.out_layers[i][0]
                      SCREAMING_SNAKE_CASE_      :  List[Any]    						=		self.out_layers[i][1]

                  if i == 0:
                      SCREAMING_SNAKE_CASE_      :  Union[str, Any]    						=		tf.concat([cur_W, self.cluster_weight]						,	0)
                      SCREAMING_SNAKE_CASE_      :  List[Any]    						=		tf.concat([cur_b, self.cluster_bias]						,	0)

                      SCREAMING_SNAKE_CASE_      :  Optional[Any]    						=		self._logit(lowercase_						,	lowercase_						,	lowercase_						,	self.out_projs[0])
                      SCREAMING_SNAKE_CASE_      :  Union[str, Any]    						=		tf.nn.log_softmax(lowercase_)
                      out.append(head_logprob[..., : self.cutoffs[0]])
                      if target is not None:
                          SCREAMING_SNAKE_CASE_      :  Any    						=		tf.boolean_mask(lowercase_						,	lowercase_)
                          SCREAMING_SNAKE_CASE_      :  List[Any]    						=		self._gather_logprob(lowercase_						,	lowercase_)
                  else:
                      SCREAMING_SNAKE_CASE_      :  List[Any]    						=		self._logit(lowercase_						,	lowercase_						,	lowercase_						,	self.out_projs[i])
                      SCREAMING_SNAKE_CASE_      :  Optional[Any]    						=		tf.nn.log_softmax(lowercase_)
                      SCREAMING_SNAKE_CASE_      :  Dict    						=		self.cutoffs[0] + i - 1  # No probability for the head cluster
                      SCREAMING_SNAKE_CASE_      :  Any    						=		head_logprob[..., cluster_prob_idx, None] + tail_logprob
                      out.append(lowercase_)
                      if target is not None:
                          SCREAMING_SNAKE_CASE_      :  List[str]    						=		tf.boolean_mask(lowercase_						,	lowercase_)
                          SCREAMING_SNAKE_CASE_      :  int    						=		tf.boolean_mask(lowercase_						,	lowercase_)
                          SCREAMING_SNAKE_CASE_      :  int    						=		self._gather_logprob(lowercase_						,	lowercase_)
                          cur_logprob += cur_head_logprob[:, self.cutoff_ends[1] + i - 1]
                  if target is not None:
                      loss += tf.scatter_nd(lowercase_						,	-cur_logprob						,	shape_list(lowercase_))
              SCREAMING_SNAKE_CASE_      :  List[str]    						=		tf.concat(lowercase_						,	axis=-1)

          if target is not None:
              if return_mean:
                  SCREAMING_SNAKE_CASE_      :  Optional[int]    						=		tf.reduce_mean(lowercase_)
              # Add the training-time loss value to the layer using `self.add_loss()`.
              self.add_loss(lowercase_)

              # Log the loss as a metric (we could log arbitrary metrics,
              # including different metrics for training and inference.
              self.add_metric(lowercase_						,	name=self.name						,	aggregation='''mean''' if return_mean else '''''')

          return out






"""simple docstring"""



import random
import unittest

import numpy as np
import torch

from diffusers import (
    DPMSolverMultistepScheduler,
    EulerAncestralDiscreteScheduler,
    EulerDiscreteScheduler,
    LMSDiscreteScheduler,
    OnnxStableDiffusionUpscalePipeline,
    PNDMScheduler,
)
from diffusers.utils import floats_tensor
from diffusers.utils.testing_utils import (
    is_onnx_available,
    load_image,
    nightly,
    require_onnxruntime,
    require_torch_gpu,
)

from ..test_pipelines_onnx_common import OnnxPipelineTesterMixin


if is_onnx_available():
  import onnxruntime as ort







class _lowerCAmelCase				(      snake_case_    ,							unittest.TestCase	):
   # TODO: is there an appropriate internal test set?
   __UpperCAmelCase :       Optional[Any]							  =  '''ssube/stable-diffusion-x4-upscaler-onnx'''







   def  lowerCamelCase		(	self						,	UpperCamelCase__=0						)	->		Tuple:

       '''simple docstring'''




       snake_case    :      Union[str, Any]      					=    floats_tensor((1, 3, 128, 128)						,	rng=random.Random(UpperCamelCase__						)						)
       snake_case    :      Any      					=    torch.manual_seed(UpperCamelCase__						)
       snake_case    :      List[str]      					=    {
           "prompt": "A painting of a squirrel eating a burger",
           "image": image,
           "generator": generator,
           "num_inference_steps": 3,
           "guidance_scale": 7.5,
           "output_type": "numpy",
       }
       return inputs







   def  lowerCamelCase		(	self						)	->		Tuple:

       '''simple docstring'''




       snake_case    :      Optional[Any]      					=    OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint						,	provider="CPUExecutionProvider"						)
       pipe.set_progress_bar_config(disable=UpperCamelCase__						)

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

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







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

       '''simple docstring'''




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

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

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







   def  lowerCamelCase		(	self						)	->		Any:

       '''simple docstring'''




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

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

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

       '''simple docstring'''




       snake_case    :      Union[str, Any]      					=    OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint						,	provider="CPUExecutionProvider"						)
       snake_case    :      Any      					=    EulerDiscreteScheduler.from_config(pipe.scheduler.config						)
       pipe.set_progress_bar_config(disable=UpperCamelCase__						)

       snake_case    :      Dict      					=    self.get_dummy_inputs()
       snake_case    :      Union[str, Any]      					=    pipe(**UpperCamelCase__						).images
       snake_case    :      str      					=    image[0, -3:, -3:, -1]

       assert image.shape == (1, 512, 512, 3)
       snake_case    :      Union[str, 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.flatten() - expected_slice						).max() < 1e-1







   def  lowerCamelCase		(	self						)	->		str:

       '''simple docstring'''




       snake_case    :      Optional[Any]      					=    OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint						,	provider="CPUExecutionProvider"						)
       snake_case    :      Optional[Any]      					=    EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config						)
       pipe.set_progress_bar_config(disable=UpperCamelCase__						)

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

       assert image.shape == (1, 512, 512, 3)
       snake_case    :      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 _lowerCAmelCase				(      unittest.TestCase	):







   @property
   def  lowerCamelCase		(	self						)	->		Tuple:

       '''simple docstring'''




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







   @property
   def  lowerCamelCase		(	self						)	->		List[Any]:

       '''simple docstring'''




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







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

       '''simple docstring'''




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

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

       snake_case    :      Union[str, Any]      					=    torch.manual_seed(0						)
       snake_case    :      int      					=    pipe(
           prompt=UpperCamelCase__						,	image=UpperCamelCase__						,	guidance_scale=7.5						,	num_inference_steps=10						,	generator=UpperCamelCase__						,	output_type="np"						,	)
       snake_case    :      str      					=    output.images
       snake_case    :      Tuple      					=    images[0, 255:258, 383:386, -1]

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

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







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

       '''simple docstring'''




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

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

       snake_case    :      Tuple      					=    torch.manual_seed(0						)
       snake_case    :      str      					=    pipe(
           prompt=UpperCamelCase__						,	image=UpperCamelCase__						,	guidance_scale=7.5						,	num_inference_steps=20						,	generator=UpperCamelCase__						,	output_type="np"						,	)
       snake_case    :      List[Any]      					=    output.images
       snake_case    :      List[Any]      					=    images[0, 255:258, 383:386, -1]

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

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





