





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    __snake_case	(							_UpperCAmelCase       ,      unittest.TestCase      ):





       '''simple docstring'''







       lowerCAmelCase__						=					"""hf-internal-testing/tiny-random-OnnxStableDiffusionPipeline"""




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

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

         assert image.shape == (1, 128, 128, 3)
         __snake_case:						List[str]			  =						np.array([0.6_5072, 0.5_8492, 0.4_8219, 0.5_5521, 0.5_3180, 0.5_5939, 0.5_0697, 0.3_9800, 0.4_6455]		)

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




       def 							UpperCAmelCase__			(  self			:   Dict		):
         __snake_case:						Optional[Any]			  =						OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint							,    provider="""CPUExecutionProvider"""		)
         __snake_case:						Tuple			  =						PNDMScheduler.from_config(pipe.scheduler.config							,    skip_prk_steps=lowercase_		)
         pipe.set_progress_bar_config(disable=lowercase_		)

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

         assert image.shape == (1, 128, 128, 3)
         __snake_case:						Dict			  =						np.array([0.6_5863, 0.5_9425, 0.4_9326, 0.5_6313, 0.5_3875, 0.5_6627, 0.5_1065, 0.3_9777, 0.4_6330]		)

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




       def 							UpperCAmelCase__			(  self			:   Tuple		):
         __snake_case:						Optional[int]			  =						OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint							,    provider="""CPUExecutionProvider"""		)
         __snake_case:						Optional[Any]			  =						LMSDiscreteScheduler.from_config(pipe.scheduler.config		)
         pipe.set_progress_bar_config(disable=lowercase_		)

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

         assert image.shape == (1, 128, 128, 3)
         __snake_case:						Any			  =						np.array([0.5_3755, 0.6_0786, 0.4_7402, 0.4_9488, 0.5_1869, 0.4_9819, 0.4_7985, 0.3_8957, 0.4_4279]		)

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




       def 							UpperCAmelCase__			(  self			:   Union[str, Any]		):
         __snake_case:						List[str]			  =						OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint							,    provider="""CPUExecutionProvider"""		)
         __snake_case:						Optional[int]			  =						EulerDiscreteScheduler.from_config(pipe.scheduler.config		)
         pipe.set_progress_bar_config(disable=lowercase_		)

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

         assert image.shape == (1, 128, 128, 3)
         __snake_case:						Any			  =						np.array([0.5_3755, 0.6_0786, 0.4_7402, 0.4_9488, 0.5_1869, 0.4_9819, 0.4_7985, 0.3_8957, 0.4_4279]		)

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




       def 							UpperCAmelCase__			(  self			:   Optional[int]		):
         __snake_case:						Optional[int]			  =						OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint							,    provider="""CPUExecutionProvider"""		)
         __snake_case:						Union[str, Any]			  =						EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config		)
         pipe.set_progress_bar_config(disable=lowercase_		)

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

         assert image.shape == (1, 128, 128, 3)
         __snake_case:						List[Any]			  =						np.array([0.5_3817, 0.6_0812, 0.4_7384, 0.4_9530, 0.5_1894, 0.4_9814, 0.4_7984, 0.3_8958, 0.4_4271]		)

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




       def 							UpperCAmelCase__			(  self			:   List[Any]		):
         __snake_case:						Optional[int]			  =						OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint							,    provider="""CPUExecutionProvider"""		)
         __snake_case:						Dict			  =						DPMSolverMultistepScheduler.from_config(pipe.scheduler.config		)
         pipe.set_progress_bar_config(disable=lowercase_		)

         __snake_case:						Optional[Any]			  =						self.get_dummy_inputs()
         __snake_case:						Dict			  =						pipe(**lowercase_		).images
         __snake_case:						Any			  =						image[0, -3:, -3:, -1]

         assert image.shape == (1, 128, 128, 3)
         __snake_case:						Union[str, Any]			  =						np.array([0.5_3895, 0.6_0808, 0.4_7933, 0.4_9608, 0.5_1886, 0.4_9950, 0.4_8053, 0.3_8957, 0.4_4200]		)

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




       def 							UpperCAmelCase__			(  self			:   Union[str, Any]		):
         __snake_case:						Dict			  =						OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint							,    provider="""CPUExecutionProvider"""		)
         pipe.set_progress_bar_config(disable=lowercase_		)

         __snake_case:						Tuple			  =						self.get_dummy_inputs()
         __snake_case:						Optional[int]			  =						3 * [inputs["""prompt"""]]

         # forward
         __snake_case:						Optional[int]			  =						pipe(**lowercase_		)
         __snake_case:						List[Any]			  =						output.images[0, -3:, -3:, -1]

         __snake_case:						Union[str, Any]			  =						self.get_dummy_inputs()
         __snake_case:						str			  =						3 * [inputs.pop("""prompt"""		)]

         __snake_case:						Optional[Any]			  =						pipe.tokenizer(
             lowercase_							,    padding="""max_length"""							,    max_length=pipe.tokenizer.model_max_length							,    truncation=lowercase_							,    return_tensors="""np"""							,    )
         __snake_case:						Optional[int]			  =						text_inputs["""input_ids"""]

         __snake_case:						Any			  =						pipe.text_encoder(input_ids=text_inputs.astype(np.intaa		)		)[0]

         __snake_case:						int			  =						prompt_embeds

         # forward
         __snake_case:						Union[str, Any]			  =						pipe(**lowercase_		)
         __snake_case:						List[Any]			  =						output.images[0, -3:, -3:, -1]

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






       def 							UpperCAmelCase__			(  self			:   Union[str, Any]		):
         __snake_case:						Optional[int]			  =						OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint							,    provider="""CPUExecutionProvider"""		)
         pipe.set_progress_bar_config(disable=lowercase_		)

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

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

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

         __snake_case:						Any			  =						[]
         for p in [prompt, negative_prompt]:
           __snake_case:						str			  =						pipe.tokenizer(
               lowercase_							,    padding="""max_length"""							,    max_length=pipe.tokenizer.model_max_length							,    truncation=lowercase_							,    return_tensors="""np"""							,    )
           __snake_case:						Optional[Any]			  =						text_inputs["""input_ids"""]

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

         __snake_case:						Optional[Any]			  =						embeds

         # forward
         __snake_case:						Union[str, Any]			  =						pipe(**lowercase_		)
         __snake_case:						Optional[int]			  =						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    __snake_case	(							unittest.TestCase      ):





       '''simple docstring'''







       @property
       def 							UpperCAmelCase__			(  self			:   Union[str, Any]		):
         return (
             "CUDAExecutionProvider",
             {
                 "gpu_mem_limit": "15000000000",  # 15GB
                 "arena_extend_strategy": "kSameAsRequested",
             },
         )




       @property
       def 							UpperCAmelCase__			(  self			:   Tuple		):
         __snake_case:						Dict			  =						ort.SessionOptions()
         __snake_case:						List[str]			  =						False
         return options




       def 							UpperCAmelCase__			(  self			:   Optional[Any]		):
         # using the PNDM scheduler by default
         __snake_case:						Tuple			  =						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_		)

         __snake_case:						Optional[Any]			  =						"""A painting of a squirrel eating a burger"""
         np.random.seed(0		)
         __snake_case:						Optional[int]			  =						sd_pipe([prompt]							,    guidance_scale=6.0							,    num_inference_steps=10							,    output_type="""np"""		)
         __snake_case:						List[Any]			  =						output.images

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

         assert image.shape == (1, 512, 512, 3)
         __snake_case:						Tuple			  =						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 							UpperCAmelCase__			(  self			:   Optional[Any]		):
         __snake_case:						Any			  =						DDIMScheduler.from_pretrained(
             """runwayml/stable-diffusion-v1-5"""							,    subfolder="""scheduler"""							,    revision="""onnx"""		)
         __snake_case:						int			  =						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_		)

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

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

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

         assert image.shape == (1, 512, 512, 3)
         __snake_case:						Tuple			  =						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 							UpperCAmelCase__			(  self			:   Tuple		):
         __snake_case:						Tuple			  =						0

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

         __snake_case:						Dict			  =						False

         __snake_case:						Optional[Any]			  =						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_		)

         __snake_case:						Tuple			  =						"""Andromeda galaxy in a bottle"""

         __snake_case:						List[Any]			  =						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 							UpperCAmelCase__			(  self			:   Optional[int]		):
         __snake_case:						Any			  =						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

         __snake_case:						str			  =						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_		)
           __snake_case:						List[str]			  =						OnnxStableDiffusionPipeline.from_pretrained(lowercase_		)

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













from ...utils import (
    OptionalDependencyNotAvailable,
    is_torch_available,
    is_transformers_available,
    is_transformers_version,
)


try:
      if not (is_transformers_available() and is_torch_available()):
            raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
      from ...utils.dummy_torch_and_transformers_objects import (
          ImageTextPipelineOutput,
          UniDiffuserPipeline,
      )
else:
      from .modeling_text_decoder import UniDiffuserTextDecoder
      from .modeling_uvit import UniDiffuserModel, UTransformeraDModel
      from .pipeline_unidiffuser import ImageTextPipelineOutput, UniDiffuserPipeline






