




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     snake_case__							(_UpperCamelCase    ,       unittest.TestCase   ):
				"""simple docstring"""



				SCREAMING_SNAKE_CASE_       :							Any							 =					"""hf-internal-testing/tiny-random-OnnxStableDiffusionPipeline"""





				def 							__UpperCAmelCase							(      self				:     Tuple					,      __lowerCamelCase				:     Any=0							)					->     Any:
							a    =			np.random.RandomState(lowerCamelCase_							)
							a    =			{
							    "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 							__UpperCAmelCase							(      self				:     Any							)					->     Optional[Any]:
							a    =			OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint					,      provider="CPUExecutionProvider"							)
							pipe.set_progress_bar_config(disable=lowerCamelCase_							)

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

							assert image.shape == (1, 1_28, 1_28, 3)
							a    =			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 							__UpperCAmelCase							(      self				:     int							)					->     Union[str, Any]:
							a    =			OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint					,      provider="CPUExecutionProvider"							)
							a    =			PNDMScheduler.from_config(pipe.scheduler.config					,      skip_prk_steps=lowerCamelCase_							)
							pipe.set_progress_bar_config(disable=lowerCamelCase_							)

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

							assert image.shape == (1, 1_28, 1_28, 3)
							a    =			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 							__UpperCAmelCase							(      self				:     Dict							)					->     Optional[int]:
							a    =			OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint					,      provider="CPUExecutionProvider"							)
							a    =			LMSDiscreteScheduler.from_config(pipe.scheduler.config							)
							pipe.set_progress_bar_config(disable=lowerCamelCase_							)

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

							assert image.shape == (1, 1_28, 1_28, 3)
							a    =			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 							__UpperCAmelCase							(      self				:     Optional[int]							)					->     Dict:
							a    =			OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint					,      provider="CPUExecutionProvider"							)
							a    =			EulerDiscreteScheduler.from_config(pipe.scheduler.config							)
							pipe.set_progress_bar_config(disable=lowerCamelCase_							)

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

							assert image.shape == (1, 1_28, 1_28, 3)
							a    =			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 							__UpperCAmelCase							(      self				:     str							)					->     Tuple:
							a    =			OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint					,      provider="CPUExecutionProvider"							)
							a    =			EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config							)
							pipe.set_progress_bar_config(disable=lowerCamelCase_							)

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

							assert image.shape == (1, 1_28, 1_28, 3)
							a    =			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 							__UpperCAmelCase							(      self				:     Tuple							)					->     Union[str, Any]:
							a    =			OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint					,      provider="CPUExecutionProvider"							)
							a    =			DPMSolverMultistepScheduler.from_config(pipe.scheduler.config							)
							pipe.set_progress_bar_config(disable=lowerCamelCase_							)

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

							assert image.shape == (1, 1_28, 1_28, 3)
							a    =			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 							__UpperCAmelCase							(      self				:     List[str]							)					->     Any:
							a    =			OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint					,      provider="CPUExecutionProvider"							)
							pipe.set_progress_bar_config(disable=lowerCamelCase_							)

							a    =			self.get_dummy_inputs()
							a    =			3 * [inputs["prompt"]]

							# forward
							a    =			pipe(**lowerCamelCase_							)
							a    =			output.images[0, -3:, -3:, -1]

							a    =			self.get_dummy_inputs()
							a    =			3 * [inputs.pop("prompt"							)]

							a    =			pipe.tokenizer(
							    lowerCamelCase_					,      padding="max_length"					,      max_length=pipe.tokenizer.model_max_length					,      truncation=lowerCamelCase_					,      return_tensors="np"					,      )
							a    =			text_inputs["input_ids"]

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

							a    =			prompt_embeds

							# forward
							a    =			pipe(**lowerCamelCase_							)
							a    =			output.images[0, -3:, -3:, -1]

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





				def 							__UpperCAmelCase							(      self				:     Tuple							)					->     str:
							a    =			OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint					,      provider="CPUExecutionProvider"							)
							pipe.set_progress_bar_config(disable=lowerCamelCase_							)

							a    =			self.get_dummy_inputs()
							a    =			3 * ["this is a negative prompt"]
							a    =			negative_prompt
							a    =			3 * [inputs["prompt"]]

							# forward
							a    =			pipe(**lowerCamelCase_							)
							a    =			output.images[0, -3:, -3:, -1]

							a    =			self.get_dummy_inputs()
							a    =			3 * [inputs.pop("prompt"							)]

							a    =			[]
							for p in [prompt, negative_prompt]:
										a    =			pipe.tokenizer(
										    lowerCamelCase_					,      padding="max_length"					,      max_length=pipe.tokenizer.model_max_length					,      truncation=lowerCamelCase_					,      return_tensors="np"					,      )
										a    =			text_inputs["input_ids"]

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

							a    ,		a    =			embeds

							# forward
							a    =			pipe(**lowerCamelCase_							)
							a    =			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     snake_case__							(unittest.TestCase   ):
				"""simple docstring"""



				@property
				def 							__UpperCAmelCase							(      self				:     List[Any]							)					->     Any:
							return (
							    "CUDAExecutionProvider",
							    {
							        "gpu_mem_limit": "15000000000",  # 15GB
							        "arena_extend_strategy": "kSameAsRequested",
							    },
							)





				@property
				def 							__UpperCAmelCase							(      self				:     Optional[int]							)					->     int:
							a    =			ort.SessionOptions()
							a    =			False
							return options





				def 							__UpperCAmelCase							(      self				:     Optional[int]							)					->     List[Any]:
							a    =			OnnxStableDiffusionPipeline.from_pretrained(
							    "CompVis/stable-diffusion-v1-4"					,      revision="onnx"					,      safety_checker=lowerCamelCase_					,      feature_extractor=lowerCamelCase_					,      provider=self.gpu_provider					,      sess_options=self.gpu_options					,      )
							sd_pipe.set_progress_bar_config(disable=lowerCamelCase_							)

							a    =			"A painting of a squirrel eating a burger"
							np.random.seed(0							)
							a    =			sd_pipe([prompt]					,      guidance_scale=6.0					,      num_inference_steps=10					,      output_type="np"							)
							a    =			output.images

							a    =			image[0, -3:, -3:, -1]

							assert image.shape == (1, 5_12, 5_12, 3)
							a    =			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 							__UpperCAmelCase							(      self				:     Optional[Any]							)					->     Dict:
							a    =			DDIMScheduler.from_pretrained(
							    "runwayml/stable-diffusion-v1-5"					,      subfolder="scheduler"					,      revision="onnx"							)
							a    =			OnnxStableDiffusionPipeline.from_pretrained(
							    "runwayml/stable-diffusion-v1-5"					,      revision="onnx"					,      scheduler=lowerCamelCase_					,      safety_checker=lowerCamelCase_					,      feature_extractor=lowerCamelCase_					,      provider=self.gpu_provider					,      sess_options=self.gpu_options					,      )
							sd_pipe.set_progress_bar_config(disable=lowerCamelCase_							)

							a    =			"open neural network exchange"
							a    =			np.random.RandomState(0							)
							a    =			sd_pipe([prompt]					,      guidance_scale=7.5					,      num_inference_steps=10					,      generator=lowerCamelCase_					,      output_type="np"							)
							a    =			output.images
							a    =			image[0, -3:, -3:, -1]

							assert image.shape == (1, 5_12, 5_12, 3)
							a    =			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 							__UpperCAmelCase							(      self				:     int							)					->     Optional[int]:
							a    =			LMSDiscreteScheduler.from_pretrained(
							    "runwayml/stable-diffusion-v1-5"					,      subfolder="scheduler"					,      revision="onnx"							)
							a    =			OnnxStableDiffusionPipeline.from_pretrained(
							    "runwayml/stable-diffusion-v1-5"					,      revision="onnx"					,      scheduler=lowerCamelCase_					,      safety_checker=lowerCamelCase_					,      feature_extractor=lowerCamelCase_					,      provider=self.gpu_provider					,      sess_options=self.gpu_options					,      )
							sd_pipe.set_progress_bar_config(disable=lowerCamelCase_							)

							a    =			"open neural network exchange"
							a    =			np.random.RandomState(0							)
							a    =			sd_pipe([prompt]					,      guidance_scale=7.5					,      num_inference_steps=10					,      generator=lowerCamelCase_					,      output_type="np"							)
							a    =			output.images
							a    =			image[0, -3:, -3:, -1]

							assert image.shape == (1, 5_12, 5_12, 3)
							a    =			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 							__UpperCAmelCase							(      self				:     Union[str, Any]							)					->     int:
							a    =			0

							def test_callback_fn(__lowerCamelCase				:     Any					,      __lowerCamelCase				:     List[str]					,      __lowerCamelCase				:     Optional[int]							) -> None:
										a    =			True
										nonlocal number_of_steps
										number_of_steps += 1
										if step == 0:
													assert latents.shape == (1, 4, 64, 64)
													a    =			latents[0, -3:, -3:, -1]
													a    =			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, 64, 64)
													a    =			latents[0, -3:, -3:, -1]
													a    =			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    =			False

							a    =			OnnxStableDiffusionPipeline.from_pretrained(
							    "runwayml/stable-diffusion-v1-5"					,      revision="onnx"					,      safety_checker=lowerCamelCase_					,      feature_extractor=lowerCamelCase_					,      provider=self.gpu_provider					,      sess_options=self.gpu_options					,      )
							pipe.set_progress_bar_config(disable=lowerCamelCase_							)

							a    =			"Andromeda galaxy in a bottle"

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





				def 							__UpperCAmelCase							(      self				:     Any							)					->     List[Any]:
							a    =			OnnxStableDiffusionPipeline.from_pretrained(
							    "runwayml/stable-diffusion-v1-5"					,      revision="onnx"					,      safety_checker=lowerCamelCase_					,      feature_extractor=lowerCamelCase_					,      provider=self.gpu_provider					,      sess_options=self.gpu_options					,      )
							assert isinstance(lowerCamelCase_					,      lowerCamelCase_							)
							assert pipe.safety_checker is None

							a    =			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(lowerCamelCase_							)
										a    =			OnnxStableDiffusionPipeline.from_pretrained(lowerCamelCase_							)

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











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


_lowercase: str      			=		logging.get_logger(__name__)

_lowercase: List[str]      			=		{
    "microsoft/trocr-base-handwritten": (
        "https://huggingface.co/microsoft/trocr-base-handwritten/resolve/main/config.json"
    ),
    # See all TrOCR models at https://huggingface.co/models?filter=trocr
}



class 	_lowercase							(   lowerCAmelCase	):





	"""simple docstring"""
	__A							=							"trocr"
	__A							=							["past_key_values"]
	__A							=							{
	    "num_attention_heads": "decoder_attention_heads",
	    "hidden_size": "d_model",
	    "num_hidden_layers": "decoder_layers",
	}




	def __init__(self		,       lowerCamelCase_=50265		,       lowerCamelCase_=1024		,       lowerCamelCase_=12		,       lowerCamelCase_=16		,       lowerCamelCase_=4096		,       lowerCamelCase_="gelu"		,       lowerCamelCase_=512		,       lowerCamelCase_=0.1		,       lowerCamelCase_=0.0		,       lowerCamelCase_=0.0		,       lowerCamelCase_=2		,       lowerCamelCase_=0.02		,       lowerCamelCase_=0.0		,       lowerCamelCase_=True		,       lowerCamelCase_=False		,       lowerCamelCase_=True		,       lowerCamelCase_=True		,       lowerCamelCase_=1		,       lowerCamelCase_=0		,       lowerCamelCase_=2		,       **lowerCamelCase_		,       ):
						"""simple docstring"""

						a					     =   vocab_size
						a					     =   d_model
						a					     =   decoder_layers
						a					     =   decoder_attention_heads
						a					     =   decoder_ffn_dim
						a					     =   activation_function
						a					     =   max_position_embeddings
						a					     =   dropout
						a					     =   attention_dropout
						a					     =   activation_dropout
						a					     =   init_std
						a					     =   decoder_layerdrop
						a					     =   use_cache
						a					     =   scale_embedding
						a					     =   use_learned_position_embeddings
						a					     =   layernorm_embedding

						super().__init__(
						    pad_token_id=lowerCamelCase_		,       bos_token_id=lowerCamelCase_		,       eos_token_id=lowerCamelCase_		,       decoder_start_token_id=lowerCamelCase_		,       **lowerCamelCase_		,       )
