


'''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	):
    """simple docstring"""
    a													=   "hf-internal-testing/tiny-random-OnnxStableDiffusionPipeline"


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

           _A					:     List[Any]     = self.get_dummy_inputs()
           _A					:     List[str]     = pipe(**SCREAMING_SNAKE_CASE).images
           _A					:     Optional[int]     = image[0, -3:, -3:, -1]

           assert image.shape == (1, 128, 128, 3)
           _A					:     str     = np.array([0.6_5072, 0.5_8492, 0.4_8219, 0.5_5521, 0.5_3180, 0.5_5939, 0.5_0697, 0.3_9800, 0.4_6455])

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


    def        A   (  self      :							Optional[Any]):
           _A					:     Any     = OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint				,							provider='CPUExecutionProvider')
           _A					:     Union[str, Any]     = PNDMScheduler.from_config(pipe.scheduler.config				,							skip_prk_steps=SCREAMING_SNAKE_CASE)
           pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE)

           _A					:     Union[str, Any]     = self.get_dummy_inputs()
           _A					:     List[Any]     = pipe(**SCREAMING_SNAKE_CASE).images
           _A					:     List[str]     = image[0, -3:, -3:, -1]

           assert image.shape == (1, 128, 128, 3)
           _A					:     Dict     = np.array([0.6_5863, 0.5_9425, 0.4_9326, 0.5_6313, 0.5_3875, 0.5_6627, 0.5_1065, 0.3_9777, 0.4_6330])

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


    def        A   (  self      :							Any):
           _A					:     Optional[int]     = OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint				,							provider='CPUExecutionProvider')
           _A					:     Tuple     = LMSDiscreteScheduler.from_config(pipe.scheduler.config)
           pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE)

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

           assert image.shape == (1, 128, 128, 3)
           _A					:     Optional[int]     = np.array([0.5_3755, 0.6_0786, 0.4_7402, 0.4_9488, 0.5_1869, 0.4_9819, 0.4_7985, 0.3_8957, 0.4_4279])

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


    def        A   (  self      :							List[str]):
           _A					:     List[Any]     = OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint				,							provider='CPUExecutionProvider')
           _A					:     Union[str, Any]     = EulerDiscreteScheduler.from_config(pipe.scheduler.config)
           pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE)

           _A					:     int     = self.get_dummy_inputs()
           _A					:     Union[str, Any]     = pipe(**SCREAMING_SNAKE_CASE).images
           _A					:     Optional[int]     = image[0, -3:, -3:, -1]

           assert image.shape == (1, 128, 128, 3)
           _A					:     str     = np.array([0.5_3755, 0.6_0786, 0.4_7402, 0.4_9488, 0.5_1869, 0.4_9819, 0.4_7985, 0.3_8957, 0.4_4279])

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


    def        A   (  self      :							str):
           _A					:     List[str]     = OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint				,							provider='CPUExecutionProvider')
           _A					:     Optional[int]     = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
           pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE)

           _A					:     Optional[Any]     = self.get_dummy_inputs()
           _A					:     Optional[int]     = pipe(**SCREAMING_SNAKE_CASE).images
           _A					:     int     = image[0, -3:, -3:, -1]

           assert image.shape == (1, 128, 128, 3)
           _A					:     List[Any]     = np.array([0.5_3817, 0.6_0812, 0.4_7384, 0.4_9530, 0.5_1894, 0.4_9814, 0.4_7984, 0.3_8958, 0.4_4271])

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


    def        A   (  self      :							Tuple):
           _A					:     Dict     = OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint				,							provider='CPUExecutionProvider')
           _A					:     Optional[Any]     = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
           pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE)

           _A					:     Any     = self.get_dummy_inputs()
           _A					:     int     = pipe(**SCREAMING_SNAKE_CASE).images
           _A					:     str     = image[0, -3:, -3:, -1]

           assert image.shape == (1, 128, 128, 3)
           _A					:     int     = np.array([0.5_3895, 0.6_0808, 0.4_7933, 0.4_9608, 0.5_1886, 0.4_9950, 0.4_8053, 0.3_8957, 0.4_4200])

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


    def        A   (  self      :							List[Any]):
           _A					:     List[Any]     = OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint				,							provider='CPUExecutionProvider')
           pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE)

           _A					:     Union[str, Any]     = self.get_dummy_inputs()
           _A					:     Dict     = 3 * [inputs['prompt']]

           # forward
           _A					:     int     = pipe(**SCREAMING_SNAKE_CASE)
           _A					:     Dict     = output.images[0, -3:, -3:, -1]

           _A					:     List[Any]     = self.get_dummy_inputs()
           _A					:     Dict     = 3 * [inputs.pop('prompt')]

           _A					:     Dict     = pipe.tokenizer(
               SCREAMING_SNAKE_CASE				,							padding='max_length'				,							max_length=pipe.tokenizer.model_max_length				,							truncation=SCREAMING_SNAKE_CASE				,							return_tensors='np'				,							)
           _A					:     List[Any]     = text_inputs['input_ids']

           _A					:     List[Any]     = pipe.text_encoder(input_ids=text_inputs.astype(np.intaa))[0]

           _A					:     Any     = prompt_embeds

           # forward
           _A					:     int     = pipe(**SCREAMING_SNAKE_CASE)
           _A					:     Any     = output.images[0, -3:, -3:, -1]

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


    def        A   (  self      :							Optional[Any]):
           _A					:     int     = OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint				,							provider='CPUExecutionProvider')
           pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE)

           _A					:     List[Any]     = self.get_dummy_inputs()
           _A					:     List[Any]     = 3 * ['this is a negative prompt']
           _A					:     Union[str, Any]     = negative_prompt
           _A					:     str     = 3 * [inputs['prompt']]

           # forward
           _A					:     Dict     = pipe(**SCREAMING_SNAKE_CASE)
           _A					:     Union[str, Any]     = output.images[0, -3:, -3:, -1]

           _A					:     Optional[Any]     = self.get_dummy_inputs()
           _A					:     List[Any]     = 3 * [inputs.pop('prompt')]

           _A					:     List[str]     = []
           for p in [prompt, negative_prompt]:
                  _A					:     Union[str, Any]     = pipe.tokenizer(
                      SCREAMING_SNAKE_CASE				,							padding='max_length'				,							max_length=pipe.tokenizer.model_max_length				,							truncation=SCREAMING_SNAKE_CASE				,							return_tensors='np'				,							)
                  _A					:     Optional[int]     = text_inputs['input_ids']

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

           _A					:     List[str]     = embeds

           # forward
           _A					:     Optional[int]     = pipe(**SCREAMING_SNAKE_CASE)
           _A					:     List[Any]     = 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	):
    """simple docstring"""


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


    @property
    def        A   (  self      :							Tuple):
           _A					:     List[Any]     = ort.SessionOptions()
           _A					:     Optional[int]     = False
           return options


    def        A   (  self      :							Dict):
           # using the PNDM scheduler by default
           _A					:     str     = OnnxStableDiffusionPipeline.from_pretrained(
               'CompVis/stable-diffusion-v1-4'				,							revision='onnx'				,							safety_checker=SCREAMING_SNAKE_CASE				,							feature_extractor=SCREAMING_SNAKE_CASE				,							provider=self.gpu_provider				,							sess_options=self.gpu_options				,							)
           sd_pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE)

           _A					:     List[str]     = 'A painting of a squirrel eating a burger'
           np.random.seed(0)
           _A					:     Union[str, Any]     = sd_pipe([prompt]				,							guidance_scale=6.0				,							num_inference_steps=10				,							output_type='np')
           _A					:     List[Any]     = output.images

           _A					:     Dict     = image[0, -3:, -3:, -1]

           assert image.shape == (1, 512, 512, 3)
           _A					:     Optional[int]     = np.array([0.0452, 0.0390, 0.0087, 0.0350, 0.0617, 0.0364, 0.0544, 0.0523, 0.0720])

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


    def        A   (  self      :							str):
           _A					:     List[Any]     = DDIMScheduler.from_pretrained(
               'runwayml/stable-diffusion-v1-5'				,							subfolder='scheduler'				,							revision='onnx')
           _A					:     List[Any]     = OnnxStableDiffusionPipeline.from_pretrained(
               'runwayml/stable-diffusion-v1-5'				,							revision='onnx'				,							scheduler=SCREAMING_SNAKE_CASE				,							safety_checker=SCREAMING_SNAKE_CASE				,							feature_extractor=SCREAMING_SNAKE_CASE				,							provider=self.gpu_provider				,							sess_options=self.gpu_options				,							)
           sd_pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE)

           _A					:     Optional[int]     = 'open neural network exchange'
           _A					:     Optional[Any]     = np.random.RandomState(0)
           _A					:     List[Any]     = sd_pipe([prompt]				,							guidance_scale=7.5				,							num_inference_steps=10				,							generator=SCREAMING_SNAKE_CASE				,							output_type='np')
           _A					:     Dict     = output.images
           _A					:     List[Any]     = image[0, -3:, -3:, -1]

           assert image.shape == (1, 512, 512, 3)
           _A					:     Tuple     = np.array([0.2867, 0.1974, 0.1481, 0.7294, 0.7251, 0.6667, 0.4194, 0.5642, 0.6486])

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


    def        A   (  self      :							Optional[Any]):
           _A					:     List[Any]     = LMSDiscreteScheduler.from_pretrained(
               'runwayml/stable-diffusion-v1-5'				,							subfolder='scheduler'				,							revision='onnx')
           _A					:     int     = OnnxStableDiffusionPipeline.from_pretrained(
               'runwayml/stable-diffusion-v1-5'				,							revision='onnx'				,							scheduler=SCREAMING_SNAKE_CASE				,							safety_checker=SCREAMING_SNAKE_CASE				,							feature_extractor=SCREAMING_SNAKE_CASE				,							provider=self.gpu_provider				,							sess_options=self.gpu_options				,							)
           sd_pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE)

           _A					:     Optional[int]     = 'open neural network exchange'
           _A					:     Optional[Any]     = np.random.RandomState(0)
           _A					:     str     = sd_pipe([prompt]				,							guidance_scale=7.5				,							num_inference_steps=10				,							generator=SCREAMING_SNAKE_CASE				,							output_type='np')
           _A					:     List[str]     = output.images
           _A					:     Optional[int]     = image[0, -3:, -3:, -1]

           assert image.shape == (1, 512, 512, 3)
           _A					:     int     = np.array([0.2306, 0.1959, 0.1593, 0.6549, 0.6394, 0.5408, 0.5065, 0.6010, 0.6161])

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


    def        A   (  self      :							str):
           _A					:     Any     = 0

           def test_callback_fn(SCREAMING_SNAKE_CASE      :							int				,							SCREAMING_SNAKE_CASE      :							int				,							SCREAMING_SNAKE_CASE      :							np.ndarray) -> None:
                  _A					:     Dict     = True
                  nonlocal number_of_steps
                  number_of_steps += 1
                  if step == 0:
                         assert latents.shape == (1, 4, 64, 64)
                         _A					:     Union[str, Any]     = latents[0, -3:, -3:, -1]
                         _A					:     int     = np.array(
                             [-0.6772, -0.3835, -1.2456, 0.1905, -1.0974, 0.6967, -1.9353, 0.0178, 1.0167])

                         assert np.abs(latents_slice.flatten() - expected_slice).max() < 1e-3
                  elif step == 5:
                         assert latents.shape == (1, 4, 64, 64)
                         _A					:     int     = latents[0, -3:, -3:, -1]
                         _A					:     str     = np.array(
                             [-0.3351, 0.2241, -0.1837, -0.2325, -0.6577, 0.3393, -0.0241, 0.5899, 1.3875])

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

           _A					:     Dict     = False

           _A					:     Dict     = OnnxStableDiffusionPipeline.from_pretrained(
               'runwayml/stable-diffusion-v1-5'				,							revision='onnx'				,							safety_checker=SCREAMING_SNAKE_CASE				,							feature_extractor=SCREAMING_SNAKE_CASE				,							provider=self.gpu_provider				,							sess_options=self.gpu_options				,							)
           pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE)

           _A					:     str     = 'Andromeda galaxy in a bottle'

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


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

           _A					:     Dict     = 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(SCREAMING_SNAKE_CASE)
                  _A					:     Any     = OnnxStableDiffusionPipeline.from_pretrained(SCREAMING_SNAKE_CASE)

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










'''simple docstring'''







import operator as op


A				:					List[str]    						=			'''scaler.pt'''
A				:					Optional[int]    						=			'''pytorch_model'''
A				:					Union[str, Any]    						=			'''random_states'''
A				:					List[Any]    						=			'''optimizer'''
A				:					Optional[Any]    						=			'''scheduler'''
A				:					Optional[int]    						=			'''pytorch_model.bin'''
A				:					List[str]    						=			'''pytorch_model.bin.index.json'''
A				:					List[Any]    						=			'''model.safetensors'''
A				:					int    						=			'''model.safetensors.index.json'''
A				:					List[str]    						=			'''1.10.2'''
A				:					str    						=			'''py38'''
A				:					int    						=			'''4.17.0'''
A				:					Tuple    						=			['''ml.p3.16xlarge''', '''ml.p3dn.24xlarge''', '''ml.p4dn.24xlarge''']
A				:					Union[str, Any]    						=			['''FULL_SHARD''', '''SHARD_GRAD_OP''', '''NO_SHARD''', '''HYBRID_SHARD''', '''HYBRID_SHARD_ZERO2''']
A				:					Optional[Any]    						=			['''TRANSFORMER_BASED_WRAP''', '''SIZE_BASED_WRAP''', '''NO_WRAP''']
A				:					Dict    						=			['''BACKWARD_PRE''', '''BACKWARD_POST''', '''NO_PREFETCH''']
A				:					Any    						=			['''FULL_STATE_DICT''', '''LOCAL_STATE_DICT''', '''SHARDED_STATE_DICT''']
A				:					Tuple    						=			'''2.0.1'''
A				:					List[str]    						=			['''pdsh''', '''standard''', '''openmpi''', '''mvapich''']
A				:					Union[str, Any]    						=			['''default''', '''reduce-overhead''', '''max-autotune''']

A				:					List[Any]    						=			{'''>''': op.gt, '''>=''': op.ge, '''==''': op.eq, '''!=''': op.ne, '''<=''': op.le, '''<''': op.lt}

# These are the args for `torch.distributed.launch` for pytorch < 1.9
A				:					Tuple    						=			[
    '''nnodes''',
    '''nproc_per_node''',
    '''rdzv_backend''',
    '''rdzv_endpoint''',
    '''rdzv_id''',
    '''rdzv_conf''',
    '''standalone''',
    '''max_restarts''',
    '''monitor_interval''',
    '''start_method''',
    '''role''',
    '''module''',
    '''m''',
    '''no_python''',
    '''run_path''',
    '''log_dir''',
    '''r''',
    '''redirects''',
    '''t''',
    '''tee''',
    '''node_rank''',
    '''master_addr''',
    '''master_port''',
]

A				:					List[str]    						=			['''DEEPSPEED''', '''MULTI_GPU''', '''FSDP''', '''MEGATRON_LM''']
A				:					Any    						=			['''DEEPSPEED''', '''MULTI_XPU''', '''FSDP''']






