import io import re import threading import numpy as np import soundfile as sf import torch import hashlib import random import tempfile import os import uuid from fastapi import FastAPI, Form, UploadFile, File, BackgroundTasks from fastapi.responses import Response, FileResponse from fastapi.middleware.cors import CORSMiddleware from huggingface_hub import snapshot_download from voxcpm import VoxCPM import whisper import uvicorn # ---------------------------------------------------- # 1. Download and Auto-Convert Safetensors # ---------------------------------------------------- print("⏳ Pre-downloading VoxCPM2 weights...") try: # Download snapshot manually snapshot_dir = snapshot_download(repo_id="openbmb/VoxCPM2") print(f"✅ Snapshot downloaded to: {snapshot_dir}") bin_path = os.path.join(snapshot_dir, "pytorch_model.bin") safetensors_path = os.path.join(snapshot_dir, "model.safetensors") # If the library's required .bin file is missing, convert it from .safetensors if not os.path.exists(bin_path) and os.path.exists(safetensors_path): print("🔄 Converting model.safetensors to pytorch_model.bin for compatibility...") from safetensors.torch import load_file state_dict = load_file(safetensors_path) # Wrap the weights in the 'state_dict' key that the library expects torch.save({"state_dict": state_dict}, bin_path) print("✅ Conversion complete!") except Exception as e: print(f"⚠️ Warning during weight pre-download/conversion: {e}") # ---------------------------------------------------- # 2. Load the AI Models into Memory # ---------------------------------------------------- print("⏳ Loading VoxCPM2 Voice Clone Model...") model = VoxCPM.from_pretrained( "openbmb/VoxCPM2", load_denoiser=False, ) print("⏳ Loading Whisper Base model...") whisper_model = whisper.load_model("base") print("✅ AI Models Loaded successfully!") app = FastAPI() app.add_middleware( CORSMiddleware, allow_origins=["*"], allow_credentials=True, allow_methods=["*"], allow_headers=["*"], ) lock = threading.Lock() MAX_CHUNK_CHARS = 180 TASKS = {} def split_text(text, max_chars=MAX_CHUNK_CHARS): text = re.sub(r"\s+", " ", text.strip()) parts = re.split(r"(?<=[\u17d4.!?])\s*", text) chunks = [] current = "" for part in parts: if not part: continue if len(current) + len(part) + 1 <= max_chars: current = (current + " " + part).strip() else: if current: chunks.append(current) current = part if current: chunks.append(current) return chunks @app.get("/") def root(): return { "ok": True, "model": "VoxCPM2", "mode": "Balanced Async", "cuda": torch.cuda.is_available(), } def process_tts(task_id, text, voice_id, reference_wav_path, cfg_value, inference_timesteps): try: if voice_id: seed = int(hashlib.md5(voice_id.encode('utf-8')).hexdigest()[:8], 16) else: seed = 42 random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed) chunks = split_text(text) wavs = [] total_chunks = len(chunks) with lock: for i, chunk in enumerate(chunks): final_text = chunk if reference_wav_path else (f"({voice_id}){chunk}" if voice_id else chunk) print(f"Generating chunk {i+1}/{total_chunks}: {final_text}") wav = model.generate( text=final_text, reference_wav_path=reference_wav_path, cfg_value=cfg_value, inference_timesteps=inference_timesteps, normalize=True, denoise=False, retry_badcase=True, ) wavs.append(np.asarray(wav, dtype=np.float32)) TASKS[task_id]["progress"] = round(((i + 1) / total_chunks) * 100, 2) full_wav = np.concatenate(wavs) temp_result = tempfile.NamedTemporaryFile(delete=False, suffix=".wav") sf.write( temp_result.name, full_wav, model.tts_model.sample_rate, format="WAV" ) temp_result.close() TASKS[task_id]["status"] = "completed" TASKS[task_id]["result_path"] = temp_result.name except Exception as e: print(f"Error: {e}") TASKS[task_id]["status"] = "failed" TASKS[task_id]["error"] = str(e) finally: if reference_wav_path and os.path.exists(reference_wav_path): os.remove(reference_wav_path) @app.post("/tts_job") async def tts_job( background_tasks: BackgroundTasks, text: str = Form(...), voice_id: str = Form(""), gender: str = Form("Auto"), ref_audio: UploadFile = File(None), cfg_value: float = Form(1.8), inference_timesteps: int = Form(8), ): task_id = str(uuid.uuid4()) if gender and gender != "Auto": voice_id = f"{gender}_{voice_id}" reference_wav_path = None if ref_audio is not None: temp_wav = tempfile.NamedTemporaryFile(delete=False, suffix=".wav") content = await ref_audio.read() temp_wav.write(content) temp_wav.close() reference_wav_path = temp_wav.name TASKS[task_id] = { "status": "processing", "progress": 0.0, "result_path": None, "error": None } background_tasks.add_task( process_tts, task_id, text, voice_id, reference_wav_path, cfg_value, inference_timesteps ) return {"task_id": task_id} @app.post("/transcribe") async def transcribe(audio: UploadFile = File(...)): print(f"🎙️ Received Dictation Audio: {audio.filename}") temp_audio = tempfile.NamedTemporaryFile(delete=False, suffix=".m4a") content = await audio.read() temp_audio.write(content) temp_audio.close() try: result = whisper_model.transcribe(temp_audio.name) text = result["text"].strip() print(f"📝 Transcribed: {text}") return {"text": text} except Exception as e: print(f"Error transcribing: {e}") return {"error": str(e)} finally: if os.path.exists(temp_audio.name): os.remove(temp_audio.name) @app.get("/status/{task_id}") def get_status(task_id: str): if task_id not in TASKS: return {"error": "Task not found"} return TASKS[task_id] @app.get("/result/{task_id}") def get_result(task_id: str): if task_id not in TASKS or TASKS[task_id]["status"] != "completed": return {"error": "Result not ready or task failed"} path = TASKS[task_id]["result_path"] return FileResponse(path, media_type="audio/wav") if __name__ == "__main__": uvicorn.run(app, host="0.0.0.0", port=7860)