fix: fix the bug in the continuation mode of moss_tts_local; update the readme of moss_tts_local.

README.md

@@ -231,6 +231,27 @@ torch.backends.cuda.enable_flash_sdp(True)
 torch.backends.cuda.enable_mem_efficient_sdp(True)
 torch.backends.cuda.enable_math_sdp(True)
 
+class DelayGenerationConfig(GenerationConfig):
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        self.layers = kwargs.get("layers", [{} for _ in range(32)])
+        self.do_samples = kwargs.get("do_samples", None)
+        self.n_vq_for_inference = 32
+        
+def initial_config(tokenizer, model_name_or_path):
+    generation_config = DelayGenerationConfig.from_pretrained(model_name_or_path)
+    generation_config.pad_token_id = tokenizer.pad_token_id
+    generation_config.eos_token_id = 151653
+    generation_config.max_new_tokens = 1000000
+    generation_config.temperature = 1.0
+    generation_config.top_p = 0.95
+    generation_config.top_k = 100
+    generation_config.repetition_penalty = 1.1
+    generation_config.use_cache = True
+    generation_config.do_sample = False
+    return generation_config
+
+
 pretrained_model_name_or_path = "OpenMOSS-Team/MOSS-TTS-Local-Transformer"
 device = "cuda" if torch.cuda.is_available() else "cpu"
 dtype = torch.bfloat16 if device == "cuda" else torch.float32
@@ -325,6 +346,25 @@ model = AutoModel.from_pretrained(
 ).to(device)
 model.eval()
 
+generation_config = initial_config(processor.tokenizer, pretrained_model_name_or_path)
+generation_config.n_vq_for_inference = model.channels - 1
+generation_config.do_samples = [True] * model.channels
+generation_config.layers = [
+    {
+        "repetition_penalty": 1.0, 
+        "temperature": 1.5, 
+        "top_p": 1.0, 
+        "top_k": 50
+    }
+] + [ 
+    {
+        "repetition_penalty": 1.1, 
+        "temperature": 1.0, 
+        "top_p": 0.95,
+        "top_k": 50
+    }
+] * (model.channels - 1) 
+
 batch_size = 1
 
 save_dir = Path(f"inference_root_moss_tts_local_transformer_generation")
@@ -340,7 +380,7 @@ with torch.no_grad():
         outputs = model.generate(
             input_ids=input_ids,
             attention_mask=attention_mask,
-            max_new_tokens=4096,
+            generation_config=generation_config
         )
 
         for message in processor.decode(outputs):
@@ -348,6 +388,7 @@ with torch.no_grad():
             out_path = save_dir / f"sample{sample_idx}.wav"
             sample_idx += 1
             torchaudio.save(out_path, audio.unsqueeze(0), processor.model_config.sampling_rate)
+
 ```
 
 ### Continuation + Voice Cloning (Prefix Audio + Text)
@@ -367,6 +408,27 @@ torch.backends.cuda.enable_flash_sdp(True)
 torch.backends.cuda.enable_mem_efficient_sdp(True)
 torch.backends.cuda.enable_math_sdp(True)
 
+class DelayGenerationConfig(GenerationConfig):
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        self.layers = kwargs.get("layers", [{} for _ in range(32)])
+        self.do_samples = kwargs.get("do_samples", None)
+        self.n_vq_for_inference = 32
+        
+def initial_config(tokenizer, model_name_or_path):
+    generation_config = DelayGenerationConfig.from_pretrained(model_name_or_path)
+    generation_config.pad_token_id = tokenizer.pad_token_id
+    generation_config.eos_token_id = 151653
+    generation_config.max_new_tokens = 1000000
+    generation_config.temperature = 1.0
+    generation_config.top_p = 0.95
+    generation_config.top_k = 100
+    generation_config.repetition_penalty = 1.1
+    generation_config.use_cache = True
+    generation_config.do_sample = False
+    return generation_config
+
+
 pretrained_model_name_or_path = "OpenMOSS-Team/MOSS-TTS-Local-Transformer"
 device = "cuda" if torch.cuda.is_available() else "cpu"
 dtype = torch.bfloat16 if device == "cuda" else torch.float32
@@ -433,6 +495,25 @@ model = AutoModel.from_pretrained(
 ).to(device)
 model.eval()
 
+generation_config = initial_config(processor.tokenizer, pretrained_model_name_or_path)
+generation_config.n_vq_for_inference = model.channels - 1
+generation_config.do_samples = [True] * model.channels
+generation_config.layers = [
+    {
+        "repetition_penalty": 1.0, 
+        "temperature": 1.5, 
+        "top_p": 1.0, 
+        "top_k": 50
+    }
+] + [ 
+    {
+        "repetition_penalty": 1.1, 
+        "temperature": 1.0, 
+        "top_p": 0.95,
+        "top_k": 50
+    }
+] * (model.channels - 1) 
+
 batch_size = 1
 
 save_dir = Path("inference_root_moss_tts_local_transformer_continuation")
@@ -448,7 +529,7 @@ with torch.no_grad():
         outputs = model.generate(
             input_ids=input_ids,
             attention_mask=attention_mask,
-            max_new_tokens=4096,
+            generation_config=generation_config
         )
 
         for message in processor.decode(outputs):
@@ -456,6 +537,7 @@ with torch.no_grad():
             out_path = save_dir / f"sample{sample_idx}.wav"
             sample_idx += 1
             torchaudio.save(out_path, audio.unsqueeze(0), processor.model_config.sampling_rate)
+
 ```
 
 

modeling_moss_tts.py

@@ -616,109 +616,6 @@ class MossTTSDelayModel(MosiTTSPretrainedModel, CustomMixin):
     def can_generate(self):
         return True
 
-    def _build_generation_config(
-        self,
-        generation_config: Optional[GenerationConfig] = None,
-        max_new_tokens: Optional[int] = None,
-        text_temperature: Optional[float] = None,
-        text_top_p: Optional[float] = None,
-        text_top_k: Optional[int] = None,
-        text_repetition_penalty: Optional[float] = None,
-        audio_temperature: Optional[float] = None,
-        audio_top_p: Optional[float] = None,
-        audio_top_k: Optional[int] = None,
-        audio_repetition_penalty: Optional[float] = None,
-        n_vq_for_inference: Optional[int] = None,
-    ) -> GenerationConfig:
-        config = copy.deepcopy(generation_config or self.generation_config)
-
-        text_temperature = 1.5 if text_temperature is None else float(text_temperature)
-        text_top_p = 1.0 if text_top_p is None else float(text_top_p)
-        text_top_k = 50 if text_top_k is None else int(text_top_k)
-        text_repetition_penalty = 1.0 if text_repetition_penalty is None else float(text_repetition_penalty)
-        audio_temperature = 1.0 if audio_temperature is None else float(audio_temperature)
-        audio_top_p = 0.95 if audio_top_p is None else float(audio_top_p)
-        audio_top_k = 50 if audio_top_k is None else int(audio_top_k)
-        audio_repetition_penalty = 1.1 if audio_repetition_penalty is None else float(audio_repetition_penalty)
-
-        text_do_sample = text_temperature > 0
-        if not text_do_sample:
-            text_temperature = 1.0
-        audio_do_sample = audio_temperature > 0
-        if not audio_do_sample:
-            audio_temperature = 1.0
-
-        if max_new_tokens is not None:
-            config.max_new_tokens = int(max_new_tokens)
-        elif getattr(config, "max_new_tokens", None) is None:
-            config.max_new_tokens = 100000 # about 2.2 hours , can be overridden by user input, you can set to a smaller value for faster generation during debugging
-
-        if getattr(config, "pad_token_id", None) is None:
-            config.pad_token_id = self.config.pad_token_id
-        config.eos_token_id = self.config.audio_end_token_id
-        config.use_cache = True
-        config.do_sample = text_do_sample or audio_do_sample
-
-        resolved_n_vq = self.channels - 1 if n_vq_for_inference is None else int(n_vq_for_inference)
-        resolved_n_vq = max(1, min(self.channels - 1, resolved_n_vq))
-        config.n_vq_for_inference = resolved_n_vq
-        config.do_samples = [text_do_sample] + [audio_do_sample] * (self.channels - 1)
-        config.layers = [
-            {
-                "repetition_penalty": text_repetition_penalty,
-                "temperature": text_temperature,
-                "top_p": text_top_p,
-                "top_k": text_top_k,
-            }
-        ] + [
-            {
-                "repetition_penalty": audio_repetition_penalty,
-                "temperature": audio_temperature,
-                "top_p": audio_top_p,
-                "top_k": audio_top_k,
-            }
-            for _ in range(self.channels - 1)
-        ]
-        return config
-
-    @torch.inference_mode()
-    def generate(
-        self,
-        input_ids: torch.LongTensor,
-        attention_mask: Optional[torch.Tensor] = None,
-        generation_config: Optional[GenerationConfig] = None,
-        max_new_tokens: Optional[int] = None,
-        text_temperature: Optional[float] = None,
-        text_top_p: Optional[float] = None,
-        text_top_k: Optional[int] = None,
-        text_repetition_penalty: Optional[int] = None,
-        audio_temperature: Optional[float] = None,
-        audio_top_p: Optional[float] = None,
-        audio_top_k: Optional[int] = None,
-        audio_repetition_penalty: Optional[float] = None,
-        n_vq_for_inference: Optional[int] = None,
-        **kwargs,
-    ):
-        resolved_generation_config = self._build_generation_config(
-            generation_config=generation_config,
-            max_new_tokens=max_new_tokens,
-            text_temperature=text_temperature,
-            text_top_p=text_top_p,
-            text_top_k=text_top_k,
-            text_repetition_penalty=text_repetition_penalty,
-            audio_temperature=audio_temperature,
-            audio_top_p=audio_top_p,
-            audio_top_k=audio_top_k,
-            audio_repetition_penalty=audio_repetition_penalty,
-            n_vq_for_inference=n_vq_for_inference,
-        )
-        return super().generate(
-            input_ids=input_ids,
-            attention_mask=attention_mask,
-            generation_config=resolved_generation_config,
-            **kwargs,
-        )
-
     # def tie_weights(self):
     #     ...
         # for i in range(self.config.channels):

processing_moss_tts.py

@@ -621,7 +621,7 @@ class MossTTSDelayProcessor(ProcessorMixin):
                 prefix_idx = audio_end_idx
 
             if truncation:
-                ...
+                raise RuntimeError("Truncation generation is not supported at present")
             else:
                 last_audio_end_idx = int(audio_end_indices[-1].item())
                 pad_codes = torch.full(