






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    _lowerCAmelCase			:



				"""simple docstring"""



				def __init__(      self    :    Tuple			,		UpperCamelCase__    :    List[Any]			,		UpperCamelCase__    :    Optional[int]=1_3			,		UpperCamelCase__    :    Union[str, Any]=7			,		UpperCamelCase__    :    Optional[int]=6			,		UpperCamelCase__    :    Dict=1_7			,		UpperCamelCase__    :    Union[str, Any]=2_3			,		UpperCamelCase__    :    int=1_1			,		UpperCamelCase__    :    Optional[int]=True			,		):






											'''simple docstring'''







											snake_case__	  =       parent
											snake_case__	  =       batch_size
											snake_case__	  =       seq_length
											snake_case__	  =       act_dim
											snake_case__	  =       state_dim
											snake_case__	  =       hidden_size
											snake_case__	  =       max_length
											snake_case__	  =       is_training



				def 					__magic_name__							(      self    :    List[Any]):






											'''simple docstring'''







											snake_case__	  =       floats_tensor((self.batch_size, self.seq_length, self.state_dim))
											snake_case__	  =       floats_tensor((self.batch_size, self.seq_length, self.act_dim))
											snake_case__	  =       floats_tensor((self.batch_size, self.seq_length, 1))
											snake_case__	  =       floats_tensor((self.batch_size, self.seq_length, 1))
											snake_case__	  =       ids_tensor((self.batch_size, self.seq_length)			,		vocab_size=1_0_0_0)
											snake_case__	  =       random_attention_mask((self.batch_size, self.seq_length))

											snake_case__	  =       self.get_config()

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



				def 					__magic_name__							(      self    :    str):






											'''simple docstring'''







											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 					__magic_name__							(      self    :    Any			,		UpperCamelCase__    :    int			,		UpperCamelCase__    :    List[str]			,		UpperCamelCase__    :    Optional[Any]			,		UpperCamelCase__    :    Union[str, Any]			,		UpperCamelCase__    :    str			,		UpperCamelCase__    :    Tuple			,		UpperCamelCase__    :    str			,		):






											'''simple docstring'''







											snake_case__	  =       DecisionTransformerModel(config=UpperCamelCase__)
											model.to(UpperCamelCase__)
											model.eval()
											snake_case__	  =       model(UpperCamelCase__			,		UpperCamelCase__			,		UpperCamelCase__			,		UpperCamelCase__			,		UpperCamelCase__			,		UpperCamelCase__)

											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 					__magic_name__							(      self    :    Union[str, Any]):






											'''simple docstring'''







											snake_case__	  =       self.prepare_config_and_inputs()
											(
											    (
											    snake_case__
											)  , (
											    snake_case__
											)  , (
											    snake_case__
											)  , (
											    snake_case__
											)  , (
											    snake_case__
											)  , (
											    snake_case__
											)  , (
											    snake_case__
											)  , 
											)	  =       config_and_inputs
											snake_case__	  =       {
											    """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    _lowerCAmelCase			(		lowercase_     ,      lowercase_     ,      lowercase_     ,      unittest.TestCase      ):



				"""simple docstring"""



				_lowercase				:   List[str]      =  (DecisionTransformerModel,) if is_torch_available() else ()
				_lowercase				:   Dict      =  ()
				_lowercase				:   Tuple      =  {'''feature-extraction''': DecisionTransformerModel} if is_torch_available() else {}

				# Ignoring of a failing test from GenerationTesterMixin, as the model does not use inputs_ids
				_lowercase				:   str      =  False

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



				def 					__magic_name__							(      self    :    Optional[Any]):






											'''simple docstring'''







											snake_case__	  =       DecisionTransformerModelTester(self)
											snake_case__	  =       ConfigTester(self			,		config_class=UpperCamelCase__			,		hidden_size=3_7)



				def 					__magic_name__							(      self    :    int):






											'''simple docstring'''







											self.config_tester.run_common_tests()



				def 					__magic_name__							(      self    :    Optional[int]):






											'''simple docstring'''







											snake_case__	  =       self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_model(*UpperCamelCase__)



				@slow
				def 					__magic_name__							(      self    :    str):






											'''simple docstring'''







											for model_name in DECISION_TRANSFORMER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																		snake_case__	  =       DecisionTransformerModel.from_pretrained(UpperCamelCase__)
																		self.assertIsNotNone(UpperCamelCase__)



				def 					__magic_name__							(      self    :    Optional[int]):






											'''simple docstring'''







											snake_case__  , snake_case__	  =       self.model_tester.prepare_config_and_inputs_for_common()

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

																		snake_case__	  =       [
																		    """states""",
																		    """actions""",
																		    """rewards""",
																		    """returns_to_go""",
																		    """timesteps""",
																		    """attention_mask""",
																		]

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


@require_torch
class    _lowerCAmelCase			(		unittest.TestCase      ):



				"""simple docstring"""



				@slow
				def 					__magic_name__							(      self    :    List[str]):






											'''simple docstring'''







											snake_case__	  =       2  # number of steps of autoregressive prediction we will perform
											snake_case__	  =       1_0  # defined by the RL environment, may be normalized
											snake_case__	  =       DecisionTransformerModel.from_pretrained("""edbeeching/decision-transformer-gym-hopper-expert""")
											snake_case__	  =       model.to(UpperCamelCase__)
											snake_case__	  =       model.config
											torch.manual_seed(0)
											snake_case__	  =       torch.randn(1			,		1			,		config.state_dim).to(device=UpperCamelCase__			,		dtype=torch.floataa)  # env.reset()

											snake_case__	  =       torch.tensor(
											    [[0.24_27_93, -0.28_69_30_74, 0.8_74_26_13], [0.67_81_52_74, -0.08_10_10_85, -0.12_95_21_47]]			,		device=UpperCamelCase__)

											snake_case__	  =       torch.tensor(UpperCamelCase__			,		device=UpperCamelCase__			,		dtype=torch.floataa).reshape(1			,		1			,		1)
											snake_case__	  =       state
											snake_case__	  =       torch.zeros(1			,		0			,		config.act_dim			,		device=UpperCamelCase__			,		dtype=torch.floataa)
											snake_case__	  =       torch.zeros(1			,		0			,		device=UpperCamelCase__			,		dtype=torch.floataa)
											snake_case__	  =       torch.tensor(0			,		device=UpperCamelCase__			,		dtype=torch.long).reshape(1			,		1)

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

																		snake_case__	  =       torch.ones(1			,		states.shape[1]).to(dtype=torch.long			,		device=states.device)

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

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

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








import math
from typing import Any, Callable, List, Optional, Tuple, Union

import numpy as np
import torch

from ...models import TaFilmDecoder
from ...schedulers import DDPMScheduler
from ...utils import is_onnx_available, logging, randn_tensor


if is_onnx_available():
	from ..onnx_utils import OnnxRuntimeModel

from ..pipeline_utils import AudioPipelineOutput, DiffusionPipeline
from .continous_encoder import SpectrogramContEncoder
from .notes_encoder import SpectrogramNotesEncoder


a__       =     logging.get_logger(__name__)  # pylint: disable=invalid-name

a__       =     2_5_6


class    _lowerCAmelCase			(		lowercase_      ):



				"""simple docstring"""



				_lowercase				:   str      =  ['''melgan''']



				def __init__(      self    :    int			,		UpperCamelCase__    :    SpectrogramNotesEncoder			,		UpperCamelCase__    :    SpectrogramContEncoder			,		UpperCamelCase__    :    TaFilmDecoder			,		UpperCamelCase__    :    DDPMScheduler			,		UpperCamelCase__    :    OnnxRuntimeModel if is_onnx_available() else Any			,		):






											'''simple docstring'''







											super().__init__()

											# From MELGAN
											snake_case__	  =       math.log(1E-5)  # Matches MelGAN training.
											snake_case__	  =       4.0  # Largest value for most examples
											snake_case__	  =       1_2_8

											self.register_modules(
											    notes_encoder=UpperCamelCase__			,		continuous_encoder=UpperCamelCase__			,		decoder=UpperCamelCase__			,		scheduler=UpperCamelCase__			,		melgan=UpperCamelCase__			,		)



				def 					__magic_name__							(      self    :    Tuple			,		UpperCamelCase__    :    Any			,		UpperCamelCase__    :    int=(-1.0, 1.0)			,		UpperCamelCase__    :    Tuple=False):






											'''simple docstring'''







											snake_case__  , snake_case__	  =       output_range
											if clip:
																		snake_case__	  =       torch.clip(UpperCamelCase__			,		self.min_value			,		self.max_value)
											# Scale to [0, 1].
											snake_case__	  =       (features - self.min_value) / (self.max_value - self.min_value)
											# Scale to [min_out, max_out].
											return zero_one * (max_out - min_out) + min_out



				def 					__magic_name__							(      self    :    Optional[Any]			,		UpperCamelCase__    :    str			,		UpperCamelCase__    :    List[str]=(-1.0, 1.0)			,		UpperCamelCase__    :    Optional[Any]=False):






											'''simple docstring'''







											snake_case__  , snake_case__	  =       input_range
											snake_case__	  =       torch.clip(UpperCamelCase__			,		UpperCamelCase__			,		UpperCamelCase__) if clip else outputs
											# Scale to [0, 1].
											snake_case__	  =       (outputs - min_out) / (max_out - min_out)
											# Scale to [self.min_value, self.max_value].
											return zero_one * (self.max_value - self.min_value) + self.min_value



				def 					__magic_name__							(      self    :    Optional[int]			,		UpperCamelCase__    :    Union[str, Any]			,		UpperCamelCase__    :    Optional[int]			,		UpperCamelCase__    :    Any):






											'''simple docstring'''







											snake_case__	  =       input_tokens > 0
											snake_case__  , snake_case__	  =       self.notes_encoder(
											    encoder_input_tokens=UpperCamelCase__			,		encoder_inputs_mask=UpperCamelCase__)

											snake_case__  , snake_case__	  =       self.continuous_encoder(
											    encoder_inputs=UpperCamelCase__			,		encoder_inputs_mask=UpperCamelCase__)

											return [(tokens_encoded, tokens_mask), (continuous_encoded, continuous_mask)]



				def 					__magic_name__							(      self    :    List[Any]			,		UpperCamelCase__    :    Dict			,		UpperCamelCase__    :    Any			,		UpperCamelCase__    :    int):






											'''simple docstring'''







											snake_case__	  =       noise_time
											if not torch.is_tensor(UpperCamelCase__):
																		snake_case__	  =       torch.tensor([timesteps]			,		dtype=torch.long			,		device=input_tokens.device)
											elif torch.is_tensor(UpperCamelCase__) and len(timesteps.shape) == 0:
																		snake_case__	  =       timesteps[None].to(input_tokens.device)

											# broadcast to batch dimension in a way that's compatible with ONNX/Core ML
											snake_case__	  =       timesteps * torch.ones(input_tokens.shape[0]			,		dtype=timesteps.dtype			,		device=timesteps.device)

											snake_case__	  =       self.decoder(
											    encodings_and_masks=UpperCamelCase__			,		decoder_input_tokens=UpperCamelCase__			,		decoder_noise_time=UpperCamelCase__)
											return logits



				@torch.no_grad()
				def __call__(      self    :    Dict			,		UpperCamelCase__    :    List[List[int]]			,		UpperCamelCase__    :    Optional[torch.Generator] = None			,		UpperCamelCase__    :    int = 1_0_0			,		UpperCamelCase__    :    bool = True			,		UpperCamelCase__    :    str = "numpy"			,		UpperCamelCase__    :    Optional[Callable[[int, int, torch.FloatTensor], None]] = None			,		UpperCamelCase__    :    int = 1			,		):






											'''simple docstring'''







											if (callback_steps is None) or (
											    callback_steps is not None and (not isinstance(UpperCamelCase__			,		UpperCamelCase__) or callback_steps <= 0)
											):
																		raise ValueError(
																		    F'''`callback_steps` has to be a positive integer but is {callback_steps} of type'''
																		    F''' {type(UpperCamelCase__)}.''')

											snake_case__	  =       np.zeros([1, TARGET_FEATURE_LENGTH, self.n_dims]			,		dtype=np.floataa)
											snake_case__	  =       np.zeros([1, 0, self.n_dims]			,		np.floataa)
											snake_case__	  =       torch.ones((1, TARGET_FEATURE_LENGTH)			,		dtype=UpperCamelCase__			,		device=self.device)

											for i, encoder_input_tokens in enumerate(UpperCamelCase__):
																		if i == 0:
																									snake_case__	  =       torch.from_numpy(pred_mel[:1].copy()).to(
																									    device=self.device			,		dtype=self.decoder.dtype)
																									# The first chunk has no previous context.
																									snake_case__	  =       torch.zeros((1, TARGET_FEATURE_LENGTH)			,		dtype=UpperCamelCase__			,		device=self.device)
																		else:
																									# The full song pipeline does not feed in a context feature, so the mask
																									# will be all 0s after the feature converter. Because we know we're
																									# feeding in a full context chunk from the previous prediction, set it
																									# to all 1s.
																									snake_case__	  =       ones

																		snake_case__	  =       self.scale_features(
																		    UpperCamelCase__			,		output_range=[-1.0, 1.0]			,		clip=UpperCamelCase__)

																		snake_case__	  =       self.encode(
																		    input_tokens=torch.IntTensor([encoder_input_tokens]).to(device=self.device)			,		continuous_inputs=UpperCamelCase__			,		continuous_mask=UpperCamelCase__			,		)

																		# Sample encoder_continuous_inputs shaped gaussian noise to begin loop
																		snake_case__	  =       randn_tensor(
																		    shape=encoder_continuous_inputs.shape			,		generator=UpperCamelCase__			,		device=self.device			,		dtype=self.decoder.dtype			,		)

																		# set step values
																		self.scheduler.set_timesteps(UpperCamelCase__)

																		# Denoising diffusion loop
																		for j, t in enumerate(self.progress_bar(self.scheduler.timesteps)):
																									snake_case__	  =       self.decode(
																									    encodings_and_masks=UpperCamelCase__			,		input_tokens=UpperCamelCase__			,		noise_time=t / self.scheduler.config.num_train_timesteps			,		)

																									# Compute previous output: x_t -> x_t-1
																									snake_case__	  =       self.scheduler.step(UpperCamelCase__			,		UpperCamelCase__			,		UpperCamelCase__			,		generator=UpperCamelCase__).prev_sample

																		snake_case__	  =       self.scale_to_features(UpperCamelCase__			,		input_range=[-1.0, 1.0])
																		snake_case__	  =       mel[:1]
																		snake_case__	  =       mel.cpu().float().numpy()

																		snake_case__	  =       np.concatenate([full_pred_mel, pred_mel[:1]]			,		axis=1)

																		# call the callback, if provided
																		if callback is not None and i % callback_steps == 0:
																									callback(UpperCamelCase__			,		UpperCamelCase__)

																		logger.info("""Generated segment"""			,		UpperCamelCase__)

											if output_type == "numpy" and not is_onnx_available():
																		raise ValueError(
																		    """Cannot return output in 'np' format if ONNX is not available. Make sure to have ONNX installed or set 'output_type' to 'mel'.""")
											elif output_type == "numpy" and self.melgan is None:
																		raise ValueError(
																		    """Cannot return output in 'np' format if melgan component is not defined. Make sure to define `self.melgan` or set 'output_type' to 'mel'.""")

											if output_type == "numpy":
																		snake_case__	  =       self.melgan(input_features=full_pred_mel.astype(np.floataa))
											else:
																		snake_case__	  =       full_pred_mel

											if not return_dict:
																		return (output,)

											return AudioPipelineOutput(audios=UpperCamelCase__)
