app.py

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)