Blame for adam/tools/wan_video_adapter.py

"""Run the connected Musubi environment without importing it into ADAM."""
from __future__ import annotations

import json
import os
import queue
import random
import re
import subprocess
import threading
from datetime import datetime, timezone
from pathlib import Path
from typing import Any
from uuid import uuid4

from adam.executor import ToolCancelled, ToolContext, ToolExecutionError
from adam.process_control import set_process_tree_paused, terminate_process_tree
from adam.video_lora import (
    MODEL_FILES, checkpoint_info, clips_in, dataset_errors, frames_for_seconds,
    generation_plan, inspect_clip, settings_for, setup_errors, trainer_root, training_plan,
)


def _runtime(context: ToolContext) -> tuple[Path, Path, Path]:
    root = trainer_root(context.root)
    errors = setup_errors(root)
    if errors:
        raise ToolExecutionError("\n".join(errors))
    assert root is not None
    return root, root / ".venv" / "Scripts" / "python.exe", root / "trainer-engine" / "src" / "musubi_tuner"


def _progress_counter(line: str, required_prefix: str = "") -> tuple[int, int] | None:
    """Read a tqdm counter, optionally accepting only the named primary loop."""
    if required_prefix:
        pattern = rf"\b{re.escape(required_prefix)}\s*:\s*\d+%.*?\b(\d+)\s*/\s*(\d+)\b"
    else:
        pattern = r"\d+%.*?\b(\d+)\s*/\s*(\d+)\b"
    match = re.search(pattern, line, re.IGNORECASE)
    if not match:
        return None
    return int(match[1]), int(match[2])


def run_stage(
    context: ToolContext, command: list[str], root: Path, label: str,
    start: int, end: int, *, speed: bool = False, progress_prefix: str = "",
) -> None:
    context.checkpoint()
    context.progress(start, label, reset_eta=True)
    env = os.environ.copy()
    env.update(PYTHONUTF8="1", PYTHONUNBUFFERED="1")
    if speed:
        env["TORCH_ALLOW_TF32_CUBLAS_OVERRIDE"] = "1"
    process = subprocess.Popen(
        command, cwd=str(root), env=env, stdout=subprocess.PIPE, stderr=subprocess.STDOUT,
        text=True, encoding="utf-8", errors="replace", shell=False,
        creationflags=subprocess.CREATE_NO_WINDOW if os.name == "nt" else 0,
    )
    lines: queue.Queue[str | None] = queue.Queue()

    def read_output() -> None:
        try:
            assert process.stdout is not None
            for line in process.stdout:
                lines.put(line.rstrip())
        finally:
            lines.put(None)

    reader = threading.Thread(target=read_output, daemon=True)
    reader.start()
    suspended = False
    finished_reading = False
    percent = start
    try:
        while not finished_reading or process.poll() is None:
            if context.cancel_event.is_set():
                raise ToolCancelled("Video job cancelled. Previously saved checkpoints remain in the run folder.")
            paused = not context.run_event.is_set()
            if paused != suspended and process.poll() is None:
                if not set_process_tree_paused(process, paused):
                    raise ToolExecutionError("Could not pause/resume the video worker; it has been stopped.")
                suspended = paused
            try:
                line = lines.get(timeout=0.1)
            except queue.Empty:
                continue
            if line is None:
                finished_reading = True
                continue
            if not line:
                continue
            context.log(line)
            progress = _progress_counter(line, progress_prefix)
            if progress:
                current, total = progress
                if total:
                    percent = max(percent, min(end - 1, start + round((end - start) * current / total)))
                    # Do not call checkpoint() while controlling an external process:
                    # a pause must suspend its process tree before waiting.
                    context.progress_callback(percent, label, current_step=current, total_steps=total, unit="step")
        if process.returncode != 0:
            raise ToolExecutionError(f"{label} failed (exit {process.returncode}). See the job log.")
    except BaseException:
        if suspended:
            set_process_tree_paused(process, False)
        terminate_process_tree(process)
        raise
    finally:
        reader.join(timeout=1)
        if process.stdout is not None:
            process.stdout.close()
    context.progress(end, f"{label} completed")


def _fingerprint(folder: Path) -> list[tuple[str, int, int]]:
    return [(str(path), path.stat().st_size, path.stat().st_mtime_ns)
            for clip in clips_in(folder) for path in (clip, clip.with_suffix(".txt"))]


def _write_json(path: Path, values: dict[str, Any]) -> None:
    temporary = path.with_suffix(".tmp")
    temporary.write_text(json.dumps(values, indent=2), encoding="utf-8")
    temporary.replace(path)


def train(context: ToolContext, dataset_dir: str, model_name: str, epochs: int, output_dir: str, resume_from: str = "", **values: Any) -> dict[str, Any]:
    try:
        plan = training_plan(context.root, dataset_dir, model_name, epochs, values, resume_from, output_dir)
        settings = settings_for("training", values)
    except ValueError as exc:
        raise ToolExecutionError(str(exc)) from exc
    root, python, engine = _runtime(context)
    dataset = Path(plan.steps[0].arguments["dataset_dir"])
    output = Path(plan.steps[0].arguments["output_dir"])
    minimum = min(int(value) for value in settings["target_frames"].split(","))
    errors = dataset_errors(dataset, settings["trigger_word"], minimum, inspect_media=True)
    if errors:
        raise ToolExecutionError("\n".join(errors[:12]))
    context.checkpoint()
    fingerprint = _fingerprint(dataset)
    # mkdir with exist_ok=False reserves the run, including against queued duplicates.
    output.mkdir(parents=True, exist_ok=False)
    cache = output / "cache"
    cache.mkdir()
    config = output / "dataset_config.toml"
    width, height = (int(value) for value in settings["resolution"].split("x"))
    config.write_text(
        f'[general]\nresolution = [{width}, {height}]\ncaption_extension = ".txt"\n'
        f'batch_size = {settings["batch_size"]}\nenable_bucket = true\nbucket_no_upscale = true\n\n'
        f'[[datasets]]\nvideo_directory = {json.dumps(dataset.as_posix(), ensure_ascii=False)}\n'
        f'cache_directory = {json.dumps(cache.as_posix(), ensure_ascii=False)}\n'
        f'target_frames = [{settings["target_frames"]}]\nframe_extraction = "head"\nnum_repeats = 1\n',
        encoding="utf-8",
    )
    info = {"model_type": "wan21_t2v_1_3b_lora", "model_name": model_name, "dataset_dir": str(dataset),
            "epochs": epochs, "trigger_word": settings["trigger_word"], "settings": settings,
            "dataset_files": fingerprint,
            "resume_from": resume_from, "continuation_mode": "weights_only" if resume_from else "fresh",
            "trainer_root": str(root), "job_id": context.job_id, "status": "running"}
    _write_json(output / "model_info.json", info)
    dit, vae, text_encoder = (str(root / "models" / name) for name in MODEL_FILES)
    command = [str(python), "-m", "accelerate.commands.launch", "--num_cpu_threads_per_process", "1", "--mixed_precision", "bf16",
               str(engine / "wan_train_network.py"), "--task", "t2v-1.3B", "--dit", dit,
               "--dataset_config", str(config), "--sdpa", "--mixed_precision", "bf16", "--save_precision", "bf16",
               "--fp8_base", "--optimizer_type", "adamw8bit", "--learning_rate", str(settings["learning_rate"]),
               "--gradient_checkpointing", "--blocks_to_swap", str(settings["blocks_to_swap"]),
               "--block_swap_h2d_only", "--block_swap_ring_size", "1", "--max_data_loader_n_workers", "1",
               "--network_module", "networks.lora_wan", "--network_dim", str(settings["rank"]), "--network_alpha", str(settings["alpha"]),
               "--timestep_sampling", "shift", "--discrete_flow_shift", "3.0", "--max_train_epochs", str(epochs),
               "--save_every_n_epochs", str(settings["save_every"]), "--seed", str(settings["seed"]),
               "--output_dir", str(output), "--output_name", model_name, "--log_with", "tensorboard", "--logging_dir", str(output / "logs")]
    if resume_from:
        command += ["--network_weights", resume_from, "--dim_from_weights", resume_from]
        context.log("Continuing LoRA weights with a fresh optimizer and schedule. The checkpoint supplies rank and alpha.")
    stages = [
        ([str(python), str(root / "tools" / "validate_dataset.py"), "--directory", str(dataset), "--trigger", settings["trigger_word"]], "Validate clips", 0, 3),
        ([str(python), str(engine / "wan_cache_latents.py"), "--dataset_config", str(config), "--vae", vae, "--vae_cache_cpu"], "Cache video latents", 3, 18),
        ([str(python), str(engine / "wan_cache_text_encoder_outputs.py"), "--dataset_config", str(config), "--t5", text_encoder, "--batch_size", "1", "--fp8_t5"], "Cache captions", 18, 23),
        (command, "Train Wan Video LoRA", 23, 99),
    ]
    try:
        for args, label, start, end in stages:
            if fingerprint != _fingerprint(dataset):
                raise ToolExecutionError("Dataset changed during this run. Start a new run so captions and cached tensors match.")
            run_stage(
                context, args, root, label, start, end,
                progress_prefix="steps" if label == "Train Wan Video LoRA" else "",
            )
        checkpoint = output / f"{model_name}.safetensors"
        checkpoint_info(checkpoint)
        info.update(status="finished", checkpoint=str(checkpoint))
    except BaseException as exc:
        info.update(status="cancelled" if isinstance(exc, ToolCancelled) else "failed", error=str(exc))
        _write_json(output / "model_info.json", info)
        raise
    _write_json(output / "model_info.json", info)
    context.progress(100, "Wan Video LoRA training complete")
    return {"output_folder": str(output), "assets": [
        {"kind": "dataset", "name": dataset.name, "path": str(dataset), "metadata": {"media_type": "video"}},
        {"kind": "model", "name": model_name, "path": str(checkpoint), "trainer": "wan_video", "checkpoint": str(checkpoint),
         "dataset_path": str(dataset), "epochs": epochs, "metadata": info},
    ]}


def generate(context: ToolContext, model_path: str, model_name: str = "", **values: Any) -> dict[str, Any]:
    try:
        generation_plan(context.root, model_path, model_name, values)
        settings = settings_for("generation", values)
    except ValueError as exc:
        raise ToolExecutionError(str(exc)) from exc
    root, python, engine = _runtime(context)
    if settings["randomize_seed"]:
        settings["seed"] = random.SystemRandom().randrange(2147483648)
    frames = frames_for_seconds(settings["duration"], settings["fps"])
    width, height = (832, 480) if settings["format"].startswith("Landscape") else (480, 832)
    output = context.root / "data" / "video_generations" / f"{context.job_id}_{uuid4().hex[:8]}"
    output.mkdir(parents=True, exist_ok=False)
    metadata = {"model_name": model_name or Path(model_path).stem, "model_path": model_path, "trainer": "wan_video",
                "job_id": context.job_id, "created_at": datetime.now(timezone.utc).isoformat(), "settings": settings,
                "frames": frames, "width": width, "height": height, "actual_duration": frames / settings["fps"], "status": "running"}
    _write_json(output / "generation.json", metadata)
    dit, vae, text_encoder = (str(root / "models" / name) for name in MODEL_FILES)
    command = [str(python), str(engine / "wan_generate_video.py"), "--fp8", "--task", "t2v-1.3B",
               "--video_size", str(height), str(width), "--video_length", str(frames), "--fps", str(settings["fps"]),
               "--infer_steps", str(settings["steps"]), "--prompt", settings["prompt"], "--seed", str(settings["seed"]),
               "--save_path", str(output), "--output_type", "video", "--dit", dit, "--vae", vae, "--t5", text_encoder,
               "--attn_mode", "torch", "--blocks_to_swap", str(settings["blocks_to_swap"]),
               "--lora_weight", model_path, "--lora_multiplier", str(settings["lora_strength"])]
    try:
        context.log(f"Generation seed: {settings['seed']}. Expected video: {width}×{height}, {frames} frames.")
        run_stage(
            context, command, root, "Generate Wan video", 0, 99,
            speed=settings["experimental_speed"], progress_prefix="steps",
        )
        videos = sorted(output.glob("*.mp4"))
        if len(videos) != 1:
            raise ToolExecutionError("Generation completed without exactly one MP4 in its output folder.")
        media = inspect_clip(videos[0])
        if (media["width"], media["height"], media["frames"]) != (width, height, frames) or abs(media["fps"] - settings["fps"]) > 0.1:
            raise ToolExecutionError("The saved MP4 does not match the requested dimensions, frame count or FPS.")
        metadata.update(status="finished", video=str(videos[0]))
    except BaseException as exc:
        metadata.update(status="cancelled" if isinstance(exc, ToolCancelled) else "failed", error=str(exc))
        _write_json(output / "generation.json", metadata)
        raise
    _write_json(output / "generation.json", metadata)
    context.progress(100, "Wan video saved")
    return {"output_folder": str(output), "video_path": str(videos[0]), "assets": [
        {"kind": "video", "name": metadata["model_name"], "path": str(videos[0]), "trainer": "wan_video", "metadata": metadata},
    ]}


def prepare(context: ToolContext, operation: str, dataset_dir: str, trigger_word: str, source: str = "", count: int = 32, caption: str = "") -> dict[str, Any]:
    root, python, _engine = _runtime(context)
    settings_for("training", {"trigger_word": trigger_word})
    dataset = Path(dataset_dir).expanduser().resolve()
    if not dataset_dir.strip():
        raise ToolExecutionError("Select a video dataset folder.")
    if operation == "validate":
        errors = dataset_errors(dataset, trigger_word, inspect_media=True)
        if errors:
            raise ToolExecutionError("\n".join(errors[:12]))
        command = [str(python), str(root / "tools" / "validate_dataset.py"), "--directory", str(dataset), "--trigger", trigger_word]
    elif operation == "split":
        if isinstance(count, bool) or not isinstance(count, int) or not 1 <= count <= 100:
            raise ToolExecutionError("Choose 1–100 clips to extract.")
        if not caption.strip() or not re.search(rf"(?<![\w]){re.escape(trigger_word)}(?![\w])", caption):
            raise ToolExecutionError("Write a visual caption containing the trigger word before splitting.")
        from adam.video_lora import VIDEO_EXTENSIONS
        if not Path(source).is_file() or Path(source).suffix.lower() not in VIDEO_EXTENSIONS:
            raise ToolExecutionError("Choose a readable source video.")
        # Each split gets an exclusive folder so cancelled or repeated jobs never overwrite clips.
        dataset = dataset / f"clips_{context.job_id}_{uuid4().hex[:6]}"
        dataset.mkdir(parents=True, exist_ok=False)
        caption_path = dataset / "split_caption.caption"
        caption_path.write_text(caption.strip(), encoding="utf-8")
        command = [str(python), str(root / "tools" / "split_video.py"), source, "--output-dir", str(dataset),
                   "--count", str(count), "--trigger", trigger_word, "--caption-file", str(caption_path)]
    else:
        raise ToolExecutionError("Choose validate or split.")
    run_stage(context, command, root, "Validate clips" if operation == "validate" else "Split source video", 0, 100)
    return {"output_folder": str(dataset), "assets": [{"kind": "dataset", "name": dataset.name, "path": str(dataset), "metadata": {"media_type": "video"}}]}