






import torch

from diffusers import EulerAncestralDiscreteScheduler
from diffusers.utils import torch_device

from .test_schedulers import SchedulerCommonTest




class        snake_case_					(						__UpperCamelCase			):

						"""simple docstring"""



						snake_case__   					=			(EulerAncestralDiscreteScheduler,)
						snake_case__   					=			10



						def 					UpperCAmelCase__				(self:				Any   ,				**__UpperCAmelCase:				List[Any]      )						->    Any:




											'''simple docstring'''



											__a				:	Any			=				{
											    "num_train_timesteps": 1100,
											    "beta_start": 0.00_01,
											    "beta_end": 0.02,
											    "beta_schedule": "linear",
											}

											config.update(**__UpperCAmelCase      )
											return config



						def 					UpperCAmelCase__				(self:				Any      )						->    str:




											'''simple docstring'''



											for timesteps in [10, 50, 100, 1000]:
																self.check_over_configs(num_train_timesteps=__UpperCAmelCase      )



						def 					UpperCAmelCase__				(self:				str      )						->    List[Any]:




											'''simple docstring'''



											for beta_start, beta_end in zip([0.0_00_01, 0.00_01, 0.0_01]   ,				[0.00_02, 0.0_02, 0.02]      ):
																self.check_over_configs(beta_start=__UpperCAmelCase   ,				beta_end=__UpperCAmelCase      )



						def 					UpperCAmelCase__				(self:				Optional[int]      )						->    Any:




											'''simple docstring'''



											for schedule in ["linear", "scaled_linear"]:
																self.check_over_configs(beta_schedule=__UpperCAmelCase      )



						def 					UpperCAmelCase__				(self:				str      )						->    List[Any]:




											'''simple docstring'''



											for prediction_type in ["epsilon", "v_prediction"]:
																self.check_over_configs(prediction_type=__UpperCAmelCase      )



						def 					UpperCAmelCase__				(self:				List[str]      )						->    Any:




											'''simple docstring'''



											__a				:	Optional[Any]			=				self.scheduler_classes[0]
											__a				:	List[Any]			=				self.get_scheduler_config()
											__a				:	List[Any]			=				scheduler_class(**__UpperCAmelCase      )

											scheduler.set_timesteps(self.num_inference_steps      )

											__a				:	str			=				torch.manual_seed(0      )

											__a				:	Union[str, Any]			=				self.dummy_model()
											__a				:	Optional[int]			=				self.dummy_sample_deter * scheduler.init_noise_sigma.cpu()
											__a				:	List[Any]			=				sample.to(__UpperCAmelCase      )

											for i, t in enumerate(scheduler.timesteps      ):
																__a				:	Union[str, Any]			=				scheduler.scale_model_input(__UpperCAmelCase   ,				__UpperCAmelCase      )

																__a				:	Optional[int]			=				model(__UpperCAmelCase   ,				__UpperCAmelCase      )

																__a				:	Optional[int]			=				scheduler.step(__UpperCAmelCase   ,				__UpperCAmelCase   ,				__UpperCAmelCase   ,				generator=__UpperCAmelCase      )
																__a				:	Optional[Any]			=				output.prev_sample

											__a				:	List[Any]			=				torch.sum(torch.abs(__UpperCAmelCase      )      )
											__a				:	List[Any]			=				torch.mean(torch.abs(__UpperCAmelCase      )      )

											assert abs(result_sum.item() - 1_52.31_92      ) < 1E-2
											assert abs(result_mean.item() - 0.19_83      ) < 1E-3



						def 					UpperCAmelCase__				(self:				Optional[int]      )						->    List[str]:




											'''simple docstring'''



											__a				:	Dict			=				self.scheduler_classes[0]
											__a				:	List[Any]			=				self.get_scheduler_config(prediction_type="v_prediction"      )
											__a				:	str			=				scheduler_class(**__UpperCAmelCase      )

											scheduler.set_timesteps(self.num_inference_steps      )

											__a				:	Optional[int]			=				torch.manual_seed(0      )

											__a				:	Optional[int]			=				self.dummy_model()
											__a				:	Tuple			=				self.dummy_sample_deter * scheduler.init_noise_sigma
											__a				:	int			=				sample.to(__UpperCAmelCase      )

											for i, t in enumerate(scheduler.timesteps      ):
																__a				:	Dict			=				scheduler.scale_model_input(__UpperCAmelCase   ,				__UpperCAmelCase      )

																__a				:	Dict			=				model(__UpperCAmelCase   ,				__UpperCAmelCase      )

																__a				:	List[Any]			=				scheduler.step(__UpperCAmelCase   ,				__UpperCAmelCase   ,				__UpperCAmelCase   ,				generator=__UpperCAmelCase      )
																__a				:	int			=				output.prev_sample

											__a				:	Optional[int]			=				torch.sum(torch.abs(__UpperCAmelCase      )      )
											__a				:	int			=				torch.mean(torch.abs(__UpperCAmelCase      )      )

											assert abs(result_sum.item() - 1_08.44_39      ) < 1E-2
											assert abs(result_mean.item() - 0.14_12      ) < 1E-3



						def 					UpperCAmelCase__				(self:				List[Any]      )						->    int:




											'''simple docstring'''



											__a				:	int			=				self.scheduler_classes[0]
											__a				:	int			=				self.get_scheduler_config()
											__a				:	Optional[Any]			=				scheduler_class(**__UpperCAmelCase      )

											scheduler.set_timesteps(self.num_inference_steps   ,				device=__UpperCAmelCase      )
											__a				:	Tuple			=				torch.manual_seed(0      )

											__a				:	List[str]			=				self.dummy_model()
											__a				:	List[Any]			=				self.dummy_sample_deter * scheduler.init_noise_sigma.cpu()
											__a				:	Optional[Any]			=				sample.to(__UpperCAmelCase      )

											for t in scheduler.timesteps:
																__a				:	int			=				scheduler.scale_model_input(__UpperCAmelCase   ,				__UpperCAmelCase      )

																__a				:	List[str]			=				model(__UpperCAmelCase   ,				__UpperCAmelCase      )

																__a				:	Optional[Any]			=				scheduler.step(__UpperCAmelCase   ,				__UpperCAmelCase   ,				__UpperCAmelCase   ,				generator=__UpperCAmelCase      )
																__a				:	Union[str, Any]			=				output.prev_sample

											__a				:	int			=				torch.sum(torch.abs(__UpperCAmelCase      )      )
											__a				:	List[Any]			=				torch.mean(torch.abs(__UpperCAmelCase      )      )

											assert abs(result_sum.item() - 1_52.31_92      ) < 1E-2
											assert abs(result_mean.item() - 0.19_83      ) < 1E-3













import copy

from ...configuration_utils import PretrainedConfig
from ...utils import logging
from ..auto.configuration_auto import AutoConfig


UpperCAmelCase__			 =	logging.get_logger(__name__)

UpperCAmelCase__			 =	{
    '''facebook/musicgen-small''': '''https://huggingface.co/facebook/musicgen-small/resolve/main/config.json''',
    # See all Musicgen models at https://huggingface.co/models?filter=musicgen
}




class        snake_case_					(						__UpperCamelCase			):

						"""simple docstring"""



						snake_case__   					=			"""musicgen_decoder"""
						snake_case__   					=			["""past_key_values"""]



						def __init__(self:				str   ,				__UpperCAmelCase:				Union[str, Any]=2048   ,				__UpperCAmelCase:				Union[str, Any]=2048   ,				__UpperCAmelCase:				List[Any]=24   ,				__UpperCAmelCase:				Dict=4096   ,				__UpperCAmelCase:				int=16   ,				__UpperCAmelCase:				Dict=0.0   ,				__UpperCAmelCase:				str=True   ,				__UpperCAmelCase:				Tuple="gelu"   ,				__UpperCAmelCase:				Optional[Any]=1024   ,				__UpperCAmelCase:				Optional[Any]=0.1   ,				__UpperCAmelCase:				List[Any]=0.0   ,				__UpperCAmelCase:				int=0.0   ,				__UpperCAmelCase:				List[Any]=0.02   ,				__UpperCAmelCase:				Union[str, Any]=False   ,				__UpperCAmelCase:				Any=4   ,				__UpperCAmelCase:				Tuple=2048   ,				__UpperCAmelCase:				Tuple=2048   ,				__UpperCAmelCase:				Dict=None   ,				__UpperCAmelCase:				Dict=False   ,				**__UpperCAmelCase:				Any   ,				)						->    Optional[Any]:




											'''simple docstring'''



											__a				:	Dict			=				vocab_size
											__a				:	List[str]			=				max_position_embeddings
											__a				:	int			=				hidden_size
											__a				:	Union[str, Any]			=				ffn_dim
											__a				:	List[Any]			=				num_hidden_layers
											__a				:	List[Any]			=				num_attention_heads
											__a				:	List[Any]			=				dropout
											__a				:	str			=				attention_dropout
											__a				:	Union[str, Any]			=				activation_dropout
											__a				:	Optional[int]			=				activation_function
											__a				:	List[str]			=				initializer_factor
											__a				:	Union[str, Any]			=				layerdrop
											__a				:	Optional[int]			=				use_cache
											__a				:	List[Any]			=				scale_embedding  # scale factor will be sqrt(d_model) if True
											__a				:	str			=				num_codebooks
											super().__init__(
											    pad_token_id=__UpperCAmelCase   ,				bos_token_id=__UpperCAmelCase   ,				eos_token_id=__UpperCAmelCase   ,				tie_word_embeddings=__UpperCAmelCase   ,				**__UpperCAmelCase   ,				)




class        snake_case_					(						__UpperCamelCase			):

						"""simple docstring"""



						snake_case__   					=			"""musicgen"""
						snake_case__   					=			True



						def __init__(self:				Union[str, Any]   ,				**__UpperCAmelCase:				Optional[int]      )						->    List[str]:




											'''simple docstring'''



											super().__init__(**__UpperCAmelCase      )
											if "text_encoder" not in kwargs or "audio_encoder" not in kwargs or "decoder" not in kwargs:
																raise ValueError("Config has to be initialized with text_encoder, audio_encoder and decoder config"      )

											__a				:	List[Any]			=				kwargs.pop("text_encoder"      )
											__a				:	Optional[Any]			=				text_encoder_config.pop("model_type"      )

											__a				:	Union[str, Any]			=				kwargs.pop("audio_encoder"      )
											__a				:	Dict			=				audio_encoder_config.pop("model_type"      )

											__a				:	Optional[Any]			=				kwargs.pop("decoder"      )

											__a				:	Optional[int]			=				AutoConfig.for_model(__UpperCAmelCase   ,				**__UpperCAmelCase      )
											__a				:	int			=				AutoConfig.for_model(__UpperCAmelCase   ,				**__UpperCAmelCase      )
											__a				:	List[Any]			=				MusicgenDecoderConfig(**__UpperCAmelCase      )
											__a				:	Optional[int]			=				True



						@classmethod
						def 					UpperCAmelCase__				(cls:				Union[str, Any]   ,				__UpperCAmelCase:				PretrainedConfig   ,				__UpperCAmelCase:				PretrainedConfig   ,				__UpperCAmelCase:				MusicgenDecoderConfig   ,				**__UpperCAmelCase:				Any   ,				)						->    Tuple:




											'''simple docstring'''



											return cls(
											    text_encoder=text_encoder_config.to_dict()   ,				audio_encoder=audio_encoder_config.to_dict()   ,				decoder=decoder_config.to_dict()   ,				**__UpperCAmelCase   ,				)



						def 					UpperCAmelCase__				(self:				List[str]      )						->    int:




											'''simple docstring'''



											__a				:	Dict			=				copy.deepcopy(self.__dict__      )
											__a				:	Optional[Any]			=				self.text_encoder.to_dict()
											__a				:	Optional[Any]			=				self.audio_encoder.to_dict()
											__a				:	Tuple			=				self.decoder.to_dict()
											__a				:	List[Any]			=				self.__class__.model_type
											return output





