




import collections.abc
from typing import Optional, Tuple, Union

import torch
import torch.utils.checkpoint
from torch import nn
from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss

from ...activations import ACTaFN
from ...modeling_outputs import BaseModelOutputWithNoAttention, ImageClassifierOutputWithNoAttention
from ...modeling_utils import PreTrainedModel
from ...utils import add_code_sample_docstrings, add_start_docstrings, add_start_docstrings_to_model_forward, logging
from .configuration_poolformer import PoolFormerConfig


UpperCAmelCase_						=logging.get_logger(__name__)

# General docstring
UpperCAmelCase_						="""PoolFormerConfig"""

# Base docstring
UpperCAmelCase_						="""sail/poolformer_s12"""
UpperCAmelCase_						=[1, 512, 7, 7]

# Image classification docstring
UpperCAmelCase_						="""sail/poolformer_s12"""
UpperCAmelCase_						="""tabby, tabby cat"""

UpperCAmelCase_						=[
    """sail/poolformer_s12""",
    # See all PoolFormer models at https://huggingface.co/models?filter=poolformer
]

def       UpperCAmelCase       (     _snake_case						, _snake_case = 0.0						, _snake_case = False							):
							if drop_prob == 0.0 or not training:
														return input
							lowerCAmelCase   =		1 - drop_prob
							lowerCAmelCase   =		(input.shape[0],) + (1,) * (input.ndim - 1)  # work with diff dim tensors, not just 2D ConvNets
							lowerCAmelCase   =		keep_prob + torch.rand(_snake_case						, dtype=input.dtype						, device=input.device							)
							random_tensor.floor_()  # binarize
							lowerCAmelCase   =		input.div(_snake_case							) * random_tensor
							return output


class 							__UpperCamelCase			( nn.Module  ):





				'''simple docstring'''



				def __init__( self       , UpperCAmelCase_ = None						):
											super().__init__()
											lowerCAmelCase   =		drop_prob



				def       __snake_case		( self       , UpperCAmelCase_						):
											return drop_path(UpperCAmelCase_       , self.drop_prob       , self.training						)






				def       __snake_case		( self						):
											return "p={}".format(self.drop_prob						)




class 							__UpperCamelCase			( nn.Module  ):





				'''simple docstring'''



				def __init__( self       , UpperCAmelCase_       , UpperCAmelCase_       , UpperCAmelCase_       , UpperCAmelCase_       , UpperCAmelCase_       , UpperCAmelCase_=None						):
											super().__init__()
											lowerCAmelCase   =		patch_size if isinstance(UpperCAmelCase_       , collections.abc.Iterable						) else (patch_size, patch_size)
											lowerCAmelCase   =		stride if isinstance(UpperCAmelCase_       , collections.abc.Iterable						) else (stride, stride)
											lowerCAmelCase   =		padding if isinstance(UpperCAmelCase_       , collections.abc.Iterable						) else (padding, padding)

											lowerCAmelCase   =		nn.Convad(UpperCAmelCase_       , UpperCAmelCase_       , kernel_size=UpperCAmelCase_       , stride=UpperCAmelCase_       , padding=UpperCAmelCase_						)
											lowerCAmelCase   =		norm_layer(UpperCAmelCase_						) if norm_layer else nn.Identity()






				def       __snake_case		( self       , UpperCAmelCase_						):
											lowerCAmelCase   =		self.projection(UpperCAmelCase_						)
											lowerCAmelCase   =		self.norm(UpperCAmelCase_						)
											return embeddings




class 							__UpperCamelCase			( nn.GroupNorm  ):





				'''simple docstring'''



				def __init__( self       , UpperCAmelCase_       , **UpperCAmelCase_						):
											super().__init__(1       , UpperCAmelCase_       , **UpperCAmelCase_						)




class 							__UpperCamelCase			( nn.Module  ):





				'''simple docstring'''



				def __init__( self       , UpperCAmelCase_						):
											super().__init__()
											lowerCAmelCase   =		nn.AvgPoolad(UpperCAmelCase_       , stride=1       , padding=pool_size // 2       , count_include_pad=UpperCAmelCase_						)






				def       __snake_case		( self       , UpperCAmelCase_						):
											return self.pool(UpperCAmelCase_						) - hidden_states




class 							__UpperCamelCase			( nn.Module  ):





				'''simple docstring'''



				def __init__( self       , UpperCAmelCase_       , UpperCAmelCase_       , UpperCAmelCase_       , UpperCAmelCase_						):
											super().__init__()
											lowerCAmelCase   =		nn.Convad(UpperCAmelCase_       , UpperCAmelCase_       , 1						)
											lowerCAmelCase   =		nn.Convad(UpperCAmelCase_       , UpperCAmelCase_       , 1						)
											lowerCAmelCase   =		PoolFormerDropPath(UpperCAmelCase_						)
											if isinstance(config.hidden_act       , UpperCAmelCase_						):
																		lowerCAmelCase   =		ACTaFN[config.hidden_act]
											else:
																		lowerCAmelCase   =		config.hidden_act






				def       __snake_case		( self       , UpperCAmelCase_						):
											lowerCAmelCase   =		self.conva(UpperCAmelCase_						)
											lowerCAmelCase   =		self.act_fn(UpperCAmelCase_						)
											lowerCAmelCase   =		self.drop(UpperCAmelCase_						)
											lowerCAmelCase   =		self.conva(UpperCAmelCase_						)
											lowerCAmelCase   =		self.drop(UpperCAmelCase_						)

											return hidden_states




class 							__UpperCamelCase			( nn.Module  ):





				'''simple docstring'''



				def __init__( self       , UpperCAmelCase_       , UpperCAmelCase_       , UpperCAmelCase_       , UpperCAmelCase_       , UpperCAmelCase_       , UpperCAmelCase_						):
											super().__init__()
											lowerCAmelCase   =		PoolFormerPooling(UpperCAmelCase_						)
											lowerCAmelCase   =		PoolFormerOutput(UpperCAmelCase_       , UpperCAmelCase_       , UpperCAmelCase_       , UpperCAmelCase_						)
											lowerCAmelCase   =		PoolFormerGroupNorm(UpperCAmelCase_						)
											lowerCAmelCase   =		PoolFormerGroupNorm(UpperCAmelCase_						)

											# Useful for training neural nets
											lowerCAmelCase   =		PoolFormerDropPath(UpperCAmelCase_						) if drop_path > 0.0 else nn.Identity()
											lowerCAmelCase   =		config.use_layer_scale
											if config.use_layer_scale:
																		lowerCAmelCase   =		nn.Parameter(
																		    config.layer_scale_init_value * torch.ones((UpperCAmelCase_)						)       , requires_grad=UpperCAmelCase_						)
																		lowerCAmelCase   =		nn.Parameter(
																		    config.layer_scale_init_value * torch.ones((UpperCAmelCase_)						)       , requires_grad=UpperCAmelCase_						)






				def       __snake_case		( self       , UpperCAmelCase_						):
											if self.use_layer_scale:
																		lowerCAmelCase   =		self.pooling(self.before_norm(UpperCAmelCase_						)						)
																		lowerCAmelCase   =		self.layer_scale_a.unsqueeze(-1						).unsqueeze(-1						) * pooling_output
																		# First residual connection
																		lowerCAmelCase   =		hidden_states + self.drop_path(UpperCAmelCase_						)
																		lowerCAmelCase   =		()

																		lowerCAmelCase   =		self.output(self.after_norm(UpperCAmelCase_						)						)
																		lowerCAmelCase   =		self.layer_scale_a.unsqueeze(-1						).unsqueeze(-1						) * layer_output
																		# Second residual connection
																		lowerCAmelCase   =		hidden_states + self.drop_path(UpperCAmelCase_						)

																		lowerCAmelCase   =		(output,) + outputs
																		return outputs

											else:
																		lowerCAmelCase   =		self.drop_path(self.pooling(self.before_norm(UpperCAmelCase_						)						)						)
																		# First residual connection
																		lowerCAmelCase   =		pooling_output + hidden_states
																		lowerCAmelCase   =		()

																		# Second residual connection inside the PoolFormerOutput block
																		lowerCAmelCase   =		self.drop_path(self.output(self.after_norm(UpperCAmelCase_						)						)						)
																		lowerCAmelCase   =		hidden_states + layer_output

																		lowerCAmelCase   =		(output,) + outputs
																		return outputs




class 							__UpperCamelCase			( nn.Module  ):





				'''simple docstring'''



				def __init__( self       , UpperCAmelCase_						):
											super().__init__()
											lowerCAmelCase   =		config
											# stochastic depth decay rule
											lowerCAmelCase   =		[x.item() for x in torch.linspace(0       , config.drop_path_rate       , sum(config.depths						)						)]

											# patch embeddings
											lowerCAmelCase   =		[]
											for i in range(config.num_encoder_blocks						):
																		embeddings.append(
																		    PoolFormerEmbeddings(
																		        patch_size=config.patch_sizes[i]       , stride=config.strides[i]       , padding=config.padding[i]       , num_channels=config.num_channels if i == 0 else config.hidden_sizes[i - 1]       , hidden_size=config.hidden_sizes[i]       , )						)
											lowerCAmelCase   =		nn.ModuleList(UpperCAmelCase_						)

											# Transformer blocks
											lowerCAmelCase   =		[]
											lowerCAmelCase   =		0
											for i in range(config.num_encoder_blocks						):
																		# each block consists of layers
																		lowerCAmelCase   =		[]
																		if i != 0:
																									cur += config.depths[i - 1]
																		for j in range(config.depths[i]						):
																									layers.append(
																									    PoolFormerLayer(
																									        UpperCAmelCase_       , num_channels=config.hidden_sizes[i]       , pool_size=config.pool_size       , hidden_size=config.hidden_sizes[i]       , intermediate_size=int(config.hidden_sizes[i] * config.mlp_ratio						)       , drop_path=dpr[cur + j]       , )						)
																		blocks.append(nn.ModuleList(UpperCAmelCase_						)						)

											lowerCAmelCase   =		nn.ModuleList(UpperCAmelCase_						)






				def       __snake_case		( self       , UpperCAmelCase_       , UpperCAmelCase_=False       , UpperCAmelCase_=True						):
											lowerCAmelCase   =		() if output_hidden_states else None

											lowerCAmelCase   =		pixel_values
											for idx, layers in enumerate(zip(self.patch_embeddings       , self.block						)						):
																		lowerCAmelCase   ,     lowerCAmelCase   =		layers
																		# Get patch embeddings from hidden_states
																		lowerCAmelCase   =		embedding_layer(UpperCAmelCase_						)
																		# Send the embeddings through the blocks
																		for _, blk in enumerate(UpperCAmelCase_						):
																									lowerCAmelCase   =		blk(UpperCAmelCase_						)
																									lowerCAmelCase   =		layer_outputs[0]

																		if output_hidden_states:
																									lowerCAmelCase   =		all_hidden_states + (hidden_states,)

											if not return_dict:
																		return tuple(v for v in [hidden_states, all_hidden_states] if v is not None						)

											return BaseModelOutputWithNoAttention(last_hidden_state=UpperCAmelCase_       , hidden_states=UpperCAmelCase_						)




class 							__UpperCamelCase			( __UpperCAmelCase  ):





				'''simple docstring'''



				__a							: Tuple										=PoolFormerConfig
				__a							: List[Any]										="""poolformer"""
				__a							: Dict										="""pixel_values"""
				__a							: int										=True



				def       __snake_case		( self       , UpperCAmelCase_						):
											if isinstance(UpperCAmelCase_       , (nn.Linear, nn.Convad)						):
																		module.weight.data.normal_(mean=0.0       , std=self.config.initializer_range						)
																		if module.bias is not None:
																									module.bias.data.zero_()
											elif isinstance(UpperCAmelCase_       , nn.LayerNorm						):
																		module.bias.data.zero_()
																		module.weight.data.fill_(1.0						)






				def       __snake_case		( self       , UpperCAmelCase_       , UpperCAmelCase_=False						):
											if isinstance(UpperCAmelCase_       , UpperCAmelCase_						):
																		lowerCAmelCase   =		value




UpperCAmelCase_						=R"""
    This model is a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) sub-class. Use
    it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage and
    behavior.

    Parameters:
        config ([`PoolFormerConfig`]): Model configuration class with all the parameters of the model.
            Initializing with a config file does not load the weights associated with the model, only the
            configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.
"""

UpperCAmelCase_						=R"""
    Args:
        pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):
            Pixel values. Pixel values can be obtained using [`AutoImageProcessor`]. See
            [`PoolFormerImageProcessor.__call__`] for details.
"""


@add_start_docstrings(
    """The bare PoolFormer Model transformer outputting raw hidden-states without any specific head on top."""     ,					__UpperCAmelCase     ,					)
class 							__UpperCamelCase			( __UpperCAmelCase  ):





				'''simple docstring'''



				def __init__( self       , UpperCAmelCase_						):
											super().__init__(UpperCAmelCase_						)
											lowerCAmelCase   =		config

											lowerCAmelCase   =		PoolFormerEncoder(UpperCAmelCase_						)

											# Initialize weights and apply final processing
											self.post_init()



				def       __snake_case		( self						):
											return self.embeddings.patch_embeddings






				@add_start_docstrings_to_model_forward(UpperCAmelCase_						)
				@add_code_sample_docstrings(
				    checkpoint=_CHECKPOINT_FOR_DOC       , output_type=UpperCAmelCase_       , config_class=_CONFIG_FOR_DOC       , modality='''vision'''       , expected_output=_EXPECTED_OUTPUT_SHAPE       , )
				def       __snake_case		( self       , UpperCAmelCase_ = None       , UpperCAmelCase_ = None       , UpperCAmelCase_ = None       , ):
											lowerCAmelCase   =		(
											    output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
											)
											lowerCAmelCase   =		return_dict if return_dict is not None else self.config.use_return_dict

											if pixel_values is None:
																		raise ValueError('''You have to specify pixel_values'''						)

											lowerCAmelCase   =		self.encoder(
											    UpperCAmelCase_       , output_hidden_states=UpperCAmelCase_       , return_dict=UpperCAmelCase_       , )
											lowerCAmelCase   =		encoder_outputs[0]

											if not return_dict:
																		return (sequence_output, None) + encoder_outputs[1:]

											return BaseModelOutputWithNoAttention(
											    last_hidden_state=UpperCAmelCase_       , hidden_states=encoder_outputs.hidden_states       , )




class 							__UpperCamelCase			( nn.Module  ):





				'''simple docstring'''



				def __init__( self       , UpperCAmelCase_						):
											super().__init__()
											lowerCAmelCase   =		nn.Linear(config.hidden_size       , config.hidden_size						)






				def       __snake_case		( self       , UpperCAmelCase_						):
											lowerCAmelCase   =		self.dense(UpperCAmelCase_						)
											return output




@add_start_docstrings(
    """
    PoolFormer Model transformer with an image classification head on top
    """     ,					__UpperCAmelCase     ,					)
class 							__UpperCamelCase			( __UpperCAmelCase  ):





				'''simple docstring'''



				def __init__( self       , UpperCAmelCase_						):
											super().__init__(UpperCAmelCase_						)
											lowerCAmelCase   =		config.num_labels
											lowerCAmelCase   =		PoolFormerModel(UpperCAmelCase_						)

											# Final norm
											lowerCAmelCase   =		PoolFormerGroupNorm(config.hidden_sizes[-1]						)
											# Classifier head
											lowerCAmelCase   =		(
											    nn.Linear(config.hidden_sizes[-1]       , config.num_labels						) if config.num_labels > 0 else nn.Identity()
											)

											# Initialize weights and apply final processing
											self.post_init()






				@add_start_docstrings_to_model_forward(UpperCAmelCase_						)
				@add_code_sample_docstrings(
				    checkpoint=_IMAGE_CLASS_CHECKPOINT       , output_type=UpperCAmelCase_       , config_class=_CONFIG_FOR_DOC       , expected_output=_IMAGE_CLASS_EXPECTED_OUTPUT       , )
				def       __snake_case		( self       , UpperCAmelCase_ = None       , UpperCAmelCase_ = None       , UpperCAmelCase_ = None       , UpperCAmelCase_ = None       , ):
											lowerCAmelCase   =		return_dict if return_dict is not None else self.config.use_return_dict

											lowerCAmelCase   =		self.poolformer(
											    UpperCAmelCase_       , output_hidden_states=UpperCAmelCase_       , return_dict=UpperCAmelCase_       , )

											lowerCAmelCase   =		outputs[0]

											lowerCAmelCase   =		self.classifier(self.norm(UpperCAmelCase_						).mean([-2, -1]						)						)

											lowerCAmelCase   =		None
											if labels is not None:
																		if self.config.problem_type is None:
																									if self.num_labels == 1:
																																lowerCAmelCase   =		'''regression'''
																									elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
																																lowerCAmelCase   =		'''single_label_classification'''
																									else:
																																lowerCAmelCase   =		'''multi_label_classification'''

																		if self.config.problem_type == "regression":
																									lowerCAmelCase   =		MSELoss()
																									if self.num_labels == 1:
																																lowerCAmelCase   =		loss_fct(logits.squeeze()       , labels.squeeze()						)
																									else:
																																lowerCAmelCase   =		loss_fct(UpperCAmelCase_       , UpperCAmelCase_						)
																		elif self.config.problem_type == "single_label_classification":
																									lowerCAmelCase   =		CrossEntropyLoss()
																									lowerCAmelCase   =		loss_fct(logits.view(-1       , self.num_labels						)       , labels.view(-1						)						)
																		elif self.config.problem_type == "multi_label_classification":
																									lowerCAmelCase   =		BCEWithLogitsLoss()
																									lowerCAmelCase   =		loss_fct(UpperCAmelCase_       , UpperCAmelCase_						)

											if not return_dict:
																		lowerCAmelCase   =		(logits,) + outputs[2:]
																		return ((loss,) + output) if loss is not None else output

											return ImageClassifierOutputWithNoAttention(loss=UpperCAmelCase_       , logits=UpperCAmelCase_       , hidden_states=outputs.hidden_states						)









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 							__UpperCamelCase			( __UpperCAmelCase     ,					unittest.TestCase  ):





				'''simple docstring'''



				__a							: Any										="""hf-internal-testing/tiny-random-OnnxStableDiffusionPipeline"""



				def       __snake_case		( self       , UpperCAmelCase_=0						):
											lowerCAmelCase   =		floats_tensor((1, 3, 1_28, 1_28)       , rng=random.Random(UpperCAmelCase_						)						)
											lowerCAmelCase   =		np.random.RandomState(UpperCAmelCase_						)
											lowerCAmelCase   =		{
											    '''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						):
											lowerCAmelCase   =		OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint       , provider='''CPUExecutionProvider'''						)
											pipe.set_progress_bar_config(disable=UpperCAmelCase_						)

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

											assert image.shape == (1, 1_28, 1_28, 3)
											lowerCAmelCase   =		np.array([0.69643, 0.58484, 0.50314, 0.58760, 0.55368, 0.59643, 0.51529, 0.41217, 0.49087]						)
											assert np.abs(image_slice - expected_slice						).max() < 1E-1



				def       __snake_case		( self						):
											lowerCAmelCase   =		OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint       , provider='''CPUExecutionProvider'''						)
											lowerCAmelCase   =		PNDMScheduler.from_config(pipe.scheduler.config       , skip_prk_steps=UpperCAmelCase_						)
											pipe.set_progress_bar_config(disable=UpperCAmelCase_						)

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

											assert image.shape == (1, 1_28, 1_28, 3)
											lowerCAmelCase   =		np.array([0.61737, 0.54642, 0.53183, 0.54465, 0.52742, 0.60525, 0.49969, 0.40655, 0.48154]						)

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



				def       __snake_case		( self						):
											lowerCAmelCase   =		OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint       , provider='''CPUExecutionProvider'''						)
											lowerCAmelCase   =		LMSDiscreteScheduler.from_config(pipe.scheduler.config						)
											pipe.set_progress_bar_config(disable=UpperCAmelCase_						)

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

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

											assert image.shape == (1, 1_28, 1_28, 3)
											lowerCAmelCase   =		np.array([0.52761, 0.59977, 0.49033, 0.49619, 0.54282, 0.50311, 0.47600, 0.40918, 0.45203]						)

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



				def       __snake_case		( self						):
											lowerCAmelCase   =		OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint       , provider='''CPUExecutionProvider'''						)
											lowerCAmelCase   =		EulerDiscreteScheduler.from_config(pipe.scheduler.config						)
											pipe.set_progress_bar_config(disable=UpperCAmelCase_						)

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

											assert image.shape == (1, 1_28, 1_28, 3)
											lowerCAmelCase   =		np.array([0.52911, 0.60004, 0.49229, 0.49805, 0.54502, 0.50680, 0.47777, 0.41028, 0.45304]						)

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



				def       __snake_case		( self						):
											lowerCAmelCase   =		OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint       , provider='''CPUExecutionProvider'''						)
											lowerCAmelCase   =		EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config						)
											pipe.set_progress_bar_config(disable=UpperCAmelCase_						)

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

											assert image.shape == (1, 1_28, 1_28, 3)
											lowerCAmelCase   =		np.array([0.52911, 0.60004, 0.49229, 0.49805, 0.54502, 0.50680, 0.47777, 0.41028, 0.45304]						)

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






				def       __snake_case		( self						):
											lowerCAmelCase   =		OnnxStableDiffusionImgaImgPipeline.from_pretrained(self.hub_checkpoint       , provider='''CPUExecutionProvider'''						)
											lowerCAmelCase   =		DPMSolverMultistepScheduler.from_config(pipe.scheduler.config						)
											pipe.set_progress_bar_config(disable=UpperCAmelCase_						)

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

											assert image.shape == (1, 1_28, 1_28, 3)
											lowerCAmelCase   =		np.array([0.65331, 0.58277, 0.48204, 0.56059, 0.53665, 0.56235, 0.50969, 0.40009, 0.46552]						)

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




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





				'''simple docstring'''



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



				@property
				def       __snake_case		( self						):
											lowerCAmelCase   =		ort.SessionOptions()
											lowerCAmelCase   =		False
											return options



				def       __snake_case		( self						):
											lowerCAmelCase   =		load_image(
											    '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
											    '''/img2img/sketch-mountains-input.jpg'''						)
											lowerCAmelCase   =		init_image.resize((7_68, 5_12)						)
											# using the PNDM scheduler by default
											lowerCAmelCase   =		OnnxStableDiffusionImgaImgPipeline.from_pretrained(
											    '''CompVis/stable-diffusion-v1-4'''       , revision='''onnx'''       , safety_checker=UpperCAmelCase_       , feature_extractor=UpperCAmelCase_       , provider=self.gpu_provider       , sess_options=self.gpu_options       , )
											pipe.set_progress_bar_config(disable=UpperCAmelCase_						)

											lowerCAmelCase   =		'''A fantasy landscape, trending on artstation'''

											lowerCAmelCase   =		np.random.RandomState(0						)
											lowerCAmelCase   =		pipe(
											    prompt=UpperCAmelCase_       , image=UpperCAmelCase_       , strength=0.75       , guidance_scale=7.5       , num_inference_steps=10       , generator=UpperCAmelCase_       , output_type='''np'''       , )
											lowerCAmelCase   =		output.images
											lowerCAmelCase   =		images[0, 2_55:2_58, 3_83:3_86, -1]

											assert images.shape == (1, 5_12, 7_68, 3)
											lowerCAmelCase   =		np.array([0.4909, 0.5059, 0.5372, 0.4623, 0.4876, 0.5049, 0.4820, 0.4956, 0.5019]						)
											# 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						):
											lowerCAmelCase   =		load_image(
											    '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
											    '''/img2img/sketch-mountains-input.jpg'''						)
											lowerCAmelCase   =		init_image.resize((7_68, 5_12)						)
											lowerCAmelCase   =		LMSDiscreteScheduler.from_pretrained(
											    '''runwayml/stable-diffusion-v1-5'''       , subfolder='''scheduler'''       , revision='''onnx'''						)
											lowerCAmelCase   =		OnnxStableDiffusionImgaImgPipeline.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       , )
											pipe.set_progress_bar_config(disable=UpperCAmelCase_						)

											lowerCAmelCase   =		'''A fantasy landscape, trending on artstation'''

											lowerCAmelCase   =		np.random.RandomState(0						)
											lowerCAmelCase   =		pipe(
											    prompt=UpperCAmelCase_       , image=UpperCAmelCase_       , strength=0.75       , guidance_scale=7.5       , num_inference_steps=20       , generator=UpperCAmelCase_       , output_type='''np'''       , )
											lowerCAmelCase   =		output.images
											lowerCAmelCase   =		images[0, 2_55:2_58, 3_83:3_86, -1]

											assert images.shape == (1, 5_12, 7_68, 3)
											lowerCAmelCase   =		np.array([0.8043, 0.926, 0.9581, 0.8119, 0.8954, 0.913, 0.7209, 0.7463, 0.7431]						)
											# TODO: lower the tolerance after finding the cause of onnxruntime reproducibility issues

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



