webui/app.py

@@ -45,10 +45,14 @@ 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_FLASH_CKPT = AUK_DIR / "ckpts" / "AuK-Flash" / "auk_flash.safetensors"
 AUK_QWEN_PATH = AUK_DIR / "ckpts" / "Qwen2.5-Omni-3B"
 TTS_PYTHON = TTS_DIR / ".venv" / "bin" / "python3"
 TTS_GENERATE = TTS_DIR / "generate.py"
 
+MODEL_BASE = "Base (full quality)"
+MODEL_FLASH = "Flash (fast draft, 4-step)"
+
 # The-service-station checkout on this same box - where character cards
 # live. Override with the env var if running this UI somewhere else.
 SERVICE_STATION_DIR = Path(os.environ.get("SERVICE_STATION_REPO", "/root/projects/the-service-station"))
@@ -63,6 +67,7 @@ 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"]
+VOLUME_OPTIONS = ["+5", "+10", "+15", "-5", "-10", "-15"]
 LANGUAGES = [
     "English", "Spanish", "French", "German", "Portuguese", "Italian",
     "Chinese", "Japanese", "Korean", "Russian",
@@ -248,15 +253,37 @@ ENHANCE_INSTRUCTIONS = {
 }
 
 
+def _auk_model_args(model_choice: str) -> list[str]:
+    """AuK-Flash is the same architecture distilled to 4 fixed diffusion
+    steps with CFG=0 (vendor README) - same --qwen_path, different --ckpt
+    plus --nfe/--cfg. VRAM footprint is essentially identical to Base
+    (measured in the README), so no extra offload concerns switching."""
+    if model_choice == MODEL_FLASH:
+        if not AUK_FLASH_CKPT.exists():
+            raise gr.Error(
+                f"AuK-Flash isn't downloaded yet ({AUK_FLASH_CKPT} not found) - "
+                f"run `hf download tencent/AuK-Flash --local-dir ./ckpts/AuK-Flash` "
+                f"from the auk/ directory first, or switch back to Base."
+            )
+        return ["--ckpt", str(AUK_FLASH_CKPT), "--qwen_path", str(AUK_QWEN_PATH), "--nfe", "4", "--cfg", "0"]
+    return ["--ckpt", str(AUK_CKPT), "--qwen_path", str(AUK_QWEN_PATH)]
+
+
 def apply_edit(
     current_ref: str,
     history: list,
+    model_choice: str,
     enhance: str,
     timbre_description: str,
     emotion: str,
     other_emotion: str,
     pitch: str,
     speed: str,
+    deaccent: bool,
+    volume: str,
+    whisper_direction: str,
+    nonverbal_instruction: str,
+    nonverbal_gen_seconds: float,
     custom_instruction: str,
 ):
     if not current_ref:
@@ -284,15 +311,32 @@ def apply_edit(
     if pitch and pitch != "none":
         direction = "Raise" if pitch.startswith("+") else "Lower"
         instructions.append(f"{direction} the pitch by {pitch.lstrip('+-')} semitone(s).")
-    speed_gen_seconds = None
     if speed and speed != "none":
         speed_gen_seconds = _duration(current_ref) / float(speed)
         instructions.append((f"Adjust the speech speed to {speed}x.", speed_gen_seconds))
+    if deaccent:
+        instructions.append("Remove the regional accent while preserving the speaker's voice and content.")
+    if volume and volume != "none":
+        direction = "Increase" if volume.startswith("+") else "Decrease"
+        instructions.append(f"{direction} the volume by {volume.lstrip('+-')} dB.")
+    if whisper_direction and whisper_direction != "none":
+        if whisper_direction == "To whisper":
+            instructions.append("Convert this speech into a soft whisper while preserving the speaker and content.")
+        else:
+            instructions.append("Convert this whispered speech into a normal speaking voice while preserving the speaker and content.")
+    if nonverbal_instruction and nonverbal_instruction.strip():
+        # Add/remove changes duration unpredictably (vendor cookbook 4.4 sets
+        # --gen_seconds for both directions, not just "add") - no way to
+        # auto-estimate this one like Speed, the user has to know roughly
+        # what they want. Defaults to the current reference's own duration
+        # if left blank.
+        target_seconds = float(nonverbal_gen_seconds) if nonverbal_gen_seconds else _duration(current_ref)
+        instructions.append((nonverbal_instruction.strip(), target_seconds))
     if custom_instruction and custom_instruction.strip():
         instructions.append(custom_instruction.strip())
 
     if not instructions:
-        raise gr.Error("Pick at least one edit (enhance / emotion / pitch / speed / custom instruction).")
+        raise gr.Error("Pick at least one edit.")
 
     _free_comfyui_memory()
     ref = current_ref
@@ -304,15 +348,10 @@ def apply_edit(
         if isinstance(instruction, tuple):
             instruction, gen_seconds = instruction
         out_path = out_dir / f"edit_{i}.wav"
-        cmd = [
+        cmd = (
-            str(AUK_INFER),
+            [str(AUK_INFER)] + _auk_model_args(model_choice) +
-            "--ckpt", str(AUK_CKPT),
+            ["--audio", ref, "--instruction", instruction, "--output", str(out_path), "--cpu_offload"]
-            "--qwen_path", str(AUK_QWEN_PATH),
+        )
-            "--audio", ref,
-            "--instruction", instruction,
-            "--output", str(out_path),
-            "--cpu_offload",
-        ]
         if gen_seconds is not None:
             cmd += ["--gen_seconds", str(round(gen_seconds, 2))]
         _run(cmd, cwd=AUK_DIR)
@@ -492,7 +531,7 @@ def create_character(
     )
 
 
-def generate_emotion_candidate(character_name: str, emotion_word: str):
+def generate_emotion_candidate(character_name: str, emotion_word: str, model_choice: str):
     if not character_name:
         raise gr.Error("Pick a character first.")
     if not emotion_word or not emotion_word.strip():
@@ -511,11 +550,11 @@ def generate_emotion_candidate(character_name: str, emotion_word: str):
 
     _free_comfyui_memory()
     instruction = auk_emotion_instruction(emotion_word)
-    cmd = [
+    cmd = (
-        str(AUK_INFER), "--ckpt", str(AUK_CKPT), "--qwen_path", str(AUK_QWEN_PATH),
+        [str(AUK_INFER)] + _auk_model_args(model_choice) +
-        "--audio", base_reference, "--instruction", instruction,
+        ["--audio", base_reference, "--instruction", instruction,
-        "--output", str(cache_path), "--cpu_offload",
+         "--output", str(cache_path), "--cpu_offload"]
-    ]
+    )
     _run(cmd, cwd=AUK_DIR)
 
     ref_transcript = _whisper_transcribe(str(cache_path))
@@ -535,8 +574,14 @@ def generate_emotion_candidate(character_name: str, emotion_word: str):
     )
     preview_transcript = _whisper_transcribe(str(preview_path))
 
+    model_note = (
+        " Generated with **Flash** - if this direction is right, regenerate the same "
+        "word with **Base** before promoting it, since Flash trades fidelity for speed "
+        "and this is a candidate you might lock in permanently."
+        if model_choice == MODEL_FLASH else ""
+    )
     status = (
-        f"Saved to `{cache_path.relative_to(SERVICE_STATION_DIR)}` (cached, **not approved**).\n\n"
+        f"Saved to `{cache_path.relative_to(SERVICE_STATION_DIR)}` (cached, **not approved**).{model_note}\n\n"
         f"Reference transcript: {ref_transcript}\n\n"
         f"Preview clone transcript: {preview_transcript}\n\n"
         f"If this sounds right, promote it manually per `.claude/skills/produce-episode/"
@@ -582,6 +627,10 @@ with gr.Blocks(title="Audio-demo pipeline studio") as demo:
                     current_ref_audio = gr.Audio(label="Current reference", type="filepath", interactive=False)
                     ref_transcript_box = gr.Textbox(label="Whisper transcript of current reference", interactive=False)
                     history_label = gr.Markdown("No reference set yet.")
+                    model_dd = gr.Radio(
+                        [MODEL_BASE, MODEL_FLASH], value=MODEL_BASE, label="AuK model",
+                        info="Draft with Flash (fast, 4-step, lower fidelity), confirm with Base before committing a voice or emotion you're actually going to keep.",
+                    )
                     enhance_dd = gr.Dropdown(
                         list(ENHANCE_INSTRUCTIONS.keys()), value="none",
                         label="Enhance (denoise/dereverb/restore) - runs first, before character edits",
@@ -609,10 +658,6 @@ with gr.Blocks(title="Audio-demo pipeline studio") as demo:
                     )
                     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(
                         "*Pitch editing has been unreliable in testing — even \"1 semitone\" "
                         "has overshot to ~4+ semitones on some voices. Prefer a free-form "
@@ -624,6 +669,28 @@ with gr.Blocks(title="Audio-demo pipeline studio") as demo:
                         "source completed in seconds) - trim your reference shorter first if "
                         "using Speed.*"
                     )
+                    with gr.Accordion("More AuK tasks (de-accent, volume, whisper, nonverbal)", open=False):
+                        deaccent_cb = gr.Checkbox(label="De-accent (remove regional accent, keep voice)")
+                        volume_dd = gr.Dropdown(
+                            ["none"] + VOLUME_OPTIONS, value="none", label="Volume (dB)",
+                        )
+                        whisper_dd = gr.Dropdown(
+                            ["none", "To whisper", "From whisper"], value="none",
+                            label="Whisper conversion",
+                        )
+                        nonverbal_tb = gr.Textbox(
+                            label="Nonverbal editing (advanced)",
+                            placeholder='e.g. "Add a cough before \'We tested\'" or "Remove all laughs from the audio."',
+                        )
+                        nonverbal_seconds_nb = gr.Number(
+                            label="Target duration for the nonverbal edit above (seconds)",
+                            value=None,
+                            info="AuK needs an explicit target length for add/remove nonverbal edits - it can't be auto-computed like Speed. Defaults to the current reference's own duration; increase it a little for \"add\", decrease it a little for \"remove\".",
+                        )
+                    custom_instruction = gr.Textbox(
+                        label="Custom instruction (advanced, optional)",
+                        placeholder='e.g. "Lower the pitch very slightly."',
+                    )
                     apply_btn = gr.Button("Apply edit(s)", variant="primary")
                     with gr.Row():
                         undo_btn = gr.Button("Undo last edit")
@@ -659,8 +726,9 @@ with gr.Blocks(title="Audio-demo pipeline studio") as demo:
             )
             apply_btn.click(
                 apply_edit,
-                inputs=[current_ref_audio, history_state, enhance_dd, timbre_tb, emotion_dd, other_emotion_tb,
+                inputs=[current_ref_audio, history_state, model_dd, enhance_dd, timbre_tb, emotion_dd,
-                        pitch_dd, speed_dd, custom_instruction],
+                        other_emotion_tb, pitch_dd, speed_dd, deaccent_cb, volume_dd, whisper_dd,
+                        nonverbal_tb, nonverbal_seconds_nb, custom_instruction],
                 outputs=[current_ref_audio, history_state, history_label, ref_transcript_box],
             )
             undo_btn.click(
@@ -699,6 +767,10 @@ with gr.Blocks(title="Audio-demo pipeline studio") as demo:
             with gr.Row():
                 emotion_char_dd = gr.Dropdown(label="Character", choices=character_names())
                 emotion_word_tb = gr.Textbox(label="Emotion word", placeholder="e.g. surprised, bemused, wistful")
+            emotion_model_dd = gr.Radio(
+                [MODEL_BASE, MODEL_FLASH], value=MODEL_FLASH, label="AuK model",
+                info="Flash defaults here since this is exploring candidates - regenerate the same word with Base before promoting one.",
+            )
             gen_emotion_btn = gr.Button("Generate candidate", variant="primary")
             emotion_ref_audio = gr.Audio(label="Edited reference (candidate)", type="filepath", interactive=False)
             emotion_preview_audio = gr.Audio(label="Preview clone", type="filepath", interactive=False)
@@ -711,7 +783,7 @@ with gr.Blocks(title="Audio-demo pipeline studio") as demo:
             char_dd.change(load_character, inputs=[char_dd], outputs=[char_info, char_audio])
             gen_emotion_btn.click(
                 generate_emotion_candidate,
-                inputs=[emotion_char_dd, emotion_word_tb],
+                inputs=[emotion_char_dd, emotion_word_tb, emotion_model_dd],
                 outputs=[emotion_ref_audio, emotion_preview_audio, emotion_status],
             )