Add custom handler for TTS inference

handler.py

@@ -0,0 +1,155 @@
+#!/usr/bin/env python3
+"""
+Handler for Sesame CSM-1B TTS model deployment on Hugging Face Inference Endpoints
+"""
+
+import os
+import base64
+import io
+import torch
+import numpy as np
+from typing import Dict, Any, List
+from transformers import AutoTokenizer, AutoModelForCausalLM
+import scipy.io.wavfile as wavfile
+
+# Global variables for model and tokenizer
+model = None
+tokenizer = None
+
+def init():
+    """
+    Initialize the model and tokenizer
+    This is called once when the endpoint starts
+    """
+    global model, tokenizer
+
+    print("Initializing CSM-1B model...")
+
+    # Set device
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    print(f"Using device: {device}")
+
+    try:
+        # Load tokenizer
+        tokenizer = AutoTokenizer.from_pretrained(
+            "farazmoradi98/csm-1b",  # Use your forked model
+            trust_remote_code=True
+        )
+
+        # Load model
+        model = AutoModelForCausalLM.from_pretrained(
+            "farazmoradi98/csm-1b",  # Use your forked model
+            trust_remote_code=True,
+            torch_dtype=torch.float16 if device == "cuda" else torch.float32,
+            device_map="auto"
+        )
+
+        print("✅ Model and tokenizer loaded successfully!")
+
+    except Exception as e:
+        print(f"❌ Error loading model: {e}")
+        raise
+
+def generate_speech(text: str, speaker: int = 0) -> bytes:
+    """
+    Generate speech from text using CSM-1B model
+
+    Args:
+        text (str): Input text to convert to speech
+        speaker (int): Speaker ID (0-3 for CSM-1B)
+
+    Returns:
+        bytes: WAV audio data
+    """
+    global model, tokenizer
+
+    try:
+        # Tokenize input text
+        inputs = tokenizer(text, return_tensors="pt").to(model.device)
+
+        # Generate speech
+        with torch.no_grad():
+            output = model.generate(
+                **inputs,
+                speaker=speaker,
+                max_new_tokens=1024,  # Adjust as needed
+                do_sample=True,
+                temperature=0.8,
+                top_p=0.9,
+                repetition_penalty=1.1
+            )
+
+        # Decode audio from model output
+        # CSM-1B outputs audio tokens that need to be converted to waveform
+        audio_tokens = output[0][inputs.input_ids.shape[1]:]
+        audio_array = model.decode_audio(audio_tokens)
+
+        # Convert to 16-bit PCM WAV
+        audio_array = (audio_array * 32767).astype(np.int16)
+
+        # Save to WAV buffer
+        wav_buffer = io.BytesIO()
+        wavfile.write(wav_buffer, 24000, audio_array)  # CSM-1B uses 24kHz
+        wav_buffer.seek(0)
+
+        return wav_buffer.getvalue()
+
+    except Exception as e:
+        print(f"❌ Error generating speech: {e}")
+        raise
+
+def handler(request: Dict[str, Any]) -> Dict[str, Any]:
+    """
+    Main handler function for Hugging Face Inference API
+
+    Args:
+        request (dict): Request containing input data
+
+    Returns:
+        dict: Response with base64 encoded audio
+    """
+    try:
+        # Extract inputs from request
+        inputs = request.get("inputs", {})
+
+        # Handle different input formats
+        if isinstance(inputs, str):
+            text = inputs
+            speaker = 0
+        elif isinstance(inputs, dict):
+            text = inputs.get("text", "")
+            speaker = inputs.get("speaker", 0)
+        else:
+            return {
+                "error": "Invalid input format. Expected string or dict with 'text' field."
+            }
+
+        if not text:
+            return {
+                "error": "No text provided for speech generation."
+            }
+
+        print(f"Generating speech for: '{text}' (speaker: {speaker})")
+
+        # Generate speech
+        audio_data = generate_speech(text, speaker)
+
+        # Convert to base64
+        audio_base64 = base64.b64encode(audio_data).decode('utf-8')
+
+        return {
+            "audio": audio_base64,
+            "format": "wav",
+            "sample_rate": 24000,
+            "speaker": speaker
+        }
+
+    except Exception as e:
+        print(f"❌ Handler error: {e}")
+        return {
+            "error": f"Speech generation failed: {str(e)}"
+        }
+
+# Initialize model on startup
+if __name__ != "__main__":
+    init()