ui-codec.py

#!/usr/bin/env python3
"""Simple Gradio UI for qwen-codec via codec.sh (speaker embedding extraction)."""

import os
import signal
import subprocess
import tempfile
from pathlib import Path

from dotenv import load_dotenv
import gradio as gr

# ── load .env if present ─────────────────────────────────────────────
load_dotenv(Path(__file__).parent / ".env")

SCRIPT = str(Path(__file__).parent / "codec.sh")
EMBEDDINGS_DIR = str(Path(__file__).parent / "embeddings")


def extract(ref_wav_path: str, speaker_name: str, codec_mode: str):
    """Run codec.sh to extract .spk + .rvq embeddings from a reference WAV."""

    if not (Path(__file__).parent / ".env").exists():
        raise gr.Error("Missing .env — copy .env.example to .env and configure it (see README.md).")

    if not ref_wav_path:
        raise gr.Error("Please upload an input file.")

    speaker_name = speaker_name.strip() or "speaker"

    env = os.environ.copy()
    env["QWENTTS_CODEC_MODE"] = codec_mode

    try:
        result = subprocess.run(
            ["bash", SCRIPT, ref_wav_path, EMBEDDINGS_DIR, speaker_name],
            check=True, stdout=subprocess.PIPE, text=True, env=env,
        )
        return f"✅ Extraction complete.\n{result.stdout}"
    except subprocess.CalledProcessError as e:
        raise gr.Error(f"codec.sh failed (exit {e.returncode}). Check terminal for details.")


def list_embeddings():
    """List available .spk files in the embeddings directory."""
    emb_dir = Path(EMBEDDINGS_DIR)
    if not emb_dir.exists():
        return "No embeddings directory found."
    spk_files = sorted(emb_dir.glob("*.spk"))
    if not spk_files:
        return "No embeddings extracted yet."
    lines = ["**Extracted speakers:**"]
    for f in spk_files:
        name = f.stem
        rvq = f.with_suffix(".rvq")
        rvq_info = f" (.rvq present)" if rvq.exists() else " (.rvq MISSING!)"
        lines.append(f"- `{name}`{rvq_info}")
    return "\n".join(lines)


with gr.Blocks(title="Qwen Codec - Embedding Extractor") as demo:
    gr.Markdown("# Qwen Codec — Speaker Embedding Extractor")
    gr.Markdown(
        "Upload a reference WAV to extract the speaker embedding (`.spk`) and "
        "RVQ codes (`.rvq`). The outputs go into `embeddings/` and can be used "
        "with voice cloning via `--ref-spk` / `--ref-rvq`."
    )

    with gr.Row():
        ref_wav = gr.Audio(
            label="Reference WAV",
            type="filepath",
            sources=["upload"],
        )
        speaker_name = gr.Textbox(
            label="Speaker Name",
            placeholder="e.g. freeman, alice, my_voice",
            value="speaker",
        )

    with gr.Accordion("Advanced Options", open=False):
        codec_mode = gr.Dropdown(
            label="Codec Mode",
            choices=["extract", "encode", "decode"],
            value="extract",
        )

    btn = gr.Button("Extract Embeddings", variant="primary")
    stop_btn = gr.Button("Stop", variant="stop")
    output = gr.Textbox(label="Output", lines=6)

    task = btn.click(extract, inputs=[ref_wav, speaker_name, codec_mode], outputs=output)
    stop_btn.click(None, None, None, cancels=[task])

    gr.Markdown("---")
    list_btn = gr.Button("List Saved Embeddings")
    list_output = gr.Markdown()

    list_btn.click(list_embeddings, outputs=list_output)

demo.launch(server_name="127.0.0.1", server_port=7864, share=False)