

'''simple docstring'''






import tempfile
import unittest

import numpy as np

from diffusers import (
    DDIMScheduler,
    DPMSolverMultistepScheduler,
    EulerAncestralDiscreteScheduler,
    EulerDiscreteScheduler,
    LMSDiscreteScheduler,
    OnnxStableDiffusionPipeline,
    PNDMScheduler,
)
from diffusers.utils.testing_utils import is_onnx_available, nightly, require_onnxruntime, require_torch_gpu

from ..test_pipelines_onnx_common import OnnxPipelineTesterMixin


if is_onnx_available():
 import onnxruntime as ort






class 			lowerCAmelCase__     (					lowerCAmelCase__				,			unittest.TestCase						):




     """simple docstring"""





     __UpperCamelCase									= """hf-internal-testing/tiny-random-OnnxStableDiffusionPipeline"""







     def 			__lowerCAmelCase	(		self						:	List[str]      ,						A__						:	Optional[int]=0	)       ->			int:
          '''simple docstring'''




          a__   :      str     =						np.random.RandomState(__a	)
          a__   :      str     =						{
              '''prompt''': '''A painting of a squirrel eating a burger''',
              '''generator''': generator,
              '''num_inference_steps''': 2,
              '''guidance_scale''': 7.5,
              '''output_type''': '''numpy''',
          }
          return inputs







     def 			__lowerCAmelCase	(		self						:	int	)       ->			List[Any]:
          '''simple docstring'''




          a__   :      Union[str, Any]     =						OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint      ,						provider='''CPUExecutionProvider'''	)
          pipe.set_progress_bar_config(disable=__a	)

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

          assert image.shape == (1, 1_2_8, 1_2_8, 3)
          a__   :      List[Any]     =						np.array([0.65_072, 0.58_492, 0.48_219, 0.55_521, 0.53_180, 0.55_939, 0.50_697, 0.39_800, 0.46_455]	)

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







     def 			__lowerCAmelCase	(		self						:	Dict	)       ->			str:
          '''simple docstring'''




          a__   :      List[str]     =						OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint      ,						provider='''CPUExecutionProvider'''	)
          a__   :      Optional[Any]     =						PNDMScheduler.from_config(pipe.scheduler.config      ,						skip_prk_steps=__a	)
          pipe.set_progress_bar_config(disable=__a	)

          a__   :      int     =						self.get_dummy_inputs()
          a__   :      str     =						pipe(**__a	).images
          a__   :      Dict     =						image[0, -3:, -3:, -1]

          assert image.shape == (1, 1_2_8, 1_2_8, 3)
          a__   :      str     =						np.array([0.65_863, 0.59_425, 0.49_326, 0.56_313, 0.53_875, 0.56_627, 0.51_065, 0.39_777, 0.46_330]	)

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







     def 			__lowerCAmelCase	(		self						:	List[str]	)       ->			Optional[int]:
          '''simple docstring'''




          a__   :      Dict     =						OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint      ,						provider='''CPUExecutionProvider'''	)
          a__   :      Optional[Any]     =						LMSDiscreteScheduler.from_config(pipe.scheduler.config	)
          pipe.set_progress_bar_config(disable=__a	)

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

          assert image.shape == (1, 1_2_8, 1_2_8, 3)
          a__   :      Dict     =						np.array([0.53_755, 0.60_786, 0.47_402, 0.49_488, 0.51_869, 0.49_819, 0.47_985, 0.38_957, 0.44_279]	)

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







     def 			__lowerCAmelCase	(		self						:	Optional[int]	)       ->			int:
          '''simple docstring'''




          a__   :      Tuple     =						OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint      ,						provider='''CPUExecutionProvider'''	)
          a__   :      Union[str, Any]     =						EulerDiscreteScheduler.from_config(pipe.scheduler.config	)
          pipe.set_progress_bar_config(disable=__a	)

          a__   :      Tuple     =						self.get_dummy_inputs()
          a__   :      str     =						pipe(**__a	).images
          a__   :      List[str]     =						image[0, -3:, -3:, -1]

          assert image.shape == (1, 1_2_8, 1_2_8, 3)
          a__   :      Any     =						np.array([0.53_755, 0.60_786, 0.47_402, 0.49_488, 0.51_869, 0.49_819, 0.47_985, 0.38_957, 0.44_279]	)

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







     def 			__lowerCAmelCase	(		self						:	Optional[int]	)       ->			Optional[int]:
          '''simple docstring'''




          a__   :      Optional[Any]     =						OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint      ,						provider='''CPUExecutionProvider'''	)
          a__   :      Optional[Any]     =						EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config	)
          pipe.set_progress_bar_config(disable=__a	)

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

          assert image.shape == (1, 1_2_8, 1_2_8, 3)
          a__   :      Any     =						np.array([0.53_817, 0.60_812, 0.47_384, 0.49_530, 0.51_894, 0.49_814, 0.47_984, 0.38_958, 0.44_271]	)

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







     def 			__lowerCAmelCase	(		self						:	Optional[Any]	)       ->			List[str]:
          '''simple docstring'''




          a__   :      Optional[Any]     =						OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint      ,						provider='''CPUExecutionProvider'''	)
          a__   :      Optional[int]     =						DPMSolverMultistepScheduler.from_config(pipe.scheduler.config	)
          pipe.set_progress_bar_config(disable=__a	)

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

          assert image.shape == (1, 1_2_8, 1_2_8, 3)
          a__   :      Any     =						np.array([0.53_895, 0.60_808, 0.47_933, 0.49_608, 0.51_886, 0.49_950, 0.48_053, 0.38_957, 0.44_200]	)

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







     def 			__lowerCAmelCase	(		self						:	List[str]	)       ->			str:
          '''simple docstring'''




          a__   :      Union[str, Any]     =						OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint      ,						provider='''CPUExecutionProvider'''	)
          pipe.set_progress_bar_config(disable=__a	)

          a__   :      Any     =						self.get_dummy_inputs()
          a__   :      Tuple     =						3 * [inputs['''prompt''']]

          # forward
          a__   :      List[Any]     =						pipe(**__a	)
          a__   :      Union[str, Any]     =						output.images[0, -3:, -3:, -1]

          a__   :      str     =						self.get_dummy_inputs()
          a__   :      Tuple     =						3 * [inputs.pop('''prompt'''	)]

          a__   :      Union[str, Any]     =						pipe.tokenizer(
              __a      ,						padding='''max_length'''      ,						max_length=pipe.tokenizer.model_max_length      ,						truncation=__a      ,						return_tensors='''np'''      ,						)
          a__   :      List[Any]     =						text_inputs['''input_ids''']

          a__   :      Tuple     =						pipe.text_encoder(input_ids=text_inputs.astype(np.intaa	)	)[0]

          a__   :      List[Any]     =						prompt_embeds

          # forward
          a__   :      List[Any]     =						pipe(**__a	)
          a__   :      str     =						output.images[0, -3:, -3:, -1]

          assert np.abs(image_slice_a.flatten() - image_slice_a.flatten()	).max() < 1E-4







     def 			__lowerCAmelCase	(		self						:	Optional[Any]	)       ->			Tuple:
          '''simple docstring'''




          a__   :      Tuple     =						OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint      ,						provider='''CPUExecutionProvider'''	)
          pipe.set_progress_bar_config(disable=__a	)

          a__   :      str     =						self.get_dummy_inputs()
          a__   :      str     =						3 * ['''this is a negative prompt''']
          a__   :      int     =						negative_prompt
          a__   :      int     =						3 * [inputs['''prompt''']]

          # forward
          a__   :      Tuple     =						pipe(**__a	)
          a__   :      int     =						output.images[0, -3:, -3:, -1]

          a__   :      Tuple     =						self.get_dummy_inputs()
          a__   :      Tuple     =						3 * [inputs.pop('''prompt'''	)]

          a__   :      Optional[int]     =						[]
          for p in [prompt, negative_prompt]:
               a__   :      Dict     =						pipe.tokenizer(
                   __a      ,						padding='''max_length'''      ,						max_length=pipe.tokenizer.model_max_length      ,						truncation=__a      ,						return_tensors='''np'''      ,						)
               a__   :      Any     =						text_inputs['''input_ids''']

               embeds.append(pipe.text_encoder(input_ids=text_inputs.astype(np.intaa	)	)[0]	)

          a__     ,    a__   :      Optional[Any]     =						embeds

          # forward
          a__   :      List[str]     =						pipe(**__a	)
          a__   :      int     =						output.images[0, -3:, -3:, -1]

          assert np.abs(image_slice_a.flatten() - image_slice_a.flatten()	).max() < 1E-4






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




     """simple docstring"""





     @property
     def 			__lowerCAmelCase	(		self						:	Any	)       ->			Optional[int]:
          '''simple docstring'''




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







     @property
     def 			__lowerCAmelCase	(		self						:	Any	)       ->			int:
          '''simple docstring'''




          a__   :      Union[str, Any]     =						ort.SessionOptions()
          a__   :      List[str]     =						False
          return options







     def 			__lowerCAmelCase	(		self						:	List[str]	)       ->			Optional[Any]:
          '''simple docstring'''




          a__   :      Optional[Any]     =						OnnxStableDiffusionPipeline.from_pretrained(
              '''CompVis/stable-diffusion-v1-4'''      ,						revision='''onnx'''      ,						safety_checker=__a      ,						feature_extractor=__a      ,						provider=self.gpu_provider      ,						sess_options=self.gpu_options      ,						)
          sd_pipe.set_progress_bar_config(disable=__a	)

          a__   :      Union[str, Any]     =						'''A painting of a squirrel eating a burger'''
          np.random.seed(0	)
          a__   :      Any     =						sd_pipe([prompt]      ,						guidance_scale=6.0      ,						num_inference_steps=1_0      ,						output_type='''np'''	)
          a__   :      str     =						output.images

          a__   :      Dict     =						image[0, -3:, -3:, -1]

          assert image.shape == (1, 5_1_2, 5_1_2, 3)
          a__   :      Dict     =						np.array([0.0_452, 0.0_390, 0.0_087, 0.0_350, 0.0_617, 0.0_364, 0.0_544, 0.0_523, 0.0_720]	)

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







     def 			__lowerCAmelCase	(		self						:	List[Any]	)       ->			Optional[Any]:
          '''simple docstring'''




          a__   :      Tuple     =						DDIMScheduler.from_pretrained(
              '''runwayml/stable-diffusion-v1-5'''      ,						subfolder='''scheduler'''      ,						revision='''onnx'''	)
          a__   :      Any     =						OnnxStableDiffusionPipeline.from_pretrained(
              '''runwayml/stable-diffusion-v1-5'''      ,						revision='''onnx'''      ,						scheduler=__a      ,						safety_checker=__a      ,						feature_extractor=__a      ,						provider=self.gpu_provider      ,						sess_options=self.gpu_options      ,						)
          sd_pipe.set_progress_bar_config(disable=__a	)

          a__   :      int     =						'''open neural network exchange'''
          a__   :      List[Any]     =						np.random.RandomState(0	)
          a__   :      List[str]     =						sd_pipe([prompt]      ,						guidance_scale=7.5      ,						num_inference_steps=1_0      ,						generator=__a      ,						output_type='''np'''	)
          a__   :      Optional[int]     =						output.images
          a__   :      Union[str, Any]     =						image[0, -3:, -3:, -1]

          assert image.shape == (1, 5_1_2, 5_1_2, 3)
          a__   :      Tuple     =						np.array([0.2_867, 0.1_974, 0.1_481, 0.7_294, 0.7_251, 0.6_667, 0.4_194, 0.5_642, 0.6_486]	)

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







     def 			__lowerCAmelCase	(		self						:	Union[str, Any]	)       ->			Any:
          '''simple docstring'''




          a__   :      List[str]     =						LMSDiscreteScheduler.from_pretrained(
              '''runwayml/stable-diffusion-v1-5'''      ,						subfolder='''scheduler'''      ,						revision='''onnx'''	)
          a__   :      List[Any]     =						OnnxStableDiffusionPipeline.from_pretrained(
              '''runwayml/stable-diffusion-v1-5'''      ,						revision='''onnx'''      ,						scheduler=__a      ,						safety_checker=__a      ,						feature_extractor=__a      ,						provider=self.gpu_provider      ,						sess_options=self.gpu_options      ,						)
          sd_pipe.set_progress_bar_config(disable=__a	)

          a__   :      int     =						'''open neural network exchange'''
          a__   :      List[Any]     =						np.random.RandomState(0	)
          a__   :      Dict     =						sd_pipe([prompt]      ,						guidance_scale=7.5      ,						num_inference_steps=1_0      ,						generator=__a      ,						output_type='''np'''	)
          a__   :      str     =						output.images
          a__   :      Union[str, Any]     =						image[0, -3:, -3:, -1]

          assert image.shape == (1, 5_1_2, 5_1_2, 3)
          a__   :      Optional[int]     =						np.array([0.2_306, 0.1_959, 0.1_593, 0.6_549, 0.6_394, 0.5_408, 0.5_065, 0.6_010, 0.6_161]	)

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







     def 			__lowerCAmelCase	(		self						:	List[str]	)       ->			int:
          '''simple docstring'''




          a__   :      str     =						0

          def test_callback_fn(A__						:	str      ,						A__						:	Optional[int]      ,						A__						:	str	) -> None:
               a__   :      int     =						True
               nonlocal number_of_steps
               number_of_steps += 1
               if step == 0:
                    assert latents.shape == (1, 4, 6_4, 6_4)
                    a__   :      Union[str, Any]     =						latents[0, -3:, -3:, -1]
                    a__   :      List[str]     =						np.array(
                        [-0.6_772, -0.3_835, -1.2_456, 0.1_905, -1.0_974, 0.6_967, -1.9_353, 0.0_178, 1.0_167]	)

                    assert np.abs(latents_slice.flatten() - expected_slice	).max() < 1E-3
               elif step == 5:
                    assert latents.shape == (1, 4, 6_4, 6_4)
                    a__   :      Dict     =						latents[0, -3:, -3:, -1]
                    a__   :      List[str]     =						np.array(
                        [-0.3_351, 0.2_241, -0.1_837, -0.2_325, -0.6_577, 0.3_393, -0.0_241, 0.5_899, 1.3_875]	)

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

          a__   :      Optional[int]     =						False

          a__   :      Optional[Any]     =						OnnxStableDiffusionPipeline.from_pretrained(
              '''runwayml/stable-diffusion-v1-5'''      ,						revision='''onnx'''      ,						safety_checker=__a      ,						feature_extractor=__a      ,						provider=self.gpu_provider      ,						sess_options=self.gpu_options      ,						)
          pipe.set_progress_bar_config(disable=__a	)

          a__   :      Optional[Any]     =						'''Andromeda galaxy in a bottle'''

          a__   :      Dict     =						np.random.RandomState(0	)
          pipe(
              prompt=__a      ,						num_inference_steps=5      ,						guidance_scale=7.5      ,						generator=__a      ,						callback=__a      ,						callback_steps=1      ,						)
          assert test_callback_fn.has_been_called
          assert number_of_steps == 6







     def 			__lowerCAmelCase	(		self						:	List[Any]	)       ->			Dict:
          '''simple docstring'''




          a__   :      str     =						OnnxStableDiffusionPipeline.from_pretrained(
              '''runwayml/stable-diffusion-v1-5'''      ,						revision='''onnx'''      ,						safety_checker=__a      ,						feature_extractor=__a      ,						provider=self.gpu_provider      ,						sess_options=self.gpu_options      ,						)
          assert isinstance(__a      ,						__a	)
          assert pipe.safety_checker is None

          a__   :      Union[str, Any]     =						pipe('''example prompt'''      ,						num_inference_steps=2	).images[0]
          assert image is not None

          # check that there's no error when saving a pipeline with one of the models being None
          with tempfile.TemporaryDirectory() as tmpdirname:
               pipe.save_pretrained(__a	)
               a__   :      List[str]     =						OnnxStableDiffusionPipeline.from_pretrained(__a	)

          # sanity check that the pipeline still works
          assert pipe.safety_checker is None
          a__   :      Optional[int]     =						pipe('''example prompt'''      ,						num_inference_steps=2	).images[0]
          assert image is not None




"""simple docstring"""


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

from ...tokenization_utils import AddedToken
from ...tokenization_utils_fast import PreTrainedTokenizerFast
from ...utils import is_sentencepiece_available, logging


if is_sentencepiece_available():
	from .tokenization_xlnet import XLNetTokenizer
else:
	A	:  Optional[Any]            =     None


A	:  Any            =     logging.get_logger(__name__)

A	:  Union[str, Any]            =     {"vocab_file": "spiece.model", "tokenizer_file": "tokenizer.json"}

A	:  List[Any]            =     {
    "vocab_file": {
        "xlnet-base-cased": "https://huggingface.co/xlnet-base-cased/resolve/main/spiece.model",
        "xlnet-large-cased": "https://huggingface.co/xlnet-large-cased/resolve/main/spiece.model",
    },
    "tokenizer_file": {
        "xlnet-base-cased": "https://huggingface.co/xlnet-base-cased/resolve/main/tokenizer.json",
        "xlnet-large-cased": "https://huggingface.co/xlnet-large-cased/resolve/main/tokenizer.json",
    },
}

A	:  Union[str, Any]            =     {
    "xlnet-base-cased": None,
    "xlnet-large-cased": None,
}

A	:  Optional[int]            =     "▁"

# Segments (not really needed)
A	:  str            =     0
A	:  Dict            =     1
A	:  List[str]            =     2
A	:  Tuple            =     3
A	:  Tuple            =     4
class     _UpperCamelCase    (  lowerCAmelCase__						):
							'''simple docstring'''







							__UpperCAmelCase		:       Optional[Any]							 =VOCAB_FILES_NAMES
							__UpperCAmelCase		:       Optional[int]							 =PRETRAINED_VOCAB_FILES_MAP
							__UpperCAmelCase		:       Dict							 =PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
							__UpperCAmelCase		:       List[Any]							 ="""left"""
							__UpperCAmelCase		:       List[str]							 =XLNetTokenizer

							def __init__(	self       ,  __a=None       ,  __a=None       ,  __a=False       ,  __a=True       ,  __a=False       ,  __a="<s>"       ,  __a="</s>"       ,  __a="<unk>"       ,  __a="<sep>"       ,  __a="<pad>"       ,  __a="<cls>"       ,  __a="<mask>"       ,  __a=["<eop>", "<eod>"]       ,  **__a       ,  ):
										# Mask token behave like a normal word, i.e. include the space before it
										__lowerCAmelCase							       =     AddedToken(__a       ,  lstrip=__a       ,  rstrip=__a      ) if isinstance(__a       ,  __a      ) else mask_token

										super().__init__(
										    vocab_file=__a       ,  tokenizer_file=__a       ,  do_lower_case=__a       ,  remove_space=__a       ,  keep_accents=__a       ,  bos_token=__a       ,  eos_token=__a       ,  unk_token=__a       ,  sep_token=__a       ,  pad_token=__a       ,  cls_token=__a       ,  mask_token=__a       ,  additional_special_tokens=__a       ,  **__a       ,  )

										__lowerCAmelCase							       =     3
										__lowerCAmelCase							       =     do_lower_case
										__lowerCAmelCase							       =     remove_space
										__lowerCAmelCase							       =     keep_accents
										__lowerCAmelCase							       =     vocab_file
										__lowerCAmelCase							       =     False if not self.vocab_file else True

							def      snake_case							(	self       ,  __a       ,  __a = None      ):
										__lowerCAmelCase							       =     [self.sep_token_id]
										__lowerCAmelCase							       =     [self.cls_token_id]
										if token_ids_a is None:
													return token_ids_a + sep + cls
										return token_ids_a + sep + token_ids_a + sep + cls

							def      snake_case							(	self       ,  __a       ,  __a = None      ):
										__lowerCAmelCase							       =     [self.sep_token_id]
										__lowerCAmelCase							       =     [2]

										if token_ids_a is None:
													return len(token_ids_a + sep      ) * [0] + cls_segment_id
										return len(token_ids_a + sep      ) * [0] + len(token_ids_a + sep      ) * [1] + cls_segment_id



							def      snake_case							(	self       ,  __a       ,  __a = None      ):
										if not self.can_save_slow_tokenizer:
													raise ValueError(
													    "Your fast tokenizer does not have the necessary information to save the vocabulary for a slow "
													    "tokenizer."      )

										if not os.path.isdir(__a      ):
													logger.error(f"Vocabulary path ({save_directory}) should be a directory"      )
													return
										__lowerCAmelCase							       =     os.path.join(
										    __a       ,  (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["vocab_file"]      )

										if os.path.abspath(self.vocab_file      ) != os.path.abspath(__a      ):
													copyfile(self.vocab_file       ,  __a      )

										return (out_vocab_file,)




