diff --git a/image_conveyor_server.py b/image_conveyor_server.py index 877ecae..e45b41a 100644 --- a/image_conveyor_server.py +++ b/image_conveyor_server.py @@ -9,7 +9,8 @@ import stat as stat_module import tempfile import threading import time -from contextlib import contextmanager +import uuid +from contextlib import ExitStack, contextmanager from dataclasses import dataclass from pathlib import Path, PurePosixPath, PureWindowsPath from typing import Any, Dict, Iterable, List, Optional, Sequence, Tuple @@ -39,6 +40,10 @@ class InvalidThumbnail(ValueError): pass +class InvalidPreset(ValueError): + pass + + def normalize_relative_path(value: Any, *, allow_empty: bool = False) -> str: raw = str(value or "").strip().replace("\\", "/") if not raw: @@ -70,6 +75,239 @@ def normalize_filename(value: Any) -> str: return raw +def normalize_reference_slot(value: Any, input_root: str, *, must_exist: bool = True) -> Optional[Dict[str, str]]: + if value is None: + return None + if not isinstance(value, dict): + raise InvalidPreset("Reference slots must be image records or null.") + storage_type = str(value.get("type", "input")).strip().lower() or "input" + if storage_type != "input": + raise InvalidPreset("Reference presets may only contain ComfyUI input images.") + annotated = str(value.get("annotated", "")).strip() + suffix = " [input]" + if not annotated.endswith(suffix): + raise InvalidPreset("Reference images must use an annotated ComfyUI input path.") + try: + relative_path = normalize_relative_path(annotated[:-len(suffix)]) + except InvalidInputPath as exc: + raise InvalidPreset("A reference slot contains an invalid input path.") from exc + if Path(relative_path).suffix.lower() not in SUPPORTED_IMAGE_EXTENSIONS: + raise InvalidPreset("A reference slot has an unsupported image extension.") + try: + resolve_under_root(input_root, relative_path, must_exist=must_exist) + except (InvalidInputPath, FileNotFoundError) as exc: + raise InvalidPreset(f"Reference image '{annotated}' is outside input or missing.") from exc + relative = PurePosixPath(relative_path) + parent = "" if str(relative.parent) == "." else str(relative.parent) + return { + "annotated": f"{relative_path} [input]", + "filename": relative.name, + "subfolder": parent, + "type": "input", + } + + +class PresetStore: + VERSION = 1 + SLOT_COUNT = 8 + + def __init__(self, path: str, input_root: str): + self.path = path + self.input_root = os.path.realpath(input_root) + self._lock = threading.RLock() + self._recovery_pending = False + + @staticmethod + def _empty_document() -> Dict[str, Any]: + return {"version": PresetStore.VERSION, "presets": []} + + def _normalize_name(self, value: Any) -> str: + name = " ".join(str(value or "").strip().split()) + if not name or len(name) > 120 or any(ord(character) < 32 or ord(character) == 127 for character in name): + raise InvalidPreset("Preset names must contain 1 to 120 printable characters.") + return name + + def _normalize_id(self, value: Any) -> str: + try: + return str(uuid.UUID(str(value or "").strip())) + except (ValueError, AttributeError) as exc: + raise InvalidPreset("The preset ID is invalid.") from exc + + def _normalize_slots(self, value: Any, *, must_exist: bool = True) -> List[Optional[Dict[str, str]]]: + if not isinstance(value, list) or len(value) != self.SLOT_COUNT: + raise InvalidPreset("A reference preset must contain exactly 8 slots.") + return [ + normalize_reference_slot(slot, self.input_root, must_exist=must_exist) + for slot in value + ] + + def _normalize_preset(self, value: Any, *, must_exist: bool = False) -> Dict[str, Any]: + if not isinstance(value, dict): + raise InvalidPreset("The preset document is malformed.") + return { + "id": self._normalize_id(value.get("id")), + "name": self._normalize_name(value.get("name")), + "slots": self._normalize_slots(value.get("slots"), must_exist=must_exist), + "created_at": max(0, int(value.get("created_at", 0) or 0)), + "updated_at": max(0, int(value.get("updated_at", 0) or 0)), + } + + def _quarantine(self) -> None: + if not os.path.exists(self.path): + return + quarantine = f"{self.path}.corrupt-{int(time.time() * 1000)}" + try: + os.replace(self.path, quarantine) + except OSError: + LOGGER.exception("Image Conveyor could not quarantine malformed preset storage.") + + def _load_unlocked(self) -> Dict[str, Any]: + if not os.path.isfile(self.path): + return self._empty_document() + try: + with open(self.path, "r", encoding="utf-8") as handle: + raw = json.load(handle) + if not isinstance(raw, dict) or raw.get("version") != self.VERSION: + raise InvalidPreset("Unsupported preset document version.") + values = raw.get("presets") + if not isinstance(values, list): + raise InvalidPreset("The preset document is malformed.") + presets = [self._normalize_preset(value, must_exist=False) for value in values] + ids = {preset["id"] for preset in presets} + names = {preset["name"].casefold() for preset in presets} + if len(ids) != len(presets) or len(names) != len(presets): + raise InvalidPreset("The preset document contains duplicate IDs or names.") + return {"version": self.VERSION, "presets": presets} + except (OSError, UnicodeError, json.JSONDecodeError, InvalidPreset, ValueError, TypeError): + LOGGER.warning("Image Conveyor preset storage was malformed and has been quarantined.", exc_info=True) + self._quarantine() + self._recovery_pending = True + return self._empty_document() + + def _write_unlocked(self, document: Dict[str, Any]) -> None: + parent = os.path.dirname(self.path) + os.makedirs(parent, exist_ok=True) + descriptor, temporary = tempfile.mkstemp(prefix=".reference-presets-", suffix=".tmp", dir=parent) + try: + with os.fdopen(descriptor, "w", encoding="utf-8") as handle: + json.dump(document, handle, ensure_ascii=False, separators=(",", ":")) + handle.flush() + os.fsync(handle.fileno()) + os.replace(temporary, self.path) + self._recovery_pending = False + try: + directory_fd = os.open(parent, os.O_RDONLY) + try: + os.fsync(directory_fd) + finally: + os.close(directory_fd) + except OSError: + pass + finally: + try: + os.unlink(temporary) + except OSError: + pass + + @staticmethod + def _ordered(presets: Sequence[Dict[str, Any]]) -> List[Dict[str, Any]]: + return sorted(presets, key=lambda preset: (preset["name"].casefold(), preset["name"], preset["id"])) + + def list(self) -> List[Dict[str, Any]]: + with self._lock: + return json.loads(json.dumps(self._ordered(self._load_unlocked()["presets"]))) + + def create(self, name: Any, slots: Any) -> Dict[str, Any]: + with self._lock: + document = self._load_unlocked() + normalized_name = self._normalize_name(name) + if any(preset["name"].casefold() == normalized_name.casefold() for preset in document["presets"]): + raise InvalidPreset("A preset with that name already exists.") + now = int(time.time() * 1000) + preset = { + "id": str(uuid.uuid4()), + "name": normalized_name, + "slots": self._normalize_slots(slots), + "created_at": now, + "updated_at": now, + } + document["presets"].append(preset) + document["presets"] = self._ordered(document["presets"]) + self._write_unlocked(document) + return json.loads(json.dumps(preset)) + + def update(self, preset_id: Any, *, name: Any = None, slots: Any = None) -> Dict[str, Any]: + with self._lock: + document = self._load_unlocked() + normalized_id = self._normalize_id(preset_id) + preset = next((entry for entry in document["presets"] if entry["id"] == normalized_id), None) + if preset is None: + raise KeyError(normalized_id) + if name is not None: + normalized_name = self._normalize_name(name) + if any( + entry["id"] != normalized_id and entry["name"].casefold() == normalized_name.casefold() + for entry in document["presets"] + ): + raise InvalidPreset("A preset with that name already exists.") + preset["name"] = normalized_name + if slots is not None: + preset["slots"] = self._normalize_slots(slots) + if name is None and slots is None: + raise InvalidPreset("The preset update did not contain a name or slots.") + preset["updated_at"] = int(time.time() * 1000) + document["presets"] = self._ordered(document["presets"]) + self._write_unlocked(document) + return json.loads(json.dumps(preset)) + + def delete(self, preset_id: Any) -> bool: + with self._lock: + document = self._load_unlocked() + normalized_id = self._normalize_id(preset_id) + kept = [entry for entry in document["presets"] if entry["id"] != normalized_id] + if len(kept) == len(document["presets"]): + return False + document["presets"] = kept + self._write_unlocked(document) + return True + + def relink_paths(self, replacements: Sequence[Dict[str, Any]]) -> int: + mapping: Dict[str, str] = {} + for replacement in replacements: + if not isinstance(replacement, dict): + raise InvalidPreset("A preset relink entry is malformed.") + old_path = normalize_relative_path(replacement.get("relative_path")) + keep_path = normalize_relative_path(replacement.get("keep_path")) + resolve_under_root(self.input_root, keep_path, must_exist=True) + mapping[old_path] = keep_path + if not mapping: + return 0 + with self._lock: + document = self._load_unlocked() + if self._recovery_pending: + raise InvalidPreset( + "Preset storage was malformed; duplicate cleanup is blocked until presets are repaired or saved again." + ) + changed = 0 + for preset in document["presets"]: + for index, slot in enumerate(preset["slots"]): + if slot is None: + continue + relative_path = slot["annotated"][:-len(" [input]")] + keep_path = mapping.get(relative_path) + if not keep_path: + continue + preset["slots"][index] = normalize_reference_slot( + {"annotated": f"{keep_path} [input]", "type": "input"}, + self.input_root, + ) + preset["updated_at"] = int(time.time() * 1000) + changed += 1 + if changed: + self._write_unlocked(document) + return changed + + def resolve_under_root(root: str, relative_path: str, *, must_exist: bool = False) -> str: relative = normalize_relative_path(relative_path) root_real = os.path.realpath(root) @@ -290,13 +528,23 @@ class _KeyLockEntry: class InputLibrary: - def __init__(self, input_root: str, cache_root: str, snapshot_ttl: float = SNAPSHOT_TTL_SECONDS): + def __init__( + self, + input_root: str, + cache_root: str, + snapshot_ttl: float = SNAPSHOT_TTL_SECONDS, + preset_path: Optional[str] = None, + ): self.input_root = os.path.realpath(input_root) self.cache_root = cache_root self.snapshot_ttl = snapshot_ttl self.index = ContentIndex(os.path.join(cache_root, "content-index.sqlite3")) self.thumbnail_root = os.path.join(cache_root, "thumbnails") self.upload_temp_root = os.path.join(cache_root, "uploads") + self.preset_store = PresetStore( + preset_path or os.path.join(cache_root, "reference-presets.json"), + self.input_root, + ) self._snapshot: Tuple[FileRecord, ...] = () self._snapshot_at = 0.0 self._snapshot_wall_ms = 0 @@ -724,6 +972,7 @@ class InputLibrary: plan = [] planned_duplicates = 0 + planned_paths = set() for group in groups: if not isinstance(group, dict): raise InvalidUpload("The duplicate cleanup plan is malformed.") @@ -745,9 +994,14 @@ class InputLibrary: relative_path = normalize_relative_path(duplicate.get("relative_path")) if relative_path == keep_path or not self._is_managed_path(relative_path): raise InvalidUpload("Duplicate cleanup is restricted to the image_conveyor folder.") + if relative_path in planned_paths: + raise InvalidUpload("The duplicate cleanup plan contains a repeated path.") + planned_paths.add(relative_path) relative_paths.append(relative_path) plan.append((digest, keep_path, relative_paths)) + presets_relinked = 0 + for digest, keep_path, relative_paths in plan: with self._key_lock(self._digest_locks, digest): keep = self._record_with_digest(FileRecord(keep_path, -1, -1)) @@ -761,6 +1015,7 @@ class InputLibrary: ) continue + candidates = [] for relative_path in relative_paths: if relative_path in protected: skipped.append( @@ -770,7 +1025,17 @@ class InputLibrary: } ) continue - with self._key_lock(self._destination_locks, relative_path.casefold()): + candidates.append(relative_path) + + with ExitStack() as destination_locks: + lock_keys = sorted({relative_path.casefold() for relative_path in candidates}) + for lock_key in lock_keys: + destination_locks.enter_context( + self._key_lock(self._destination_locks, lock_key) + ) + + validated = [] + for relative_path in candidates: managed = self._managed_regular_file(relative_path) if managed is None or managed[1].content_hash != digest: skipped.append( @@ -780,7 +1045,21 @@ class InputLibrary: } ) continue - lexical_path, current, expected_identity = managed + validated.append(managed) + + # The preset document is durable. Its atomic update must + # succeed while every validated destination is locked and + # before the first unlink in this digest group. + presets_relinked += self.preset_store.relink_paths([ + { + "relative_path": current.relative_path, + "keep_path": keep_path, + } + for _lexical_path, current, _expected_identity in validated + ]) + + for lexical_path, current, expected_identity in validated: + relative_path = current.relative_path try: identity = os.lstat(lexical_path) if ( @@ -813,6 +1092,7 @@ class InputLibrary: "deleted": deleted, "skipped": skipped, "reclaimed_bytes": sum(entry["size"] for entry in deleted), + "presets_relinked": presets_relinked, } def verify_image(self, path: str) -> None: @@ -899,13 +1179,31 @@ def _cache_directory(folder_paths_module) -> str: return os.path.join(root, "image_conveyor") +def _preset_path(folder_paths_module) -> str: + return os.path.join( + folder_paths_module.get_user_directory(), + "image_conveyor", + "reference-presets.json", + ) + + def get_service(folder_paths_module) -> InputLibrary: global _SERVICE input_root = os.path.realpath(folder_paths_module.get_input_directory()) cache_root = _cache_directory(folder_paths_module) + preset_path = _preset_path(folder_paths_module) with _SERVICE_LOCK: - if _SERVICE is None or _SERVICE.input_root != input_root or _SERVICE.cache_root != cache_root: - _SERVICE = InputLibrary(input_root, cache_root) + if ( + _SERVICE is None + or _SERVICE.input_root != input_root + or _SERVICE.cache_root != cache_root + or _SERVICE.preset_store.path != preset_path + ): + _SERVICE = InputLibrary( + input_root, + cache_root, + preset_path=preset_path, + ) return _SERVICE @@ -1055,14 +1353,90 @@ def register_routes() -> None: raise InvalidUpload("The duplicate cleanup plan is malformed or too large.") if not isinstance(protected_paths, list) or len(protected_paths) > 10000: raise InvalidUpload("The duplicate cleanup reservation is malformed or too large.") - result = await asyncio.to_thread(service.delete_managed_duplicates, groups, protected_paths) + result = await asyncio.to_thread( + service.delete_managed_duplicates, + groups, + protected_paths, + ) return web.json_response(result) - except (InvalidUpload, InvalidInputPath, json.JSONDecodeError, ValueError) as exc: + except (InvalidUpload, InvalidInputPath, InvalidPreset, json.JSONDecodeError, ValueError) as exc: return web.json_response({"error": str(exc)}, status=400) except Exception: LOGGER.exception("Image Conveyor failed to delete managed duplicates.") return web.json_response({"error": "Unable to delete exact duplicates."}, status=500) + @routes.get("/image-conveyor/reference-presets") + async def image_conveyor_list_reference_presets(_request): + service = get_service(folder_paths) + try: + presets = await asyncio.to_thread(service.preset_store.list) + return web.json_response({"version": PresetStore.VERSION, "presets": presets}) + except Exception: + LOGGER.exception("Image Conveyor failed to list reference presets.") + return web.json_response({"error": "Unable to read reference presets."}, status=500) + + @routes.post("/image-conveyor/reference-presets") + async def image_conveyor_create_reference_preset(request): + service = get_service(folder_paths) + try: + payload = await request.json() + if not isinstance(payload, dict): + raise InvalidPreset("The preset request is malformed.") + preset = await asyncio.to_thread( + service.preset_store.create, + payload.get("name"), + payload.get("slots"), + ) + return web.json_response({"preset": preset}, status=201) + except (InvalidPreset, InvalidInputPath, json.JSONDecodeError, ValueError) as exc: + return web.json_response({"error": str(exc)}, status=400) + except Exception: + LOGGER.exception("Image Conveyor failed to create a reference preset.") + return web.json_response({"error": "Unable to create the reference preset."}, status=500) + + @routes.put("/image-conveyor/reference-presets/{preset_id}") + async def image_conveyor_update_reference_preset(request): + service = get_service(folder_paths) + try: + payload = await request.json() + if not isinstance(payload, dict): + raise InvalidPreset("The preset request is malformed.") + kwargs = {} + if "name" in payload: + kwargs["name"] = payload.get("name") + if "slots" in payload: + kwargs["slots"] = payload.get("slots") + preset = await asyncio.to_thread( + service.preset_store.update, + request.match_info["preset_id"], + **kwargs, + ) + return web.json_response({"preset": preset}) + except KeyError: + return web.json_response({"error": "Reference preset not found."}, status=404) + except (InvalidPreset, InvalidInputPath, json.JSONDecodeError, ValueError) as exc: + return web.json_response({"error": str(exc)}, status=400) + except Exception: + LOGGER.exception("Image Conveyor failed to update a reference preset.") + return web.json_response({"error": "Unable to update the reference preset."}, status=500) + + @routes.delete("/image-conveyor/reference-presets/{preset_id}") + async def image_conveyor_delete_reference_preset(request): + service = get_service(folder_paths) + try: + deleted = await asyncio.to_thread( + service.preset_store.delete, + request.match_info["preset_id"], + ) + if not deleted: + return web.json_response({"error": "Reference preset not found."}, status=404) + return web.json_response({"deleted": True}) + except (InvalidPreset, ValueError) as exc: + return web.json_response({"error": str(exc)}, status=400) + except Exception: + LOGGER.exception("Image Conveyor failed to delete a reference preset.") + return web.json_response({"error": "Unable to delete the reference preset."}, status=500) + @routes.get("/image-conveyor/thumbnail") async def image_conveyor_thumbnail(request): service = get_service(folder_paths)