"""Safe execution bridge for the user's existing DDPM command-line trainer.""" from __future__ import annotations import json import importlib.util import math import queue import re 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 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 _safe_model_name(value: str) -> str: name = value.strip() if not name or len(name) > 96 or any(char in name for char in "<>:\\|?*\x00"): raise ToolExecutionError("Choose a short model name without filesystem-reserved characters.") return name def _parse_progress( context: ToolContext, payload: dict[str, object], epochs: int, output: Path, preview_enabled: bool, preview_every: int, preview_prompt: str, preview_seed: int, preview_steps: int, ) -> None: event = str(payload.get("event", "progress")) if event == "error": raise ToolExecutionError(str(payload.get("message", "DDPM trainer reported an error."))) if event == "step": current = int(payload.get("global_step", 0) or 0) total = int(payload.get("total_steps", 0) or 0) if total: context.progress(max(1, min(99, round(current * 100 / total))), f"Training step {current:,} of {total:,}") return if event == "epoch_end": epoch = int(payload.get("epoch", 0) or 0) context.progress(max(1, min(99, round(epoch * 100 / max(epochs, 1)))), f"Finished epoch {epoch} of {epochs}") if preview_enabled and epoch and epoch % preview_every == 0: candidate = Path(str(payload.get("preview_path", ""))) if payload.get("preview_path") else _latest_preview(output) if candidate: context.preview(candidate, epoch=epoch, next_epoch=min(epochs, epoch + preview_every), prompt=preview_prompt, seed=preview_seed, steps=preview_steps) return if event == "done": context.progress(100, "DDPM training completed") def train_ddpm( context: ToolContext, dataset_dir: str, model_name: str, epochs: int, output_dir: str, resume_from: str = "", resolution: int = 128, batch_size: int = 1, learning_rate: float = 0.0001, gradient_accumulation_steps: int = 1, dataloader_num_workers: int = 4, mixed_precision: str = "fp16", save_every: int = 10, preview_steps: int = 50, training_intensity: int = 100, preview_enabled: bool = True, preview_every: int = 5, preview_prompt: str = "", preview_seed: int = 123456789, ) -> dict[str, object]: """Run the registered DDPM project without shell interpolation or overwrites.""" trainer_root = Path(str(ConfigManager(context.root).get("tool_folders", {}).get("ddpm_trainer", ""))).expanduser() script = trainer_root / "train.py" dataset = Path(dataset_dir).expanduser().resolve() output = Path(output_dir).expanduser().resolve() model_name = _safe_model_name(model_name) if not script.is_file(): raise ToolExecutionError("DDPM train.py was not found. Re-scan the DDPM folder in Settings.") if not dataset.is_dir(): raise ToolExecutionError("The selected DDPM dataset folder no longer exists.") image_count = sum(1 for item in dataset.iterdir() if item.is_file() and item.suffix.lower() in IMAGE_EXTENSIONS) if image_count < 2: raise ToolExecutionError("The DDPM dataset needs at least two image files before training can start.") requested_epochs = int(epochs) if not 64 <= int(resolution) <= 512 or int(resolution) % 8 or not 1 <= int(batch_size) <= 64: raise ToolExecutionError("DDPM resolution must be a multiple of 8 (64–512) and batch size 1–64.") if not 1e-7 <= float(learning_rate) <= 0.1 or not 1 <= int(gradient_accumulation_steps) <= 64: raise ToolExecutionError("DDPM learning rate or gradient accumulation is outside ADAM's safe range.") if not 0 <= int(dataloader_num_workers) <= 16 or not 1 <= int(save_every) <= 1000 or not 1 <= int(preview_steps) <= 500 or not 1 <= int(preview_every) <= 100_000 or not 10 <= int(training_intensity) <= 100 or mixed_precision not in {"fp16", "no"}: raise ToolExecutionError("DDPM training options are outside ADAM's safe range.") if requested_epochs == 0: epochs = 200 if image_count <= 100 else 100 context.log( f"Adaptive epoch policy selected {epochs} epochs for {image_count} collected images." ) elif 1 <= requested_epochs <= 100_000: epochs = requested_epochs else: raise ToolExecutionError("Epoch count must be between 1 and 100000, or 0 for ADAM's adaptive policy.") missing_packages = [ package for package in ("datasets", "diffusers", "transformers", "accelerate", "torch", "torchvision") if importlib.util.find_spec(package) is None ] if missing_packages: raise ToolExecutionError( "ADAM's Python environment is missing DDPM packages: " + ", ".join(missing_packages) + ". Close ADAM and open Launch ADAM.bat; it will install the needed DDPM requirements. " f"Current Python: {sys.executable}" ) output_root = (trainer_root / "output").resolve() try: output.relative_to(output_root) except ValueError as exc: raise ToolExecutionError("DDPM outputs must stay inside the registered DDPM output folder.") from exc resume = Path(resume_from).expanduser().resolve() if resume_from else None pretrained_model: Path | None = None if resume: # Accelerate checkpoints store the training UNet below ``unet/`` plus # optimizer and scheduler state; they are not standalone pipelines. accelerate_checkpoint = ( (resume / "unet" / "diffusion_pytorch_model.safetensors").is_file() and (resume / "optimizer.bin").is_file() and (resume / "scheduler.bin").is_file() ) standalone_checkpoint = (resume / "pytorch_model.bin").is_file() or (resume / "model.safetensors").is_file() if not accelerate_checkpoint and not standalone_checkpoint: if not (output / "model_index.json").is_file(): raise ToolExecutionError("The saved DDPM model is incomplete and cannot be fine-tuned safely.") pretrained_model = output timestamp = time.strftime("%Y%m%d_%H%M%S") output = output.with_name(f"{output.name}_finetuned_{timestamp}") resume = None context.log( "The exact resume checkpoint is incomplete. Creating a new fine-tuned model " "from the saved DDPM pipeline instead." ) else: if not output.is_dir(): raise ToolExecutionError("The DDPM model folder for resume no longer exists.") try: resume.relative_to(output) except ValueError as exc: raise ToolExecutionError("The DDPM checkpoint must be inside its model folder.") from exc if not resume.is_dir() or not resume.name.startswith("checkpoint-"): raise ToolExecutionError("A valid DDPM checkpoint-* folder is required to resume.") steps_per_epoch = max(1, math.ceil(image_count / int(batch_size))) checkpoint_step = int(resume.name.rsplit("-", 1)[-1]) completed_epochs = checkpoint_step // steps_per_epoch epochs = completed_epochs + int(epochs) context.log( f"Continuing after approximately {completed_epochs} completed epochs " f"for {int(epochs) - completed_epochs} additional epochs." ) if not resume: if output.exists(): raise ToolExecutionError("The chosen DDPM output folder already exists; ADAM will not overwrite it.") output.mkdir(parents=True, exist_ok=False) # The connected trainer asks Accelerate/TensorBoard to write directly to # output/logs/train. Create it up front because its writer does not always # create the nested directory on Windows. (output / "logs" / "train").mkdir(parents=True, exist_ok=True) stop_file = output / ".adam_stop_training.flag" command = [ sys.executable, str(script), "--train_data_dir", str(dataset), "--output_dir", str(output), "--model_name", model_name, "--resolution", str(int(resolution)), "--train_batch_size", str(int(batch_size)), "--num_epochs", str(int(epochs)), "--learning_rate", str(float(learning_rate)), "--mixed_precision", mixed_precision, "--ddpm_beta_schedule", "linear", "--tf32", "true", "--save_images_epochs", str(int(preview_every) if preview_enabled else int(epochs) + 1), "--save_model_epochs", str(int(save_every)), "--training_intensity", str(int(training_intensity)), "--dataloader_num_workers", str(int(dataloader_num_workers)), "--gradient_accumulation_steps", str(int(gradient_accumulation_steps)), "--preview_num_inference_steps", str(int(preview_steps)), "--preview_sampler", "DDIM", "--pin_memory", "true", "--stop_signal_file", str(stop_file), "--checkpointing_steps", str(max(1, image_count)), "--checkpoints_total_limit", "1", "--keep_latest_resume_checkpoint", "--gui_progress", ] if resume: command.extend(["--resume_from_checkpoint", resume.name]) if pretrained_model: command.extend(["--pretrained_model_path", str(pretrained_model)]) context.log(f"Starting real DDPM training with {image_count} images at {resolution}px, batch {batch_size}, lr {learning_rate}.") context.log(f"Output folder: {output}") process = subprocess.Popen(command, cwd=str(trainer_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 DDPM trainer") stop_requested_at: float | None = None cancelled = False suspended = False 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("DDPM trainer paused safely." if suspended else "DDPM trainer resumed.") if context.cancel_event.is_set() and stop_requested_at is None: if suspended: set_process_tree_paused(process, False) suspended = False stop_file.touch(exist_ok=True) stop_requested_at = time.monotonic() cancelled = True context.log("Safe stop requested; waiting for DDPM to finish its current step.") if stop_requested_at and time.monotonic() - stop_requested_at > 75: process.terminate() context.log("DDPM did not stop in time; terminating the trainer process.") try: line = lines.get(timeout=0.12) if line: if line.startswith("PROGRESS_JSON:"): try: if not cancelled: _parse_progress(context, json.loads(line.split(":", 1)[1]), int(epochs), output, bool(preview_enabled), int(preview_every), preview_prompt, int(preview_seed), int(preview_steps)) except json.JSONDecodeError: context.log(line) else: context.log(line) except queue.Empty: pass if process.poll() is not None and lines.empty(): break if cancelled: raise ToolCancelled("DDPM training stopped by user.") if process.returncode != 0: raise ToolExecutionError(f"DDPM trainer exited with code {process.returncode}. See the job log for details.") context.progress(100, "DDPM training completed") checkpoints = sorted( output.glob("checkpoint-*"), key=lambda path: int(path.name.rsplit("-", 1)[-1]) if path.name.rsplit("-", 1)[-1].isdigit() else -1, ) latest_checkpoint = str(checkpoints[-1]) if checkpoints else "" return { "output_folder": str(output), "model_name": model_name, "assets": [ { "kind": "model", "name": model_name, "path": str(output), "trainer": "ddpm", "dataset_path": str(dataset), "checkpoint": latest_checkpoint, "epochs": int(epochs), } ], }