import errno import io import json import os import shutil import sqlite3 import tempfile import threading import zlib from collections.abc import Collection, Iterable, Iterator, Sequence from dataclasses import dataclass from pathlib import Path from queue import Queue from typing import Optional, Union from PIL import Image, PngImagePlugin from PIL.Image import Image as PILImageType from invokeai.app.services.image_files.image_files_base import ImageFileStorageBase from invokeai.app.services.image_files.image_files_common import ( ImageFileDeleteException, ImageFileNotFoundException, ImageFileSaveException, ) from invokeai.app.services.invoker import Invoker from invokeai.app.util.thumbnails import get_thumbnail_name, make_thumbnail from invokeai.backend.util.logging import InvokeAILogger _PNG_RLE_MIN_PIXELS = 512 * 512 _PNG_RLE_SAMPLE_TILE = 32 _PNG_RLE_MIN_RAW_SIZE_PERCENT = 30 _PNG_RLE_MAX_SAMPLE_SIZE_PERCENT = 102 @dataclass class _StagedDelete: """Files moved aside by ``stage_delete()``, restorable until the token is committed.""" directory: Path files: list[tuple[Path, Path]] image_name: str image_subfolder: str @dataclass class _PendingDelete: """A durable record of an intent to purge files, written before the records are deleted. Nothing is moved: the journal directory names the images whose files are about to become unreferenced. Startup recovery reconciles any journal that outlives its operation by asking the record store which of its images are really gone. """ directory: Path images: list[tuple[str, str]] def _get_png_size(image: PILImageType, compress_type: Optional[int] = None) -> int: output = io.BytesIO() options = {"compress_level": 1} if compress_type is not None: options["compress_type"] = compress_type image.save(output, "PNG", **options) return output.tell() def _should_use_png_rle(image: PILImageType) -> bool: if image.mode not in {"RGB", "RGBA"} or image.width * image.height < _PNG_RLE_MIN_PIXELS: return False # Native-resolution tiles distinguish high-entropy data from filter-friendly structured images. tile_width = min(_PNG_RLE_SAMPLE_TILE, image.width) tile_height = min(_PNG_RLE_SAMPLE_TILE, image.height) x_positions = (0, (image.width - tile_width) // 2, image.width - tile_width) y_positions = (0, (image.height - tile_height) // 2, image.height - tile_height) sample = Image.new(image.mode, (tile_width * 3, tile_height * 3)) for row, y in enumerate(y_positions): for column, x in enumerate(x_positions): with image.crop((x, y, x + tile_width, y + tile_height)) as tile: sample.paste(tile, (column * tile_width, row * tile_height)) try: raw = sample.tobytes() if len(zlib.compress(raw, level=1)) * 100 < len(raw) * _PNG_RLE_MIN_RAW_SIZE_PERCENT: return False default_size = _get_png_size(sample) rle_size = _get_png_size(sample, zlib.Z_RLE) return rle_size * 100 <= default_size * _PNG_RLE_MAX_SAMPLE_SIZE_PERCENT finally: sample.close() class DiskImageFileStorage(ImageFileStorageBase): """Stores images on disk""" def __init__(self, output_folder: Union[str, Path]): self.__cache: dict[Path, PILImageType] = {} self.__cache_ids = Queue[Path]() self.__max_cache_size = 10 # TODO: get this from config # Guards the cache structures (__cache / __cache_ids), which are read and mutated from # multiple session-processor worker threads in multi-GPU parallel mode. self.__cache_lock = threading.Lock() self.__cache_generation = 0 self.__output_folder = output_folder if isinstance(output_folder, Path) else Path(output_folder) self.__thumbnails_folder = self.__output_folder / "thumbnails" # Validate required output folders at launch self.__validate_storage_folders() def start(self, invoker: Invoker) -> None: self.__invoker = invoker self.__recover_pending_deletes() @property def image_root(self) -> Path: return self.__output_folder.resolve() @property def thumbnail_root(self) -> Path: return self.__thumbnails_folder.resolve() def iter_image_paths(self) -> Iterator[Path]: """Walk supported image files without following symlinks or entering journals/thumbnails.""" yield from self.__iter_supported_paths(self.image_root, suffixes={".png"}, skip_thumbnail_root=True) def iter_thumbnail_paths(self) -> Iterator[Path]: """Walk WEBP thumbnails without following symlinks or entering reserved journals.""" yield from self.__iter_supported_paths(self.thumbnail_root, suffixes={".webp"}) @staticmethod def __iter_supported_paths(root: Path, suffixes: set[str], skip_thumbnail_root: bool = False) -> Iterator[Path]: # Inventory and traversal state stay in a temporary SQLite database. A generated gallery # can contain enough files that retaining DirEntry objects or directory paths in Python # would make maintenance memory usage grow with the output tree. with tempfile.TemporaryDirectory(prefix="invokeai-image-path-inventory-") as inventory_dir: inventory_path = Path(inventory_dir) / "inventory.sqlite" connection = sqlite3.connect(inventory_path) try: connection.execute("PRAGMA cache_size = -2048") connection.execute("PRAGMA temp_store = FILE") connection.execute("PRAGMA journal_mode = OFF") connection.execute("PRAGMA synchronous = OFF") connection.execute("CREATE TABLE pending_directories (id INTEGER PRIMARY KEY, path BLOB NOT NULL)") connection.execute( "CREATE TABLE directory_entries (path BLOB NOT NULL, sort_key BLOB NOT NULL, is_directory INTEGER NOT NULL)" ) connection.execute("INSERT INTO pending_directories (path) VALUES (?)", (os.fsencode(root),)) connection.commit() while True: pending = connection.execute( "SELECT id, path FROM pending_directories ORDER BY id DESC LIMIT 1" ).fetchone() if pending is None: break directory_id, directory_value = pending connection.execute("DELETE FROM pending_directories WHERE id = ?", (directory_id,)) connection.execute("DELETE FROM directory_entries") directory = Path(os.fsdecode(directory_value)) with os.scandir(directory) as scan: for entry in scan: if entry.name.startswith(".delete_"): continue if directory == root and skip_thumbnail_root and entry.name == "thumbnails": continue if entry.is_symlink(): continue if entry.is_dir(follow_symlinks=False): is_directory = 1 elif entry.is_file(follow_symlinks=False) and Path(entry.name).suffix.lower() in suffixes: is_directory = 0 else: continue entry_path = os.fsencode(directory / entry.name) sort_key = entry.name.encode("utf-32-be", "surrogatepass") connection.execute( "INSERT INTO directory_entries (path, sort_key, is_directory) VALUES (?, ?, ?)", (entry_path, sort_key, is_directory), ) connection.commit() entries = connection.execute("SELECT path, is_directory FROM directory_entries ORDER BY sort_key") for entry_path, is_directory in entries: if not is_directory: yield Path(os.fsdecode(entry_path)) # Push children in reverse lexical order so the disk-backed LIFO stack visits # them in the same ascending order as the previous in-memory traversal. connection.execute( """--sql INSERT INTO pending_directories (path) SELECT path FROM directory_entries WHERE is_directory = 1 ORDER BY sort_key DESC """ ) connection.commit() finally: connection.close() def evict_cache_paths(self, paths: list[Path]) -> None: resolved_paths = [path.resolve() for path in paths] with self.__cache_lock: self.__invalidate_cache_paths_locked(resolved_paths) def __invalidate_cache_paths_locked(self, paths: list[Path]) -> None: if not paths: return self.__cache_generation += 1 for path in paths: self.__cache.pop(path, None) def get(self, image_name: str, image_subfolder: str = "") -> PILImageType: try: image_path = self.get_path(image_name, image_subfolder=image_subfolder) cache_item, cache_generation = self.__get_cache(image_path) if cache_item: return cache_item image = Image.open(image_path) # Image.open() is lazy: it reads the header but defers pixel decoding (and holds the # file handle open) until the first .load()/.copy()/.convert(). The opened object is # cached and the SAME object is handed to every caller, so in multi-GPU parallel mode # two worker threads can call .copy() on it concurrently and race on the shared file # handle and decoder state, producing "broken data stream" / "self.png is not None" # errors. Forcing the decode here makes the cached object safe for concurrent reads. image.load() self.__set_cache(image_path, image, cache_generation) return image except FileNotFoundError as e: raise ImageFileNotFoundException from e def save( self, image: PILImageType, image_name: str, metadata: Optional[str] = None, workflow: Optional[str] = None, graph: Optional[str] = None, thumbnail_size: int = 256, image_subfolder: str = "", ) -> None: image_path: Optional[Path] = None thumbnail_path: Optional[Path] = None image_existed = False thumbnail_existed = False try: self.__validate_storage_folders() image_path = self.get_path(image_name, image_subfolder=image_subfolder) image_existed = image_path.exists() # Ensure subfolder directories exist image_path.parent.mkdir(parents=True, exist_ok=True) pnginfo = PngImagePlugin.PngInfo() info_dict = {} if metadata is not None: info_dict["invokeai_metadata"] = metadata pnginfo.add_text("invokeai_metadata", metadata) if workflow is not None: info_dict["invokeai_workflow"] = workflow pnginfo.add_text("invokeai_workflow", workflow) if graph is not None: info_dict["invokeai_graph"] = graph pnginfo.add_text("invokeai_graph", graph) thumbnail_path = self.get_path(image_name, thumbnail=True, image_subfolder=image_subfolder) thumbnail_existed = thumbnail_path.exists() # Build the thumbnail before replacing image.info with Invoke metadata. PIL stores # palette transparency there, and it must remain available to make_thumbnail(). thumbnail_path.parent.mkdir(parents=True, exist_ok=True) thumbnail_image = make_thumbnail(image, thumbnail_size) # When saving the image, the image object's info field is not populated. We need to set it image.info = info_dict compress_level = self.__invoker.services.configuration.pil_compress_level save_options = {"compress_level": compress_level} if compress_level == 1 and _should_use_png_rle(image): save_options["compress_type"] = zlib.Z_RLE image.save( image_path, "PNG", pnginfo=pnginfo, **save_options, ) thumbnail_image.save(thumbnail_path) self.__set_cache(image_path, image) self.__set_cache(thumbnail_path, thumbnail_image) except Exception as e: # A thumbnail failure must not leave a full-size image with no thumbnail. The # names are normally new, but preserve any pre-existing files when save() is # used to overwrite an existing image. for path, existed in ((image_path, image_existed), (thumbnail_path, thumbnail_existed)): if path is not None and not existed: try: path.unlink(missing_ok=True) except OSError: pass self.evict_cache_paths([path for path in (image_path, thumbnail_path) if path is not None]) raise ImageFileSaveException from e def generate_thumbnail_if_missing( self, image_name: str, image_subfolder: str = "", thumbnail_size: int = 256 ) -> bool: """Create a complete thumbnail at the canonical path without touching the source image.""" image_path = self.get_path(image_name, image_subfolder=image_subfolder) thumbnail_path = self.get_path(image_name, thumbnail=True, image_subfolder=image_subfolder) try: thumbnail_path.stat() return False except FileNotFoundError: pass thumbnail_path.parent.mkdir(parents=True, exist_ok=True) temporary_path: Path | None = None thumbnail: PILImageType | None = None published_identity: tuple[int, int] | None = None try: with Image.open(image_path) as source_image: thumbnail = make_thumbnail(source_image, thumbnail_size) descriptor, temporary_name = tempfile.mkstemp( prefix=".thumbnail-", suffix=".webp", dir=thumbnail_path.parent ) os.close(descriptor) temporary_path = Path(temporary_name) thumbnail.save(temporary_path, format="WEBP") with open(temporary_path, "rb+") as completed_file: os.fsync(completed_file.fileno()) temporary_stat = temporary_path.stat() temporary_identity = (temporary_stat.st_dev, temporary_stat.st_ino) try: # Linking publishes the fully written temporary file atomically and refuses to # replace a thumbnail another writer created after our first existence check. os.link(temporary_path, thumbnail_path) except FileExistsError: return False published_identity = temporary_identity try: self.__fsync_directory(thumbnail_path.parent) except Exception: try: current_stat = thumbnail_path.lstat() if (current_stat.st_dev, current_stat.st_ino) == published_identity: thumbnail_path.unlink() self.evict_cache_paths([thumbnail_path]) try: self.__fsync_directory(thumbnail_path.parent) except OSError: pass except OSError: pass raise return True finally: if thumbnail is not None: thumbnail.close() if temporary_path is not None: temporary_path.unlink(missing_ok=True) def copy( self, source_image_name: str, image_name: str, source_subfolder: str = "", image_subfolder: str = "", thumbnail_size: int = 256, ) -> None: """Duplicate an image's files under a new name, without decoding them. `shutil.copy2` is both cheapest and most faithful: every PNG chunk travels, including ones this application does not parse and would drop on a re-encode. Nothing is cached — a copy nobody has read yet has no business occupying a slot. """ try: self.__validate_storage_folders() source_path = self.get_path(source_image_name, image_subfolder=source_subfolder) image_path = self.get_path(image_name, image_subfolder=image_subfolder) image_path.parent.mkdir(parents=True, exist_ok=True) shutil.copy2(source_path, image_path) source_thumbnail_path = self.get_path(source_image_name, thumbnail=True, image_subfolder=source_subfolder) thumbnail_path = self.get_path(image_name, thumbnail=True, image_subfolder=image_subfolder) thumbnail_path.parent.mkdir(parents=True, exist_ok=True) if source_thumbnail_path.exists(): shutil.copy2(source_thumbnail_path, thumbnail_path) else: # A source predating thumbnails, or one whose thumbnail was lost. Rare enough to be # worth a decode rather than leaving the copy with a broken gallery tile. with Image.open(image_path) as copied: make_thumbnail(copied, thumbnail_size).save(thumbnail_path) except Exception as e: raise ImageFileSaveException from e def delete(self, image_name: str, image_subfolder: str = "") -> None: token = self.stage_delete(image_name, image_subfolder) self.commit_delete(token) def stage_delete(self, image_name: str, image_subfolder: str = "") -> _StagedDelete: candidates = self.__delete_candidates(image_name, image_subfolder) staging_dir = Path(tempfile.mkdtemp(prefix=".delete_", dir=self.__output_folder)) staged: list[tuple[Path, Path]] = [] try: with open(staging_dir / "manifest.json", "w", encoding="utf-8") as manifest: manifest.write(json.dumps({"image_name": image_name, "image_subfolder": image_subfolder})) manifest.flush() os.fsync(manifest.fileno()) # The manifest has to be durable before the files move, or a crash can leave staged # files in a directory that names nothing and recovery cannot put them back. self.__persist_journal_directory(staging_dir) with self.__cache_lock: try: for index, source in enumerate(candidates): if source.exists(): destination = staging_dir / str(index) source.replace(destination) staged.append((source, destination)) finally: self.__invalidate_cache_paths_locked(candidates) return _StagedDelete( directory=staging_dir, files=staged, image_name=image_name, image_subfolder=image_subfolder ) except Exception as e: for source, destination in reversed(staged): if destination.exists(): source.parent.mkdir(parents=True, exist_ok=True) destination.replace(source) shutil.rmtree(staging_dir, ignore_errors=True) raise ImageFileDeleteException from e def begin_delete(self, images: Sequence[tuple[str, str]]) -> _PendingDelete: """Durably records the intent to purge these images' files, before their records are deleted. Callers delete the records first and purge afterwards, so the only inconsistency that can outlive a crash or a storage failure is a file nobody references. The journal written here is what makes that recoverable: ``__recover_pending_deletes()`` asks the record store about every image it names, purges the ones whose record is gone, and leaves the rest untouched. """ # Resolve every path up front. A name that cannot be turned into a path must fail here, # while the caller can still abort — not after it has deleted the records. for image_name, image_subfolder in images: self.__delete_candidates(image_name, image_subfolder) journal_dir = Path(tempfile.mkdtemp(prefix=".delete_", dir=self.__output_folder)) try: entries = [ {"image_name": image_name, "image_subfolder": image_subfolder} for image_name, image_subfolder in images ] manifest_path = journal_dir / "manifest.json" with open(manifest_path, "w", encoding="utf-8") as manifest: manifest.write(json.dumps({"version": 2, "images": entries})) manifest.flush() os.fsync(manifest.fileno()) self.__persist_journal_directory(journal_dir) return _PendingDelete(directory=journal_dir, images=[(name, subfolder) for name, subfolder in images]) except Exception as e: shutil.rmtree(journal_dir, ignore_errors=True) raise ImageFileDeleteException from e def abandon_delete(self, token: object) -> None: """Drops a pending-delete journal without purging anything. Used when the record deletion failed: the images are still live, so their files must stay. """ if not isinstance(token, _PendingDelete): raise ImageFileDeleteException("Invalid pending-delete token") shutil.rmtree(token.directory, ignore_errors=True) def commit_delete(self, token: object, image_names: Optional[Collection[str]] = None) -> None: if isinstance(token, _PendingDelete): self.__commit_pending_delete(token, image_names) return if not isinstance(token, _StagedDelete): raise ImageFileDeleteException("Invalid staged-delete token") try: # Purge the live paths as well as the staged copies. stage_delete() captures whatever # was there at that instant, so a second deleter racing the first gets an empty token — # and if that second one is the one whose record deletion succeeds, dropping its empty # staging directory alone would strand the files the first deleter restores. Committing # has to mean "no file for this image survives", whichever request moved them. self.__purge_files(token.image_name, token.image_subfolder) self.__persist_purges([(token.image_name, token.image_subfolder)]) shutil.rmtree(token.directory) except Exception as e: raise ImageFileDeleteException from e def __commit_pending_delete(self, token: _PendingDelete, image_names: Optional[Collection[str]]) -> None: # ``image_names`` narrows the purge to the records that were actually deleted; the journal # still lists every candidate, which is harmless because recovery re-checks each one against # the record store and skips any that survived. selected = image_names if image_names is None else set(image_names) purged: list[tuple[str, str]] = [] failures: list[str] = [] for image_name, image_subfolder in token.images: if selected is not None and image_name not in selected: continue try: self.__purge_files(image_name, image_subfolder) purged.append((image_name, image_subfolder)) except OSError as e: failures.append(f"{image_name}: {e}") if failures: # Leave the journal in place so the next startup retries every entry whose record is # gone. Removing it here would turn a transient storage error into a permanent orphan. raise ImageFileDeleteException(f"Failed to purge deleted image files: {'; '.join(failures)}") try: self.__persist_purges(purged) except OSError as e: # Same rule: the journal is the only thing that can redo a purge the disk lost. raise ImageFileDeleteException(f"Failed to make the purge of deleted image files durable: {e}") from e shutil.rmtree(token.directory, ignore_errors=True) def rollback_delete(self, token: object) -> None: if not isinstance(token, _StagedDelete): raise ImageFileDeleteException("Invalid staged-delete token") try: for source, destination in reversed(token.files): if destination.exists(): source.parent.mkdir(parents=True, exist_ok=True) destination.replace(source) # While these files sat in the staging directory another request may have deleted the # record; restoring them would leave files nothing references and no journal to find # them by. Re-check now that the files are back: every deleter purges an image's files # only *after* its record is committed as gone, so a record still present here cannot # have been purged before this restore, and a record already absent means the purge # either found nothing or is still to come — either way the files must go. self.__purge_if_record_absent(token.image_name, token.image_subfolder) self.__persist_purges([(token.image_name, token.image_subfolder)]) shutil.rmtree(token.directory, ignore_errors=True) except Exception as e: raise ImageFileDeleteException from e def __delete_candidates(self, image_name: str, image_subfolder: str) -> list[Path]: return [ self.get_path(image_name, image_subfolder=image_subfolder), self.get_path(image_name, thumbnail=True, image_subfolder=image_subfolder), ] def __purge_files(self, image_name: str, image_subfolder: str) -> None: """Removes an image's file and thumbnail. Missing files are not an error.""" for path in self.__delete_candidates(image_name, image_subfolder): with self.__cache_lock: try: path.unlink(missing_ok=True) finally: self.__invalidate_cache_paths_locked([path]) def __purge_if_record_absent(self, image_name: str, image_subfolder: str) -> None: try: record_exists = self.__invoker.services.image_records.exists(image_name) except Exception as e: # A storage fault must never destroy a live image's files. Keep them: a stale file is # recoverable at the next startup, a deleted one is not. InvokeAILogger.get_logger().error(f"Could not confirm whether {image_name} still exists: {e}") return if record_exists: return self.__purge_files(image_name, image_subfolder) def __persist_journal_directory(self, journal_dir: Path) -> None: """Makes a journal directory and its manifest survive a power loss. Both fsyncs are needed: the first commits ``manifest.json``'s entry inside the journal directory, the second commits the journal directory's own entry in the output folder. Without the second, the record deletion — which SQLite does fsync — can outlive the journal that is supposed to make it recoverable. A sync that fails propagates for the same reason: a journal whose durability is unknown must not license a record deletion. """ self.__fsync_directory(journal_dir) self.__fsync_directory(self.__output_folder) def __persist_purges(self, images: Iterable[tuple[str, str]]) -> None: """Makes the removal (or restoration) of these images' files survive a power loss. Runs before the journal that names them is dropped. Unlinks and renames live in the parent directories' entries, and nothing else syncs those; a filesystem that persists the journal's removal ahead of them would bring the files back with no journal left to find them by. Raises ``OSError`` when a sync fails, so the caller keeps the journal. """ parents: dict[Path, None] = {} for image_name, image_subfolder in images: for path in self.__delete_candidates(image_name, image_subfolder): parents[path.parent] = None for parent in parents: # A parent that no longer exists has no entries to make durable. self.__fsync_directory(parent, missing_ok=True) @staticmethod def __fsync_directory(directory: Path, missing_ok: bool = False) -> None: if os.name == "nt": # Windows cannot open a directory for fsync; the manifest write is all we get. return try: dir_fd = os.open(directory, os.O_RDONLY) except (FileNotFoundError, NotADirectoryError): if missing_ok: return raise try: os.fsync(dir_fd) except OSError as e: # Some filesystems cannot sync a directory at all and say so with one of these; there is # nothing more to be had from them. Anything else (EIO, ENOSPC, ...) means the entries # may not be on disk, and the caller must not proceed as if they were. if e.errno not in (errno.EINVAL, errno.ENOTSUP, errno.EOPNOTSUPP, errno.ENOSYS, errno.EBADF): raise finally: os.close(dir_fd) def get_file_size_bytes(self, image_name: str, image_subfolder: str = "") -> Optional[int]: try: size = self.get_path(image_name, image_subfolder=image_subfolder).stat().st_size except FileNotFoundError: return None try: size += self.get_path(image_name, thumbnail=True, image_subfolder=image_subfolder).stat().st_size except FileNotFoundError: pass return size def get_path(self, image_name: str, thumbnail: bool = False, image_subfolder: str = "") -> Path: base_folder = self.__thumbnails_folder if thumbnail else self.__output_folder filename = get_thumbnail_name(image_name) if thumbnail else image_name # Validate the filename itself (no path separators allowed in the filename) basename = Path(filename).name if basename != filename: raise ValueError("Invalid image name, potential directory traversal detected") # Build the full path with optional subfolder if image_subfolder: self._validate_subfolder(image_subfolder) image_path = base_folder / image_subfolder / basename else: image_path = base_folder / basename # Ensure the image path is within the base folder to prevent directory traversal resolved_base = base_folder.resolve() resolved_image_path = image_path.resolve() if not resolved_image_path.is_relative_to(resolved_base): raise ValueError("Image path outside outputs folder, potential directory traversal detected") return resolved_image_path @staticmethod def _validate_subfolder(subfolder: str) -> None: """Validates a subfolder path to prevent directory traversal while allowing controlled subdirectories.""" if not subfolder: return if "\\" in subfolder: raise ValueError("Backslashes not allowed in subfolder path") if subfolder.startswith("/"): raise ValueError("Absolute paths not allowed in subfolder path") parts = subfolder.split("/") for part in parts: if part == "..": raise ValueError("Parent directory references not allowed in subfolder path") if part == "": raise ValueError("Empty path segments not allowed in subfolder path") def validate_path(self, path: Union[str, Path]) -> bool: """Validates the path given for an image or thumbnail.""" path = path if isinstance(path, Path) else Path(path) return path.exists() def get_workflow(self, image_name: str, image_subfolder: str = "") -> str | None: image = self.get(image_name, image_subfolder=image_subfolder) workflow = image.info.get("invokeai_workflow", None) if isinstance(workflow, str): return workflow return None def get_graph(self, image_name: str, image_subfolder: str = "") -> str | None: image = self.get(image_name, image_subfolder=image_subfolder) graph = image.info.get("invokeai_graph", None) if isinstance(graph, str): return graph return None def __validate_storage_folders(self) -> None: """Checks if the required output folders exist and create them if they don't""" folders: list[Path] = [self.__output_folder, self.__thumbnails_folder] for folder in folders: folder.mkdir(parents=True, exist_ok=True) def __recover_pending_deletes(self) -> None: """Reconciles every delete journal left behind by an interrupted or failed deletion. One rule covers both journal shapes, and the record store decides it: an image whose record survives was never really deleted, so anything staged for it is put back and the journal dropped; an image whose record is gone is an orphan, so its files are purged wherever the interrupted operation left them. """ logger = InvokeAILogger.get_logger() for journal_dir in sorted(self.__output_folder.glob(".delete_*")): try: manifest_path = journal_dir / "manifest.json" if not manifest_path.is_file(): # mkdtemp() ran but the manifest never landed, so this directory names nothing # and cannot be reconciled. Drop it when it is empty; otherwise say so, because # anything inside it is a staged file that can no longer be put back. if not journal_dir.is_dir(): logger.warning(f"Ignoring unexpected entry in the outputs folder: {journal_dir}") elif not any(journal_dir.iterdir()): journal_dir.rmdir() else: logger.warning(f"Image deletion journal {journal_dir} has no manifest and cannot be recovered") continue with open(manifest_path, encoding="utf-8") as manifest: data = json.load(manifest) records = self.__invoker.services.image_records images = self.__manifest_images(data) for image_name, image_subfolder in images: if records.exists(image_name): # Put back whatever stage_delete() moved aside. Only a single-image journal # ever holds staged files, at indices 0 and 1; a pending-delete journal # moves nothing, so these lookups simply find nothing to restore. restored = False for index, source in enumerate(self.__delete_candidates(image_name, image_subfolder)): staged = journal_dir / str(index) if staged.exists(): source.parent.mkdir(parents=True, exist_ok=True) staged.replace(source) restored = True # Re-check after a restore, for the same reason rollback_delete() does: # another Invoke sharing this output folder may have deleted the record # while the files sat staged, and its own purge found nothing to remove. # Every deleter purges only after its record is committed as gone, so a # record still present cannot have been purged before this restore, and a # record now absent means the files must go — this journal is the last # thing that can find them. A record-store fault here propagates and keeps # the journal, like every other lookup in this loop. if not restored or records.exists(image_name): continue self.__purge_files(image_name, image_subfolder) # Every restore and purge above must be on disk before the journal that could redo # them is gone. A failed sync propagates and keeps the journal for the next startup. self.__persist_purges(images) shutil.rmtree(journal_dir, ignore_errors=True) except Exception as error: # Includes a record-store fault: leave the journal for the next startup rather than # guess. Retrying is always safe; both branches above are idempotent. logger.error(f"Failed to recover image deletion journal {journal_dir}: {error}") @staticmethod def __manifest_images(data: dict) -> list[tuple[str, str]]: """Reads both journal shapes: a pending delete lists many images, a staged delete names one.""" entries = data.get("images") if entries is None: return [(data["image_name"], data.get("image_subfolder", ""))] return [(entry["image_name"], entry.get("image_subfolder", "")) for entry in entries] def __get_cache(self, image_name: Path) -> tuple[Optional[PILImageType], int]: with self.__cache_lock: cache_item = None if image_name not in self.__cache else self.__cache[image_name] return cache_item, self.__cache_generation def __set_cache(self, image_name: Path, image: PILImageType, generation: int | None = None) -> None: with self.__cache_lock: if generation is not None and generation != self.__cache_generation: return if image_name not in self.__cache: self.__cache[image_name] = image self.__cache_ids.put(image_name) # TODO: this should refresh position for LRU cache if len(self.__cache) > self.__max_cache_size: cache_id = self.__cache_ids.get() if cache_id in self.__cache: del self.__cache[cache_id]