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						):
  lowercase_			=     'hf-internal-testing/tiny-random-OnnxStableDiffusionPipeline'




  def 	__lowercase(      self     :	int     , a_     :	Any=0				)->       Any:




       """simple docstring"""
       SCREAMING_SNAKE_CASE__				:				Any			=		np.random.RandomState(a_				)
       SCREAMING_SNAKE_CASE__				:				int			=		{
           '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 	__lowercase(      self     :	int				)->       List[str]:




       """simple docstring"""
       SCREAMING_SNAKE_CASE__				:				Optional[Any]			=		OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint     , provider='CPUExecutionProvider'				)
       pipe.set_progress_bar_config(disable=a_				)

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

       assert image.shape == (1, 128, 128, 3)
       SCREAMING_SNAKE_CASE__				:				Union[str, Any]			=		np.array([0.6_5072, 0.5_8492, 0.4_8219, 0.5_5521, 0.5_3180, 0.5_5939, 0.5_0697, 0.3_9800, 0.4_6455]				)

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




  def 	__lowercase(      self     :	Tuple				)->       Dict:




       """simple docstring"""
       SCREAMING_SNAKE_CASE__				:				Optional[Any]			=		OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint     , provider='CPUExecutionProvider'				)
       SCREAMING_SNAKE_CASE__				:				Union[str, Any]			=		PNDMScheduler.from_config(pipe.scheduler.config     , skip_prk_steps=a_				)
       pipe.set_progress_bar_config(disable=a_				)

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

       assert image.shape == (1, 128, 128, 3)
       SCREAMING_SNAKE_CASE__				:				Union[str, Any]			=		np.array([0.6_5863, 0.5_9425, 0.4_9326, 0.5_6313, 0.5_3875, 0.5_6627, 0.5_1065, 0.3_9777, 0.4_6330]				)

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




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




       """simple docstring"""
       SCREAMING_SNAKE_CASE__				:				Tuple			=		OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint     , provider='CPUExecutionProvider'				)
       SCREAMING_SNAKE_CASE__				:				int			=		LMSDiscreteScheduler.from_config(pipe.scheduler.config				)
       pipe.set_progress_bar_config(disable=a_				)

       SCREAMING_SNAKE_CASE__				:				Dict			=		self.get_dummy_inputs()
       SCREAMING_SNAKE_CASE__				:				Dict			=		pipe(**a_				).images
       SCREAMING_SNAKE_CASE__				:				List[str]			=		image[0, -3:, -3:, -1]

       assert image.shape == (1, 128, 128, 3)
       SCREAMING_SNAKE_CASE__				:				int			=		np.array([0.5_3755, 0.6_0786, 0.4_7402, 0.4_9488, 0.5_1869, 0.4_9819, 0.4_7985, 0.3_8957, 0.4_4279]				)

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




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




       """simple docstring"""
       SCREAMING_SNAKE_CASE__				:				Dict			=		OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint     , provider='CPUExecutionProvider'				)
       SCREAMING_SNAKE_CASE__				:				List[str]			=		EulerDiscreteScheduler.from_config(pipe.scheduler.config				)
       pipe.set_progress_bar_config(disable=a_				)

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

       assert image.shape == (1, 128, 128, 3)
       SCREAMING_SNAKE_CASE__				:				str			=		np.array([0.5_3755, 0.6_0786, 0.4_7402, 0.4_9488, 0.5_1869, 0.4_9819, 0.4_7985, 0.3_8957, 0.4_4279]				)

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




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




       """simple docstring"""
       SCREAMING_SNAKE_CASE__				:				Dict			=		OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint     , provider='CPUExecutionProvider'				)
       SCREAMING_SNAKE_CASE__				:				Any			=		EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config				)
       pipe.set_progress_bar_config(disable=a_				)

       SCREAMING_SNAKE_CASE__				:				Tuple			=		self.get_dummy_inputs()
       SCREAMING_SNAKE_CASE__				:				Optional[Any]			=		pipe(**a_				).images
       SCREAMING_SNAKE_CASE__				:				Union[str, Any]			=		image[0, -3:, -3:, -1]

       assert image.shape == (1, 128, 128, 3)
       SCREAMING_SNAKE_CASE__				:				Optional[Any]			=		np.array([0.5_3817, 0.6_0812, 0.4_7384, 0.4_9530, 0.5_1894, 0.4_9814, 0.4_7984, 0.3_8958, 0.4_4271]				)

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




  def 	__lowercase(      self     :	Dict				)->       List[Any]:




       """simple docstring"""
       SCREAMING_SNAKE_CASE__				:				List[str]			=		OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint     , provider='CPUExecutionProvider'				)
       SCREAMING_SNAKE_CASE__				:				int			=		DPMSolverMultistepScheduler.from_config(pipe.scheduler.config				)
       pipe.set_progress_bar_config(disable=a_				)

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

       assert image.shape == (1, 128, 128, 3)
       SCREAMING_SNAKE_CASE__				:				List[Any]			=		np.array([0.5_3895, 0.6_0808, 0.4_7933, 0.4_9608, 0.5_1886, 0.4_9950, 0.4_8053, 0.3_8957, 0.4_4200]				)

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




  def 	__lowercase(      self     :	List[str]				)->       Any:




       """simple docstring"""
       SCREAMING_SNAKE_CASE__				:				List[Any]			=		OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint     , provider='CPUExecutionProvider'				)
       pipe.set_progress_bar_config(disable=a_				)

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

       # forward
       SCREAMING_SNAKE_CASE__				:				Optional[int]			=		pipe(**a_				)
       SCREAMING_SNAKE_CASE__				:				Optional[Any]			=		output.images[0, -3:, -3:, -1]

       SCREAMING_SNAKE_CASE__				:				int			=		self.get_dummy_inputs()
       SCREAMING_SNAKE_CASE__				:				List[Any]			=		3 * [inputs.pop('prompt'				)]

       SCREAMING_SNAKE_CASE__				:				List[Any]			=		pipe.tokenizer(
           a_     , padding='max_length'     , max_length=pipe.tokenizer.model_max_length     , truncation=a_     , return_tensors='np'     , )
       SCREAMING_SNAKE_CASE__				:				Any			=		text_inputs['input_ids']

       SCREAMING_SNAKE_CASE__				:				int			=		pipe.text_encoder(input_ids=text_inputs.astype(np.intaa				)				)[0]

       SCREAMING_SNAKE_CASE__				:				Optional[int]			=		prompt_embeds

       # forward
       SCREAMING_SNAKE_CASE__				:				Union[str, Any]			=		pipe(**a_				)
       SCREAMING_SNAKE_CASE__				:				Optional[Any]			=		output.images[0, -3:, -3:, -1]

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




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




       """simple docstring"""
       SCREAMING_SNAKE_CASE__				:				Any			=		OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint     , provider='CPUExecutionProvider'				)
       pipe.set_progress_bar_config(disable=a_				)

       SCREAMING_SNAKE_CASE__				:				List[str]			=		self.get_dummy_inputs()
       SCREAMING_SNAKE_CASE__				:				List[Any]			=		3 * ['this is a negative prompt']
       SCREAMING_SNAKE_CASE__				:				Dict			=		negative_prompt
       SCREAMING_SNAKE_CASE__				:				Dict			=		3 * [inputs['prompt']]

       # forward
       SCREAMING_SNAKE_CASE__				:				Tuple			=		pipe(**a_				)
       SCREAMING_SNAKE_CASE__				:				Union[str, Any]			=		output.images[0, -3:, -3:, -1]

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

       SCREAMING_SNAKE_CASE__				:				Dict			=		[]
       for p in [prompt, negative_prompt]:
            SCREAMING_SNAKE_CASE__				:				Dict			=		pipe.tokenizer(
                a_     , padding='max_length'     , max_length=pipe.tokenizer.model_max_length     , truncation=a_     , return_tensors='np'     , )
            SCREAMING_SNAKE_CASE__				:				Union[str, Any]			=		text_inputs['input_ids']

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

       SCREAMING_SNAKE_CASE__			,							SCREAMING_SNAKE_CASE__				:				Any			=		embeds

       # forward
       SCREAMING_SNAKE_CASE__				:				Union[str, Any]			=		pipe(**a_				)
       SCREAMING_SNAKE_CASE__				:				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       snake_case					(      unittest.TestCase						):




  @property
  def 	__lowercase(      self     :	Tuple				)->       Tuple:




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




  @property
  def 	__lowercase(      self     :	Union[str, Any]				)->       Optional[Any]:




       """simple docstring"""
       SCREAMING_SNAKE_CASE__				:				Tuple			=		ort.SessionOptions()
       SCREAMING_SNAKE_CASE__				:				int			=		False
       return options




  def 	__lowercase(      self     :	Optional[int]				)->       Tuple:




       """simple docstring"""
       # using the PNDM scheduler by default
       SCREAMING_SNAKE_CASE__				:				Union[str, 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_				)

       SCREAMING_SNAKE_CASE__				:				Union[str, Any]			=		'A painting of a squirrel eating a burger'
       np.random.seed(0				)
       SCREAMING_SNAKE_CASE__				:				Tuple			=		sd_pipe([prompt]     , guidance_scale=6.0     , num_inference_steps=10     , output_type='np'				)
       SCREAMING_SNAKE_CASE__				:				Optional[int]			=		output.images

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

       assert image.shape == (1, 512, 512, 3)
       SCREAMING_SNAKE_CASE__				:				Any			=		np.array([0.0452, 0.0390, 0.0087, 0.0350, 0.0617, 0.0364, 0.0544, 0.0523, 0.0720]				)

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




  def 	__lowercase(      self     :	Optional[int]				)->       Optional[int]:




       """simple docstring"""
       SCREAMING_SNAKE_CASE__				:				str			=		DDIMScheduler.from_pretrained(
           'runwayml/stable-diffusion-v1-5'     , subfolder='scheduler'     , revision='onnx'				)
       SCREAMING_SNAKE_CASE__				:				Tuple			=		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_				)

       SCREAMING_SNAKE_CASE__				:				Any			=		'open neural network exchange'
       SCREAMING_SNAKE_CASE__				:				List[str]			=		np.random.RandomState(0				)
       SCREAMING_SNAKE_CASE__				:				Any			=		sd_pipe([prompt]     , guidance_scale=7.5     , num_inference_steps=10     , generator=a_     , output_type='np'				)
       SCREAMING_SNAKE_CASE__				:				Any			=		output.images
       SCREAMING_SNAKE_CASE__				:				List[str]			=		image[0, -3:, -3:, -1]

       assert image.shape == (1, 512, 512, 3)
       SCREAMING_SNAKE_CASE__				:				Optional[Any]			=		np.array([0.2867, 0.1974, 0.1481, 0.7294, 0.7251, 0.6667, 0.4194, 0.5642, 0.6486]				)

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




  def 	__lowercase(      self     :	Any				)->       Any:




       """simple docstring"""
       SCREAMING_SNAKE_CASE__				:				str			=		LMSDiscreteScheduler.from_pretrained(
           'runwayml/stable-diffusion-v1-5'     , subfolder='scheduler'     , revision='onnx'				)
       SCREAMING_SNAKE_CASE__				:				str			=		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_				)

       SCREAMING_SNAKE_CASE__				:				Optional[int]			=		'open neural network exchange'
       SCREAMING_SNAKE_CASE__				:				str			=		np.random.RandomState(0				)
       SCREAMING_SNAKE_CASE__				:				Optional[Any]			=		sd_pipe([prompt]     , guidance_scale=7.5     , num_inference_steps=10     , generator=a_     , output_type='np'				)
       SCREAMING_SNAKE_CASE__				:				Dict			=		output.images
       SCREAMING_SNAKE_CASE__				:				Any			=		image[0, -3:, -3:, -1]

       assert image.shape == (1, 512, 512, 3)
       SCREAMING_SNAKE_CASE__				:				Dict			=		np.array([0.2306, 0.1959, 0.1593, 0.6549, 0.6394, 0.5408, 0.5065, 0.6010, 0.6161]				)

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




  def 	__lowercase(      self     :	Optional[int]				)->       Optional[int]:




       """simple docstring"""
       SCREAMING_SNAKE_CASE__				:				Union[str, Any]			=		0

       def test_callback_fn(a_     :	int     , a_     :	int     , a_     :	np.ndarray				) -> None:
            SCREAMING_SNAKE_CASE__				:				Any			=		True
            nonlocal number_of_steps
            number_of_steps += 1
            if step == 0:
                 assert latents.shape == (1, 4, 64, 64)
                 SCREAMING_SNAKE_CASE__				:				Optional[Any]			=		latents[0, -3:, -3:, -1]
                 SCREAMING_SNAKE_CASE__				:				str			=		np.array(
                     [-0.6772, -0.3835, -1.2456, 0.1905, -1.0974, 0.6967, -1.9353, 0.0178, 1.0167]				)

                 assert np.abs(latents_slice.flatten() - expected_slice				).max() < 1e-3
            elif step == 5:
                 assert latents.shape == (1, 4, 64, 64)
                 SCREAMING_SNAKE_CASE__				:				List[Any]			=		latents[0, -3:, -3:, -1]
                 SCREAMING_SNAKE_CASE__				:				Tuple			=		np.array(
                     [-0.3351, 0.2241, -0.1837, -0.2325, -0.6577, 0.3393, -0.0241, 0.5899, 1.3875]				)

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

       SCREAMING_SNAKE_CASE__				:				Optional[Any]			=		False

       SCREAMING_SNAKE_CASE__				:				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_				)

       SCREAMING_SNAKE_CASE__				:				Optional[int]			=		'Andromeda galaxy in a bottle'

       SCREAMING_SNAKE_CASE__				:				Union[str, Any]			=		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 	__lowercase(      self     :	str				)->       Union[str, Any]:




       """simple docstring"""
       SCREAMING_SNAKE_CASE__				:				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     , )
       assert isinstance(a_     , a_				)
       assert pipe.safety_checker is None

       SCREAMING_SNAKE_CASE__				:				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_				)
            SCREAMING_SNAKE_CASE__				:				Dict			=		OnnxStableDiffusionPipeline.from_pretrained(a_				)

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




def     _a				(      lowercase__  :			int = 1_00_00_00 ):







     '''simple docstring'''


     SCREAMING_SNAKE_CASE__				:				str			=		[i - 1 for i in range(limit + 1 )]

     for i in range(2     ,  limit + 1 ):
          if phi[i] == i - 1:
               for j in range(2 * i     ,  limit + 1     ,  lowercase__ ):
                    phi[j] -= phi[j] // i

     return sum(phi[2 : limit + 1] )


if __name__ == "__main__":
      print(solution())



