app.py

"""
FreeVC Voice Converter - HF Space (Docker / FastAPI + Gradio), MULTI-VOICE.

  GET  /api/health    -> {"status":"ok"|"loading","voices":[...]}
  GET  /api/voices    -> {"voices":[...]}
  POST /api/convert   -> converted audio bytes
        multipart: audio=<file>, voice="GS"|"doc"|..., format="mp3"|"wav"
  /                   -> Gradio UI (record/upload, choose voice + format)

Models + voices load in a BACKGROUND THREAD so the web port binds immediately
(avoids HF's slow-start restart loop). Target voices = files in ./voices/.
"""
import os
# Cap CPU threads BEFORE torch/numpy load (HF free CPU = 2 vCPU; otherwise the
# math libs spawn host-many threads and thrash -> ~10x slower inference).
_THREADS = os.environ.get("TORCH_THREADS", "2")
for _v in ("OMP_NUM_THREADS", "MKL_NUM_THREADS", "OPENBLAS_NUM_THREADS",
           "NUMEXPR_NUM_THREADS", "VECLIB_MAXIMUM_THREADS"):
    os.environ.setdefault(_v, _THREADS)

import shutil, subprocess, tempfile, uuid, time, threading, warnings
warnings.filterwarnings("ignore")

OUT_DIR = "outputs"
VOICES_DIR = "voices"
os.makedirs(OUT_DIR, exist_ok=True)
_ready = {"ok": False, "voices": []}


def _load_everything():
    """Heavy startup work: download checkpoints, load models, register voices."""
    from huggingface_hub import hf_hub_download
    os.makedirs("checkpoints", exist_ok=True)
    os.makedirs("wavlm", exist_ok=True)
    for fn, dst in [("freevc.pth", "checkpoints/freevc.pth"),
                    ("WavLM-Large.pt", "wavlm/WavLM-Large.pt")]:
        if not os.path.exists(dst):
            print("downloading", fn, flush=True)
            shutil.copy(hf_hub_download("Pranjal4554/FreeVC", fn), dst)
    import freevc_core as core
    core.init()
    exts = (".wav", ".mp3", ".m4a", ".flac", ".ogg", ".mpeg")
    for fn in sorted(os.listdir(VOICES_DIR)):
        if fn.lower().endswith(exts):
            core.add_voice(os.path.splitext(fn)[0], os.path.join(VOICES_DIR, fn))
    _ready["voices"] = core.voices()
    _ready["ok"] = True
    print("MODELS READY. voices:", core.voices(), flush=True)


threading.Thread(target=_load_everything, daemon=True).start()

import numpy as np, soundfile as sf, gradio as gr
from fastapi import FastAPI, UploadFile, File, Form
from fastapi.responses import FileResponse, JSONResponse
from fastapi.middleware.cors import CORSMiddleware


def _wait_ready(timeout=240):
    t0 = time.time()
    while not _ready["ok"]:
        if time.time() - t0 > timeout:
            return False
        time.sleep(0.5)
    return True


def _match(a, rms, ceil=0.95):
    r = float(np.sqrt(np.mean(a ** 2))) + 1e-9
    a = a * (rms / r)
    pk = float(np.max(np.abs(a)))
    if pk > ceil:
        a = a * (ceil / pk)
    return a.astype(np.float32)


def convert_file(audio_path, voice, out_format="MP3"):
    """Used by both the API and the Gradio UI. Returns output file path."""
    if not audio_path:
        return None
    if not _wait_ready():
        raise RuntimeError("models still loading; try again shortly")
    import freevc_core as core
    if not voice or voice not in core.voices():
        voice = core.voices()[0]
    audio, sr = core.convert(audio_path, voice)
    audio = _match(audio, core.target_rms(voice))
    uid = uuid.uuid4().hex[:8]
    wav = os.path.join(OUT_DIR, f"c_{uid}.wav")
    sf.write(wav, audio, sr)
    fmt = str(out_format).upper()
    if fmt == "WAV":
        return wav
    if fmt in ("OPUS", "OGG"):
        # OGG/Opus mono -> WhatsApp shows it as a real voice note
        opus = os.path.join(OUT_DIR, f"c_{uid}.opus")
        try:
            subprocess.run(["ffmpeg", "-y", "-i", wav, "-c:a", "libopus",
                            "-b:a", "24k", "-ar", "48000", "-ac", "1",
                            "-application", "voip", opus],
                           capture_output=True, check=True)
            return opus
        except Exception:
            return wav
    mp3 = os.path.join(OUT_DIR, f"c_{uid}.mp3")
    try:
        subprocess.run(["ffmpeg", "-y", "-i", wav, "-codec:a", "libmp3lame",
                        "-q:a", "2", mp3], capture_output=True, check=True)
        return mp3
    except Exception:
        return wav


# ---------------- FastAPI (REST for the mobile app) ----------------
app = FastAPI(title="FreeVC Voice Converter")
app.add_middleware(CORSMiddleware, allow_origins=["*"],
                   allow_methods=["*"], allow_headers=["*"])


@app.get("/api/health")
def health():
    return {"status": "ok" if _ready["ok"] else "loading", "voices": _ready["voices"]}


@app.get("/api/voices")
def list_voices():
    return {"voices": _ready["voices"]}


@app.post("/api/convert")
async def api_convert(audio: UploadFile = File(...),
                      voice: str = Form(None), format: str = Form("mp3")):
    if not _ready["ok"] and not _wait_ready(5):
        return JSONResponse({"status": "loading",
                             "error": "server warming up, retry in ~30s"}, status_code=503)
    if voice and voice not in _ready["voices"]:
        return JSONResponse({"error": f"unknown voice '{voice}'",
                             "voices": _ready["voices"]}, status_code=400)
    uid = uuid.uuid4().hex[:8]
    raw = os.path.join(tempfile.gettempdir(), f"in_{uid}")
    with open(raw, "wb") as f:
        f.write(await audio.read())
    src = os.path.join(tempfile.gettempdir(), f"src_{uid}.wav")
    try:
        subprocess.run(["ffmpeg", "-y", "-i", raw, "-ac", "1", "-ar", "16000",
                        "-c:a", "pcm_s16le", src], capture_output=True, check=True)
    except Exception as e:
        return JSONResponse({"error": "could not decode audio: " + str(e)[-200:]},
                            status_code=400)
    chk, _ = sf.read(src, dtype="float32")
    if len(chk) < 16000 * 0.4:
        return JSONResponse({"error": "recording too short - speak >= 1s"}, status_code=400)
    try:
        out = convert_file(src, voice, format.upper())
    except Exception as e:
        return JSONResponse({"error": f"{type(e).__name__}: {e}"}, status_code=500)
    if out.endswith(".mp3"):
        media = "audio/mpeg"
    elif out.endswith(".opus"):
        media = "audio/ogg"
    else:
        media = "audio/wav"
    return FileResponse(out, media_type=media, filename=os.path.basename(out),
                        headers={"X-Voice": voice or (_ready["voices"][0] if _ready["voices"] else "")})


# ---------------- Gradio UI (browser), mounted at / ----------------
def _ui_convert(audio_path, voice, fmt):
    return convert_file(audio_path, voice, fmt)


with gr.Blocks(title="FreeVC Voice Converter") as demo:
    gr.Markdown("# FreeVC Voice Converter\nRecord/upload speech, pick a target "
                "voice and format. Any language. ~5-10s on CPU.")
    with gr.Row():
        inp = gr.Audio(sources=["microphone", "upload"], type="filepath",
                       label="Your voice (record or upload)")
        with gr.Column():
            voice_dd = gr.Dropdown(choices=[], label="Target voice")
            fmt = gr.Radio(["MP3", "WAV", "OPUS"], value="MP3",
                           label="Format (OPUS = WhatsApp voice note)")
            btn = gr.Button("Convert", variant="primary")
    out = gr.Audio(type="filepath", label="Converted to target voice", autoplay=True)
    btn.click(_ui_convert, inputs=[inp, voice_dd, fmt], outputs=out)

    def _populate():
        _wait_ready()
        import freevc_core as core
        vs = core.voices()
        return gr.update(choices=vs, value=(vs[0] if vs else None))
    demo.load(_populate, outputs=voice_dd)

app = gr.mount_gradio_app(app, demo, path="/")