from __future__ import annotations import json import re import sqlite3 from dataclasses import dataclass from datetime import datetime, timezone from pathlib import Path from typing import Any from adam.models import Job, JobStatus, SystemSnapshot def _utc_now() -> str: return datetime.now(timezone.utc).isoformat() def _json(value: Any) -> str: return json.dumps(value, sort_keys=True) def _safe_json(value: str, fallback: Any) -> Any: try: parsed = json.loads(value or "") except (TypeError, ValueError, json.JSONDecodeError): return fallback return parsed if isinstance(parsed, type(fallback)) else fallback def _safe_int(value: Any, default: int = 0) -> int: try: if isinstance(value, bool): return default return int(value) except (TypeError, ValueError): return default def _safe_float(value: Any, default: float = 0.0) -> float: try: if isinstance(value, bool): return default return float(value) except (TypeError, ValueError): return default def _loss_from_logs(logs: list[str]) -> float | None: for line in reversed(logs): match = re.search(r"\bloss(?:\s*[:=]\s*|\s+)(-?\d+(?:\.\d+)?(?:e[+-]?\d+)?)", line, re.I) if match: try: return float(match.group(1)) except ValueError: return None return None def _duration_seconds(job: Job) -> int: if not job.started_at: return 0 try: start = datetime.fromisoformat(job.started_at) end = datetime.fromisoformat(job.ended_at) if job.ended_at else datetime.now(timezone.utc) return max(0, int((end - start).total_seconds())) except ValueError: return 0 def _image_count(path: str) -> int: folder = Path(path).expanduser() if not folder.is_dir(): return 0 try: return sum( 1 for item in folder.rglob("*") if item.is_file() and item.suffix.casefold() in {".png", ".jpg", ".jpeg", ".webp", ".bmp"} ) except OSError: return 0 @dataclass(slots=True) class ExperimentRun: id: str job_id: str timestamp: str model_architecture: str model_name: str trigger_word: str base_model: str dataset_path: str dataset_name: str dataset_item_count: int epochs: int batch_size: int learning_rate: float optimizer: str scheduler: str resolution: int seed: int status: str training_time_seconds: int final_loss: float | None output_folder: str checkpoint_paths: list[str] preview_images: list[str] peak_vram_gb: float | None hardware: dict[str, Any] settings: dict[str, Any] generation_settings: dict[str, Any] notes: str = "" quality_score: int | None = None @classmethod def from_row(cls, row: sqlite3.Row) -> "ExperimentRun": payload = dict(row) for key in ("checkpoint_paths", "preview_images"): payload[key] = _safe_json(payload.get(key, "[]"), []) for key in ("hardware", "settings", "generation_settings"): payload[key] = _safe_json(payload.get(key, "{}"), {}) payload["quality_score"] = ( _safe_int(payload["quality_score"]) if payload.get("quality_score") is not None else None ) return cls(**payload) class ExperimentStore: def __init__(self, root: Path) -> None: self.root = root.resolve() self.path = self.root / "data" / "experiments.sqlite3" self.path.parent.mkdir(parents=True, exist_ok=True) self._init_db() def connect(self) -> sqlite3.Connection: connection = sqlite3.connect(self.path) connection.row_factory = sqlite3.Row return connection def _init_db(self) -> None: with self.connect() as db: db.execute( """ CREATE TABLE IF NOT EXISTS experiments ( id TEXT PRIMARY KEY, job_id TEXT UNIQUE NOT NULL, timestamp TEXT NOT NULL, model_architecture TEXT NOT NULL, model_name TEXT NOT NULL, trigger_word TEXT NOT NULL DEFAULT '', base_model TEXT NOT NULL, dataset_path TEXT NOT NULL, dataset_name TEXT NOT NULL, dataset_item_count INTEGER NOT NULL, epochs INTEGER NOT NULL, batch_size INTEGER NOT NULL, learning_rate REAL NOT NULL, optimizer TEXT NOT NULL, scheduler TEXT NOT NULL, resolution INTEGER NOT NULL, seed INTEGER NOT NULL, status TEXT NOT NULL, training_time_seconds INTEGER NOT NULL, final_loss REAL, output_folder TEXT NOT NULL, checkpoint_paths TEXT NOT NULL, preview_images TEXT NOT NULL, peak_vram_gb REAL, hardware TEXT NOT NULL, settings TEXT NOT NULL, generation_settings TEXT NOT NULL, notes TEXT NOT NULL DEFAULT '', quality_score INTEGER ) """ ) self._migrate_columns(db) @staticmethod def _migrate_columns(db: sqlite3.Connection) -> None: existing = {row["name"] for row in db.execute("PRAGMA table_info(experiments)").fetchall()} columns = { "id": "TEXT PRIMARY KEY", "job_id": "TEXT NOT NULL DEFAULT ''", "timestamp": "TEXT NOT NULL DEFAULT ''", "model_architecture": "TEXT NOT NULL DEFAULT ''", "model_name": "TEXT NOT NULL DEFAULT ''", "trigger_word": "TEXT NOT NULL DEFAULT ''", "base_model": "TEXT NOT NULL DEFAULT ''", "dataset_path": "TEXT NOT NULL DEFAULT ''", "dataset_name": "TEXT NOT NULL DEFAULT ''", "dataset_item_count": "INTEGER NOT NULL DEFAULT 0", "epochs": "INTEGER NOT NULL DEFAULT 0", "batch_size": "INTEGER NOT NULL DEFAULT 0", "learning_rate": "REAL NOT NULL DEFAULT 0", "optimizer": "TEXT NOT NULL DEFAULT ''", "scheduler": "TEXT NOT NULL DEFAULT ''", "resolution": "INTEGER NOT NULL DEFAULT 0", "seed": "INTEGER NOT NULL DEFAULT 0", "status": "TEXT NOT NULL DEFAULT ''", "training_time_seconds": "INTEGER NOT NULL DEFAULT 0", "final_loss": "REAL", "output_folder": "TEXT NOT NULL DEFAULT ''", "checkpoint_paths": "TEXT NOT NULL DEFAULT '[]'", "preview_images": "TEXT NOT NULL DEFAULT '[]'", "peak_vram_gb": "REAL", "hardware": "TEXT NOT NULL DEFAULT '{}'", "settings": "TEXT NOT NULL DEFAULT '{}'", "generation_settings": "TEXT NOT NULL DEFAULT '{}'", "notes": "TEXT NOT NULL DEFAULT ''", "quality_score": "INTEGER", } for name, definition in columns.items(): if name not in existing and name != "id": db.execute(f"ALTER TABLE experiments ADD COLUMN {name} {definition}") def record_job(self, job: Job, snapshot: SystemSnapshot | None = None) -> ExperimentRun | None: training_steps = [step for step in job.plan.steps if step.tool_id.endswith("_trainer")] if not training_steps: return None step = training_steps[-1] args = dict(step.arguments) architecture = step.tool_id.removesuffix("_trainer") dataset_path = str(args.get("dataset_dir", "")) output_folder = str(job.output_folder or args.get("output_dir", "")) preview_images = [job.preview_path] if job.preview_path else [] checkpoints = [] if output_folder: folder = Path(output_folder) if folder.is_dir(): try: checkpoints = [ str(path) for path in sorted(folder.rglob("*")) if path.is_file() and path.suffix.casefold() in {".safetensors", ".ckpt", ".pt", ".bin"} ][-10:] discovered_previews = [ str(path) for path in sorted(folder.rglob("*")) if path.is_file() and path.suffix.casefold() in {".png", ".jpg", ".jpeg", ".webp", ".bmp"} and any(token in path.name.casefold() for token in ("preview", "sample", "epoch")) ][-12:] preview_images = list(dict.fromkeys([*preview_images, *discovered_previews])) except OSError: checkpoints = [] hardware = {} peak_vram = None if snapshot is not None: hardware = { "gpu_name": snapshot.gpu_name, "gpu_percent": snapshot.gpu_percent, "vram_used_gb": snapshot.vram_used_gb, "vram_total_gb": snapshot.vram_total_gb, "memory_used_gb": snapshot.memory_used_gb, "memory_total_gb": snapshot.memory_total_gb, "cpu_percent": snapshot.cpu_percent, "gpu_temperature": snapshot.gpu_temperature, } peak_vram = snapshot.vram_used_gb or None run = ExperimentRun( id=f"EXP-{job.id}", job_id=job.id, timestamp=job.ended_at or _utc_now(), model_architecture=architecture, model_name=str(args.get("model_name", job.plan.project_name)), trigger_word=str(args.get("trigger_word", "")), base_model=str(args.get("base_model", args.get("base_model_path", ""))), dataset_path=dataset_path, dataset_name=Path(dataset_path).name if dataset_path else "", dataset_item_count=_image_count(dataset_path), epochs=_safe_int(args.get("epochs")), batch_size=_safe_int(args.get("batch_size")), learning_rate=_safe_float(args.get("learning_rate")), optimizer=str(args.get("optimizer", "")), scheduler=str(args.get("scheduler", args.get("sampler", ""))), resolution=_safe_int(args.get("resolution")), seed=_safe_int(args.get("seed", args.get("preview_seed", 0))), status=job.status.value, training_time_seconds=_duration_seconds(job), final_loss=_loss_from_logs(job.logs), output_folder=output_folder, checkpoint_paths=checkpoints, preview_images=preview_images, peak_vram_gb=peak_vram, hardware=hardware, settings=args, generation_settings={}, ) with self.connect() as db: db.execute( """ INSERT INTO experiments ( id, job_id, timestamp, model_architecture, model_name, trigger_word, base_model, dataset_path, dataset_name, dataset_item_count, epochs, batch_size, learning_rate, optimizer, scheduler, resolution, seed, status, training_time_seconds, final_loss, output_folder, checkpoint_paths, preview_images, peak_vram_gb, hardware, settings, generation_settings ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) ON CONFLICT(job_id) DO UPDATE SET timestamp=excluded.timestamp, trigger_word=excluded.trigger_word, status=excluded.status, training_time_seconds=excluded.training_time_seconds, final_loss=excluded.final_loss, output_folder=excluded.output_folder, checkpoint_paths=excluded.checkpoint_paths, preview_images=excluded.preview_images, peak_vram_gb=excluded.peak_vram_gb, hardware=excluded.hardware, settings=excluded.settings """, ( run.id, run.job_id, run.timestamp, run.model_architecture, run.model_name, run.trigger_word, run.base_model, run.dataset_path, run.dataset_name, run.dataset_item_count, run.epochs, run.batch_size, run.learning_rate, run.optimizer, run.scheduler, run.resolution, run.seed, run.status, run.training_time_seconds, run.final_loss, run.output_folder, _json(run.checkpoint_paths), _json(run.preview_images), run.peak_vram_gb, _json(run.hardware), _json(run.settings), _json(run.generation_settings), ), ) return run def list_runs(self, search: str = "", architecture: str = "", dataset: str = "", limit: int = 200) -> list[ExperimentRun]: clauses = [] params: list[Any] = [] if search: clauses.append("(id LIKE ? OR job_id LIKE ? OR model_name LIKE ? OR dataset_name LIKE ? OR dataset_path LIKE ? OR output_folder LIKE ? OR notes LIKE ? OR status LIKE ?)") term = f"%{search}%" params.extend([term, term, term, term, term, term, term, term]) if architecture: clauses.append("model_architecture = ?") params.append(architecture) if dataset: clauses.append("dataset_name LIKE ?") params.append(f"%{dataset}%") where = " WHERE " + " AND ".join(clauses) if clauses else "" with self.connect() as db: rows = db.execute( "SELECT * FROM experiments" + where + " ORDER BY timestamp DESC LIMIT ?", [*params, int(limit)], ).fetchall() return [ExperimentRun.from_row(row) for row in rows] def get(self, run_id: str) -> ExperimentRun | None: with self.connect() as db: row = db.execute("SELECT * FROM experiments WHERE id = ?", (run_id,)).fetchone() return ExperimentRun.from_row(row) if row else None def update_notes(self, run_id: str, notes: str, quality_score: int | None) -> None: with self.connect() as db: db.execute( "UPDATE experiments SET notes = ?, quality_score = ? WHERE id = ?", (notes, quality_score, run_id), ) def compare(self, run_ids: list[str]) -> list[dict[str, Any]]: runs = [run for run_id in run_ids if (run := self.get(run_id)) is not None] fields = [ "model_architecture", "epochs", "final_loss", "training_time_seconds", "resolution", "batch_size", "learning_rate", "scheduler", "peak_vram_gb", "dataset_name", "quality_score", ] rows = [] for field in fields: values = {run.id: getattr(run, field) for run in runs} comparable = {str(value) for value in values.values()} rows.append({"field": field, "changed": len(comparable) > 1, **values}) return rows def clone_request(self, run_id: str) -> str: run = self.get(run_id) if run is None: return "" options = { key: value for key, value in run.settings.items() if key not in {"dataset_dir", "model_name", "epochs", "output_dir", "resume_from"} } return ( f"From the {run.dataset_name or run.dataset_path} dataset, train a " f"{run.model_architecture.upper()} model for {run.epochs} epochs. " f"Name the model {run.model_name} Clone. " "[ADAM_TRAINING_OPTIONS:" + json.dumps(options, sort_keys=True) + "] " "[ADAM_TRAINER:" + run.model_architecture + "]" )