models/neural_cellular_automata/trainer.py

from __future__ import annotations

import json
import random
from pathlib import Path

import numpy as np
import torch
from PIL import Image
from torch.nn import functional as F

from .model import create_model, create_seed
from .image_preprocessing import image_to_tensor, prepare_image, randomly_flip


EXTENSIONS = {".png", ".jpg", ".jpeg", ".webp"}


def _device(config: dict) -> torch.device:
    requested = config.get("device", "auto")
    if requested == "cuda" and not torch.cuda.is_available():
        raise RuntimeError("CUDA was selected, but PyTorch cannot access an NVIDIA CUDA device.")
    return torch.device("cuda" if requested == "cuda" or (requested == "auto" and torch.cuda.is_available()) else "cpu")


def _load_targets(
    folder: str,
    resolution: int,
    fit_mode: str,
    background: str,
    training_mode: str = "Dataset average (experimental)",
    target_image: str = "",
) -> tuple[list[torch.Tensor], list[Path]]:
    if training_mode == "Single target":
        selected = Path(target_image).expanduser() if target_image else None
        if selected is None or not selected.is_file():
            raise ValueError("Single Target Growth requires an existing target image. Choose one on the Dataset tab or in Target Image.")
        if selected.suffix.lower() not in EXTENSIONS:
            raise ValueError("The selected target must be a PNG, JPG, JPEG, or WEBP image.")
        paths = [selected.resolve()]
    elif training_mode == "Dataset average (experimental)":
        paths = sorted(file for file in Path(folder).rglob("*") if file.suffix.lower() in EXTENSIONS)
    else:
        raise ValueError(f"Unsupported NCA training mode: {training_mode}")
    targets = []
    loaded_paths = []
    for path in paths:
        try:
            with Image.open(path) as image:
                prepared = prepare_image(image, resolution, fit_mode, background)
                targets.append(image_to_tensor(prepared))
                loaded_paths.append(path)
        except OSError:
            continue
    if not targets:
        raise ValueError("No readable PNG, JPG, JPEG, or WEBP images were found in the dataset folder.")
    return targets, loaded_paths


def _state_image(state: torch.Tensor) -> Image.Image:
    rgba = state[0, :4].detach().float().clamp(0, 1).permute(1, 2, 0).cpu().numpy()
    return Image.fromarray((rgba * 255).astype(np.uint8), "RGBA")


def _save_checkpoint(path, model, optimizer, scaler, config, losses, completed_step, best_smoothed_loss):
    temporary = path.with_suffix(".pt.tmp")
    torch.save(
        {
            "format_version": 2,
            "model_state": model.state_dict(),
            "optimizer_state": optimizer.state_dict(),
            "scaler_state": scaler.state_dict(),
            "config": dict(config),
            "loss_history": list(losses),
            "completed_step": completed_step,
            "best_smoothed_loss": best_smoothed_loss,
            "python_random_state": random.getstate(),
            "numpy_random_state": np.random.get_state(),
            "torch_random_state": torch.get_rng_state(),
        },
        temporary,
    )
    temporary.replace(path)


def train_model(config, dataset_path, output_path, progress_callback=None, preview_callback=None, stop_event=None):
    output = Path(output_path)
    output.mkdir(parents=True, exist_ok=True)
    (output / "previews").mkdir(exist_ok=True)

    seed = int(config.get("random_seed", 42))
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    device = _device(config)
    if device.type == "cuda":
        torch.cuda.manual_seed_all(seed)

    resolution = int(config["resolution"])
    channels = int(config["cell_channels"])
    batch_size = int(config["batch_size"])
    total = int(config["epochs"])
    min_steps = int(config["min_growth_steps"])
    max_steps = int(config["max_growth_steps"])
    if min_steps > max_steps:
        raise ValueError("Minimum growth steps cannot exceed maximum growth steps.")
    stable_growth = bool(config.get("stable_growth", False))
    stable_max_steps = int(config.get("stable_max_growth_steps", max_steps))
    if stable_growth and stable_max_steps < max_steps:
        raise ValueError("Stable maximum growth steps cannot be less than Maximum Growth Steps.")
    training_max_steps = stable_max_steps if stable_growth else max_steps

    training_mode = str(config.get("training_mode", "Dataset average (experimental)"))
    targets, target_paths = _load_targets(
        dataset_path,
        resolution,
        str(config.get("fit_mode", "contain")),
        str(config.get("background", "transparent")),
        training_mode,
        str(config.get("target_image", "")),
    )
    model = create_model(config).to(device)
    optimizer = torch.optim.Adam(model.parameters(), lr=float(config["learning_rate"]))
    use_amp = bool(config.get("mixed_precision", True)) and device.type == "cuda"
    scaler = torch.amp.GradScaler("cuda", enabled=use_amp)
    losses: list[float] = []
    start_step = 1
    best_smoothed_loss = float("inf")
    resume_path = str(config.get("resume_checkpoint", "")).strip()
    if resume_path:
        checkpoint = torch.load(resume_path, map_location=device, weights_only=False)
        old_config = checkpoint.get("config", {})
        for key in ("resolution", "cell_channels", "hidden_size"):
            if int(old_config.get(key, config[key])) != int(config[key]):
                raise ValueError(f"Cannot resume because {key} differs from the checkpoint.")
        model.load_state_dict(checkpoint["model_state"])
        if "optimizer_state" in checkpoint:
            optimizer.load_state_dict(checkpoint["optimizer_state"])
        if "scaler_state" in checkpoint:
            scaler.load_state_dict(checkpoint["scaler_state"])
        losses = list(checkpoint.get("loss_history", []))
        best_smoothed_loss = float(checkpoint.get("best_smoothed_loss", float("inf")))
        start_step = int(checkpoint.get("completed_step", len(losses))) + 1
        if "python_random_state" in checkpoint:
            random.setstate(checkpoint["python_random_state"])
        if "numpy_random_state" in checkpoint:
            np.random.set_state(checkpoint["numpy_random_state"])
        if "torch_random_state" in checkpoint:
            torch.set_rng_state(checkpoint["torch_random_state"].cpu())
    log_path = output / "training_log.txt"
    checkpoint_path = output / "checkpoint.pt"
    best_checkpoint_path = output / "best_checkpoint.pt"
    completed_step = start_step - 1

    with log_path.open("a", encoding="utf-8", buffering=1) as log:
        log.write(f"Training on {device}; mode={training_mode}; {len(targets)} target(s); mixed precision={use_amp}\n")
        log.write(f"Growth rollout range: {min_steps}-{training_max_steps}; stable growth={stable_growth}\n")
        if training_mode == "Single target":
            log.write(f"Target image: {target_paths[0]}\n")
        if resume_path:
            log.write(f"Resumed from {resume_path} at step {start_step}.\n")
        for step in range(start_step, total + 1):
            if stop_event is not None and stop_event.is_set():
                log.write("Training stopped by user.\n")
                break
            selected = random.choices(targets, k=batch_size)
            target = torch.stack(selected).to(device)
            if randomly_flip(bool(config.get("horizontal_flip", False))):
                target = torch.flip(target, dims=(3,))
            state = create_seed(batch_size, channels, resolution, device)
            growth_steps = random.randint(min_steps, training_max_steps)

            optimizer.zero_grad(set_to_none=True)
            with torch.amp.autocast(device_type=device.type, enabled=use_amp):
                for _ in range(growth_steps):
                    state = model(state)
                loss = F.mse_loss(state[:, :4], target)
            scaler.scale(loss).backward()
            scaler.unscale_(optimizer)
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            scaler.step(optimizer)
            scaler.update()

            value = float(loss.detach().cpu())
            losses.append(value)
            smoothing_window = max(1, int(config.get("loss_smoothing_window", 25)))
            smoothed_loss = sum(losses[-smoothing_window:]) / min(len(losses), smoothing_window)
            completed_step = step
            update = {"step": step, "total_steps": total, "epoch": step, "total_epochs": total, "loss": value, "smoothed_loss": smoothed_loss, "device": str(device), "growth_steps": growth_steps}
            if progress_callback:
                progress_callback(update)
            log.write(json.dumps(update) + "\n")

            preview_every = max(1, int(config.get("preview_every", 10)))
            if step == 1 or step % preview_every == 0 or step == total:
                preview = _state_image(state)
                preview.save(output / "previews" / f"step_{step:06d}.png")
                if preview_callback:
                    preview_callback(preview.copy())

            checkpoint_every = max(1, int(config.get("checkpoint_every", 50)))
            if step % checkpoint_every == 0:
                _save_checkpoint(checkpoint_path, model, optimizer, scaler, config, losses, completed_step, best_smoothed_loss)
                log.write(f"Saved recovery checkpoint at step {step}.\n")
                if smoothed_loss < best_smoothed_loss:
                    best_smoothed_loss = smoothed_loss
                    _save_checkpoint(best_checkpoint_path, model, optimizer, scaler, config, losses, completed_step, best_smoothed_loss)
                    log.write(f"Saved best checkpoint at step {step}; smoothed loss={smoothed_loss:.6f}.\n")

    final_smoothed = sum(losses[-max(1, int(config.get("loss_smoothing_window", 25))):]) / min(
        len(losses), max(1, int(config.get("loss_smoothing_window", 25)))
    ) if losses else float("inf")
    _save_checkpoint(checkpoint_path, model, optimizer, scaler, config, losses, completed_step, best_smoothed_loss)
    if not best_checkpoint_path.exists() or final_smoothed < best_smoothed_loss:
        best_smoothed_loss = final_smoothed
        _save_checkpoint(best_checkpoint_path, model, optimizer, scaler, config, losses, completed_step, best_smoothed_loss)
    return {
        "checkpoint_path": str(checkpoint_path),
        "best_checkpoint_path": str(best_checkpoint_path),
        "best_smoothed_loss": best_smoothed_loss,
        "loss_history": losses,
        "completed_step": completed_step,
        "stopped": bool(stop_event and stop_event.is_set()),
    }


def train(
    context,
    dataset_dir: str,
    model_name: str,
    epochs: int,
    output_dir: str,
    resume_from: str = "",
    **settings,
):
    from .manifest import TRAINING_SETTINGS

    config = {
        key: spec.get("default")
        for key, spec in TRAINING_SETTINGS.items()
        if isinstance(spec, dict) and "default" in spec
    }
    config.update(settings)
    config["epochs"] = int(epochs)
    if resume_from:
        config["resume_checkpoint"] = resume_from

    output = Path(output_dir)

    def progress(update):
        step = int(update.get("step", update.get("epoch", 0)) or 0)
        total = int(update.get("total_steps", update.get("total_epochs", epochs)) or epochs)
        loss = update.get("loss")
        message = f"NCA step {step:,} of {total:,}"
        if loss is not None:
            message += f" · loss {float(loss):.5f}"
        context.progress(max(1, min(99, round(step * 100 / max(total, 1)))), message)

    result = train_model(
        config,
        dataset_dir,
        output,
        progress_callback=progress,
        stop_event=context.cancel_event,
    )
    checkpoint = str(result.get("checkpoint_path") or output / "checkpoint.pt")
    preview = output / "previews"
    if preview.is_dir():
        latest = max(preview.glob("*.png"), key=lambda path: path.stat().st_mtime, default=None)
        if latest:
            context.preview(latest, epoch=int(result.get("completed_step", epochs) or epochs))
    context.progress(100, "NCA training completed")
    return {
        "output_folder": str(output),
        "assets": [
            {
                "kind": "model",
                "name": model_name,
                "path": str(output),
                "trainer": "neural_cellular_automata",
                "dataset_path": str(Path(dataset_dir).resolve()),
                "checkpoint": checkpoint,
                "epochs": int(epochs),
            }
        ],
    }