






import random
import unittest

import numpy as np
import torch

from diffusers import (
    DPMSolverMultistepScheduler,
    EulerAncestralDiscreteScheduler,
    EulerDiscreteScheduler,
    LMSDiscreteScheduler,
    OnnxStableDiffusionUpscalePipeline,
    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        lowercase__      (							snake_case__,       unittest.TestCase       ):
			# TODO: is there an appropriate internal test set?
			_UpperCAmelCase		:List[Any]             =			'ssube/stable-diffusion-x4-upscaler-onnx'


			def      UpperCAmelCase__						(				self					:  List[str]       ,	snake_case__					:  Any=0 ):
										lowerCamelCase_	:							Dict									=floats_tensor((1, 3, 128, 128)       ,	rng=random.Random(__lowerCAmelCase ) )
										lowerCamelCase_	:							Union[str, Any]									=torch.manual_seed(__lowerCAmelCase )
										lowerCamelCase_	:							Union[str, Any]									={
										    "prompt": "A painting of a squirrel eating a burger",
										    "image": image,
										    "generator": generator,
										    "num_inference_steps": 3,
										    "guidance_scale": 7.5,
										    "output_type": "numpy",
										}
										return inputs


			def      UpperCAmelCase__						(				self					:  Tuple ):
										lowerCamelCase_	:							List[Any]									=OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint       ,	provider="CPUExecutionProvider" )
										pipe.set_progress_bar_config(disable=__lowerCAmelCase )

										lowerCamelCase_	:							Optional[Any]									=self.get_dummy_inputs()
										lowerCamelCase_	:							Optional[Any]									=pipe(**__lowerCAmelCase ).images
										lowerCamelCase_	:							List[Any]									=image[0, -3:, -3:, -1].flatten()

										# started as 128, should now be 512
										assert image.shape == (1, 512, 512, 3)
										lowerCamelCase_	:							List[str]									=np.array(
										    [0.6_974_782, 0.68_902_093, 0.70_135_885, 0.7_583_618, 0.7_804_545, 0.7_854_912, 0.78_667_426, 0.78_743_863, 0.78_070_223] )
										assert np.abs(image_slice - expected_slice ).max() < 1E-1


			def      UpperCAmelCase__						(				self					:  str ):
										lowerCamelCase_	:							str									=OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint       ,	provider="CPUExecutionProvider" )
										lowerCamelCase_	:							Union[str, Any]									=PNDMScheduler.from_config(pipe.scheduler.config       ,	skip_prk_steps=__lowerCAmelCase )
										pipe.set_progress_bar_config(disable=__lowerCAmelCase )

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

										assert image.shape == (1, 512, 512, 3)
										lowerCamelCase_	:							str									=np.array(
										    [0.6_898_892, 0.59_240_556, 0.52_499_527, 0.58_866_215, 0.52_258_235, 0.52_572_715, 0.62_414_473, 0.6_174_387, 0.6_214_964] )
										assert np.abs(image_slice.flatten() - expected_slice ).max() < 1E-1


			def      UpperCAmelCase__						(				self					:  Dict ):
										lowerCamelCase_	:							List[str]									=OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint       ,	provider="CPUExecutionProvider" )
										lowerCamelCase_	:							Dict									=DPMSolverMultistepScheduler.from_config(pipe.scheduler.config )
										pipe.set_progress_bar_config(disable=__lowerCAmelCase )

										lowerCamelCase_	:							Dict									=self.get_dummy_inputs()
										lowerCamelCase_	:							List[Any]									=pipe(**__lowerCAmelCase ).images
										lowerCamelCase_	:							Dict									=image[0, -3:, -3:, -1]

										assert image.shape == (1, 512, 512, 3)
										lowerCamelCase_	:							Dict									=np.array(
										    [0.7_659_278, 0.76_437_664, 0.75_579_107, 0.7_691_116, 0.77_666_986, 0.7_727_672, 0.7_758_664, 0.7_812_226, 0.76_942_515] )

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


			def      UpperCAmelCase__						(				self					:  Dict ):
										lowerCamelCase_	:							List[Any]									=OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint       ,	provider="CPUExecutionProvider" )
										lowerCamelCase_	:							Tuple									=EulerDiscreteScheduler.from_config(pipe.scheduler.config )
										pipe.set_progress_bar_config(disable=__lowerCAmelCase )

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

										assert image.shape == (1, 512, 512, 3)
										lowerCamelCase_	:							int									=np.array(
										    [0.6_974_782, 0.68_902_093, 0.70_135_885, 0.7_583_618, 0.7_804_545, 0.7_854_912, 0.78_667_426, 0.78_743_863, 0.78_070_223] )
										assert np.abs(image_slice.flatten() - expected_slice ).max() < 1E-1


			def      UpperCAmelCase__						(				self					:  Dict ):
										lowerCamelCase_	:							List[str]									=OnnxStableDiffusionUpscalePipeline.from_pretrained(self.hub_checkpoint       ,	provider="CPUExecutionProvider" )
										lowerCamelCase_	:							int									=EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config )
										pipe.set_progress_bar_config(disable=__lowerCAmelCase )

										lowerCamelCase_	:							str									=self.get_dummy_inputs()
										lowerCamelCase_	:							str									=pipe(**__lowerCAmelCase ).images
										lowerCamelCase_	:							Optional[Any]									=image[0, -3:, -3:, -1]

										assert image.shape == (1, 512, 512, 3)
										lowerCamelCase_	:							Union[str, Any]									=np.array(
										    [0.77_424_496, 0.773_601, 0.7_645_288, 0.7_769_598, 0.7_772_739, 0.7_738_688, 0.78_187_233, 0.77_879_584, 0.767_043] )

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







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


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


			@property
			def      UpperCAmelCase__						(				self					:  List[str] ):
										lowerCamelCase_	:							List[Any]									=ort.SessionOptions()
										lowerCamelCase_	:							Any									=False
										return options


			def      UpperCAmelCase__						(				self					:  Optional[int] ):
										lowerCamelCase_	:							Union[str, Any]									=load_image(
										    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
										    "/img2img/sketch-mountains-input.jpg" )
										lowerCamelCase_	:							Optional[Any]									=init_image.resize((128, 128) )
										# using the PNDM scheduler by default
										lowerCamelCase_	:							List[str]									=OnnxStableDiffusionUpscalePipeline.from_pretrained(
										    "ssube/stable-diffusion-x4-upscaler-onnx"       ,	provider=self.gpu_provider       ,	sess_options=self.gpu_options       ,	)
										pipe.set_progress_bar_config(disable=__lowerCAmelCase )

										lowerCamelCase_	:							str									="A fantasy landscape, trending on artstation"

										lowerCamelCase_	:							Dict									=torch.manual_seed(0 )
										lowerCamelCase_	:							int									=pipe(
										    prompt=__lowerCAmelCase       ,	image=__lowerCAmelCase       ,	guidance_scale=7.5       ,	num_inference_steps=10       ,	generator=__lowerCAmelCase       ,	output_type="np"       ,	)
										lowerCamelCase_	:							str									=output.images
										lowerCamelCase_	:							Any									=images[0, 255:258, 383:386, -1]

										assert images.shape == (1, 512, 512, 3)
										lowerCamelCase_	:							List[Any]									=np.array([0.4_883, 0.4_947, 0.4_980, 0.4_975, 0.4_982, 0.4_980, 0.5_000, 0.5_006, 0.4_972] )
										# TODO: lower the tolerance after finding the cause of onnxruntime reproducibility issues

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


			def      UpperCAmelCase__						(				self					:  str ):
										lowerCamelCase_	:							str									=load_image(
										    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
										    "/img2img/sketch-mountains-input.jpg" )
										lowerCamelCase_	:							Tuple									=init_image.resize((128, 128) )
										lowerCamelCase_	:							List[Any]									=LMSDiscreteScheduler.from_pretrained(
										    "ssube/stable-diffusion-x4-upscaler-onnx"       ,	subfolder="scheduler" )
										lowerCamelCase_	:							int									=OnnxStableDiffusionUpscalePipeline.from_pretrained(
										    "ssube/stable-diffusion-x4-upscaler-onnx"       ,	scheduler=__lowerCAmelCase       ,	provider=self.gpu_provider       ,	sess_options=self.gpu_options       ,	)
										pipe.set_progress_bar_config(disable=__lowerCAmelCase )

										lowerCamelCase_	:							Optional[Any]									="A fantasy landscape, trending on artstation"

										lowerCamelCase_	:							str									=torch.manual_seed(0 )
										lowerCamelCase_	:							List[str]									=pipe(
										    prompt=__lowerCAmelCase       ,	image=__lowerCAmelCase       ,	guidance_scale=7.5       ,	num_inference_steps=20       ,	generator=__lowerCAmelCase       ,	output_type="np"       ,	)
										lowerCamelCase_	:							List[Any]									=output.images
										lowerCamelCase_	:							Optional[Any]									=images[0, 255:258, 383:386, -1]

										assert images.shape == (1, 512, 512, 3)
										lowerCamelCase_	:							Optional[int]									=np.array(
										    [0.50_173_753, 0.50_223_356, 0.502_039, 0.50_233_036, 0.5_023_725, 0.5_022_601, 0.5_018_758, 0.50_234_085, 0.50_241_566] )
										# TODO: lower the tolerance after finding the cause of onnxruntime reproducibility issues

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











"""simple docstring"""



import inspect
import unittest

from transformers import DecisionTransformerConfig, is_torch_available
from transformers.testing_utils import require_torch, slow, torch_device

from ...generation.test_utils import GenerationTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
						import torch

						from transformers import DecisionTransformerModel
						from transformers.models.decision_transformer.modeling_decision_transformer import (
						    DECISION_TRANSFORMER_PRETRAINED_MODEL_ARCHIVE_LIST,
						)







class        lowercase__      :


			def __init__(				self					:  List[Any]       ,	snake_case__					:  Optional[Any]       ,	snake_case__					:  Tuple=13       ,	snake_case__					:  str=7       ,	snake_case__					:  Union[str, Any]=6       ,	snake_case__					:  str=17       ,	snake_case__					:  Any=23       ,	snake_case__					:  int=11       ,	snake_case__					:  Tuple=True       ,	):
										lowerCamelCase_	:							str									=parent
										lowerCamelCase_	:							Union[str, Any]									=batch_size
										lowerCamelCase_	:							List[Any]									=seq_length
										lowerCamelCase_	:							Union[str, Any]									=act_dim
										lowerCamelCase_	:							Optional[Any]									=state_dim
										lowerCamelCase_	:							Optional[Any]									=hidden_size
										lowerCamelCase_	:							Tuple									=max_length
										lowerCamelCase_	:							List[Any]									=is_training


			def      UpperCAmelCase__						(				self					:  Dict ):
										lowerCamelCase_	:							Optional[Any]									=floats_tensor((self.batch_size, self.seq_length, self.state_dim) )
										lowerCamelCase_	:							Optional[Any]									=floats_tensor((self.batch_size, self.seq_length, self.act_dim) )
										lowerCamelCase_	:							List[Any]									=floats_tensor((self.batch_size, self.seq_length, 1) )
										lowerCamelCase_	:							Optional[Any]									=floats_tensor((self.batch_size, self.seq_length, 1) )
										lowerCamelCase_	:							List[Any]									=ids_tensor((self.batch_size, self.seq_length)       ,	vocab_size=1000 )
										lowerCamelCase_	:							Optional[int]									=random_attention_mask((self.batch_size, self.seq_length) )

										lowerCamelCase_	:							List[str]									=self.get_config()

										return (
										    config,
										    states,
										    actions,
										    rewards,
										    returns_to_go,
										    timesteps,
										    attention_mask,
										)


			def      UpperCAmelCase__						(				self					:  Any ):
										return DecisionTransformerConfig(
										    batch_size=self.batch_size       ,	seq_length=self.seq_length       ,	act_dim=self.act_dim       ,	state_dim=self.state_dim       ,	hidden_size=self.hidden_size       ,	max_length=self.max_length       ,	)


			def      UpperCAmelCase__						(				self					:  Optional[Any]       ,	snake_case__					:  Tuple       ,	snake_case__					:  Dict       ,	snake_case__					:  int       ,	snake_case__					:  Dict       ,	snake_case__					:  Optional[int]       ,	snake_case__					:  List[str]       ,	snake_case__					:  List[str]       ,	):
										lowerCamelCase_	:							Tuple									=DecisionTransformerModel(config=snake_case__ )
										model.to(snake_case__ )
										model.eval()
										lowerCamelCase_	:							str									=model(snake_case__       ,	snake_case__       ,	snake_case__       ,	snake_case__       ,	snake_case__       ,	snake_case__ )

										self.parent.assertEqual(result.state_preds.shape       ,	states.shape )
										self.parent.assertEqual(result.action_preds.shape       ,	actions.shape )
										self.parent.assertEqual(result.return_preds.shape       ,	returns_to_go.shape )
										self.parent.assertEqual(
										    result.last_hidden_state.shape       ,	(self.batch_size, self.seq_length * 3, self.hidden_size) )  # seq length *3 as there are 3 modelities: states, returns and actions


			def      UpperCAmelCase__						(				self					:  List[str] ):
										lowerCamelCase_	:							List[str]									=self.prepare_config_and_inputs()
										(
										    (
										    lowerCamelCase_
										)							,						(
										    lowerCamelCase_
										)							,						(
										    lowerCamelCase_
										)							,						(
										    lowerCamelCase_
										)							,						(
										    lowerCamelCase_
										)							,						(
										    lowerCamelCase_
										)							,						(
										    lowerCamelCase_
										)							,						
										)	:							Optional[int]									=config_and_inputs
										lowerCamelCase_	:							Optional[int]									={
										    "states": states,
										    "actions": actions,
										    "rewards": rewards,
										    "returns_to_go": returns_to_go,
										    "timesteps": timesteps,
										    "attention_mask": attention_mask,
										}
										return config, inputs_dict







@require_torch
class        lowercase__      (							snake_case__,       snake_case__,       snake_case__,       unittest.TestCase       ):
			_UpperCAmelCase		:Optional[Any]             =			(DecisionTransformerModel,) if is_torch_available() else ()
			_UpperCAmelCase		:int             =			()
			_UpperCAmelCase		:int             =			{"feature-extraction": DecisionTransformerModel} if is_torch_available() else {}

			# Ignoring of a failing test from GenerationTesterMixin, as the model does not use inputs_ids
			_UpperCAmelCase		:Union[str, Any]             =			False

			# Ignoring of a failing tests from ModelTesterMixin, as the model does not implement these features
			_UpperCAmelCase		:Optional[Any]             =			False
			_UpperCAmelCase		:Tuple             =			False
			_UpperCAmelCase		:Tuple             =			False
			_UpperCAmelCase		:List[Any]             =			False
			_UpperCAmelCase		:Dict             =			False
			_UpperCAmelCase		:Any             =			False
			_UpperCAmelCase		:List[Any]             =			False
			_UpperCAmelCase		:int             =			False
			_UpperCAmelCase		:str             =			False


			def      UpperCAmelCase__						(				self					:  str ):
										lowerCamelCase_	:							Dict									=DecisionTransformerModelTester(self )
										lowerCamelCase_	:							str									=ConfigTester(self       ,	config_class=snake_case__       ,	hidden_size=37 )


			def      UpperCAmelCase__						(				self					:  Union[str, Any] ):
										self.config_tester.run_common_tests()


			def      UpperCAmelCase__						(				self					:  int ):
										lowerCamelCase_	:							Tuple									=self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_model(*snake_case__ )


			@slow
			def      UpperCAmelCase__						(				self					:  List[str] ):
										for model_name in DECISION_TRANSFORMER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																	lowerCamelCase_	:							str									=DecisionTransformerModel.from_pretrained(snake_case__ )
																	self.assertIsNotNone(snake_case__ )


			def      UpperCAmelCase__						(				self					:  str ):
										lowerCamelCase_							,						lowerCamelCase_	:							Dict									=self.model_tester.prepare_config_and_inputs_for_common()

										for model_class in self.all_model_classes:
																	lowerCamelCase_	:							List[Any]									=model_class(snake_case__ )
																	lowerCamelCase_	:							int									=inspect.signature(model.forward )
																	# signature.parameters is an OrderedDict => so arg_names order is deterministic
																	lowerCamelCase_	:							List[Any]									=[*signature.parameters.keys()]

																	lowerCamelCase_	:							List[str]									=[
																	    "states",
																	    "actions",
																	    "rewards",
																	    "returns_to_go",
																	    "timesteps",
																	    "attention_mask",
																	]

																	self.assertListEqual(arg_names[: len(snake_case__ )]       ,	snake_case__ )







@require_torch
class        lowercase__      (							unittest.TestCase       ):


			@slow
			def      UpperCAmelCase__						(				self					:  Any ):
										lowerCamelCase_	:							Optional[int]									=2  # number of steps of autoregressive prediction we will perform
										lowerCamelCase_	:							int									=10  # defined by the RL environment, may be normalized
										lowerCamelCase_	:							List[Any]									=DecisionTransformerModel.from_pretrained("edbeeching/decision-transformer-gym-hopper-expert" )
										lowerCamelCase_	:							Union[str, Any]									=model.to(snake_case__ )
										lowerCamelCase_	:							Any									=model.config
										torch.manual_seed(0 )
										lowerCamelCase_	:							Optional[Any]									=torch.randn(1       ,	1       ,	config.state_dim ).to(device=snake_case__       ,	dtype=torch.floataa )  # env.reset()

										lowerCamelCase_	:							Optional[Any]									=torch.tensor(
										    [[0.242_793, -0.28_693_074, 0.8_742_613], [0.67_815_274, -0.08_101_085, -0.12_952_147]]       ,	device=snake_case__ )

										lowerCamelCase_	:							int									=torch.tensor(snake_case__       ,	device=snake_case__       ,	dtype=torch.floataa ).reshape(1       ,	1       ,	1 )
										lowerCamelCase_	:							str									=state
										lowerCamelCase_	:							Optional[int]									=torch.zeros(1       ,	0       ,	config.act_dim       ,	device=snake_case__       ,	dtype=torch.floataa )
										lowerCamelCase_	:							int									=torch.zeros(1       ,	0       ,	device=snake_case__       ,	dtype=torch.floataa )
										lowerCamelCase_	:							Tuple									=torch.tensor(0       ,	device=snake_case__       ,	dtype=torch.long ).reshape(1       ,	1 )

										for step in range(snake_case__ ):
																	lowerCamelCase_	:							str									=torch.cat([actions, torch.zeros(1       ,	1       ,	config.act_dim       ,	device=snake_case__ )]       ,	dim=1 )
																	lowerCamelCase_	:							Union[str, Any]									=torch.cat([rewards, torch.zeros(1       ,	1       ,	device=snake_case__ )]       ,	dim=1 )

																	lowerCamelCase_	:							Optional[int]									=torch.ones(1       ,	states.shape[1] ).to(dtype=torch.long       ,	device=states.device )

																	with torch.no_grad():
																								lowerCamelCase_							,						lowerCamelCase_							,						lowerCamelCase_	:							Dict									=model(
																								    states=snake_case__       ,	actions=snake_case__       ,	rewards=snake_case__       ,	returns_to_go=snake_case__       ,	timesteps=snake_case__       ,	attention_mask=snake_case__       ,	return_dict=snake_case__       ,	)

																	self.assertEqual(action_pred.shape       ,	actions.shape )
																	self.assertTrue(torch.allclose(action_pred[0, -1]       ,	expected_outputs[step]       ,	atol=1E-4 ) )
																	lowerCamelCase_							,						lowerCamelCase_							,						lowerCamelCase_							,						lowerCamelCase_	:							Optional[Any]									=(  # env.step(action)
																	    torch.randn(1       ,	1       ,	config.state_dim ).to(device=snake_case__       ,	dtype=torch.floataa ),
																	    1.0,
																	    False,
																	    {},
																	)

																	lowerCamelCase_	:							str									=action_pred[0, -1]
																	lowerCamelCase_	:							Optional[int]									=torch.cat([states, state]       ,	dim=1 )
																	lowerCamelCase_	:							Optional[Any]									=returns_to_go[0, -1] - reward
																	lowerCamelCase_	:							str									=torch.cat([returns_to_go, pred_return.reshape(1       ,	1       ,	1 )]       ,	dim=1 )
																	lowerCamelCase_	:							int									=torch.cat(
																	    [timesteps, torch.ones((1, 1)       ,	device=snake_case__       ,	dtype=torch.long ) * (step + 1)]       ,	dim=1 )



