ltx_trainer_gui.py

#!/usr/bin/env python3
"""
LTX-2 / LTX-2.3 LoRA Trainer GUI
================================

A Gradio front-end for the official Lightricks LTX-2 trainer
(https://github.com/Lightricks/LTX-2/tree/main/packages/ltx-trainer).

Features
--------
* Every setting of `LtxTrainerConfig` exposed in the UI (model, LoRA, strategy,
  optimization, acceleration, data, validation, checkpoints, hub, W&B, flow matching).
* Defaults match the official low-VRAM config (`configs/ltx2_av_lora_low_vram.yaml`):
  rank-16 LoRA, adamw8bit, int8-quanto quantization, 8-bit text encoder,
  gradient checkpointing — fits 32 GB GPUs (RTX 5090).
* Save / load / preview training YAML configs.
* Dataset tools: scene splitting, auto-captioning, preprocessing (latent computation).
* One-click model downloads (LTX-2.3 checkpoints + Gemma text encoder).
* Start/stop training (single GPU or multi-GPU via accelerate), live log streaming,
  step/loss/LR progress, loss chart, validation video browser, checkpoint list.

Designed for vast.ai instances but works on any Linux + CUDA box.

Usage:
    python3 ltx_trainer_gui.py [--host 0.0.0.0] [--port 7860] [--share]

Requires: pip install "gradio>=5" pyyaml pandas
"""

from __future__ import annotations

import argparse
import json
import os
import re
import shlex
import shutil
import signal
import subprocess
import threading
import time
from collections import deque
from pathlib import Path

import pandas as pd
import yaml

import gradio as gr

# =============================================================================
# Defaults (vast.ai-friendly paths; override base dir with LTX_WORKSPACE env var)
# =============================================================================

WORKSPACE = os.environ.get("LTX_WORKSPACE", "/workspace")

D = {
    "repo_dir": f"{WORKSPACE}/LTX-2",
    "models_dir": f"{WORKSPACE}/models",
    "model_path": f"{WORKSPACE}/models/ltx-2.3/ltx-2.3-22b-dev.safetensors",
    "text_encoder_path": f"{WORKSPACE}/models/gemma-3-12b-it-qat-q4_0-unquantized",
    "dataset_dir": f"{WORKSPACE}/dataset",
    "dataset_json": f"{WORKSPACE}/dataset/dataset.json",
    "data_root": f"{WORKSPACE}/dataset/.precomputed",
    "output_dir": f"{WORKSPACE}/outputs/my_lora",
    "config_path": f"{WORKSPACE}/outputs/gui_config.yaml",
}

MODEL_PRESETS = {
    "LTX-2.3 dev (recommended for LoRA training)": ("Lightricks/LTX-2.3", "ltx-2.3-22b-dev.safetensors"),
    "LTX-2.3 distilled 1.1": ("Lightricks/LTX-2.3", "ltx-2.3-22b-distilled-1.1.safetensors"),
    "LTX-2.3 dev fp8": ("Lightricks/LTX-2.3-fp8", "ltx-2.3-22b-dev-fp8.safetensors"),
    "LTX-2 (19B) dev": ("Lightricks/LTX-2", "ltx-2-19b-dev.safetensors"),
    "Custom (edit fields below)": ("", ""),
}
GEMMA_REPO = "google/gemma-3-12b-it-qat-q4_0-unquantized"

TARGET_MODULE_CHOICES = [
    "to_k", "to_q", "to_v", "to_out.0",
    "ff.net.0.proj", "ff.net.2",
    "audio_ff.net.0.proj", "audio_ff.net.2",
    "attn1.to_k", "attn1.to_q", "attn1.to_v", "attn1.to_out.0",
    "attn2.to_k", "attn2.to_q", "attn2.to_v", "attn2.to_out.0",
]

DEFAULT_PROMPTS = (
    "A woman with long brown hair sits at a wooden desk in a cozy home office, typing on a laptop "
    "while occasionally glancing at notes beside her. Soft natural light streams through a large "
    "window, casting warm shadows across the room. She pauses to take a sip from a ceramic mug, "
    "then continues working with focused concentration. The audio captures the gentle clicking of "
    "keyboard keys, the soft rustle of papers, and ambient room tone with occasional distant bird "
    "chirps from outside.\n"
    "A chef in a white uniform stands in a professional kitchen, carefully plating a gourmet dish "
    "with precise movements. Steam rises from freshly cooked vegetables as he arranges them with "
    "tweezers. The stainless steel surfaces gleam under bright overhead lights, and various pots "
    "simmer on the stove behind him. The audio features the sizzling of pans, the clinking of "
    "utensils against plates, and the ambient hum of kitchen ventilation."
)

ACCEL_CONFIGS = {
    "Single GPU": None,
    "Multi-GPU DDP": "configs/accelerate/ddp.yaml",
    "Multi-GPU DDP + torch.compile": "configs/accelerate/ddp_compile.yaml",
    "Multi-GPU FSDP": "configs/accelerate/fsdp.yaml",
    "Multi-GPU FSDP + torch.compile": "configs/accelerate/fsdp_compile.yaml",
}

# =============================================================================
# Job manager — background subprocesses with live logs & progress parsing
# =============================================================================

RE_STEP = re.compile(r"Step\s+(\d+)\s*/\s*(\d+)\s*-\s*Loss:\s*([0-9.eE+\-]+),\s*LR:\s*([0-9.eE+\-]+)")
RE_RESUME = re.compile(r"Resuming (?:training )?from step (\d+)")
RE_CKPT = re.compile(r"weights for step (\d+) saved in (\S+)")


class Job:
    def __init__(self, name: str):
        self.name = name
        self.proc: subprocess.Popen | None = None
        self.log: deque[str] = deque(maxlen=5000)
        self.cmd_str = ""
        self.started_at: float | None = None
        self.ended_at: float | None = None
        self.returncode: int | None = None
        # training progress
        self.step = 0
        self.total = 0
        self.loss: float | None = None
        self.lr: float | None = None
        self.history: list[tuple[int, float, float]] = []  # (step, loss, lr)
        self.last_ckpt = ""
        self.lock = threading.Lock()

    @property
    def running(self) -> bool:
        return self.proc is not None and self.proc.poll() is None

    def reset_progress(self, total: int = 0):
        self.step, self.total, self.loss, self.lr = 0, total, None, None
        self.history, self.last_ckpt = [], ""

    def tail(self, n: int = 300) -> str:
        with self.lock:
            return "\n".join(list(self.log)[-n:])


JOBS: dict[str, Job] = {}


def get_job(name: str) -> Job:
    if name not in JOBS:
        JOBS[name] = Job(name)
    return JOBS[name]


def _reader(job: Job):
    proc = job.proc
    assert proc is not None and proc.stdout is not None
    for raw in proc.stdout:
        line = raw.rstrip("\n")
        with job.lock:
            job.log.append(line)
        if job.name == "train":
            m = RE_STEP.search(line)
            if m:
                job.step, job.total = int(m.group(1)), int(m.group(2))
                try:
                    job.loss, job.lr = float(m.group(3)), float(m.group(4))
                    job.history.append((job.step, job.loss, job.lr))
                    if len(job.history) > 4000:
                        job.history = job.history[::2]
                except ValueError:
                    pass
            m = RE_CKPT.search(line)
            if m:
                job.last_ckpt = f"step {m.group(1)} → {m.group(2)}"
            m = RE_RESUME.search(line)
            if m:
                job.step = int(m.group(1))
    proc.wait()
    job.returncode = proc.returncode
    job.ended_at = time.time()
    with job.lock:
        job.log.append(f"--- process exited with code {proc.returncode} ---")


def build_env(hf_token: str = "", wandb_key: str = "", cuda_devices: str = "",
              enable_hf_transfer: bool = False) -> dict:
    env = os.environ.copy()
    env["PYTHONUNBUFFERED"] = "1"
    env["COLUMNS"] = "200"          # stop rich from wrapping log lines we parse
    env["NO_COLOR"] = "1"
    env["PATH"] = f"{Path.home()}/.local/bin:" + env.get("PATH", "")
    if hf_token.strip():
        env["HF_TOKEN"] = hf_token.strip()
        env["HUGGING_FACE_HUB_TOKEN"] = hf_token.strip()
    if wandb_key.strip():
        env["WANDB_API_KEY"] = wandb_key.strip()
    if cuda_devices.strip():
        env["CUDA_VISIBLE_DEVICES"] = cuda_devices.strip()
    if enable_hf_transfer:
        # Only for `hf download` jobs, which run in THIS python env where
        # hf_transfer may actually be installed.
        try:
            import importlib.util
            available = importlib.util.find_spec("hf_transfer") is not None
        except Exception:
            available = False
        env["HF_HUB_ENABLE_HF_TRANSFER"] = "1" if available else "0"
    else:
        # uv-run jobs execute in the trainer's own environment, which usually
        # lacks the hf_transfer package. A leaked HF_HUB_ENABLE_HF_TRANSFER=1
        # (from this GUI or a vast.ai template) makes every HF download inside
        # those jobs crash with "hf_transfer package is not available".
        env["HF_HUB_ENABLE_HF_TRANSFER"] = "0"
    return env


def start_job(name: str, cmd: list[str], cwd: str | None, env: dict) -> str:
    job = get_job(name)
    if job.running:
        return f"⚠️ A '{name}' job is already running. Stop it first."
    if cwd and not Path(cwd).is_dir():
        return f"❌ Working directory not found: {cwd}"
    job.log.clear()
    job.returncode = None
    job.started_at = time.time()
    job.ended_at = None
    job.cmd_str = " ".join(shlex.quote(c) for c in cmd)
    with job.lock:
        job.log.append(f"$ {job.cmd_str}")
        job.log.append(f"(cwd: {cwd or os.getcwd()})")
    try:
        job.proc = subprocess.Popen(
            cmd, cwd=cwd, env=env,
            stdout=subprocess.PIPE, stderr=subprocess.STDOUT,
            text=True, errors="replace", bufsize=1,
            start_new_session=True,
        )
    except FileNotFoundError as e:
        return f"❌ Command not found: {e}"
    except Exception as e:  # noqa: BLE001
        return f"❌ Failed to start: {e}"
    threading.Thread(target=_reader, args=(job,), daemon=True).start()
    return f"🚀 Started: {job.cmd_str}"


def stop_job(name: str) -> str:
    job = get_job(name)
    if not job.running:
        return f"No running '{name}' job."
    try:
        pgid = os.getpgid(job.proc.pid)
        os.killpg(pgid, signal.SIGTERM)
        for _ in range(20):
            if job.proc.poll() is not None:
                return "🛑 Stopped (SIGTERM)."
            time.sleep(0.5)
        os.killpg(pgid, signal.SIGKILL)
        return "🛑 Force-killed (SIGKILL)."
    except Exception as e:  # noqa: BLE001
        return f"⚠️ Error stopping job: {e}"


# =============================================================================
# System info helpers
# =============================================================================

def gpu_info() -> str:
    if not shutil.which("nvidia-smi"):
        return "nvidia-smi not found (no NVIDIA GPU visible)"
    try:
        out = subprocess.run(
            ["nvidia-smi", "--query-gpu=index,name,memory.used,memory.total,utilization.gpu,temperature.gpu",
             "--format=csv,noheader"],
            capture_output=True, text=True, timeout=10,
        ).stdout.strip()
        lines = []
        for row in out.splitlines():
            p = [x.strip() for x in row.split(",")]
            if len(p) >= 6:
                lines.append(f"GPU {p[0]} · {p[1]} · VRAM {p[2]} / {p[3]} · util {p[4]} · {p[5]}°C")
        return "\n".join(lines) or out
    except Exception as e:  # noqa: BLE001
        return f"nvidia-smi error: {e}"


def disk_info(path: str = WORKSPACE) -> str:
    try:
        target = path if Path(path).exists() else "/"
        u = shutil.disk_usage(target)
        return f"Disk ({target}): {u.free / 1e9:.0f} GB free of {u.total / 1e9:.0f} GB"
    except Exception as e:  # noqa: BLE001
        return f"disk check error: {e}"


def sys_info_md() -> str:
    return f"```\n{gpu_info()}\n{disk_info()}\n```"


def trainer_dir(repo_dir: str) -> Path:
    return Path(repo_dir) / "packages" / "ltx-trainer"


def hf_cli() -> str:
    return shutil.which("hf") or shutil.which("huggingface-cli") or "hf"


def env_check(repo_dir: str, model_path: str, te_path: str, data_root: str) -> str:
    checks = []

    def add(ok: bool, label: str, extra: str = ""):
        checks.append(f"{'✅' if ok else '❌'} {label}" + (f" — {extra}" if extra else ""))

    add(shutil.which("git") is not None, "git installed")
    uv = shutil.which("uv") or Path.home().joinpath(".local/bin/uv").exists()
    add(bool(uv), "uv installed", "curl -LsSf https://astral.sh/uv/install.sh | sh")
    add(shutil.which("nvidia-smi") is not None, "NVIDIA GPU visible")
    add(shutil.which("ffmpeg") is not None, "ffmpeg installed (needed for dataset prep)")
    td = trainer_dir(repo_dir)
    add(td.is_dir(), f"trainer repo at {td}", "use 'Clone repo + install deps'")
    add((Path(repo_dir) / ".venv").exists() or (td / ".venv").exists(),
        "uv environment synced (.venv)", "run 'Clone repo + install deps'")
    mp = Path(model_path)
    add(mp.is_file(), f"model checkpoint: {model_path}",
        f"{mp.stat().st_size / 1e9:.1f} GB" if mp.is_file() else "download in section 3")
    add(Path(te_path).is_dir(), f"text encoder dir: {te_path}", "download in section 3")
    add(Path(data_root).is_dir(), f"preprocessed data root: {data_root}",
        "created by the Dataset tab preprocessing step")
    checks.append("")
    checks.append(gpu_info())
    checks.append(disk_info())
    return "```\n" + "\n".join(checks) + "\n```"


# =============================================================================
# Config building  (mirrors LtxTrainerConfig in ltx_trainer/config.py)
# =============================================================================

CONFIG_KEYS: list[str] = []   # ordered registry, filled while building the UI
C: dict[str, object] = {}     # key -> gradio component


def _lines(s) -> list[str]:
    return [ln.strip() for ln in (s or "").splitlines() if ln.strip()]


def _csv(s) -> list[str]:
    return [x.strip() for x in (s or "").split(",") if x.strip()]


def _int(v, default=0) -> int:
    try:
        return int(float(v))
    except (TypeError, ValueError):
        return default


def _parse_json(s: str, label: str, warnings: list[str]) -> dict:
    s = (s or "").strip()
    if not s:
        return {}
    try:
        d = json.loads(s)
        if not isinstance(d, dict):
            raise ValueError("must be a JSON object")
        return d
    except Exception as e:  # noqa: BLE001
        warnings.append(f"⚠️ {label}: invalid JSON ({e}) — using {{}}")
        return {}


def gather_config(v: dict) -> tuple[dict, list[str]]:
    """Build the trainer YAML dict from UI values. Returns (config, warnings)."""
    w: list[str] = []
    cfg: dict = {}

    # ---- model ----
    cfg["model"] = {
        "model_path": v["model.model_path"].strip(),
        "text_encoder_path": v["model.text_encoder_path"].strip() or None,
        "training_mode": v["model.training_mode"],
        "load_checkpoint": v["model.load_checkpoint"].strip() or None,
    }
    if not cfg["model"]["model_path"]:
        w.append("❌ model.model_path is empty")
    elif not Path(cfg["model"]["model_path"]).is_file():
        w.append(f"⚠️ model checkpoint not found on this machine: {cfg['model']['model_path']}")
    if cfg["model"]["text_encoder_path"] and not Path(cfg["model"]["text_encoder_path"]).is_dir():
        w.append(f"⚠️ text encoder dir not found: {cfg['model']['text_encoder_path']}")

    # ---- lora ----
    if v["model.training_mode"] == "lora":
        targets = list(v["_lora.target_modules"] or [])
        for extra in _csv(v["_lora.target_extra"]):
            if extra not in targets:
                targets.append(extra)
        if not targets:
            w.append("❌ LoRA target_modules is empty")
        cfg["lora"] = {
            "rank": _int(v["lora.rank"], 16),
            "alpha": _int(v["lora.alpha"], 16),
            "dropout": float(v["lora.dropout"] or 0.0),
            "target_modules": targets,
        }

    # ---- training strategy ----
    override = (v["_strategy.override"] or "").strip()
    if override:
        try:
            cfg["training_strategy"] = yaml.safe_load(override)
            w.append("ℹ️ Using raw training_strategy YAML override")
        except Exception as e:  # noqa: BLE001
            w.append(f"❌ strategy override YAML invalid: {e}")
            cfg["training_strategy"] = {"name": "text_to_video"}
    elif v["strategy.name"] == "text_to_video":
        cfg["training_strategy"] = {
            "name": "text_to_video",
            "first_frame_conditioning_p": float(v["strategy.first_frame_conditioning_p"]),
            "with_audio": bool(v["strategy.with_audio"]),
            "audio_latents_dir": v["strategy.audio_latents_dir"].strip() or "audio_latents",
        }
    else:  # video_to_video (IC-LoRA)
        cfg["training_strategy"] = {
            "name": "video_to_video",
            "first_frame_conditioning_p": float(v["strategy.first_frame_conditioning_p"]),
            "reference_latents_dir": v["strategy.reference_latents_dir"].strip() or "reference_latents",
        }
        if v["model.training_mode"] != "lora":
            w.append("❌ video_to_video (IC-LoRA) requires training_mode = lora")

    # ---- optimization ----
    sched_params = _parse_json(v["_optimization.scheduler_params"], "scheduler_params", w)
    cfg["optimization"] = {
        "learning_rate": float(v["optimization.learning_rate"]),
        "steps": _int(v["optimization.steps"], 2000),
        "batch_size": _int(v["optimization.batch_size"], 1),
        "gradient_accumulation_steps": _int(v["optimization.gradient_accumulation_steps"], 1),
        "max_grad_norm": float(v["optimization.max_grad_norm"]),
        "optimizer_type": v["optimization.optimizer_type"],
        "scheduler_type": v["optimization.scheduler_type"],
        "scheduler_params": sched_params,
        "enable_gradient_checkpointing": bool(v["optimization.enable_gradient_checkpointing"]),
    }

    # ---- acceleration ----
    quant = v["_acceleration.quantization"]
    quant = None if quant in (None, "", "none (off)") else quant
    cfg["acceleration"] = {
        "mixed_precision_mode": v["acceleration.mixed_precision_mode"],
        "quantization": quant,
        "load_text_encoder_in_8bit": bool(v["acceleration.load_text_encoder_in_8bit"]),
        "offload_optimizer_during_validation": bool(v["acceleration.offload_optimizer_during_validation"]),
    }
    if quant and v["model.training_mode"] == "full":
        w.append("❌ Quantization is not supported with full fine-tuning — set it to none or use LoRA")

    # ---- data ----
    cfg["data"] = {
        "preprocessed_data_root": v["data.preprocessed_data_root"].strip(),
        "num_dataloader_workers": _int(v["data.num_dataloader_workers"], 2),
    }
    if not Path(cfg["data"]["preprocessed_data_root"]).is_dir():
        w.append(f"⚠️ preprocessed data root not found: {cfg['data']['preprocessed_data_root']} "
                 "(run preprocessing in the Dataset tab first)")

    # ---- validation ----
    interval = _int(v["_validation.interval"], 0)
    dims = [_int(v["_validation.dims_w"], 576), _int(v["_validation.dims_h"], 576), _int(v["_validation.dims_f"], 49)]
    if dims[0] % 32 or dims[1] % 32:
        w.append(f"❌ validation video width/height must be divisible by 32 (got {dims[0]}x{dims[1]})")
    if dims[2] % 8 != 1:
        w.append(f"❌ validation frames must satisfy frames % 8 == 1 (got {dims[2]}; try 49, 89, 121)")
    val: dict = {
        "prompts": _lines(v["_validation.prompts"]),
        "negative_prompt": v["validation.negative_prompt"].strip(),
        "video_dims": dims,
        "frame_rate": float(v["validation.frame_rate"]),
        "seed": _int(v["validation.seed"], 42),
        "inference_steps": _int(v["validation.inference_steps"], 30),
        "interval": interval if interval > 0 else None,
        "guidance_scale": float(v["validation.guidance_scale"]),
        "stg_scale": float(v["validation.stg_scale"]),
        "stg_mode": v["validation.stg_mode"],
        "generate_audio": bool(v["validation.generate_audio"]),
        "generate_video": bool(v["validation.generate_video"]),
        "skip_initial_validation": bool(v["validation.skip_initial_validation"]),
    }
    stg_blocks = _csv(v["_validation.stg_blocks"])
    val["stg_blocks"] = [_int(x) for x in stg_blocks] if stg_blocks else None
    images = _lines(v["_validation.images"])
    if images:
        val["images"] = images
        if len(images) != len(val["prompts"]):
            w.append(f"❌ validation images ({len(images)}) must match number of prompts ({len(val['prompts'])})")
    refs = _lines(v["_validation.reference_videos"])
    if refs:
        val["reference_videos"] = refs
        if len(refs) != len(val["prompts"]):
            w.append(f"❌ reference videos ({len(refs)}) must match number of prompts ({len(val['prompts'])})")
    samples_override = (v["_validation.samples_override"] or "").strip()
    if samples_override:
        try:
            val["samples"] = yaml.safe_load(samples_override)
            val["prompts"] = []
            val.pop("images", None)
            val.pop("reference_videos", None)
            w.append("ℹ️ Using advanced validation samples YAML (prompts list ignored)")
        except Exception as e:  # noqa: BLE001
            w.append(f"❌ validation samples YAML invalid: {e}")
    if interval > 0 and not val["prompts"] and not samples_override:
        w.append("⚠️ validation enabled but no prompts given")
    if cfg["training_strategy"].get("name") == "video_to_video" and interval > 0 and not refs and not samples_override:
        w.append("❌ video_to_video strategy requires validation reference_videos (or samples) when validation is on")
    if not val["generate_video"] and not val["generate_audio"]:
        w.append("❌ at least one of generate_video / generate_audio must be enabled")
    cfg["validation"] = val

    # ---- checkpoints ----
    ck_int = _int(v["_checkpoints.interval"], 0)
    cfg["checkpoints"] = {
        "interval": ck_int if ck_int > 0 else None,
        "keep_last_n": _int(v["checkpoints.keep_last_n"], -1),
        "precision": v["checkpoints.precision"],
        "no_resume": bool(v["checkpoints.no_resume"]),
        "save_training_state": v["checkpoints.save_training_state"],
    }

    # ---- flow matching ----
    cfg["flow_matching"] = {
        "timestep_sampling_mode": v["flow_matching.timestep_sampling_mode"],
        "timestep_sampling_params": _parse_json(v["_flow_matching.params"], "timestep_sampling_params", w),
    }

    # ---- hub ----
    cfg["hub"] = {
        "push_to_hub": bool(v["hub.push_to_hub"]),
        "hub_model_id": v["hub.hub_model_id"].strip() or None,
    }
    if cfg["hub"]["push_to_hub"] and not cfg["hub"]["hub_model_id"]:
        w.append("❌ hub_model_id required when push_to_hub is enabled")

    # ---- wandb ----
    cfg["wandb"] = {
        "enabled": bool(v["wandb.enabled"]),
        "project": v["wandb.project"].strip() or "ltx-2-trainer",
        "entity": v["wandb.entity"].strip() or None,
        "tags": _csv(v["_wandb.tags"]),
        "log_validation_videos": bool(v["wandb.log_validation_videos"]),
    }

    # ---- general ----
    cfg["seed"] = _int(v["seed"], 42)
    cfg["output_dir"] = v["output_dir"].strip()
    if not cfg["output_dir"]:
        w.append("❌ output_dir is empty")

    if cfg["optimization"]["batch_size"] > 1:
        w.append("ℹ️ batch_size > 1: only valid if ALL samples share one resolution bucket")

    return cfg, w


def config_to_yaml(cfg: dict) -> str:
    return yaml.safe_dump(cfg, sort_keys=False, allow_unicode=True, width=100)


def apply_config(cfg: dict) -> dict:
    """Map a loaded YAML dict back onto UI values. Returns {key: new_value} (missing = unchanged)."""
    out: dict = {}

    def put(key: str, *path, conv=None):
        node = cfg
        for p in path:
            if not isinstance(node, dict) or p not in node:
                return
            node = node[p]
        out[key] = conv(node) if conv else node

    def s(x):  # none-safe str
        return "" if x is None else str(x)

    put("model.model_path", "model", "model_path", conv=s)
    put("model.text_encoder_path", "model", "text_encoder_path", conv=s)
    put("model.training_mode", "model", "training_mode")
    put("model.load_checkpoint", "model", "load_checkpoint", conv=s)

    lora = cfg.get("lora") or {}
    if lora:
        out["lora.rank"] = lora.get("rank", 16)
        out["lora.alpha"] = lora.get("alpha", 16)
        out["lora.dropout"] = lora.get("dropout", 0.0)
        mods = lora.get("target_modules", [])
        out["_lora.target_modules"] = [m for m in mods if m in TARGET_MODULE_CHOICES]
        out["_lora.target_extra"] = ", ".join(m for m in mods if m not in TARGET_MODULE_CHOICES)

    strat = cfg.get("training_strategy") or {}
    name = strat.get("name", "text_to_video")
    if name in ("text_to_video", "video_to_video"):
        out["strategy.name"] = name
        out["strategy.first_frame_conditioning_p"] = strat.get("first_frame_conditioning_p", 0.5)
        if name == "text_to_video":
            out["strategy.with_audio"] = strat.get("with_audio", False)
            out["strategy.audio_latents_dir"] = strat.get("audio_latents_dir", "audio_latents")
        else:
            out["strategy.reference_latents_dir"] = strat.get("reference_latents_dir", "reference_latents")
        out["_strategy.override"] = ""
    else:  # e.g. "flexible" — dump raw
        out["_strategy.override"] = yaml.safe_dump(strat, sort_keys=False)

    o = cfg.get("optimization") or {}
    for k in ("learning_rate", "steps", "batch_size", "gradient_accumulation_steps", "max_grad_norm",
              "optimizer_type", "scheduler_type", "enable_gradient_checkpointing"):
        if k in o:
            out[f"optimization.{k}"] = o[k]
    if "scheduler_params" in o:
        out["_optimization.scheduler_params"] = json.dumps(o["scheduler_params"] or {})

    a = cfg.get("acceleration") or {}
    if "mixed_precision_mode" in a:
        out["acceleration.mixed_precision_mode"] = a["mixed_precision_mode"]
    if "quantization" in a:
        out["_acceleration.quantization"] = a["quantization"] or "none (off)"
    for k in ("load_text_encoder_in_8bit", "offload_optimizer_during_validation"):
        if k in a:
            out[f"acceleration.{k}"] = a[k]

    d = cfg.get("data") or {}
    if "preprocessed_data_root" in d:
        out["data.preprocessed_data_root"] = s(d["preprocessed_data_root"])
    if "num_dataloader_workers" in d:
        out["data.num_dataloader_workers"] = d["num_dataloader_workers"]

    val = cfg.get("validation") or {}
    if "prompts" in val:
        out["_validation.prompts"] = "\n".join(val["prompts"] or [])
    if "negative_prompt" in val:
        out["validation.negative_prompt"] = s(val["negative_prompt"])
    if "images" in val:
        out["_validation.images"] = "\n".join(val.get("images") or [])
    if "reference_videos" in val:
        out["_validation.reference_videos"] = "\n".join(val.get("reference_videos") or [])
    if "video_dims" in val and val["video_dims"]:
        dims = list(val["video_dims"])
        out["_validation.dims_w"], out["_validation.dims_h"], out["_validation.dims_f"] = dims[0], dims[1], dims[2]
    for k in ("frame_rate", "seed", "inference_steps", "guidance_scale", "stg_scale", "stg_mode",
              "generate_audio", "generate_video", "skip_initial_validation"):
        if k in val:
            out[f"validation.{k}"] = val[k]
    if "interval" in val:
        out["_validation.interval"] = val["interval"] or 0
    if "stg_blocks" in val:
        out["_validation.stg_blocks"] = ", ".join(str(x) for x in (val["stg_blocks"] or []))
    if "samples" in val and val["samples"]:
        out["_validation.samples_override"] = yaml.safe_dump(val["samples"], sort_keys=False)

    ck = cfg.get("checkpoints") or {}
    if "interval" in ck:
        out["_checkpoints.interval"] = ck["interval"] or 0
    for k in ("keep_last_n", "precision", "no_resume", "save_training_state"):
        if k in ck:
            out[f"checkpoints.{k}"] = ck[k]

    fm = cfg.get("flow_matching") or {}
    if "timestep_sampling_mode" in fm:
        out["flow_matching.timestep_sampling_mode"] = fm["timestep_sampling_mode"]
    if "timestep_sampling_params" in fm:
        out["_flow_matching.params"] = json.dumps(fm["timestep_sampling_params"] or {})

    h = cfg.get("hub") or {}
    if "push_to_hub" in h:
        out["hub.push_to_hub"] = h["push_to_hub"]
    if "hub_model_id" in h:
        out["hub.hub_model_id"] = s(h["hub_model_id"])

    wb = cfg.get("wandb") or {}
    for k in ("enabled", "project", "log_validation_videos"):
        if k in wb:
            out[f"wandb.{k}"] = wb[k]
    if "entity" in wb:
        out["wandb.entity"] = s(wb["entity"])
    if "tags" in wb:
        out["_wandb.tags"] = ", ".join(wb["tags"] or [])

    if "seed" in cfg:
        out["seed"] = cfg["seed"]
    if "output_dir" in cfg:
        out["output_dir"] = s(cfg["output_dir"])
    return out


# =============================================================================
# Tab actions
# =============================================================================

def do_save_config(config_path: str, *vals) -> tuple[str, str]:
    v = dict(zip(CONFIG_KEYS, vals))
    cfg, warnings = gather_config(v)
    text = config_to_yaml(cfg)
    msg = []
    path = Path(config_path.strip() or D["config_path"])
    try:
        path.parent.mkdir(parents=True, exist_ok=True)
        path.write_text(text)
        msg.append(f"💾 Saved to {path}")
    except Exception as e:  # noqa: BLE001
        msg.append(f"❌ Could not write {path}: {e}")
    msg.extend(warnings)
    return "\n".join(msg), text


def do_preview_config(*vals) -> tuple[str, str]:
    v = dict(zip(CONFIG_KEYS, vals))
    cfg, warnings = gather_config(v)
    return "\n".join(warnings) or "✅ No issues found", config_to_yaml(cfg)


def do_load_config(file) -> list:
    """Populate UI from an uploaded YAML. Returns updates for [status, *CONFIG_KEYS]."""
    empty = [gr.update() for _ in CONFIG_KEYS]
    if file is None:
        return ["⚠️ No file selected"] + empty
    try:
        cfg = yaml.safe_load(Path(file).read_text())
        assert isinstance(cfg, dict)
    except Exception as e:  # noqa: BLE001
        return [f"❌ Could not parse YAML: {e}"] + empty
    values = apply_config(cfg)
    updates = [gr.update(value=values[k]) if k in values else gr.update() for k in CONFIG_KEYS]
    return [f"✅ Loaded {Path(file).name} — {len(values)} fields applied"] + updates


def do_start_training(repo_dir, config_path, launch_mode, num_processes, cuda_devices,
                      use_uv, hf_token, wandb_key, *vals):
    v = dict(zip(CONFIG_KEYS, vals))
    cfg, warnings = gather_config(v)
    errors = [x for x in warnings if x.startswith("❌")]
    if errors:
        return "Fix these before training:\n" + "\n".join(errors), ""

    td = trainer_dir(repo_dir)
    if not (td / "scripts" / "train.py").is_file():
        return f"❌ trainer not found at {td} — clone the repo in the Setup tab", ""

    path = Path(config_path.strip() or D["config_path"])
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(config_to_yaml(cfg))

    prefix = ["uv", "run"] if use_uv else []
    accel_cfg = ACCEL_CONFIGS.get(launch_mode)
    if accel_cfg is None:
        cmd = [*prefix, "python", "scripts/train.py", str(path), "--disable-progress-bars"]
    else:
        cmd = [*prefix, "accelerate", "launch", "--config_file", accel_cfg,
               "--num_processes", str(_int(num_processes, 2)),
               "scripts/train.py", str(path), "--disable-progress-bars"]

    job = get_job("train")
    job.reset_progress(total=cfg["optimization"]["steps"])
    env = build_env(hf_token, wandb_key, cuda_devices)
    msg = start_job("train", cmd, cwd=str(td), env=env)
    notes = [x for x in warnings if not x.startswith("❌")]
    if notes:
        msg += "\n" + "\n".join(notes)
    msg += ("\nℹ️ First minutes are model loading + quantization (quiet log). "
            "Progress lines appear every 20 steps once training starts.")
    return msg, job.cmd_str


def do_validate_config(repo_dir, config_path, use_uv, hf_token, *vals):
    """Run the trainer's own Pydantic validation on the current settings."""
    v = dict(zip(CONFIG_KEYS, vals))
    cfg, _ = gather_config(v)
    path = Path(config_path.strip() or D["config_path"])
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(config_to_yaml(cfg))
    td = trainer_dir(repo_dir)
    code = (
        "import sys, yaml\n"
        "from ltx_trainer.config import LtxTrainerConfig\n"
        "cfg = yaml.safe_load(open(sys.argv[1]))\n"
        "LtxTrainerConfig(**cfg)\n"
        "print('CONFIG OK - trainer accepts this configuration')\n"
    )
    prefix = ["uv", "run"] if use_uv else []
    cmd = [*prefix, "python", "-c", code, str(path)]
    return start_job("train", cmd, cwd=str(td), env=build_env(hf_token))


def do_stop_training():
    return stop_job("train")


def training_status() -> tuple[str, str]:
    """Returns (progress_html, markdown_line)."""
    job = get_job("train")
    if job.started_at is None:
        return _progress_html(0, 0), "No training started yet."
    pct = (job.step / job.total * 100) if job.total else 0
    state = "🟢 running" if job.running else (
        "✅ finished" if job.returncode == 0 else f"🔴 exited (code {job.returncode})")
    elapsed = (job.ended_at or time.time()) - job.started_at
    eta = ""
    if job.running and job.step > 0 and job.total:
        rate = job.step / max(elapsed, 1)
        remaining = (job.total - job.step) / max(rate, 1e-9)
        eta = f" · ETA {remaining / 3600:.0f}h {(remaining % 3600) / 60:.0f}m"
    parts = [
        f"**{state}**",
        f"step **{job.step} / {job.total}**" if job.total else "",
        f"loss **{job.loss:.4f}**" if job.loss is not None else "",
        f"lr {job.lr:.2e}" if job.lr is not None else "",
        f"elapsed {elapsed / 3600:.0f}h {(elapsed % 3600) / 60:.0f}m{eta}",
        f"last checkpoint: {job.last_ckpt}" if job.last_ckpt else "",
    ]
    return _progress_html(pct, job.step, job.total), " · ".join(p for p in parts if p)


def _progress_html(pct: float, step: int = 0, total: int = 0) -> str:
    label = f"{pct:.1f}% ({step}/{total})" if total else "—"
    return f"""
    <div style="width:100%;background:#e5e7eb;border-radius:8px;height:26px;position:relative;overflow:hidden">
      <div style="width:{min(pct, 100):.1f}%;background:linear-gradient(90deg,#6366f1,#8b5cf6);height:100%;
                  border-radius:8px;transition:width .5s"></div>
      <div style="position:absolute;inset:0;display:flex;align-items:center;justify-content:center;
                  font-size:13px;font-weight:600;color:#111">{label}</div>
    </div>"""


def loss_dataframe() -> pd.DataFrame:
    job = get_job("train")
    if not job.history:
        return pd.DataFrame({"step": [], "loss": []})
    df = pd.DataFrame(job.history, columns=["step", "loss", "lr"])
    if len(df) > 20:  # light smoothing for readability
        df["loss"] = df["loss"].rolling(window=5, min_periods=1).mean()
    return df[["step", "loss"]]


def list_validation_videos(output_dir: str):
    root = Path(output_dir) / "validation_samples"
    vids = sorted(root.rglob("*.mp4"), key=lambda p: p.stat().st_mtime, reverse=True) if root.is_dir() else []
    choices = [str(p) for p in vids[:200]]
    return gr.update(choices=choices, value=choices[0] if choices else None)


def checkpoints_table(output_dir: str) -> pd.DataFrame:
    root = Path(output_dir) / "checkpoints"
    rows = []
    if root.is_dir():
        for p in sorted(root.glob("*.safetensors"), key=lambda p: p.stat().st_mtime, reverse=True):
            st = p.stat()
            rows.append({
                "file": str(p),
                "size (MB)": round(st.st_size / 1e6, 1),
                "modified": time.strftime("%Y-%m-%d %H:%M", time.localtime(st.st_mtime)),
            })
    return pd.DataFrame(rows, columns=["file", "size (MB)", "modified"])


# ---- setup tab actions --------------------------------------------------------

def do_clone_repo(repo_dir: str, hf_token: str):
    repo_dir = repo_dir.strip()
    script = (
        "set -e\n"
        'export PATH="$HOME/.local/bin:$PATH"\n'
        "command -v uv >/dev/null 2>&1 || (curl -LsSf https://astral.sh/uv/install.sh | sh)\n"
        'export PATH="$HOME/.local/bin:$PATH"\n'
        f"if [ ! -d {shlex.quote(repo_dir)} ]; then git clone https://github.com/Lightricks/LTX-2 {shlex.quote(repo_dir)}; fi\n"
        f"cd {shlex.quote(repo_dir)}\n"
        "uv sync\n"
        "echo '--- DONE: repo ready, dependencies installed ---'\n"
    )
    return start_job("setup", ["bash", "-c", script], cwd=None, env=build_env(hf_token))


def do_download_model(preset: str, repo: str, filename: str, dest_dir: str, hf_token: str):
    if not hf_token.strip() and not os.environ.get("HF_TOKEN"):
        return ("⚠️ No HF token set — LTX-2.3 and Gemma are gated repos. Add your token above "
                "(and accept each model's license on huggingface.co first).")
    repo, filename = repo.strip(), filename.strip()
    if not repo:
        return "❌ Model repo is empty"
    cmd = [hf_cli(), "download", repo]
    if filename:
        cmd.append(filename)
    cmd += ["--local-dir", dest_dir.strip()]
    return start_job("download_model", cmd, cwd=None, env=build_env(hf_token, enable_hf_transfer=True))


def do_download_te(dest_dir: str, hf_token: str):
    if not hf_token.strip() and not os.environ.get("HF_TOKEN"):
        return "⚠️ No HF token set — Gemma is a gated repo (accept its license on huggingface.co first)."
    cmd = [hf_cli(), "download", GEMMA_REPO, "--local-dir", dest_dir.strip()]
    return start_job("download_te", cmd, cwd=None, env=build_env(hf_token, enable_hf_transfer=True))


# ---- dataset tab actions ------------------------------------------------------

def do_split_scenes(repo_dir, input_video, out_dir, min_len, use_uv, hf_token):
    td = trainer_dir(repo_dir)
    prefix = ["uv", "run", "python"] if use_uv else ["python"]
    cmd = [*prefix, "scripts/split_scenes.py", input_video.strip(), out_dir.strip()]
    if min_len.strip():
        cmd += ["--filter-shorter-than", min_len.strip()]
    return start_job("dataset", cmd, cwd=str(td), env=build_env(hf_token))


def do_caption(repo_dir, media_dir, out_json, captioner, recursive, override,
               instruction, num_workers, api_key, use_uv, hf_token):
    td = trainer_dir(repo_dir)
    prefix = ["uv", "run", "python"] if use_uv else ["python"]
    cmd = [*prefix, "scripts/caption_videos.py", media_dir.strip(),
           "--output", out_json.strip(), "--captioner-type", captioner]
    if recursive:
        cmd.append("--recursive")
    if override:
        cmd.append("--override")
    if instruction.strip():
        cmd += ["--instruction", instruction.strip()]
    if captioner == "gemini_flash" and _int(num_workers, 1) > 1:
        cmd += ["--num-workers", str(_int(num_workers, 1))]
    if api_key.strip():
        cmd += ["--api-key", api_key.strip()]
    return start_job("dataset", cmd, cwd=str(td), env=build_env(hf_token))


def do_capserver_start(repo_dir, use_uv, hf_token):
    """Launch the local vLLM captioner server (required for qwen_omni)."""
    td = trainer_dir(repo_dir)
    prefix = ["uv", "run", "python"] if use_uv else ["python"]
    cmd = [*prefix, "scripts/serve_captioner.py"]
    return start_job("capserver", cmd, cwd=str(td), env=build_env(hf_token))


def do_preprocess(repo_dir, dataset_file, buckets, model_path, te_path, out_dir,
                  trigger, batch, vae_tiling, decode, skip_audio, te_8bit,
                  remove_prefixes, overwrite, num_gpus, use_uv, hf_token):
    td = trainer_dir(repo_dir)
    n = _int(num_gpus, 1)
    prefix = ["uv", "run"] if use_uv else []
    if n > 1:
        base = [*prefix, "accelerate", "launch", "--num_processes", str(n), "scripts/process_dataset.py"]
    else:
        base = [*prefix, "python", "scripts/process_dataset.py"]
    cmd = [*base, dataset_file.strip(),
           "--resolution-buckets", buckets.strip(),
           "--model-path", model_path.strip(),
           "--text-encoder-path", te_path.strip()]
    if out_dir.strip():
        cmd += ["--output-dir", out_dir.strip()]
    if trigger.strip():
        cmd += ["--lora-trigger", trigger.strip()]
    if _int(batch, 1) != 1:
        cmd += ["--batch-size", str(_int(batch, 1))]
    if vae_tiling:
        cmd.append("--vae-tiling")
    if decode:
        cmd.append("--decode")
    if skip_audio:
        cmd.append("--skip-audio")
    if te_8bit:
        cmd.append("--load-text-encoder-in-8bit")
    if remove_prefixes:
        cmd.append("--remove-llm-prefixes")
    if overwrite:
        cmd.append("--overwrite")
    return start_job("dataset", cmd, cwd=str(td), env=build_env(hf_token))


def dataset_data_root(dataset_file: str, out_dir: str) -> str:
    if out_dir.strip():
        return out_dir.strip()
    return str(Path(dataset_file.strip()).parent / ".precomputed")


# =============================================================================
# vast.ai onstart script (also shipped as vast_onstart.sh next to this file)
# =============================================================================

ONSTART_SCRIPT = r"""#!/bin/bash
# ============ vast.ai on-start script — LTX-2.3 LoRA Trainer GUI ============
# Paste this into the "On-start Script" field of your vast.ai template.
# Optional template env vars:
#   HF_TOKEN   – HuggingFace token (gated models: LTX-2.3 + Gemma)
#   GUI_URL    – direct URL to raw ltx_trainer_gui.py (auto-download)
#   GUI_SHARE  – set to 1 to also expose a public gradio.live link
exec > >(tee -a /workspace/onstart.log) 2>&1
set -x
export DEBIAN_FRONTEND=noninteractive
mkdir -p /workspace && cd /workspace

apt-get update -y || true
apt-get install -y git curl ffmpeg tmux || true

# uv — manages the trainer's python environment
if ! command -v uv >/dev/null 2>&1; then
  curl -LsSf https://astral.sh/uv/install.sh | sh
fi
export PATH="$HOME/.local/bin:$PATH"

# trainer repo + deps
if [ ! -d /workspace/LTX-2 ]; then
  git clone https://github.com/Lightricks/LTX-2 /workspace/LTX-2
fi
cd /workspace/LTX-2 && uv sync && cd /workspace

# GUI deps (system python — separate from the trainer's uv env)
python3 -m pip install --upgrade pip
python3 -m pip install "gradio>=5" pyyaml pandas "huggingface_hub[cli]" hf_transfer

# fetch the GUI file if hosted somewhere and not already present
if [ -n "${GUI_URL:-}" ] && [ ! -f /workspace/ltx_trainer_gui.py ]; then
  curl -fsSL "$GUI_URL" -o /workspace/ltx_trainer_gui.py
fi

# launch GUI on port 7860
if [ -f /workspace/ltx_trainer_gui.py ]; then
  EXTRA=""
  [ "${GUI_SHARE:-0}" = "1" ] && EXTRA="--share"
  pkill -f ltx_trainer_gui.py || true
  nohup python3 /workspace/ltx_trainer_gui.py --host 0.0.0.0 --port 7860 $EXTRA \
    > /workspace/gui.log 2>&1 &
  echo "GUI starting — log: /workspace/gui.log"
else
  echo "NOTE: /workspace/ltx_trainer_gui.py not found — upload it (Jupyter/scp) and run:"
  echo "  python3 /workspace/ltx_trainer_gui.py --host 0.0.0.0 --port 7860"
fi
"""

VAST_GUIDE = f"""
## Running this GUI on vast.ai

### 1 · Template settings

| Setting | Value |
|---|---|
| **Docker image** | vast.ai **"PyTorch (Vast)"** template, or `nvidia/cuda:12.8.1-devel-ubuntu22.04` (RTX 5090 needs CUDA ≥ 12.8; the trainer installs its own PyTorch via `uv sync`) |
| **Launch mode** | Jupyter + SSH (Jupyter makes uploading files & videos easy) |
| **Disk** | **250 GB+** (model 46 GB + Gemma ~25 GB + latents + checkpoints) |
| **Ports** | add `7860` (Docker options: `-p 7860:7860`) |
| **Environment vars** | `HF_TOKEN=hf_...` · optional `GUI_URL=<raw url to ltx_trainer_gui.py>` · optional `GUI_SHARE=1` |
| **On-start script** | paste `vast_onstart.sh` (shown below) |
| **GPU** | ≥ 32 GB VRAM (RTX 5090) for the default low-VRAM config · 48–80 GB (L40S / A100 / H100) lets you raise rank / resolution / disable quantization |

### 2 · Before renting
1. Accept the gated licenses on HuggingFace while logged in: **Lightricks/LTX-2.3** and **{GEMMA_REPO}**.
2. Create a **Read** token at huggingface.co/settings/tokens (enable "read gated repos" for fine-grained tokens).

### 3 · Access the GUI
- Instance card → **Open Ports** → click the mapping for `7860` (e.g. `http://<ip>:<external-port>`).
- Or set `GUI_SHARE=1` and grab the `*.gradio.live` link from `/workspace/gui.log`.
- If you didn't use `GUI_URL`, upload `ltx_trainer_gui.py` to `/workspace` via Jupyter, then:
  `python3 /workspace/ltx_trainer_gui.py --host 0.0.0.0 --port 7860`

### 4 · Typical workflow
1. **Setup tab** → paste HF token → *Clone repo + install deps* → download LTX-2.3 dev + Gemma → *Run environment check*.
2. Upload training videos to `/workspace/dataset/videos/` (Jupyter drag & drop, or `vastai copy` / `scp` / `rclone`).
3. **Dataset tab** → caption (optional) → preprocess (computes latents; audio auto-extracted unless skipped).
4. **Training Settings tab** → set trigger word already baked into captions, adjust steps/rank → *Save config*.
5. **Train tab** → *Start training* → watch progress, loss chart, validation videos.
6. Grab `lora_weights.safetensors` from the checkpoints list (ComfyUI-compatible).

> 💰 Use *interruptible* instances at ~half price for long runs — checkpoints + `load_checkpoint`
> let you resume. Keep checkpoint interval (default 250) and `save_training_state: minimal`.
"""


# =============================================================================
# UI
# =============================================================================

def reg(key: str, comp):
    """Register a config component."""
    CONFIG_KEYS.append(key)
    C[key] = comp
    return comp


# Gradio 6 moved theme/css from Blocks() to launch()
_GRADIO_MAJOR = int(gr.__version__.split(".")[0])
_STYLE = {"theme": gr.themes.Soft(), "css": "footer {display: none !important}"}
_BLOCKS_KW = {"title": "LTX-2.3 LoRA Trainer"}
_LAUNCH_KW = {}
if _GRADIO_MAJOR >= 6:
    _LAUNCH_KW.update(_STYLE)
else:
    _BLOCKS_KW.update(_STYLE)


def create_demo() -> gr.Blocks:
    with gr.Blocks(**_BLOCKS_KW) as demo:
        gr.Markdown(
            "# 🎬 LTX-2.3 LoRA Trainer\n"
            "GUI for the official [Lightricks LTX-2 trainer]"
            "(https://github.com/Lightricks/LTX-2/tree/main/packages/ltx-trainer) — "
            "defaults follow the **low-VRAM config** (32 GB GPUs, e.g. RTX 5090)."
        )

        # ---- shared/global fields ----
        with gr.Row():
            repo_dir = gr.Textbox(label="LTX-2 repo directory", value=D["repo_dir"], scale=2)
            hf_token = gr.Textbox(label="HuggingFace token (gated models)", type="password",
                                  value=os.environ.get("HF_TOKEN", ""), scale=2)
            wandb_key = gr.Textbox(label="W&B API key (optional)", type="password", scale=1)
            use_uv = gr.Checkbox(label="Run via `uv run` (recommended)", value=True, scale=1)

        with gr.Tabs():
            # =================================================================
            # SETUP TAB
            # =================================================================
            with gr.Tab("🛠️ Setup"):
                gr.Markdown("### 1 · Install trainer\nClones `Lightricks/LTX-2` and runs `uv sync`.")
                with gr.Row():
                    clone_btn = gr.Button("📦 Clone repo + install deps", variant="primary")
                    envcheck_btn = gr.Button("🔍 Run environment check")
                setup_status = gr.Textbox(label="Status", interactive=False)
                setup_log = gr.Textbox(label="Setup log", lines=10, max_lines=10,
                                       interactive=False, autoscroll=True)
                envcheck_md = gr.Markdown("")

                gr.Markdown("### 2 · Hardware")
                sysinfo_md = gr.Markdown(sys_info_md())

                gr.Markdown(
                    "### 3 · Download models\n"
                    "⚠️ **Gated repos** — accept the licenses of `Lightricks/LTX-2.3` and "
                    f"`{GEMMA_REPO}` on huggingface.co, then paste a Read token above."
                )
                with gr.Row():
                    model_preset = gr.Dropdown(label="Checkpoint preset", choices=list(MODEL_PRESETS),
                                               value="LTX-2.3 dev (recommended for LoRA training)", scale=2)
                    model_repo = gr.Textbox(label="HF repo", value="Lightricks/LTX-2.3", scale=1)
                    model_file = gr.Textbox(label="File", value="ltx-2.3-22b-dev.safetensors", scale=1)
                    model_dest = gr.Textbox(label="Destination dir", value=f"{D['models_dir']}/ltx-2.3", scale=1)
                    dl_model_btn = gr.Button("⬇️ Download checkpoint", variant="primary", scale=1)
                dl_model_status = gr.Textbox(label="Download status", interactive=False)
                dl_model_log = gr.Textbox(label="Checkpoint download log", lines=6, max_lines=6,
                                          interactive=False, autoscroll=True)
                with gr.Row():
                    te_dest = gr.Textbox(label="Text encoder destination (Gemma)",
                                         value=D["text_encoder_path"], scale=3)
                    dl_te_btn = gr.Button("⬇️ Download Gemma text encoder", variant="primary", scale=1)
                dl_te_status = gr.Textbox(label="Download status", interactive=False)
                dl_te_log = gr.Textbox(label="Text encoder download log", lines=6, max_lines=6,
                                       interactive=False, autoscroll=True)

            # =================================================================
            # DATASET TAB
            # =================================================================
            with gr.Tab("🎞️ Dataset"):
                gr.Markdown(
                    "Prepare your dataset: *(optional)* split long videos → *(optional)* auto-caption → "
                    "**preprocess** (encodes videos/audio/captions into latents the trainer consumes).\n\n"
                    "Upload videos to the instance first (Jupyter, `scp`, `vastai copy`, `rclone`…)."
                )
                with gr.Accordion("✂️ Split long videos into scenes (optional)", open=False):
                    with gr.Row():
                        split_input = gr.Textbox(label="Input video file", value=f"{D['dataset_dir']}/raw.mp4", scale=2)
                        split_out = gr.Textbox(label="Output dir", value=f"{D['dataset_dir']}/videos", scale=2)
                        split_minlen = gr.Textbox(label="Drop scenes shorter than", value="5s", scale=1)
                        split_btn = gr.Button("Run scene split", scale=1)

                with gr.Accordion("📝 Auto-caption videos & images (optional)", open=False):
                    gr.Markdown(
                        "*Works on videos **and** images (jpg/png). `qwen_omni` runs locally via the captioner "
                        "server below — but that downloads **Qwen3-Omni-30B (~60 GB)** and its weights don't fit "
                        "a 32 GB GPU. On a single consumer card use `gemini_flash` (API key, fast, cheap) or "
                        "write your own captions into `dataset.json`. Review captions afterwards — they can "
                        "hallucinate. The trigger word is added at preprocessing, not here.*"
                    )
                    with gr.Row():
                        cap_dir = gr.Textbox(label="Media dir (videos/images)",
                                             value=f"{D['dataset_dir']}/media", scale=2)
                        cap_out = gr.Textbox(label="Output dataset JSON", value=D["dataset_json"], scale=2)
                        cap_type = gr.Dropdown(label="Captioner", choices=["qwen_omni", "gemini_flash"],
                                               value="qwen_omni", scale=1)
                    with gr.Row():
                        cap_recursive = gr.Checkbox(label="Recurse subdirs", value=False)
                        cap_override = gr.Checkbox(label="Re-caption existing", value=False)
                        cap_workers = gr.Number(label="Workers (gemini only)", value=1, precision=0)
                        cap_apikey = gr.Textbox(label="Gemini API key (gemini_flash only)", type="password", scale=2)
                        cap_btn = gr.Button("Run captioning", scale=1)
                    cap_instruction = gr.Textbox(label="Custom caption instruction (optional)", value="")
                    with gr.Row():
                        capserver_start_btn = gr.Button("▶️ Start captioner server (needed for qwen_omni)")
                        capserver_stop_btn = gr.Button("🛑 Stop captioner server", variant="stop")
                    capserver_log = gr.Textbox(label="Captioner server log", lines=5, max_lines=5,
                                               interactive=False, autoscroll=True)

                gr.Markdown("### ⚡ Preprocess dataset (required)")
                with gr.Row():
                    pp_dataset = gr.Textbox(label="Dataset file (JSON/JSONL/CSV with caption + video columns)",
                                            value=D["dataset_json"], scale=2)
                    pp_buckets = gr.Textbox(label='Resolution buckets "WxHxF" (`;`-separated, W/H %32==0, F%8==1)',
                                            value="576x576x49", scale=2)
                    pp_trigger = gr.Textbox(label="LoRA trigger word (prepended to captions)", value="", scale=1)
                with gr.Row():
                    pp_model = gr.Textbox(label="Model checkpoint", value=D["model_path"], scale=2)
                    pp_te = gr.Textbox(label="Text encoder dir", value=D["text_encoder_path"], scale=2)
                    pp_outdir = gr.Textbox(label="Output dir (blank = <dataset dir>/.precomputed)", value="", scale=1)
                with gr.Row():
                    pp_te8bit = gr.Checkbox(label="Load text encoder in 8-bit", value=True)
                    pp_skipaudio = gr.Checkbox(label="Skip audio extraction", value=False,
                                               info="audio is auto-extracted from videos by default")
                    pp_tiling = gr.Checkbox(label="VAE tiling (large resolutions)", value=False)
                    pp_decode = gr.Checkbox(label="Decode latents for verification", value=False)
                    pp_removepfx = gr.Checkbox(label="Remove LLM caption prefixes", value=False)
                    pp_overwrite = gr.Checkbox(label="Overwrite existing outputs", value=False)
                with gr.Row():
                    pp_batch = gr.Number(label="Batch size", value=1, precision=0, scale=1)
                    pp_gpus = gr.Number(label="GPUs (accelerate)", value=1, precision=0, scale=1)
                    pp_btn = gr.Button("🚀 Run preprocessing", variant="primary", scale=2)
                    pp_use_btn = gr.Button("→ use result as training data root", scale=2)
                dataset_status = gr.Textbox(label="Status", interactive=False)
                dataset_log = gr.Textbox(label="Dataset job log", lines=16, max_lines=16,
                                         interactive=False, autoscroll=True)
                dataset_stop_btn = gr.Button("🛑 Stop dataset job", variant="stop")

            # =================================================================
            # TRAINING SETTINGS TAB
            # =================================================================
            with gr.Tab("⚙️ Training Settings"):
                gr.Markdown("Defaults = official **low-VRAM** config (32 GB GPU). "
                            "Every `LtxTrainerConfig` option is exposed below.")

                with gr.Row():
                    config_path = gr.Textbox(label="Config YAML path", value=D["config_path"], scale=3)
                    save_cfg_btn = gr.Button("💾 Save config", variant="primary", scale=1)
                    preview_cfg_btn = gr.Button("👁️ Preview YAML", scale=1)
                with gr.Row():
                    load_cfg_file = gr.File(label="Load existing config YAML", file_types=[".yaml", ".yml"], scale=3)
                    load_cfg_btn = gr.Button("📂 Load into UI", scale=1)
                with gr.Row():
                    img_preset_btn = gr.Button("📷 Apply image-dataset preset (training on stills)", scale=1)
                    gr.Markdown("*Sets image-appropriate values here **and** in the Dataset tab: 1-frame "
                                "buckets, no audio, first-frame conditioning off, video-only LoRA targets.*")
                cfg_status = gr.Textbox(label="Config status / warnings", interactive=False, lines=3)
                cfg_preview = gr.Code(label="Generated YAML", language="yaml", lines=14)

                with gr.Accordion("🧠 Model", open=True):
                    with gr.Row():
                        reg("model.model_path", gr.Textbox(label="Model checkpoint (.safetensors)",
                                                           value=D["model_path"], scale=2))
                        reg("model.text_encoder_path", gr.Textbox(label="Gemma text encoder dir",
                                                                  value=D["text_encoder_path"], scale=2))
                    with gr.Row():
                        reg("model.training_mode", gr.Radio(label="Training mode", choices=["lora", "full"],
                                                            value="lora"))
                        reg("model.load_checkpoint", gr.Textbox(
                            label="Resume checkpoint (file or dir, blank = fresh)", value="", scale=2))
                        resume_latest_btn = gr.Button("↩︎ point at this run's checkpoints", scale=1)

                with gr.Accordion("🎯 LoRA", open=True):
                    with gr.Row():
                        reg("lora.rank", gr.Slider(label="Rank", minimum=2, maximum=256, step=2, value=16,
                                                   info="low-VRAM default 16 · more capacity: 32–64"))
                        reg("lora.alpha", gr.Slider(label="Alpha", minimum=1, maximum=256, step=1, value=16,
                                                    info="usually = rank"))
                        reg("lora.dropout", gr.Slider(label="Dropout", minimum=0.0, maximum=1.0, step=0.01,
                                                      value=0.0))
                    reg("_lora.target_modules", gr.CheckboxGroup(
                        label="Target modules", choices=TARGET_MODULE_CHOICES,
                        value=["to_k", "to_q", "to_v", "to_out.0"],
                        info="short patterns (to_k…) match video + audio + cross-modal attention — "
                             "recommended for audio-video training. attn1/attn2 = video-only. "
                             "Add ff.net for extra capacity."))
                    reg("_lora.target_extra", gr.Textbox(label="Extra target modules (comma-separated)", value=""))

                with gr.Accordion("🎬 Training strategy", open=True):
                    with gr.Row():
                        strat_name = reg("strategy.name", gr.Radio(
                            label="Strategy", choices=["text_to_video", "video_to_video"], value="text_to_video",
                            info="video_to_video = IC-LoRA (needs reference latents)"))
                        reg("strategy.first_frame_conditioning_p", gr.Slider(
                            label="First-frame conditioning probability", minimum=0.0, maximum=1.0, step=0.05,
                            value=0.5, info="higher → better image-to-video behaviour"))
                    with gr.Row():
                        strat_audio = reg("strategy.with_audio", gr.Checkbox(
                            label="Train audio too (joint AV)", value=True))
                        strat_audio_dir = reg("strategy.audio_latents_dir", gr.Textbox(
                            label="Audio latents dir", value="audio_latents"))
                        strat_ref_dir = reg("strategy.reference_latents_dir", gr.Textbox(
                            label="Reference latents dir (IC-LoRA)", value="reference_latents", visible=False))
                    reg("_strategy.override", gr.Code(
                        label="Advanced: raw training_strategy YAML override (e.g. 'flexible' strategy — "
                              "leave empty normally)", language="yaml", value="", lines=3))

                with gr.Accordion("📈 Optimization", open=True):
                    with gr.Row():
                        reg("optimization.learning_rate", gr.Number(label="Learning rate", value=1e-4,
                                                                    info="LoRA typical: 5e-5 – 2e-4"))
                        reg("optimization.steps", gr.Number(label="Training steps", value=2000, precision=0))
                        reg("optimization.batch_size", gr.Number(label="Batch size / GPU", value=1, precision=0,
                                                                 info="keep 1 with multiple buckets"))
                        reg("optimization.gradient_accumulation_steps", gr.Number(
                            label="Grad accumulation", value=1, precision=0))
                    with gr.Row():
                        reg("optimization.optimizer_type", gr.Radio(label="Optimizer",
                                                                    choices=["adamw8bit", "adamw"],
                                                                    value="adamw8bit",
                                                                    info="8-bit saves ~75% optimizer VRAM"))
                        reg("optimization.scheduler_type", gr.Dropdown(
                            label="LR scheduler",
                            choices=["linear", "constant", "cosine", "cosine_with_restarts", "polynomial", "step"],
                            value="linear"))
                        reg("optimization.max_grad_norm", gr.Number(label="Max grad norm", value=1.0))
                        reg("optimization.enable_gradient_checkpointing", gr.Checkbox(
                            label="Gradient checkpointing", value=True, info="essential for low VRAM"))
                    reg("_optimization.scheduler_params", gr.Textbox(
                        label='Scheduler params (JSON, e.g. {"end_factor": 0.1})', value="{}"))

                with gr.Accordion("⚡ Acceleration / memory", open=True):
                    with gr.Row():
                        reg("acceleration.mixed_precision_mode", gr.Radio(
                            label="Mixed precision", choices=["bf16", "fp16", "no"], value="bf16"))
                        reg("_acceleration.quantization", gr.Dropdown(
                            label="Base model quantization",
                            choices=["int8-quanto", "none (off)", "int4-quanto", "int2-quanto",
                                     "fp8-quanto", "fp8uz-quanto"],
                            value="int8-quanto", info="int8 ≈ −50% model VRAM · LoRA mode only"))
                        reg("acceleration.load_text_encoder_in_8bit", gr.Checkbox(
                            label="Text encoder in 8-bit", value=True))
                        reg("acceleration.offload_optimizer_during_validation", gr.Checkbox(
                            label="Offload optimizer during validation", value=False,
                            info="enable if validation OOMs"))

                with gr.Accordion("💾 Data", open=True):
                    with gr.Row():
                        reg("data.preprocessed_data_root", gr.Textbox(
                            label="Preprocessed data root (contains latents/, conditions/…)",
                            value=D["data_root"], scale=3))
                        reg("data.num_dataloader_workers", gr.Number(label="Dataloader workers", value=2,
                                                                     precision=0, scale=1))

                with gr.Accordion("🎥 Validation (sample videos during training)", open=False):
                    with gr.Row():
                        reg("_validation.interval", gr.Number(label="Every N steps (0 = off)", value=100,
                                                              precision=0))
                        reg("validation.inference_steps", gr.Number(label="Inference steps", value=30, precision=0))
                        reg("validation.seed", gr.Number(label="Seed", value=42, precision=0))
                        reg("validation.frame_rate", gr.Number(label="FPS", value=25.0))
                        reg("validation.skip_initial_validation", gr.Checkbox(label="Skip step-0 validation",
                                                                              value=False))
                    reg("_validation.prompts", gr.Textbox(
                        label="Validation prompts (one per line — include your trigger word!)",
                        value=DEFAULT_PROMPTS, lines=5))
                    reg("validation.negative_prompt", gr.Textbox(
                        label="Negative prompt",
                        value="worst quality, inconsistent motion, blurry, jittery, distorted"))
                    with gr.Row():
                        reg("_validation.dims_w", gr.Number(label="Width (÷32)", value=576, precision=0))
                        reg("_validation.dims_h", gr.Number(label="Height (÷32)", value=576, precision=0))
                        reg("_validation.dims_f", gr.Number(label="Frames (%8==1)", value=49, precision=0))
                        reg("validation.guidance_scale", gr.Number(label="CFG scale", value=4.0))
                        reg("validation.stg_scale", gr.Number(label="STG scale (0=off)", value=1.0))
                    with gr.Row():
                        reg("_validation.stg_blocks", gr.Textbox(label="STG blocks (csv, blank = all)", value="29"))
                        reg("validation.stg_mode", gr.Radio(label="STG mode", choices=["stg_av", "stg_v"],
                                                            value="stg_av"))
                        reg("validation.generate_audio", gr.Checkbox(label="Generate audio", value=True))
                        reg("validation.generate_video", gr.Checkbox(label="Generate video", value=True))
                    reg("_validation.images", gr.Textbox(
                        label="Image-to-video: first-frame image paths (one per line, must match prompt count)",
                        value="", lines=2))
                    reg("_validation.reference_videos", gr.Textbox(
                        label="IC-LoRA: reference video paths (one per line, must match prompt count)",
                        value="", lines=2))
                    reg("_validation.samples_override", gr.Code(
                        label="Advanced: validation `samples` YAML (overrides prompts; supports prefix/suffix/"
                              "mask/crop conditions — leave empty normally)", language="yaml", value="", lines=3))

                with gr.Accordion("🗃️ Checkpoints", open=False):
                    with gr.Row():
                        reg("_checkpoints.interval", gr.Number(label="Save every N steps (0 = only at end)",
                                                               value=250, precision=0))
                        reg("checkpoints.keep_last_n", gr.Number(label="Keep last N (−1 = all)", value=-1,
                                                                 precision=0))
                        reg("checkpoints.precision", gr.Radio(label="Save precision",
                                                              choices=["bfloat16", "float32"], value="bfloat16"))
                    with gr.Row():
                        reg("checkpoints.save_training_state", gr.Radio(
                            label="Save training state (for resume)", choices=["minimal", "full", "off"],
                            value="minimal", info="minimal = a few KB, enough for LoRA resume"))
                        reg("checkpoints.no_resume", gr.Checkbox(
                            label="no_resume (ignore saved state, start at step 0)", value=False))

                with gr.Accordion("🌊 Flow matching · 🤗 Hub · 📊 W&B · general", open=False):
                    with gr.Row():
                        reg("flow_matching.timestep_sampling_mode", gr.Radio(
                            label="Timestep sampling", choices=["shifted_logit_normal", "uniform"],
                            value="shifted_logit_normal"))
                        reg("_flow_matching.params", gr.Textbox(label="Sampling params (JSON)", value="{}"))
                    with gr.Row():
                        reg("hub.push_to_hub", gr.Checkbox(label="Push to HF Hub when done", value=False))
                        reg("hub.hub_model_id", gr.Textbox(label="Hub repo id (user/name)", value=""))
                    with gr.Row():
                        reg("wandb.enabled", gr.Checkbox(label="Enable W&B", value=False))
                        reg("wandb.project", gr.Textbox(label="W&B project", value="ltx-2-trainer"))
                        reg("wandb.entity", gr.Textbox(label="W&B entity", value=""))
                        reg("_wandb.tags", gr.Textbox(label="Tags (csv)", value="ltx2, lora"))
                        reg("wandb.log_validation_videos", gr.Checkbox(label="Log validation videos", value=True))
                    with gr.Row():
                        reg("seed", gr.Number(label="Global seed", value=42, precision=0))
                        reg("output_dir", gr.Textbox(label="Output directory", value=D["output_dir"], scale=3))

            # =================================================================
            # TRAIN TAB
            # =================================================================
            with gr.Tab("🚀 Train & Monitor"):
                with gr.Row():
                    launch_mode = gr.Dropdown(label="Launch mode", choices=list(ACCEL_CONFIGS),
                                              value="Single GPU", scale=2)
                    num_processes = gr.Number(label="# GPUs (multi-GPU modes)", value=2, precision=0, scale=1)
                    cuda_devices = gr.Textbox(label="CUDA_VISIBLE_DEVICES (blank = all)", value="", scale=1)
                with gr.Row():
                    start_btn = gr.Button("▶️ Start training", variant="primary", scale=2)
                    stop_btn = gr.Button("🛑 Stop training", variant="stop", scale=1)
                    validate_btn = gr.Button("🧪 Validate config with trainer", scale=1)
                train_action_status = gr.Textbox(label="Status", interactive=False, lines=3)
                train_cmd_box = gr.Textbox(label="Command", interactive=False)

                progress_html = gr.HTML(_progress_html(0))
                status_md = gr.Markdown("No training started yet.")
                loss_plot = gr.LinePlot(value=loss_dataframe(), x="step", y="loss",
                                        title="Training loss (smoothed)", height=260)
                train_log = gr.Textbox(label="Training log", lines=18, max_lines=18,
                                       interactive=False, autoscroll=True)

                gr.Markdown("### 🎥 Validation samples")
                with gr.Row():
                    vid_refresh_btn = gr.Button("🔄 Refresh videos", scale=1)
                    vid_dropdown = gr.Dropdown(label="Validation videos (newest first)", choices=[], scale=3)
                vid_player = gr.Video(label="Preview", height=360)

                gr.Markdown("### 🗃️ Checkpoints (LoRA .safetensors — ComfyUI-compatible)")
                ckpt_refresh_btn = gr.Button("🔄 Refresh checkpoints")
                ckpt_table = gr.Dataframe(value=checkpoints_table(D["output_dir"]), interactive=False)

            # =================================================================
            # VAST.AI TAB
            # =================================================================
            with gr.Tab("☁️ vast.ai Setup"):
                gr.Markdown(VAST_GUIDE)
                gr.Markdown("### On-start script (`vast_onstart.sh`)")
                gr.Code(value=ONSTART_SCRIPT, language="shell", lines=25)

        # =====================================================================
        # Event wiring
        # =====================================================================
        cfg_inputs = [C[k] for k in CONFIG_KEYS]

        # setup tab
        clone_btn.click(do_clone_repo, [repo_dir, hf_token], setup_status)
        envcheck_btn.click(env_check,
                           [repo_dir, C["model.model_path"], C["model.text_encoder_path"],
                            C["data.preprocessed_data_root"]],
                           envcheck_md)

        def _apply_preset(name):
            repo, fname = MODEL_PRESETS.get(name, ("", ""))
            if not repo:
                return gr.update(), gr.update()
            return gr.update(value=repo), gr.update(value=fname)

        model_preset.change(_apply_preset, model_preset, [model_repo, model_file])
        dl_model_btn.click(do_download_model, [model_preset, model_repo, model_file, model_dest, hf_token],
                           dl_model_status)
        dl_te_btn.click(do_download_te, [te_dest, hf_token], dl_te_status)

        # dataset tab
        split_btn.click(do_split_scenes, [repo_dir, split_input, split_out, split_minlen, use_uv, hf_token],
                        dataset_status)
        cap_btn.click(do_caption,
                      [repo_dir, cap_dir, cap_out, cap_type, cap_recursive, cap_override,
                       cap_instruction, cap_workers, cap_apikey, use_uv, hf_token],
                      dataset_status)
        capserver_start_btn.click(do_capserver_start, [repo_dir, use_uv, hf_token], dataset_status)
        capserver_stop_btn.click(lambda: stop_job("capserver"), None, dataset_status)
        pp_btn.click(do_preprocess,
                     [repo_dir, pp_dataset, pp_buckets, pp_model, pp_te, pp_outdir, pp_trigger, pp_batch,
                      pp_tiling, pp_decode, pp_skipaudio, pp_te8bit, pp_removepfx, pp_overwrite, pp_gpus,
                      use_uv, hf_token],
                     dataset_status)
        pp_use_btn.click(dataset_data_root, [pp_dataset, pp_outdir], C["data.preprocessed_data_root"])
        dataset_stop_btn.click(lambda: stop_job("dataset"), None, dataset_status)

        # settings tab
        save_cfg_btn.click(do_save_config, [config_path, *cfg_inputs], [cfg_status, cfg_preview])
        preview_cfg_btn.click(do_preview_config, cfg_inputs, [cfg_status, cfg_preview])
        load_cfg_btn.click(do_load_config, load_cfg_file, [cfg_status, *cfg_inputs])
        resume_latest_btn.click(lambda out: str(Path(out) / "checkpoints"),
                                C["output_dir"], C["model.load_checkpoint"])

        def _strategy_vis(name):
            t2v = name == "text_to_video"
            return (gr.update(visible=t2v), gr.update(visible=t2v), gr.update(visible=not t2v))

        strat_name.change(_strategy_vis, strat_name, [strat_audio, strat_audio_dir, strat_ref_dir])

        def _image_preset():
            video_targets = ["attn1.to_k", "attn1.to_q", "attn1.to_v", "attn1.to_out.0",
                             "attn2.to_k", "attn2.to_q", "attn2.to_v", "attn2.to_out.0"]
            msg = (
                "📷 Image-dataset preset applied:\n"
                "• Dataset tab: resolution buckets → 1024x1024x1 (F=1 = stills), skip audio extraction\n"
                "• with_audio off · first_frame_conditioning_p → 0.0 (with 1-frame samples it would "
                "condition the whole image and zero the loss) · video-only LoRA target modules · "
                "no validation audio · batch size 2 (single bucket, so batching is allowed)\n"
                "• Validation still renders 576x576x49 videos to check motion — set frames to 1 for stills"
            )
            return (gr.update(value="1024x1024x1"), gr.update(value=True), gr.update(value=False),
                    gr.update(value=0.0), gr.update(value=video_targets), gr.update(value=False),
                    gr.update(value=2), msg)

        img_preset_btn.click(_image_preset, None,
                             [pp_buckets, pp_skipaudio, C["strategy.with_audio"],
                              C["strategy.first_frame_conditioning_p"], C["_lora.target_modules"],
                              C["validation.generate_audio"], C["optimization.batch_size"], cfg_status])

        # train tab
        start_btn.click(do_start_training,
                        [repo_dir, config_path, launch_mode, num_processes, cuda_devices,
                         use_uv, hf_token, wandb_key, *cfg_inputs],
                        [train_action_status, train_cmd_box])
        stop_btn.click(do_stop_training, None, train_action_status)
        validate_btn.click(do_validate_config, [repo_dir, config_path, use_uv, hf_token, *cfg_inputs],
                           train_action_status)
        vid_refresh_btn.click(list_validation_videos, C["output_dir"], vid_dropdown)
        vid_dropdown.change(lambda p: p, vid_dropdown, vid_player)
        ckpt_refresh_btn.click(checkpoints_table, C["output_dir"], ckpt_table)

        # periodic refresh
        timer = gr.Timer(2.0)

        def on_tick():
            p_html, s_md = training_status()
            return (
                get_job("setup").tail(150),
                get_job("download_model").tail(80),
                get_job("download_te").tail(80),
                get_job("dataset").tail(300),
                get_job("capserver").tail(60),
                get_job("train").tail(400),
                p_html,
                s_md,
                loss_dataframe(),
                sys_info_md(),
            )

        timer.tick(on_tick, None,
                   [setup_log, dl_model_log, dl_te_log, dataset_log, capserver_log, train_log,
                    progress_html, status_md, loss_plot, sysinfo_md])

    return demo


def main():
    ap = argparse.ArgumentParser(description="LTX-2.3 LoRA Trainer GUI")
    ap.add_argument("--host", default="0.0.0.0")
    ap.add_argument("--port", type=int, default=7860)
    ap.add_argument("--share", action="store_true", help="create a public gradio.live link")
    args = ap.parse_args()
    demo = create_demo()
    demo.queue().launch(server_name=args.host, server_port=args.port, share=args.share,
                        allowed_paths=[WORKSPACE, str(Path.home())], **_LAUNCH_KW)


if __name__ == "__main__":
    main()