audio-demo/webui/app.py

"""
Pipeline studio: Demucs (strip) -> AuK (edit voice, with preview/undo) ->
MOSS-TTS (generate speech in the edited voice).

Orchestrates across the three existing isolated venvs (demucs/, auk/, tts/)
via subprocess, rather than merging their conflicting dependencies
(MOSS-TTS needs transformers 5.0.0, AuK needs 4.57.6) into one environment.

Run with: webui/.venv/bin/python webui/app.py
"""

import json
import os
import subprocess
import tempfile
import time
import urllib.error
import urllib.request
from pathlib import Path

import gradio as gr

REPO_ROOT = Path(__file__).resolve().parent.parent
DEMUCS_DIR = REPO_ROOT / "demucs"
AUK_DIR = REPO_ROOT / "auk"
TTS_DIR = REPO_ROOT / "tts"

DEMUCS_PYTHON = DEMUCS_DIR / ".venv" / "bin" / "python3"
AUK_INFER = AUK_DIR / ".venv" / "bin" / "auk-infer"
AUK_CKPT = AUK_DIR / "ckpts" / "AuK" / "auk_base.safetensors"
AUK_QWEN_PATH = AUK_DIR / "ckpts" / "Qwen2.5-Omni-3B"
TTS_PYTHON = TTS_DIR / ".venv" / "bin" / "python3"
TTS_GENERATE = TTS_DIR / "generate.py"

WORKDIR = Path(tempfile.mkdtemp(prefix="audio-demo-webui-"))

EMOTIONS = ["happy", "angry", "sad", "fearful", "surprised", "disgusted", "calm", "excited"]
PITCH_OPTIONS = ["-3", "-2", "-1", "+1", "+2", "+3"]
SPEED_OPTIONS = ["0.5", "0.75", "1.25", "1.5", "2.0"]
LANGUAGES = [
    "English", "Spanish", "French", "German", "Portuguese", "Italian",
    "Chinese", "Japanese", "Korean", "Russian",
]

COMFYUI_BASE = "http://localhost:8188"


def _comfyui_queue_empty() -> bool:
    try:
        with urllib.request.urlopen(f"{COMFYUI_BASE}/queue", timeout=3) as resp:
            data = json.loads(resp.read())
            return not data.get("queue_running") and not data.get("queue_pending")
    except (urllib.error.URLError, OSError):
        return True  # ComfyUI not reachable -> nothing to worry about


def _free_comfyui_memory() -> None:
    if not _comfyui_queue_empty():
        return
    try:
        req = urllib.request.Request(
            f"{COMFYUI_BASE}/free",
            data=json.dumps({"unload_models": True, "free_memory": True}).encode(),
            headers={"Content-Type": "application/json"},
            method="POST",
        )
        urllib.request.urlopen(req, timeout=5)
    except (urllib.error.URLError, OSError):
        pass


def _run(cmd: list[str], cwd: Path, env: dict | None = None) -> None:
    result = subprocess.run(
        cmd, cwd=cwd, env=env, capture_output=True, text=True, timeout=1800
    )
    if result.returncode != 0:
        tail = "\n".join(result.stderr.strip().splitlines()[-25:])
        raise gr.Error(f"Command failed ({cmd[0]}):\n{tail}")


def _ld_library_path_for_venv(venv_dir: Path) -> str:
    site_packages = next(venv_dir.glob(".venv/lib/python*/site-packages"))
    nvidia_dir = site_packages / "nvidia"
    if not nvidia_dir.is_dir():
        return ""
    lib_dirs = sorted(p for p in nvidia_dir.glob("*/lib") if p.is_dir())
    return ":".join(str(p) for p in lib_dirs)


TTS_LD_LIBRARY_PATH = _ld_library_path_for_venv(TTS_DIR)


# ---------------------------------------------------------------------------
# Stage 1: input / demucs stripping
# ---------------------------------------------------------------------------


def strip_vocals(input_audio: str):
    if not input_audio:
        raise gr.Error("Upload or record some audio first.")
    _free_comfyui_memory()
    out_dir = WORKDIR / f"demucs_{int(time.time())}"
    out_dir.mkdir(parents=True, exist_ok=True)
    _run(
        [
            str(DEMUCS_PYTHON), "-m", "demucs",
            "--two-stems=vocals", "-n", "htdemucs",
            "-o", str(out_dir), input_audio,
        ],
        cwd=DEMUCS_DIR,
    )
    stem_name = Path(input_audio).stem
    vocals = out_dir / "htdemucs" / stem_name / "vocals.wav"
    no_vocals = out_dir / "htdemucs" / stem_name / "no_vocals.wav"
    return str(vocals), str(no_vocals)


# ---------------------------------------------------------------------------
# Stage 2: AuK voice editing, with history for undo/reset
# ---------------------------------------------------------------------------


def use_as_reference(audio_path: str, history: list):
    if not audio_path:
        raise gr.Error("No audio to use yet.")
    history = [audio_path]
    return audio_path, history, _history_label(history)


def _history_label(history: list) -> str:
    if not history:
        return "No reference set yet."
    return f"Edit history: {len(history) - 1} edit(s) applied on top of the original reference."


def apply_edit(current_ref: str, history: list, emotion: str, pitch: str, speed: str, custom_instruction: str):
    if not current_ref:
        raise gr.Error("Set a voice reference first (upload, strip with Demucs, or use a previous edit).")

    instructions = []
    if emotion and emotion != "none":
        instructions.append(f"Change the emotion to {emotion}.")
    if pitch and pitch != "none":
        direction = "Raise" if pitch.startswith("+") else "Lower"
        instructions.append(f"{direction} the pitch by {pitch.lstrip('+-')} semitone(s).")
    if speed and speed != "none":
        instructions.append(f"Adjust the speech speed to {speed}x.")
    if custom_instruction and custom_instruction.strip():
        instructions.append(custom_instruction.strip())

    if not instructions:
        raise gr.Error("Pick at least one edit (emotion / pitch / speed / custom instruction).")

    _free_comfyui_memory()
    ref = current_ref
    out_dir = WORKDIR / f"auk_{int(time.time())}"
    out_dir.mkdir(parents=True, exist_ok=True)
    for i, instruction in enumerate(instructions):
        out_path = out_dir / f"edit_{i}.wav"
        cmd = [
            str(AUK_INFER),
            "--ckpt", str(AUK_CKPT),
            "--qwen_path", str(AUK_QWEN_PATH),
            "--audio", ref,
            "--instruction", instruction,
            "--output", str(out_path),
            "--cpu_offload",
        ]
        _run(cmd, cwd=AUK_DIR)
        ref = str(out_path)

    history = history + [ref]
    return ref, history, _history_label(history)


def undo_edit(history: list):
    if len(history) <= 1:
        raise gr.Error("Nothing to undo.")
    history = history[:-1]
    return history[-1], history, _history_label(history)


def reset_to_original(history: list):
    if not history:
        raise gr.Error("No reference set yet.")
    history = [history[0]]
    return history[0], history, _history_label(history)


# ---------------------------------------------------------------------------
# Stage 3: MOSS-TTS generation
# ---------------------------------------------------------------------------


def generate_speech(reference: str, text: str, language: str):
    if not reference:
        raise gr.Error("Set a voice reference first (upload, strip, and/or edit it above).")
    if not text or not text.strip():
        raise gr.Error("Enter some text to synthesize.")

    _free_comfyui_memory()
    out_path = WORKDIR / f"tts_{int(time.time())}.wav"
    env = None
    if TTS_LD_LIBRARY_PATH:
        env = os.environ.copy()
        env["LD_LIBRARY_PATH"] = TTS_LD_LIBRARY_PATH

    _run(
        [
            str(TTS_PYTHON), str(TTS_GENERATE),
            "--text", text,
            "--reference", reference,
            "--language", language,
            "--output", str(out_path),
        ],
        cwd=TTS_DIR,
        env=env,
    )
    return str(out_path)


# ---------------------------------------------------------------------------
# UI
# ---------------------------------------------------------------------------

with gr.Blocks(title="Audio-demo pipeline studio") as demo:
    gr.Markdown(
        "# Pipeline studio\n"
        "Strip a voice with **Demucs**, edit it with **AuK**, then generate new "
        "speech in that voice with **MOSS-TTS**. Each stage shells out to its own "
        "isolated venv, so this box's GPU is the real constraint — heavy edits "
        "(long clips) can OOM AuK; keep reference clips short (under ~25s)."
    )

    reference_state = gr.State(None)
    history_state = gr.State([])

    with gr.Row():
        with gr.Column():
            gr.Markdown("## 1. Input")
            input_audio = gr.Audio(label="Upload or record", type="filepath")
            strip_btn = gr.Button("Strip vocals with Demucs")
            with gr.Row():
                vocals_out = gr.Audio(label="Vocals (isolated)", type="filepath")
                no_vocals_out = gr.Audio(label="Background (no vocals)", type="filepath")
            with gr.Row():
                use_vocals_btn = gr.Button("Use vocals as reference")
                use_no_vocals_btn = gr.Button("Use background as reference")
                use_original_btn = gr.Button("Use uploaded audio as-is")

        with gr.Column():
            gr.Markdown("## 2. Voice editor (AuK)")
            current_ref_audio = gr.Audio(label="Current reference", type="filepath", interactive=False)
            history_label = gr.Markdown("No reference set yet.")
            emotion_dd = gr.Dropdown(["none"] + EMOTIONS, value="none", label="Emotion")
            pitch_dd = gr.Dropdown(["none"] + PITCH_OPTIONS, value="none", label="Pitch (semitones)")
            speed_dd = gr.Dropdown(["none"] + SPEED_OPTIONS, value="none", label="Speed")
            custom_instruction = gr.Textbox(
                label="Custom instruction (advanced, optional)",
                placeholder='e.g. "Lower the pitch very slightly."',
            )
            gr.Markdown(
                "*Note: pitch editing has been unreliable in testing — even "
                "\"1 semitone\" has overshot to ~4+ semitones on some voices. "
                "Preview before committing.*"
            )
            apply_btn = gr.Button("Apply edit(s)", variant="primary")
            with gr.Row():
                undo_btn = gr.Button("Undo last edit")
                reset_btn = gr.Button("Reset to original")

        with gr.Column():
            gr.Markdown("## 3. Generate (MOSS-TTS)")
            gen_text = gr.Textbox(label="Text to speak", lines=6)
            gen_language = gr.Dropdown(LANGUAGES, value="English", label="Language")
            generate_btn = gr.Button("Generate speech", variant="primary")
            generated_audio = gr.Audio(label="Generated speech", type="filepath")

    strip_btn.click(strip_vocals, inputs=[input_audio], outputs=[vocals_out, no_vocals_out])
    use_vocals_btn.click(
        use_as_reference, inputs=[vocals_out, history_state],
        outputs=[current_ref_audio, history_state, history_label],
    )
    use_no_vocals_btn.click(
        use_as_reference, inputs=[no_vocals_out, history_state],
        outputs=[current_ref_audio, history_state, history_label],
    )
    use_original_btn.click(
        use_as_reference, inputs=[input_audio, history_state],
        outputs=[current_ref_audio, history_state, history_label],
    )
    apply_btn.click(
        apply_edit,
        inputs=[current_ref_audio, history_state, emotion_dd, pitch_dd, speed_dd, custom_instruction],
        outputs=[current_ref_audio, history_state, history_label],
    )
    undo_btn.click(
        undo_edit, inputs=[history_state],
        outputs=[current_ref_audio, history_state, history_label],
    )
    reset_btn.click(
        reset_to_original, inputs=[history_state],
        outputs=[current_ref_audio, history_state, history_label],
    )
    generate_btn.click(
        generate_speech, inputs=[current_ref_audio, gen_text, gen_language],
        outputs=[generated_audio],
    )

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)