


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







							__UpperCAmelCase		:       Optional[Any]							 ="""hf-internal-testing/tiny-random-OnnxStableDiffusionPipeline"""

							def      snake_case							(	self       ,  __a=0      ):
										__lowerCAmelCase							       =     np.random.RandomState(__a      )
										__lowerCAmelCase							       =     {
										    "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      snake_case							(	self      ):
										__lowerCAmelCase							       =     OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint       ,  provider="CPUExecutionProvider"      )
										pipe.set_progress_bar_config(disable=__a      )

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

										assert image.shape == (1, 1_28, 1_28, 3)
										__lowerCAmelCase							       =     np.array([0.6_5_0_7_2, 0.5_8_4_9_2, 0.4_8_2_1_9, 0.5_5_5_2_1, 0.5_3_1_8_0, 0.5_5_9_3_9, 0.5_0_6_9_7, 0.3_9_8_0_0, 0.4_6_4_5_5]      )

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

							def      snake_case							(	self      ):
										__lowerCAmelCase							       =     OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint       ,  provider="CPUExecutionProvider"      )
										__lowerCAmelCase							       =     PNDMScheduler.from_config(pipe.scheduler.config       ,  skip_prk_steps=__a      )
										pipe.set_progress_bar_config(disable=__a      )

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

										assert image.shape == (1, 1_28, 1_28, 3)
										__lowerCAmelCase							       =     np.array([0.6_5_8_6_3, 0.5_9_4_2_5, 0.4_9_3_2_6, 0.5_6_3_1_3, 0.5_3_8_7_5, 0.5_6_6_2_7, 0.5_1_0_6_5, 0.3_9_7_7_7, 0.4_6_3_3_0]      )

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

							def      snake_case							(	self      ):
										__lowerCAmelCase							       =     OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint       ,  provider="CPUExecutionProvider"      )
										__lowerCAmelCase							       =     LMSDiscreteScheduler.from_config(pipe.scheduler.config      )
										pipe.set_progress_bar_config(disable=__a      )

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

										assert image.shape == (1, 1_28, 1_28, 3)
										__lowerCAmelCase							       =     np.array([0.5_3_7_5_5, 0.6_0_7_8_6, 0.4_7_4_0_2, 0.4_9_4_8_8, 0.5_1_8_6_9, 0.4_9_8_1_9, 0.4_7_9_8_5, 0.3_8_9_5_7, 0.4_4_2_7_9]      )

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

							def      snake_case							(	self      ):
										__lowerCAmelCase							       =     OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint       ,  provider="CPUExecutionProvider"      )
										__lowerCAmelCase							       =     EulerDiscreteScheduler.from_config(pipe.scheduler.config      )
										pipe.set_progress_bar_config(disable=__a      )

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

										assert image.shape == (1, 1_28, 1_28, 3)
										__lowerCAmelCase							       =     np.array([0.5_3_7_5_5, 0.6_0_7_8_6, 0.4_7_4_0_2, 0.4_9_4_8_8, 0.5_1_8_6_9, 0.4_9_8_1_9, 0.4_7_9_8_5, 0.3_8_9_5_7, 0.4_4_2_7_9]      )

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

							def      snake_case							(	self      ):
										__lowerCAmelCase							       =     OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint       ,  provider="CPUExecutionProvider"      )
										__lowerCAmelCase							       =     EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config      )
										pipe.set_progress_bar_config(disable=__a      )

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

										assert image.shape == (1, 1_28, 1_28, 3)
										__lowerCAmelCase							       =     np.array([0.5_3_8_1_7, 0.6_0_8_1_2, 0.4_7_3_8_4, 0.4_9_5_3_0, 0.5_1_8_9_4, 0.4_9_8_1_4, 0.4_7_9_8_4, 0.3_8_9_5_8, 0.4_4_2_7_1]      )

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

							def      snake_case							(	self      ):
										__lowerCAmelCase							       =     OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint       ,  provider="CPUExecutionProvider"      )
										__lowerCAmelCase							       =     DPMSolverMultistepScheduler.from_config(pipe.scheduler.config      )
										pipe.set_progress_bar_config(disable=__a      )

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

										assert image.shape == (1, 1_28, 1_28, 3)
										__lowerCAmelCase							       =     np.array([0.5_3_8_9_5, 0.6_0_8_0_8, 0.4_7_9_3_3, 0.4_9_6_0_8, 0.5_1_8_8_6, 0.4_9_9_5_0, 0.4_8_0_5_3, 0.3_8_9_5_7, 0.4_4_2_0_0]      )

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

							def      snake_case							(	self      ):
										__lowerCAmelCase							       =     OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint       ,  provider="CPUExecutionProvider"      )
										pipe.set_progress_bar_config(disable=__a      )

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

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

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

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

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

										__lowerCAmelCase							       =     prompt_embeds

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

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



							def      snake_case							(	self      ):
										__lowerCAmelCase							       =     OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint       ,  provider="CPUExecutionProvider"      )
										pipe.set_progress_bar_config(disable=__a      )

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

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

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

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

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

										__lowerCAmelCase       ,      __lowerCAmelCase							       =     embeds

										# forward
										__lowerCAmelCase							       =     pipe(**__a      )
										__lowerCAmelCase							       =     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     _UpperCamelCase    (  unittest.TestCase						):
							'''simple docstring'''







							@property
							def      snake_case							(	self      ):
										return (
										    "CUDAExecutionProvider",
										    {
										        "gpu_mem_limit": "15000000000",  # 15GB
										        "arena_extend_strategy": "kSameAsRequested",
										    },
										)

							@property
							def      snake_case							(	self      ):
										__lowerCAmelCase							       =     ort.SessionOptions()
										__lowerCAmelCase							       =     False
										return options

							def      snake_case							(	self      ):
										# using the PNDM scheduler by default
										__lowerCAmelCase							       =     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      )

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

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

										assert image.shape == (1, 5_12, 5_12, 3)
										__lowerCAmelCase							       =     np.array([0.0_4_5_2, 0.0_3_9_0, 0.0_0_8_7, 0.0_3_5_0, 0.0_6_1_7, 0.0_3_6_4, 0.0_5_4_4, 0.0_5_2_3, 0.0_7_2_0]      )

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

							def      snake_case							(	self      ):
										__lowerCAmelCase							       =     DDIMScheduler.from_pretrained(
										    "runwayml/stable-diffusion-v1-5"       ,  subfolder="scheduler"       ,  revision="onnx"      )
										__lowerCAmelCase							       =     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      )

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

										assert image.shape == (1, 5_12, 5_12, 3)
										__lowerCAmelCase							       =     np.array([0.2_8_6_7, 0.1_9_7_4, 0.1_4_8_1, 0.7_2_9_4, 0.7_2_5_1, 0.6_6_6_7, 0.4_1_9_4, 0.5_6_4_2, 0.6_4_8_6]      )

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

							def      snake_case							(	self      ):
										__lowerCAmelCase							       =     LMSDiscreteScheduler.from_pretrained(
										    "runwayml/stable-diffusion-v1-5"       ,  subfolder="scheduler"       ,  revision="onnx"      )
										__lowerCAmelCase							       =     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      )

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

										assert image.shape == (1, 5_12, 5_12, 3)
										__lowerCAmelCase							       =     np.array([0.2_3_0_6, 0.1_9_5_9, 0.1_5_9_3, 0.6_5_4_9, 0.6_3_9_4, 0.5_4_0_8, 0.5_0_6_5, 0.6_0_1_0, 0.6_1_6_1]      )

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

							def      snake_case							(	self      ):
										__lowerCAmelCase							       =     0

										def test_callback_fn(__a       ,  __a       ,  __a      ) -> None:
													__lowerCAmelCase							       =     True
													nonlocal number_of_steps
													number_of_steps += 1
													if step == 0:
																assert latents.shape == (1, 4, 64, 64)
																__lowerCAmelCase							       =     latents[0, -3:, -3:, -1]
																__lowerCAmelCase							       =     np.array(
																    [-0.6_7_7_2, -0.3_8_3_5, -1.2_4_5_6, 0.1_9_0_5, -1.0_9_7_4, 0.6_9_6_7, -1.9_3_5_3, 0.0_1_7_8, 1.0_1_6_7]      )

																assert np.abs(latents_slice.flatten() - expected_slice      ).max() < 1e-3
													elif step == 5:
																assert latents.shape == (1, 4, 64, 64)
																__lowerCAmelCase							       =     latents[0, -3:, -3:, -1]
																__lowerCAmelCase							       =     np.array(
																    [-0.3_3_5_1, 0.2_2_4_1, -0.1_8_3_7, -0.2_3_2_5, -0.6_5_7_7, 0.3_3_9_3, -0.0_2_4_1, 0.5_8_9_9, 1.3_8_7_5]      )

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

										__lowerCAmelCase							       =     False

										__lowerCAmelCase							       =     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      )

										__lowerCAmelCase							       =     "Andromeda galaxy in a bottle"

										__lowerCAmelCase							       =     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      snake_case							(	self      ):
										__lowerCAmelCase							       =     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

										__lowerCAmelCase							       =     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      )
													__lowerCAmelCase							       =     OnnxStableDiffusionPipeline.from_pretrained(__a      )

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








"""simple docstring"""


from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available


A	:  Union[str, Any]            =     {
    "configuration_timesformer": ["TIMESFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP", "TimesformerConfig"],
}

try:
	if not is_torch_available():
		raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
	pass
else:
	A	:  Optional[int]            =     [
	    "TIMESFORMER_PRETRAINED_MODEL_ARCHIVE_LIST",
	    "TimesformerModel",
	    "TimesformerForVideoClassification",
	    "TimesformerPreTrainedModel",
	]

if TYPE_CHECKING:
	from .configuration_timesformer import TIMESFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP, TimesformerConfig

	try:
		if not is_torch_available():
			raise OptionalDependencyNotAvailable()
	except OptionalDependencyNotAvailable:
		pass
	else:
		from .modeling_timesformer import (
		    TIMESFORMER_PRETRAINED_MODEL_ARCHIVE_LIST,
		    TimesformerForVideoClassification,
		    TimesformerModel,
		    TimesformerPreTrainedModel,
		)

else:
	import sys

	A	:  List[str]            =     _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)




