"""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), } ], }