fixes and compat (MoE-fying an existing model and retraining from there just ruins it after a second of audio...)

vall_e/config.py

@@ -184,7 +184,7 @@ class Model:
 			name.append(self.size)
 
 		if self.arch_type != "transformer":
-			if self.experts:
+			if self.experts > 1:
 				name.append(f'{self.experts}x'+self.arch_type.replace("/", "-"))
 			else:
 				name.append(self.arch_type.replace("/", "-"))

vall_e/engines/__init__.py

@@ -13,6 +13,7 @@ from .base import Engines, TrainFeeder, default_feeder, Engine as _Engine
 from ..models import get_models
 from ..utils import wrapper as ml
 import torch
+import re
 
 deepspeed_available = False
 try:
@@ -90,6 +91,22 @@ def load_engines():
 			if "module" in state:
 				state = state["module"]
 
+			# maintain compat if I change variable names
+			insert = {}
+			erase = []
+
+			for k in state.keys():
+				key = re.sub(r'^retnet\.', "model.", k)
+				if k != key:
+					insert[key] = state[k]
+					erase.append(k)
+	
+			for k in insert.keys():
+				state[k] = insert[k]
+
+			for k in erase:
+				del state[k]
+
 			model.load_state_dict(state, strict=cfg.trainer.strict_loading)
 
 		# deepspeed inferencing

vall_e/inference.py

@@ -175,11 +175,17 @@ class TTS():
 		mirostat_eta=0.1,
 		out_path=None
 	):
+		lines = text.split("\n")
+
+		wavs = []
+		sr = None
+
+		for line in lines:
 			if out_path is None:
 				out_path = f"./data/{cfg.start_time}.wav"
 
 			prom = self.encode_audio( references, trim_length=input_prompt_length )
-		phns = self.encode_text( text, language=language )
+			phns = self.encode_text( line, language=language )
 			lang = self.encode_lang( language )
 
 			prom = to_device(prom, self.device).to(torch.int16)
@@ -209,6 +215,7 @@ class TTS():
 				)
 
 			wav, sr = qnt.decode_to_file(resps_list[0], out_path, device=self.device)
+			wavs.append(wav)
 		
-		return (wav, sr)
+		return (torch.concat(wavs, dim=-1), sr)
 		

vall_e/models/ar_nar.py

@@ -327,7 +327,6 @@ def example_usage():
 	proms_list = proms_list[:1]
 	resps_list = resps_list[:1]
 
-	"""
 	kwargs = {
 		'n_tokens': 1024,
 		'd_model': 1024, # 256, # 1024, # 1536
@@ -343,6 +342,7 @@ def example_usage():
 		'n_layers': 12,
 		'n_experts': 8,
 	}
+	"""
 	
 	"""
 	try:

vall_e/models/base.py

@@ -308,7 +308,7 @@ class Base(nn.Module):
 					num_experts_per_tok=min(2, n_experts),
 				))
 		elif self.arch_type == "retnet":
-			self.model = RetNetDecoder(RetNetConfig(
+			kwargs = dict(
 				vocab_size=n_resp_tokens,
 				decoder_embed_dim=d_model,
 				decoder_value_embed_dim =d_model * 2,
@@ -328,14 +328,18 @@ class Base(nn.Module):
 				decoder_normalize_before=True,
 
 				rotary_embedding_base=self.rotary_embedding_base, # 10000
+			)
 
-				# MoE
+			if n_experts > 1:
-				use_xmoe=n_experts>1,
+				kwargs.update(dict(
+					use_xmoe=True,
 					moe_freq=1,
 					moe_expert_count=n_experts,
 					moe_gating_use_fp32=False,
 				))
 
+			self.model = RetNetDecoder(RetNetConfig(**kwargs))
+
 		self.classifier = nn.Linear(d_model, n_resp_tokens)
 
 		self.accuracy_metric = MulticlassAccuracy(
@@ -422,7 +426,7 @@ class Base(nn.Module):
 		elif self.arch_type == "retnet":
 			# pass our inputs through the RetNet
 			x, _ = self.model(x, incremental_state=state, token_embeddings=x, features_only=True)
-			if _ is not None and "l_aux" in _:
+			if _ is not None and "l_aux" in _ and self.n_experts > 1:
 				aux_loss = torch.sum(torch.stack([ t for t in _["l_aux"] if t is not None])) * 0.001
 		# output projection layer with masking
 		x = self.classifier(x) * m