Updated handler.py with claude

handler.py

@@ -1,20 +1,176 @@
+import os
 import torch
 import numpy as np
-import soundfile as sf
-import io
-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
+from snac import SNAC
 
-    for i in range(num_groups):
-        idx = 7 * i
-        try:
+class EndpointHandler:
+    def __init__(self, path=""):
+        # Load the Orpheus model and tokenizer
+        self.model_name = "hypaai/Hypa_Orpheus-3b-0.1-ft-unsloth-merged_16bit"
+        self.model = AutoModelForCausalLM.from_pretrained(
+            self.model_name, 
+            torch_dtype=torch.bfloat16
+        )
+        
+        # Move model to GPU if available
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        self.model.to(self.device)
+        
+        # Load tokenizer
+        self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)
+        
+        # Load SNAC model for audio decoding
+        self.snac_model = SNAC.from_pretrained("hubertsiuzdak/snac_24khz")
+        self.snac_model.to(self.device)
+        
+        # Special tokens
+        self.start_token = torch.tensor([[128259]], dtype=torch.int64)  # Start of human
+        self.end_tokens = torch.tensor([[128009, 128260]], dtype=torch.int64)  # End of text, End of human
+        self.padding_token = 128263
+        self.start_audio_token = 128257  # Start of Audio token
+        self.end_audio_token = 128258  # End of Audio token
+        
+        print(f"Model loaded on {self.device}")
+        
+    def preprocess(self, data):
+        """
+        Preprocess input data before inference
+        """
+        inputs = data.pop("inputs", data)
+        
+        # Extract parameters from request
+        text = inputs.get("text", "")
+        voice = inputs.get("voice", "tara")
+        temperature = float(inputs.get("temperature", 0.6))
+        top_p = float(inputs.get("top_p", 0.95))
+        max_new_tokens = int(inputs.get("max_new_tokens", 1200))
+        repetition_penalty = float(inputs.get("repetition_penalty", 1.1))
+        
+        # Format prompt with voice
+        prompt = f"{voice}: {text}"
+        
+        # Tokenize
+        input_ids = self.tokenizer(prompt, return_tensors="pt").input_ids
+        
+        # Add special tokens
+        modified_input_ids = torch.cat([self.start_token, input_ids, self.end_tokens], dim=1)
+        
+        # No need for padding as we're processing a single sequence
+        input_ids = modified_input_ids.to(self.device)
+        attention_mask = torch.ones_like(input_ids)
+        
+        return {
+            "input_ids": input_ids,
+            "attention_mask": attention_mask,
+            "temperature": temperature,
+            "top_p": top_p,
+            "max_new_tokens": max_new_tokens,
+            "repetition_penalty": repetition_penalty
+        }
+    
+    def inference(self, inputs):
+        """
+        Run model inference on the preprocessed inputs
+        """
+        # Extract parameters
+        input_ids = inputs["input_ids"]
+        attention_mask = inputs["attention_mask"]
+        temperature = inputs["temperature"]
+        top_p = inputs["top_p"]
+        max_new_tokens = inputs["max_new_tokens"]
+        repetition_penalty = inputs["repetition_penalty"]
+        
+        # Generate output tokens
+        with torch.no_grad():
+            generated_ids = self.model.generate(
+                input_ids=input_ids,
+                attention_mask=attention_mask,
+                max_new_tokens=max_new_tokens,
+                do_sample=True,
+                temperature=temperature,
+                top_p=top_p,
+                repetition_penalty=repetition_penalty,
+                num_return_sequences=1,
+                eos_token_id=self.end_audio_token,
+            )
+        
+        return generated_ids
+    
+    def postprocess(self, generated_ids):
+        """
+        Process generated tokens into audio
+        """
+        # Find Start of Audio token
+        token_indices = (generated_ids == self.start_audio_token).nonzero(as_tuple=True)
+        
+        if len(token_indices[1]) > 0:
+            last_occurrence_idx = token_indices[1][-1].item()
+            cropped_tensor = generated_ids[:, last_occurrence_idx+1:]
+        else:
+            cropped_tensor = generated_ids
+            
+        # Remove End of Audio tokens
+        processed_rows = []
+        for row in cropped_tensor:
+            masked_row = row[row != self.end_audio_token]
+            processed_rows.append(masked_row)
+            
+        # Prepare audio codes
+        code_lists = []
+        for row in processed_rows:
+            row_length = row.size(0)
+            # Ensure length is multiple of 7 for SNAC
+            new_length = (row_length // 7) * 7
+            trimmed_row = row[:new_length]
+            trimmed_row = [t.item() - 128266 for t in trimmed_row]  # Adjust token values
+            code_lists.append(trimmed_row)
+            
+        # Generate audio from codes
+        audio_samples = []
+        for code_list in code_lists:
+            audio = self.redistribute_codes(code_list)
+            audio_samples.append(audio)
+            
+        # Return first (and only) audio sample
+        audio_sample = audio_samples[0].detach().squeeze().cpu().numpy()
+        
+        # Convert to base64 for transmission
+        import base64
+        import io
+        import wave
+        
+        # Convert float32 array to int16 for WAV format
+        audio_int16 = (audio_sample * 32767).astype(np.int16)
+        
+        # Create WAV in memory
+        with io.BytesIO() as wav_io:
+            with wave.open(wav_io, 'wb') as wav_file:
+                wav_file.setnchannels(1)  # Mono
+                wav_file.setsampwidth(2)  # 16-bit
+                wav_file.setframerate(24000)  # 24kHz
+                wav_file.writeframes(audio_int16.tobytes())
+            wav_data = wav_io.getvalue()
+        
+        # Encode as base64
+        audio_b64 = base64.b64encode(wav_data).decode('utf-8')
+        
+        return {
+            "audio_b64": audio_b64,
+            "sample_rate": 24000
+        }
+    
+    def redistribute_codes(self, code_list):
+        """
+        Reorganize tokens for SNAC decoding
+        """
+        layer_1 = []  # Coarsest layer
+        layer_2 = []  # Intermediate layer
+        layer_3 = []  # Finest layer
+        
+        num_groups = len(code_list) // 7
+        for i in range(num_groups):
+            idx = 7 * i
             layer_1.append(code_list[idx])
             layer_2.append(code_list[idx + 1] - 4096)
             layer_3.append(code_list[idx + 2] - (2 * 4096))
@@ -22,185 +178,22 @@ def redistribute_codes_static(code_list):
             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 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 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
-
-        try:
-            # 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)
-            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 during model loading: {e}")
-            raise RuntimeError(f"Failed to load models.", e)
-
-    def __call__(self, data: dict) -> bytes:
+        
+        codes = [
+            torch.tensor(layer_1).unsqueeze(0).to(self.device),
+            torch.tensor(layer_2).unsqueeze(0).to(self.device),
+            torch.tensor(layer_3).unsqueeze(0).to(self.device)
+        ]
+        
+        # Decode audio
+        audio_hat = self.snac_model.decode(codes)
+        return audio_hat
+    
+    def __call__(self, data):
         """
-        Handles incoming API requests for TTS inference.
-        Expects data['inputs'] (text) and optionally data['parameters']
+        Main entry point for the handler
         """
-        try:
-            # --- Get Inputs & Parameters ---
-            text = data.pop("inputs", None)
-            if text is None:
-                raise ValueError("Missing 'inputs' key in request data")
-
-            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 ---
-            prompt = f"{voice}: {text}"
-            input_ids = self.tokenizer(prompt, return_tensors="pt").input_ids
-
-            # 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)
-
-            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:
-                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]
-
-            # 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 ---
-            buffer = io.BytesIO()
-            sf.write(buffer, audio_waveform, self.sampling_rate, format='WAV')
-            buffer.seek(0)
-            wav_bytes = buffer.read()
-
-            print(f"Generated {len(wav_bytes)} bytes of WAV audio.")
-            return wav_bytes
-
-        except Exception as e:
-            print(f"Error during inference call: {e}")
-            # Re-raise for endpoint framework
-            raise RuntimeError(f"Inference failed: {e}")
+        preprocessed_inputs = self.preprocess(data)
+        model_outputs = self.inference(preprocessed_inputs)
+        response = self.postprocess(model_outputs)
+        return response