"""Execution bridge for the connected Rectified Flow image trainer.""" from __future__ import annotations import json import queue import subprocess import sys import threading import time from pathlib import Path from adam.config import ConfigManager from adam.executor import ToolCancelled, ToolContext, ToolExecutionError from adam.process_control import set_process_tree_paused, terminate_process_tree IMAGE_EXTENSIONS = {".jpg", ".jpeg", ".png", ".webp", ".bmp"} FORCE_STOP_TIMEOUT_SECONDS = 30 def _latest_preview(folder: Path) -> Path | None: try: images = [path for path in folder.rglob("*") if path.is_file() and path.suffix.lower() in IMAGE_EXTENSIONS and any(token in path.name.lower() for token in ("preview", "sample", "epoch"))] return max(images, key=lambda path: path.stat().st_mtime) if images else None except OSError: return None def train_flow( context: ToolContext, dataset_dir: str, model_name: str, epochs: int, output_dir: str, resume_from: str = "", resolution: int = 256, batch_size: int = 8, learning_rate: float = 0.0002, gradient_accumulation: int = 1, workers: int = 4, mixed_precision: str = "fp16", save_every: int = 10, preview_every: int = 10, preview_steps: int = 10, gradient_checkpointing: bool = False, preview_enabled: bool = True, preview_prompt: str = "", preview_seed: int = 123456789, ) -> dict[str, object]: """Launch the user's Flow Matching worker and relay its structured progress.""" root = Path(str(ConfigManager(context.root).get("tool_folders", {}).get("flow_trainer", ""))).expanduser() script = root / "flow_matching_app.py" dataset = Path(dataset_dir).expanduser().resolve() output = Path(output_dir).expanduser().resolve() if not script.is_file(): raise ToolExecutionError("Flow Matching flow_matching_app.py was not found. Re-scan its folder in Settings.") if not dataset.is_dir(): raise ToolExecutionError("The selected Flow Matching dataset folder no longer exists.") if sum(1 for path in dataset.iterdir() if path.is_file() and path.suffix.lower() in IMAGE_EXTENSIONS) < 2: raise ToolExecutionError("The Flow Matching dataset needs at least two image files before training can start.") if not 1 <= int(epochs) <= 100_000: raise ToolExecutionError("Epoch count must be between 1 and 100000.") if not 64 <= int(resolution) <= 512 or int(resolution) % 16 or not 1 <= int(batch_size) <= 64 or not 1e-7 <= float(learning_rate) <= 0.1 or not 1 <= int(gradient_accumulation) <= 64 or not 0 <= int(workers) <= 16 or mixed_precision not in {"fp16", "no"} or min(int(save_every), int(preview_every), int(preview_steps)) < 1: raise ToolExecutionError("Flow training options are outside ADAM's safe range.") safe_name = model_name.strip() if not safe_name or len(safe_name) > 96 or any(char in safe_name for char in "<>:\\|?*\x00"): raise ToolExecutionError("Choose a short model name without filesystem-reserved characters.") output_root = (root / "output_flow_models").resolve() try: output.relative_to(output_root) except ValueError as exc: raise ToolExecutionError("Flow Matching outputs must stay inside output_flow_models.") from exc if output.exists(): raise ToolExecutionError("The chosen Flow Matching output folder already exists; ADAM will not overwrite it.") resume = Path(resume_from).expanduser().resolve() if resume_from else None if resume: try: metadata = json.loads((resume / "flow_model_info.json").read_text(encoding="utf-8")) if metadata.get("model_type") != "rectified_flow" or not (resume / "unet" / "config.json").is_file(): raise ValueError saved_resolution = int(metadata.get("resolution", 0) or 0) except (OSError, ValueError, TypeError, json.JSONDecodeError) as exc: raise ToolExecutionError("Choose a valid completed Flow Matching model to continue.") from exc if saved_resolution != int(resolution): raise ToolExecutionError( f"The selected Flow model is {saved_resolution}px; continuation must use the same resolution." ) output.parent.mkdir(parents=True, exist_ok=True) command = [ sys.executable, str(script), "--train-worker", "--data-dir", str(dataset), "--output-dir", str(output), "--model-name", safe_name, "--epochs", str(int(epochs)), "--resolution", str(int(resolution)), "--batch-size", str(int(batch_size)), "--learning-rate", str(float(learning_rate)), "--workers", str(int(workers)), "--gradient-accumulation", str(int(gradient_accumulation)), "--mixed-precision", mixed_precision, "--save-every", str(int(save_every)), "--preview-every", str(int(preview_every) if preview_enabled else int(epochs) + 1), "--preview-steps", str(int(preview_steps)), "--tf32", ] if gradient_checkpointing: command.append("--gradient-checkpointing") if resume: command.extend(["--continue-model", str(resume)]) context.log(f"Starting real Flow Matching training. Output folder: {output}") process = subprocess.Popen(command, cwd=str(root), stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True, encoding="utf-8", errors="replace", shell=False) lines: queue.Queue[str | None] = queue.Queue() def read_output() -> None: assert process.stdout is not None for line in process.stdout: lines.put(line.rstrip()) lines.put(None) threading.Thread(target=read_output, daemon=True).start() context.progress(1, "Starting Flow Matching trainer") stopped = False stop_requested_at: float | None = None force_stop_sent = False suspended = False stop_file = output / "stop_flow_training.flag" while True: should_pause = not context.run_event.is_set() if should_pause != suspended: if set_process_tree_paused(process, should_pause): suspended = should_pause context.log("Flow Matching trainer paused safely." if suspended else "Flow Matching trainer resumed.") if context.cancel_event.is_set() and not stopped: if suspended: set_process_tree_paused(process, False) suspended = False stop_file.touch(exist_ok=True) stopped = True stop_requested_at = time.monotonic() context.log("Safe stop requested; waiting for Flow Matching to finish its current batch.") if stop_requested_at and not force_stop_sent and time.monotonic() - stop_requested_at > FORCE_STOP_TIMEOUT_SECONDS: terminate_process_tree(process, timeout=3) force_stop_sent = True context.log("Flow Matching did not stop in time; terminating the trainer process.") try: line = lines.get(timeout=0.15) if line and line.startswith("FLOW_EVENT:"): event = json.loads(line.split(":", 1)[1]) if event.get("type") == "progress" and not stopped: current = int(event.get("epoch", 0) or 0) context.progress(max(1, min(99, round(current * 100 / int(epochs)))), f"Finished epoch {current} of {epochs}") if preview_enabled and current and current % int(preview_every) == 0: candidate = Path(str(event.get("preview_path", ""))) if event.get("preview_path") else _latest_preview(output) if candidate: context.preview(candidate, epoch=current, next_epoch=min(int(epochs), current + int(preview_every)), prompt=preview_prompt, seed=int(preview_seed), steps=int(preview_steps)) elif event.get("type") == "warning": context.log(str(event.get("message", "Flow trainer warning."))) elif line: context.log(line) except queue.Empty: pass if process.poll() is not None and lines.empty(): break if stopped: raise ToolCancelled("Flow Matching training stopped by user.") if process.returncode != 0: raise ToolExecutionError(f"Flow Matching trainer exited with code {process.returncode}. See the job log for details.") context.progress(100, "Flow Matching training completed") return {"output_folder": str(output), "model_name": safe_name, "assets": [{ "kind": "model", "name": safe_name, "path": str(output), "trainer": "flow", "dataset_path": str(dataset), "checkpoint": str(output), "epochs": int(epochs), }]}