#!/usr/bin/env python3 """Standalone TinyAya HTTP server with a minimal browser UI.""" import argparse import asyncio import io import re import wave from pathlib import Path import torch from fastapi import FastAPI, HTTPException from fastapi.responses import HTMLResponse, Response from huggingface_hub import snapshot_download from pydantic import BaseModel, Field from transformers import AutoFeatureExtractor, AutoModelForCausalLM, AutoTokenizer, MimiModel import uvicorn AUDIO_RE = re.compile(r"^<(\d+)_(\d+)>$") SPEAKERS = ("Ira", "Aisha", "Siya", "Zoya", "Silver") class SpeechRequest(BaseModel): input: str = Field(min_length=1) speaker: str = "Ira" temperature: float = 0.8 top_k: int = 30 max_new_tokens: int = 2048 class TinyAya: def __init__(self, repo_id: str, device: str): self.device = device self.lock = asyncio.Lock() source = Path(repo_id) root = source if source.exists() else Path(snapshot_download(repo_id)) dtype = torch.bfloat16 if device.startswith("cuda") else torch.float32 self.tokenizer = AutoTokenizer.from_pretrained(root, trust_remote_code=True) self.model = AutoModelForCausalLM.from_pretrained( root, trust_remote_code=True, dtype=dtype, attn_implementation="sdpa", ).eval().to(device) self.mimi = MimiModel.from_pretrained(root / "codec", dtype=dtype).eval().to(device) self.sample_rate = int(AutoFeatureExtractor.from_pretrained(root / "codec").sampling_rate) vocab = self.tokenizer.get_vocab() self.audio_end_id = int(vocab[""]) self.mapping = { int(token_id): (int(match.group(1)), int(match.group(2))) for token, token_id in vocab.items() if (match := AUDIO_RE.match(token)) } self.allowed_ids = torch.tensor( sorted([*self.mapping, self.audio_end_id]), device=device ) @torch.inference_mode() def synthesize(self, req: SpeechRequest) -> bytes: if req.speaker not in SPEAKERS: raise ValueError(f"speaker must be one of: {', '.join(SPEAKERS)}") prompt = f'{req.speaker}: {req.input}