Spaces:

MBZUAI
/

artst-tts-demo

Runtime error

App Files Files Community

artst-tts-demo / artst /models /artst.py

herwoww

first upload

1547a56 7 months ago

raw history blame contribute delete

No virus

64.6 kB

	# --------------------------------------------------------
	# ArTST: Arabic Text and Speech Transformer (https://arxiv.org/abs/2310.16621)
	# Github source: https://github.com/mbzuai-nlp/ArTST
	# Based on speecht5, fairseq and espnet code bases
	# https://github.com/microsoft/SpeechT5/tree/main/SpeechT5; https://github.com/pytorch/fairseq; https://github.com/espnet/espnet
	# --------------------------------------------------------

	import logging
	from ast import literal_eval
	from typing import Dict, List, Optional, Tuple

	import torch
	import torch.nn.functional as F
	from fairseq import utils
	from fairseq.models import (
	FairseqEncoderDecoderModel,
	FairseqIncrementalDecoder,
	register_model,
	register_model_architecture,
	)
	from .modules.text_encoder_prenet import TextEncoderPrenet
	from .modules.text_decoder_prenet import TextDecoderPrenet
	from .modules.text_decoder_postnet import TextDecoderPostnet
	from .modules.speech_encoder_prenet import SpeechEncoderPrenet
	from .modules.speech_encoder_postnet import SpeechEncoderPostnet
	from .modules.speech_decoder_prenet import SpeechDecoderPrenet
	from .modules.speech_decoder_postnet import SpeechDecoderPostnet
	from .modules.speaker_decoder_postnet import SpeakerDecoderPostnet
	from .modules.encoder import TransformerEncoder
	from .modules.decoder import TransformerDecoder
	from fairseq.modules.transformer_sentence_encoder import init_bert_params
	from fairseq.models.transformer import Embedding
	from fairseq.modules import (
	GumbelVectorQuantizer,
	)
	from torch import Tensor


	logger = logging.getLogger(__name__)

	DEFAULT_MAX_TEXT_POSITIONS = 450
	DEFAULT_MAX_SPEECH_POSITIONS = 4000


	@register_model("artst_transformer")
	class ArTSTTransformerModel(FairseqEncoderDecoderModel):
	"""Adapted Transformer model (https://arxiv.org/abs/1706.03762) for
	speech-to-text tasks. The Transformer encoder/decoder remains the same.
	A trainable input subsampler is prepended to the Transformer encoder to
	project inputs into the encoder dimension as well as downsample input
	sequence for computational efficiency."""

	def __init__(
	self,
	args,
	encoder, decoder,
	text_encoder_prenet, speech_encoder_prenet,
	text_decoder_prenet, speech_decoder_prenet,
	text_decoder_postnet, speech_decoder_postnet,
	speaker_decoder_postnet, speech_encoder_postnet,
	):
	super().__init__(encoder, decoder)

	self.encoder = encoder
	self.decoder = decoder

	self.text_encoder_prenet = text_encoder_prenet
	self.speech_encoder_prenet = speech_encoder_prenet

	self.text_decoder_prenet = text_decoder_prenet
	self.speech_decoder_prenet = speech_decoder_prenet

	self.text_decoder_postnet = text_decoder_postnet
	self.speech_decoder_postnet = speech_decoder_postnet
	self.speaker_decoder_postnet = speaker_decoder_postnet

	self.hubert_layer = speech_encoder_postnet

	self.reduction_factor = args.reduction_factor
	self.spk_embed_dim = args.spk_embed_dim

	# define projection layer
	self.spk_embed_integration_type = args.spk_embed_integration_type
	if self.spk_embed_dim is not None and self.spk_embed_integration_type != 'pre':
	if self.spk_embed_integration_type == "add":
	self.projection = torch.nn.Linear(self.spk_embed_dim, args.decoder_embed_dim)
	else:
	self.projection = torch.nn.Linear(
	args.decoder_embed_dim + self.spk_embed_dim, args.decoder_embed_dim
	)

	# Hawau: here we can add language embedding integration

	self.use_codebook = args.use_codebook
	self.codebook_prob = getattr(args, "codebook_prob", 0.5) # args.codebook_prob
	if self.use_codebook:
	vq_dim = args.latent_dim if args.latent_dim > 0 else args.encoder_embed_dim
	self.quantizer = GumbelVectorQuantizer(
	dim=args.encoder_embed_dim,
	num_vars=args.latent_vars,
	temp=args.latent_temp,
	groups=args.latent_groups,
	combine_groups=False,
	vq_dim=vq_dim,
	time_first=True,
	weight_proj_depth=args.quantizer_depth,
	weight_proj_factor=args.quantizer_factor,
	)

	self.num_updates = 0

	# # Follow BERT's random weight initialization (for BART)
	if args.bert_init:
	self.apply(init_bert_params)
	self.args = args
	self.prune_modules(args.modules_filter)

	@staticmethod
	def add_args(parser):
	"""Add model-specific arguments to the parser."""
	# Transformer
	parser.add_argument(
	"--activation-fn",
	type=str,
	choices=utils.get_available_activation_fns(),
	help="activation function to use",
	)
	parser.add_argument(
	"--dropout", type=float, metavar="D", help="dropout probability"
	)
	parser.add_argument(
	"--attention-dropout",
	type=float,
	metavar="D",
	help="dropout probability for attention weights",
	)
	parser.add_argument(
	"--activation-dropout",
	"--relu-dropout",
	type=float,
	metavar="D",
	help="dropout probability after activation in FFN.",
	)
	parser.add_argument(
	"--encoder-embed-dim",
	type=int,
	metavar="N",
	help="encoder embedding dimension",
	)
	parser.add_argument(
	"--encoder-ffn-embed-dim",
	type=int,
	metavar="N",
	help="encoder embedding dimension for FFN",
	)
	parser.add_argument(
	"--encoder-layers", type=int, metavar="N", help="num encoder layers"
	)
	parser.add_argument(
	"--encoder-attention-heads",
	type=int,
	metavar="N",
	help="num encoder attention heads",
	)
	parser.add_argument(
	"--encoder-normalize-before",
	action="store_true",
	help="apply layernorm before each encoder block",
	)
	parser.add_argument(
	"--decoder-normalize-before",
	action="store_true",
	help="apply layernorm before each decoder block",
	)
	parser.add_argument(
	"--decoder-embed-dim",
	type=int,
	metavar="N",
	help="decoder embedding dimension",
	)
	parser.add_argument(
	"--decoder-ffn-embed-dim",
	type=int,
	metavar="N",
	help="decoder embedding dimension for FFN",
	)
	parser.add_argument(
	"--decoder-layers", type=int, metavar="N", help="num decoder layers"
	)
	parser.add_argument(
	"--decoder-attention-heads",
	type=int,
	metavar="N",
	help="num decoder attention heads",
	)
	parser.add_argument(
	"--reduction-factor",
	type=int,
	help="reduction factor for decoder",
	)
	parser.add_argument(
	"--spk-embed-dim",
	type=int,
	help="speaker embedding dimension",
	)
	parser.add_argument(
	"--layernorm-embedding",
	action="store_true",
	help="add layernorm to embedding",
	)
	parser.add_argument(
	"--load-pretrained-encoder-from",
	type=str,
	metavar="STR",
	help="model to take encoder weights from (for initialization)",
	)
	parser.add_argument(
	'--freeze-encoder-updates',
	type=int,
	help='number of steps to freeze encoder before finetune'
	)
	parser.add_argument(
	'--freeze-decoder-updates',
	type=int,
	help='number of steps to freeze decoder before finetune'
	)
	parser.add_argument(
	'--no-freeze-encoder-layer',
	type=str,
	help='which encoder layer not freeze during finetune'
	)
	parser.add_argument(
	"--share-input-output-embed",
	action="store_true",
	help="share decoder input and output embeddings",
	)
	parser.add_argument(
	"--share-ctc-embed",
	action="store_true",
	help="share ctc embed and decoder embed",
	)
	parser.add_argument(
	"--encoder-sliding-window-attn",
	default=None,
	type=int,
	help="If not None but a even number, set sliding window attention to encoder's attn_mask, e.g., 4, 10, and 20",
	)

	# Convolutional subsampler
	parser.add_argument(
	"--encoder-speech-prenet",
	default="conv",
	type=str,
	choices=["conv", "linear"],
	help="The type of encoder speech prenet, e.g., conv or linear."
	)
	parser.add_argument(
	"--conv-kernel-sizes",
	default="5,5",
	type=str,
	help="The layer of convolution of encoder speech prenet."
	)
	parser.add_argument(
	"--conv-channels",
	default=1024,
	type=int,
	help="The channels of encoder speech prenet."
	)
	parser.add_argument(
	"--subsample-stride",
	default="2,2",
	type=str,
	help="The subsample stride for conv1dsubsample."
	)
	parser.add_argument(
	"--spk-embed-integration-type",
	type=str,
	choices=["pre", "add"],
	help="speaker embedding integration type"
	)
	parser.add_argument(
	"--dprenet-dropout-rate",
	default=0.5,
	type=float,
	help="The dropout rate of decoder speech prenet."
	)

	## SE
	parser.add_argument(
	"--se-predict",
	default=None,
	choices=["masking", "target", "delta"],
	help="If set, source speech inputs decoder to predict the masking/target/delta of corresponding inputs."
	+ "masking is [0, 1], target is predicted output, delta is difference between inputs and outputs",
	)
	parser.add_argument(
	"--se-decoder-input",
	type=str,
	default="previous_target",
	choices=["previous_target", "source"],
	)

	## SID
	parser.add_argument(
	"--modules-filter",
	default=None,
	type=str,
	help="Remove unused modules for, e.g., SID.",
	)
	parser.add_argument(
	"--sid-pad-prenet",
	action="store_true",
	help="If set, the size of text dictionary is as small as for <pad> token.",
	)
	parser.add_argument(
	"--encoder-attn-branch",
	type=str,
	default="identity,full",
	help="encoder attention branch sliding window, e.g., 'identity,0,2,4,full'",
	)
	parser.add_argument(
	"--encoder-block-branch",
	type=str,
	help="average the output of encoder, e.g., '4,5,6'",
	)
	parser.add_argument(
	"--sid-encoder-cls",
	default=None,
	choices=["encoder"],
	help="If set, add cls vector to the encoder input, e.g., constant vector.",
	)
	parser.add_argument(
	"--sid-shuffle-encoder-input",
	action="store_true",
	help="If set, shuffle encoder input in time.",
	)
	parser.add_argument(
	"--sid-decoder-speaker",
	action="store_true",
	help="If set, apply speaker decoder as transformer decoder.",
	)
	parser.add_argument(
	"--sid-decoder-attn-dim",
	default=128,
	type=int,
	help="Attention dimension in attensive statistics pooling of speaker decoder.",
	)
	parser.add_argument(
	"--sid-t5-postnet",
	action="store_true",
	help="If set, apply TextDecoderPostnet as speaker classification.",
	)
	parser.add_argument(
	"--sid-embed-dim",
	default=128,
	type=int,
	help="Embedding dimension in speaker postnet for speaker identification if embed postnet.",
	)
	parser.add_argument(
	"--sid-pooling-layer",
	default="decoder",
	type=str,
	choices=["decoder-las", "decoder", "encoder", "encoder-cls", "encoder-speaker"],
	help="The output of decoder or encoder uses as SID pooling layer over temporal dimension.",
	)
	parser.add_argument(
	"--sid-no-pooling-bn",
	action="store_true",
	help="If set, not attention batchnorm.",
	)
	parser.add_argument(
	"--sid-no-embed-postnet",
	action="store_true",
	help="If set, no layer between decoder output and classification layer.",
	)
	parser.add_argument(
	"--sid-normalize-postnet",
	action="store_true",
	help="If set, normalize input and weight in postnet/classifier.",
	)
	parser.add_argument(
	"--sid-softmax-type",
	default="softmax",
	choices=["softmax", "amsoftmax", "aamsoftmax"],
	help="If using amsoftmax or aamsoftmax, the target should be given.",
	)
	parser.add_argument(
	"--softmax-scale",
	default=1.0,
	type=float,
	help="Scale for AMSoftmax or AAMSoftmax.",
	)
	parser.add_argument(
	"--softmax-margin",
	default=0.0,
	type=float,
	help="Margin for AMSoftmax or AAMSoftmax.",
	)
	parser.add_argument(
	"--softmax-easy-margin",
	action="store_true",
	help="Enable easy margin for AAMSoftmax.",
	)
	parser.add_argument(
	"--encoder-layerdrop",
	type=float,
	metavar="D",
	help="LayerDrop probability for encoder",
	)
	parser.add_argument(
	"--decoder-layerdrop",
	type=float,
	metavar="D",
	help="LayerDrop probability for decoder",
	)

	## Hubert
	parser.add_argument(
	'--feature-grad-mult',
	type=float,
	help='multiply feature extractor var grads by this'
	)
	parser.add_argument(
	'--logit-temp',
	type=float,
	help='temperature to divide logits by'
	)
	parser.add_argument(
	'--final-dim',
	type=int,
	help="project final representations and targets to this many "
	"dimensions. set to encoder_embed_dim is <= 0"
	)

	# mask
	parser.add_argument(
	'--hubert-mask-length',
	type=int,
	help='mask length'
	)
	parser.add_argument(
	'--mask-prob',
	type=float,
	help='probability of replacing a token with mask'
	)
	parser.add_argument(
	"--mask-selection",
	choices=["static", "uniform", "normal", "poisson"],
	help="how to choose mask length",
	)
	parser.add_argument(
	'--mask-other',
	type=float,
	help="secondary mask argument "
	"(used for more complex distributions), "
	"see help in compute_mask_indices"
	)
	parser.add_argument(
	'--mask-min-space',
	type=int,
	help='min space between spans (if no overlap is enabled)'
	)

	# channel masking
	parser.add_argument(
	'--mask-channel-length',
	type=int,
	help='length of the mask for features (channels)'
	)
	parser.add_argument(
	'--mask-channel-prob',
	type=float,
	help="probability of replacing a feature with 0"
	)
	parser.add_argument(
	"--mask-channel-selection",
	choices=["static", "uniform", "normal", "poisson"],
	help="how to choose mask length for channel masking",
	)
	parser.add_argument(
	'--mask-channel-other',
	type=float,
	help="secondary mask argument "
	"(used for more complex distributions), "
	"see help in compute_mask_indices"
	)
	parser.add_argument(
	'--mask-channel-min-space',
	type=int,
	help='min space between spans (if no overlap is enabled)'
	)

	# abs positional embeddings
	parser.add_argument(
	'--conv-pos',
	type=int,
	help='number of filters for convolutional positional embeddings'
	)
	parser.add_argument(
	'--conv-pos-groups',
	type=int,
	help='number of groups for convolutional positional embedding'
	)

	# codebook related
	parser.add_argument(
	"--use-codebook",
	action="store_true",
	help="whether to use codebook",
	)
	parser.add_argument(
	"--codebook-prob",
	type=float,
	help="probability to use codebook",
	)
	parser.add_argument(
	"--latent-vars",
	type=int,
	help="number of latent variables V in each group of the codebook",
	)
	parser.add_argument(
	"--latent-groups",
	type=int,
	help="number of groups G of latent variables in the codebook",
	)
	parser.add_argument(
	"--latent-dim",
	type=int,
	help="if > 0, uses this dimensionality for latent variables. "
	"otherwise uses final_dim / latent_groups",
	)
	parser.add_argument(
	"--latent-temp",
	type=literal_eval,
	help="temperature for latent variable sampling. "
	"can be tuple of 3 values (start, end, decay)",
	)
	parser.add_argument(
	"--quantizer-depth",
	type=int,
	help="number of quantizer layers",
	)
	parser.add_argument(
	"--quantizer-factor",
	type=int,
	help="number of quantizer layers",
	)
	parser.add_argument(
	"--get-code-distribution",
	action='store_true',
	help="whether to get the code distribution (for test)",
	)

	# relative pos enc
	parser.add_argument(
	"--relative-position-embedding",
	action='store_true',
	help="whether to use relative position embedding",
	)
	parser.add_argument(
	"--num-buckets",
	type=int,
	default=320,
	help="num of buckets for relative position embedding",
	)
	parser.add_argument(
	"--max-distance",
	type=int,
	default=1280,
	help="max distance for relative position embedding",
	)
	parser.add_argument(
	"--encoder-max-relative-position",
	type=int,
	help="max distance for relative position embedding in encoder",
	)
	parser.add_argument(
	"--decoder-max-relative-position",
	type=int,
	help="max distance for relative position embedding in decoder",
	)

	# hubert feature extractor
	parser.add_argument(
	"--conv-feature-layers",
	type=str,
	help= "string describing convolutional feature extraction "
	"layers in form of a python list that contains "
	"[(dim, kernel_size, stride), ...]",
	)
	parser.add_argument(
	"--conv-bias",
	action='store_true',
	help="include bias in conv encoder",
	)
	parser.add_argument(
	"--extractor-mode",
	choices=["default", "layer_norm"],
	help="mode for feature extractor. default has a single group "
	"norm with d groups in the first conv block, whereas layer_norm "
	"has layer norms in every block (meant to use with normalize=True)"
	)

	# others
	parser.add_argument(
	"--bert-init",
	action='store_true',
	help="initilize as bert",
	)
	parser.add_argument(
	"--unb-enc-layer",
	type=int,
	default=-1,
	help="which layer's output is used as the input of decoder",
	)

	# Encoder, Decoder
	@classmethod
	def build_encoder(cls, args, dictionary=None, embed_tokens=None):
	return TransformerEncoder(args, dictionary, embed_tokens)

	@classmethod
	def build_decoder(cls, args):
	return TransformerDecoder(args)

	# Encoder Prenet
	@classmethod
	def build_text_encoder_prenet(cls, embed_tokens, args):
	return TextEncoderPrenet(embed_tokens, args)

	@classmethod
	def build_speech_encoder_prenet(cls, args):
	return SpeechEncoderPrenet(args)

	# Decoder Prenet
	@classmethod
	def build_text_decoder_prenet(cls, embed_tokens, args):
	return TextDecoderPrenet(embed_tokens, args)

	@classmethod
	def build_speech_decoder_prenet(cls, odim, args):
	return SpeechDecoderPrenet(odim, args)

	# Decoder Postnet
	@classmethod
	def build_text_decoder_postnet(cls, embed_tokens, dictionary, args):
	return TextDecoderPostnet(embed_tokens, dictionary, args)

	@classmethod
	def build_speaker_decoder_postnet(cls, embed_dim, class_num, args):
	return SpeakerDecoderPostnet(embed_dim, class_num, args)

	@classmethod
	def build_speech_decoder_postnet(cls, odim, args):
	return SpeechDecoderPostnet(odim, args)

	@classmethod
	def build_speech_encoder_postnet(cls, dictionaries, args):
	return SpeechEncoderPostnet(dictionaries, args)

	@classmethod
	def build_model(cls, args, task):
	"""Build a new model instance."""

	# make sure all arguments are present in older models
	base_architecture(args)

	def build_embedding(dictionary, embed_dim, max_num_embeddings=None):
	num_embeddings = len(dictionary)
	if max_num_embeddings is not None and isinstance(max_num_embeddings, int):
	num_embeddings = min(num_embeddings, max_num_embeddings)
	padding_idx = dictionary.pad()
	return Embedding(num_embeddings, embed_dim, padding_idx)

	if hasattr(args, "sid_pad_prenet") and args.sid_pad_prenet:
	max_num_embeddings = 3 # <pad> at index 2
	else:
	max_num_embeddings = None

	text_decoder_embed_tokens = build_embedding(
	task.dicts["text"], args.decoder_embed_dim, max_num_embeddings
	)

	if args.share_input_output_embed:
	text_encoder_embed_tokens = text_decoder_embed_tokens
	else:
	text_encoder_embed_tokens = build_embedding(
	task.dicts["text"], args.encoder_embed_dim
	)

	speech_odim = args.speech_odim
	if "text" in task.dicts:
	encoder = cls.build_encoder(args, task.dicts["text"], text_encoder_embed_tokens)
	else:
	encoder = cls.build_encoder(args)
	decoder = cls.build_decoder(args)

	text_encoder_prenet = cls.build_text_encoder_prenet(text_encoder_embed_tokens, args)
	speech_encoder_prenet = cls.build_speech_encoder_prenet(args)

	text_decoder_prenet = cls.build_text_decoder_prenet(text_decoder_embed_tokens, args)
	if getattr(args, "sid_pooling_layer", None) == "decoder-las":
	speech_decoder_prenet = cls.build_speech_encoder_prenet(args)
	else:
	speech_decoder_prenet = cls.build_speech_decoder_prenet(speech_odim, args)

	text_decoder_postnet = cls.build_text_decoder_postnet(text_decoder_embed_tokens, task.dicts['text'], args)
	speech_decoder_postnet = cls.build_speech_decoder_postnet(speech_odim, args)

	if getattr(args, "sid_t5_postnet", False):
	speaker_decoder_postnet = None
	else:
	if task.t5_task == "s2c":
	speaker_decoder_postnet = cls.build_speaker_decoder_postnet(args.sid_embed_dim, len(task.dicts['text']), args)
	else:
	speaker_decoder_postnet = None

	if "hubert" in task.dicts:
	speech_encoder_postnet = cls.build_speech_encoder_postnet(task.dicts['hubert'], args)
	else:
	speech_encoder_postnet = None

	return cls(
	args,
	encoder, decoder,
	text_encoder_prenet, speech_encoder_prenet,
	text_decoder_prenet, speech_decoder_prenet,
	text_decoder_postnet, speech_decoder_postnet,
	speaker_decoder_postnet, speech_encoder_postnet,
	)

	def get_normalized_probs(
	self,
	net_output: Tuple[Tensor, Optional[Dict[str, List[Optional[Tensor]]]]],
	log_probs: bool,
	sample: Optional[Dict[str, Tensor]] = None,
	):
	# net_output['encoder_out'] is a (B, T, D) tensor
	lprobs = self.get_normalized_probs_scriptable(net_output, log_probs, sample)
	lprobs.batch_first = True
	return lprobs

	def get_normalized_probs_for_ctc(self, net_output, log_probs):
	"""Get normalized probabilities (or log probs) from a net's output."""

	logits = net_output["encoder_out_for_ctc"][0]
	if log_probs:
	return utils.log_softmax(logits.float(), dim=-1)
	else:
	return utils.softmax(logits.float(), dim=-1)

	def get_logits(self, net_output, is_masked=True):
	if is_masked:
	logits_list = net_output["logit_m_list"]
	else:
	logits_list = net_output["logit_u_list"]
	logits_list = [x.float() for x in logits_list if x is not None]
	return logits_list

	def get_targets(self, sample, net_output, is_masked=True):
	if "logit_m_list" in net_output:
	logits_list = self.get_logits(net_output, is_masked)
	targets_list = [
	x.new_zeros(x.size(0), dtype=torch.long) for x in logits_list
	]
	return targets_list
	else:
	return sample["target"]

	def get_extra_losses(self, net_output):
	extra_losses = []
	names = []

	if "features_pen" in net_output:
	extra_losses.append(net_output["features_pen"])
	names.append("features_pen")

	if "prob_perplexity" in net_output:
	extra_losses.append(
	(net_output["num_vars"] - net_output["prob_perplexity"])
	/ net_output["num_vars"]
	)
	names.append("prob_perplexity")

	return extra_losses, names

	def forward(self, source=None, src_tokens=None, src_lengths=None, prev_output_tokens=None, tgt_lengths=None, spkembs=None, target_list=None, task_name=None, padding_mask=None, only_hubert=False, only_ctc=False, feature_only=False, tgt_enc_layer=None, mask=True):
	"""
	The forward method inherited from the base class has a **kwargs
	argument in its input, which is not supported in torchscript. This
	method overwrites the forward method definition without **kwargs.
	"""
	assert source is not None or src_tokens is not None
	# padding_mask is not none only when input is waveform
	if source is None and padding_mask is None and not feature_only:
	input_type = 'text'
	else:
	input_type = 'speech'

	if prev_output_tokens is not None and len(prev_output_tokens.size()) == 2:
	output_type = 'text'
	codebook_out = {}
	else:
	output_type = 'speech'

	if task_name is not None and task_name == "s2c":
	if target_list is not None and target_list.size(1) == 1 and not getattr(self.args, "sid_t5_postnet", False):
	sid_target = F.one_hot(target_list.squeeze(1), num_classes=self.speaker_decoder_postnet.class_num)
	else:
	sid_target = None
	target_list = None

	# Encoder Prenet
	if input_type == 'text':
	encoder_input, encoder_padding_mask = self.text_encoder_prenet(src_tokens)
	else:
	if target_list is not None:
	encoder_input, encoder_padding_mask = self.speech_encoder_prenet(source, require_feat_pen=True, target_list=target_list, padding_mask=padding_mask, mask=mask)
	encoder_input, features_pen, mask_indices, target_list = encoder_input
	else:
	encoder_input, encoder_padding_mask = self.speech_encoder_prenet(source, padding_mask=padding_mask, mask=self.training)
	# shuffle a batch of inputs of encoder
	if self.training and hasattr(self.args, "sid_shuffle_encoder_input") and getattr(self.args, "sid_shuffle_encoder_input", False):
	shuffle_index = torch.randperm(encoder_padding_mask.size(1), device=encoder_padding_mask.device)
	encoder_input = torch.index_select(encoder_input, 1, shuffle_index)
	encoder_padding_mask = torch.index_select(encoder_padding_mask, 1, shuffle_index)
	if getattr(self.args, "sid_encoder_cls", None) == "encoder":
	prev_output_tokens = torch.zeros_like(prev_output_tokens)
	encoder_input, encoder_padding_mask = self._integrate_with_speaker_cls(prev_output_tokens, encoder_input, encoder_padding_mask)

	# Encoder: T x B x C
	encoder_output = self.encoder(encoder_input, encoder_padding_mask, tgt_layer=tgt_enc_layer)

	if task_name is not None and task_name == 'speech_pretrain' and feature_only:
	return encoder_output["encoder_out"][0].transpose(0, 1)

	if task_name is not None and task_name == 's2c':
	if self.args.sid_pooling_layer == "encoder":
	return self.speaker_decoder_postnet(encoder_output["encoder_out"][0].transpose(0, 1).mean(1), sid_target), None
	elif self.args.sid_pooling_layer == "encoder-cls":
	return self.speaker_decoder_postnet(encoder_output["encoder_out"][0].transpose(0, 1)[:,0], sid_target), None
	elif self.args.sid_pooling_layer == "encoder-speaker" or getattr(self.args, "sid_decoder_speaker", False):
	return self.speaker_decoder_postnet(encoder_output["encoder_out"][0].transpose(0, 1), sid_target), None

	if target_list is not None:
	hubert_results = self.hubert_layer(
	encoder_output["encoder_out"][0].transpose(0, 1),
	encoder_padding_mask,
	mask_indices,
	target_list
	)

	hubert_results['features_pen'] = features_pen

	if "decoder_input" in encoder_output and encoder_output["decoder_input"][0] is not None:
	# Change the encoder output to decoder input once set unb-enc-layer
	encoder_output["encoder_out"] = encoder_output["decoder_input"]

	if self.use_codebook:
	q = self.quantizer(encoder_output["encoder_out"][0].transpose(0, 1))

	# q["x"]: B x T x C
	# Sample indexs according to the codebook prob
	random_idx = torch.randperm(q["x"].size(1))[:int(q["x"].size(1) * self.codebook_prob)]
	# Make weight for q
	q_w = q["x"].new_zeros(q["x"].size(1))
	q_w[random_idx] = 1.0
	# Combine quantized codes and encoder output
	encoder_output["encoder_out"][0] = (
	q_w.view(-1, 1) * q["x"] + (- q_w + 1).view(-1, 1) * encoder_output["encoder_out"][0].transpose(0, 1)
	).transpose(0, 1)

	# encoder_output["encoder_out"][0] = q["x"].transpose(0, 1)
	if output_type == 'speech':
	hubert_results["prob_perplexity"] = q["prob_perplexity"]
	hubert_results["code_perplexity"] = q["code_perplexity"]
	hubert_results["num_vars"] = q["num_vars"]
	hubert_results["temp"] = q["temp"]
	elif output_type == 'text':
	codebook_out["prob_perplexity"] = q["prob_perplexity"]
	codebook_out["code_perplexity"] = q["code_perplexity"]
	codebook_out["num_vars"] = q["num_vars"]
	codebook_out["temp"] = q["temp"]

	if only_hubert and target_list is not None:
	return hubert_results, None

	if only_ctc and task_name is not None and task_name == "s2t":
	return None, encoder_output
	elif not self.training and prev_output_tokens is None and task_name == "s2t" and task_name is not None:
	return encoder_output

	# Decoder Prenet
	if output_type == 'text':
	# _ is the incremental state
	prev_output_tokens, tgt_mask, _ = self.text_decoder_prenet(prev_output_tokens)
	if task_name is not None and task_name == 's2c':
	prev_output_tokens = torch.zeros_like(prev_output_tokens)
	else:
	# integrate speaker embedding
	if self.spk_embed_integration_type == "pre" and self.spk_embed_dim is not None:
	# Decoder Prenet
	prev_output_tokens, tgt_mask = self.speech_decoder_prenet(prev_output_tokens, tgt_lengths, spkembs)
	else:
	if self.spk_embed_dim is not None:
	encoder_output["encoder_out"] = [self._integrate_with_spk_embed(
	encoder_output["encoder_out"][0].transpose(0, 1), spkembs
	).transpose(0, 1)]

	prev_output_tokens, tgt_mask = self.speech_decoder_prenet(prev_output_tokens, tgt_lengths)

	# BART Sequence Classification: cat <pad> + feature before decoder
	if task_name is not None and task_name == 's2c' and self.args.sid_pooling_layer == "decoder-las":
	decoder_feat_input, decoder_feat_mask = self.speech_decoder_prenet(src_tokens, src_lengths)
	prev_output_tokens, tgt_mask = self._integrate_with_speaker_cls((prev_output_tokens, tgt_mask), decoder_feat_input, decoder_feat_mask, cls_first=False)

	# SE predict masking to corresponding inputs and source speech replaces the prev_output_tokens as the input of decoder
	if task_name is not None and task_name == "s2s" and getattr(self.args, "se_decoder_input", "previous_target") == "source":
	prev_output_tokens, tgt_mask = self.speech_decoder_prenet(src_tokens, src_lengths)

	# Decoder
	decoder_output, extra = self.decoder(prev_output_tokens, tgt_mask, encoder_output,
	full_context_alignment=getattr(self.args, "decoder_full_context_alignment", False),
	alignment_layer=(-1 if target_list is None and output_type == 'speech' else None))
	# Decoder Postnet
	if task_name is not None and task_name == 's2c':
	if not getattr(self.args, "sid_t5_postnet", False):
	if self.args.sid_pooling_layer == "decoder":
	return self.speaker_decoder_postnet(decoder_output.mean(1), sid_target), None
	elif self.args.sid_pooling_layer == "decoder-las":
	indices = (tgt_mask.eq(False).float().sum(1) - 1.0).type(torch.int64)
	indices = indices.unsqueeze(1).unsqueeze(2).expand(-1, -1, decoder_output.size(2))
	return self.speaker_decoder_postnet(decoder_output.gather(1, indices), sid_target), None
	else:
	return (self.text_decoder_postnet(decoder_output), None), encoder_output

	# SE predict: masking, target, delta. Ensure reduction factor 1
	if task_name is not None and task_name == 's2s' and getattr(self.args, "se_predict", None) is not None:
	assert self.reduction_factor == 1, f"{self.reduction_factor} != 1"
	before_outs, after_outs, logits = self.speech_decoder_postnet(decoder_output)
	se_predict = getattr(self.args, "se_predict")
	if se_predict == "masking":
	before_outs = torch.sigmoid(before_outs) * src_tokens
	after_outs = torch.sigmoid(after_outs) * src_tokens
	return before_outs, after_outs, logits, extra['attn'][0]
	elif se_predict == "target":
	return before_outs, after_outs, logits, extra['attn'][0]
	elif se_predict == "delta":
	before_outs = before_outs - src_tokens
	after_outs = after_outs - src_tokens
	return before_outs, after_outs, logits, extra['attn'][0]
	else:
	raise ValueError(f"{se_predict} not in [masking, target, delta]")

	if task_name is not None and task_name == 's2t':
	#return self.text_decoder_postnet(decoder_output), None
	return (self.text_decoder_postnet(decoder_output), None), encoder_output
	if output_type == 'text':
	return (self.text_decoder_postnet(decoder_output), None), codebook_out, encoder_output
	else:
	if target_list is not None:
	return hubert_results, (self.speech_decoder_postnet(decoder_output) + (extra['attn'][0],))
	else:
	return self.speech_decoder_postnet(decoder_output) + (extra['attn'][0],)

	def _integrate_with_speaker_cls(self, pad_input, encoder_input, encoder_padding_mask=None, cls_first=True):
	"""
	encoder_input: [B, T, C]
	encoder_padding_mask: [B, T]
	"""
	if hasattr(self, "text_decoder_prenet"):
	if isinstance(pad_input, tuple):
	repeat_cls_vector, repeat_cls_mask = pad_input
	else:
	repeat_cls_vector, repeat_cls_mask, _ = self.text_decoder_prenet(pad_input)

	if encoder_padding_mask is not None:
	bsz = encoder_input.size(0)
	tsz = encoder_input.size(1)
	encoder_padding_mask = encoder_input.new_zeros((bsz, tsz)) == 1.0
	if repeat_cls_mask is None:
	mask_size = (encoder_padding_mask.size(0), 1)
	mask_type = encoder_padding_mask.dtype
	repeat_cls_mask = encoder_padding_mask.new_zeros(mask_size) == 1.0
	ret_encoder_padding_mask = torch.cat([repeat_cls_mask, encoder_padding_mask], dim=1)

	if cls_first:
	ret_encoder_input = torch.cat([repeat_cls_vector, encoder_input], dim=1)
	else:
	ret_encoder_input = torch.cat([encoder_input, encoder_input[:,-1:,:]], dim=1)
	mask_size = (encoder_padding_mask.size(0), 1)
	mask_type = encoder_padding_mask.dtype
	repeat_cls_mask_ = encoder_padding_mask.new_ones(mask_size) == 1.0
	encoder_padding_mask_ = torch.cat([encoder_padding_mask, repeat_cls_mask_], dim=1)
	indices = encoder_padding_mask.eq(False).float().sum(1).type(torch.int64).unsqueeze(1)
	indices_mask = torch.zeros_like(ret_encoder_padding_mask).scatter(1, indices, 1.0)
	ret_encoder_input = ret_encoder_input * (1.0 - encoder_padding_mask_.type(ret_encoder_input.dtype).unsqueeze(2)) \
	+ repeat_cls_vector * indices_mask.type(repeat_cls_vector.dtype).unsqueeze(2)

	return ret_encoder_input, ret_encoder_padding_mask

	def _integrate_with_spk_embed(self, hs, spembs):
	"""Integrate speaker embedding with hidden states.
	Args:
	hs (Tensor): Batch of hidden state sequences (B, Tmax, adim).
	spembs (Tensor): Batch of speaker embeddings (B, spk_embed_dim).
	Returns:
	Tensor: Batch of integrated hidden state sequences (B, Tmax, adim)
	"""
	if self.spk_embed_integration_type == "add":
	# apply projection and then add to hidden states
	spembs = self.projection(F.normalize(spembs))
	hs = hs + spembs.unsqueeze(1)
	elif self.spk_embed_integration_type == "concat":
	# concat hidden states with spk embeds and then apply projection
	spembs = F.normalize(spembs).unsqueeze(1).expand(-1, hs.size(1), -1)
	hs = self.projection(torch.cat([hs, spembs], dim=-1))
	else:
	raise NotImplementedError("support only add or concat.")

	return hs

	def load_state_dict(
	self,
	state_dict,
	strict=True,
	model_cfg=None,
	args=None,
	):
	"""NOT STRICT Copies parameters and buffers from state_dict into this module and
	its descendants.

	Overrides the method in :class:`nn.Module`. Compared with that method
	this additionally "upgrades" state_dicts from old checkpoints.
	"""
	# self.prune_modules(model_cfg.modules_filter)
	model_dict_size = self.text_decoder_postnet.output_projection.out_features
	ckpt_dict_size = state_dict["text_decoder_postnet.output_projection.weight"].size(0)
	if model_dict_size != ckpt_dict_size:
	# reset dictionary-related modules, such as embedding table and encoder ctc embed
	logger.warn(f"not equal dictionary between model and checkpoint: {model_dict_size} vs {ckpt_dict_size}")
	logger.info(f"reset model dictionary with size of {model_dict_size}")
	removed_keys = [
	key for key in state_dict.keys() if any(
	key.startswith(previ) for previ in [
	"encoder.proj", "text_encoder_prenet", "text_decoder_prenet", "text_decoder_postnet"
	]
	)
	]
	for key in removed_keys:
	state_dict.pop(key, None)
	logger.info(f"removed loaded checkpoint: {key}")
	for m in self._modules.keys():
	m_state_dict = {
	key.replace(f"{m}.", ""): value for key, value in state_dict.items() if key.startswith(f"{m}.")
	}
	if hasattr(self, m):
	self._modules[m].load_state_dict(m_state_dict, False)
	return self

	def prune_modules(self, modules_filter=None):
	"""Prune unused modules for specific tasks."""
	if modules_filter is None:
	return
	elif modules_filter == "s2c":
	if hasattr(self, "text_encoder_prenet"): del self.text_encoder_prenet
	if hasattr(self, "speech_decoder_prenet") and getattr(self.args, "sid_pooling_layer", None) != "decoder-las":
	del self.speech_decoder_prenet
	if hasattr(self, "speech_decoder_postnet"): del self.speech_decoder_postnet
	if hasattr(self, "text_decoder_postnet"): del self.text_decoder_postnet
	if hasattr(self, "speech_encoder_postnet"): del self.speech_encoder_postnet
	if hasattr(self.encoder, "proj"): self.encoder.proj = None
	if hasattr(self, "projection"): del self.projection
	if hasattr(self, "quantizer"): del self.quantizer
	if getattr(self.args, "sid_pooling_layer", "decoder").startswith("encoder") or getattr(self.args, "sid_decoder_speaker", False):
	if hasattr(self.decoder, "dropout_module"): del self.decoder.dropout_module
	if hasattr(self.decoder, "layers"): del self.decoder.layers
	if hasattr(self.decoder, "layer_norm"): del self.decoder.layer_norm
	if hasattr(self, "text_decoder_prenet"): del self.text_decoder_prenet
	elif modules_filter == "s2s":
	if hasattr(self, "speaker_decoder_postnet"): del self.speaker_decoder_postnet
	if hasattr(self, "text_encoder_prenet"): del self.text_encoder_prenet
	if hasattr(self, "text_decoder_prenet"): del self.text_decoder_prenet
	if hasattr(self, "text_decoder_postnet"): del self.text_decoder_postnet
	if hasattr(self, "speech_encoder_postnet"): del self.speech_encoder_postnet
	if hasattr(self.encoder, "proj"): self.encoder.proj = None
	if hasattr(self, "projection"): del self.projection
	if hasattr(self, "quantizer"): del self.quantizer
	elif modules_filter == "t2s":
	if hasattr(self, "speaker_decoder_postnet"): del self.speaker_decoder_postnet
	if hasattr(self, "speech_encoder_prenet"): del self.speech_encoder_prenet
	if hasattr(self, "text_decoder_prenet"): del self.text_decoder_prenet
	if hasattr(self, "text_decoder_postnet"): del self.text_decoder_postnet
	if hasattr(self, "speech_encoder_postnet"): del self.speech_encoder_postnet
	if hasattr(self.encoder, "proj"): self.encoder.proj = None
	if hasattr(self, "projection"): del self.projection
	if hasattr(self, "quantizer"): del self.quantizer
	elif modules_filter == "s3prl":
	# remain the encoder and the pre/post net
	if hasattr(self.decoder, "dropout_module"): del self.decoder.dropout_module
	if hasattr(self.decoder, "layers"): del self.decoder.layers
	if hasattr(self.decoder, "layer_norm"): del self.decoder.layer_norm
	if hasattr(self, "speaker_decoder_postnet"): del self.speaker_decoder_postnet
	if hasattr(self, "text_decoder_prenet"): del self.text_decoder_prenet
	if hasattr(self, "text_decoder_postnet"): del self.text_decoder_postnet
	if hasattr(self, "speech_decoder_prenet"): del self.speech_decoder_prenet
	if hasattr(self, "speech_decoder_postnet"): del self.speech_decoder_postnet
	if hasattr(self, "speech_encoder_postnet"): del self.speech_encoder_postnet
	if hasattr(self.encoder, "proj"): self.encoder.proj = None
	if hasattr(self, "projection"): del self.projection
	if hasattr(self, "quantizer"): del self.quantizer

	def forward_encoder_torchscript(self, net_input: Dict[str, Tensor]):
	"""A TorchScript-compatible version of forward.

	Encoders which use additional arguments may want to override
	this method for TorchScript compatibility.
	"""
	if torch.jit.is_scripting():
	return self.forward_encoder(
	source=net_input["source"],
	padding_mask=net_input["padding_mask"]
	)
	else:
	return self.forward_encoder_non_torchscript(net_input)

	@torch.jit.unused
	def forward_encoder_non_torchscript(self, net_input: Dict[str, Tensor]):
	encoder_input = {
	k: v for k, v in net_input.items() if k != "prev_output_tokens" and k != "task_name"
	}
	return self.forward_encoder(**encoder_input)

	def forward_encoder(self, source, padding_mask=None):
	# Encoder Prenet
	encoder_input, encoder_padding_mask = self.speech_encoder_prenet(source, padding_mask=padding_mask, mask=False)

	# Encoder
	encoder_output = self.encoder(encoder_input, encoder_padding_mask)

	return encoder_output

	def forward_text_encoder(self, src_tokens):
	# Text Encoder Prenet
	encoder_input, encoder_padding_mask = self.text_encoder_prenet(src_tokens)

	# Encoder
	encoder_output = self.encoder(encoder_input, encoder_padding_mask)

	return encoder_output

	def forward_decoder(self, tokens, encoder_out, incremental_state):
	# Decoder Prenet
	prev_output_tokens, tgt_mask, incremental_state = self.text_decoder_prenet(tokens, incremental_state)

	# Decoder
	decoder_output, extra = self.decoder(
	prev_output_tokens,
	tgt_mask,
	encoder_out=encoder_out,
	incremental_state=incremental_state,
	)

	# Decoder Postnet
	return self.text_decoder_postnet(decoder_output), extra

	def set_num_updates(self, num_updates):
	"""Set the number of parameters updates."""
	super().set_num_updates(num_updates)
	self.num_updates = num_updates

	def generate_class(self, source, prev_output_tokens, **kwargs):
	encoder_out = self.forward_encoder(source, padding_mask=kwargs["padding_mask"])

	prev_output_tokens, tgt_mask, _ = self.text_decoder_prenet(prev_output_tokens, {})
	prev_output_tokens = torch.zeros_like(prev_output_tokens) # s2c use zero vector as [CLS]

	decoder_output, extra = self.decoder(
	prev_output_tokens,
	tgt_mask,
	encoder_out=encoder_out,
	)

	decoder_out, embed = self.speaker_decoder_postnet(decoder_output.mean(1))

	pred_class = decoder_out.argmax(1)
	return pred_class

	def generate_speech(self, source=None, src_tokens=None, spkembs=None, **kwargs):
	assert source is not None or src_tokens is not None

	threshold = kwargs.get("threshold", 0.5)
	minlenratio = kwargs.get("threshold", 0.0)

	if source is None:
	assert src_tokens.size(0) == 1
	encoder_out = self.forward_text_encoder(src_tokens)
	maxlenratio = kwargs.get("threshold", 20.0)
	else:
	assert source.size(0) == 1
	encoder_out = self.forward_encoder(source, padding_mask=kwargs["padding_mask"])
	maxlenratio = kwargs.get("threshold", 10.0)

	if spkembs is not None and self.spk_embed_integration_type != "pre":
	encoder_out["encoder_out"] = [self._integrate_with_spk_embed(
	encoder_out["encoder_out"][0].transpose(0, 1), spkembs
	).transpose(0, 1)]
	spkembs = None

	maxlen = int(encoder_out["encoder_out"][0].size(0) * maxlenratio / self.reduction_factor)
	minlen = int(encoder_out["encoder_out"][0].size(0) * minlenratio / self.reduction_factor)

	idx = 0
	ys = encoder_out["encoder_out"][0].new_zeros(1, 1, self.speech_decoder_postnet.odim)
	outs, probs = [], []

	# forward decoder step-by-step
	if isinstance(self.decoder, FairseqIncrementalDecoder):
	incremental_states = {}
	else:
	incremental_states = None
	attns = []
	while True:
	# update index
	idx += 1
	# calculate output and stop prob at idx-th step
	decoder_in, _ = self.speech_decoder_prenet(ys, spkembs=spkembs)
	z, extra = self.decoder(decoder_in[:,-1:], None, encoder_out, incremental_states, alignment_layer=-1)
	outs += [self.speech_decoder_postnet.feat_out(z[0, -1]).view(self.reduction_factor, self.speech_decoder_postnet.odim)] # [(r, odim), ...]
	probs += [torch.sigmoid(self.speech_decoder_postnet.prob_out(z[0, -1]))] # [(r), ...]

	# update next inputs
	ys = torch.cat((ys, outs[-1][-1].view(1, 1, self.speech_decoder_postnet.odim)), dim=1) # (1, idx + 1, odim)
	attns.append(torch.stack([att_l[0] for att_l in extra['attn'][0]], dim=0))
	# check whether to finish generation
	if int(sum(probs[-1] >= threshold)) > 0 or idx >= maxlen:
	# check mininum length
	if idx < minlen:
	continue
	outs = (torch.cat(outs, dim=0).unsqueeze(0).transpose(1, 2)) # (L, odim) -> (1, L, odim) -> (1, odim, L)
	if self.speech_decoder_postnet.postnet is not None:
	outs = outs + self.speech_decoder_postnet.postnet(outs) # (1, odim, L)
	outs = outs.transpose(2, 1).squeeze(0) # (L, odim)
	probs = torch.cat(probs, dim=0)
	attn = torch.cat(attns, dim=2)
	break

	if outs.size(0) == maxlen:
	logging.warning("output length reaches maximum length")
	return outs, probs, attn


	@register_model_architecture(model_name="artst_transformer", arch_name="artst_transformer")
	def base_architecture(args):
	# Transformer
	args.bert_init = getattr(args, "bert_init", False)
	args.encoder_embed_dim = getattr(args, "encoder_embed_dim", 768)
	args.encoder_ffn_embed_dim = getattr(args, "encoder_ffn_embed_dim", 768 * 4)
	args.encoder_layers = getattr(args, "encoder_layers", 12)
	args.encoder_attention_heads = getattr(args, "encoder_attention_heads", 12)
	args.encoder_normalize_before = getattr(args, "encoder_normalize_before", False)
	args.decoder_embed_dim = getattr(args, "decoder_embed_dim", args.encoder_embed_dim)
	args.decoder_ffn_embed_dim = getattr(
	args, "decoder_ffn_embed_dim", args.encoder_ffn_embed_dim
	)
	args.decoder_layers = getattr(args, "decoder_layers", 6)
	args.decoder_attention_heads = getattr(args, "decoder_attention_heads", 12)
	args.decoder_normalize_before = getattr(args, "decoder_normalize_before", False)
	args.dropout = getattr(args, "dropout", 0.1)
	args.attention_dropout = getattr(args, "attention_dropout", args.dropout)
	args.activation_dropout = getattr(args, "activation_dropout", args.dropout)
	args.activation_fn = getattr(args, "activation_fn", "gelu")
	args.decoder_layerdrop = getattr(args, "decoder_layerdrop", 0.0)
	args.decoder_output_dim = getattr(
	args, "decoder_output_dim", args.decoder_embed_dim
	)
	args.decoder_input_dim = getattr(args, "decoder_input_dim", args.decoder_embed_dim)
	args.encoder_layerdrop = getattr(args, "encoder_layerdrop", 0)
	args.decoder_layerdrop = getattr(args, "decoder_layerdrop", 0)
	args.max_text_positions = getattr(args, "max_text_positions", DEFAULT_MAX_TEXT_POSITIONS)
	args.max_speech_positions = getattr(args, "max_speech_positions", DEFAULT_MAX_SPEECH_POSITIONS)

	# Espnet related, including prenet, postnet
	args.eprenet_conv_layers = getattr(args, "eprenet_conv_layers", 0)
	args.eprenet_conv_filts = getattr(args, "eprenet_conv_filts", 0)
	args.eprenet_conv_chans = getattr(args, "eprenet_conv_chans", 0)
	args.use_batch_norm = getattr(args, "use_batch_norm", True)
	args.eprenet_dropout_rate = getattr(args, "eprenet_dropout_rate", 0.0)
	args.enc_use_scaled_pos_enc = getattr(args, "enc_use_scaled_pos_enc", True)
	args.dec_use_scaled_pos_enc = getattr(args, "dec_use_scaled_pos_enc", True)
	args.postnet_layers = getattr(args, "postnet_layers", 5)
	args.postnet_chans = getattr(args, "postnet_chans", 256)
	args.postnet_filts = getattr(args, "postnet_filts", 5)
	args.postnet_dropout_rate = getattr(args, "postnet_dropout_rate", 0.5)
	args.dprenet_dropout_rate = getattr(args, "dprenet_dropout_rate", 0.5)
	args.dprenet_layers = getattr(args, "dprenet_layers", 2)
	args.dprenet_units = getattr(args, "dprenet_units", 256)
	args.initial_encoder_alpha = getattr(args, "initial_encoder_alpha", 1.0)
	args.initial_decoder_alpha = getattr(args, "initial_decoder_alpha", 1.0)
	args.spk_embed_integration_type = getattr(args, "spk_embed_integration_type", "pre")
	args.spk_embed_dim = getattr(args, "spk_embed_dim", 512)
	args.encoder_reduction_factor = getattr(args, "encoder_reduction_factor", 1)
	args.reduction_factor = getattr(args, "reduction_factor", 2)
	args.transformer_enc_positional_dropout_rate = getattr(args, "transformer_enc_positional_dropout_rate", 0.1)
	args.transformer_dec_positional_dropout_rate = getattr(args, "transformer_dec_positional_dropout_rate", 0.1)
	args.layer_norm_eps = getattr(args, "layer_norm_eps", 1e-5)
	args.no_scale_embedding = getattr(args, "no_scale_embedding", True)
	# Convolutional subsampler
	args.encoder_speech_prenet = getattr(args, "encoder_speech_prenet", "conv")
	args.conv_kernel_sizes = getattr(args, "conv_kernel_sizes", "5,5")
	args.conv_channels = getattr(args, "conv_channels", 1024)
	args.quant_noise_pq = getattr(args, "quant_noise_pq", 0)

	args.adaptive_softmax_cutoff = getattr(args, "adaptive_softmax_cutoff", None)
	args.adaptive_softmax_dropout = getattr(args, "adaptive_softmax_dropout", 0)
	args.no_token_positional_embeddings = getattr(
	args, "no_token_positional_embeddings", False
	)
	args.adaptive_input = getattr(args, "adaptive_input", False)
	args.decoder_learned_pos = getattr(args, "decoder_learned_pos", False)
	args.share_input_output_embed = getattr(args, "share_input_output_embed", False)
	args.share_ctc_embed = getattr(args, "share_ctc_embed", False)
	args.freeze_encoder_updates = getattr(args, "freeze_encoder_updates", 0)
	args.freeze_decoder_updates = getattr(args, "freeze_decoder_updates", 0)
	args.no_freeze_encoder_layer = getattr(args, "no_freeze_encoder_layer", None)

	## sid
	args.sid_embed_dim = getattr(args, "sid_embed_dim", 128)
	args.sid_pooling_layer = getattr(args, "sid_pooling_layer", "decoder")
	args.softmax_scale = getattr(args, "softmax_scale", 1)
	args.softmax_margin = getattr(args, "softmax_margin", 0)
	args.softmax_easy_margin = getattr(args, "softmax_easy_margin", False)
	args.modules_filter = getattr(args, "modules_filter", None)

	## Hubert
	args.conv_pos = getattr(args, "conv_pos", 128)
	args.conv_pos_groups = getattr(args, "conv_pos_groups", 16)
	args.target_glu = getattr(args, "target_glu", False)
	args.logit_temp = getattr(args, "logit_temp", 0.1)
	args.final_dim = getattr(args, "final_dim", 256)
	args.untie_final_proj = getattr(args, "untie_final_proj", True)
	args.feature_grad_mult = getattr(args, "feature_grad_mult", 0.1)
	args.use_sent_enc_layer = getattr(args, "use_sent_enc_layer", True)
	# hubert feature extractor
	args.extractor_mode = getattr(args, "extractor_mode", "default")
	args.conv_feature_layers = getattr(args, "conv_feature_layers", "[(512,10,5)] + [(512,3,2)] * 4 + [(512,2,2)] * 2")
	args.conv_bias = getattr(args, "conv_bias", False)
	# mask
	args.hubert_mask_length = getattr(args, "hubert_mask_length", 10)
	args.mask_prob = getattr(args, "mask_prob", 0.0)
	args.mask_selection = getattr(args, "mask_selection", "static")
	args.mask_other = getattr(args, "mask_other", 0)
	args.no_mask_overlap = getattr(args, "no_mask_overlap", False)
	args.mask_min_space = getattr(args, "mask_min_space", 1)
	# channel mask
	args.mask_channel_length = getattr(args, "mask_channel_length", 10)
	args.mask_channel_prob = getattr(args, "mask_channel_prob", 0.0)
	args.mask_channel_selection = getattr(args, "mask_channel_selection", "static")
	args.mask_channel_other = getattr(args, "mask_channel_other", 0)
	args.no_mask_channel_overlap = getattr(args, "no_mask_channel_overlap", False)
	args.mask_channel_min_space = getattr(args, "mask_channel_min_space", 1)
	# loss computation
	args.skip_masked = getattr(args, "skip_masked", False)
	args.skip_nomask = getattr(args, "skip_nomask", False)
	# conv Pos
	args.use_conv_pos = getattr(args, "use_conv_pos", False)
	args.use_sinc_pos = getattr(args, "use_sinc_pos", False)

	# codebook
	args.use_codebook = getattr(args, "use_codebook", False)
	args.latent_vars = getattr(args, "latent_vars", 100)
	args.latent_groups = getattr(args, "latent_groups", 2)
	args.latent_dim = getattr(args, "latent_dim", 0)
	args.latent_temp = getattr(args, "latent_temp", (2, 0.5, 0.999995))
	args.quantizer_depth = getattr(args, "quantizer_depth", 1)
	args.quantizer_factor = getattr(args, "quantizer_factor", 3)
	args.codebook_prob = getattr(args, "codebook_prob", 0.5)

	# Relative pos embed
	args.relative_position_embedding = getattr(args, "relative_position_embedding", False)
	args.num_buckets = getattr(args, "num_buckets", 320)
	args.max_distance = getattr(args, "max_distance", 1280)
	args.encoder_max_relative_position = getattr(args, "encoder_max_relative_position", 160)
	args.decoder_max_relative_position = getattr(args, "decoder_max_relative_position", 160)

	@register_model_architecture("artst_transformer", "artst_transformer_base")
	def artst_transformer_base(args):
	args.use_conv_pos = getattr(args, "use_conv_pos", True)
	args.use_sinc_pos = getattr(args, "use_sinc_pos", True)
	args.layernorm_embedding = getattr(args, "layernorm_embedding", False)
	args.encoder_normalize_before = getattr(args, "encoder_normalize_before", False)
	args.decoder_normalize_before = getattr(args, "decoder_normalize_before", False)
	args.layer_norm_first = getattr(args, "layer_norm_first", False)
	args.relative_position_embedding = getattr(args, "relative_position_embedding", True)
	args.dropout = getattr(args, "dropout", 0.1)
	args.activation_dropout = getattr(args, "activation_dropout", 0.0)
	args.attention_dropout = getattr(args, "attention_dropout", 0.1)
	args.encoder_layerdrop = getattr(args, "encoder_layerdrop", 0.05)
	args.decoder_layerdrop = getattr(args, "decoder_layerdrop", 0.05)
	args.mask_prob = getattr(args, "mask_prob", 0.80)
	base_architecture(args)

	@register_model_architecture("artst_transformer", "artst_transformer_large")
	def artst_transformer_large(args):
	args.use_conv_pos = getattr(args, "use_conv_pos", True)
	args.use_sinc_pos = getattr(args, "use_sinc_pos", True)
	args.encoder_normalize_before = getattr(args, "encoder_normalize_before", False)
	args.decoder_normalize_before = getattr(args, "decoder_normalize_before", True)
	args.layer_norm_first = getattr(args, "layer_norm_first", True)
	args.relative_position_embedding = getattr(args, "relative_position_embedding", True)
	args.dropout = getattr(args, "dropout", 0.0)
	args.activation_dropout = getattr(args, "activation_dropout", 0.0)
	args.attention_dropout = getattr(args, "attention_dropout", 0.0)
	args.encoder_layerdrop = getattr(args, "encoder_layerdrop", 0.0)
	args.decoder_layerdrop = getattr(args, "decoder_layerdrop", 0.0)
	args.encoder_embed_dim = getattr(args, "encoder_embed_dim", 1024)
	args.encoder_layers = getattr(args, "encoder_layers", 24)
	args.decoder_layers = getattr(args, "decoder_layers", 6)
	args.encoder_ffn_embed_dim = getattr(args, "encoder_ffn_embed_dim", 4096)
	args.encoder_attention_heads = getattr(args, "encoder_attention_heads", 16)
	args.decoder_attention_heads = getattr(args, "decoder_attention_heads", 16)
	args.feature_grad_mult = getattr(args, "feature_grad_mult", 1.0)
	args.extractor_mode = getattr(args, "extractor_mode", "layer_norm")
	args.final_dim = getattr(args, "final_dim", 768)
	args.mask_prob = getattr(args, "mask_prob", 0.80)
	base_architecture(args)

	@register_model_architecture("artst_transformer", "artst_transformer_base_asr")
	def artst_transformer_base_asr(args):
	args.use_conv_pos = getattr(args, "use_conv_pos", True)
	args.use_sinc_pos = getattr(args, "use_sinc_pos", True)
	args.encoder_normalize_before = getattr(args, "encoder_normalize_before", False)
	args.decoder_normalize_before = getattr(args, "decoder_normalize_before", False)
	args.layer_norm_first = getattr(args, "layer_norm_first", False)
	args.relative_position_embedding = getattr(args, "relative_position_embedding", True)
	args.dropout = getattr(args, "dropout", 0.1)
	args.activation_dropout = getattr(args, "activation_dropout", 0.1)
	args.attention_dropout = getattr(args, "attention_dropout", 0.1)
	args.feature_grad_mult = getattr(args, "feature_grad_mult", 0.0)
	args.encoder_layerdrop = getattr(args, "encoder_layerdrop", 0.1)
	args.decoder_layerdrop = getattr(args, "decoder_layerdrop", 0.1)
	args.mask_prob = getattr(args, "mask_prob", 0.75)
	args.mask_selection = getattr(args, "mask_selection", "static")
	args.mask_channel_length = getattr(args, "mask_channel_length", 64)
	args.mask_channel_prob = getattr(args, "mask_channel_prob", 0.5)
	args.mask_channel_selection = getattr(args, "mask_channel_selection", "static")
	args.max_text_positions = getattr(args, "max_text_positions", 600)
	base_architecture(args)