Update orpheus-tts/engine_class.py

orpheus-tts/engine_class.py

@@ -86,33 +86,25 @@ class OrpheusModel:
             if voice not in self.engine.available_voices:
                 raise ValueError(f"Voice {voice} is not available for model {self.model_name}")
     
-    def _format_prompt(self, prompt, voice="tara", model_type="larger"):
-        if model_type == "smaller":
-            if voice:
-                return f"<custom_token_3>{prompt}[{voice}]<custom_token_4><custom_token_5>"
-            else:
-                return f"<custom_token_3>{prompt}<custom_token_4><custom_token_5>"
+    def _format_prompt(self, prompt, voice="Sophie", model_type="larger"):
+        # Use Kartoffel model format based on documentation
+        if voice:
+            full_prompt = f"{voice}: {prompt}"
         else:
-            if voice:
-                adapted_prompt = f"{voice}: {prompt}"
-                prompt_tokens = self.tokenizer(adapted_prompt, return_tensors="pt")
-                start_token = torch.tensor([[ 128259]], dtype=torch.int64)
-                end_tokens = torch.tensor([[128009, 128260, 128261, 128257]], dtype=torch.int64)
-                all_input_ids = torch.cat([start_token, prompt_tokens.input_ids, end_tokens], dim=1)
-                prompt_string = self.tokenizer.decode(all_input_ids[0])
-                return prompt_string
-            else:
-                prompt_tokens = self.tokenizer(prompt, return_tensors="pt")
-                start_token = torch.tensor([[ 128259]], dtype=torch.int64)
-                end_tokens = torch.tensor([[128009, 128260, 128261, 128257]], dtype=torch.int64)
-                all_input_ids = torch.cat([start_token, prompt_tokens.input_ids, end_tokens], dim=1)
-                prompt_string = self.tokenizer.decode(all_input_ids[0])
-                return prompt_string
+            full_prompt = prompt
+            
+        # Kartoffel model token format
+        start_token = torch.tensor([[128259]], dtype=torch.int64)
+        end_tokens = torch.tensor([[128009, 128260]], dtype=torch.int64)
+        
+        input_ids = self.tokenizer(full_prompt, return_tensors="pt").input_ids
+        modified_input_ids = torch.cat([start_token, input_ids, end_tokens], dim=1)
+        
+        prompt_string = self.tokenizer.decode(modified_input_ids[0])
+        return prompt_string
 
  
-
-
-    def generate_tokens_sync(self, prompt, voice=None, request_id="req-001", temperature=0.6, top_p=0.8, max_tokens=1200, stop_token_ids = [128258], repetition_penalty=1.3):
+    def generate_tokens_sync(self, prompt, voice=None, request_id="req-001", temperature=0.6, top_p=0.95, max_tokens=4000, stop_token_ids = [128258], repetition_penalty=1.1):
         prompt_string = self._format_prompt(prompt, voice)
         print(f"DEBUG: Original prompt: {prompt}")
         print(f"DEBUG: Formatted prompt: {prompt_string}")