File size: 8,379 Bytes
e0265b9 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 | """Execution bridge for the connected Rectified Flow image trainer."""
from __future__ import annotations
import json
import queue
import subprocess
import sys
import threading
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
IMAGE_EXTENSIONS = {".jpg", ".jpeg", ".png", ".webp", ".bmp"}
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
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
context.log("Safe stop requested; waiting for Flow Matching to finish its current batch.")
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),
}]}
|