maya1 / handler.py

Add custom inference handler for Maya1 TTS

1e1da0c 11 months ago

15.5 kB

import base64
import io
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=""):
        # allow overriding device and dtype via environment variables for local CPU testing
        # `MAYA_DEVICE`: 'cpu' or 'auto' (default 'auto')
        # `MAYA_TORCH_DTYPE`: 'bf16' or 'fp32' (default 'bf16')
        # `MAYA_USE_FAKE`: if '1', use a tiny fake pipeline for smoke testing (no HF downloads)
        device_override = os.getenv("MAYA_DEVICE", "auto")
        dtype_override = os.getenv("MAYA_TORCH_DTYPE", "bf16")
        use_fake = os.getenv("MAYA_USE_FAKE", "0") == "1"

        # map string to torch dtype will be set after importing torch

        if use_fake:
            # Minimal fake components for quick local smoke tests. Generates a short sine tone.
            self.model = None
            self.tokenizer = None
            # keep device as plain string for fake mode
            self.device = "cpu"
            self.snac = None
            # no external libs required for fake mode
            self.sf = None
            return

        # import heavy inference dependencies lazily so fake-mode doesn't require them
        try:
            import soundfile as sf
            import torch
            from snac import SNAC
            from transformers import AutoModelForCausalLM, AutoTokenizer
        except Exception as e:
            raise RuntimeError(
                f"Failed to import inference dependencies: {e}.\n"
                "Install the packages listed in `inference_endpoints/requirements.txt`."
            )

        # map string to torch dtype
        torch_dtype = torch.bfloat16 if dtype_override == "bf16" else torch.float32

        # load Maya1 text-to-voice model
        # force CPU device_map when requested to avoid trying to use GPUs
        device_map_arg = "auto"
        if device_override == "cpu":
            device_map_arg = "cpu"

        # Allow loading from HuggingFace Hub via environment variable
        # MAYA_MODEL_ID: HuggingFace model ID (e.g., "maya-research/maya1" or similar Maya model)
        model_id = os.getenv("MAYA_MODEL_ID", "")

        # Determine model path: use env var if set, otherwise local path
        local_path = path if path else "/repository"

        # Check if trying to load from self (circular reference)
        # Only block if it's exactly "binhqd/maya1" - other repos ending in /maya1 are fine
        if model_id == "binhqd/maya1":
            print(f"⚠️ MAYA_MODEL_ID is set to {model_id} which is this repository!")
            print("This creates a circular reference. Checking local files instead...")
            model_id = ""  # Force local loading

        if model_id:
            # Load from HuggingFace Hub
            model_path = model_id
            print(f"Loading Maya1 model from HuggingFace Hub: {model_id}")
        else:
            # Load from local repository
            model_path = local_path

            # Check if this is a valid model directory
            config_path = os.path.join(model_path, "config.json")
            if not os.path.exists(config_path):
                raise RuntimeError(
                    f"❌ Model configuration not found at: {config_path}\n\n"
                    f"The repository appears to be missing model weights and configuration.\n\n"
                    f"To fix this, you have 3 options:\n\n"
                    f"1. 🧪 TESTING MODE (Quick Start):\n"
                    f"   Set: MAYA_USE_FAKE=1\n"
                    f"   This generates test audio (sine wave) without requiring model files.\n\n"
                    f"2. 🌐 USE EXISTING MODEL FROM HUGGINGFACE:\n"
                    f"   Set: MAYA_MODEL_ID=<actual-model-id>\n"
                    f"   Example: MAYA_MODEL_ID=maya-research/maya1\n"
                    f"   (Must be a valid HuggingFace model repository with weights)\n\n"
                    f"3. 📦 UPLOAD YOUR MODEL FILES:\n"
                    f"   Upload these files to your repository (binhqd/maya1):\n"
                    f"   - config.json\n"
                    f"   - model.safetensors (or pytorch_model.bin)\n"
                    f"   - tokenizer.json\n"
                    f"   - tokenizer_config.json\n"
                    f"   - generation_config.json (optional)\n\n"
                    f"Current path: {model_path}\n"
                    f"Files found: {os.listdir(model_path) if os.path.exists(model_path) else 'path does not exist'}\n\n"
                    f"⚠️ Note: MAYA_MODEL_ID=binhqd/maya1 won't work (this repository)\n"
                    f"   because it doesn't contain model weights.\n"
                )

        try:
            self.model = AutoModelForCausalLM.from_pretrained(
                model_path, torch_dtype=torch_dtype, device_map=device_map_arg
            )
            self.tokenizer = AutoTokenizer.from_pretrained(model_path)
        except Exception as e:
            raise RuntimeError(
                f"Failed to load Maya1 model from {model_path}.\n"
                f"Error: {e}\n\n"
                f"Please ensure:\n"
                f"1. The model repository contains valid model files, OR\n"
                f"2. Set MAYA_MODEL_ID to a valid HuggingFace model ID, OR\n"
                f"3. Set MAYA_USE_FAKE=1 for testing without a real model"
            )

        # determine device from model parameters (safer than using `model.device`)
        try:
            self.device = next(self.model.parameters()).device
        except StopIteration:
            # fallback to CPU if model has no parameters
            self.device = torch.device("cpu")

        # load SNAC model (audio decoder) 24 kHz
        self.snac = (
            SNAC.from_pretrained("hubertsiuzdak/snac_24khz").eval().to(self.device)
        )
        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:
        {
          "inputs": "text to speak" or {"text": "...", "description": "..."},
          "description": "... optional voice description ...",
          "generation_args": { optional dict for text generation params }
        }
        Returns dict with base64 audio:
        {
          "audio_base64": "<base64-encoded WAV data>",
          "sampling_rate": 24000
        }
        """
        # 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
        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}"}
        
        # 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:
            # generate a 1-second 24kHz sine wave and write a 16-bit WAV using stdlib
            sr = 24000
            duration = 1.0
            n_samples = int(sr * duration)
            freq = 220.0
            # generate samples without numpy
            waveform = [
                int(
                    0.1
                    * 32767
                    * __import__("math").sin(
                        2 * __import__("math").pi * freq * (i / sr)
                    )
                )
                for i in range(n_samples)
            ]

            buf = io.BytesIO()
            with wave.open(buf, "wb") as wf:
                wf.setnchannels(1)
                wf.setsampwidth(2)  # 16-bit
                wf.setframerate(sr)
                # pack samples as little-endian signed 16-bit
                frames = struct.pack("<" + ("h" * len(waveform)), *waveform)
                wf.writeframes(frames)
            wav_bytes = buf.getvalue()
            b64 = base64.b64encode(wav_bytes).decode("utf-8")

            return [{"blob": b64, "content-type": "audio/wav"}]

        # 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_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)

        # Generate tokens
        with torch.inference_mode():
            outputs = self.model.generate(**tokenizer_inputs, **default_gen_args)

        # Extract generated tokens (everything after the input prompt)
        generated_ids = outputs[0, tokenizer_inputs["input_ids"].shape[1]:].tolist()

        # Extract SNAC audio tokens
        snac_tokens = self.extract_snac_codes(generated_ids)
        
        if len(snac_tokens) < 7:
            return {"error": f"Not enough SNAC tokens generated: {len(snac_tokens)}"}

        # 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, audio, 24000, format="WAV")
        wav_bytes = buf.getvalue()
        b64 = base64.b64encode(wav_bytes).decode("utf-8")

        # Return in HuggingFace Inference Endpoints format
        # The endpoint expects a list with blob data
        return [{"blob": b64, "content-type": "audio/wav"}]
# 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.
_HANDLER = None


def init(model_id: str = None):
    """Initialize the module-level handler.
    model_id: optional path or model identifier to pass to EndpointHandler.
    If omitted the handler will look for environment variables `HF_MODEL_ID` or
    `MODEL_ID` and otherwise instantiate with a blank path (which may be
    appropriate when the model files are bundled with the repo).
    """
    global _HANDLER
    if _HANDLER is not None:
        return

    model_path = model_id or os.getenv("HF_MODEL_ID") or os.getenv("MODEL_ID") or ""
    _HANDLER = EndpointHandler(path=model_path)


def predict(payload):
    """Predict/predict wrapper for hosted endpoints.
    Accepts a dict (recommended) or a JSON-ish payload. If a list is passed,
    the first element is used. If a plain string is provided it is treated as
    the `text` field.
    Returns the same dict structure produced by EndpointHandler.__call__.
    """
    global _HANDLER
    if _HANDLER is None:
        init()

    data = payload
    # handle list payloads (common in some platform wrappers)
    if isinstance(payload, (list, tuple)) and len(payload) > 0:
        data = payload[0]

    # allow raw JSON string -> dict
    if isinstance(data, str):
        try:
            import json

            data = json.loads(data)
        except Exception:
            # treat plain string as the text to synthesize
            data = {"text": data}

    if not isinstance(data, dict):
        # best-effort normalization
        data = {"text": str(data)}

    try:
        return _HANDLER(data)
    except Exception as e:
        return {"error": str(e)}


def run(payload):
    """Alias for predict for platforms that expect `run` entrypoint."""
    return predict(payload)

	import base64
	import io
	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=""):
	# allow overriding device and dtype via environment variables for local CPU testing
	# `MAYA_DEVICE`: 'cpu' or 'auto' (default 'auto')
	# `MAYA_TORCH_DTYPE`: 'bf16' or 'fp32' (default 'bf16')
	# `MAYA_USE_FAKE`: if '1', use a tiny fake pipeline for smoke testing (no HF downloads)
	device_override = os.getenv("MAYA_DEVICE", "auto")
	dtype_override = os.getenv("MAYA_TORCH_DTYPE", "bf16")
	use_fake = os.getenv("MAYA_USE_FAKE", "0") == "1"

	# map string to torch dtype will be set after importing torch

	if use_fake:
	# Minimal fake components for quick local smoke tests. Generates a short sine tone.
	self.model = None
	self.tokenizer = None
	# keep device as plain string for fake mode
	self.device = "cpu"
	self.snac = None
	# no external libs required for fake mode
	self.sf = None
	return

	# import heavy inference dependencies lazily so fake-mode doesn't require them
	try:
	import soundfile as sf
	import torch
	from snac import SNAC
	from transformers import AutoModelForCausalLM, AutoTokenizer
	except Exception as e:
	raise RuntimeError(
	f"Failed to import inference dependencies: {e}.\n"
	"Install the packages listed in `inference_endpoints/requirements.txt`."
	)

	# map string to torch dtype
	torch_dtype = torch.bfloat16 if dtype_override == "bf16" else torch.float32

	# load Maya1 text-to-voice model
	# force CPU device_map when requested to avoid trying to use GPUs
	device_map_arg = "auto"
	if device_override == "cpu":
	device_map_arg = "cpu"

	# Allow loading from HuggingFace Hub via environment variable
	# MAYA_MODEL_ID: HuggingFace model ID (e.g., "maya-research/maya1" or similar Maya model)
	model_id = os.getenv("MAYA_MODEL_ID", "")

	# Determine model path: use env var if set, otherwise local path
	local_path = path if path else "/repository"

	# Check if trying to load from self (circular reference)
	# Only block if it's exactly "binhqd/maya1" - other repos ending in /maya1 are fine
	if model_id == "binhqd/maya1":
	print(f"⚠️ MAYA_MODEL_ID is set to {model_id} which is this repository!")
	print("This creates a circular reference. Checking local files instead...")
	model_id = "" # Force local loading

	if model_id:
	# Load from HuggingFace Hub
	model_path = model_id
	print(f"Loading Maya1 model from HuggingFace Hub: {model_id}")
	else:
	# Load from local repository
	model_path = local_path

	# Check if this is a valid model directory
	config_path = os.path.join(model_path, "config.json")
	if not os.path.exists(config_path):
	raise RuntimeError(
	f"❌ Model configuration not found at: {config_path}\n\n"
	f"The repository appears to be missing model weights and configuration.\n\n"
	f"To fix this, you have 3 options:\n\n"
	f"1. 🧪 TESTING MODE (Quick Start):\n"
	f" Set: MAYA_USE_FAKE=1\n"
	f" This generates test audio (sine wave) without requiring model files.\n\n"
	f"2. 🌐 USE EXISTING MODEL FROM HUGGINGFACE:\n"
	f" Set: MAYA_MODEL_ID=<actual-model-id>\n"
	f" Example: MAYA_MODEL_ID=maya-research/maya1\n"
	f" (Must be a valid HuggingFace model repository with weights)\n\n"
	f"3. 📦 UPLOAD YOUR MODEL FILES:\n"
	f" Upload these files to your repository (binhqd/maya1):\n"
	f" - config.json\n"
	f" - model.safetensors (or pytorch_model.bin)\n"
	f" - tokenizer.json\n"
	f" - tokenizer_config.json\n"
	f" - generation_config.json (optional)\n\n"
	f"Current path: {model_path}\n"
	f"Files found: {os.listdir(model_path) if os.path.exists(model_path) else 'path does not exist'}\n\n"
	f"⚠️ Note: MAYA_MODEL_ID=binhqd/maya1 won't work (this repository)\n"
	f" because it doesn't contain model weights.\n"
	)

	try:
	self.model = AutoModelForCausalLM.from_pretrained(
	model_path, torch_dtype=torch_dtype, device_map=device_map_arg
	)
	self.tokenizer = AutoTokenizer.from_pretrained(model_path)
	except Exception as e:
	raise RuntimeError(
	f"Failed to load Maya1 model from {model_path}.\n"
	f"Error: {e}\n\n"
	f"Please ensure:\n"
	f"1. The model repository contains valid model files, OR\n"
	f"2. Set MAYA_MODEL_ID to a valid HuggingFace model ID, OR\n"
	f"3. Set MAYA_USE_FAKE=1 for testing without a real model"
	)

	# determine device from model parameters (safer than using `model.device`)
	try:
	self.device = next(self.model.parameters()).device
	except StopIteration:
	# fallback to CPU if model has no parameters
	self.device = torch.device("cpu")

	# load SNAC model (audio decoder) 24 kHz
	self.snac = (
	SNAC.from_pretrained("hubertsiuzdak/snac_24khz").eval().to(self.device)
	)
	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[i7:(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:
	{
	"inputs": "text to speak" or {"text": "...", "description": "..."},
	"description": "... optional voice description ...",
	"generation_args": { optional dict for text generation params }
	}

	Returns dict with base64 audio:
	{
	"audio_base64": "<base64-encoded WAV data>",
	"sampling_rate": 24000
	}
	"""
	# 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
	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}"}

	# 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:
	# generate a 1-second 24kHz sine wave and write a 16-bit WAV using stdlib
	sr = 24000
	duration = 1.0
	n_samples = int(sr * duration)
	freq = 220.0
	# generate samples without numpy
	waveform = [
	int(
	0.1
	* 32767
	* __import__("math").sin(
	2 * __import__("math").pi * freq * (i / sr)
	)
	)
	for i in range(n_samples)
	]

	buf = io.BytesIO()
	with wave.open(buf, "wb") as wf:
	wf.setnchannels(1)
	wf.setsampwidth(2) # 16-bit
	wf.setframerate(sr)
	# pack samples as little-endian signed 16-bit
	frames = struct.pack("<" + ("h" * len(waveform)), *waveform)
	wf.writeframes(frames)
	wav_bytes = buf.getvalue()
	b64 = base64.b64encode(wav_bytes).decode("utf-8")

	return [{"blob": b64, "content-type": "audio/wav"}]

	# 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_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)

	# Generate tokens
	with torch.inference_mode():
	outputs = self.model.generate(tokenizer_inputs, default_gen_args)

	# Extract generated tokens (everything after the input prompt)
	generated_ids = outputs[0, tokenizer_inputs["input_ids"].shape[1]:].tolist()

	# Extract SNAC audio tokens
	snac_tokens = self.extract_snac_codes(generated_ids)

	if len(snac_tokens) < 7:
	return {"error": f"Not enough SNAC tokens generated: {len(snac_tokens)}"}

	# 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, audio, 24000, format="WAV")
	wav_bytes = buf.getvalue()
	b64 = base64.b64encode(wav_bytes).decode("utf-8")

	# Return in HuggingFace Inference Endpoints format
	# The endpoint expects a list with blob data
	return [{"blob": b64, "content-type": "audio/wav"}]
	# 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.
	_HANDLER = None


	def init(model_id: str = None):
	"""Initialize the module-level handler.

	model_id: optional path or model identifier to pass to EndpointHandler.
	If omitted the handler will look for environment variables `HF_MODEL_ID` or
	`MODEL_ID` and otherwise instantiate with a blank path (which may be
	appropriate when the model files are bundled with the repo).
	"""
	global _HANDLER
	if _HANDLER is not None:
	return

	model_path = model_id or os.getenv("HF_MODEL_ID") or os.getenv("MODEL_ID") or ""
	_HANDLER = EndpointHandler(path=model_path)


	def predict(payload):
	"""Predict/predict wrapper for hosted endpoints.

	Accepts a dict (recommended) or a JSON-ish payload. If a list is passed,
	the first element is used. If a plain string is provided it is treated as
	the `text` field.
	Returns the same dict structure produced by EndpointHandler.__call__.
	"""
	global _HANDLER
	if _HANDLER is None:
	init()

	data = payload
	# handle list payloads (common in some platform wrappers)
	if isinstance(payload, (list, tuple)) and len(payload) > 0:
	data = payload[0]

	# allow raw JSON string -> dict
	if isinstance(data, str):
	try:
	import json

	data = json.loads(data)
	except Exception:
	# treat plain string as the text to synthesize
	data = {"text": data}

	if not isinstance(data, dict):
	# best-effort normalization
	data = {"text": str(data)}

	try:
	return _HANDLER(data)
	except Exception as e:
	return {"error": str(e)}


	def run(payload):
	"""Alias for predict for platforms that expect `run` entrypoint."""
	return predict(payload)