




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 		lowercase							(     UpperCamelCase__,unittest.TestCase	):
     _a											=		"hf-internal-testing/tiny-random-OnnxStableDiffusionPipeline"

     def       a__			(						self	,  _a=0       )		->							List[str]:
           _A							:  Dict    =		np.random.RandomState(_a       )
           _A							:  Any    =		{
               """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       a__			(						self       )		->							Optional[int]:
           _A							:  Dict    =		OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint	,  provider="""CPUExecutionProvider"""       )
           pipe.set_progress_bar_config(disable=_a       )

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

           assert image.shape == (1, 128, 128, 3)
           _A							:  Any    =		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       a__			(						self       )		->							List[str]:
           _A							:  List[Any]    =		OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint	,  provider="""CPUExecutionProvider"""       )
           _A							:  List[str]    =		PNDMScheduler.from_config(pipe.scheduler.config	,  skip_prk_steps=_a       )
           pipe.set_progress_bar_config(disable=_a       )

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

           assert image.shape == (1, 128, 128, 3)
           _A							:  List[Any]    =		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       a__			(						self       )		->							List[str]:
           _A							:  Union[str, Any]    =		OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint	,  provider="""CPUExecutionProvider"""       )
           _A							:  List[str]    =		LMSDiscreteScheduler.from_config(pipe.scheduler.config       )
           pipe.set_progress_bar_config(disable=_a       )

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

           assert image.shape == (1, 128, 128, 3)
           _A							:  Optional[int]    =		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       a__			(						self       )		->							Union[str, Any]:
           _A							:  Dict    =		OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint	,  provider="""CPUExecutionProvider"""       )
           _A							:  Optional[Any]    =		EulerDiscreteScheduler.from_config(pipe.scheduler.config       )
           pipe.set_progress_bar_config(disable=_a       )

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

           assert image.shape == (1, 128, 128, 3)
           _A							:  Any    =		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       a__			(						self       )		->							Tuple:
           _A							:  List[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							:  Any    =		self.get_dummy_inputs()
           _A							:  int    =		pipe(**_a       ).images
           _A							:  int    =		image[0, -3:, -3:, -1]

           assert image.shape == (1, 128, 128, 3)
           _A							:  Tuple    =		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       a__			(						self       )		->							str:
           _A							:  int    =		OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint	,  provider="""CPUExecutionProvider"""       )
           _A							:  Any    =		DPMSolverMultistepScheduler.from_config(pipe.scheduler.config       )
           pipe.set_progress_bar_config(disable=_a       )

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

           assert image.shape == (1, 128, 128, 3)
           _A							:  Optional[int]    =		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       a__			(						self       )		->							Dict:
           _A							:  List[Any]    =		OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint	,  provider="""CPUExecutionProvider"""       )
           pipe.set_progress_bar_config(disable=_a       )

           _A							:  str    =		self.get_dummy_inputs()
           _A							:  List[str]    =		3 * [inputs["""prompt"""]]

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

           _A							:  Optional[int]    =		self.get_dummy_inputs()
           _A							:  Optional[Any]    =		3 * [inputs.pop("""prompt"""       )]

           _A							:  Dict    =		pipe.tokenizer(
               _a	,  padding="""max_length"""	,  max_length=pipe.tokenizer.model_max_length	,  truncation=_a	,  return_tensors="""np"""	,  )
           _A							:  Union[str, Any]    =		text_inputs["""input_ids"""]

           _A							:  Union[str, Any]    =		pipe.text_encoder(input_ids=text_inputs.astype(np.intaa       )       )[0]

           _A							:  int    =		prompt_embeds

           # forward
           _A							:  Dict    =		pipe(**_a       )
           _A							:  Optional[Any]    =		output.images[0, -3:, -3:, -1]

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



     def       a__			(						self       )		->							List[str]:
           _A							:  Optional[Any]    =		OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint	,  provider="""CPUExecutionProvider"""       )
           pipe.set_progress_bar_config(disable=_a       )

           _A							:  Optional[int]    =		self.get_dummy_inputs()
           _A							:  Dict    =		3 * ["""this is a negative prompt"""]
           _A							:  List[Any]    =		negative_prompt
           _A							:  int    =		3 * [inputs["""prompt"""]]

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

           _A							:  Optional[int]    =		self.get_dummy_inputs()
           _A							:  List[Any]    =		3 * [inputs.pop("""prompt"""       )]

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

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

           _A    ,     _A							:  List[Any]    =		embeds

           # forward
           _A							:  Tuple    =		pipe(**_a       )
           _A							:  Optional[Any]    =		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 		lowercase							(     unittest.TestCase	):

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

     @property
     def       a__			(						self       )		->							Dict:
           _A							:  Any    =		ort.SessionOptions()
           _A							:  int    =		False
           return options

     def       a__			(						self       )		->							List[Any]:
           # using the PNDM scheduler by default
           _A							:  str    =		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							:  Optional[Any]    =		"""A painting of a squirrel eating a burger"""
           np.random.seed(0       )
           _A							:  Union[str, Any]    =		sd_pipe([prompt]	,  guidance_scale=6.0	,  num_inference_steps=10	,  output_type="""np"""       )
           _A							:  Union[str, Any]    =		output.images

           _A							:  Union[str, Any]    =		image[0, -3:, -3:, -1]

           assert image.shape == (1, 512, 512, 3)
           _A							:  int    =		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       a__			(						self       )		->							Union[str, Any]:
           _A							:  Optional[int]    =		DDIMScheduler.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							:  Optional[int]    =		"""open neural network exchange"""
           _A							:  Optional[Any]    =		np.random.RandomState(0       )
           _A							:  Any    =		sd_pipe([prompt]	,  guidance_scale=7.5	,  num_inference_steps=10	,  generator=_a	,  output_type="""np"""       )
           _A							:  Dict    =		output.images
           _A							:  str    =		image[0, -3:, -3:, -1]

           assert image.shape == (1, 512, 512, 3)
           _A							:  Union[str, 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       a__			(						self       )		->							List[str]:
           _A							:  str    =		LMSDiscreteScheduler.from_pretrained(
               """runwayml/stable-diffusion-v1-5"""	,  subfolder="""scheduler"""	,  revision="""onnx"""       )
           _A							:  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       )

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

           assert image.shape == (1, 512, 512, 3)
           _A							:  Union[str, Any]    =		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       a__			(						self       )		->							int:
           _A							:  List[Any]    =		0

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

           _A							:  List[Any]    =		False

           _A							:  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							:  List[Any]    =		"""Andromeda galaxy in a bottle"""

           _A							:  Optional[int]    =		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       a__			(						self       )		->							Union[str, Any]:
           _A							:  Dict    =		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							:  Dict    =		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[Any]    =		OnnxStableDiffusionPipeline.from_pretrained(_a       )

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







from collections.abc import Callable
from math import pi, sqrt
from random import uniform
from statistics import mean







def        lowerCAmelCase_			(  snake_case_  ):

      # A local function to see if a dot lands in the circle.
      def is_in_circle(snake_case_,snake_case_  ) -> bool:
            _A							:  List[str]    =		sqrt((x**2) + (y**2)  )
            # Our circle has a radius of 1, so a distance
            # greater than 1 would land outside the circle.
            return distance_from_centre <= 1

      # The proportion of guesses that landed in the circle
      _A							:  Optional[int]    =		mean(
          int(is_in_circle(uniform(-1.0,1.0  ),uniform(-1.0,1.0  )  )  )
          for _ in range(snake_case_  )  )
      # The ratio of the area for circle to square is pi/4.
      _A							:  List[str]    =		proportion * 4
      print(f'''The estimated value of pi is {pi_estimate}'''  )
      print(f'''The numpy value of pi is {pi}'''  )
      print(f'''The total error is {abs(pi - pi_estimate  )}'''  )







def        lowerCAmelCase_			(  snake_case_,snake_case_,snake_case_ = 0.0,snake_case_ = 1.0,):
      return mean(
          function_to_integrate(uniform(snake_case_,snake_case_  )  ) for _ in range(snake_case_  )  ) * (max_value - min_value)







def        lowerCAmelCase_			(  snake_case_,snake_case_ = 0.0,snake_case_ = 1.0  ):

      def identity_function(snake_case_  ) -> float:
            return x

      _A							:  Any    =		area_under_curve_estimator(
          snake_case_,snake_case_,snake_case_,snake_case_  )
      _A							:  Tuple    =		(max_value * max_value - min_value * min_value) / 2

      print("""******************"""  )
      print(f'''Estimating area under y=x where x varies from {min_value} to {max_value}'''  )
      print(f'''Estimated value is {estimated_value}'''  )
      print(f'''Expected value is {expected_value}'''  )
      print(f'''Total error is {abs(estimated_value - expected_value  )}'''  )
      print("""******************"""  )







def        lowerCAmelCase_			(  snake_case_  ):

      def function_to_integrate(snake_case_  ) -> float:
            return sqrt(4.0 - x * x  )

      _A							:  Optional[int]    =		area_under_curve_estimator(
          snake_case_,snake_case_,0.0,2.0  )

      print("""******************"""  )
      print("""Estimating pi using area_under_curve_estimator"""  )
      print(f'''Estimated value is {estimated_value}'''  )
      print(f'''Expected value is {pi}'''  )
      print(f'''Total error is {abs(estimated_value - pi  )}'''  )
      print("""******************"""  )


if __name__ == "__main__":
  import doctest

  doctest.testmod()

