



"""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    SCREAMING_SNAKE_CASE__ (							__UpperCAmelCase      ,  unittest.TestCase   ):
    lowercase__      =   "hf-internal-testing/tiny-random-OnnxStableDiffusionPipeline"

    def 					_UpperCAmelCase	(				self			:				List[str] ,  lowerCAmelCase_			:				Dict=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			:				Dict):







      """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_2_8, 1_2_8, 3)
      lowercase_								= 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 					_UpperCAmelCase	(				self			:				List[Any]):







      """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_2_8, 1_2_8, 3)
      lowercase_								= 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 					_UpperCAmelCase	(				self			:				Dict):







      """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_2_8, 1_2_8, 3)
      lowercase_								= 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 					_UpperCAmelCase	(				self			:				Any):







      """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_2_8, 1_2_8, 3)
      lowercase_								= 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 					_UpperCAmelCase	(				self			:				Dict):







      """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_2_8, 1_2_8, 3)
      lowercase_								= 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 					_UpperCAmelCase	(				self			:				str):







      """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_2_8, 1_2_8, 3)
      lowercase_								= 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 					_UpperCAmelCase	(				self			:				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			:				Optional[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    SCREAMING_SNAKE_CASE__ (							unittest.TestCase   ):

    @property
    def 					_UpperCAmelCase	(				self			:				Optional[Any]):







      """simple docstring"""







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

    @property
    def 					_UpperCAmelCase	(				self			:				Optional[Any]):







      """simple docstring"""







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

    def 					_UpperCAmelCase	(				self			:				Dict):







      """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=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_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 					_UpperCAmelCase	(				self			:				List[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=1_0 ,  generator=lowerCAmelCase_ ,  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_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 					_UpperCAmelCase	(				self			:				Optional[Any]):







      """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=1_0 ,  generator=lowerCAmelCase_ ,  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_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 					_UpperCAmelCase	(				self			:				Optional[Any]):







      """simple docstring"""







      lowercase_								= 0

      def test_callback_fn(lowerCAmelCase_			:				int ,  lowerCAmelCase_			:				int ,  lowerCAmelCase_			:				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_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, 6_4, 6_4)
          lowercase_								= latents[0, -3:, -3:, -1]
          lowercase_								= 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

      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			:				int):







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






"""simple docstring"""
from ...configuration_utils import PretrainedConfig
from ...utils import logging


UpperCAmelCase						:		List[Any]			= logging.get_logger(__name__)

UpperCAmelCase						:		Optional[Any]			= {
    "alibaba-damo/mgp-str-base": "https://huggingface.co/alibaba-damo/mgp-str-base/resolve/main/config.json",
}


class    SCREAMING_SNAKE_CASE__ (							__UpperCAmelCase   ):
    lowercase__      =   "mgp-str"

    def __init__(				self			:				Optional[Any] ,  lowerCAmelCase_			:				Optional[Any]=[3_2, 1_2_8] ,  lowerCAmelCase_			:				Union[str, Any]=4 ,  lowerCAmelCase_			:				Dict=3 ,  lowerCAmelCase_			:				List[str]=2_7 ,  lowerCAmelCase_			:				str=3_8 ,  lowerCAmelCase_			:				Optional[int]=5_0_2_5_7 ,  lowerCAmelCase_			:				Any=3_0_5_2_2 ,  lowerCAmelCase_			:				str=7_6_8 ,  lowerCAmelCase_			:				Any=1_2 ,  lowerCAmelCase_			:				Optional[Any]=1_2 ,  lowerCAmelCase_			:				Tuple=4.0 ,  lowerCAmelCase_			:				str=True ,  lowerCAmelCase_			:				List[str]=False ,  lowerCAmelCase_			:				List[Any]=1E-5 ,  lowerCAmelCase_			:				Any=0.0 ,  lowerCAmelCase_			:				int=0.0 ,  lowerCAmelCase_			:				str=0.0 ,  lowerCAmelCase_			:				Dict=False ,  lowerCAmelCase_			:				List[str]=0.02 ,  **lowerCAmelCase_			:				str ,  ):







      """simple docstring"""







      super().__init__(**lowerCAmelCase_)

      lowercase_								= image_size
      lowercase_								= patch_size
      lowercase_								= num_channels
      lowercase_								= max_token_length
      lowercase_								= num_character_labels
      lowercase_								= num_bpe_labels
      lowercase_								= num_wordpiece_labels
      lowercase_								= hidden_size
      lowercase_								= num_hidden_layers
      lowercase_								= num_attention_heads
      lowercase_								= mlp_ratio
      lowercase_								= distilled
      lowercase_								= layer_norm_eps
      lowercase_								= drop_rate
      lowercase_								= qkv_bias
      lowercase_								= attn_drop_rate
      lowercase_								= drop_path_rate
      lowercase_								= output_aa_attentions
      lowercase_								= initializer_range

