Add custom inference handler for Maya1 TTS

handler.py

@@ -4,6 +4,20 @@ import os
 import struct
 import wave
 
+# SNAC token constants for Maya1
+CODE_START_TOKEN_ID = 128257
+CODE_END_TOKEN_ID = 128258
+CODE_TOKEN_OFFSET = 128266
+SNAC_MIN_ID = 128266
+SNAC_MAX_ID = 156937
+SNAC_TOKENS_PER_FRAME = 7
+
+SOH_ID = 128259
+EOH_ID = 128260
+SOA_ID = 128261
+BOS_ID = 128000
+TEXT_EOT_ID = 128009
+
 
 class EndpointHandler:
     def __init__(self, path=""):
@@ -126,6 +140,67 @@ class EndpointHandler:
         )
         self.sf = sf
 
+    def build_prompt(self, description: str, text: str) -> str:
+        """Build formatted prompt for Maya1."""
+        soh_token = self.tokenizer.decode([SOH_ID])
+        eoh_token = self.tokenizer.decode([EOH_ID])
+        soa_token = self.tokenizer.decode([SOA_ID])
+        sos_token = self.tokenizer.decode([CODE_START_TOKEN_ID])
+        eot_token = self.tokenizer.decode([TEXT_EOT_ID])
+        bos_token = self.tokenizer.bos_token
+        
+        formatted_text = f'<description="{description}"> {text}'
+        
+        prompt = (
+            soh_token + bos_token + formatted_text + eot_token +
+            eoh_token + soa_token + sos_token
+        )
+        
+        return prompt
+
+    def extract_snac_codes(self, token_ids: list) -> list:
+        """Extract SNAC codes from generated tokens."""
+        try:
+            eos_idx = token_ids.index(CODE_END_TOKEN_ID)
+        except ValueError:
+            eos_idx = len(token_ids)
+        
+        snac_codes = [
+            token_id for token_id in token_ids[:eos_idx]
+            if SNAC_MIN_ID <= token_id <= SNAC_MAX_ID
+        ]
+        
+        return snac_codes
+
+    def unpack_snac_from_7(self, snac_tokens: list) -> list:
+        """Unpack 7-token SNAC frames to 3 hierarchical levels."""
+        if snac_tokens and snac_tokens[-1] == CODE_END_TOKEN_ID:
+            snac_tokens = snac_tokens[:-1]
+        
+        frames = len(snac_tokens) // SNAC_TOKENS_PER_FRAME
+        snac_tokens = snac_tokens[:frames * SNAC_TOKENS_PER_FRAME]
+        
+        if frames == 0:
+            return [[], [], []]
+        
+        l1, l2, l3 = [], [], []
+        
+        for i in range(frames):
+            slots = snac_tokens[i*7:(i+1)*7]
+            l1.append((slots[0] - CODE_TOKEN_OFFSET) % 4096)
+            l2.extend([
+                (slots[1] - CODE_TOKEN_OFFSET) % 4096,
+                (slots[4] - CODE_TOKEN_OFFSET) % 4096,
+            ])
+            l3.extend([
+                (slots[2] - CODE_TOKEN_OFFSET) % 4096,
+                (slots[3] - CODE_TOKEN_OFFSET) % 4096,
+                (slots[5] - CODE_TOKEN_OFFSET) % 4096,
+                (slots[6] - CODE_TOKEN_OFFSET) % 4096,
+            ])
+        
+        return [l1, l2, l3]
+
     def __call__(self, data):
         """
         HF Endpoints format:
@@ -164,7 +239,10 @@ class EndpointHandler:
 
         if not text:
             return {"error": f"No text provided. Received: {data}"}
-        prompt = description + "\n" + text
+        
+        # Use default description if not provided
+        if not description:
+            description = "Realistic male voice in the 30s age with american accent. Normal pitch, warm timbre, conversational pacing."
 
         # If running in fake mode (quick smoke test), synthesize a sine tone
         if getattr(self, "snac", None) is None and getattr(self, "model", None) is None:
@@ -198,64 +276,64 @@ class EndpointHandler:
 
             return {"audio_base64": b64, "sampling_rate": sr}
 
-        # Tokenize and generate
-        tokenizer_inputs = self.tokenizer(prompt, return_tensors="pt").to(self.device)
+        # Build properly formatted prompt for Maya1
+        prompt = self.build_prompt(description, text)
+
+        # Tokenize
+        import torch
+        tokenizer_inputs = self.tokenizer(prompt, return_tensors="pt")
+        if torch.cuda.is_available():
+            tokenizer_inputs = {k: v.to(self.device) for k, v in tokenizer_inputs.items()}
+        else:
+            tokenizer_inputs = tokenizer_inputs.to(self.device)
 
         # Set default generation args if not provided
         default_gen_args = {
-            "max_length": 2048,
-            "do_sample": True,
-            "temperature": 0.7,
+            "max_new_tokens": 2048,
+            "min_new_tokens": 28,
+            "temperature": 0.4,
             "top_p": 0.9,
+            "repetition_penalty": 1.1,
+            "do_sample": True,
+            "eos_token_id": CODE_END_TOKEN_ID,
+            "pad_token_id": self.tokenizer.pad_token_id,
         }
         default_gen_args.update(generation_args)
 
-        outputs = self.model.generate(**tokenizer_inputs, **default_gen_args)
-        token_ids = outputs[0]
-
-        # 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
+        # Generate tokens
+        with torch.inference_mode():
+            outputs = self.model.generate(**tokenizer_inputs, **default_gen_args)
 
-        snac_codes = generated_ids.unsqueeze(0)  # Add batch dimension
+        # Extract generated tokens (everything after the input prompt)
+        generated_ids = outputs[0, tokenizer_inputs["input_ids"].shape[1]:].tolist()
 
-        # 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
-        # SNAC outputs shape (batch, samples) or (batch, 1, samples)
-        # Safely remove all size-1 dimensions
-        waveform = waveform.squeeze()  # Remove all dimensions of size 1
+        # Extract SNAC audio tokens
+        snac_tokens = self.extract_snac_codes(generated_ids)
         
-        # Ensure we have a 1D tensor
-        if waveform.dim() > 1:
-            waveform = waveform[0]  # Take first item if still multi-dimensional
+        if len(snac_tokens) < 7:
+            return {"error": f"Not enough SNAC tokens generated: {len(snac_tokens)}"}
 
-        waveform = waveform.cpu().numpy()
-
-        # convert waveform to bytes (e.g. WAV) using soundfile loaded into self.sf
+        # Unpack SNAC tokens to 3 hierarchical levels
+        levels = self.unpack_snac_from_7(snac_tokens)
+        
+        # Convert to tensors
+        codes_tensor = [
+            torch.tensor(level, dtype=torch.long, device=self.device).unsqueeze(0)
+            for level in levels
+        ]
+
+        # Generate final audio with SNAC decoder
+        with torch.inference_mode():
+            z_q = self.snac.quantizer.from_codes(codes_tensor)
+            audio = self.snac.decoder(z_q)[0, 0].cpu().numpy()
+
+        # Trim warmup samples (first 2048 samples)
+        if len(audio) > 2048:
+            audio = audio[2048:]
+
+        # Save audio to WAV
         buf = io.BytesIO()
-        self.sf.write(buf, waveform, 24000, format="WAV")
+        self.sf.write(buf, audio, 24000, format="WAV")
         wav_bytes = buf.getvalue()
         b64 = base64.b64encode(wav_bytes).decode("utf-8")
 
@@ -263,8 +341,6 @@ class EndpointHandler:
             "audio_base64": b64,
             "sampling_rate": 24000,
         }
-
-
 # Module-level convenience functions for hosting platforms (Hugging Face Endpoints)
 # The platform typically expects top-level `init` and `predict` (or `run`) callables
 # so we provide thin wrappers around the EndpointHandler class.