handler.py

import os
import torch
import numpy as np
from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
from snac import SNAC
import logging
import json
import base64
import io
import wave
from threading import Thread

# Set up logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

class EndpointHandler:
    def __init__(self, path=""):
        logger.info("Initializing Orpheus TTS handler")
        # 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)
        logger.info(f"Model loaded on {self.device}")

        # Load tokenizer
        self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)
        logger.info("Tokenizer loaded")

        # Load SNAC model for audio decoding
        try:
            self.snac_model = SNAC.from_pretrained("hubertsiuzdak/snac_24khz")
            self.snac_model.to(self.device)
            logger.info("SNAC model loaded")
        except Exception as e:
            logger.error(f"Error loading SNAC: {str(e)}")
            raise

        # Special tokens
        self.start_token = torch.tensor([[128259]], dtype=torch.int64)
        self.end_tokens = torch.tensor([[128009, 128260]], dtype=torch.int64)
        self.start_audio_token = 128257
        self.end_audio_token = 128258

        self._warmed_up = False
        logger.info("Handler initialization complete")

    def preprocess(self, data):
        """Preprocess input data before inference."""
        logger.info(f"Preprocessing data: {type(data)}")

        # Handle health check
        if data == "ping" or (isinstance(data, dict) and data.get("inputs") == "ping"):
            return {"health_check": True}

        if isinstance(data, dict) and "inputs" in data:
            text = data["inputs"]
            parameters = data.get("parameters", {})
        else:
            text = data
            parameters = {}

        # Extract parameters
        voice = parameters.get("voice", "tara")
        temperature = float(parameters.get("temperature", 0.6))
        top_p = float(parameters.get("top_p", 0.95))
        max_new_tokens = int(parameters.get("max_new_tokens", 1200))
        repetition_penalty = float(parameters.get("repetition_penalty", 1.1))
        stream = parameters.get("stream", False)  # Check for stream parameter

        prompt = f"{voice}: {text}"
        logger.info(f"Formatted prompt with voice {voice}")

        input_ids = self.tokenizer(prompt, return_tensors="pt").input_ids
        modified_input_ids = torch.cat([self.start_token, input_ids, self.end_tokens], dim=1)
        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,
            "stream": stream,
            "health_check": False
        }

    def inference(self, inputs):
        """Run model inference (non-streaming)."""
        if inputs.get("health_check", False):
            return {"status": "ok"}

        input_ids = inputs["input_ids"]
        attention_mask = inputs["attention_mask"]

        logger.info(f"Running non-streaming inference with max_new_tokens={inputs['max_new_tokens']}")
        with torch.no_grad():
            generated_ids = self.model.generate(
                input_ids=input_ids,
                attention_mask=attention_mask,
                max_new_tokens=inputs['max_new_tokens'],
                do_sample=True,
                temperature=inputs['temperature'],
                top_p=inputs['top_p'],
                repetition_penalty=inputs['repetition_penalty'],
                num_return_sequences=1,
                eos_token_id=self.end_audio_token,
            )
        logger.info(f"Generation complete, output shape: {generated_ids.shape}")
        return generated_ids

    def postprocess(self, generated_ids):
        """Process generated tokens into a single audio file (non-streaming)."""
        if isinstance(generated_ids, dict) and "status" in generated_ids:
            return generated_ids

        logger.info("Postprocessing generated tokens for non-streaming output")
        
        # Isolate audio tokens after the start_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
            logger.warning("No start audio token found in non-streaming output")

        # Get list of token integers, removing the end token
        code_list = [t.item() for t in cropped_tensor.squeeze() if t.item() != self.end_audio_token]

        # Decode and encode the full audio
        audio_b64 = self._decode_audio_chunk(code_list)
        if not audio_b64:
             return {"error": "No audio samples generated"}
        
        logger.info(f"Audio encoded as base64, length: {len(audio_b64)}")
        return {
            "audio_b64": audio_b64,
            "sample_rate": 24000
        }

    def _decode_audio_chunk(self, code_list):
        """Decodes a list of token codes into a base64 WAV string."""
        if not code_list:
            return None

        # Ensure length is a multiple of 7 for SNAC
        new_length = (len(code_list) // 7) * 7
        if new_length == 0:
            return None
        
        trimmed_list = code_list[:new_length]
        
        # Adjust token values for SNAC
        adjusted_codes = [t - 128266 for t in trimmed_list]

        # Redistribute codes into layers for SNAC model
        audio = self.redistribute_codes(adjusted_codes)
        audio_sample = audio.detach().squeeze().cpu().numpy()
        
        # Convert float32 array to int16 for WAV format
        audio_int16 = (audio_sample * 32767).astype(np.int16)
        
        # Create WAV in memory and encode as base64
        with io.BytesIO() as wav_io:
            with wave.open(wav_io, 'wb') as wav_file:
                wav_file.setnchannels(1)
                wav_file.setsampwidth(2)
                wav_file.setframerate(24000)
                wav_file.writeframes(audio_int16.tobytes())
            wav_data = wav_io.getvalue()
        
        return base64.b64encode(wav_data).decode('utf-8')

    def _stream_inference(self, inputs):
        """Generator function for streaming inference."""
        streamer = TextIteratorStreamer(self.tokenizer, skip_prompt=True, skip_special_tokens=True)

        generation_kwargs = dict(
            input_ids=inputs["input_ids"],
            attention_mask=inputs["attention_mask"],
            streamer=streamer,
            max_new_tokens=inputs['max_new_tokens'],
            do_sample=True,
            temperature=inputs['temperature'],
            top_p=inputs['top_p'],
            repetition_penalty=inputs['repetition_penalty'],
            eos_token_id=self.end_audio_token,
        )

        # Run generation in a separate thread
        thread = Thread(target=self.model.generate, kwargs=generation_kwargs)
        thread.start()

        logger.info("Starting streaming inference...")
        token_buffer = []
        found_start_audio = False

        for new_text in streamer:
            # The streamer decodes tokens to text, so we must re-encode to get token IDs.
            # This is a workaround for HF streamer design.
            new_tokens = self.tokenizer.encode(new_text, add_special_tokens=False)
            
            for token in new_tokens:
                if not found_start_audio:
                    if token == self.start_audio_token:
                        found_start_audio = True
                    continue # Skip until start_audio_token is found

                if token == self.end_audio_token:
                    # Process any remaining tokens in buffer and stop
                    if len(token_buffer) >= 7:
                        audio_b64 = self._decode_audio_chunk(token_buffer)
                        if audio_b64:
                             yield f"data: {json.dumps({'audio_b64': audio_b64})}\n\n"
                    logger.info("End of audio token found. Stopping stream.")
                    return

                token_buffer.append(token)
                
                # If buffer has enough tokens for a multiple of 7, process and yield
                if len(token_buffer) >= 7:
                    process_len = (len(token_buffer) // 7) * 7
                    codes_to_process = token_buffer[:process_len]
                    token_buffer = token_buffer[process_len:] # Keep the remainder

                    audio_b64 = self._decode_audio_chunk(codes_to_process)
                    if audio_b64:
                        logger.info(f"Yielding audio chunk of {len(codes_to_process)} tokens")
                        yield f"data: {json.dumps({'audio_b64': audio_b64})}\n\n"
        
        # Process any final tokens left in the buffer after the loop finishes
        if token_buffer:
            audio_b64 = self._decode_audio_chunk(token_buffer)
            if audio_b64:
                yield f"data: {json.dumps({'audio_b64': audio_b64})}\n\n"
        
        logger.info("Streaming complete.")


    def redistribute_codes(self, code_list):
        """Reorganize tokens for SNAC decoding."""
        layer_1, layer_2, layer_3 = [], [], []
        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))
            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))
        
        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)
        ]
        return self.snac_model.decode(codes)

    def __call__(self, data):
        """Main entry point for the handler."""
        if not self._warmed_up:
            self._warmup()
        
        try:
            # Handle health check separately to avoid preprocessing
            if data == "ping" or (isinstance(data, dict) and data.get("inputs") == "ping"):
                logger.info("Processing health check request")
                return {"status": "ok"}

            preprocessed_inputs = self.preprocess(data)

            # Route to streaming or non-streaming path
            if preprocessed_inputs.get("stream"):
                return self._stream_inference(preprocessed_inputs)
            else:
                model_outputs = self.inference(preprocessed_inputs)
                return self.postprocess(model_outputs)

        except Exception as e:
            import traceback
            logger.error(f"Error processing request: {str(e)}\n{traceback.format_exc()}")
            return {"error": str(e), "traceback": traceback.format_exc()}

    def _warmup(self):
        try:
            logger.info("Warming up model...")
            dummy_prompt = "tara: Hello"
            input_ids = self.tokenizer(dummy_prompt, return_tensors="pt").input_ids.to(self.device)
            _ = self.model.generate(input_ids=input_ids, max_new_tokens=10)
            self._warmed_up = True
            logger.info("Warmup complete.")
        except Exception as e:
            logger.error(f"[WARMUP ERROR] {str(e)}")