Blame for adam/tools/ddpm_adapter.py

"""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 ToolAdjustmentRequested, 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 _saved_unet_resolution(path: Path) -> int | None:
    config_path = path / "unet" / "config.json"
    if not config_path.is_file():
        return None
    try:
        data = json.loads(config_path.read_text(encoding="utf-8"))
    except (OSError, json.JSONDecodeError):
        return None
    sample_size = data.get("sample_size")
    if isinstance(sample_size, list):
        sample_size = sample_size[0] if sample_size else None
    try:
        return int(sample_size) if sample_size is not None else None
    except (TypeError, ValueError):
        return None


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:,}",
                current_step=current,
                total_steps=total,
                epoch=int(payload.get("epoch", 0) or 0),
                total_epochs=epochs,
                unit="step",
            )
        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}",
            epoch=epoch,
            total_epochs=epochs,
            unit="epoch",
        )
        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,
    completed_epochs: int = 0,
) -> 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.")
            checkpoint_resolution = _saved_unet_resolution(resume)
            if checkpoint_resolution and checkpoint_resolution != int(resolution) and (output / "model_index.json").is_file():
                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(
                    f"Changing DDPM resolution from {checkpoint_resolution}px to {int(resolution)}px. "
                    "Starting a fresh fine-tune from the saved model weights instead of resuming the old optimizer schedule."
                )
            if resume is not None:
                steps_per_epoch = max(1, math.ceil(image_count / int(batch_size)))
                checkpoint_step = int(resume.name.rsplit("-", 1)[-1])
                prior_epochs = int(completed_epochs) if int(completed_epochs) > 0 else checkpoint_step // steps_per_epoch
                epochs = prior_epochs + int(epochs)
                context.log(
                    f"Continuing after approximately {prior_epochs} completed epochs "
                    f"for {int(epochs) - prior_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, math.ceil(image_count / int(batch_size)))),
        "--checkpoints_total_limit", "1", "--keep_latest_resume_checkpoint", "--gui_progress",
    ]
    if resume:
        command.extend(["--resume_from_checkpoint", resume.name])
        if int(completed_epochs) > 0:
            command.extend(["--resume_completed_epochs", str(int(completed_epochs))])
    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
    force_stop_sent = False
    suspended = False
    adjustment_requested = False
    stopped_details: dict[str, object] = {}
    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.adjustment_event and context.adjustment_event.is_set() and not adjustment_requested:
            if suspended:
                set_process_tree_paused(process, False)
                suspended = False
            stop_file.write_text("after_epoch", encoding="utf-8")
            adjustment_requested = True
            context.log("Settings change accepted; finishing this epoch and saving a resume checkpoint.")
        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 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("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:
                        payload = json.loads(line.split(":", 1)[1])
                        if str(payload.get("event", "")) == "stopped":
                            stopped_details = payload
                        if not cancelled:
                            _parse_progress(context, payload, 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 adjustment_requested:
        checkpoints = sorted(output.glob("checkpoint-*"), key=lambda path: int(path.name.rsplit("-", 1)[-1]))
        if not checkpoints:
            raise ToolExecutionError("Training stopped for adjustment, but no complete resume checkpoint was found.")
        raise ToolAdjustmentRequested(
            "DDPM training reached a safe epoch boundary.",
            {
                "checkpoint": str(checkpoints[-1]),
                "completed_epochs": int(stopped_details.get("completed_epochs", 0) or 0),
                "updates": dict(context.adjustment_request or {}),
            },
        )
    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),
            }
        ],
    }