

'''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 			lowercase_	(lowerCamelCase__				,       unittest.TestCase      ):




  """simple docstring"""




  SCREAMING_SNAKE_CASE    :  Optional[Any]							=			'hf-internal-testing/tiny-random-OnnxStableDiffusionPipeline'






  def 					SCREAMING_SNAKE_CASE      (      self   :     Optional[Any]					,lowercase__   :     Tuple=0			):
   __lowercase							  = np.random.RandomState(lowercase__			)
   __lowercase							  = {
       '''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   :     Optional[Any]			):
   __lowercase							  = OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint					,provider='''CPUExecutionProvider'''			)
   pipe.set_progress_bar_config(disable=lowercase__			)

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

   assert image.shape == (1, 1_2_8, 1_2_8, 3)
   __lowercase							  = 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 					SCREAMING_SNAKE_CASE      (      self   :     Any			):
   __lowercase							  = OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint					,provider='''CPUExecutionProvider'''			)
   __lowercase							  = PNDMScheduler.from_config(pipe.scheduler.config					,skip_prk_steps=lowercase__			)
   pipe.set_progress_bar_config(disable=lowercase__			)

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

   assert image.shape == (1, 1_2_8, 1_2_8, 3)
   __lowercase							  = 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 					SCREAMING_SNAKE_CASE      (      self   :     Any			):
   __lowercase							  = OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint					,provider='''CPUExecutionProvider'''			)
   __lowercase							  = LMSDiscreteScheduler.from_config(pipe.scheduler.config			)
   pipe.set_progress_bar_config(disable=lowercase__			)

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

   assert image.shape == (1, 1_2_8, 1_2_8, 3)
   __lowercase							  = 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 					SCREAMING_SNAKE_CASE      (      self   :     Optional[int]			):
   __lowercase							  = OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint					,provider='''CPUExecutionProvider'''			)
   __lowercase							  = EulerDiscreteScheduler.from_config(pipe.scheduler.config			)
   pipe.set_progress_bar_config(disable=lowercase__			)

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

   assert image.shape == (1, 1_2_8, 1_2_8, 3)
   __lowercase							  = 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 					SCREAMING_SNAKE_CASE      (      self   :     Tuple			):
   __lowercase							  = OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint					,provider='''CPUExecutionProvider'''			)
   __lowercase							  = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config			)
   pipe.set_progress_bar_config(disable=lowercase__			)

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

   assert image.shape == (1, 1_2_8, 1_2_8, 3)
   __lowercase							  = 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 					SCREAMING_SNAKE_CASE      (      self   :     List[Any]			):
   __lowercase							  = OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint					,provider='''CPUExecutionProvider'''			)
   __lowercase							  = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config			)
   pipe.set_progress_bar_config(disable=lowercase__			)

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

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

   __lowercase							  = self.get_dummy_inputs()
   __lowercase							  = 3 * [inputs['''prompt''']]

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

   __lowercase							  = self.get_dummy_inputs()
   __lowercase							  = 3 * [inputs.pop('''prompt'''			)]

   __lowercase							  = pipe.tokenizer(
       lowercase__					,padding='''max_length'''					,max_length=pipe.tokenizer.model_max_length					,truncation=lowercase__					,return_tensors='''np'''					,)
   __lowercase							  = text_inputs['''input_ids''']

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

   __lowercase							  = prompt_embeds

   # forward
   __lowercase							  = pipe(**lowercase__			)
   __lowercase							  = 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   :     Optional[Any]			):
   __lowercase							  = OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint					,provider='''CPUExecutionProvider'''			)
   pipe.set_progress_bar_config(disable=lowercase__			)

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

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

   __lowercase							  = self.get_dummy_inputs()
   __lowercase							  = 3 * [inputs.pop('''prompt'''			)]

   __lowercase							  = []
   for p in [prompt, negative_prompt]:
    __lowercase							  = pipe.tokenizer(
        lowercase__					,padding='''max_length'''					,max_length=pipe.tokenizer.model_max_length					,truncation=lowercase__					,return_tensors='''np'''					,)
    __lowercase							  = text_inputs['''input_ids''']

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

   __lowercase  ,	__lowercase							  = embeds

   # forward
   __lowercase							  = pipe(**lowercase__			)
   __lowercase							  = 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      ):




  """simple docstring"""




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






  @property
  def 					SCREAMING_SNAKE_CASE      (      self   :     int			):
   __lowercase							  = ort.SessionOptions()
   __lowercase							  = False
   return options






  def 					SCREAMING_SNAKE_CASE      (      self   :     str			):
   # using the PNDM scheduler by default
   __lowercase							  = OnnxStableDiffusionPipeline.from_pretrained(
       '''CompVis/stable-diffusion-v1-4'''					,revision='''onnx'''					,safety_checker=lowercase__					,feature_extractor=lowercase__					,provider=self.gpu_provider					,sess_options=self.gpu_options					,)
   sd_pipe.set_progress_bar_config(disable=lowercase__			)

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

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

   assert image.shape == (1, 5_1_2, 5_1_2, 3)
   __lowercase							  = 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 					SCREAMING_SNAKE_CASE      (      self   :     Tuple			):
   __lowercase							  = DDIMScheduler.from_pretrained(
       '''runwayml/stable-diffusion-v1-5'''					,subfolder='''scheduler'''					,revision='''onnx'''			)
   __lowercase							  = OnnxStableDiffusionPipeline.from_pretrained(
       '''runwayml/stable-diffusion-v1-5'''					,revision='''onnx'''					,scheduler=lowercase__					,safety_checker=lowercase__					,feature_extractor=lowercase__					,provider=self.gpu_provider					,sess_options=self.gpu_options					,)
   sd_pipe.set_progress_bar_config(disable=lowercase__			)

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

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

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

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

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

   __lowercase							  = False

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

   __lowercase							  = '''Andromeda galaxy in a bottle'''

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






  def 					SCREAMING_SNAKE_CASE      (      self   :     Optional[Any]			):
   __lowercase							  = OnnxStableDiffusionPipeline.from_pretrained(
       '''runwayml/stable-diffusion-v1-5'''					,revision='''onnx'''					,safety_checker=lowercase__					,feature_extractor=lowercase__					,provider=self.gpu_provider					,sess_options=self.gpu_options					,)
   assert isinstance(lowercase__					,lowercase__			)
   assert pipe.safety_checker is None

   __lowercase							  = 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(lowercase__			)
    __lowercase							  = OnnxStableDiffusionPipeline.from_pretrained(lowercase__			)

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





'''simple docstring'''
def 					a__       (					lowerCAmelCase__						)    ->	int:
    if a < 0:
        raise ValueError('''Input value must be a positive integer'''						)
    elif isinstance(lowerCAmelCase__  ,					lowerCAmelCase__						):
        raise TypeError('''Input value must be a \'int\' type'''						)
    return bin(lowerCAmelCase__						).count('''1'''						)


if __name__ == "__main__":
  import doctest

  doctest.testmod()
