🐛 fix ensure_non_empty issues #99

lenML · Jul 12, 2024 · bfc3789 · bfc3789
1 parent ee07d82
commit bfc3789
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 1 deletion.
diff --git a/modules/ChatTTS/ChatTTS/core.py b/modules/ChatTTS/ChatTTS/core.py
@@ -647,6 +647,7 @@ def _infer_code(
             stream=stream,
             context=self.context,
             stream_chunk_size=params.stream_chunk_size,
+            ensure_non_empty=params.ensure_non_empty,
         )
 
         del emb, input_ids

diff --git a/modules/ChatTTS/ChatTTS/model/gpt.py b/modules/ChatTTS/ChatTTS/model/gpt.py
@@ -506,6 +506,10 @@ def generate(
 
                 del logits
 
+                if i == 0:
+                    # when i == 0, we want to ensure that the first token is not eos_token
+                    scores[:, eos_token] = 0
+
                 idx_next = torch.multinomial(scores, num_samples=1).to(finish.device)
 
                 if not infer_text:
@@ -561,7 +565,7 @@ def generate(
                         )
                         for result in new_gen:
                             yield result
-                    return
+                        return
 
                 del inputs_ids
                 inputs_ids = inputs_ids_tmp