Update orpheus-tts/engine_class.py

orpheus-tts/engine_class.py

@@ -14,9 +14,11 @@ class OrpheusModel:
         self.dtype = dtype
         self.engine_kwargs = engine_kwargs  # vLLM engine kwargs
         self.engine = self._setup_engine()
-        # Available voices for German Kartoffel model
-        if "german" in model_name.lower() or "kartoffel" in model_name.lower():
+        # Available voices based on model type
+        if "kartoffel" in model_name.lower():
             self.available_voices = ["Jakob", "Anton", "Julian", "Sophie", "Marie", "Mia"]
+        elif "3b-de-ft" in model_name.lower():
+            self.available_voices = ["jana", "thomas", "max"]
         else:
             # Original English voices as fallback
             self.available_voices = ["zoe", "zac", "jess", "leo", "mia", "julia", "leah", "tara"]
@@ -57,7 +59,7 @@ class OrpheusModel:
             #     "repo_id": "canopylabs/orpheus-tts-0.1-finetune-prod",
             # },
             "medium-3b":{
-                "repo_id": "canopylabs/orpheus-tts-0.1-finetune-prod",
+                "repo_id": "canopylabs/3b-de-ft-research_release",
             },
         }
         unsupported_models = ["nano-150m", "micro-400m", "small-1b"]
@@ -88,31 +90,60 @@ class OrpheusModel:
                 raise ValueError(f"Voice {voice} is not available for model {self.model_name}")
     
     def _format_prompt(self, prompt, voice="Jakob", model_type="larger"):
-        # Use Kartoffel model format based on documentation
-        if voice:
-            full_prompt = f"{voice}: {prompt}"
-        else:
-            full_prompt = prompt
+        # Unterschiedliche Formate für verschiedene Modelle
+        print(f"DEBUG: Model name for format check: {self.model_name}")
+        if "kartoffel" in self.model_name.lower():
+            print("DEBUG: Using Kartoffel format")
+            # Kartoffel model format
+            if voice:
+                full_prompt = f"{voice}: {prompt}"
+            else:
+                full_prompt = prompt
+                
+            # Kartoffel model token format - direkt die Token-IDs einfügen
+            start_token_id = 128259
+            end_token_ids = [128009, 128260]
             
-        # Kartoffel model token format - direkt die Token-IDs einfügen
-        start_token_id = 128259
-        end_token_ids = [128009, 128260]
-        
-        # Text tokenisieren
-        input_ids = self.tokenizer(full_prompt, return_tensors="pt").input_ids[0].tolist()
-        
-        # Token-IDs zusammenfügen
-        all_token_ids = [start_token_id] + input_ids + end_token_ids
-        
-        # Zurück zu String dekodieren - aber die speziellen Token-IDs bleiben erhalten
-        prompt_string = self.tokenizer.decode(all_token_ids, skip_special_tokens=False)
-        
-        return prompt_string
+            # Text tokenisieren
+            input_ids = self.tokenizer(full_prompt, return_tensors="pt").input_ids[0].tolist()
+            
+            # Token-IDs zusammenfügen
+            all_token_ids = [start_token_id] + input_ids + end_token_ids
+            
+            # Zurück zu String dekodieren
+            prompt_string = self.tokenizer.decode(all_token_ids, skip_special_tokens=False)
+            
+            return prompt_string
+        else:
+            # Original Orpheus format (für Canopy-Deutsch und English)
+            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>"
+            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
 
  
     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: Voice: {voice}")
+        print(f"DEBUG: Model name: {self.model_name}")
+        prompt_string = self._format_prompt(prompt, voice)
         print(f"DEBUG: Formatted prompt: {prompt_string}")
         
         sampling_params = SamplingParams(
@@ -171,12 +202,15 @@ class OrpheusModel:
             token_generator = self.generate_tokens_sync(**kwargs)
             print("DEBUG: Token generator created successfully")
             
-            # Verwende Kartoffel-Decoder für deutsche Modelle
-            if "german" in self.model_name.lower() or "kartoffel" in self.model_name.lower():
-                print("DEBUG: Using Kartoffel decoder for German model")
+            # Verwende Kartoffel-Decoder nur für Kartoffel-Modell, Original-Decoder für Canopy-Deutsch
+            if "kartoffel" in self.model_name.lower():
+                print("DEBUG: Using Kartoffel decoder for Kartoffel model")
                 audio_generator = tokens_decoder_kartoffel_sync(token_generator, self.tokenizer)
+            elif "3b-de-ft" in self.model_name.lower() or "german" in self.model_name.lower():
+                print("DEBUG: Using original decoder for Canopy German model")
+                audio_generator = tokens_decoder_sync(token_generator)
             else:
-                print("DEBUG: Using original decoder")
+                print("DEBUG: Using original decoder for English model")
                 audio_generator = tokens_decoder_sync(token_generator)
             
             print("DEBUG: Audio decoder called successfully")