lora_cache.py

from __future__ import annotations

import hashlib
import os
import threading
import time
from collections import OrderedDict
from dataclasses import dataclass
from pathlib import Path

from safetensors import safe_open

MAX_FILE_BYTES = 500 * 1024 * 1024
MAX_CACHE_BYTES = 2 * 1024 * 1024 * 1024


@dataclass(frozen=True)
class CachedLora:
    name: str
    path: Path
    sha256: str
    size: int


class LoraCache:
    def __init__(self, root: Path, capacity: int = MAX_CACHE_BYTES):
        self.root = root
        self.capacity = capacity
        self.root.mkdir(parents=True, exist_ok=True)
        self._entries: OrderedDict[str, CachedLora] = OrderedDict()
        self._uploads_by_identity: dict[tuple[str, int], CachedLora] = {}
        self._lock = threading.RLock()

    @staticmethod
    def _digest(path: Path) -> tuple[str, int]:
        digest = hashlib.sha256()
        size = 0
        with path.open("rb") as stream:
            for chunk in iter(lambda: stream.read(4 * 1024 * 1024), b""):
                size += len(chunk)
                if size > MAX_FILE_BYTES:
                    raise ValueError("Each LoRA must be 500 MiB or smaller.")
                digest.update(chunk)
        return digest.hexdigest(), size

    @staticmethod
    def _validate_safetensors(path: Path) -> None:
        if path.suffix.lower() != ".safetensors":
            raise ValueError("Only .safetensors LoRA files are accepted.")
        try:
            with safe_open(path, framework="numpy") as handle:
                keys = [key.lower() for key in handle.keys()]
        except Exception as exc:
            raise ValueError("The uploaded file is not a valid safetensors model.") from exc
        if not keys:
            raise ValueError("The safetensors file contains no tensors.")
        lora_markers = ("lora", "lycoris", "hada_", "lokr_", "oft_")
        if not any(any(marker in key for marker in lora_markers) for key in keys):
            raise ValueError("The safetensors file does not contain recognizable LoRA tensors.")

    def add_with_status(self, upload: str | os.PathLike[str]) -> tuple[CachedLora, bool]:
        source = Path(upload)
        digest, size = self._digest(source)
        self._validate_safetensors(source)
        with self._lock:
            if digest in self._entries:
                entry = self._entries.pop(digest)
                self._entries[digest] = entry
                self._uploads_by_identity[(source.name, size)] = entry
                return entry, True
            target = self.root / f"user-{digest}.safetensors"
            temp = target.with_suffix(f".{time.time_ns()}.tmp")
            try:
                os.link(source, temp)
            except OSError:
                import shutil
                shutil.copyfile(source, temp)
            temp.replace(target)
            entry = CachedLora(target.name, target, digest, size)
            self._entries[digest] = entry
            self._uploads_by_identity[(source.name, size)] = entry
            self._evict()
            return entry, False

    def prepare_with_status(self, upload: str | os.PathLike[str]) -> tuple[CachedLora, bool]:
        source = Path(upload)
        if source.suffix.lower() != ".safetensors":
            raise ValueError("Only .safetensors LoRA files are accepted.")
        size = source.stat().st_size
        if size > MAX_FILE_BYTES:
            raise ValueError("Each LoRA must be 500 MiB or smaller.")
        identity = (source.name, size)
        with self._lock:
            entry = self._uploads_by_identity.get(identity)
            if entry is not None and entry.path.exists() and entry.sha256 in self._entries:
                self._entries.move_to_end(entry.sha256)
                return entry, True
        return self.add_with_status(source)

    def add(self, upload: str | os.PathLike[str]) -> CachedLora:
        entry, _already_present = self.add_with_status(upload)
        return entry

    def _evict(self) -> None:
        total = sum(item.size for item in self._entries.values())
        while total > self.capacity and self._entries:
            _, item = self._entries.popitem(last=False)
            total -= item.size
            item.path.unlink(missing_ok=True)
            self._uploads_by_identity = {
                identity: entry for identity, entry in self._uploads_by_identity.items()
                if entry.sha256 != item.sha256
            }

    def clear(self) -> None:
        with self._lock:
            for item in self._entries.values():
                item.path.unlink(missing_ok=True)
            self._entries.clear()
            self._uploads_by_identity.clear()