


from copy import deepcopy

import torch
import torch.nn.functional as F
from torch.optim import AdamW
from torch.optim.lr_scheduler import LambdaLR
from torch.utils.data import DataLoader

from accelerate.accelerator import Accelerator
from accelerate.state import GradientState
from accelerate.test_utils import RegressionDataset, RegressionModel
from accelerate.utils import DistributedType, is_torch_version, set_seed



def 							UpperCamelCase_( lowerCamelCase_				,					lowerCamelCase_				,					lowerCamelCase_				,					lowerCamelCase_					)       ->		Dict:
    for param, grad_param in zip(model_a.parameters()				,					model_b.parameters()					):
        if not param.requires_grad:
            continue
        if not did_step:
            # Grads should not be in sync
            assert (
                torch.allclose(param.grad				,					grad_param.grad					) is False
            ), F'''Gradients in sync when they should not be at iteration {iteration}:\nmodel_a grad ({param.grad}) == model_b grad ({grad_param.grad})'''
        else:
            # Grads should be in sync
            assert (
                torch.allclose(param.grad				,					grad_param.grad					) is True
            ), F'''Gradients not in sync when they should be at iteration {iteration}:\nmodel_a grad ({param.grad}) != model_b grad ({grad_param.grad})'''



def 							UpperCamelCase_( lowerCamelCase_				,					lowerCamelCase_				,					lowerCamelCase_				,					lowerCamelCase_				,					lowerCamelCase_=True					)       ->		Optional[int]:
    model.train()
    _lowercase			:				List[Any] 						=						model(lowerCamelCase_					)
    _lowercase			:				str 						=						F.mse_loss(lowerCamelCase_				,					target.to(output.device					)					)
    if not do_backward:
        loss /= accelerator.gradient_accumulation_steps
        loss.backward()
    else:
        accelerator.backward(lowerCamelCase_					)



def 							UpperCamelCase_( lowerCamelCase_				,					lowerCamelCase_=False					)       ->		Tuple:
    set_seed(42					)
    _lowercase			:				List[str] 						=						RegressionModel()
    _lowercase			:				Union[str, Any] 						=						deepcopy(lowerCamelCase_					)
    _lowercase			:				str 						=						RegressionDataset(length=80					)
    _lowercase			:				Dict 						=						DataLoader(lowerCamelCase_				,					batch_size=16					)
    model.to(accelerator.device					)
    if sched:
        _lowercase			:				int 						=						AdamW(params=model.parameters()				,					lr=1e-3					)
        _lowercase			:				Any 						=						AdamW(params=ddp_model.parameters()				,					lr=1e-3					)
        _lowercase			:				Optional[int] 						=						LambdaLR(lowerCamelCase_				,					lr_lambda=lambda lowerCamelCase_					: epoch**0.65					)
        _lowercase			:				int 						=						LambdaLR(lowerCamelCase_				,					lr_lambda=lambda lowerCamelCase_					: epoch**0.65					)
    # Make a copy of `model`
    if sched:
        _lowercase						,							_lowercase						,							_lowercase						,							_lowercase			:				int 						=						accelerator.prepare(lowerCamelCase_				,					lowerCamelCase_				,					lowerCamelCase_				,					lowerCamelCase_					)
    else:
        _lowercase						,							_lowercase			:				str 						=						accelerator.prepare(lowerCamelCase_				,					lowerCamelCase_					)
    if sched:
        return (model, opt, sched, dataloader, ddp_model, ddp_opt, ddp_sched)
    return model, ddp_model, dataloader



def 							UpperCamelCase_( lowerCamelCase_					)       ->		Optional[int]:
    # Test when on a single CPU or GPU that the context manager does nothing
    _lowercase						,							_lowercase						,							_lowercase			:				Optional[int] 						=						get_training_setup(lowerCamelCase_					)
    # Use a single batch
    _lowercase						,							_lowercase			:				int 						=						next(iter(lowerCamelCase_					)					).values()
    for iteration in range(3					):
        # Gather the distributed inputs and targs for the base model
        _lowercase						,							_lowercase			:				Optional[int] 						=						accelerator.gather((ddp_input, ddp_target)					)
        _lowercase						,							_lowercase			:				Optional[Any] 						=						input.to(accelerator.device					), target.to(accelerator.device					)
        # Perform our initial ground truth step in non "DDP"
        step_model(lowerCamelCase_				,					lowerCamelCase_				,					lowerCamelCase_				,					lowerCamelCase_					)
        # Do "gradient accumulation" (noop)
        if iteration % 2 == 0:
            # Accumulate grads locally
            with accelerator.no_sync(lowerCamelCase_					):
                step_model(lowerCamelCase_				,					lowerCamelCase_				,					lowerCamelCase_				,					lowerCamelCase_					)
        else:
            # Sync grads
            step_model(lowerCamelCase_				,					lowerCamelCase_				,					lowerCamelCase_				,					lowerCamelCase_					)

        # Since `no_sync` is a noop, `ddp_model` and `model` grads should always be in sync
        check_model_parameters(lowerCamelCase_				,					lowerCamelCase_				,					lowerCamelCase_				,					lowerCamelCase_					)
        for param, ddp_param in zip(model.parameters()				,					ddp_model.parameters()					):
            if not param.requires_grad:
                continue
            assert torch.allclose(
                param.grad				,					ddp_param.grad					), F'''Gradients not in sync when they should be:\nModel grad ({param.grad}) != DDP grad ({ddp_param.grad})'''

        # Shuffle ddp_input on each iteration
        torch.manual_seed(1337 + iteration					)
        _lowercase			:				str 						=						ddp_input[torch.randperm(len(lowerCamelCase_					)					)]



def 							UpperCamelCase_( lowerCamelCase_					)       ->		str:
    # Test on distributed setup that context manager behaves properly
    _lowercase						,							_lowercase						,							_lowercase			:				str 						=						get_training_setup(lowerCamelCase_					)
    # Use a single batch
    _lowercase						,							_lowercase			:				Tuple 						=						next(iter(lowerCamelCase_					)					).values()
    for iteration in range(3					):
        # Gather the distributed inputs and targs for the base model
        _lowercase						,							_lowercase			:				Any 						=						accelerator.gather((ddp_input, ddp_target)					)
        _lowercase						,							_lowercase			:				Dict 						=						input.to(accelerator.device					), target.to(accelerator.device					)
        # Perform our initial ground truth step in non "DDP"
        step_model(lowerCamelCase_				,					lowerCamelCase_				,					lowerCamelCase_				,					lowerCamelCase_					)
        # Do "gradient accumulation" (noop)
        if iteration % 2 == 0:
            # Accumulate grads locally
            with accelerator.no_sync(lowerCamelCase_					):
                step_model(lowerCamelCase_				,					lowerCamelCase_				,					lowerCamelCase_				,					lowerCamelCase_					)
        else:
            # Sync grads
            step_model(lowerCamelCase_				,					lowerCamelCase_				,					lowerCamelCase_				,					lowerCamelCase_					)

        # DDP model and model should only be in sync when not (iteration % 2 == 0)
        for param, ddp_param in zip(model.parameters()				,					ddp_model.parameters()					):
            if not param.requires_grad:
                continue
            if iteration % 2 == 0:
                # Grads should not be in sync
                assert (
                    torch.allclose(param.grad				,					ddp_param.grad					) is False
                ), F'''Gradients in sync when they should not be:\nModel grad ({param.grad}) == DDP grad ({ddp_param.grad})'''
            else:
                # Grads should be in sync
                assert (
                    torch.allclose(param.grad				,					ddp_param.grad					) is True
                ), F'''Gradients not in sync when they should be:\nModel grad ({param.grad}) != DDP grad ({ddp_param.grad})'''

        # Shuffle ddp_input on each iteration
        torch.manual_seed(1337 + iteration					)
        _lowercase			:				Tuple 						=						ddp_input[torch.randperm(len(lowerCamelCase_					)					)]



def 							UpperCamelCase_( lowerCamelCase_=False				,					lowerCamelCase_=False					)       ->		Optional[Any]:
    _lowercase			:				Dict 						=						Accelerator(
        split_batches=lowerCamelCase_				,					dispatch_batches=lowerCamelCase_				,					gradient_accumulation_steps=2					)
    # Test that context manager behaves properly
    _lowercase						,							_lowercase						,							_lowercase			:				Optional[Any] 						=						get_training_setup(lowerCamelCase_					)
    for iteration, batch in enumerate(lowerCamelCase_					):
        _lowercase						,							_lowercase			:				Any 						=						batch.values()
        # Gather the distributed inputs and targs for the base model
        _lowercase						,							_lowercase			:				int 						=						accelerator.gather((ddp_input, ddp_target)					)
        _lowercase						,							_lowercase			:				int 						=						input.to(accelerator.device					), target.to(accelerator.device					)
        # Perform our initial ground truth step in non "DDP"
        step_model(lowerCamelCase_				,					lowerCamelCase_				,					lowerCamelCase_				,					lowerCamelCase_				,					lowerCamelCase_					)
        # Do "gradient accumulation" (noop)
        with accelerator.accumulate(lowerCamelCase_					):
            step_model(lowerCamelCase_				,					lowerCamelCase_				,					lowerCamelCase_				,					lowerCamelCase_					)

        # DDP model and model should only be in sync when not (iteration % 2 == 0)
        for param, ddp_param in zip(model.parameters()				,					ddp_model.parameters()					):
            if not param.requires_grad:
                continue
            if ((iteration + 1) % 2 == 0) or (iteration == len(lowerCamelCase_					) - 1):
                # Grads should be in sync
                assert (
                    torch.allclose(param.grad				,					ddp_param.grad					) is True
                ), F'''Gradients not in sync when they should be at iteration {iteration}:\nModel grad ({param.grad}) != DDP grad ({ddp_param.grad})'''
            else:
                # Grads should not be in sync
                assert (
                    torch.allclose(param.grad				,					ddp_param.grad					) is False
                ), F'''Gradients in sync when they should not be at iteration {iteration}:\nModel grad ({param.grad}) == DDP grad ({ddp_param.grad})'''

        # Shuffle ddp_input on each iteration
        torch.manual_seed(1337 + iteration					)
        _lowercase			:				List[str] 						=						ddp_input[torch.randperm(len(lowerCamelCase_					)					)]
    GradientState._reset_state()



def 							UpperCamelCase_( lowerCamelCase_=False				,					lowerCamelCase_=False					)       ->		Dict:
    _lowercase			:				Optional[int] 						=						Accelerator(
        split_batches=lowerCamelCase_				,					dispatch_batches=lowerCamelCase_				,					gradient_accumulation_steps=2					)
    # Test that context manager behaves properly
    _lowercase						,							_lowercase						,							_lowercase						,							_lowercase						,							_lowercase						,							_lowercase						,							_lowercase			:				List[Any] 						=						get_training_setup(lowerCamelCase_				,					lowerCamelCase_					)
    for iteration, batch in enumerate(lowerCamelCase_					):
        _lowercase						,							_lowercase			:				Any 						=						batch.values()
        # Gather the distributed inputs and targs for the base model
        _lowercase						,							_lowercase			:				Any 						=						accelerator.gather((ddp_input, ddp_target)					)
        _lowercase						,							_lowercase			:				str 						=						input.to(accelerator.device					), target.to(accelerator.device					)
        # Perform our initial ground truth step in non "DDP"
        model.train()
        ddp_model.train()
        step_model(lowerCamelCase_				,					lowerCamelCase_				,					lowerCamelCase_				,					lowerCamelCase_				,					lowerCamelCase_					)
        opt.step()

        if ((iteration + 1) % 2 == 0) or ((iteration + 1) == len(lowerCamelCase_					)):
            if split_batches:
                sched.step()
            else:
                for _ in range(accelerator.num_processes					):
                    sched.step()
        opt.zero_grad()
        # Perform gradient accumulation under wrapper
        with accelerator.accumulate(lowerCamelCase_					):
            step_model(lowerCamelCase_				,					lowerCamelCase_				,					lowerCamelCase_				,					lowerCamelCase_					)
            ddp_opt.step()
            ddp_sched.step()
            ddp_opt.zero_grad()

        # Learning rates should be the same
        assert (
            opt.param_groups[0]["lr"] == ddp_opt.param_groups[0]["lr"]
        ), F'''Learning rates found in each optimizer did not align\nopt: {opt.param_groups[0]["lr"]}\nDDP opt: {ddp_opt.param_groups[0]["lr"]}\n'''
        _lowercase			:				List[Any] 						=						(((iteration + 1) % 2) == 0) or ((iteration + 1) == len(lowerCamelCase_					))
        if accelerator.num_processes > 1:
            check_model_parameters(lowerCamelCase_				,					lowerCamelCase_				,					lowerCamelCase_				,					lowerCamelCase_					)
        # Shuffle ddp_input on each iteration
        torch.manual_seed(1337 + iteration					)
    GradientState._reset_state()



def 							UpperCamelCase_( )       ->		Dict:
    _lowercase			:				Any 						=						Accelerator()

    _lowercase			:				Dict 						=						RegressionDataset(length=80					)
    _lowercase			:				Any 						=						DataLoader(lowerCamelCase_				,					batch_size=16					)
    _lowercase			:				Union[str, Any] 						=						RegressionDataset(length=96					)
    _lowercase			:				Optional[Any] 						=						DataLoader(lowerCamelCase_				,					batch_size=16					)
    _lowercase						,							_lowercase			:				Tuple 						=						accelerator.prepare(lowerCamelCase_				,					lowerCamelCase_					)
    assert accelerator.gradient_state.active_dataloader is None
    for iteration, _ in enumerate(lowerCamelCase_					):
        assert id(accelerator.gradient_state.active_dataloader					) == id(lowerCamelCase_					)
        if iteration < len(lowerCamelCase_					) - 1:
            assert not accelerator.gradient_state.end_of_dataloader
            if iteration == 1:
                for batch_num, _ in enumerate(lowerCamelCase_					):
                    assert id(accelerator.gradient_state.active_dataloader					) == id(lowerCamelCase_					)
                    if batch_num < len(lowerCamelCase_					) - 1:
                        assert not accelerator.gradient_state.end_of_dataloader
                    else:
                        assert accelerator.gradient_state.end_of_dataloader
        else:
            assert accelerator.gradient_state.end_of_dataloader
    assert accelerator.gradient_state.active_dataloader is None



def 							UpperCamelCase_( )       ->		Tuple:
    _lowercase			:				List[str] 						=						Accelerator()
    _lowercase			:				List[Any] 						=						accelerator.state
    if state.local_process_index == 0:
        print('**Test `accumulate` gradient accumulation with dataloader break**'					)
    test_dataloader_break()
    if state.distributed_type == DistributedType.NO:
        if state.local_process_index == 0:
            print('**Test NOOP `no_sync` context manager**'					)
        test_noop_sync(lowerCamelCase_					)
    if state.distributed_type in (DistributedType.MULTI_GPU, DistributedType.MULTI_CPU):
        if state.local_process_index == 0:
            print('**Test Distributed `no_sync` context manager**'					)
        test_distributed_sync(lowerCamelCase_					)
    if state.distributed_type == DistributedType.MULTI_GPU:
        for split_batch in [True, False]:
            for dispatch_batches in [True, False]:
                if state.local_process_index == 0:
                    print(
                        '**Test `accumulate` gradient accumulation, '				,					F'''`split_batches={split_batch}` and `dispatch_batches={dispatch_batches}`**'''				,					)
                test_gradient_accumulation(lowerCamelCase_				,					lowerCamelCase_					)

    # Currently will break on torch 2.0 +, need to investigate why
    if is_torch_version('<'				,					'2.0'					) or state.distributed_type == DistributedType.NO:
        if state.local_process_index == 0:
            print(
                '**Test `accumulate` gradient accumulation with optimizer and scheduler, '				,					'`split_batches=False`, `dispatch_batches=False`**'				,					)
        test_gradient_accumulation_with_opt_and_scheduler()
        if state.distributed_type == DistributedType.MULTI_GPU:
            for split_batch in [True, False]:
                for dispatch_batches in [True, False]:
                    if not split_batch and not dispatch_batches:
                        continue
                    if state.local_process_index == 0:
                        print(
                            '**Test `accumulate` gradient accumulation with optimizer and scheduler, '				,					F'''`split_batches={split_batch}` and `dispatch_batches={dispatch_batches}`**'''				,					)
                    test_gradient_accumulation_with_opt_and_scheduler(lowerCamelCase_				,					lowerCamelCase_					)



def 							UpperCamelCase_( lowerCamelCase_					)       ->		Any:
    # For xla_spawn (TPUs)
    main()


if __name__ == "__main__":
 main()






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  ):
       lowercase_					:      Optional[int]       					=     """hf-internal-testing/tiny-random-OnnxStableDiffusionPipeline"""
       def      UpperCamelCase	(   self, lowerCamelCase=0)			->  str:




           """simple docstring"""

           _lowercase			:				Optional[int] 						=						np.random.RandomState(lowerCamelCase)
           _lowercase			:				Union[str, Any] 						=						{
               '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      UpperCamelCase	(   self)			->  List[str]:




           """simple docstring"""

           _lowercase			:				int 						=						OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint, provider='CPUExecutionProvider')
           pipe.set_progress_bar_config(disable=lowerCamelCase)

           _lowercase			:				List[str] 						=						self.get_dummy_inputs()
           _lowercase			:				Tuple 						=						pipe(**lowerCamelCase).images
           _lowercase			:				Optional[Any] 						=						image[0, -3:, -3:, -1]

           assert image.shape == (1, 1_28, 1_28, 3)
           _lowercase			:				List[Any] 						=						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      UpperCamelCase	(   self)			->  List[Any]:




           """simple docstring"""

           _lowercase			:				str 						=						OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint, provider='CPUExecutionProvider')
           _lowercase			:				List[Any] 						=						PNDMScheduler.from_config(pipe.scheduler.config, skip_prk_steps=lowerCamelCase)
           pipe.set_progress_bar_config(disable=lowerCamelCase)

           _lowercase			:				Dict 						=						self.get_dummy_inputs()
           _lowercase			:				Optional[int] 						=						pipe(**lowerCamelCase).images
           _lowercase			:				Dict 						=						image[0, -3:, -3:, -1]

           assert image.shape == (1, 1_28, 1_28, 3)
           _lowercase			:				List[Any] 						=						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      UpperCamelCase	(   self)			->  str:




           """simple docstring"""

           _lowercase			:				Optional[int] 						=						OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint, provider='CPUExecutionProvider')
           _lowercase			:				Optional[int] 						=						LMSDiscreteScheduler.from_config(pipe.scheduler.config)
           pipe.set_progress_bar_config(disable=lowerCamelCase)

           _lowercase			:				Dict 						=						self.get_dummy_inputs()
           _lowercase			:				Union[str, Any] 						=						pipe(**lowerCamelCase).images
           _lowercase			:				str 						=						image[0, -3:, -3:, -1]

           assert image.shape == (1, 1_28, 1_28, 3)
           _lowercase			:				Union[str, Any] 						=						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      UpperCamelCase	(   self)			->  Optional[int]:




           """simple docstring"""

           _lowercase			:				Union[str, Any] 						=						OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint, provider='CPUExecutionProvider')
           _lowercase			:				Any 						=						EulerDiscreteScheduler.from_config(pipe.scheduler.config)
           pipe.set_progress_bar_config(disable=lowerCamelCase)

           _lowercase			:				List[Any] 						=						self.get_dummy_inputs()
           _lowercase			:				Any 						=						pipe(**lowerCamelCase).images
           _lowercase			:				Tuple 						=						image[0, -3:, -3:, -1]

           assert image.shape == (1, 1_28, 1_28, 3)
           _lowercase			:				Union[str, Any] 						=						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      UpperCamelCase	(   self)			->  Dict:




           """simple docstring"""

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

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

           assert image.shape == (1, 1_28, 1_28, 3)
           _lowercase			:				List[str] 						=						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      UpperCamelCase	(   self)			->  List[Any]:




           """simple docstring"""

           _lowercase			:				List[Any] 						=						OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint, provider='CPUExecutionProvider')
           _lowercase			:				Any 						=						DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
           pipe.set_progress_bar_config(disable=lowerCamelCase)

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

           assert image.shape == (1, 1_28, 1_28, 3)
           _lowercase			:				Any 						=						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      UpperCamelCase	(   self)			->  int:




           """simple docstring"""

           _lowercase			:				Dict 						=						OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint, provider='CPUExecutionProvider')
           pipe.set_progress_bar_config(disable=lowerCamelCase)

           _lowercase			:				str 						=						self.get_dummy_inputs()
           _lowercase			:				Any 						=						3 * [inputs['prompt']]

           # forward
           _lowercase			:				int 						=						pipe(**lowerCamelCase)
           _lowercase			:				Optional[int] 						=						output.images[0, -3:, -3:, -1]

           _lowercase			:				int 						=						self.get_dummy_inputs()
           _lowercase			:				Union[str, Any] 						=						3 * [inputs.pop('prompt')]

           _lowercase			:				Union[str, Any] 						=						pipe.tokenizer(
               lowerCamelCase, padding='max_length', max_length=pipe.tokenizer.model_max_length, truncation=lowerCamelCase, return_tensors='np', )
           _lowercase			:				Tuple 						=						text_inputs['input_ids']

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

           _lowercase			:				List[Any] 						=						prompt_embeds

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

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




       def      UpperCamelCase	(   self)			->  Optional[int]:




           """simple docstring"""

           _lowercase			:				Optional[Any] 						=						OnnxStableDiffusionPipeline.from_pretrained(self.hub_checkpoint, provider='CPUExecutionProvider')
           pipe.set_progress_bar_config(disable=lowerCamelCase)

           _lowercase			:				Optional[Any] 						=						self.get_dummy_inputs()
           _lowercase			:				Any 						=						3 * ['this is a negative prompt']
           _lowercase			:				str 						=						negative_prompt
           _lowercase			:				Optional[int] 						=						3 * [inputs['prompt']]

           # forward
           _lowercase			:				int 						=						pipe(**lowerCamelCase)
           _lowercase			:				str 						=						output.images[0, -3:, -3:, -1]

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

           _lowercase			:				Optional[int] 						=						[]
           for p in [prompt, negative_prompt]:
               _lowercase			:				Tuple 						=						pipe.tokenizer(
                   lowerCamelCase, padding='max_length', max_length=pipe.tokenizer.model_max_length, truncation=lowerCamelCase, return_tensors='np', )
               _lowercase			:				Dict 						=						text_inputs['input_ids']

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

           _lowercase						,							_lowercase			:				str 						=						embeds

           # forward
           _lowercase			:				Dict 						=						pipe(**lowerCamelCase)
           _lowercase			:				Tuple 						=						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      UpperCamelCase	(   self)			->  Union[str, Any]:




           """simple docstring"""

           return (
               "CUDAExecutionProvider",
               {
                   "gpu_mem_limit": "15000000000",  # 15GB
                   "arena_extend_strategy": "kSameAsRequested",
               },
           )
       @property
       def      UpperCamelCase	(   self)			->  str:




           """simple docstring"""

           _lowercase			:				int 						=						ort.SessionOptions()
           _lowercase			:				str 						=						False
           return options
       def      UpperCamelCase	(   self)			->  Optional[Any]:




           """simple docstring"""

           _lowercase			:				Optional[int] 						=						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			:				List[Any] 						=						'A painting of a squirrel eating a burger'
           np.random.seed(0)
           _lowercase			:				Union[str, Any] 						=						sd_pipe([prompt], guidance_scale=6.0, num_inference_steps=10, output_type='np')
           _lowercase			:				Optional[Any] 						=						output.images

           _lowercase			:				Optional[int] 						=						image[0, -3:, -3:, -1]

           assert image.shape == (1, 5_12, 5_12, 3)
           _lowercase			:				Union[str, Any] 						=						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      UpperCamelCase	(   self)			->  Optional[int]:




           """simple docstring"""

           _lowercase			:				str 						=						DDIMScheduler.from_pretrained(
               'runwayml/stable-diffusion-v1-5', subfolder='scheduler', revision='onnx')
           _lowercase			:				str 						=						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			:				List[Any] 						=						'open neural network exchange'
           _lowercase			:				List[Any] 						=						np.random.RandomState(0)
           _lowercase			:				Optional[Any] 						=						sd_pipe([prompt], guidance_scale=7.5, num_inference_steps=10, generator=lowerCamelCase, output_type='np')
           _lowercase			:				Optional[Any] 						=						output.images
           _lowercase			:				List[Any] 						=						image[0, -3:, -3:, -1]

           assert image.shape == (1, 5_12, 5_12, 3)
           _lowercase			:				Optional[int] 						=						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      UpperCamelCase	(   self)			->  List[str]:




           """simple docstring"""

           _lowercase			:				Dict 						=						LMSDiscreteScheduler.from_pretrained(
               'runwayml/stable-diffusion-v1-5', subfolder='scheduler', revision='onnx')
           _lowercase			:				Dict 						=						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			:				Tuple 						=						'open neural network exchange'
           _lowercase			:				str 						=						np.random.RandomState(0)
           _lowercase			:				Dict 						=						sd_pipe([prompt], guidance_scale=7.5, num_inference_steps=10, generator=lowerCamelCase, output_type='np')
           _lowercase			:				Optional[Any] 						=						output.images
           _lowercase			:				int 						=						image[0, -3:, -3:, -1]

           assert image.shape == (1, 5_12, 5_12, 3)
           _lowercase			:				Optional[Any] 						=						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      UpperCamelCase	(   self)			->  int:




           """simple docstring"""

           _lowercase			:				List[Any] 						=						0

           def test_callback_fn(lowerCamelCase, lowerCamelCase, lowerCamelCase) -> None:
               _lowercase			:				List[str] 						=						True
               nonlocal number_of_steps
               number_of_steps += 1
               if step == 0:
                   assert latents.shape == (1, 4, 64, 64)
                   _lowercase			:				Any 						=						latents[0, -3:, -3:, -1]
                   _lowercase			:				Tuple 						=						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			:				List[Any] 						=						latents[0, -3:, -3:, -1]
                   _lowercase			:				str 						=						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			:				Any 						=						False

           _lowercase			:				int 						=						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			:				Any 						=						'Andromeda galaxy in a bottle'

           _lowercase			:				str 						=						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      UpperCamelCase	(   self)			->  Tuple:




           """simple docstring"""

           _lowercase			:				List[str] 						=						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			:				Optional[int] 						=						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			:				Any 						=						OnnxStableDiffusionPipeline.from_pretrained(lowerCamelCase)

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


