Update handler.py

handler.py

@@ -1,107 +1,199 @@
 import torch
-from transformers import AutoProcessor, AutoModelForTextToSpeech
+import numpy as np
 import soundfile as sf
 import io
-import numpy as np
 import os
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from snac import SNAC # Assuming SNAC is installed via requirements.txt
+
+# --- Helper Function (can be outside or inside the class) ---
+def redistribute_codes_static(code_list):
+    """ Reorganizes the flattened token list into three separate layers for SNAC. """
+    layer_1, layer_2, layer_3 = [], [], []
+    num_groups = len(code_list) // 7 # Use floor division
 
+    for i in range(num_groups):
+        idx = 7 * i
+        try:
+            layer_1.append(code_list[idx])
+            layer_2.append(code_list[idx + 1] - 4096)
+            layer_3.append(code_list[idx + 2] - (2 * 4096))
+            layer_3.append(code_list[idx + 3] - (3 * 4096))
+            layer_2.append(code_list[idx + 4] - (4 * 4096))
+            layer_3.append(code_list[idx + 5] - (5 * 4096))
+            layer_3.append(code_list[idx + 6] - (6 * 4096))
+        except IndexError:
+            print(f"Warning: Index out of range during code redistribution at group {i}. Code list length: {len(code_list)}")
+            break # Stop processing if indices go out of bounds
+
+    # Ensure tensors are created even if empty
+    codes = [
+        torch.tensor(layer_1 or [0]).unsqueeze(0).long(), # Use long dtype, provide default if empty
+        torch.tensor(layer_2 or [0]).unsqueeze(0).long(),
+        torch.tensor(layer_3 or [0]).unsqueeze(0).long()
+    ]
+    # Remove the default [0] if lists were actually populated
+    if layer_1: codes[0] = torch.tensor(layer_1).unsqueeze(0).long()
+    if layer_2: codes[1] = torch.tensor(layer_2).unsqueeze(0).long()
+    if layer_3: codes[2] = torch.tensor(layer_3).unsqueeze(0).long()
+
+    return codes
+
+# --- Endpoint Handler Class ---
 class EndpointHandler():
     def __init__(self, path=""):
         """
-        Initializes the handler. Loads the model and processor.
-        'path' is the directory where the model files are located.
+        Initializes the handler. Loads both Orpheus LLM and SNAC Vocoder.
+        'path' points to the directory containing the Orpheus model files specified in the endpoint config.
         """
         self.device = "cuda" if torch.cuda.is_available() else "cpu"
         print(f"Using device: {self.device}")
 
-        # Define the model path explicitly if needed, or rely on 'path'
-        model_path = path if path else os.getenv("HF_MODEL_DIR", ".") # HF Endpoints provide model dir via env var or path arg
+        # Define Model Names/Paths
+        # Orpheus LLM path is determined by the endpoint configuration ('path' variable)
+        orpheus_model_path = path if path else "hypaai/Hypa_Orpheus-3b-0.1-ft-unsloth-merged_16bit"
+        snac_model_name = "hubertsiuzdak/snac_24khz"
+
+        # Define Special Token IDs (matching your script)
+        self.start_human_token_id = 128259
+        self.end_text_token_id = 128009
+        self.end_human_token_id = 128260
+        # self.padding_token_id = 128263 # Not needed for single sequence generation
+        self.start_audio_token_id = 128257
+        self.end_audio_token_id = 128258
+        self.audio_code_offset = 128266
+
+        # Define sampling rate
+        self.sampling_rate = 24000
 
-        # --- Load Model Components ---
-        # Adjust these lines based on the specific classes your Orpheus model needs
-        # It might be AutoModelForSpeechSeq2Seq, VitsModel, BarkModel, etc.
-        # Ensure you use the correct class names from the transformers library or
-        # the library your model relies on.
         try:
-            print(f"Loading processor from: {model_path}")
-            self.processor = AutoProcessor.from_pretrained(model_path)
-            print(f"Loading model from: {model_path}")
-            self.model = AutoModelForTextToSpeech.from_pretrained(model_path)
+            # Load Orpheus LLM and Tokenizer
+            print(f"Loading Orpheus tokenizer from: {orpheus_model_path}")
+            self.tokenizer = AutoTokenizer.from_pretrained(orpheus_model_path)
+            print(f"Loading Orpheus model from: {orpheus_model_path}")
+            self.model = AutoModelForCausalLM.from_pretrained(
+                orpheus_model_path,
+                torch_dtype=torch.bfloat16 # Use bfloat16 as in your script
+            )
             self.model.to(self.device)
-            print("Model and processor loaded successfully.")
-
-            # --- Get Sampling Rate ---
-            # Try to get sampling rate from config, provide a default if not found
-            # Common locations: model.config.sampling_rate or processor.feature_extractor.sampling_rate
-            # Adjust this based on your specific model architecture!
-            self.sampling_rate = getattr(self.model.config, 'sampling_rate', None)
-            if self.sampling_rate is None and hasattr(self.processor, 'feature_extractor'):
-                 self.sampling_rate = getattr(self.processor.feature_extractor, 'sampling_rate', 16000) # Default fallback
-            elif self.sampling_rate is None:
-                 self.sampling_rate = 16000 # Default fallback if no config found
-            print(f"Using sampling rate: {self.sampling_rate}")
+            self.model.eval() # Set model to evaluation mode
+            print("Orpheus model and tokenizer loaded successfully.")
 
+            # Load SNAC Vocoder
+            print(f"Loading SNAC model from: {snac_model_name}")
+            self.snac_model = SNAC.from_pretrained(snac_model_name)
+            self.snac_model.to(self.device) # Move SNAC to the same device
+            self.snac_model.eval() # Set model to evaluation mode
+            print("SNAC model loaded successfully.")
 
         except Exception as e:
-            print(f"Error loading model or processor: {e}")
-            raise RuntimeError(f"Failed to load model/processor from {model_path}", e)
-
+            print(f"Error during model loading: {e}")
+            raise RuntimeError(f"Failed to load models.", e)
 
     def __call__(self, data: dict) -> bytes:
         """
-        Runs inference on the input data.
-        'data' is the dictionary parsed from the incoming JSON request payload.
-        Should return raw audio bytes (e.g., WAV format).
+        Handles incoming API requests for TTS inference.
+        Expects data['inputs'] (text) and optionally data['parameters']
         """
         try:
-            # --- Get Inputs ---
-            # Extract text input - adjust key 'inputs' if necessary
-            inputs_text = data.pop("inputs", None)
-            if inputs_text is None:
+            # --- Get Inputs & Parameters ---
+            text = data.pop("inputs", None)
+            if text is None:
                 raise ValueError("Missing 'inputs' key in request data")
 
-            # Optional: handle other parameters passed in the request
             parameters = data.pop("parameters", {})
+            # Default voice if not provided
+            voice = parameters.get("voice", "Eniola")
+            # Default generation parameters (merge with provided ones)
+            gen_params = {
+                "max_new_tokens": 1200,
+                "do_sample": True,
+                "temperature": 0.6,
+                "top_p": 0.95,
+                "repetition_penalty": 1.1,
+                "num_return_sequences": 1,
+                "eos_token_id": self.end_audio_token_id,
+                **parameters # Overwrite defaults with user params
+            }
+            # Remove non-generate params if they were passed
+            gen_params.pop("voice", None)
+
+            print(f"Received request: text='{text[:50]}...', voice='{voice}', params={gen_params}")
 
             # --- Preprocess Text ---
-            # Use the processor to prepare model inputs
-            processed_inputs = self.processor(text=inputs_text, return_tensors="pt").to(self.device)
+            prompt = f"{voice}: {text}"
+            input_ids = self.tokenizer(prompt, return_tensors="pt").input_ids
 
-            # --- Generate Speech ---
-            # Adjust generation parameters if needed (e.g., speaker embeddings)
-            # The output key might vary ('waveform', 'audio', 'speech', etc.)
-            with torch.no_grad():
-                # If your model needs specific args like speaker_embeddings, handle them here
-                # Example: speaker_embeddings = self.load_speaker_embedding(...)
-                # output = self.model.generate(**processed_inputs, speaker_embeddings=speaker_embeddings, **parameters)
-                output = self.model.generate(**processed_inputs, **parameters)
+            # Add special tokens: SOH + Input Tokens + EOT + EOH
+            start_token_tensor = torch.tensor([[self.start_human_token_id]], dtype=torch.int64)
+            end_tokens_tensor = torch.tensor([[self.end_text_token_id, self.end_human_token_id]], dtype=torch.int64)
+            processed_input_ids = torch.cat([start_token_tensor, input_ids, end_tokens_tensor], dim=1).to(self.device)
 
+            # Create attention mask (all ones for single, unpadded sequence)
+            attention_mask = torch.ones_like(processed_input_ids).to(self.device)
 
-            # --- Postprocess Audio ---
-            # Ensure output is on CPU and convert to numpy array
-            # The exact processing depends on the model output format
-            # Assuming output tensor contains the waveform
-            if isinstance(output, torch.Tensor):
-                speech_waveform = output.cpu().numpy().squeeze()
-            # Handle cases where output might be in a dictionary (common with pipelines)
-            elif isinstance(output, dict) and 'audio' in output:
-                 speech_waveform = output['audio'].cpu().numpy().squeeze()
-            elif isinstance(output, dict) and 'waveform' in output:
-                 speech_waveform = output['waveform'].cpu().numpy().squeeze()
+            print(f"Processed input shape: {processed_input_ids.shape}")
+
+            # --- Generate Audio Codes (LLM Inference) ---
+            with torch.no_grad():
+                generated_ids = self.model.generate(
+                    input_ids=processed_input_ids,
+                    attention_mask=attention_mask,
+                    **gen_params
+                )
+            print(f"Generated IDs shape: {generated_ids.shape}")
+
+            # --- Process Generated Tokens (Extract Audio Codes) ---
+            # Find the last Start of Audio token
+            soa_indices = (generated_ids[0] == self.start_audio_token_id).nonzero(as_tuple=True)[0]
+            if len(soa_indices) == 0:
+                print("Warning: Start of Audio token (128257) not found in generated sequence!")
+                # Handle this case: maybe return error, or try processing from start?
+                # For now, let's assume it might still contain codes and try processing all generated *new* tokens
+                start_idx = processed_input_ids.shape[1] # Start after the input prompt
             else:
-                 # Add handling for other potential output types if needed
-                 raise TypeError(f"Unexpected model output type: {type(output)}")
+                start_idx = soa_indices[-1].item() + 1 # Start after the last SOA token
+
+            # Extract potential audio codes (after last SOA or after input)
+            cropped_tokens = generated_ids[0, start_idx:]
+
+            # Remove End of Audio tokens
+            audio_codes_raw = cropped_tokens[cropped_tokens != self.end_audio_token_id]
+            print(f"Extracted raw audio codes count: {len(audio_codes_raw)}")
+
+            if len(audio_codes_raw) == 0:
+                 raise ValueError("No audio codes generated or extracted after processing.")
 
+            # --- Prepare Codes for SNAC Vocoder ---
+            # Adjust token values
+            adjusted_codes = [t.item() - self.audio_code_offset for t in audio_codes_raw]
 
-            # Normalize if necessary (some models output in range [-1, 1], others don't)
-            # If the waveform isn't in [-1, 1], soundfile might require normalization
-            # Example check: if np.max(np.abs(speech_waveform)) > 1.0:
-            #    speech_waveform = speech_waveform / np.max(np.abs(speech_waveform))
+            # Trim to multiple of 7
+            num_codes = len(adjusted_codes)
+            valid_length = (num_codes // 7) * 7
+            if valid_length == 0:
+                raise ValueError(f"Not enough audio codes ({num_codes}) to form a multiple of 7 after processing.")
+            trimmed_codes = adjusted_codes[:valid_length]
+            print(f"Trimmed adjusted audio codes count: {len(trimmed_codes)}")
+
+            # --- Redistribute Codes ---
+            # Use static method or instance method, ensure tensors are on correct device
+            snac_input_codes = redistribute_codes_static(trimmed_codes)
+            snac_input_codes = [layer.to(self.device) for layer in snac_input_codes]
+
+            # --- Decode Audio (SNAC Inference) ---
+            print("Decoding audio with SNAC...")
+            with torch.no_grad():
+                audio_hat = self.snac_model.decode(snac_input_codes)
+            print(f"Decoded audio tensor shape: {audio_hat.shape}") # Should be [1, 1, num_samples]
+
+            # --- Postprocess Audio ---
+            # Move to CPU, remove batch/channel dims, convert to numpy
+            audio_waveform = audio_hat.detach().squeeze().cpu().numpy()
 
             # --- Convert to WAV Bytes ---
-            # Use an in-memory buffer to store the WAV file
             buffer = io.BytesIO()
-            sf.write(buffer, speech_waveform, self.sampling_rate, format='WAV')
+            sf.write(buffer, audio_waveform, self.sampling_rate, format='WAV')
             buffer.seek(0)
             wav_bytes = buffer.read()
 
@@ -109,6 +201,6 @@ class EndpointHandler():
             return wav_bytes
 
         except Exception as e:
-            print(f"Error during inference: {e}")
-            # Re-raise the exception so the endpoint framework knows it failed
+            print(f"Error during inference call: {e}")
+            # Re-raise for endpoint framework
             raise RuntimeError(f"Inference failed: {e}")