


import os
from shutil import copyfile
from typing import Any, Dict, List, Optional, Tuple

import sentencepiece as spm

from ...tokenization_utils import PreTrainedTokenizer
from ...utils import logging


SCREAMING_SNAKE_CASE       :      Optional[Any]												= "▁"

SCREAMING_SNAKE_CASE       :      Dict												= {"vocab_file": "spiece.model"}

SCREAMING_SNAKE_CASE       :      int												= {
    "vocab_file": {"google/pegasus-xsum": "https://huggingface.co/google/pegasus-xsum/resolve/main/spiece.model"}
}

SCREAMING_SNAKE_CASE       :      Optional[Any]												= {
    "google/pegasus-xsum": 512,
}


SCREAMING_SNAKE_CASE       :      Tuple												= logging.get_logger(__name__)






class _lowerCamelCase(						_a   ):
       lowercase_ :	List[Any]     =			VOCAB_FILES_NAMES

       lowercase_ :	Any     =			VOCAB_FILES_NAMES
       lowercase_ :	List[Any]     =			PRETRAINED_VOCAB_FILES_MAP
       lowercase_ :	List[Any]     =			PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
       lowercase_ :	int     =			["""input_ids""", """attention_mask"""]
       def __init__(					self,      lowerCamelCase,      lowerCamelCase="<pad>",      lowerCamelCase="</s>",      lowerCamelCase="<unk>",      lowerCamelCase="<mask_2>",      lowerCamelCase="<mask_1>",      lowerCamelCase=None,      lowerCamelCase=1_03,      lowerCamelCase = None,      **lowerCamelCase,      )       ->   None:




           """simple docstring"""

           _lowercase       :							Any      		=		offset
           if additional_special_tokens is not None:
               if not isinstance(lowerCamelCase,      lowerCamelCase):
                   raise TypeError(
                       F'''additional_special_tokens should be of type {type(lowerCamelCase)}, but is'''
                       F''' {type(lowerCamelCase)}''')

               _lowercase       :							Optional[Any]      		=		(
                   ([mask_token_sent] + additional_special_tokens)
                   if mask_token_sent not in additional_special_tokens and mask_token_sent is not None
                   else additional_special_tokens
               )
               # fill additional tokens with ..., <unk_token_102> in case not all additional tokens are already taken
               additional_special_tokens_extended += [
                   F'''<unk_{i}>''' for i in range(len(lowerCamelCase),      self.offset - 1)
               ]

               if len(set(lowerCamelCase)) != len(lowerCamelCase):
                   raise ValueError(
                       'Please make sure that the provided additional_special_tokens do not contain an incorrectly'
                       F''' shifted list of <unk_x> tokens. Found {additional_special_tokens_extended}.''')
               _lowercase       :							int      		=		additional_special_tokens_extended
           else:
               _lowercase       :							List[Any]      		=		[mask_token_sent] if mask_token_sent is not None else []
               additional_special_tokens += [F'''<unk_{i}>''' for i in range(2,      self.offset)]

           _lowercase       :							List[Any]      		=		{} if sp_model_kwargs is None else sp_model_kwargs

           super().__init__(
               eos_token=lowerCamelCase,      unk_token=lowerCamelCase,      mask_token=lowerCamelCase,      pad_token=lowerCamelCase,      mask_token_sent=lowerCamelCase,      offset=lowerCamelCase,      additional_special_tokens=lowerCamelCase,      sp_model_kwargs=self.sp_model_kwargs,      **lowerCamelCase,      )
           _lowercase       :							Optional[int]      		=		mask_token_sent
           _lowercase       :							Optional[int]      		=		vocab_file
           _lowercase       :							int      		=		spm.SentencePieceProcessor(**self.sp_model_kwargs)
           self.sp_model.Load(lowerCamelCase)

           # add special tokens to encoder dict
           _lowercase       :							Dict[int, str]      		=		{
               0: self.pad_token,
               1: self.eos_token,
           }

           if self.mask_token_sent is not None:
               self.encoder.update(
                   {
                       2: self.mask_token_sent,
                       3: self.mask_token,
                   })

           if self.offset > 0:
               # entries 2-104 are only used for pretraining and called <mask_1>, <mask_2>, unk_2, ...unk_102
               # mask_token_sent is already added to list -> so start at 1
               self.encoder.update({i + 3: additional_special_tokens[i] for i in range(1,      self.offset - 1)})

           _lowercase       :							Dict[str, int]      		=		{v: k for k, v in self.encoder.items()}
       @property
       def 			UpperCamelCase     (					self)       ->   int:




           """simple docstring"""

           return len(self.sp_model) + self.offset
       def 			UpperCamelCase     (					self)       ->   Dict[str, int]:




           """simple docstring"""

           _lowercase       :							Union[str, Any]      		=		{self.convert_ids_to_tokens(lowerCamelCase): i for i in range(self.vocab_size)}
           vocab.update(self.added_tokens_encoder)
           return vocab
       def __getstate__(					self)       ->   Union[str, Any]:




           """simple docstring"""

           _lowercase       :							List[Any]      		=		self.__dict__.copy()
           _lowercase       :							List[str]      		=		None
           return state
       def __setstate__(					self,      lowerCamelCase)       ->   Dict:




           """simple docstring"""

           _lowercase       :							Union[str, Any]      		=		d

           # for backward compatibility
           if not hasattr(self,      'sp_model_kwargs'):
               _lowercase       :							Optional[Any]      		=		{}

           _lowercase       :							Optional[int]      		=		spm.SentencePieceProcessor(**self.sp_model_kwargs)
           self.sp_model.Load(self.vocab_file)
       def 			UpperCamelCase     (					self,      lowerCamelCase)       ->   List[str]:




           """simple docstring"""

           return self.sp_model.encode(lowerCamelCase,      out_type=lowerCamelCase)
       def 			UpperCamelCase     (					self,      lowerCamelCase)       ->   int:




           """simple docstring"""

           if token in self.decoder:
               return self.decoder[token]
           elif token in self.added_tokens_decoder:
               return self.added_tokens_decoder[token]
           _lowercase       :							str      		=		self.sp_model.piece_to_id(lowerCamelCase)
           return sp_id + self.offset
       def 			UpperCamelCase     (					self,      lowerCamelCase)       ->   str:




           """simple docstring"""

           if index in self.encoder:
               return self.encoder[index]
           elif index in self.added_tokens_encoder:
               return self.added_tokens_encoder[index]
           else:
               _lowercase       :							Optional[int]      		=		self.sp_model.IdToPiece(index - self.offset)
           return token
       def 			UpperCamelCase     (					self,      lowerCamelCase)       ->   Any:




           """simple docstring"""

           _lowercase       :							List[str]      		=		[]
           _lowercase       :							Union[str, Any]      		=		''
           for token in tokens:
               # make sure that special tokens are not decoded using sentencepiece model
               if token in self.all_special_tokens:
                   out_string += self.sp_model.decode(lowerCamelCase) + token
                   _lowercase       :							str      		=		[]
               else:
                   current_sub_tokens.append(lowerCamelCase)
           out_string += self.sp_model.decode(lowerCamelCase)
           return out_string.strip()
       def 			UpperCamelCase     (					self,      lowerCamelCase=False)       ->   Dict:




           """simple docstring"""

           return 1
       def 			UpperCamelCase     (					self,      lowerCamelCase)       ->   str:




           """simple docstring"""

           _lowercase       :							str      		=		set(self.all_special_ids)  # call it once instead of inside list comp
           all_special_ids.remove(self.unk_token_id)  # <unk> is only sometimes special

           return [1 if x in all_special_ids else 0 for x in seq]
       def 			UpperCamelCase     (					self,      lowerCamelCase,      lowerCamelCase = None,      lowerCamelCase = False)       ->   List[int]:




           """simple docstring"""

           if already_has_special_tokens:
               return self._special_token_mask(lowerCamelCase)
           elif token_ids_a is None:
               return self._special_token_mask(lowerCamelCase) + [1]
           else:
               return self._special_token_mask(token_ids_a + token_ids_a) + [1]
       def 			UpperCamelCase     (					self,      lowerCamelCase,      lowerCamelCase=None)       ->   List[int]:




           """simple docstring"""

           if token_ids_a is None:
               return token_ids_a + [self.eos_token_id]
           # We don't expect to process pairs, but leave the pair logic for API consistency
           return token_ids_a + token_ids_a + [self.eos_token_id]




       def 			UpperCamelCase     (					self,      lowerCamelCase,      lowerCamelCase = None)       ->   Tuple[str]:




           """simple docstring"""

           if not os.path.isdir(lowerCamelCase):
               logger.error(F'''Vocabulary path ({save_directory}) should be a directory''')
               return
           _lowercase       :							Any      		=		os.path.join(
               lowerCamelCase,      (filename_prefix + '-' if filename_prefix else '') + VOCAB_FILES_NAMES['vocab_file'])

           if os.path.abspath(self.vocab_file) != os.path.abspath(lowerCamelCase) and os.path.isfile(self.vocab_file):
               copyfile(self.vocab_file,      lowerCamelCase)
           elif not os.path.isfile(self.vocab_file):
               with open(lowerCamelCase,      'wb') as fi:
                   _lowercase       :							int      		=		self.sp_model.serialized_model_proto()
                   fi.write(lowerCamelCase)

           return (out_vocab_file,)






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


