File size: 7,079 Bytes
e0265b9 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 | """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),
}
],
}
|