ADAM October 2026 source release: PixelRow, INRFlow, Wan Video, Oasis player and field guide
f8c73f9 verified | """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"}}]} | |