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
import json
from pathlib import Path
import os
import shutil
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
from src.translation import (
    DEFAULT_LM_STUDIO_API_KEY,
    DEFAULT_LM_STUDIO_BASE_URL,
    DEFAULT_LM_STUDIO_MODEL,
)


BASE_DIR = Path(__file__).resolve().parent
LOG_DIR = BASE_DIR / "logs" / "gradio"
SETTINGS_FILE = BASE_DIR / ".cache" / "web_settings.json"
UPLOAD_DIR = BASE_DIR / ".cache" / "uploads"


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

    id: str
    command: list[str]
    log_path: Path
    env_overrides: dict[str, str] = field(default_factory=dict)
    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 _default_translation_settings() -> dict[str, str]:
    return {
        "base_url": os.getenv("LM_STUDIO_BASE_URL") or DEFAULT_LM_STUDIO_BASE_URL,
        "api_key": os.getenv("LM_STUDIO_API_KEY") or DEFAULT_LM_STUDIO_API_KEY,
        "model": os.getenv("LM_STUDIO_MODEL") or DEFAULT_LM_STUDIO_MODEL,
    }


def load_translation_settings() -> dict[str, str]:
    """Load saved OpenAI-compatible translation settings."""
    settings = _default_translation_settings()
    if not SETTINGS_FILE.exists():
        return settings

    try:
        payload = json.loads(SETTINGS_FILE.read_text(encoding="utf-8"))
    except (OSError, json.JSONDecodeError):
        return settings

    if not isinstance(payload, dict):
        return settings

    for key in settings:
        value = payload.get(key)
        if isinstance(value, str) and value.strip():
            settings[key] = value.strip()
    return settings


def save_translation_settings(base_url: str, api_key: str, model: str) -> tuple[str, str, str, str]:
    """Persist OpenAI-compatible endpoint settings for future web jobs."""
    settings = {
        "base_url": (base_url or "").strip() or DEFAULT_LM_STUDIO_BASE_URL,
        "api_key": (api_key or "").strip() or DEFAULT_LM_STUDIO_API_KEY,
        "model": (model or "").strip() or DEFAULT_LM_STUDIO_MODEL,
    }
    SETTINGS_FILE.parent.mkdir(parents=True, exist_ok=True)
    SETTINGS_FILE.write_text(json.dumps(settings, indent=2), encoding="utf-8")
    return (
        settings["base_url"],
        settings["api_key"],
        settings["model"],
        f"Saved settings to {SETTINGS_FILE}",
    )


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"),
        "--lang",
        args.lang,
        "--mix-mode",
        args.mix_mode,
    ]
    if args.url:
        command.insert(2, args.url)
    if args.input_file:
        command.extend(["--input-file", args.input_file])
    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()
    input_file = (form.get("input_file") or "").strip()
    if not url and not input_file:
        raise ValueError("A YouTube URL or uploaded MP4 is required.")
    if url and input_file:
        raise ValueError("Use either a YouTube URL or uploaded MP4, not both.")

    cli_args = [url] if url else []
    if input_file:
        cli_args.extend(["--input-file", input_file])
    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 _stage_uploaded_mp4(uploaded_file: str | None) -> str:
    if not uploaded_file:
        return ""

    source_path = Path(uploaded_file)
    if source_path.suffix.lower() != ".mp4":
        raise ValueError("Only MP4 uploads are supported.")
    if not source_path.exists():
        raise FileNotFoundError(f"Uploaded file not found: {source_path}")

    safe_stem = "".join(char if char.isalnum() or char in {"-", "_"} else "_" for char in source_path.stem)
    staged_name = f"{uuid.uuid4().hex[:12]}_{safe_stem or 'upload'}.mp4"
    UPLOAD_DIR.mkdir(parents=True, exist_ok=True)
    staged_path = UPLOAD_DIR / staged_name
    shutil.copy2(source_path, staged_path)
    return str(staged_path)


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"
    env.update(job.env_overrides)

    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,
    uploaded_mp4: str | None,
    lang: str,
    whisper_model: str,
    mix_mode: str,
    browser: str,
    cookies: str,
    lmstudio_base_url: str,
    lmstudio_api_key: str,
    lmstudio_model: str,
    gpu: bool,
) -> tuple[str, str, str, gr.Dropdown]:
    saved_settings = load_translation_settings()
    base_url = (lmstudio_base_url or "").strip() or saved_settings["base_url"]
    api_key = (lmstudio_api_key or "").strip() or saved_settings["api_key"]
    model = (lmstudio_model or "").strip() or saved_settings["model"]
    try:
        input_file = _stage_uploaded_mp4(uploaded_mp4)
    except (OSError, ValueError) as exc:
        message = str(exc) or "Invalid uploaded MP4."
        return "", message, message, gr.update(choices=_output_choices())

    form = {
        "url": url,
        "input_file": input_file,
        "lang": lang,
        "whisper_model": whisper_model,
        "mix_mode": mix_mode,
        "browser": browser,
        "cookies": cookies,
        "translation_backend": "lmstudio",
        "lmstudio_base_url": base_url,
        "lmstudio_model": 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",
        env_overrides={
            "LM_STUDIO_BASE_URL": base_url,
            "LM_STUDIO_API_KEY": api_key,
            "LM_STUDIO_MODEL": model,
        },
    )

    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."""
    saved_settings = load_translation_settings()
    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=...")
                uploaded_mp4 = gr.File(
                    label="Upload MP4",
                    file_types=[".mp4"],
                    type="filepath",
                )
                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("OpenAI-Compatible Settings", open=False):
                    lmstudio_base_url = gr.Textbox(
                        label="Endpoint",
                        value=saved_settings["base_url"],
                        placeholder=DEFAULT_LM_STUDIO_BASE_URL,
                    )
                    lmstudio_api_key = gr.Textbox(
                        label="API Key",
                        value=saved_settings["api_key"],
                        type="password",
                    )
                    lmstudio_model = gr.Textbox(
                        label="Model",
                        value=saved_settings["model"],
                        placeholder=DEFAULT_LM_STUDIO_MODEL,
                    )
                    with gr.Row():
                        save_settings = gr.Button("Save Settings")
                        settings_status = gr.Textbox(
                            label="Settings Status",
                            value=f"Loaded from {SETTINGS_FILE if SETTINGS_FILE.exists() else 'environment defaults'}",
                            interactive=False,
                        )
                    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,
            uploaded_mp4,
            lang,
            whisper_model,
            mix_mode,
            browser,
            cookies,
            lmstudio_base_url,
            lmstudio_api_key,
            lmstudio_model,
            gpu,
        ]
        save_settings.click(
            save_translation_settings,
            inputs=[lmstudio_base_url, lmstudio_api_key, lmstudio_model],
            outputs=[lmstudio_base_url, lmstudio_api_key, lmstudio_model, settings_status],
        )
        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__":
    server_name = os.getenv("GRADIO_SERVER_NAME", "127.0.0.1")
    server_port = int(os.getenv("PORT", "7860"))
    app.launch(server_name=server_name, server_port=server_port)