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