
import random
import unittest

import numpy as np

from diffusers import (
    DPMSolverMultistepScheduler,
    EulerAncestralDiscreteScheduler,
    EulerDiscreteScheduler,
    LMSDiscreteScheduler,
    OnnxStableDiffusionImgaImgPipeline,
    PNDMScheduler,
)
from diffusers.utils import floats_tensor
from diffusers.utils.testing_utils import (
    is_onnx_available,
    load_image,
    nightly,
    require_onnxruntime,
    require_torch_gpu,
)

from ..test_pipelines_onnx_common import OnnxPipelineTesterMixin


if is_onnx_available():
   import onnxruntime as ort







class   lowerCamelCase	(						A_							,				unittest.TestCase						):
     UpperCAmelCase__       :							List[str]     =	"hf-internal-testing/tiny-random-OnnxStableDiffusionPipeline"





     def 	UpperCAmelCase(self   :					Tuple       ,					_A   :					Optional[int]=0     )		->		int:
         snake_case				  =      floats_tensor((1, 3, 1_2_8, 1_2_8)       ,					rng=random.Random(_A     )     )
         snake_case				  =      np.random.RandomState(_A     )
         snake_case				  =      {
             "prompt": "A painting of a squirrel eating a burger",
             "image": image,
             "generator": generator,
             "num_inference_steps": 3,
             "strength": 0.75,
             "guidance_scale": 7.5,
             "output_type": "numpy",
         }
         return inputs





     def 	UpperCAmelCase(self   :					Tuple     )		->		Any:
         snake_case				  =      OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint       ,					provider="CPUExecutionProvider"     )
         pipe.set_progress_bar_config(disable=_A     )

         snake_case				  =      self.get_dummy_inputs()
         snake_case				  =      pipe(**_A     ).images
         snake_case				  =      image[0, -3:, -3:, -1].flatten()

         assert image.shape == (1, 1_2_8, 1_2_8, 3)
         snake_case				  =      np.array([0.6_96_43, 0.5_84_84, 0.5_03_14, 0.5_87_60, 0.5_53_68, 0.5_96_43, 0.5_15_29, 0.4_12_17, 0.4_90_87]     )
         assert np.abs(image_slice - expected_slice     ).max() < 1E-1





     def 	UpperCAmelCase(self   :					List[Any]     )		->		int:
         snake_case				  =      OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint       ,					provider="CPUExecutionProvider"     )
         snake_case				  =      PNDMScheduler.from_config(pipe.scheduler.config       ,					skip_prk_steps=_A     )
         pipe.set_progress_bar_config(disable=_A     )

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

         assert image.shape == (1, 1_2_8, 1_2_8, 3)
         snake_case				  =      np.array([0.6_17_37, 0.5_46_42, 0.5_31_83, 0.5_44_65, 0.5_27_42, 0.6_05_25, 0.4_99_69, 0.4_06_55, 0.4_81_54]     )

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





     def 	UpperCAmelCase(self   :					Dict     )		->		int:
         snake_case				  =      OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint       ,					provider="CPUExecutionProvider"     )
         snake_case				  =      LMSDiscreteScheduler.from_config(pipe.scheduler.config     )
         pipe.set_progress_bar_config(disable=_A     )

         # warmup pass to apply optimizations
         snake_case				  =      pipe(**self.get_dummy_inputs()     )

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

         assert image.shape == (1, 1_2_8, 1_2_8, 3)
         snake_case				  =      np.array([0.5_27_61, 0.5_99_77, 0.4_90_33, 0.4_96_19, 0.5_42_82, 0.5_03_11, 0.4_76_00, 0.4_09_18, 0.4_52_03]     )

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





     def 	UpperCAmelCase(self   :					List[str]     )		->		Dict:
         snake_case				  =      OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint       ,					provider="CPUExecutionProvider"     )
         snake_case				  =      EulerDiscreteScheduler.from_config(pipe.scheduler.config     )
         pipe.set_progress_bar_config(disable=_A     )

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

         assert image.shape == (1, 1_2_8, 1_2_8, 3)
         snake_case				  =      np.array([0.5_29_11, 0.6_00_04, 0.4_92_29, 0.4_98_05, 0.5_45_02, 0.5_06_80, 0.4_77_77, 0.4_10_28, 0.4_53_04]     )

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





     def 	UpperCAmelCase(self   :					str     )		->		int:
         snake_case				  =      OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint       ,					provider="CPUExecutionProvider"     )
         snake_case				  =      EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config     )
         pipe.set_progress_bar_config(disable=_A     )

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

         assert image.shape == (1, 1_2_8, 1_2_8, 3)
         snake_case				  =      np.array([0.5_29_11, 0.6_00_04, 0.4_92_29, 0.4_98_05, 0.5_45_02, 0.5_06_80, 0.4_77_77, 0.4_10_28, 0.4_53_04]     )

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





     def 	UpperCAmelCase(self   :					int     )		->		Any:
         snake_case				  =      OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint       ,					provider="CPUExecutionProvider"     )
         snake_case				  =      DPMSolverMultistepScheduler.from_config(pipe.scheduler.config     )
         pipe.set_progress_bar_config(disable=_A     )

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

         assert image.shape == (1, 1_2_8, 1_2_8, 3)
         snake_case				  =      np.array([0.6_53_31, 0.5_82_77, 0.4_82_04, 0.5_60_59, 0.5_36_65, 0.5_62_35, 0.5_09_69, 0.4_00_09, 0.4_65_52]     )

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







@nightly
@require_onnxruntime
@require_torch_gpu
class   lowerCamelCase	(						unittest.TestCase						):





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





     @property
     def 	UpperCAmelCase(self   :					str     )		->		List[str]:
         snake_case				  =      ort.SessionOptions()
         snake_case				  =      False
         return options





     def 	UpperCAmelCase(self   :					List[Any]     )		->		Dict:
         snake_case				  =      load_image(
             "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
             "/img2img/sketch-mountains-input.jpg"     )
         snake_case				  =      init_image.resize((7_6_8, 5_1_2)     )
         # using the PNDM scheduler by default
         snake_case				  =      OnnxStableDiffusionImgaImgPipeline.from_pretrained(
             "CompVis/stable-diffusion-v1-4"       ,					revision="onnx"       ,					safety_checker=_A       ,					feature_extractor=_A       ,					provider=self.gpu_provider       ,					sess_options=self.gpu_options       ,					)
         pipe.set_progress_bar_config(disable=_A     )

         snake_case				  =      "A fantasy landscape, trending on artstation"

         snake_case				  =      np.random.RandomState(0     )
         snake_case				  =      pipe(
             prompt=_A       ,					image=_A       ,					strength=0.75       ,					guidance_scale=7.5       ,					num_inference_steps=1_0       ,					generator=_A       ,					output_type="np"       ,					)
         snake_case				  =      output.images
         snake_case				  =      images[0, 2_5_5:2_5_8, 3_8_3:3_8_6, -1]

         assert images.shape == (1, 5_1_2, 7_6_8, 3)
         snake_case				  =      np.array([0.49_09, 0.50_59, 0.53_72, 0.46_23, 0.48_76, 0.50_49, 0.48_20, 0.49_56, 0.50_19]     )
         # TODO: lower the tolerance after finding the cause of onnxruntime reproducibility issues

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





     def 	UpperCAmelCase(self   :					Optional[int]     )		->		Optional[Any]:
         snake_case				  =      load_image(
             "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
             "/img2img/sketch-mountains-input.jpg"     )
         snake_case				  =      init_image.resize((7_6_8, 5_1_2)     )
         snake_case				  =      LMSDiscreteScheduler.from_pretrained(
             "runwayml/stable-diffusion-v1-5"       ,					subfolder="scheduler"       ,					revision="onnx"     )
         snake_case				  =      OnnxStableDiffusionImgaImgPipeline.from_pretrained(
             "runwayml/stable-diffusion-v1-5"       ,					revision="onnx"       ,					scheduler=_A       ,					safety_checker=_A       ,					feature_extractor=_A       ,					provider=self.gpu_provider       ,					sess_options=self.gpu_options       ,					)
         pipe.set_progress_bar_config(disable=_A     )

         snake_case				  =      "A fantasy landscape, trending on artstation"

         snake_case				  =      np.random.RandomState(0     )
         snake_case				  =      pipe(
             prompt=_A       ,					image=_A       ,					strength=0.75       ,					guidance_scale=7.5       ,					num_inference_steps=2_0       ,					generator=_A       ,					output_type="np"       ,					)
         snake_case				  =      output.images
         snake_case				  =      images[0, 2_5_5:2_5_8, 3_8_3:3_8_6, -1]

         assert images.shape == (1, 5_1_2, 7_6_8, 3)
         snake_case				  =      np.array([0.80_43, 0.9_26, 0.95_81, 0.81_19, 0.89_54, 0.9_13, 0.72_09, 0.74_63, 0.74_31]     )
         # TODO: lower the tolerance after finding the cause of onnxruntime reproducibility issues

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



import unittest

from transformers import is_torch_available
from transformers.testing_utils import require_sentencepiece, require_tokenizers, require_torch, slow, torch_device


if is_torch_available():
   from transformers import AutoModelForSeqaSeqLM, AutoTokenizer







@require_torch
@require_sentencepiece
@require_tokenizers
class   lowerCamelCase	(						unittest.TestCase						):





     @slow
     def 	UpperCAmelCase(self   :					List[str]     )		->		Optional[int]:
         snake_case				  =      AutoModelForSeqaSeqLM.from_pretrained("google/mt5-small"       ,					return_dict=_A     ).to(_A     )
         snake_case				  =      AutoTokenizer.from_pretrained("google/mt5-small"     )

         snake_case				  =      tokenizer("Hello there"       ,					return_tensors="pt"     ).input_ids
         snake_case				  =      tokenizer("Hi I am"       ,					return_tensors="pt"     ).input_ids

         snake_case				  =      model(input_ids.to(_A     )       ,					labels=labels.to(_A     )     ).loss
         snake_case				  =      -(labels.shape[-1] * loss.item())

         snake_case				  =      -84.91_27
         self.assertTrue(abs(mtf_score - EXPECTED_SCORE     ) < 1E-4     )

