Spaces:
Running on Zero
Running on Zero
File size: 4,677 Bytes
aa3132a | 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 | 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()
|