

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







   '''simple docstring'''




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





   def      _snake_case		(       self					:			List[Any]     , __lowerCamelCase					:			Dict=0			):
       SCREAMING_SNAKE_CASE										=   floats_tensor((1, 3, 128, 128)     , rng=random.Random(_lowercase			)			)
       SCREAMING_SNAKE_CASE										=   np.random.RandomState(_lowercase			)
       SCREAMING_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      _snake_case		(       self					:			int			):
       SCREAMING_SNAKE_CASE										=   OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint     , provider="CPUExecutionProvider"			)
       pipe.set_progress_bar_config(disable=_lowercase			)

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

       assert image.shape == (1, 128, 128, 3)
       SCREAMING_SNAKE_CASE										=   np.array([0.69_643, 0.58_484, 0.50_314, 0.58_760, 0.55_368, 0.59_643, 0.51_529, 0.41_217, 0.49_087]			)
       assert np.abs(image_slice - expected_slice			).max() < 1e-1





   def      _snake_case		(       self					:			Dict			):
       SCREAMING_SNAKE_CASE										=   OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint     , provider="CPUExecutionProvider"			)
       SCREAMING_SNAKE_CASE										=   PNDMScheduler.from_config(pipe.scheduler.config     , skip_prk_steps=_lowercase			)
       pipe.set_progress_bar_config(disable=_lowercase			)

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

       assert image.shape == (1, 128, 128, 3)
       SCREAMING_SNAKE_CASE										=   np.array([0.61_737, 0.54_642, 0.53_183, 0.54_465, 0.52_742, 0.60_525, 0.49_969, 0.40_655, 0.48_154]			)

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





   def      _snake_case		(       self					:			Any			):
       SCREAMING_SNAKE_CASE										=   OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint     , provider="CPUExecutionProvider"			)
       SCREAMING_SNAKE_CASE										=   LMSDiscreteScheduler.from_config(pipe.scheduler.config			)
       pipe.set_progress_bar_config(disable=_lowercase			)

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

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

       assert image.shape == (1, 128, 128, 3)
       SCREAMING_SNAKE_CASE										=   np.array([0.52_761, 0.59_977, 0.49_033, 0.49_619, 0.54_282, 0.50_311, 0.47_600, 0.40_918, 0.45_203]			)

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





   def      _snake_case		(       self					:			List[str]			):
       SCREAMING_SNAKE_CASE										=   OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint     , provider="CPUExecutionProvider"			)
       SCREAMING_SNAKE_CASE										=   EulerDiscreteScheduler.from_config(pipe.scheduler.config			)
       pipe.set_progress_bar_config(disable=_lowercase			)

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

       assert image.shape == (1, 128, 128, 3)
       SCREAMING_SNAKE_CASE										=   np.array([0.52_911, 0.60_004, 0.49_229, 0.49_805, 0.54_502, 0.50_680, 0.47_777, 0.41_028, 0.45_304]			)

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





   def      _snake_case		(       self					:			Optional[int]			):
       SCREAMING_SNAKE_CASE										=   OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint     , provider="CPUExecutionProvider"			)
       SCREAMING_SNAKE_CASE										=   EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config			)
       pipe.set_progress_bar_config(disable=_lowercase			)

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

       assert image.shape == (1, 128, 128, 3)
       SCREAMING_SNAKE_CASE										=   np.array([0.52_911, 0.60_004, 0.49_229, 0.49_805, 0.54_502, 0.50_680, 0.47_777, 0.41_028, 0.45_304]			)

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





   def      _snake_case		(       self					:			List[Any]			):
       SCREAMING_SNAKE_CASE										=   OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint     , provider="CPUExecutionProvider"			)
       SCREAMING_SNAKE_CASE										=   DPMSolverMultistepScheduler.from_config(pipe.scheduler.config			)
       pipe.set_progress_bar_config(disable=_lowercase			)

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

       assert image.shape == (1, 128, 128, 3)
       SCREAMING_SNAKE_CASE										=   np.array([0.65_331, 0.58_277, 0.48_204, 0.56_059, 0.53_665, 0.56_235, 0.50_969, 0.40_009, 0.46_552]			)

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





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







   '''simple docstring'''




   @property
   def      _snake_case		(       self					:			Optional[int]			):
       return (
           "CUDAExecutionProvider",
           {
               "gpu_mem_limit": "15000000000",  # 15GB
               "arena_extend_strategy": "kSameAsRequested",
           },
       )





   @property
   def      _snake_case		(       self					:			Any			):
       SCREAMING_SNAKE_CASE										=   ort.SessionOptions()
       SCREAMING_SNAKE_CASE										=   False
       return options





   def      _snake_case		(       self					:			int			):
       SCREAMING_SNAKE_CASE										=   load_image(
           "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
           "/img2img/sketch-mountains-input.jpg"			)
       SCREAMING_SNAKE_CASE										=   init_image.resize((768, 512)			)
       # using the PNDM scheduler by default
       SCREAMING_SNAKE_CASE										=   OnnxStableDiffusionImgaImgPipeline.from_pretrained(
           "CompVis/stable-diffusion-v1-4"     , revision="onnx"     , safety_checker=_lowercase     , feature_extractor=_lowercase     , provider=self.gpu_provider     , sess_options=self.gpu_options     , )
       pipe.set_progress_bar_config(disable=_lowercase			)

       SCREAMING_SNAKE_CASE										=   "A fantasy landscape, trending on artstation"

       SCREAMING_SNAKE_CASE										=   np.random.RandomState(0			)
       SCREAMING_SNAKE_CASE										=   pipe(
           prompt=_lowercase     , image=_lowercase     , strength=0.75     , guidance_scale=7.5     , num_inference_steps=10     , generator=_lowercase     , output_type="np"     , )
       SCREAMING_SNAKE_CASE										=   output.images
       SCREAMING_SNAKE_CASE										=   images[0, 255:258, 383:386, -1]

       assert images.shape == (1, 512, 768, 3)
       SCREAMING_SNAKE_CASE										=   np.array([0.4_909, 0.5_059, 0.5_372, 0.4_623, 0.4_876, 0.5_049, 0.4_820, 0.4_956, 0.5_019]			)
       # TODO: lower the tolerance after finding the cause of onnxruntime reproducibility issues

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





   def      _snake_case		(       self					:			Tuple			):
       SCREAMING_SNAKE_CASE										=   load_image(
           "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
           "/img2img/sketch-mountains-input.jpg"			)
       SCREAMING_SNAKE_CASE										=   init_image.resize((768, 512)			)
       SCREAMING_SNAKE_CASE										=   LMSDiscreteScheduler.from_pretrained(
           "runwayml/stable-diffusion-v1-5"     , subfolder="scheduler"     , revision="onnx"			)
       SCREAMING_SNAKE_CASE										=   OnnxStableDiffusionImgaImgPipeline.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     , )
       pipe.set_progress_bar_config(disable=_lowercase			)

       SCREAMING_SNAKE_CASE										=   "A fantasy landscape, trending on artstation"

       SCREAMING_SNAKE_CASE										=   np.random.RandomState(0			)
       SCREAMING_SNAKE_CASE										=   pipe(
           prompt=_lowercase     , image=_lowercase     , strength=0.75     , guidance_scale=7.5     , num_inference_steps=20     , generator=_lowercase     , output_type="np"     , )
       SCREAMING_SNAKE_CASE										=   output.images
       SCREAMING_SNAKE_CASE										=   images[0, 255:258, 383:386, -1]

       assert images.shape == (1, 512, 768, 3)
       SCREAMING_SNAKE_CASE										=   np.array([0.8_043, 0.926, 0.9_581, 0.8_119, 0.8_954, 0.913, 0.7_209, 0.7_463, 0.7_431]			)
       # TODO: lower the tolerance after finding the cause of onnxruntime reproducibility issues

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





'''simple docstring'''
from argparse import ArgumentParser

from .add_new_model import AddNewModelCommand
from .add_new_model_like import AddNewModelLikeCommand
from .convert import ConvertCommand
from .download import DownloadCommand
from .env import EnvironmentCommand
from .lfs import LfsCommands
from .pt_to_tf import PTtoTFCommand
from .run import RunCommand
from .serving import ServeCommand
from .user import UserCommands





def 			lowerCamelCase(							)		->		int:
			A_									= ArgumentParser('Transformers CLI tool'       ,usage='transformers-cli <command> [<args>]'     )
			A_									= parser.add_subparsers(help='transformers-cli command helpers'     )

			# Register commands
			ConvertCommand.register_subcommand(SCREAMING_SNAKE_CASE_     )
			DownloadCommand.register_subcommand(SCREAMING_SNAKE_CASE_     )
			EnvironmentCommand.register_subcommand(SCREAMING_SNAKE_CASE_     )
			RunCommand.register_subcommand(SCREAMING_SNAKE_CASE_     )
			ServeCommand.register_subcommand(SCREAMING_SNAKE_CASE_     )
			UserCommands.register_subcommand(SCREAMING_SNAKE_CASE_     )
			AddNewModelCommand.register_subcommand(SCREAMING_SNAKE_CASE_     )
			AddNewModelLikeCommand.register_subcommand(SCREAMING_SNAKE_CASE_     )
			LfsCommands.register_subcommand(SCREAMING_SNAKE_CASE_     )
			PTtoTFCommand.register_subcommand(SCREAMING_SNAKE_CASE_     )

			# Let's go
			A_									= parser.parse_args()

			if not hasattr(SCREAMING_SNAKE_CASE_       ,'func'     ):
						parser.print_help()
						exit(1     )

			# Run
			A_									= args.func(SCREAMING_SNAKE_CASE_     )
			service.run()


if __name__ == "__main__":
							main()




