adam/tools/lora_adapter.py

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

from __future__ import annotations

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

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


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


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


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 train_lora(
    context: ToolContext,
    dataset_dir: str,
    model_name: str,
    epochs: int,
    output_dir: str,
    base_model: str,
    resume_from: str = "",
    preview_enabled: bool = True, preview_every: int = 5,
    preview_prompt: str = "", preview_seed: int = 123456789,
) -> 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)
    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

    sys.path.insert(0, str(source_root))
    try:
        from loratrainer.models.training_config import TrainingConfig
        from loratrainer.trainer.diffusers_sdxl_lora_backend import (
            DiffusersSDXLLoRABackend,
        )
    except Exception as exc:
        raise ToolExecutionError(f"Could not load the connected LoRA backend: {exc}") from exc

    valid_fields = set(TrainingConfig.__dataclass_fields__)
    overrides = {
        key: value
        for key, value in settings.items()
        if key in valid_fields
        and key not in {"dataset_dir", "base_model_path", "output_dir", "resume_checkpoint"}
    }
    overrides["trigger_word"] = name
    overrides["epochs"] = int(epochs)
    config = TrainingConfig(
        dataset_dir=dataset,
        base_model_path=base,
        output_dir=output,
        resume_checkpoint=resume,
        **overrides,
    )
    control = _Control(context)

    def progress(update: Any) -> None:
        if context.cancel_event.is_set():
            return
        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", epochs) or epochs)
        percent = (
            round(step * 100 / total_steps)
            if total_steps
            else round(epoch * 100 / max(total_epochs, 1))
        )
        message = str(getattr(update, "message", "") or f"LoRA epoch {epoch}/{total_epochs}")
        context.progress(max(1, min(percent, 99)), message)
        preview_path = str(getattr(update, "preview_path", "") or "")
        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}")
    try:
        final_path = Path(
            DiffusersSDXLLoRABackend().train(config, control, progress)
        ).resolve()
    except Exception as exc:
        if context.cancel_event.is_set():
            raise ToolCancelled("LoRA training stopped by user.") from exc
        raise ToolExecutionError(f"LoRA trainer failed: {exc}") from exc

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