adam/studio.py

from __future__ import annotations

import hashlib
import json
import shutil
from dataclasses import asdict, dataclass, field
from datetime import datetime, timezone
from pathlib import Path
from typing import Any
from uuid import uuid4


IMAGE_SUFFIXES = {".png", ".jpg", ".jpeg", ".webp", ".bmp", ".gif"}
CHECKPOINT_SUFFIXES = {".safetensors", ".ckpt", ".pt", ".pth"}


def _now() -> str:
    return datetime.now(timezone.utc).isoformat()


def _atomic_json(path: Path, payload: dict[str, Any]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    temporary = path.with_suffix(path.suffix + ".tmp")
    temporary.write_text(json.dumps(payload, indent=2), encoding="utf-8")
    temporary.replace(path)


def image_files(folder: str | Path, *, limit: int = 2500) -> list[Path]:
    root = Path(folder).expanduser()
    if not root.is_dir():
        return []
    return [
        path
        for path in sorted(root.rglob("*"))
        if path.is_file() and path.suffix.casefold() in IMAGE_SUFFIXES
    ][:limit]


def preview_files(folder: str | Path, *, limit: int = 500) -> list[Path]:
    return image_files(folder, limit=limit)


def checkpoint_files(folder: str | Path, *, limit: int = 500) -> list[Path]:
    root = Path(folder).expanduser()
    if not root.exists():
        return []
    if root.is_file():
        return [root] if root.suffix.casefold() in CHECKPOINT_SUFFIXES else []
    candidates = [
        path
        for path in root.rglob("*")
        if (
            path.is_file()
            and path.suffix.casefold() in CHECKPOINT_SUFFIXES
        )
        or (path.is_dir() and path.name.casefold().startswith("checkpoint-"))
    ]
    return sorted(candidates, key=lambda path: path.stat().st_mtime, reverse=True)[:limit]


def caption_path(image: Path) -> Path:
    return image.with_suffix(".txt")


def exact_duplicate_groups(paths: list[Path]) -> list[list[str]]:
    """Return exact duplicate groups without retaining image bytes in memory."""
    by_size: dict[int, list[Path]] = {}
    for path in paths:
        try:
            by_size.setdefault(path.stat().st_size, []).append(path)
        except OSError:
            continue
    groups: list[list[str]] = []
    for same_size in by_size.values():
        if len(same_size) < 2:
            continue
        hashes: dict[str, list[str]] = {}
        for path in same_size:
            try:
                digest = hashlib.sha256(path.read_bytes()).hexdigest()
            except OSError:
                continue
            hashes.setdefault(digest, []).append(str(path))
        groups.extend(values for values in hashes.values() if len(values) > 1)
    return groups


@dataclass(slots=True)
class DatasetReview:
    dataset_path: str
    decisions: dict[str, str] = field(default_factory=dict)
    notes: dict[str, str] = field(default_factory=dict)
    reviewed_at: str = ""


@dataclass(slots=True)
class TrainingRecipe:
    name: str
    trainer: str
    epochs: int
    image_count: int = 60
    base_model: str = ""
    preview_prompt: str = ""
    notes: str = ""
    id: str = field(default_factory=lambda: uuid4().hex[:10])
    created_at: str = field(default_factory=_now)

    @classmethod
    def from_dict(cls, payload: dict[str, Any]) -> "TrainingRecipe":
        return cls(
            id=str(payload.get("id") or uuid4().hex[:10]),
            name=str(payload.get("name", "Untitled recipe")),
            trainer=str(payload.get("trainer", "lora")),
            epochs=int(payload.get("epochs", 100) or 100),
            image_count=int(payload.get("image_count", 60) or 60),
            base_model=str(payload.get("base_model", "")),
            preview_prompt=str(payload.get("preview_prompt", "")),
            notes=str(payload.get("notes", "")),
            created_at=str(payload.get("created_at") or _now()),
        )


@dataclass(slots=True)
class PreviewEvaluation:
    model_id: str
    checkpoint: str
    prompt: str
    seed: int
    rating: int
    notes: str = ""
    id: str = field(default_factory=lambda: uuid4().hex[:10])
    created_at: str = field(default_factory=_now)

    @classmethod
    def from_dict(cls, payload: dict[str, Any]) -> "PreviewEvaluation":
        return cls(
            id=str(payload.get("id") or uuid4().hex[:10]),
            model_id=str(payload.get("model_id", "")),
            checkpoint=str(payload.get("checkpoint", "")),
            prompt=str(payload.get("prompt", "")),
            seed=int(payload.get("seed", 0) or 0),
            rating=max(0, min(5, int(payload.get("rating", 0) or 0))),
            notes=str(payload.get("notes", "")),
            created_at=str(payload.get("created_at") or _now()),
        )


class StudioStore:
    """Small, durable store for reviews, recipes, and model evaluations."""

    def __init__(self, root: Path) -> None:
        self.path = root.resolve() / "data" / "studio.json"
        self.dataset_reviews: dict[str, DatasetReview] = {}
        self.recipes: list[TrainingRecipe] = []
        self.evaluations: list[PreviewEvaluation] = []
        self.best_models: set[str] = set()
        self.load()

    def load(self) -> None:
        try:
            payload = json.loads(self.path.read_text(encoding="utf-8"))
        except (OSError, ValueError, TypeError, json.JSONDecodeError):
            payload = {}
        reviews = payload.get("dataset_reviews", {})
        if isinstance(reviews, dict):
            self.dataset_reviews = {
                key: DatasetReview(
                    dataset_path=str(value.get("dataset_path", key)),
                    decisions=dict(value.get("decisions", {})),
                    notes=dict(value.get("notes", {})),
                    reviewed_at=str(value.get("reviewed_at", "")),
                )
                for key, value in reviews.items()
                if isinstance(value, dict)
            }
        self.recipes = [
            TrainingRecipe.from_dict(item)
            for item in payload.get("recipes", [])
            if isinstance(item, dict)
        ]
        self.evaluations = [
            PreviewEvaluation.from_dict(item)
            for item in payload.get("evaluations", [])
            if isinstance(item, dict)
        ]
        self.best_models = {
            str(value) for value in payload.get("best_models", []) if value
        }

    def save(self) -> None:
        _atomic_json(
            self.path,
            {
                "dataset_reviews": {
                    key: asdict(value) for key, value in self.dataset_reviews.items()
                },
                "recipes": [asdict(value) for value in self.recipes],
                "evaluations": [asdict(value) for value in self.evaluations[-500:]],
                "best_models": sorted(self.best_models),
            },
        )

    def review(self, dataset_path: str) -> DatasetReview:
        key = str(Path(dataset_path).expanduser().resolve())
        if key not in self.dataset_reviews:
            self.dataset_reviews[key] = DatasetReview(key)
        return self.dataset_reviews[key]

    def set_decision(self, dataset_path: str, image_path: str, decision: str) -> None:
        review = self.review(dataset_path)
        if decision not in {"keep", "reject", "unreviewed"}:
            raise ValueError("Unknown review decision.")
        key = str(Path(image_path).expanduser().resolve())
        if decision == "unreviewed":
            review.decisions.pop(key, None)
        else:
            review.decisions[key] = decision
        review.reviewed_at = _now()
        self.save()

    def set_all_decisions(
        self, dataset_path: str, image_paths: list[str | Path], decision: str
    ) -> int:
        """Apply one review decision to every supplied image with a single save."""
        if decision not in {"keep", "reject", "unreviewed"}:
            raise ValueError("Unknown review decision.")
        review = self.review(dataset_path)
        changed = 0
        for image_path in image_paths:
            key = str(Path(image_path).expanduser().resolve())
            previous = review.decisions.get(key, "unreviewed")
            if previous == decision:
                continue
            if decision == "unreviewed":
                review.decisions.pop(key, None)
            else:
                review.decisions[key] = decision
            changed += 1
        if changed:
            review.reviewed_at = _now()
            self.save()
        return changed

    def apply_decisions(self, dataset_path: str, decisions: dict[str, str]) -> int:
        """Apply mixed Keep/Reject/Unreviewed decisions with one durable save."""
        review = self.review(dataset_path)
        changed = 0
        for image_path, decision in decisions.items():
            if decision not in {"keep", "reject", "unreviewed"}:
                raise ValueError("Unknown review decision.")
            key = str(Path(image_path).expanduser().resolve())
            previous = review.decisions.get(key, "unreviewed")
            if previous == decision:
                continue
            if decision == "unreviewed":
                review.decisions.pop(key, None)
            else:
                review.decisions[key] = decision
            changed += 1
        if changed:
            review.reviewed_at = _now()
            self.save()
        return changed

    def apply_rejections(self, dataset_path: str) -> int:
        """Move rejected images and captions to ADAM's recoverable quarantine."""
        dataset = Path(dataset_path).expanduser().resolve()
        review = self.review(str(dataset))
        key = hashlib.sha1(str(dataset).encode("utf-8")).hexdigest()[:12]
        quarantine = self.path.parent / "dataset_quarantine" / key
        quarantine.mkdir(parents=True, exist_ok=True)
        manifest_path = quarantine / "manifest.json"
        try:
            manifest = json.loads(manifest_path.read_text(encoding="utf-8"))
        except (OSError, ValueError, TypeError, json.JSONDecodeError):
            manifest = {"dataset_path": str(dataset), "files": []}
        moved = 0
        for raw_path, decision in list(review.decisions.items()):
            source = Path(raw_path)
            if decision != "reject" or not source.is_file():
                continue
            try:
                relative = source.resolve().relative_to(dataset)
            except ValueError:
                continue
            destinations = []
            for original in (source, caption_path(source)):
                if not original.is_file():
                    continue
                destination = quarantine / relative.parent / original.name
                destination.parent.mkdir(parents=True, exist_ok=True)
                if destination.exists():
                    destination = destination.with_name(
                        f"{destination.stem}_{uuid4().hex[:6]}{destination.suffix}"
                    )
                shutil.move(str(original), str(destination))
                destinations.append({"original": str(original), "quarantine": str(destination)})
            if destinations:
                manifest["files"].extend(destinations)
                review.decisions.pop(raw_path, None)
                moved += 1
        _atomic_json(manifest_path, manifest)
        review.reviewed_at = _now()
        self.save()
        return moved

    def restore_rejections(self, dataset_path: str) -> int:
        """Restore quarantined files to their original dataset when possible."""
        dataset = Path(dataset_path).expanduser().resolve()
        key = hashlib.sha1(str(dataset).encode("utf-8")).hexdigest()[:12]
        manifest_path = self.path.parent / "dataset_quarantine" / key / "manifest.json"
        try:
            manifest = json.loads(manifest_path.read_text(encoding="utf-8"))
        except (OSError, ValueError, TypeError, json.JSONDecodeError):
            return 0
        remaining = []
        restored_images = 0
        for entry in manifest.get("files", []):
            source = Path(str(entry.get("quarantine", "")))
            destination = Path(str(entry.get("original", "")))
            if not source.is_file() or destination.exists():
                remaining.append(entry)
                continue
            destination.parent.mkdir(parents=True, exist_ok=True)
            shutil.move(str(source), str(destination))
            restored_images += int(destination.suffix.casefold() in IMAGE_SUFFIXES)
        manifest["files"] = remaining
        _atomic_json(manifest_path, manifest)
        return restored_images

    def add_recipe(self, recipe: TrainingRecipe) -> TrainingRecipe:
        existing = next((item for item in self.recipes if item.id == recipe.id), None)
        if existing:
            self.recipes[self.recipes.index(existing)] = recipe
        else:
            self.recipes.insert(0, recipe)
        self.save()
        return recipe

    def add_evaluation(self, evaluation: PreviewEvaluation) -> None:
        self.evaluations.append(evaluation)
        self.save()

    def toggle_best(self, model_id: str) -> bool:
        if model_id in self.best_models:
            self.best_models.remove(model_id)
            selected = False
        else:
            self.best_models.add(model_id)
            selected = True
        self.save()
        return selected