youtube-auto-dub/web_app.py

#!/usr/bin/env python3
"""Gradio web UI for launching YouTube Auto Dub jobs."""

from __future__ import annotations

from dataclasses import dataclass, field
from datetime import datetime, timezone
from pathlib import Path
import os
import subprocess
import sys
import threading
import uuid

import gradio as gr

from main import build_parser
from src.audio_separation import DEFAULT_MIX_MODE
from src.engines import OUTPUT_DIR


BASE_DIR = Path(__file__).resolve().parent
LOG_DIR = BASE_DIR / "logs" / "gradio"


@dataclass
class DubJob:
    """Runtime state for a web-launched dub job."""

    id: str
    command: list[str]
    log_path: Path
    created_at: datetime = field(default_factory=lambda: datetime.now(timezone.utc))
    status: str = "queued"
    returncode: int | None = None
    completed_at: datetime | None = None


JOBS: dict[str, DubJob] = {}
JOBS_LOCK = threading.Lock()


def _utc_iso(value: datetime | None) -> str | None:
    if value is None:
        return None
    return value.astimezone(timezone.utc).isoformat()


def build_pipeline_command(form: dict[str, str | bool]) -> list[str]:
    """Build a validated command for the existing CLI pipeline."""
    parser = build_parser()
    args = parser.parse_args(_form_to_cli_args(form))
    command = [
        sys.executable,
        str(BASE_DIR / "main.py"),
        args.url,
        "--lang",
        args.lang,
        "--mix-mode",
        args.mix_mode,
    ]
    if args.translation_backend:
        command.extend(["--translation-backend", args.translation_backend])

    optional_flags = {
        "--browser": args.browser,
        "--cookies": args.cookies,
        "--whisper_model": args.whisper_model,
        "--lmstudio-base-url": args.lmstudio_base_url,
        "--lmstudio-model": args.lmstudio_model,
    }
    for flag, value in optional_flags.items():
        if value:
            command.extend([flag, value])

    if args.gpu:
        command.append("--gpu")

    return command


def _form_to_cli_args(form: dict[str, str | bool]) -> list[str]:
    url = (form.get("url") or "").strip()
    if not url:
        raise ValueError("A YouTube URL is required.")

    cli_args = [url]
    field_flags = {
        "lang": "--lang",
        "browser": "--browser",
        "cookies": "--cookies",
        "whisper_model": "--whisper_model",
        "mix_mode": "--mix-mode",
        "translation_backend": "--translation-backend",
        "lmstudio_base_url": "--lmstudio-base-url",
        "lmstudio_model": "--lmstudio-model",
    }

    defaults = {
        "lang": "es",
        "mix_mode": DEFAULT_MIX_MODE,
        "translation_backend": "lmstudio",
    }

    for field_name, flag in field_flags.items():
        value = (form.get(field_name) or defaults.get(field_name) or "").strip()
        if value:
            cli_args.extend([flag, value])

    gpu_value = form.get("gpu")
    if gpu_value is True or str(gpu_value).lower() in {"1", "true", "on", "yes"}:
        cli_args.append("--gpu")

    return cli_args


def _format_job_status(job: DubJob | None) -> str:
    if job is None:
        return "Ready"

    lines = [
        f"Job: {job.id}",
        f"Status: {job.status}",
        f"Created: {_utc_iso(job.created_at)}",
    ]
    if job.completed_at:
        lines.append(f"Completed: {_utc_iso(job.completed_at)}")
    if job.returncode is not None:
        lines.append(f"Return code: {job.returncode}")
    return "\n".join(lines)


def _read_log_tail(log_path: Path, max_chars: int = 20000) -> str:
    if not log_path.exists():
        return ""
    text = log_path.read_text(encoding="utf-8", errors="replace")
    return text[-max_chars:]


def _run_job(job: DubJob) -> None:
    with JOBS_LOCK:
        job.status = "running"

    env = os.environ.copy()
    env["PYTHONUNBUFFERED"] = "1"

    with job.log_path.open("w", encoding="utf-8", errors="replace") as log_file:
        log_file.write("Gradio started a YouTube Auto Dub job.\n")
        log_file.write(f"Command: {' '.join(job.command)}\n\n")
        log_file.flush()

        process = subprocess.Popen(
            job.command,
            cwd=BASE_DIR,
            env=env,
            stdout=log_file,
            stderr=subprocess.STDOUT,
            text=True,
        )
        returncode = process.wait()

    with JOBS_LOCK:
        job.returncode = returncode
        job.completed_at = datetime.now(timezone.utc)
        job.status = "succeeded" if returncode == 0 else "failed"


def _list_outputs() -> list[Path]:
    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
    return sorted(
        (path for path in OUTPUT_DIR.glob("*") if path.is_file()),
        key=lambda path: path.stat().st_mtime,
        reverse=True,
    )

def _output_choices() -> list[str]:
    return [path.name for path in _list_outputs()[:20]]


def _start_job(
    url: str,
    lang: str,
    whisper_model: str,
    mix_mode: str,
    browser: str,
    cookies: str,
    lmstudio_base_url: str,
    lmstudio_model: str,
    gpu: bool,
) -> tuple[str, str, str, gr.Dropdown]:
    form = {
        "url": url,
        "lang": lang,
        "whisper_model": whisper_model,
        "mix_mode": mix_mode,
        "browser": browser,
        "cookies": cookies,
        "translation_backend": "lmstudio",
        "lmstudio_base_url": lmstudio_base_url,
        "lmstudio_model": lmstudio_model,
        "gpu": gpu,
    }

    try:
        command = build_pipeline_command(form)
    except (SystemExit, ValueError) as exc:
        message = str(exc) or "Invalid job options."
        return "", message, message, gr.update(choices=_output_choices())

    LOG_DIR.mkdir(parents=True, exist_ok=True)
    job_id = uuid.uuid4().hex[:12]
    job = DubJob(
        id=job_id,
        command=command,
        log_path=LOG_DIR / f"{job_id}.log",
    )

    with JOBS_LOCK:
        JOBS[job.id] = job

    thread = threading.Thread(target=_run_job, args=(job,), daemon=True)
    thread.start()
    return job.id, _format_job_status(job), _read_log_tail(job.log_path), gr.update(choices=_output_choices())


def _refresh_job(job_id: str) -> tuple[str, str, gr.Dropdown]:
    with JOBS_LOCK:
        job = JOBS.get(job_id)

    if job is None:
        return "Ready", "No job selected.", gr.update(choices=_output_choices())

    return _format_job_status(job), _read_log_tail(job.log_path), gr.update(choices=_output_choices())


def _select_output(filename: str | None) -> str | None:
    if not filename:
        return None

    output_path = OUTPUT_DIR / filename
    if not output_path.exists() or not output_path.is_file():
        return None
    return str(output_path)


def create_app() -> gr.Blocks:
    """Create the Gradio app."""
    with gr.Blocks(title="Gradio YouTube Auto Dub") as demo:
        gr.Markdown(
            """
            # YouTube Auto Dub
            Start local dubbing jobs, watch the pipeline log, and collect finished videos.
            """
        )
        job_id = gr.State("")

        with gr.Row():
            with gr.Column(scale=5):
                url = gr.Textbox(label="YouTube URL", placeholder="https://www.youtube.com/watch?v=...")
                with gr.Row():
                    lang = gr.Textbox(label="Target Language", value="es", max_lines=1)
                    whisper_model = gr.Dropdown(
                        label="Whisper Model",
                        choices=["", "tiny", "base", "small", "medium", "large-v3"],
                        value="",
                    )
                with gr.Row():
                    mix_mode = gr.Dropdown(
                        label="Mix Mode",
                        choices=[DEFAULT_MIX_MODE, "dub-only", "original-audio"],
                        value=DEFAULT_MIX_MODE,
                    )
                    browser = gr.Dropdown(
                        label="Browser Cookies",
                        choices=["", "chrome", "edge", "firefox", "brave"],
                        value="",
                    )
                cookies = gr.Textbox(label="Cookies File", placeholder=r"C:\path\to\cookies.txt")

                with gr.Accordion("Translation Settings", open=False):
                    lmstudio_base_url = gr.Textbox(
                        label="LM Studio URL",
                        placeholder="http://127.0.0.1:1234/v1",
                    )
                    lmstudio_model = gr.Textbox(label="Model", placeholder="gemma-3-4b-it")
                    gpu = gr.Checkbox(label="Prefer GPU", value=False)

                start = gr.Button("Start Dub", variant="primary")

            with gr.Column(scale=7):
                status = gr.Textbox(label="Job Status", value="Ready", lines=5, interactive=False)
                log = gr.Textbox(
                    label="Run Log",
                    value="No jobs yet.",
                    lines=20,
                    interactive=False,
                )
                refresh = gr.Button("Refresh")

        with gr.Row():
            output_choice = gr.Dropdown(label="Finished Outputs", choices=_output_choices(), interactive=True)
            output_file = gr.File(label="Download Selected Output", interactive=False)

        inputs = [
            url,
            lang,
            whisper_model,
            mix_mode,
            browser,
            cookies,
            lmstudio_base_url,
            lmstudio_model,
            gpu,
        ]
        start.click(
            _start_job,
            inputs=inputs,
            outputs=[job_id, status, log, output_choice],
        )
        refresh.click(_refresh_job, inputs=[job_id], outputs=[status, log, output_choice])
        output_choice.change(_select_output, inputs=[output_choice], outputs=[output_file])

    return demo


app = create_app()


if __name__ == "__main__":
    app.launch(server_name="127.0.0.1", server_port=7860)