

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

			def        SCREAMING_SNAKE_CASE					(       self					:    Optional[int]     ,					__UpperCAmelCase					:    List[Any]=0							) ->     Any:
						SCREAMING_SNAKE_CASE__										=   np.random.RandomState(__UpperCAmelCase							)
						SCREAMING_SNAKE_CASE__										=   {
						    """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        SCREAMING_SNAKE_CASE					(       self					:    List[str]							) ->     Dict:
						SCREAMING_SNAKE_CASE__										=   OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint     ,					provider="""CPUExecutionProvider"""							)
						pipe.set_progress_bar_config(disable=__UpperCAmelCase							)

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

						assert image.shape == (1, 1_2_8, 1_2_8, 3)
						SCREAMING_SNAKE_CASE__										=   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        SCREAMING_SNAKE_CASE					(       self					:    int							) ->     Optional[int]:
						SCREAMING_SNAKE_CASE__										=   OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint     ,					provider="""CPUExecutionProvider"""							)
						SCREAMING_SNAKE_CASE__										=   PNDMScheduler.from_config(pipe.scheduler.config     ,					skip_prk_steps=__UpperCAmelCase							)
						pipe.set_progress_bar_config(disable=__UpperCAmelCase							)

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

						assert image.shape == (1, 1_2_8, 1_2_8, 3)
						SCREAMING_SNAKE_CASE__										=   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        SCREAMING_SNAKE_CASE					(       self					:    Any							) ->     Union[str, Any]:
						SCREAMING_SNAKE_CASE__										=   OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint     ,					provider="""CPUExecutionProvider"""							)
						SCREAMING_SNAKE_CASE__										=   LMSDiscreteScheduler.from_config(pipe.scheduler.config							)
						pipe.set_progress_bar_config(disable=__UpperCAmelCase							)

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

						assert image.shape == (1, 1_2_8, 1_2_8, 3)
						SCREAMING_SNAKE_CASE__										=   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        SCREAMING_SNAKE_CASE					(       self					:    List[str]							) ->     str:
						SCREAMING_SNAKE_CASE__										=   OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint     ,					provider="""CPUExecutionProvider"""							)
						SCREAMING_SNAKE_CASE__										=   EulerDiscreteScheduler.from_config(pipe.scheduler.config							)
						pipe.set_progress_bar_config(disable=__UpperCAmelCase							)

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

						assert image.shape == (1, 1_2_8, 1_2_8, 3)
						SCREAMING_SNAKE_CASE__										=   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        SCREAMING_SNAKE_CASE					(       self					:    List[str]							) ->     Any:
						SCREAMING_SNAKE_CASE__										=   OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint     ,					provider="""CPUExecutionProvider"""							)
						SCREAMING_SNAKE_CASE__										=   EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config							)
						pipe.set_progress_bar_config(disable=__UpperCAmelCase							)

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

						assert image.shape == (1, 1_2_8, 1_2_8, 3)
						SCREAMING_SNAKE_CASE__										=   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        SCREAMING_SNAKE_CASE					(       self					:    Dict							) ->     List[str]:
						SCREAMING_SNAKE_CASE__										=   OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint     ,					provider="""CPUExecutionProvider"""							)
						SCREAMING_SNAKE_CASE__										=   DPMSolverMultistepScheduler.from_config(pipe.scheduler.config							)
						pipe.set_progress_bar_config(disable=__UpperCAmelCase							)

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

						assert image.shape == (1, 1_2_8, 1_2_8, 3)
						SCREAMING_SNAKE_CASE__										=   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        SCREAMING_SNAKE_CASE					(       self					:    Any							) ->     Optional[Any]:
						SCREAMING_SNAKE_CASE__										=   OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint     ,					provider="""CPUExecutionProvider"""							)
						pipe.set_progress_bar_config(disable=__UpperCAmelCase							)

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

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

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

						SCREAMING_SNAKE_CASE__										=   pipe.tokenizer(
						    __UpperCAmelCase     ,					padding="""max_length"""     ,					max_length=pipe.tokenizer.model_max_length     ,					truncation=__UpperCAmelCase     ,					return_tensors="""np"""     ,					)
						SCREAMING_SNAKE_CASE__										=   text_inputs["""input_ids"""]

						SCREAMING_SNAKE_CASE__										=   pipe.text_encoder(input_ids=text_inputs.astype(np.intaa							)							)[0]

						SCREAMING_SNAKE_CASE__										=   prompt_embeds

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

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





			def        SCREAMING_SNAKE_CASE					(       self					:    List[Any]							) ->     int:
						SCREAMING_SNAKE_CASE__										=   OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint     ,					provider="""CPUExecutionProvider"""							)
						pipe.set_progress_bar_config(disable=__UpperCAmelCase							)

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

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

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

						SCREAMING_SNAKE_CASE__										=   []
						for p in [prompt, negative_prompt]:
									SCREAMING_SNAKE_CASE__										=   pipe.tokenizer(
									    __UpperCAmelCase     ,					padding="""max_length"""     ,					max_length=pipe.tokenizer.model_max_length     ,					truncation=__UpperCAmelCase     ,					return_tensors="""np"""     ,					)
									SCREAMING_SNAKE_CASE__										=   text_inputs["""input_ids"""]

									embeds.append(pipe.text_encoder(input_ids=text_inputs.astype(np.intaa							)							)[0]							)

						SCREAMING_SNAKE_CASE__				,		SCREAMING_SNAKE_CASE__										=   embeds

						# forward
						SCREAMING_SNAKE_CASE__										=   pipe(**__UpperCAmelCase							)
						SCREAMING_SNAKE_CASE__										=   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        SCREAMING_SNAKE_CASE					(       self					:    List[Any]							) ->     Optional[int]:
						return (
						    "CUDAExecutionProvider",
						    {
						        "gpu_mem_limit": "15000000000",  # 15GB
						        "arena_extend_strategy": "kSameAsRequested",
						    },
						)

			@property
			def        SCREAMING_SNAKE_CASE					(       self					:    Tuple							) ->     Any:
						SCREAMING_SNAKE_CASE__										=   ort.SessionOptions()
						SCREAMING_SNAKE_CASE__										=   False
						return options

			def        SCREAMING_SNAKE_CASE					(       self					:    List[Any]							) ->     Any:
						# using the PNDM scheduler by default
						SCREAMING_SNAKE_CASE__										=   OnnxStableDiffusionPipeline.from_pretrained(
						    """CompVis/stable-diffusion-v1-4"""     ,					revision="""onnx"""     ,					safety_checker=__UpperCAmelCase     ,					feature_extractor=__UpperCAmelCase     ,					provider=self.gpu_provider     ,					sess_options=self.gpu_options     ,					)
						sd_pipe.set_progress_bar_config(disable=__UpperCAmelCase							)

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

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

						assert image.shape == (1, 5_1_2, 5_1_2, 3)
						SCREAMING_SNAKE_CASE__										=   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        SCREAMING_SNAKE_CASE					(       self					:    Union[str, Any]							) ->     int:
						SCREAMING_SNAKE_CASE__										=   DDIMScheduler.from_pretrained(
						    """runwayml/stable-diffusion-v1-5"""     ,					subfolder="""scheduler"""     ,					revision="""onnx"""							)
						SCREAMING_SNAKE_CASE__										=   OnnxStableDiffusionPipeline.from_pretrained(
						    """runwayml/stable-diffusion-v1-5"""     ,					revision="""onnx"""     ,					scheduler=__UpperCAmelCase     ,					safety_checker=__UpperCAmelCase     ,					feature_extractor=__UpperCAmelCase     ,					provider=self.gpu_provider     ,					sess_options=self.gpu_options     ,					)
						sd_pipe.set_progress_bar_config(disable=__UpperCAmelCase							)

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

						assert image.shape == (1, 5_1_2, 5_1_2, 3)
						SCREAMING_SNAKE_CASE__										=   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        SCREAMING_SNAKE_CASE					(       self					:    List[str]							) ->     Any:
						SCREAMING_SNAKE_CASE__										=   LMSDiscreteScheduler.from_pretrained(
						    """runwayml/stable-diffusion-v1-5"""     ,					subfolder="""scheduler"""     ,					revision="""onnx"""							)
						SCREAMING_SNAKE_CASE__										=   OnnxStableDiffusionPipeline.from_pretrained(
						    """runwayml/stable-diffusion-v1-5"""     ,					revision="""onnx"""     ,					scheduler=__UpperCAmelCase     ,					safety_checker=__UpperCAmelCase     ,					feature_extractor=__UpperCAmelCase     ,					provider=self.gpu_provider     ,					sess_options=self.gpu_options     ,					)
						sd_pipe.set_progress_bar_config(disable=__UpperCAmelCase							)

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

						assert image.shape == (1, 5_1_2, 5_1_2, 3)
						SCREAMING_SNAKE_CASE__										=   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        SCREAMING_SNAKE_CASE					(       self					:    List[Any]							) ->     Union[str, Any]:
						SCREAMING_SNAKE_CASE__										=   0

						def test_callback_fn(__UpperCAmelCase					:    int     ,					__UpperCAmelCase					:    int     ,					__UpperCAmelCase					:    np.ndarray							) -> None:
									SCREAMING_SNAKE_CASE__										=   True
									nonlocal number_of_steps
									number_of_steps += 1
									if step == 0:
												assert latents.shape == (1, 4, 6_4, 6_4)
												SCREAMING_SNAKE_CASE__										=   latents[0, -3:, -3:, -1]
												SCREAMING_SNAKE_CASE__										=   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)
												SCREAMING_SNAKE_CASE__										=   latents[0, -3:, -3:, -1]
												SCREAMING_SNAKE_CASE__										=   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

						SCREAMING_SNAKE_CASE__										=   False

						SCREAMING_SNAKE_CASE__										=   OnnxStableDiffusionPipeline.from_pretrained(
						    """runwayml/stable-diffusion-v1-5"""     ,					revision="""onnx"""     ,					safety_checker=__UpperCAmelCase     ,					feature_extractor=__UpperCAmelCase     ,					provider=self.gpu_provider     ,					sess_options=self.gpu_options     ,					)
						pipe.set_progress_bar_config(disable=__UpperCAmelCase							)

						SCREAMING_SNAKE_CASE__										=   """Andromeda galaxy in a bottle"""

						SCREAMING_SNAKE_CASE__										=   np.random.RandomState(0							)
						pipe(
						    prompt=__UpperCAmelCase     ,					num_inference_steps=5     ,					guidance_scale=7.5     ,					generator=__UpperCAmelCase     ,					callback=__UpperCAmelCase     ,					callback_steps=1     ,					)
						assert test_callback_fn.has_been_called
						assert number_of_steps == 6





			def        SCREAMING_SNAKE_CASE					(       self					:    List[Any]							) ->     str:
						SCREAMING_SNAKE_CASE__										=   OnnxStableDiffusionPipeline.from_pretrained(
						    """runwayml/stable-diffusion-v1-5"""     ,					revision="""onnx"""     ,					safety_checker=__UpperCAmelCase     ,					feature_extractor=__UpperCAmelCase     ,					provider=self.gpu_provider     ,					sess_options=self.gpu_options     ,					)
						assert isinstance(__UpperCAmelCase     ,					__UpperCAmelCase							)
						assert pipe.safety_checker is None

						SCREAMING_SNAKE_CASE__										=   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(__UpperCAmelCase							)
									SCREAMING_SNAKE_CASE__										=   OnnxStableDiffusionPipeline.from_pretrained(__UpperCAmelCase							)

						# sanity check that the pipeline still works
						assert pipe.safety_checker is None
						SCREAMING_SNAKE_CASE__										=   pipe("""example prompt"""     ,					num_inference_steps=2							).images[0]
						assert image is not None





"""simple docstring"""




import logging
from dataclasses import dataclass, field
from typing import Optional

from seqaseq_trainer import arg_to_scheduler

from transformers import TrainingArguments


A_				:      Optional[Any]           =			logging.getLogger(__name__)





@dataclass
class  lowerCamelCase    (A__      ):
			lowerCamelCase__ :		Optional[float]       		=				field(
			    default=0.0      ,metadata={'help': 'The label smoothing epsilon to apply (if not zero).'}      )
			lowerCamelCase__ :		bool       		=				field(default=A__      ,metadata={'help': 'Whether to SortishSamler or not.'}      )
			lowerCamelCase__ :		bool       		=				field(
			    default=A__      ,metadata={'help': 'Whether to use generate to calculate generative metrics (ROUGE, BLEU).'}      )
			lowerCamelCase__ :		bool       		=				field(default=A__      ,metadata={'help': 'whether to use adafactor'}      )
			lowerCamelCase__ :		Optional[float]       		=				field(
			    default=A__      ,metadata={'help': 'Encoder layer dropout probability. Goes into model.config.'}      )
			lowerCamelCase__ :		Optional[float]       		=				field(
			    default=A__      ,metadata={'help': 'Decoder layer dropout probability. Goes into model.config.'}      )
			lowerCamelCase__ :		Optional[float]       		=				field(default=A__      ,metadata={'help': 'Dropout probability. Goes into model.config.'}      )
			lowerCamelCase__ :		Optional[float]       		=				field(
			    default=A__      ,metadata={'help': 'Attention dropout probability. Goes into model.config.'}      )
			lowerCamelCase__ :		Optional[str]       		=				field(
			    default='linear'      ,metadata={'help': f"Which lr scheduler to use. Selected in {sorted(arg_to_scheduler.keys()      )}"}      ,)


