





'''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   A_	(  lowerCAmelCase_							,						unittest.TestCase			):
 _lowerCamelCase  : str           =							"""hf-internal-testing/tiny-random-OnnxStableDiffusionPipeline"""




 def      lowercase	(					self       :    str	,			snake_case_       :    List[str]=0    ):
      _UpperCAmelCase     							=   np.random.RandomState(snake_case_    )
      _UpperCAmelCase     							=   {
          "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       :    Any    ):
      _UpperCAmelCase     							=   OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint	,			provider="CPUExecutionProvider"    )
      pipe.set_progress_bar_config(disable=snake_case_    )

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

      assert image.shape == (1, 1_2_8, 1_2_8, 3)
      _UpperCAmelCase     							=   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      lowercase	(					self       :    int    ):
      _UpperCAmelCase     							=   OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint	,			provider="CPUExecutionProvider"    )
      _UpperCAmelCase     							=   PNDMScheduler.from_config(pipe.scheduler.config	,			skip_prk_steps=snake_case_    )
      pipe.set_progress_bar_config(disable=snake_case_    )

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

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

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

      assert image.shape == (1, 1_2_8, 1_2_8, 3)
      _UpperCAmelCase     							=   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      lowercase	(					self       :    List[Any]    ):
      _UpperCAmelCase     							=   OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint	,			provider="CPUExecutionProvider"    )
      _UpperCAmelCase     							=   EulerDiscreteScheduler.from_config(pipe.scheduler.config    )
      pipe.set_progress_bar_config(disable=snake_case_    )

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

      assert image.shape == (1, 1_2_8, 1_2_8, 3)
      _UpperCAmelCase     							=   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      lowercase	(					self       :    Any    ):
      _UpperCAmelCase     							=   OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint	,			provider="CPUExecutionProvider"    )
      _UpperCAmelCase     							=   EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config    )
      pipe.set_progress_bar_config(disable=snake_case_    )

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

      assert image.shape == (1, 1_2_8, 1_2_8, 3)
      _UpperCAmelCase     							=   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      lowercase	(					self       :    Optional[int]    ):
      _UpperCAmelCase     							=   OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint	,			provider="CPUExecutionProvider"    )
      _UpperCAmelCase     							=   DPMSolverMultistepScheduler.from_config(pipe.scheduler.config    )
      pipe.set_progress_bar_config(disable=snake_case_    )

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

      assert image.shape == (1, 1_2_8, 1_2_8, 3)
      _UpperCAmelCase     							=   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      lowercase	(					self       :    Optional[int]    ):
      _UpperCAmelCase     							=   OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint	,			provider="CPUExecutionProvider"    )
      pipe.set_progress_bar_config(disable=snake_case_    )

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

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

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

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

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

      _UpperCAmelCase     							=   prompt_embeds

      # forward
      _UpperCAmelCase     							=   pipe(**snake_case_    )
      _UpperCAmelCase     							=   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]    ):
      _UpperCAmelCase     							=   OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint	,			provider="CPUExecutionProvider"    )
      pipe.set_progress_bar_config(disable=snake_case_    )

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

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

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

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

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

      _UpperCAmelCase							,  _UpperCAmelCase     							=   embeds

      # forward
      _UpperCAmelCase     							=   pipe(**snake_case_    )
      _UpperCAmelCase     							=   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   A_	(  unittest.TestCase			):




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




 @property
 def      lowercase	(					self       :    int    ):
      _UpperCAmelCase     							=   ort.SessionOptions()
      _UpperCAmelCase     							=   False
      return options




 def      lowercase	(					self       :    Any    ):
      # using the PNDM scheduler by default
      _UpperCAmelCase     							=   OnnxStableDiffusionPipeline.from_pretrained(
          "CompVis/stable-diffusion-v1-4"	,			revision="onnx"	,			safety_checker=snake_case_	,			feature_extractor=snake_case_	,			provider=self.gpu_provider	,			sess_options=self.gpu_options	,			)
      sd_pipe.set_progress_bar_config(disable=snake_case_    )

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

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

      assert image.shape == (1, 5_1_2, 5_1_2, 3)
      _UpperCAmelCase     							=   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      lowercase	(					self       :    List[str]    ):
      _UpperCAmelCase     							=   DDIMScheduler.from_pretrained(
          "runwayml/stable-diffusion-v1-5"	,			subfolder="scheduler"	,			revision="onnx"    )
      _UpperCAmelCase     							=   OnnxStableDiffusionPipeline.from_pretrained(
          "runwayml/stable-diffusion-v1-5"	,			revision="onnx"	,			scheduler=snake_case_	,			safety_checker=snake_case_	,			feature_extractor=snake_case_	,			provider=self.gpu_provider	,			sess_options=self.gpu_options	,			)
      sd_pipe.set_progress_bar_config(disable=snake_case_    )

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

      assert image.shape == (1, 5_1_2, 5_1_2, 3)
      _UpperCAmelCase     							=   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      lowercase	(					self       :    List[str]    ):
      _UpperCAmelCase     							=   LMSDiscreteScheduler.from_pretrained(
          "runwayml/stable-diffusion-v1-5"	,			subfolder="scheduler"	,			revision="onnx"    )
      _UpperCAmelCase     							=   OnnxStableDiffusionPipeline.from_pretrained(
          "runwayml/stable-diffusion-v1-5"	,			revision="onnx"	,			scheduler=snake_case_	,			safety_checker=snake_case_	,			feature_extractor=snake_case_	,			provider=self.gpu_provider	,			sess_options=self.gpu_options	,			)
      sd_pipe.set_progress_bar_config(disable=snake_case_    )

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

      assert image.shape == (1, 5_1_2, 5_1_2, 3)
      _UpperCAmelCase     							=   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      lowercase	(					self       :    Optional[int]    ):
      _UpperCAmelCase     							=   0

      def test_callback_fn(snake_case_       :    int	,			snake_case_       :    int	,			snake_case_       :    np.ndarray    ) -> None:
           _UpperCAmelCase     							=   True
           nonlocal number_of_steps
           number_of_steps += 1
           if step == 0:
                assert latents.shape == (1, 4, 6_4, 6_4)
                _UpperCAmelCase     							=   latents[0, -3:, -3:, -1]
                _UpperCAmelCase     							=   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, 6_4, 6_4)
                _UpperCAmelCase     							=   latents[0, -3:, -3:, -1]
                _UpperCAmelCase     							=   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

      _UpperCAmelCase     							=   False

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

      _UpperCAmelCase     							=   "Andromeda galaxy in a bottle"

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




 def      lowercase	(					self       :    Dict    ):
      _UpperCAmelCase     							=   OnnxStableDiffusionPipeline.from_pretrained(
          "runwayml/stable-diffusion-v1-5"	,			revision="onnx"	,			safety_checker=snake_case_	,			feature_extractor=snake_case_	,			provider=self.gpu_provider	,			sess_options=self.gpu_options	,			)
      assert isinstance(snake_case_	,			snake_case_    )
      assert pipe.safety_checker is None

      _UpperCAmelCase     							=   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(snake_case_    )
           _UpperCAmelCase     							=   OnnxStableDiffusionPipeline.from_pretrained(snake_case_    )

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






'''simple docstring'''







import unittest

import numpy as np
import torch

from diffusers import PNDMPipeline, PNDMScheduler, UNetaDModel
from diffusers.utils.testing_utils import enable_full_determinism, require_torch, slow, torch_device


enable_full_determinism()







class UpperCAmelCase__	( unittest.TestCase	):





  """simple docstring"""

  @property
  def   __lowercase		(		self      :     List[str]       ):


     '''simple docstring'''






     torch.manual_seed(0       )
     _a				:	int			=   UNetaDModel(
         block_out_channels=(32, 64)							,layers_per_block=2							,sample_size=32							,in_channels=3							,out_channels=3							,down_block_types=('DownBlock2D', 'AttnDownBlock2D')							,up_block_types=('AttnUpBlock2D', 'UpBlock2D')							,)
     return model


  def   __lowercase		(		self      :     Optional[Any]       ):


     '''simple docstring'''






     _a				:	str			=   self.dummy_uncond_unet
     _a				:	int			=   PNDMScheduler()

     _a				:	str			=   PNDMPipeline(unet=_a							,scheduler=_a       )
     pndm.to(_a       )
     pndm.set_progress_bar_config(disable=_a       )

     _a				:	Optional[int]			=   torch.manual_seed(0       )
     _a				:	Optional[Any]			=   pndm(generator=_a							,num_inference_steps=20							,output_type='numpy'       ).images

     _a				:	List[str]			=   torch.manual_seed(0       )
     _a				:	Any			=   pndm(generator=_a							,num_inference_steps=20							,output_type='numpy'							,return_dict=_a       )[0]

     _a				:	List[Any]			=   image[0, -3:, -3:, -1]
     _a				:	Any			=   image_from_tuple[0, -3:, -3:, -1]

     assert image.shape == (1, 32, 32, 3)
     _a				:	List[Any]			=   np.array([1.0, 1.0, 0.0, 1.0, 0.0, 1.0, 0.0, 0.0, 0.0]       )

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







@slow
@require_torch
class UpperCAmelCase__	( unittest.TestCase	):





  """simple docstring"""

  def   __lowercase		(		self      :     Tuple       ):


     '''simple docstring'''






     _a				:	List[str]			=   'google/ddpm-cifar10-32'

     _a				:	str			=   UNetaDModel.from_pretrained(_a       )
     _a				:	Union[str, Any]			=   PNDMScheduler()

     _a				:	Tuple			=   PNDMPipeline(unet=_a							,scheduler=_a       )
     pndm.to(_a       )
     pndm.set_progress_bar_config(disable=_a       )
     _a				:	str			=   torch.manual_seed(0       )
     _a				:	Optional[Any]			=   pndm(generator=_a							,output_type='numpy'       ).images

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

     assert image.shape == (1, 32, 32, 3)
     _a				:	Tuple			=   np.array([0.1564, 0.1_4645, 0.1406, 0.1_4715, 0.1_2425, 0.1_4045, 0.1_3115, 0.1_2175, 0.125]       )

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



