Add custom inference handler for Maya1 TTS

handler.py

@@ -75,7 +75,7 @@ class EndpointHandler:
           "description": "... optional voice description ...",
           "generation_args": { optional dict for text generation params }
         }
-        
+
         Returns dict with base64 audio:
         {
           "audio_base64": "<base64-encoded WAV data>",
@@ -84,23 +84,25 @@ class EndpointHandler:
         """
         # Extract inputs (HF always provides this key)
         inputs = data.get("inputs", "")
-        
+
         # Get additional parameters from top level
         description = data.get("description", "")
         generation_args = data.get("generation_args", {})
-        
+
         # Parse inputs
         if isinstance(inputs, dict):
             # inputs is a dict with text and description
             text = inputs.get("text", "")
-            description = inputs.get("description", description)  # override if in inputs
+            description = inputs.get(
+                "description", description
+            )  # override if in inputs
         elif isinstance(inputs, str):
             # inputs is just the text string
             text = inputs
         else:
             # Try to convert to string
             text = str(inputs) if inputs else ""
-        
+
         if not text:
             return {"error": f"No text provided. Received: {data}"}
         prompt = description + "\n" + text
@@ -139,17 +141,55 @@ class EndpointHandler:
 
         # Tokenize and generate
         tokenizer_inputs = self.tokenizer(prompt, return_tensors="pt").to(self.device)
-        outputs = self.model.generate(**tokenizer_inputs, **generation_args)
+
+        # Set default generation args if not provided
+        default_gen_args = {
+            "max_length": 2048,
+            "do_sample": True,
+            "temperature": 0.7,
+            "top_p": 0.9,
+        }
+        default_gen_args.update(generation_args)
+
+        outputs = self.model.generate(**tokenizer_inputs, **default_gen_args)
         token_ids = outputs[0]
 
-        # decode tokens to intermediate representation (for Maya1)
-        # assuming model outputs token ids for audio generation — adjust as per model spec
-        audio_feats = (
-            token_ids  # may need further decoding depending on how Maya1 works
-        )
+        # Maya1 generates SNAC audio codes in the output tokens
+        # Extract only the generated tokens (excluding input prompt tokens)
+        generated_ids = token_ids[tokenizer_inputs["input_ids"].shape[1] :]
+
+        # Convert token IDs to SNAC codes
+        # Maya1 outputs are structured as SNAC token IDs that need to be reshaped
+        # SNAC expects codes in shape (batch, num_codebooks, seq_len)
+        # Assuming Maya1 outputs tokens sequentially for each codebook
+        import torch
+
+        snac_codes = generated_ids.unsqueeze(0)  # Add batch dimension
+
+        # SNAC 24kHz uses 7 codebooks - reshape accordingly
+        # The model outputs interleaved codes, so we need to split them
+        num_codebooks = 7
+        seq_len = len(generated_ids) // num_codebooks
+
+        if len(generated_ids) >= num_codebooks:
+            # Reshape to (1, num_codebooks, seq_len) for SNAC decoder
+            snac_codes = (
+                generated_ids[: seq_len * num_codebooks]
+                .reshape(num_codebooks, seq_len)
+                .unsqueeze(0)
+            )
+
+        # Decode using SNAC to synthesize waveform
+        with torch.no_grad():
+            waveform = self.snac.decode(snac_codes.to(self.device))
+
+        # Extract audio and convert to numpy
+        if waveform.dim() == 3:  # (batch, channels, samples)
+            waveform = waveform.squeeze(0).squeeze(0)  # Remove batch and channel dims
+        elif waveform.dim() == 2:  # (batch, samples)
+            waveform = waveform.squeeze(0)
 
-        # pass features to SNAC to synthesize waveform
-        waveform = self.snac.decode(audio_feats).cpu().numpy()  # shape (n_samples,)
+        waveform = waveform.cpu().numpy()
 
         # convert waveform to bytes (e.g. WAV) using soundfile loaded into self.sf
         buf = io.BytesIO()