


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  lowerCAmelCase__(    snake_case__    ,    unittest.TestCase						):






    '''simple docstring'''





    A_  :     Union[str, Any]			= 'hf-internal-testing/tiny-random-OnnxStableDiffusionPipeline'





    def 					_lowerCamelCase (	self				:   Optional[Any]       ,   __snake_case				:   Tuple=0	):





       '''simple docstring'''
       UpperCAmelCase_							:      Union[str, Any]						=   np.random.RandomState(_lowerCAmelCase	)
       UpperCAmelCase_							:      Tuple						=   {
           '''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 					_lowerCamelCase (	self				:   List[Any]	):





       '''simple docstring'''
       UpperCAmelCase_							:      int						=   OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint       ,   provider='''CPUExecutionProvider'''	)
       pipe.set_progress_bar_config(disable=_lowerCAmelCase	)

       UpperCAmelCase_							:      str						=   self.get_dummy_inputs()
       UpperCAmelCase_							:      str						=   pipe(**_lowerCAmelCase	).images
       UpperCAmelCase_							:      List[Any]						=   image[0, -3:, -3:, -1]

       assert image.shape == (1, 128, 128, 3)
       UpperCAmelCase_							:      Optional[int]						=   np.array([0.65_072, 0.58_492, 0.48_219, 0.55_521, 0.53_180, 0.55_939, 0.50_697, 0.39_800, 0.46_455]	)

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





    def 					_lowerCamelCase (	self				:   Union[str, Any]	):





       '''simple docstring'''
       UpperCAmelCase_							:      Optional[Any]						=   OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint       ,   provider='''CPUExecutionProvider'''	)
       UpperCAmelCase_							:      int						=   PNDMScheduler.from_config(pipe.scheduler.config       ,   skip_prk_steps=_lowerCAmelCase	)
       pipe.set_progress_bar_config(disable=_lowerCAmelCase	)

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

       assert image.shape == (1, 128, 128, 3)
       UpperCAmelCase_							:      Tuple						=   np.array([0.65_863, 0.59_425, 0.49_326, 0.56_313, 0.53_875, 0.56_627, 0.51_065, 0.39_777, 0.46_330]	)

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





    def 					_lowerCamelCase (	self				:   int	):





       '''simple docstring'''
       UpperCAmelCase_							:      Optional[Any]						=   OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint       ,   provider='''CPUExecutionProvider'''	)
       UpperCAmelCase_							:      Optional[int]						=   LMSDiscreteScheduler.from_config(pipe.scheduler.config	)
       pipe.set_progress_bar_config(disable=_lowerCAmelCase	)

       UpperCAmelCase_							:      Tuple						=   self.get_dummy_inputs()
       UpperCAmelCase_							:      Union[str, Any]						=   pipe(**_lowerCAmelCase	).images
       UpperCAmelCase_							:      Union[str, Any]						=   image[0, -3:, -3:, -1]

       assert image.shape == (1, 128, 128, 3)
       UpperCAmelCase_							:      Optional[Any]						=   np.array([0.53_755, 0.60_786, 0.47_402, 0.49_488, 0.51_869, 0.49_819, 0.47_985, 0.38_957, 0.44_279]	)

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





    def 					_lowerCamelCase (	self				:   Any	):





       '''simple docstring'''
       UpperCAmelCase_							:      Tuple						=   OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint       ,   provider='''CPUExecutionProvider'''	)
       UpperCAmelCase_							:      Optional[Any]						=   EulerDiscreteScheduler.from_config(pipe.scheduler.config	)
       pipe.set_progress_bar_config(disable=_lowerCAmelCase	)

       UpperCAmelCase_							:      Dict						=   self.get_dummy_inputs()
       UpperCAmelCase_							:      Union[str, Any]						=   pipe(**_lowerCAmelCase	).images
       UpperCAmelCase_							:      List[Any]						=   image[0, -3:, -3:, -1]

       assert image.shape == (1, 128, 128, 3)
       UpperCAmelCase_							:      Union[str, Any]						=   np.array([0.53_755, 0.60_786, 0.47_402, 0.49_488, 0.51_869, 0.49_819, 0.47_985, 0.38_957, 0.44_279]	)

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





    def 					_lowerCamelCase (	self				:   List[Any]	):





       '''simple docstring'''
       UpperCAmelCase_							:      Optional[int]						=   OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint       ,   provider='''CPUExecutionProvider'''	)
       UpperCAmelCase_							:      Union[str, Any]						=   EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config	)
       pipe.set_progress_bar_config(disable=_lowerCAmelCase	)

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

       assert image.shape == (1, 128, 128, 3)
       UpperCAmelCase_							:      Optional[int]						=   np.array([0.53_817, 0.60_812, 0.47_384, 0.49_530, 0.51_894, 0.49_814, 0.47_984, 0.38_958, 0.44_271]	)

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





    def 					_lowerCamelCase (	self				:   List[Any]	):





       '''simple docstring'''
       UpperCAmelCase_							:      Optional[Any]						=   OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint       ,   provider='''CPUExecutionProvider'''	)
       UpperCAmelCase_							:      Dict						=   DPMSolverMultistepScheduler.from_config(pipe.scheduler.config	)
       pipe.set_progress_bar_config(disable=_lowerCAmelCase	)

       UpperCAmelCase_							:      List[str]						=   self.get_dummy_inputs()
       UpperCAmelCase_							:      Tuple						=   pipe(**_lowerCAmelCase	).images
       UpperCAmelCase_							:      Any						=   image[0, -3:, -3:, -1]

       assert image.shape == (1, 128, 128, 3)
       UpperCAmelCase_							:      List[str]						=   np.array([0.53_895, 0.60_808, 0.47_933, 0.49_608, 0.51_886, 0.49_950, 0.48_053, 0.38_957, 0.44_200]	)

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





    def 					_lowerCamelCase (	self				:   Optional[Any]	):





       '''simple docstring'''
       UpperCAmelCase_							:      Union[str, Any]						=   OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint       ,   provider='''CPUExecutionProvider'''	)
       pipe.set_progress_bar_config(disable=_lowerCAmelCase	)

       UpperCAmelCase_							:      List[Any]						=   self.get_dummy_inputs()
       UpperCAmelCase_							:      str						=   3 * [inputs['''prompt''']]

       # forward
       UpperCAmelCase_							:      Tuple						=   pipe(**_lowerCAmelCase	)
       UpperCAmelCase_							:      List[str]						=   output.images[0, -3:, -3:, -1]

       UpperCAmelCase_							:      Optional[Any]						=   self.get_dummy_inputs()
       UpperCAmelCase_							:      str						=   3 * [inputs.pop('''prompt'''	)]

       UpperCAmelCase_							:      int						=   pipe.tokenizer(
           _lowerCAmelCase       ,   padding='''max_length'''       ,   max_length=pipe.tokenizer.model_max_length       ,   truncation=_lowerCAmelCase       ,   return_tensors='''np'''       ,   )
       UpperCAmelCase_							:      List[str]						=   text_inputs['''input_ids''']

       UpperCAmelCase_							:      Dict						=   pipe.text_encoder(input_ids=text_inputs.astype(np.intaa	)	)[0]

       UpperCAmelCase_							:      Union[str, Any]						=   prompt_embeds

       # forward
       UpperCAmelCase_							:      Union[str, Any]						=   pipe(**_lowerCAmelCase	)
       UpperCAmelCase_							:      Optional[int]						=   output.images[0, -3:, -3:, -1]

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





    def 					_lowerCamelCase (	self				:   Tuple	):





       '''simple docstring'''
       UpperCAmelCase_							:      List[Any]						=   OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint       ,   provider='''CPUExecutionProvider'''	)
       pipe.set_progress_bar_config(disable=_lowerCAmelCase	)

       UpperCAmelCase_							:      int						=   self.get_dummy_inputs()
       UpperCAmelCase_							:      Tuple						=   3 * ['''this is a negative prompt''']
       UpperCAmelCase_							:      int						=   negative_prompt
       UpperCAmelCase_							:      Union[str, Any]						=   3 * [inputs['''prompt''']]

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

       UpperCAmelCase_							:      str						=   self.get_dummy_inputs()
       UpperCAmelCase_							:      str						=   3 * [inputs.pop('''prompt'''	)]

       UpperCAmelCase_							:      Any						=   []
       for p in [prompt, negative_prompt]:
          UpperCAmelCase_							:      Optional[int]						=   pipe.tokenizer(
              _lowerCAmelCase       ,   padding='''max_length'''       ,   max_length=pipe.tokenizer.model_max_length       ,   truncation=_lowerCAmelCase       ,   return_tensors='''np'''       ,   )
          UpperCAmelCase_							:      Union[str, Any]						=   text_inputs['''input_ids''']

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

       UpperCAmelCase_		,      UpperCAmelCase_							:      Dict						=   embeds

       # forward
       UpperCAmelCase_							:      str						=   pipe(**_lowerCAmelCase	)
       UpperCAmelCase_							:      List[str]						=   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  lowerCAmelCase__(    unittest.TestCase						):






    '''simple docstring'''





    @property
    def 					_lowerCamelCase (	self				:   Any	):





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





    @property
    def 					_lowerCamelCase (	self				:   str	):





       '''simple docstring'''
       UpperCAmelCase_							:      List[Any]						=   ort.SessionOptions()
       UpperCAmelCase_							:      Optional[int]						=   False
       return options





    def 					_lowerCamelCase (	self				:   int	):





       '''simple docstring'''
       # using the PNDM scheduler by default
       UpperCAmelCase_							:      Tuple						=   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_							:      Optional[int]						=   '''A painting of a squirrel eating a burger'''
       np.random.seed(0	)
       UpperCAmelCase_							:      Optional[int]						=   sd_pipe([prompt]       ,   guidance_scale=6.0       ,   num_inference_steps=10       ,   output_type='''np'''	)
       UpperCAmelCase_							:      Tuple						=   output.images

       UpperCAmelCase_							:      List[str]						=   image[0, -3:, -3:, -1]

       assert image.shape == (1, 512, 512, 3)
       UpperCAmelCase_							:      List[str]						=   np.array([0.0_452, 0.0_390, 0.0_087, 0.0_350, 0.0_617, 0.0_364, 0.0_544, 0.0_523, 0.0_720]	)

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





    def 					_lowerCamelCase (	self				:   Dict	):





       '''simple docstring'''
       UpperCAmelCase_							:      int						=   DDIMScheduler.from_pretrained(
           '''runwayml/stable-diffusion-v1-5'''       ,   subfolder='''scheduler'''       ,   revision='''onnx'''	)
       UpperCAmelCase_							:      List[Any]						=   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_							:      List[str]						=   '''open neural network exchange'''
       UpperCAmelCase_							:      Any						=   np.random.RandomState(0	)
       UpperCAmelCase_							:      Optional[Any]						=   sd_pipe([prompt]       ,   guidance_scale=7.5       ,   num_inference_steps=10       ,   generator=_lowerCAmelCase       ,   output_type='''np'''	)
       UpperCAmelCase_							:      Tuple						=   output.images
       UpperCAmelCase_							:      Union[str, Any]						=   image[0, -3:, -3:, -1]

       assert image.shape == (1, 512, 512, 3)
       UpperCAmelCase_							:      Any						=   np.array([0.2_867, 0.1_974, 0.1_481, 0.7_294, 0.7_251, 0.6_667, 0.4_194, 0.5_642, 0.6_486]	)

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





    def 					_lowerCamelCase (	self				:   Any	):





       '''simple docstring'''
       UpperCAmelCase_							:      Optional[Any]						=   LMSDiscreteScheduler.from_pretrained(
           '''runwayml/stable-diffusion-v1-5'''       ,   subfolder='''scheduler'''       ,   revision='''onnx'''	)
       UpperCAmelCase_							:      Dict						=   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_							:      str						=   '''open neural network exchange'''
       UpperCAmelCase_							:      Optional[Any]						=   np.random.RandomState(0	)
       UpperCAmelCase_							:      Optional[Any]						=   sd_pipe([prompt]       ,   guidance_scale=7.5       ,   num_inference_steps=10       ,   generator=_lowerCAmelCase       ,   output_type='''np'''	)
       UpperCAmelCase_							:      int						=   output.images
       UpperCAmelCase_							:      Dict						=   image[0, -3:, -3:, -1]

       assert image.shape == (1, 512, 512, 3)
       UpperCAmelCase_							:      Any						=   np.array([0.2_306, 0.1_959, 0.1_593, 0.6_549, 0.6_394, 0.5_408, 0.5_065, 0.6_010, 0.6_161]	)

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





    def 					_lowerCamelCase (	self				:   Optional[Any]	):





       '''simple docstring'''
       UpperCAmelCase_							:      Union[str, Any]						=   0

       def test_callback_fn(__snake_case				:   Tuple       ,   __snake_case				:   Any       ,   __snake_case				:   Optional[Any]	) -> None:
          UpperCAmelCase_							:      Optional[Any]						=   True
          nonlocal number_of_steps
          number_of_steps += 1
          if step == 0:
             assert latents.shape == (1, 4, 64, 64)
             UpperCAmelCase_							:      Tuple						=   latents[0, -3:, -3:, -1]
             UpperCAmelCase_							:      List[Any]						=   np.array(
                 [-0.6_772, -0.3_835, -1.2_456, 0.1_905, -1.0_974, 0.6_967, -1.9_353, 0.0_178, 1.0_167]	)

             assert np.abs(latents_slice.flatten() - expected_slice	).max() < 1E-3
          elif step == 5:
             assert latents.shape == (1, 4, 64, 64)
             UpperCAmelCase_							:      Tuple						=   latents[0, -3:, -3:, -1]
             UpperCAmelCase_							:      Any						=   np.array(
                 [-0.3_351, 0.2_241, -0.1_837, -0.2_325, -0.6_577, 0.3_393, -0.0_241, 0.5_899, 1.3_875]	)

             assert np.abs(latents_slice.flatten() - expected_slice	).max() < 1E-3

       UpperCAmelCase_							:      Tuple						=   False

       UpperCAmelCase_							:      int						=   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_							:      Optional[int]						=   '''Andromeda galaxy in a bottle'''

       UpperCAmelCase_							:      Union[str, Any]						=   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 					_lowerCamelCase (	self				:   Optional[Any]	):





       '''simple docstring'''
       UpperCAmelCase_							:      Tuple						=   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_							:      Union[str, Any]						=   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_							:      List[Any]						=   OnnxStableDiffusionPipeline.from_pretrained(_lowerCAmelCase	)

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








import copy

from ...configuration_utils import PretrainedConfig
from ...utils import logging


__UpperCamelCase				:  Tuple      					=							logging.get_logger(__name__)





class  lowerCAmelCase__(    snake_case__						):






    '''simple docstring'''





    A_  :     Union[str, Any]			= 'encoder-decoder'
    A_  :     Optional[int]			= True





    def __init__(	self				:   Dict       ,   **__snake_case				:   Union[str, Any]	):





       '''simple docstring'''
       super().__init__(**__snake_case	)
       assert (
           "encoder" in kwargs and "decoder" in kwargs
       ), "Config has to be initialized with encoder and decoder config"
       UpperCAmelCase_							:      int						=   kwargs.pop('''encoder'''	)
       UpperCAmelCase_							:      List[Any]						=   encoder_config.pop('''model_type'''	)
       UpperCAmelCase_							:      int						=   kwargs.pop('''decoder'''	)
       UpperCAmelCase_							:      Union[str, Any]						=   decoder_config.pop('''model_type'''	)

       from ..auto.configuration_auto import AutoConfig

       UpperCAmelCase_							:      Optional[int]						=   AutoConfig.for_model(__snake_case       ,   **__snake_case	)
       UpperCAmelCase_							:      Any						=   AutoConfig.for_model(__snake_case       ,   **__snake_case	)
       UpperCAmelCase_							:      List[Any]						=   True





    @classmethod
    def 					_lowerCamelCase (	cls				:   Optional[Any]       ,   __snake_case				:   PretrainedConfig       ,   __snake_case				:   PretrainedConfig       ,   **__snake_case				:   Union[str, Any]	):





       '''simple docstring'''
       logger.info('''Set `config.is_decoder=True` and `config.add_cross_attention=True` for decoder_config'''	)
       UpperCAmelCase_							:      Dict						=   True
       UpperCAmelCase_							:      str						=   True

       return cls(encoder=encoder_config.to_dict()       ,   decoder=decoder_config.to_dict()       ,   **__snake_case	)





    def 					_lowerCamelCase (	self				:   Dict	):





       '''simple docstring'''
       UpperCAmelCase_							:      int						=   copy.deepcopy(self.__dict__	)
       UpperCAmelCase_							:      Tuple						=   self.encoder.to_dict()
       UpperCAmelCase_							:      Tuple						=   self.decoder.to_dict()
       UpperCAmelCase_							:      Tuple						=   self.__class__.model_type
       return output