"""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 					_A		(			lowerCAmelCase    ,       unittest.TestCase      ):
      snake_case__     :  List[str]				    =      'hf-internal-testing/tiny-random-OnnxStableDiffusionPipeline'

      def 							A__			(		self						,			__lowerCAmelCase=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 							A__			(		self		):







            """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, 128, 128, 3)
            lowercase								=      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		):







            """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, 128, 128, 3)
            lowercase								=      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		):







            """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, 128, 128, 3)
            lowercase								=      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		):







            """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, 128, 128, 3)
            lowercase								=      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		):







            """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, 128, 128, 3)
            lowercase								=      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		):







            """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, 128, 128, 3)
            lowercase								=      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		):







            """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 							A__			(		self		):







            """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 					_A		(			unittest.TestCase      ):

      @property
      def 							A__			(		self		):







            """simple docstring"""



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

      @property
      def 							A__			(		self		):







            """simple docstring"""



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

      def 							A__			(		self		):







            """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, 512, 512, 3)
            lowercase								=      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		):







            """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, 512, 512, 3)
            lowercase								=      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		):







            """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, 512, 512, 3)
            lowercase								=      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		):







            """simple docstring"""



            lowercase								=      0

            def test_callback_fn(__lowerCAmelCase						,			__lowerCAmelCase						,			__lowerCAmelCase		) -> 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.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, 64, 64)
                        lowercase								=      latents[0, -3:, -3:, -1]
                        lowercase								=      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

            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 							A__			(		self		):







            """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 dataclasses import dataclass, field
from typing import ClassVar, Dict

from ..features import Features, Sequence, Value
from .base import TaskTemplate





@dataclass(frozen=lowerCAmelCase      )
class 					_A		(			lowerCAmelCase      ):
      # `task` is not a ClassVar since we want it to be part of the `asdict` output for JSON serialization
      snake_case__     :  str				    =      field(default='question-answering-extractive'    ,       metadata={'include_in_asdict_even_if_is_default': True}      )
      snake_case__     :  ClassVar[Features]				    =      Features({'question': Value('string'      ), 'context': Value('string'      )}      )
      snake_case__     :  ClassVar[Features]				    =      Features(
          {
              'answers': Sequence(
                  {
                      'text': Value('string'      ),
                      'answer_start': Value('int32'      ),
                  }      )
          }      )
      snake_case__     :  str				    =      "question"
      snake_case__     :  str				    =      "context"
      snake_case__     :  str				    =      "answers"

      @property
      def 							A__			(		self		):







            """simple docstring"""



            return {self.question_column: "question", self.context_column: "context", self.answers_column: "answers"}



