




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'''

     __snake_case   :     List[str]		     =      "hf-internal-testing/tiny-random-OnnxStableDiffusionPipeline"






     def     SCREAMING_SNAKE_CASE__  (							self						:   Union[str, Any]						,lowerCamelCase__						:   List[str]=0	)					->		Optional[int]:




           '''simple docstring'''


           SCREAMING_SNAKE_CASE		=						np.random.RandomState(_a	)
           SCREAMING_SNAKE_CASE		=						{
               "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     SCREAMING_SNAKE_CASE__  (							self						:   Any	)					->		str:




           '''simple docstring'''


           SCREAMING_SNAKE_CASE		=						OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint						,provider="""CPUExecutionProvider"""	)
           pipe.set_progress_bar_config(disable=_a	)

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

           assert image.shape == (1, 128, 128, 3)
           SCREAMING_SNAKE_CASE		=						np.array([0.65072, 0.58492, 0.48219, 0.55521, 0.53180, 0.55939, 0.50697, 0.39800, 0.46455]	)

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






     def     SCREAMING_SNAKE_CASE__  (							self						:   str	)					->		Tuple:




           '''simple docstring'''


           SCREAMING_SNAKE_CASE		=						OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint						,provider="""CPUExecutionProvider"""	)
           SCREAMING_SNAKE_CASE		=						PNDMScheduler.from_config(pipe.scheduler.config						,skip_prk_steps=_a	)
           pipe.set_progress_bar_config(disable=_a	)

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

           assert image.shape == (1, 128, 128, 3)
           SCREAMING_SNAKE_CASE		=						np.array([0.65863, 0.59425, 0.49326, 0.56313, 0.53875, 0.56627, 0.51065, 0.39777, 0.46330]	)

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






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




           '''simple docstring'''


           SCREAMING_SNAKE_CASE		=						OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint						,provider="""CPUExecutionProvider"""	)
           SCREAMING_SNAKE_CASE		=						LMSDiscreteScheduler.from_config(pipe.scheduler.config	)
           pipe.set_progress_bar_config(disable=_a	)

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

           assert image.shape == (1, 128, 128, 3)
           SCREAMING_SNAKE_CASE		=						np.array([0.53755, 0.60786, 0.47402, 0.49488, 0.51869, 0.49819, 0.47985, 0.38957, 0.44279]	)

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






     def     SCREAMING_SNAKE_CASE__  (							self						:   Tuple	)					->		Tuple:




           '''simple docstring'''


           SCREAMING_SNAKE_CASE		=						OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint						,provider="""CPUExecutionProvider"""	)
           SCREAMING_SNAKE_CASE		=						EulerDiscreteScheduler.from_config(pipe.scheduler.config	)
           pipe.set_progress_bar_config(disable=_a	)

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

           assert image.shape == (1, 128, 128, 3)
           SCREAMING_SNAKE_CASE		=						np.array([0.53755, 0.60786, 0.47402, 0.49488, 0.51869, 0.49819, 0.47985, 0.38957, 0.44279]	)

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






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




           '''simple docstring'''


           SCREAMING_SNAKE_CASE		=						OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint						,provider="""CPUExecutionProvider"""	)
           SCREAMING_SNAKE_CASE		=						EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config	)
           pipe.set_progress_bar_config(disable=_a	)

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

           assert image.shape == (1, 128, 128, 3)
           SCREAMING_SNAKE_CASE		=						np.array([0.53817, 0.60812, 0.47384, 0.49530, 0.51894, 0.49814, 0.47984, 0.38958, 0.44271]	)

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






     def     SCREAMING_SNAKE_CASE__  (							self						:   Tuple	)					->		List[Any]:




           '''simple docstring'''


           SCREAMING_SNAKE_CASE		=						OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint						,provider="""CPUExecutionProvider"""	)
           SCREAMING_SNAKE_CASE		=						DPMSolverMultistepScheduler.from_config(pipe.scheduler.config	)
           pipe.set_progress_bar_config(disable=_a	)

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

           assert image.shape == (1, 128, 128, 3)
           SCREAMING_SNAKE_CASE		=						np.array([0.53895, 0.60808, 0.47933, 0.49608, 0.51886, 0.49950, 0.48053, 0.38957, 0.44200]	)

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






     def     SCREAMING_SNAKE_CASE__  (							self						:   Dict	)					->		Optional[Any]:




           '''simple docstring'''


           SCREAMING_SNAKE_CASE		=						OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint						,provider="""CPUExecutionProvider"""	)
           pipe.set_progress_bar_config(disable=_a	)

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

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

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

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

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

           SCREAMING_SNAKE_CASE		=						prompt_embeds

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

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






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




           '''simple docstring'''


           SCREAMING_SNAKE_CASE		=						OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint						,provider="""CPUExecutionProvider"""	)
           pipe.set_progress_bar_config(disable=_a	)

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

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

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

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

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

           SCREAMING_SNAKE_CASE		=						embeds

           # forward
           SCREAMING_SNAKE_CASE		=						pipe(**_a	)
           SCREAMING_SNAKE_CASE		=						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     SCREAMING_SNAKE_CASE__  (							self						:   Optional[Any]	)					->		int:




           '''simple docstring'''


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






     @property
     def     SCREAMING_SNAKE_CASE__  (							self						:   Tuple	)					->		Dict:




           '''simple docstring'''


           SCREAMING_SNAKE_CASE		=						ort.SessionOptions()
           SCREAMING_SNAKE_CASE		=						False
           return options






     def     SCREAMING_SNAKE_CASE__  (							self						:   List[str]	)					->		Dict:




           '''simple docstring'''


           SCREAMING_SNAKE_CASE		=						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		=						"A painting of a squirrel eating a burger"
           np.random.seed(0	)
           SCREAMING_SNAKE_CASE		=						sd_pipe([prompt]						,guidance_scale=6.0						,num_inference_steps=10						,output_type="""np"""	)
           SCREAMING_SNAKE_CASE		=						output.images

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

           assert image.shape == (1, 512, 512, 3)
           SCREAMING_SNAKE_CASE		=						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     SCREAMING_SNAKE_CASE__  (							self						:   Optional[Any]	)					->		Union[str, Any]:




           '''simple docstring'''


           SCREAMING_SNAKE_CASE		=						DDIMScheduler.from_pretrained(
               """runwayml/stable-diffusion-v1-5"""						,subfolder="""scheduler"""						,revision="""onnx"""	)
           SCREAMING_SNAKE_CASE		=						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		=						"open neural network exchange"
           SCREAMING_SNAKE_CASE		=						np.random.RandomState(0	)
           SCREAMING_SNAKE_CASE		=						sd_pipe([prompt]						,guidance_scale=7.5						,num_inference_steps=10						,generator=_a						,output_type="""np"""	)
           SCREAMING_SNAKE_CASE		=						output.images
           SCREAMING_SNAKE_CASE		=						image[0, -3:, -3:, -1]

           assert image.shape == (1, 512, 512, 3)
           SCREAMING_SNAKE_CASE		=						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     SCREAMING_SNAKE_CASE__  (							self						:   int	)					->		Dict:




           '''simple docstring'''


           SCREAMING_SNAKE_CASE		=						LMSDiscreteScheduler.from_pretrained(
               """runwayml/stable-diffusion-v1-5"""						,subfolder="""scheduler"""						,revision="""onnx"""	)
           SCREAMING_SNAKE_CASE		=						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		=						"open neural network exchange"
           SCREAMING_SNAKE_CASE		=						np.random.RandomState(0	)
           SCREAMING_SNAKE_CASE		=						sd_pipe([prompt]						,guidance_scale=7.5						,num_inference_steps=10						,generator=_a						,output_type="""np"""	)
           SCREAMING_SNAKE_CASE		=						output.images
           SCREAMING_SNAKE_CASE		=						image[0, -3:, -3:, -1]

           assert image.shape == (1, 512, 512, 3)
           SCREAMING_SNAKE_CASE		=						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     SCREAMING_SNAKE_CASE__  (							self						:   Optional[Any]	)					->		Optional[Any]:




           '''simple docstring'''


           SCREAMING_SNAKE_CASE		=						0

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

           SCREAMING_SNAKE_CASE		=						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		=						"Andromeda galaxy in a bottle"

           SCREAMING_SNAKE_CASE		=						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     SCREAMING_SNAKE_CASE__  (							self						:   str	)					->		List[str]:




           '''simple docstring'''


           SCREAMING_SNAKE_CASE		=						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		=						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		=						OnnxStableDiffusionPipeline.from_pretrained(_a	)

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









from __future__ import annotations



def        __lowercase       (				_SCREAMING_SNAKE_CASE     )			->       list[int]:


      '''simple docstring'''


      SCREAMING_SNAKE_CASE		=						2
      SCREAMING_SNAKE_CASE		=						[]
      while i * i <= n:
            if n % i:
                  i += 1
            else:
                  n //= i
                  factors.append(_SCREAMING_SNAKE_CASE     )
      if n > 1:
            factors.append(_SCREAMING_SNAKE_CASE     )
      return factors


if __name__ == "__main__":
      import doctest

      doctest.testmod()


