
"""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      _UpperCAmelCase							(		_lowerCAmelCase     ,			unittest.TestCase							):
      __SCREAMING_SNAKE_CASE   :							Union[str, Any]							      =	'hf-internal-testing/tiny-random-OnnxStableDiffusionPipeline'




      def     a_      (							self   ,       lowercase_=0							)						->					str:
          UpperCAmelCase									=					np.random.RandomState(_lowerCAmelCase							)
          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     a_      (							self							)						->					List[Any]:
          UpperCAmelCase									=					OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint   ,       provider='CPUExecutionProvider'							)
          pipe.set_progress_bar_config(disable=_lowerCAmelCase							)

          UpperCAmelCase									=					self.get_dummy_inputs()
          UpperCAmelCase									=					pipe(**_lowerCAmelCase							).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     a_      (							self							)						->					str:
          UpperCAmelCase									=					OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint   ,       provider='CPUExecutionProvider'							)
          UpperCAmelCase									=					PNDMScheduler.from_config(pipe.scheduler.config   ,       skip_prk_steps=_lowerCAmelCase							)
          pipe.set_progress_bar_config(disable=_lowerCAmelCase							)

          UpperCAmelCase									=					self.get_dummy_inputs()
          UpperCAmelCase									=					pipe(**_lowerCAmelCase							).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     a_      (							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=_lowerCAmelCase							)

          UpperCAmelCase									=					self.get_dummy_inputs()
          UpperCAmelCase									=					pipe(**_lowerCAmelCase							).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     a_      (							self							)						->					int:
          UpperCAmelCase									=					OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint   ,       provider='CPUExecutionProvider'							)
          UpperCAmelCase									=					EulerDiscreteScheduler.from_config(pipe.scheduler.config							)
          pipe.set_progress_bar_config(disable=_lowerCAmelCase							)

          UpperCAmelCase									=					self.get_dummy_inputs()
          UpperCAmelCase									=					pipe(**_lowerCAmelCase							).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     a_      (							self							)						->					Union[str, Any]:
          UpperCAmelCase									=					OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint   ,       provider='CPUExecutionProvider'							)
          UpperCAmelCase									=					EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config							)
          pipe.set_progress_bar_config(disable=_lowerCAmelCase							)

          UpperCAmelCase									=					self.get_dummy_inputs()
          UpperCAmelCase									=					pipe(**_lowerCAmelCase							).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     a_      (							self							)						->					List[str]:
          UpperCAmelCase									=					OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint   ,       provider='CPUExecutionProvider'							)
          UpperCAmelCase									=					DPMSolverMultistepScheduler.from_config(pipe.scheduler.config							)
          pipe.set_progress_bar_config(disable=_lowerCAmelCase							)

          UpperCAmelCase									=					self.get_dummy_inputs()
          UpperCAmelCase									=					pipe(**_lowerCAmelCase							).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     a_      (							self							)						->					List[str]:
          UpperCAmelCase									=					OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint   ,       provider='CPUExecutionProvider'							)
          pipe.set_progress_bar_config(disable=_lowerCAmelCase							)

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

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

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

          UpperCAmelCase									=					pipe.tokenizer(
              _lowerCAmelCase   ,       padding='max_length'   ,       max_length=pipe.tokenizer.model_max_length   ,       truncation=_lowerCAmelCase   ,       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(**_lowerCAmelCase							)
          UpperCAmelCase									=					output.images[0, -3:, -3:, -1]

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




      def     a_      (							self							)						->					str:
          UpperCAmelCase									=					OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint   ,       provider='CPUExecutionProvider'							)
          pipe.set_progress_bar_config(disable=_lowerCAmelCase							)

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

          # forward
          UpperCAmelCase									=					pipe(**_lowerCAmelCase							)
          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(
                  _lowerCAmelCase   ,       padding='max_length'   ,       max_length=pipe.tokenizer.model_max_length   ,       truncation=_lowerCAmelCase   ,       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(**_lowerCAmelCase							)
          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      _UpperCAmelCase							(		unittest.TestCase							):




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




      @property
      def     a_      (							self							)						->					Optional[Any]:
          UpperCAmelCase									=					ort.SessionOptions()
          UpperCAmelCase									=					False
          return options




      def     a_      (							self							)						->					Union[str, Any]:
          UpperCAmelCase									=					OnnxStableDiffusionPipeline.from_pretrained(
              'CompVis/stable-diffusion-v1-4'   ,       revision='onnx'   ,       safety_checker=_lowerCAmelCase   ,       feature_extractor=_lowerCAmelCase   ,       provider=self.gpu_provider   ,       sess_options=self.gpu_options   ,       )
          sd_pipe.set_progress_bar_config(disable=_lowerCAmelCase							)

          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     a_      (							self							)						->					Any:
          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=_lowerCAmelCase   ,       safety_checker=_lowerCAmelCase   ,       feature_extractor=_lowerCAmelCase   ,       provider=self.gpu_provider   ,       sess_options=self.gpu_options   ,       )
          sd_pipe.set_progress_bar_config(disable=_lowerCAmelCase							)

          UpperCAmelCase									=					'open neural network exchange'
          UpperCAmelCase									=					np.random.RandomState(0							)
          UpperCAmelCase									=					sd_pipe([prompt]   ,       guidance_scale=7.5   ,       num_inference_steps=1_0   ,       generator=_lowerCAmelCase   ,       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     a_      (							self							)						->					Dict:
          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=_lowerCAmelCase   ,       safety_checker=_lowerCAmelCase   ,       feature_extractor=_lowerCAmelCase   ,       provider=self.gpu_provider   ,       sess_options=self.gpu_options   ,       )
          sd_pipe.set_progress_bar_config(disable=_lowerCAmelCase							)

          UpperCAmelCase									=					'open neural network exchange'
          UpperCAmelCase									=					np.random.RandomState(0							)
          UpperCAmelCase									=					sd_pipe([prompt]   ,       guidance_scale=7.5   ,       num_inference_steps=1_0   ,       generator=_lowerCAmelCase   ,       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     a_      (							self							)						->					List[str]:
          UpperCAmelCase									=					0

          def test_callback_fn(lowercase_   ,       lowercase_   ,       lowercase_							) -> 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=_lowerCAmelCase   ,       feature_extractor=_lowerCAmelCase   ,       provider=self.gpu_provider   ,       sess_options=self.gpu_options   ,       )
          pipe.set_progress_bar_config(disable=_lowerCAmelCase							)

          UpperCAmelCase									=					'Andromeda galaxy in a bottle'

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




      def     a_      (							self							)						->					Tuple:
          UpperCAmelCase									=					OnnxStableDiffusionPipeline.from_pretrained(
              'runwayml/stable-diffusion-v1-5'   ,       revision='onnx'   ,       safety_checker=_lowerCAmelCase   ,       feature_extractor=_lowerCAmelCase   ,       provider=self.gpu_provider   ,       sess_options=self.gpu_options   ,       )
          assert isinstance(_lowerCAmelCase   ,       _lowerCAmelCase							)
          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(_lowerCAmelCase							)
              UpperCAmelCase									=					OnnxStableDiffusionPipeline.from_pretrained(_lowerCAmelCase							)

          # 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"""

def     lowercase__						(					lowerCAmelCase				:					int ,				lowerCAmelCase				:					int ) ->      int:






    """simple docstring"""




    return number | (1 << position)




def     lowercase__						(					lowerCAmelCase				:					int ,				lowerCAmelCase				:					int ) ->      int:






    """simple docstring"""




    return number & ~(1 << position)




def     lowercase__						(					lowerCAmelCase				:					int ,				lowerCAmelCase				:					int ) ->      int:






    """simple docstring"""




    return number ^ (1 << position)




def     lowercase__						(					lowerCAmelCase				:					int ,				lowerCAmelCase				:					int ) ->      bool:






    """simple docstring"""




    return ((number >> position) & 1) == 1




def     lowercase__						(					lowerCAmelCase				:					int ,				lowerCAmelCase				:					int ) ->      int:






    """simple docstring"""




    return int((number & (1 << position)) != 0 )


if __name__ == "__main__":
   import doctest

   doctest.testmod()



