Upload 3 files

README.md

@@ -231,27 +231,6 @@ 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
@@ -346,25 +325,6 @@ 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")
@@ -380,7 +340,7 @@ with torch.no_grad():
         outputs = model.generate(
             input_ids=input_ids,
             attention_mask=attention_mask,
-            generation_config=generation_config
+            max_new_tokens=4096,
         )
 
         for message in processor.decode(outputs):
@@ -388,7 +348,6 @@ 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)
@@ -408,27 +367,6 @@ 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
@@ -495,25 +433,6 @@ 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")
@@ -529,7 +448,7 @@ with torch.no_grad():
         outputs = model.generate(
             input_ids=input_ids,
             attention_mask=attention_mask,
-            generation_config=generation_config
+            max_new_tokens=4096,
         )
 
         for message in processor.decode(outputs):
@@ -537,7 +456,6 @@ 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,6 +616,109 @@ 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(