Add pipeline studio web UI (demucs -> AuK -> MOSS-TTS)

Gradio app that orchestrates the three existing isolated venvs via
subprocess (their dependencies conflict too much to merge into one
environment - MOSS-TTS needs transformers 5.0.0, AuK needs 4.57.6).

Lets you strip a voice, iteratively edit it (emotion/pitch/speed/custom
instruction) with preview and undo/reset after each step, then generate
new speech in the resulting voice - addressing the trial-and-error
workflow the pitch-overshoot investigation needed.

Also extracts tts/generate.py as a reusable, general-purpose MOSS-TTS
cloning CLI (text/reference/language/output args) rather than one-off
personal scripts.

Verified all three stages end to end via direct function calls before
wiring up the UI.

Co-Authored-By: Claude Sonnet 5 <[email protected]>
Claude-Session: https://claude.ai/code/session_01Y1oQ4h4WE4R9KRvMQLr5Ey

README.md

@@ -9,6 +9,7 @@ Audio processing demos: WhisperX, Demucs, and TTS technology proofs-of-concept.
 - `tts/` — text-to-speech service demos (voice cloning via MOSS-TTS)
 - `auk/` — voice *editing* demos (emotion/pitch/speed/content editing on existing recordings,
   as opposed to `tts/`'s clone-and-generate approach)
+- `webui/` — interactive pipeline studio chaining demucs -> AuK -> MOSS-TTS with previews
 
 Each subdirectory has its own README with setup and run instructions.
 

tts/generate.py

@@ -0,0 +1,93 @@
+"""
+Generic MOSS-TTS zero-shot cloning CLI: clone the voice in --reference and
+say --text in it. Used directly, and as the final stage of webui/app.py.
+"""
+
+import argparse
+import importlib.util
+
+import soundfile as sf
+import torch
+from transformers import AutoModel, AutoProcessor
+
+PRETRAINED = "OpenMOSS-Team/MOSS-TTS-Local-Transformer-v1.5"
+
+
+def resolve_attn_implementation(device: str, dtype: torch.dtype) -> str:
+    if (
+        device == "cuda"
+        and importlib.util.find_spec("flash_attn") is not None
+        and dtype in {torch.float16, torch.bfloat16}
+    ):
+        major, _ = torch.cuda.get_device_capability()
+        if major >= 8:
+            return "flash_attention_2"
+    return "sdpa" if device == "cuda" else "eager"
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="MOSS-TTS zero-shot voice cloning")
+    parser.add_argument("--text", required=True, help="Text to synthesize")
+    parser.add_argument("--reference", required=True, help="Reference audio path")
+    parser.add_argument("--output", required=True, help="Output wav path")
+    parser.add_argument("--language", default="English", help="Language tag")
+    parser.add_argument("--audio_temperature", type=float, default=1.7)
+    parser.add_argument("--audio_top_p", type=float, default=0.8)
+    parser.add_argument("--audio_top_k", type=int, default=25)
+    parser.add_argument("--audio_repetition_penalty", type=float, default=1.0)
+    parser.add_argument("--max_new_tokens", type=int, default=4096)
+    args = parser.parse_args()
+
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    dtype = torch.bfloat16 if device == "cuda" else torch.float32
+
+    torch.backends.cuda.enable_cudnn_sdp(False)
+    torch.backends.cuda.enable_flash_sdp(True)
+    torch.backends.cuda.enable_mem_efficient_sdp(True)
+    torch.backends.cuda.enable_math_sdp(True)
+
+    attn_implementation = resolve_attn_implementation(device, dtype)
+    print(f"[INFO] attn_implementation={attn_implementation}")
+
+    processor = AutoProcessor.from_pretrained(PRETRAINED, trust_remote_code=True)
+    processor.audio_tokenizer = processor.audio_tokenizer.to(device)
+
+    model = AutoModel.from_pretrained(
+        PRETRAINED,
+        trust_remote_code=True,
+        attn_implementation=attn_implementation,
+        dtype=dtype,
+    ).to(device)
+    model.eval()
+
+    conversation = [
+        processor.build_user_message(
+            text=args.text,
+            language=args.language,
+            reference=[args.reference],
+        )
+    ]
+
+    with torch.no_grad():
+        batch = processor([conversation], mode="generation")
+        outputs = model.generate(
+            input_ids=batch["input_ids"].to(device),
+            attention_mask=batch["attention_mask"].to(device),
+            max_new_tokens=args.max_new_tokens,
+            do_sample=True,
+            audio_temperature=args.audio_temperature,
+            audio_top_p=args.audio_top_p,
+            audio_top_k=args.audio_top_k,
+            audio_repetition_penalty=args.audio_repetition_penalty,
+        )
+        message = processor.decode(outputs)[0]
+        audio = message.audio_codes_list[0]
+        data = audio.detach().cpu().to(torch.float32).numpy()
+        if data.ndim == 2:
+            data = data.T
+        sf.write(args.output, data, processor.model_config.sampling_rate)
+        print(f"[INFO] Saved {args.output}")
+
+
+if __name__ == "__main__":
+    main()

webui/README.md

@@ -0,0 +1,56 @@
+# webui
+
+A pipeline studio: strip a voice with **demucs**, edit it with **AuK** (emotion, pitch, speed,
+free-form instructions — with preview and undo/reset), then generate new speech in that voice
+with **MOSS-TTS**.
+
+## Why this exists
+
+Manually chaining demucs → AuK → MOSS-TTS across three separate venvs (via SSH one-off scripts)
+worked but made it slow to experiment — especially with AuK's pitch/emotion controls, where you
+really want to preview a result and back out if it's wrong (see the pitch-overshoot findings in
+`auk/README.md`) rather than committing to a specific edit blind.
+
+## Architecture
+
+A single lightweight Gradio app (`app.py`) that **shells out** to each tool's own isolated venv
+via subprocess, rather than merging their dependencies into one environment — MOSS-TTS pins
+`transformers==5.0.0`, AuK needs `4.57.6`; they can't coexist in one venv. This app's own venv
+only needs `gradio` itself.
+
+Each stage's output becomes the next stage's input by filesystem path; a `history` list backs
+the undo/reset buttons in the voice editor.
+
+## Setup
+
+```bash
+python3 -m venv .venv && . .venv/bin/activate
+pip install gradio
+```
+
+Assumes `demucs/`, `auk/`, and `tts/` (siblings in this repo) already have their own working
+venvs and, for `auk/`, downloaded model weights — see their READMEs.
+
+## Running
+
+```bash
+python3 app.py
+```
+
+Serves on `0.0.0.0:7860`. The app automatically calls ComfyUI's `/free` endpoint (if its queue
+is empty) before each heavy step, to reduce GPU contention on a shared box — see the DNS/routing
+and GPU-contention notes earlier in this project's history for why that matters here.
+
+## Known limitations (inherited from the underlying tools)
+
+- **AuK's real usable clip length is ~25-30s**, not several minutes — feeding it a long recording
+  will OOM even with `--cpu_offload` (see `auk/README.md`). This app doesn't chunk long inputs
+  automatically; keep reference clips short.
+- **AuK's pitch-editing instruction-following is unreliable** — even "1 semitone" has measured
+  out at ~4+ semitones of actual shift in testing. The UI surfaces a warning; preview before
+  trusting a pitch edit.
+- Each stage is a blocking subprocess call (model load + inference each time) — the UI will
+  appear to hang for anywhere from ~10s to a couple of minutes per click, with no progress bar
+  yet.
+- No queueing/locking across browser tabs or concurrent users — this is a single-GPU box, so two
+  simultaneous generations will contend for VRAM.

webui/app.py

@@ -0,0 +1,311 @@
+"""
+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)