"""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"(?