adam/tools/lora_adapter.py

"""Execution bridge for the connected Local SDXL LoRA Trainer project."""

from __future__ import annotations

import json
import multiprocessing
import queue
import sys
import time
from pathlib import Path
from typing import Any

from adam.config import ConfigManager
from adam.executor import ToolCancelled, ToolContext, ToolExecutionError
from adam.process_control import terminate_process_tree


IMAGE_EXTENSIONS = {".jpg", ".jpeg", ".png", ".webp", ".bmp"}
FORCE_STOP_TIMEOUT_SECONDS = 30


class _Control:
    def __init__(self, context: ToolContext) -> None:
        self.context = context

    @property
    def cancel_requested(self) -> bool:
        return self.context.cancel_event.is_set()

    @property
    def pause_requested(self) -> bool:
        return not self.context.run_event.is_set()

    def wait_if_paused(self) -> None:
        try:
            self.context.checkpoint()
        except ToolCancelled:
            # The connected backend sees this flag on its next control check and
            # writes its normal cancelled snapshot before stopping.
            return


class _ProcessControl:
    def __init__(self, cancel_event: Any, run_event: Any) -> None:
        self._cancel_event = cancel_event
        self._run_event = run_event

    @property
    def cancel_requested(self) -> bool:
        return self._cancel_event.is_set()

    @property
    def pause_requested(self) -> bool:
        return not self._run_event.is_set()

    def wait_if_paused(self) -> None:
        while not self._run_event.wait(timeout=0.15):
            if self._cancel_event.is_set():
                return


def _safe_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 LoRA name without reserved characters.")
    return name


def _run_lora_training_worker(
    payload: dict[str, Any],
    events: Any,
    cancel_event: Any,
    run_event: Any,
) -> None:
    try:
        source_root = Path(str(payload["source_root"]))
        sys.path.insert(0, str(source_root))
        from loratrainer.models.training_config import TrainingConfig
        from loratrainer.trainer.diffusers_sdxl_lora_backend import (
            DiffusersSDXLLoRABackend,
        )

        valid_fields = set(TrainingConfig.__dataclass_fields__)
        blocked = {"dataset_dir", "base_model_path", "output_dir", "resume_checkpoint"}
        overrides = {
            key: value
            for key, value in dict(payload.get("settings", {})).items()
            if key in valid_fields and key not in blocked
        }
        overrides.update(
            {
                key: value
                for key, value in dict(payload.get("training_overrides", {})).items()
                if key in valid_fields and key not in {*blocked, "epochs"}
            }
        )
        overrides["trigger_word"] = str(payload.get("trigger_word") or payload["model_name"])
        overrides["epochs"] = int(payload["epochs"])
        config = TrainingConfig(
            dataset_dir=Path(str(payload["dataset"])),
            base_model_path=Path(str(payload["base"])),
            output_dir=Path(str(payload["output"])),
            resume_checkpoint=Path(str(payload["resume"])) if payload.get("resume") else None,
            **overrides,
        )
        control = _ProcessControl(cancel_event, run_event)

        def progress(update: Any) -> None:
            if cancel_event.is_set():
                return
            events.put(
                {
                    "type": "progress",
                    "total_steps": int(getattr(update, "total_steps", 0) or 0),
                    "step": int(getattr(update, "step", 0) or 0),
                    "epoch": int(getattr(update, "epoch", 0) or 0),
                    "total_epochs": int(getattr(update, "total_epochs", payload["epochs"]) or payload["epochs"]),
                    "message": str(getattr(update, "message", "") or ""),
                    "preview_path": str(getattr(update, "preview_path", "") or ""),
                }
            )

        final_path = Path(
            DiffusersSDXLLoRABackend().train(config, control, progress)
        ).resolve()
        events.put({"type": "result", "final_path": str(final_path)})
    except Exception as exc:
        if cancel_event.is_set():
            events.put({"type": "cancelled", "message": "LoRA training stopped by user."})
            return
        events.put(
            {
                "type": "error",
                "message": str(exc),
                "exception": type(exc).__name__,
            }
        )


def train_lora(
    context: ToolContext,
    dataset_dir: str,
    model_name: str,
    epochs: int,
    output_dir: str,
    base_model: str,
    trigger_word: str = "",
    resume_from: str = "",
    preview_enabled: bool = True, preview_every: int = 5,
    preview_prompt: str = "", preview_seed: int = 123456789,
    **training_overrides: Any,
) -> dict[str, Any]:
    folders = ConfigManager(context.root).get("tool_folders", {})
    trainer_root = Path(str(folders.get("lora_trainer", ""))).expanduser().resolve()
    source_root = trainer_root / "src"
    backend_file = source_root / "loratrainer" / "trainer" / "diffusers_sdxl_lora_backend.py"
    if not backend_file.is_file():
        raise ToolExecutionError(
            "The connected LoRA folder does not contain its native Diffusers backend."
        )

    dataset = Path(dataset_dir).expanduser().resolve()
    base = Path(base_model).expanduser().resolve()
    output = Path(output_dir).expanduser().resolve()
    resume = Path(resume_from).expanduser().resolve() if resume_from else None
    name = _safe_name(model_name)
    trigger = str(trigger_word or name).strip()
    if not trigger or len(trigger) > 128 or any(char in trigger for char in '<>:"/\\|?*\x00'):
        raise ToolExecutionError("Choose a short LoRA trigger word without reserved characters.")
    if not dataset.is_dir():
        raise ToolExecutionError("The selected LoRA dataset folder no longer exists.")
    images = [
        item for item in dataset.iterdir()
        if item.is_file() and item.suffix.casefold() in IMAGE_EXTENSIONS
    ]
    if len(images) < 2:
        raise ToolExecutionError("The LoRA dataset needs at least two images.")
    missing_captions = [item for item in images if not item.with_suffix(".txt").is_file()]
    if missing_captions:
        raise ToolExecutionError(
            f"The LoRA dataset is missing captions for {len(missing_captions)} image(s)."
        )
    if not base.is_file():
        raise ToolExecutionError("The selected SDXL base model does not exist.")
    if resume and not resume.is_file():
        raise ToolExecutionError("The selected LoRA checkpoint does not exist.")
    if not 1 <= int(epochs) <= 100_000:
        raise ToolExecutionError("LoRA epochs must be between 1 and 100000.")
    output_root = (trainer_root / "output").resolve()
    try:
        output.relative_to(output_root)
    except ValueError as exc:
        raise ToolExecutionError(
            "LoRA outputs must stay inside the connected trainer's output folder."
        ) from exc
    if resume:
        output = output.with_name(
            f"{output.name}_finetuned_{time.strftime('%Y%m%d_%H%M%S')}"
        )
    if output.exists():
        raise ToolExecutionError("The chosen LoRA output already exists; ADAM will not overwrite it.")
    output.mkdir(parents=True)

    settings: dict[str, Any] = {}
    settings_path = trainer_root / "config" / "app_settings.json"
    try:
        payload = json.loads(settings_path.read_text(encoding="utf-8"))
        settings = dict(payload.get("training_settings", {}))
    except (OSError, ValueError, TypeError, json.JSONDecodeError):
        pass

    def relay_progress(update: dict[str, Any]) -> None:
        total_steps = int(getattr(update, "total_steps", 0) or 0)
        if isinstance(update, dict):
            total_steps = int(update.get("total_steps", 0) or 0)
            step = int(update.get("step", 0) or 0)
            epoch = int(update.get("epoch", 0) or 0)
            total_epochs = int(update.get("total_epochs", epochs) or epochs)
            message = str(update.get("message", "") or f"LoRA epoch {epoch}/{total_epochs}")
            preview_path = str(update.get("preview_path", "") or "")
        else:
            step = int(getattr(update, "step", 0) or 0)
            epoch = int(getattr(update, "epoch", 0) or 0)
            total_epochs = int(getattr(update, "total_epochs", epochs) or epochs)
            message = str(getattr(update, "message", "") or f"LoRA epoch {epoch}/{total_epochs}")
            preview_path = str(getattr(update, "preview_path", "") or "")
        percent = (
            round(step * 100 / total_steps)
            if total_steps
            else round(epoch * 100 / max(total_epochs, 1))
        )
        context.progress(
            max(1, min(percent, 99)),
            message,
            current_step=step,
            total_steps=total_steps,
            epoch=epoch,
            total_epochs=total_epochs,
            unit="step" if total_steps else "epoch",
        )
        if preview_enabled and preview_path and epoch and epoch % max(1, int(preview_every)) == 0:
            context.preview(preview_path, epoch=epoch,
                            next_epoch=min(total_epochs, epoch + max(1, int(preview_every))),
                            prompt=preview_prompt, seed=int(preview_seed))

    context.log(f"Starting real LoRA training with {len(images)} captioned images.")
    context.log(f"Base model: {base}")
    if resume:
        context.log(f"Continuing from LoRA: {resume}")
    payload = {
        "source_root": str(source_root),
        "dataset": str(dataset),
        "base": str(base),
        "output": str(output),
        "resume": str(resume) if resume else "",
        "model_name": name,
        "trigger_word": trigger,
        "epochs": int(epochs),
        "settings": settings,
        # ``preview_prompt`` is a named argument, so Python removes it from
        # ``training_overrides``.  Explicitly include it here so a new ADAM
        # run cannot inherit the connected trainer's last saved prompt.
        "training_overrides": {
            **training_overrides,
            "preview_prompt": str(preview_prompt),
            "preview_interval_epochs": max(1, int(preview_every)),
        },
    }
    mp_context = multiprocessing.get_context("spawn")
    events = mp_context.Queue()
    process_cancel = mp_context.Event()
    process_run = mp_context.Event()
    process_run.set()
    process = mp_context.Process(
        target=_run_lora_training_worker,
        args=(payload, events, process_cancel, process_run),
        daemon=True,
    )
    process.start()
    final_path: Path | None = None
    cancelled = False
    stop_requested_at: float | None = None
    force_stop_sent = False
    process_paused = False

    def handle_child_event(event: dict[str, Any]) -> None:
        nonlocal final_path, cancelled
        event_type = str(event.get("type", ""))
        if event_type == "progress" and not cancelled:
            relay_progress(event)
        elif event_type == "result":
            final_path = Path(str(event.get("final_path", ""))).resolve()
        elif event_type == "cancelled":
            cancelled = True
        elif event_type == "error":
            message = str(event.get("message", "LoRA trainer failed."))
            terminate_process_tree(process, timeout=3)
            process.join(timeout=1)
            raise ToolExecutionError(f"LoRA trainer failed: {message}")

    while True:
        if context.cancel_event.is_set() and stop_requested_at is None:
            process_cancel.set()
            process_run.set()
            stop_requested_at = time.monotonic()
            cancelled = True
            context.log("Safe stop requested; waiting for LoRA training to finish its current step.")
        should_pause = not context.run_event.is_set() and not cancelled
        if should_pause != process_paused:
            if should_pause:
                process_run.clear()
                context.log("LoRA trainer paused safely.")
            else:
                process_run.set()
                context.log("LoRA trainer resumed.")
            process_paused = should_pause
        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("LoRA trainer did not stop in time; terminating the trainer process.")
        try:
            event = events.get(timeout=0.15)
        except queue.Empty:
            event = None
        if event:
            handle_child_event(event)
        if not process.is_alive() and event is None:
            break

    process.join(timeout=1)
    while True:
        try:
            handle_child_event(events.get_nowait())
        except queue.Empty:
            break
    if cancelled:
        raise ToolCancelled("LoRA training stopped by user.")
    if process.exitcode not in {0, None}:
        raise ToolExecutionError(f"LoRA trainer exited with code {process.exitcode}.")
    if final_path is None:
        raise ToolExecutionError("LoRA trainer finished without reporting a checkpoint.")

    context.progress(100, "LoRA training completed")
    return {
        "output_folder": str(output),
        "model_name": name,
        "trigger_word": trigger,
        "assets": [
            {
                "kind": "model",
                "name": name,
                "path": str(output),
                "trainer": "lora",
                "dataset_path": str(dataset),
                "checkpoint": str(final_path),
                "epochs": int(epochs),
                "metadata": {"trigger_word": trigger},
                "trigger_word": trigger,
            }
        ],
    }