invoke-ai--invokeai
wehub-resource-sync
cddb07a176
docs / deploy (push) Has been cancelled
docs / changes (push) Has been cancelled
docs / check-and-build (push) Has been cancelled
build container image / cpu (push) Has been cancelled
build container image / cuda (push) Has been cancelled
build container image / rocm (push) Has been cancelled
frontend checks / frontend-checks (push) Has been cancelled
frontend tests / frontend-tests (push) Has been cancelled
lfs checks / lfs-check (push) Has been cancelled
python checks / python-checks (push) Has been cancelled
python tests / py3.12: macos-default (push) Has been cancelled
python tests / py3.11: windows-cpu (push) Has been cancelled
python tests / py3.12: windows-cpu (push) Has been cancelled
python tests / py3.11: linux-cpu (push) Has been cancelled
typegen checks / typegen-checks (push) Has been cancelled
uv lock checks / uv-lock-checks (push) Has been cancelled
openapi checks / openapi-checks (push) Has been cancelled
python tests / py3.11: macos-default (push) Has been cancelled
python tests / py3.12: linux-cpu (push) Has been cancelled
84 行
2.7 KiB
Python
84 行
2.7 KiB
Python
from pathlib import Path
|
|
|
|
from PIL import Image
|
|
from PIL.Image import Image as PILImageType
|
|
|
|
from invokeai.app.services.invoker import Invoker
|
|
from invokeai.app.services.model_images.model_images_base import ModelImageFileStorageBase
|
|
from invokeai.app.services.model_images.model_images_common import (
|
|
ModelImageFileDeleteException,
|
|
ModelImageFileNotFoundException,
|
|
ModelImageFileSaveException,
|
|
)
|
|
from invokeai.app.util.misc import uuid_string
|
|
from invokeai.app.util.thumbnails import make_thumbnail
|
|
|
|
|
|
class ModelImageFileStorageDisk(ModelImageFileStorageBase):
|
|
"""Stores images on disk"""
|
|
|
|
def __init__(self, model_images_folder: Path):
|
|
self._model_images_folder = model_images_folder
|
|
self._validate_storage_folders()
|
|
|
|
def start(self, invoker: Invoker) -> None:
|
|
self._invoker = invoker
|
|
|
|
def get(self, model_key: str) -> PILImageType:
|
|
try:
|
|
path = self.get_path(model_key)
|
|
|
|
if not self._validate_path(path):
|
|
raise ModelImageFileNotFoundException
|
|
|
|
return Image.open(path)
|
|
except FileNotFoundError as e:
|
|
raise ModelImageFileNotFoundException from e
|
|
|
|
def save(self, image: PILImageType, model_key: str) -> None:
|
|
try:
|
|
self._validate_storage_folders()
|
|
image_path = self._model_images_folder / (model_key + ".webp")
|
|
thumbnail = make_thumbnail(image, 256)
|
|
thumbnail.save(image_path, format="webp")
|
|
|
|
except Exception as e:
|
|
raise ModelImageFileSaveException from e
|
|
|
|
def get_path(self, model_key: str) -> Path:
|
|
path = self._model_images_folder / (model_key + ".webp")
|
|
|
|
return path
|
|
|
|
def get_url(self, model_key: str) -> str | None:
|
|
path = self.get_path(model_key)
|
|
if not self._validate_path(path):
|
|
return
|
|
|
|
url = self._invoker.services.urls.get_model_image_url(model_key)
|
|
|
|
# The image URL never changes, so we must add random query string to it to prevent caching
|
|
url += f"?{uuid_string()}"
|
|
|
|
return url
|
|
|
|
def delete(self, model_key: str) -> None:
|
|
try:
|
|
path = self.get_path(model_key)
|
|
|
|
if not self._validate_path(path):
|
|
raise ModelImageFileNotFoundException
|
|
|
|
path.unlink()
|
|
|
|
except Exception as e:
|
|
raise ModelImageFileDeleteException from e
|
|
|
|
def _validate_path(self, path: Path) -> bool:
|
|
"""Validates the path given for an image."""
|
|
return path.exists()
|
|
|
|
def _validate_storage_folders(self) -> None:
|
|
"""Checks if the required folders exist and create them if they don't"""
|
|
self._model_images_folder.mkdir(parents=True, exist_ok=True)
|