





"""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  lowerCAmelCase							(       lowercase_							,							unittest.TestCase ):
  __lowerCamelCase		=      'hf-internal-testing/tiny-random-OnnxStableDiffusionPipeline'







  def   UpperCAmelCase					(      self  :Tuple   ,							_lowercase  :Union[str, Any]=0 ):






     '''simple docstring'''





     lowercase__			       =  np.random.RandomState(_lowerCAmelCase )
     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   UpperCAmelCase					(      self  :List[str] ):






     '''simple docstring'''





     lowercase__			       =  OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint   ,							provider="CPUExecutionProvider" )
     pipe.set_progress_bar_config(disable=_lowerCAmelCase )

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

     assert image.shape == (1, 1_28, 1_28, 3)
     lowercase__			       =  np.array([0.65072, 0.58492, 0.48219, 0.55521, 0.53180, 0.55939, 0.50697, 0.39800, 0.46455] )

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







  def   UpperCAmelCase					(      self  :Optional[int] ):






     '''simple docstring'''





     lowercase__			       =  OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint   ,							provider="CPUExecutionProvider" )
     lowercase__			       =  PNDMScheduler.from_config(pipe.scheduler.config   ,							skip_prk_steps=_lowerCAmelCase )
     pipe.set_progress_bar_config(disable=_lowerCAmelCase )

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

     assert image.shape == (1, 1_28, 1_28, 3)
     lowercase__			       =  np.array([0.65863, 0.59425, 0.49326, 0.56313, 0.53875, 0.56627, 0.51065, 0.39777, 0.46330] )

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







  def   UpperCAmelCase					(      self  :Optional[int] ):






     '''simple docstring'''





     lowercase__			       =  OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint   ,							provider="CPUExecutionProvider" )
     lowercase__			       =  LMSDiscreteScheduler.from_config(pipe.scheduler.config )
     pipe.set_progress_bar_config(disable=_lowerCAmelCase )

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

     assert image.shape == (1, 1_28, 1_28, 3)
     lowercase__			       =  np.array([0.53755, 0.60786, 0.47402, 0.49488, 0.51869, 0.49819, 0.47985, 0.38957, 0.44279] )

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







  def   UpperCAmelCase					(      self  :Dict ):






     '''simple docstring'''





     lowercase__			       =  OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint   ,							provider="CPUExecutionProvider" )
     lowercase__			       =  EulerDiscreteScheduler.from_config(pipe.scheduler.config )
     pipe.set_progress_bar_config(disable=_lowerCAmelCase )

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

     assert image.shape == (1, 1_28, 1_28, 3)
     lowercase__			       =  np.array([0.53755, 0.60786, 0.47402, 0.49488, 0.51869, 0.49819, 0.47985, 0.38957, 0.44279] )

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







  def   UpperCAmelCase					(      self  :int ):






     '''simple docstring'''





     lowercase__			       =  OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint   ,							provider="CPUExecutionProvider" )
     lowercase__			       =  EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config )
     pipe.set_progress_bar_config(disable=_lowerCAmelCase )

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

     assert image.shape == (1, 1_28, 1_28, 3)
     lowercase__			       =  np.array([0.53817, 0.60812, 0.47384, 0.49530, 0.51894, 0.49814, 0.47984, 0.38958, 0.44271] )

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







  def   UpperCAmelCase					(      self  :List[Any] ):






     '''simple docstring'''





     lowercase__			       =  OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint   ,							provider="CPUExecutionProvider" )
     lowercase__			       =  DPMSolverMultistepScheduler.from_config(pipe.scheduler.config )
     pipe.set_progress_bar_config(disable=_lowerCAmelCase )

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

     assert image.shape == (1, 1_28, 1_28, 3)
     lowercase__			       =  np.array([0.53895, 0.60808, 0.47933, 0.49608, 0.51886, 0.49950, 0.48053, 0.38957, 0.44200] )

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







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






     '''simple docstring'''





     lowercase__			       =  OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint   ,							provider="CPUExecutionProvider" )
     pipe.set_progress_bar_config(disable=_lowerCAmelCase )

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

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

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

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

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







  def   UpperCAmelCase					(      self  :int ):






     '''simple docstring'''





     lowercase__			       =  OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint   ,							provider="CPUExecutionProvider" )
     pipe.set_progress_bar_config(disable=_lowerCAmelCase )

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

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







  @property
  def   UpperCAmelCase					(      self  :Tuple ):






     '''simple docstring'''





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







  @property
  def   UpperCAmelCase					(      self  :List[Any] ):






     '''simple docstring'''





     lowercase__			       =  ort.SessionOptions()
     lowercase__			       =  False
     return options







  def   UpperCAmelCase					(      self  :Optional[int] ):






     '''simple docstring'''





     lowercase__			       =  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 )

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

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

     assert image.shape == (1, 5_12, 5_12, 3)
     lowercase__			       =  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  :Union[str, Any] ):






     '''simple docstring'''





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

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

     assert image.shape == (1, 5_12, 5_12, 3)
     lowercase__			       =  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  :List[str] ):






     '''simple docstring'''





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

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

     assert image.shape == (1, 5_12, 5_12, 3)
     lowercase__			       =  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  :Any ):






     '''simple docstring'''





     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, 64, 64)
           lowercase__			       =  latents[0, -3:, -3:, -1]
           lowercase__			       =  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)
           lowercase__			       =  latents[0, -3:, -3:, -1]
           lowercase__			       =  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

     lowercase__			       =  False

     lowercase__			       =  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 )

     lowercase__			       =  "Andromeda galaxy in a bottle"

     lowercase__			       =  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   UpperCAmelCase					(      self  :Dict ):






     '''simple docstring'''





     lowercase__			       =  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

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

     # 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













import os
import socket
from contextlib import contextmanager

import torch

from ..commands.config.default import write_basic_config  # noqa: F401
from ..state import PartialState
from .dataclasses import DistributedType
from .imports import is_deepspeed_available, is_tpu_available
from .transformer_engine import convert_model
from .versions import is_torch_version


if is_deepspeed_available():
  from deepspeed import DeepSpeedEngine

if is_tpu_available(check_device=False):
  import torch_xla.core.xla_model as xm


def _A       (   __magic_name__	):
   if is_torch_version("<"  , "2.0.0"	) or not hasattr(__magic_name__  , "_dynamo"	):
      return False
   return isinstance(__magic_name__  , torch._dynamo.eval_frame.OptimizedModule	)


def _A       (   __magic_name__  , __magic_name__ = True	):
   lowercase__			       =  (torch.nn.parallel.DistributedDataParallel, torch.nn.DataParallel)

   lowercase__			       =  is_compiled_module(__magic_name__	)
   if is_compiled:
      lowercase__			       =  model
      lowercase__			       =  model._orig_mod

   if is_deepspeed_available():
      options += (DeepSpeedEngine,)

   while isinstance(__magic_name__  , __magic_name__	):
      lowercase__			       =  model.module

   if not keep_fpaa_wrapper:
      lowercase__			       =  getattr(__magic_name__  , "forward"	)
      lowercase__			       =  model.__dict__.pop("_original_forward"  , __magic_name__	)
      if original_forward is not None:
         while hasattr(__magic_name__  , "__wrapped__"	):
            lowercase__			       =  forward.__wrapped__
            if forward == original_forward:
               break
         lowercase__			       =  forward
      if getattr(__magic_name__  , "_converted_to_transformer_engine"  , __magic_name__	):
         convert_model(__magic_name__  , to_transformer_engine=__magic_name__	)

   if is_compiled:
      lowercase__			       =  model
      lowercase__			       =  compiled_model

   return model


def _A       (   ):
   PartialState().wait_for_everyone()


def _A       (   __magic_name__  , __magic_name__	):
   if PartialState().distributed_type == DistributedType.TPU:
      xm.save(__magic_name__  , __magic_name__	)
   elif PartialState().local_process_index == 0:
      torch.save(__magic_name__  , __magic_name__	)


@contextmanager
def _A       (   **__magic_name__	):
   for key, value in kwargs.items():
      lowercase__			       =  str(__magic_name__	)

   yield

   for key in kwargs:
      if key.upper() in os.environ:
         del os.environ[key.upper()]


def _A       (   __magic_name__	):
   if not hasattr(__magic_name__  , "__qualname__"	) and not hasattr(__magic_name__  , "__name__"	):
      lowercase__			       =  getattr(__magic_name__  , "__class__"  , __magic_name__	)
   if hasattr(__magic_name__  , "__qualname__"	):
      return obj.__qualname__
   if hasattr(__magic_name__  , "__name__"	):
      return obj.__name__
   return str(__magic_name__	)


def _A       (   __magic_name__  , __magic_name__	):
   for key, value in source.items():
      if isinstance(__magic_name__  , __magic_name__	):
         lowercase__			       =  destination.setdefault(__magic_name__  , {}	)
         merge_dicts(__magic_name__  , __magic_name__	)
      else:
         lowercase__			       =  value

   return destination







def _A       (   __magic_name__ = None	):
   if port is None:
      lowercase__			       =  2_9500
   with socket.socket(socket.AF_INET  , socket.SOCK_STREAM	) as s:
      return s.connect_ex(("localhost", port)	) == 0






