adam/tools/flow_adapter.py

"""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),
    }]}