rumik-oss-1

#!/usr/bin/env python3 """Standalone rumik-oss 1 HTTP server with a minimal browser UI.""" import argparse import asyncio import io 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 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 RumikOSS: 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) # Layout comes from the config now; no vocabulary scan. self.speakers = tuple(self.model.config.speakers) @torch.inference_mode() def synthesize(self, req: SpeechRequest) -> bytes: if req.speaker not in self.speakers: raise ValueError(f"speaker must be one of: {', '.join(self.speakers)}") prompt = f'{req.speaker}: {req.input}