chore: import upstream snapshot with attribution
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
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
This commit is contained in:
@@ -0,0 +1,5 @@
|
||||
"""App settings service exports."""
|
||||
|
||||
from invokeai.app.services.app_settings.app_settings_service import AppSettingsService
|
||||
|
||||
__all__ = ["AppSettingsService"]
|
||||
@@ -0,0 +1,74 @@
|
||||
"""Service for managing application-level settings stored in the database."""
|
||||
|
||||
from typing import Optional
|
||||
|
||||
from invokeai.app.services.shared.sqlite.sqlite_database import SqliteDatabase
|
||||
|
||||
|
||||
class AppSettingsService:
|
||||
"""Service for accessing application-level settings from the database.
|
||||
|
||||
This service provides a simple key-value store for application-level configuration
|
||||
that needs to be persisted across restarts, such as JWT secrets.
|
||||
"""
|
||||
|
||||
def __init__(self, db: SqliteDatabase) -> None:
|
||||
"""Initialize the app settings service.
|
||||
|
||||
Args:
|
||||
db: The SQLite database instance
|
||||
"""
|
||||
self._db = db
|
||||
|
||||
def get(self, key: str) -> Optional[str]:
|
||||
"""Get a setting value by key.
|
||||
|
||||
Args:
|
||||
key: The setting key
|
||||
|
||||
Returns:
|
||||
The setting value if found, None otherwise
|
||||
"""
|
||||
try:
|
||||
with self._db.transaction() as cursor:
|
||||
cursor.execute("SELECT value FROM app_settings WHERE key = ?;", (key,))
|
||||
row = cursor.fetchone()
|
||||
return row[0] if row else None
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
def set(self, key: str, value: str) -> None:
|
||||
"""Set a setting value.
|
||||
|
||||
Args:
|
||||
key: The setting key
|
||||
value: The setting value
|
||||
"""
|
||||
with self._db.transaction() as cursor:
|
||||
cursor.execute(
|
||||
"""
|
||||
INSERT INTO app_settings (key, value)
|
||||
VALUES (?, ?)
|
||||
ON CONFLICT(key) DO UPDATE SET
|
||||
value = excluded.value,
|
||||
updated_at = STRFTIME('%Y-%m-%d %H:%M:%f', 'NOW');
|
||||
""",
|
||||
(key, value),
|
||||
)
|
||||
|
||||
def get_jwt_secret(self) -> str:
|
||||
"""Get the JWT secret key from the database.
|
||||
|
||||
Returns:
|
||||
The JWT secret key
|
||||
|
||||
Raises:
|
||||
RuntimeError: If the JWT secret is not found in the database
|
||||
"""
|
||||
secret = self.get("jwt_secret")
|
||||
if secret is None:
|
||||
raise RuntimeError(
|
||||
"JWT secret not found in database. This should have been created during database migration. "
|
||||
"Please ensure database migrations have been run successfully."
|
||||
)
|
||||
return secret
|
||||
@@ -0,0 +1 @@
|
||||
"""Authentication service module."""
|
||||
@@ -0,0 +1,113 @@
|
||||
"""Password hashing and validation utilities."""
|
||||
|
||||
from typing import Literal, cast
|
||||
|
||||
from passlib.context import CryptContext
|
||||
|
||||
# Configure bcrypt context - set truncate_error=False to allow passwords >72 bytes
|
||||
# without raising an error. They will be automatically truncated by bcrypt to 72 bytes.
|
||||
pwd_context = CryptContext(
|
||||
schemes=["bcrypt"],
|
||||
deprecated="auto",
|
||||
bcrypt__truncate_error=False,
|
||||
)
|
||||
|
||||
|
||||
def hash_password(password: str) -> str:
|
||||
"""Hash a password using bcrypt.
|
||||
|
||||
bcrypt has a maximum password length of 72 bytes. Longer passwords
|
||||
are automatically truncated to comply with this limit.
|
||||
|
||||
Args:
|
||||
password: The plain text password to hash
|
||||
|
||||
Returns:
|
||||
The hashed password
|
||||
"""
|
||||
# bcrypt has a 72 byte limit - encode and truncate if necessary
|
||||
password_bytes = password.encode("utf-8")
|
||||
if len(password_bytes) > 72:
|
||||
# Truncate to 72 bytes and decode back, dropping incomplete UTF-8 sequences
|
||||
password = password_bytes[:72].decode("utf-8", errors="ignore")
|
||||
return cast(str, pwd_context.hash(password))
|
||||
|
||||
|
||||
def verify_password(plain_password: str, hashed_password: str) -> bool:
|
||||
"""Verify a password against a hash.
|
||||
|
||||
bcrypt has a maximum password length of 72 bytes. Longer passwords
|
||||
are automatically truncated to match hash_password behavior.
|
||||
|
||||
Args:
|
||||
plain_password: The plain text password to verify
|
||||
hashed_password: The hashed password to verify against
|
||||
|
||||
Returns:
|
||||
True if the password matches the hash, False otherwise
|
||||
"""
|
||||
try:
|
||||
# bcrypt has a 72 byte limit - encode and truncate if necessary to match hash_password
|
||||
password_bytes = plain_password.encode("utf-8")
|
||||
if len(password_bytes) > 72:
|
||||
# Truncate to 72 bytes and decode back, dropping incomplete UTF-8 sequences
|
||||
plain_password = password_bytes[:72].decode("utf-8", errors="ignore")
|
||||
return cast(bool, pwd_context.verify(plain_password, hashed_password))
|
||||
except Exception:
|
||||
# Invalid hash format or other error - return False
|
||||
return False
|
||||
|
||||
|
||||
def validate_password_strength(password: str) -> tuple[bool, str]:
|
||||
"""Validate password meets minimum security requirements.
|
||||
|
||||
Password requirements:
|
||||
- At least 8 characters long
|
||||
- Contains at least one uppercase letter
|
||||
- Contains at least one lowercase letter
|
||||
- Contains at least one digit
|
||||
|
||||
Args:
|
||||
password: The password to validate
|
||||
|
||||
Returns:
|
||||
A tuple of (is_valid, error_message). If valid, error_message is empty.
|
||||
"""
|
||||
if len(password) < 8:
|
||||
return False, "Password must be at least 8 characters long"
|
||||
|
||||
has_upper = any(c.isupper() for c in password)
|
||||
has_lower = any(c.islower() for c in password)
|
||||
has_digit = any(c.isdigit() for c in password)
|
||||
|
||||
if not (has_upper and has_lower and has_digit):
|
||||
return False, "Password must contain uppercase, lowercase, and numbers"
|
||||
|
||||
return True, ""
|
||||
|
||||
|
||||
def get_password_strength(password: str) -> Literal["weak", "moderate", "strong"]:
|
||||
"""Determine the strength of a password.
|
||||
|
||||
Strength levels:
|
||||
- weak: less than 8 characters
|
||||
- moderate: 8+ characters but missing at least one of uppercase, lowercase, or digit
|
||||
- strong: 8+ characters with uppercase, lowercase, and digit
|
||||
|
||||
Args:
|
||||
password: The password to evaluate
|
||||
|
||||
Returns:
|
||||
One of "weak", "moderate", or "strong"
|
||||
"""
|
||||
if len(password) < 8:
|
||||
return "weak"
|
||||
|
||||
has_upper = any(c.isupper() for c in password)
|
||||
has_lower = any(c.islower() for c in password)
|
||||
has_digit = any(c.isdigit() for c in password)
|
||||
|
||||
if not (has_upper and has_lower and has_digit):
|
||||
return "moderate"
|
||||
|
||||
return "strong"
|
||||
@@ -0,0 +1,106 @@
|
||||
"""JWT token generation and validation."""
|
||||
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from typing import cast
|
||||
|
||||
from jose import JWTError, jwt
|
||||
from pydantic import BaseModel
|
||||
|
||||
ALGORITHM = "HS256"
|
||||
DEFAULT_EXPIRATION_HOURS = 24
|
||||
|
||||
# Module-level variable to store the JWT secret. This is set during application initialization
|
||||
# by calling set_jwt_secret(). The secret is loaded from the database where it is stored
|
||||
# securely after being generated during database migration.
|
||||
_jwt_secret: str | None = None
|
||||
|
||||
|
||||
class TokenData(BaseModel):
|
||||
"""Data stored in JWT token."""
|
||||
|
||||
user_id: str
|
||||
email: str
|
||||
is_admin: bool
|
||||
remember_me: bool = False
|
||||
|
||||
|
||||
def set_jwt_secret(secret: str) -> None:
|
||||
"""Set the JWT secret key for token signing and verification.
|
||||
|
||||
This should be called once during application initialization with the secret
|
||||
loaded from the database.
|
||||
|
||||
Args:
|
||||
secret: The JWT secret key
|
||||
"""
|
||||
global _jwt_secret
|
||||
_jwt_secret = secret
|
||||
|
||||
|
||||
def get_jwt_secret() -> str:
|
||||
"""Get the JWT secret key.
|
||||
|
||||
Returns:
|
||||
The JWT secret key
|
||||
|
||||
Raises:
|
||||
RuntimeError: If the secret has not been initialized
|
||||
"""
|
||||
if _jwt_secret is None:
|
||||
raise RuntimeError("JWT secret has not been initialized. Call set_jwt_secret() during application startup.")
|
||||
return _jwt_secret
|
||||
|
||||
|
||||
def create_access_token(data: TokenData, expires_delta: timedelta | None = None) -> str:
|
||||
"""Create a JWT access token.
|
||||
|
||||
Args:
|
||||
data: The token data to encode
|
||||
expires_delta: Optional expiration time delta. Defaults to 24 hours.
|
||||
|
||||
Returns:
|
||||
The encoded JWT token
|
||||
"""
|
||||
to_encode = data.model_dump()
|
||||
expire = datetime.now(timezone.utc) + (expires_delta or timedelta(hours=DEFAULT_EXPIRATION_HOURS))
|
||||
to_encode.update({"exp": expire})
|
||||
return cast(str, jwt.encode(to_encode, get_jwt_secret(), algorithm=ALGORITHM))
|
||||
|
||||
|
||||
def verify_token(token: str) -> TokenData | None:
|
||||
"""Verify and decode a JWT token.
|
||||
|
||||
Args:
|
||||
token: The JWT token to verify
|
||||
|
||||
Returns:
|
||||
TokenData if valid, None if invalid or expired
|
||||
"""
|
||||
try:
|
||||
# python-jose 3.5.0 has a bug where exp verification doesn't work properly
|
||||
# We need to manually check expiration, but MUST verify signature first
|
||||
# to prevent accepting tokens with valid payloads but invalid signatures
|
||||
|
||||
# First, verify the signature - this will raise JWTError if signature is invalid
|
||||
# Note: python-jose won't reject expired tokens here due to the bug
|
||||
payload = jwt.decode(
|
||||
token,
|
||||
get_jwt_secret(),
|
||||
algorithms=[ALGORITHM],
|
||||
)
|
||||
|
||||
# Now manually check expiration (because python-jose 3.5.0 doesn't do this properly)
|
||||
if "exp" in payload:
|
||||
exp_timestamp = payload["exp"]
|
||||
current_timestamp = datetime.now(timezone.utc).timestamp()
|
||||
if current_timestamp >= exp_timestamp:
|
||||
# Token is expired
|
||||
return None
|
||||
|
||||
return TokenData(**payload)
|
||||
except JWTError:
|
||||
# Token is invalid (bad signature, malformed, etc.)
|
||||
return None
|
||||
except Exception:
|
||||
# Catch any other exceptions (e.g., Pydantic validation errors)
|
||||
return None
|
||||
@@ -0,0 +1,59 @@
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import Optional
|
||||
|
||||
from invokeai.app.services.image_records.image_records_common import ImageCategory
|
||||
|
||||
|
||||
class BoardImageRecordStorageBase(ABC):
|
||||
"""Abstract base class for the one-to-many board-image relationship record storage."""
|
||||
|
||||
@abstractmethod
|
||||
def add_image_to_board(
|
||||
self,
|
||||
board_id: str,
|
||||
image_name: str,
|
||||
) -> None:
|
||||
"""Adds an image to a board."""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def remove_image_from_board(
|
||||
self,
|
||||
image_name: str,
|
||||
) -> None:
|
||||
"""Removes an image from a board."""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def get_all_board_image_names_for_board(
|
||||
self,
|
||||
board_id: str,
|
||||
categories: list[ImageCategory] | None,
|
||||
is_intermediate: bool | None,
|
||||
) -> list[str]:
|
||||
"""Gets all board images for a board, as a list of the image names."""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def get_board_for_image(
|
||||
self,
|
||||
image_name: str,
|
||||
) -> Optional[str]:
|
||||
"""Gets an image's board id, if it has one."""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def get_image_count_for_board(
|
||||
self,
|
||||
board_id: str,
|
||||
) -> int:
|
||||
"""Gets the number of images for a board."""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def get_asset_count_for_board(
|
||||
self,
|
||||
board_id: str,
|
||||
) -> int:
|
||||
"""Gets the number of assets for a board."""
|
||||
pass
|
||||
@@ -0,0 +1,190 @@
|
||||
import sqlite3
|
||||
from typing import Optional, cast
|
||||
|
||||
from invokeai.app.services.board_image_records.board_image_records_base import BoardImageRecordStorageBase
|
||||
from invokeai.app.services.image_records.image_records_common import (
|
||||
ASSETS_CATEGORIES,
|
||||
IMAGE_CATEGORIES,
|
||||
ImageCategory,
|
||||
ImageRecord,
|
||||
deserialize_image_record,
|
||||
)
|
||||
from invokeai.app.services.shared.pagination import OffsetPaginatedResults
|
||||
from invokeai.app.services.shared.sqlite.sqlite_database import SqliteDatabase
|
||||
|
||||
|
||||
class SqliteBoardImageRecordStorage(BoardImageRecordStorageBase):
|
||||
def __init__(self, db: SqliteDatabase) -> None:
|
||||
super().__init__()
|
||||
self._db = db
|
||||
|
||||
def add_image_to_board(
|
||||
self,
|
||||
board_id: str,
|
||||
image_name: str,
|
||||
) -> None:
|
||||
with self._db.transaction() as cursor:
|
||||
cursor.execute(
|
||||
"""--sql
|
||||
INSERT INTO board_images (board_id, image_name)
|
||||
VALUES (?, ?)
|
||||
ON CONFLICT (image_name) DO UPDATE SET board_id = ?;
|
||||
""",
|
||||
(board_id, image_name, board_id),
|
||||
)
|
||||
|
||||
def remove_image_from_board(
|
||||
self,
|
||||
image_name: str,
|
||||
) -> None:
|
||||
with self._db.transaction() as cursor:
|
||||
cursor.execute(
|
||||
"""--sql
|
||||
DELETE FROM board_images
|
||||
WHERE image_name = ?;
|
||||
""",
|
||||
(image_name,),
|
||||
)
|
||||
|
||||
def get_images_for_board(
|
||||
self,
|
||||
board_id: str,
|
||||
offset: int = 0,
|
||||
limit: int = 10,
|
||||
) -> OffsetPaginatedResults[ImageRecord]:
|
||||
with self._db.transaction() as cursor:
|
||||
cursor.execute(
|
||||
"""--sql
|
||||
SELECT images.*
|
||||
FROM board_images
|
||||
INNER JOIN images ON board_images.image_name = images.image_name
|
||||
WHERE board_images.board_id = ?
|
||||
ORDER BY board_images.updated_at DESC;
|
||||
""",
|
||||
(board_id,),
|
||||
)
|
||||
result = cast(list[sqlite3.Row], cursor.fetchall())
|
||||
images = [deserialize_image_record(dict(r)) for r in result]
|
||||
|
||||
cursor.execute(
|
||||
"""--sql
|
||||
SELECT COUNT(*) FROM images WHERE 1=1;
|
||||
"""
|
||||
)
|
||||
count = cast(int, cursor.fetchone()[0])
|
||||
|
||||
return OffsetPaginatedResults(items=images, offset=offset, limit=limit, total=count)
|
||||
|
||||
def get_all_board_image_names_for_board(
|
||||
self,
|
||||
board_id: str,
|
||||
categories: list[ImageCategory] | None,
|
||||
is_intermediate: bool | None,
|
||||
) -> list[str]:
|
||||
with self._db.transaction() as cursor:
|
||||
params: list[str | bool] = []
|
||||
|
||||
# Base query is a join between images and board_images
|
||||
stmt = """
|
||||
SELECT images.image_name
|
||||
FROM images
|
||||
LEFT JOIN board_images ON board_images.image_name = images.image_name
|
||||
WHERE 1=1
|
||||
"""
|
||||
|
||||
# Handle board_id filter
|
||||
if board_id == "none":
|
||||
stmt += """--sql
|
||||
AND board_images.board_id IS NULL
|
||||
"""
|
||||
else:
|
||||
stmt += """--sql
|
||||
AND board_images.board_id = ?
|
||||
"""
|
||||
params.append(board_id)
|
||||
|
||||
# Add the category filter
|
||||
if categories is not None:
|
||||
# Convert the enum values to unique list of strings
|
||||
category_strings = [c.value for c in set(categories)]
|
||||
# Create the correct length of placeholders
|
||||
placeholders = ",".join("?" * len(category_strings))
|
||||
stmt += f"""--sql
|
||||
AND images.image_category IN ( {placeholders} )
|
||||
"""
|
||||
|
||||
# Unpack the included categories into the query params
|
||||
for c in category_strings:
|
||||
params.append(c)
|
||||
|
||||
# Add the is_intermediate filter
|
||||
if is_intermediate is not None:
|
||||
stmt += """--sql
|
||||
AND images.is_intermediate = ?
|
||||
"""
|
||||
params.append(is_intermediate)
|
||||
|
||||
# Put a ring on it
|
||||
stmt += ";"
|
||||
|
||||
cursor.execute(stmt, params)
|
||||
|
||||
result = cast(list[sqlite3.Row], cursor.fetchall())
|
||||
image_names = [r[0] for r in result]
|
||||
return image_names
|
||||
|
||||
def get_board_for_image(
|
||||
self,
|
||||
image_name: str,
|
||||
) -> Optional[str]:
|
||||
with self._db.transaction() as cursor:
|
||||
cursor.execute(
|
||||
"""--sql
|
||||
SELECT board_id
|
||||
FROM board_images
|
||||
WHERE image_name = ?;
|
||||
""",
|
||||
(image_name,),
|
||||
)
|
||||
result = cursor.fetchone()
|
||||
if result is None:
|
||||
return None
|
||||
return cast(str, result[0])
|
||||
|
||||
def get_image_count_for_board(self, board_id: str) -> int:
|
||||
with self._db.transaction() as cursor:
|
||||
# Convert the enum values to unique list of strings
|
||||
category_strings = [c.value for c in set(IMAGE_CATEGORIES)]
|
||||
# Create the correct length of placeholders
|
||||
placeholders = ",".join("?" * len(category_strings))
|
||||
cursor.execute(
|
||||
f"""--sql
|
||||
SELECT COUNT(*)
|
||||
FROM board_images
|
||||
INNER JOIN images ON board_images.image_name = images.image_name
|
||||
WHERE images.is_intermediate = FALSE AND images.image_category IN ( {placeholders} )
|
||||
AND board_images.board_id = ?;
|
||||
""",
|
||||
(*category_strings, board_id),
|
||||
)
|
||||
count = cast(int, cursor.fetchone()[0])
|
||||
return count
|
||||
|
||||
def get_asset_count_for_board(self, board_id: str) -> int:
|
||||
with self._db.transaction() as cursor:
|
||||
# Convert the enum values to unique list of strings
|
||||
category_strings = [c.value for c in set(ASSETS_CATEGORIES)]
|
||||
# Create the correct length of placeholders
|
||||
placeholders = ",".join("?" * len(category_strings))
|
||||
cursor.execute(
|
||||
f"""--sql
|
||||
SELECT COUNT(*)
|
||||
FROM board_images
|
||||
INNER JOIN images ON board_images.image_name = images.image_name
|
||||
WHERE images.is_intermediate = FALSE AND images.image_category IN ( {placeholders} )
|
||||
AND board_images.board_id = ?;
|
||||
""",
|
||||
(*category_strings, board_id),
|
||||
)
|
||||
count = cast(int, cursor.fetchone()[0])
|
||||
return count
|
||||
@@ -0,0 +1,43 @@
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import Optional
|
||||
|
||||
from invokeai.app.services.image_records.image_records_common import ImageCategory
|
||||
|
||||
|
||||
class BoardImagesServiceABC(ABC):
|
||||
"""High-level service for board-image relationship management."""
|
||||
|
||||
@abstractmethod
|
||||
def add_image_to_board(
|
||||
self,
|
||||
board_id: str,
|
||||
image_name: str,
|
||||
) -> None:
|
||||
"""Adds an image to a board."""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def remove_image_from_board(
|
||||
self,
|
||||
image_name: str,
|
||||
) -> None:
|
||||
"""Removes an image from a board."""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def get_all_board_image_names_for_board(
|
||||
self,
|
||||
board_id: str,
|
||||
categories: list[ImageCategory] | None,
|
||||
is_intermediate: bool | None,
|
||||
) -> list[str]:
|
||||
"""Gets all board images for a board, as a list of the image names."""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def get_board_for_image(
|
||||
self,
|
||||
image_name: str,
|
||||
) -> Optional[str]:
|
||||
"""Gets an image's board id, if it has one."""
|
||||
pass
|
||||
@@ -0,0 +1,8 @@
|
||||
from pydantic import Field
|
||||
|
||||
from invokeai.app.util.model_exclude_null import BaseModelExcludeNull
|
||||
|
||||
|
||||
class BoardImage(BaseModelExcludeNull):
|
||||
board_id: str = Field(description="The id of the board")
|
||||
image_name: str = Field(description="The name of the image")
|
||||
@@ -0,0 +1,44 @@
|
||||
from typing import Optional
|
||||
|
||||
from invokeai.app.services.board_images.board_images_base import BoardImagesServiceABC
|
||||
from invokeai.app.services.image_records.image_records_common import ImageCategory
|
||||
from invokeai.app.services.invoker import Invoker
|
||||
|
||||
|
||||
class BoardImagesService(BoardImagesServiceABC):
|
||||
__invoker: Invoker
|
||||
|
||||
def start(self, invoker: Invoker) -> None:
|
||||
self.__invoker = invoker
|
||||
|
||||
def add_image_to_board(
|
||||
self,
|
||||
board_id: str,
|
||||
image_name: str,
|
||||
) -> None:
|
||||
self.__invoker.services.board_image_records.add_image_to_board(board_id, image_name)
|
||||
|
||||
def remove_image_from_board(
|
||||
self,
|
||||
image_name: str,
|
||||
) -> None:
|
||||
self.__invoker.services.board_image_records.remove_image_from_board(image_name)
|
||||
|
||||
def get_all_board_image_names_for_board(
|
||||
self,
|
||||
board_id: str,
|
||||
categories: list[ImageCategory] | None,
|
||||
is_intermediate: bool | None,
|
||||
) -> list[str]:
|
||||
return self.__invoker.services.board_image_records.get_all_board_image_names_for_board(
|
||||
board_id,
|
||||
categories,
|
||||
is_intermediate,
|
||||
)
|
||||
|
||||
def get_board_for_image(
|
||||
self,
|
||||
image_name: str,
|
||||
) -> Optional[str]:
|
||||
board_id = self.__invoker.services.board_image_records.get_board_for_image(image_name)
|
||||
return board_id
|
||||
@@ -0,0 +1,66 @@
|
||||
from abc import ABC, abstractmethod
|
||||
|
||||
from invokeai.app.services.board_records.board_records_common import BoardChanges, BoardRecord, BoardRecordOrderBy
|
||||
from invokeai.app.services.shared.pagination import OffsetPaginatedResults
|
||||
from invokeai.app.services.shared.sqlite.sqlite_common import SQLiteDirection
|
||||
|
||||
|
||||
class BoardRecordStorageBase(ABC):
|
||||
"""Low-level service responsible for interfacing with the board record store."""
|
||||
|
||||
@abstractmethod
|
||||
def delete(self, board_id: str) -> None:
|
||||
"""Deletes a board record."""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def save(
|
||||
self,
|
||||
board_name: str,
|
||||
user_id: str,
|
||||
) -> BoardRecord:
|
||||
"""Saves a board record for a specific user."""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def get(
|
||||
self,
|
||||
board_id: str,
|
||||
) -> BoardRecord:
|
||||
"""Gets a board record."""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def update(
|
||||
self,
|
||||
board_id: str,
|
||||
changes: BoardChanges,
|
||||
) -> BoardRecord:
|
||||
"""Updates a board record."""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def get_many(
|
||||
self,
|
||||
user_id: str,
|
||||
is_admin: bool,
|
||||
order_by: BoardRecordOrderBy,
|
||||
direction: SQLiteDirection,
|
||||
offset: int = 0,
|
||||
limit: int = 10,
|
||||
include_archived: bool = False,
|
||||
) -> OffsetPaginatedResults[BoardRecord]:
|
||||
"""Gets many board records for a specific user, including shared boards. Admin users see all boards."""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def get_all(
|
||||
self,
|
||||
user_id: str,
|
||||
is_admin: bool,
|
||||
order_by: BoardRecordOrderBy,
|
||||
direction: SQLiteDirection,
|
||||
include_archived: bool = False,
|
||||
) -> list[BoardRecord]:
|
||||
"""Gets all board records for a specific user, including shared boards. Admin users see all boards."""
|
||||
pass
|
||||
@@ -0,0 +1,113 @@
|
||||
from datetime import datetime
|
||||
from enum import Enum
|
||||
from typing import Optional, Union
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
from invokeai.app.util.metaenum import MetaEnum
|
||||
from invokeai.app.util.misc import get_iso_timestamp
|
||||
from invokeai.app.util.model_exclude_null import BaseModelExcludeNull
|
||||
|
||||
|
||||
class BoardVisibility(str, Enum, metaclass=MetaEnum):
|
||||
"""The visibility options for a board."""
|
||||
|
||||
Private = "private"
|
||||
"""Only the board owner (and admins) can see and modify this board."""
|
||||
Shared = "shared"
|
||||
"""All users can view this board, but only the owner (and admins) can modify it."""
|
||||
Public = "public"
|
||||
"""All users can view this board; only the owner (and admins) can modify its structure."""
|
||||
|
||||
|
||||
class BoardRecord(BaseModelExcludeNull):
|
||||
"""Deserialized board record."""
|
||||
|
||||
board_id: str = Field(description="The unique ID of the board.")
|
||||
"""The unique ID of the board."""
|
||||
board_name: str = Field(description="The name of the board.")
|
||||
"""The name of the board."""
|
||||
user_id: str = Field(description="The user ID of the board owner.")
|
||||
"""The user ID of the board owner."""
|
||||
created_at: Union[datetime, str] = Field(description="The created timestamp of the board.")
|
||||
"""The created timestamp of the image."""
|
||||
updated_at: Union[datetime, str] = Field(description="The updated timestamp of the board.")
|
||||
"""The updated timestamp of the image."""
|
||||
deleted_at: Optional[Union[datetime, str]] = Field(default=None, description="The deleted timestamp of the board.")
|
||||
"""The updated timestamp of the image."""
|
||||
cover_image_name: Optional[str] = Field(default=None, description="The name of the cover image of the board.")
|
||||
"""The name of the cover image of the board."""
|
||||
archived: bool = Field(description="Whether or not the board is archived.")
|
||||
"""Whether or not the board is archived."""
|
||||
board_visibility: BoardVisibility = Field(
|
||||
default=BoardVisibility.Private, description="The visibility of the board."
|
||||
)
|
||||
"""The visibility of the board (private, shared, or public)."""
|
||||
|
||||
|
||||
def deserialize_board_record(board_dict: dict) -> BoardRecord:
|
||||
"""Deserializes a board record."""
|
||||
|
||||
# Retrieve all the values, setting "reasonable" defaults if they are not present.
|
||||
|
||||
board_id = board_dict.get("board_id", "unknown")
|
||||
board_name = board_dict.get("board_name", "unknown")
|
||||
# Default to 'system' for backwards compatibility with boards created before multiuser support
|
||||
user_id = board_dict.get("user_id", "system")
|
||||
cover_image_name = board_dict.get("cover_image_name", "unknown")
|
||||
created_at = board_dict.get("created_at", get_iso_timestamp())
|
||||
updated_at = board_dict.get("updated_at", get_iso_timestamp())
|
||||
deleted_at = board_dict.get("deleted_at", get_iso_timestamp())
|
||||
archived = board_dict.get("archived", False)
|
||||
board_visibility_raw = board_dict.get("board_visibility", BoardVisibility.Private.value)
|
||||
try:
|
||||
board_visibility = BoardVisibility(board_visibility_raw)
|
||||
except ValueError:
|
||||
board_visibility = BoardVisibility.Private
|
||||
|
||||
return BoardRecord(
|
||||
board_id=board_id,
|
||||
board_name=board_name,
|
||||
user_id=user_id,
|
||||
cover_image_name=cover_image_name,
|
||||
created_at=created_at,
|
||||
updated_at=updated_at,
|
||||
deleted_at=deleted_at,
|
||||
archived=archived,
|
||||
board_visibility=board_visibility,
|
||||
)
|
||||
|
||||
|
||||
class BoardChanges(BaseModel, extra="forbid"):
|
||||
board_name: Optional[str] = Field(default=None, description="The board's new name.", max_length=300)
|
||||
cover_image_name: Optional[str] = Field(default=None, description="The name of the board's new cover image.")
|
||||
archived: Optional[bool] = Field(default=None, description="Whether or not the board is archived")
|
||||
board_visibility: Optional[BoardVisibility] = Field(default=None, description="The visibility of the board.")
|
||||
|
||||
|
||||
class BoardRecordOrderBy(str, Enum, metaclass=MetaEnum):
|
||||
"""The order by options for board records"""
|
||||
|
||||
CreatedAt = "created_at"
|
||||
Name = "board_name"
|
||||
|
||||
|
||||
class BoardRecordNotFoundException(Exception):
|
||||
"""Raised when an board record is not found."""
|
||||
|
||||
def __init__(self, message="Board record not found"):
|
||||
super().__init__(message)
|
||||
|
||||
|
||||
class BoardRecordSaveException(Exception):
|
||||
"""Raised when an board record cannot be saved."""
|
||||
|
||||
def __init__(self, message="Board record not saved"):
|
||||
super().__init__(message)
|
||||
|
||||
|
||||
class BoardRecordDeleteException(Exception):
|
||||
"""Raised when an board record cannot be deleted."""
|
||||
|
||||
def __init__(self, message="Board record not deleted"):
|
||||
super().__init__(message)
|
||||
@@ -0,0 +1,290 @@
|
||||
import sqlite3
|
||||
from typing import Union, cast
|
||||
|
||||
from invokeai.app.services.board_records.board_records_base import BoardRecordStorageBase
|
||||
from invokeai.app.services.board_records.board_records_common import (
|
||||
BoardChanges,
|
||||
BoardRecord,
|
||||
BoardRecordDeleteException,
|
||||
BoardRecordNotFoundException,
|
||||
BoardRecordOrderBy,
|
||||
BoardRecordSaveException,
|
||||
deserialize_board_record,
|
||||
)
|
||||
from invokeai.app.services.shared.pagination import OffsetPaginatedResults
|
||||
from invokeai.app.services.shared.sqlite.sqlite_common import SQLiteDirection
|
||||
from invokeai.app.services.shared.sqlite.sqlite_database import SqliteDatabase
|
||||
from invokeai.app.util.misc import uuid_string
|
||||
|
||||
|
||||
class SqliteBoardRecordStorage(BoardRecordStorageBase):
|
||||
def __init__(self, db: SqliteDatabase) -> None:
|
||||
super().__init__()
|
||||
self._db = db
|
||||
|
||||
def delete(self, board_id: str) -> None:
|
||||
with self._db.transaction() as cursor:
|
||||
try:
|
||||
cursor.execute(
|
||||
"""--sql
|
||||
DELETE FROM boards
|
||||
WHERE board_id = ?;
|
||||
""",
|
||||
(board_id,),
|
||||
)
|
||||
except Exception as e:
|
||||
raise BoardRecordDeleteException from e
|
||||
|
||||
def save(
|
||||
self,
|
||||
board_name: str,
|
||||
user_id: str,
|
||||
) -> BoardRecord:
|
||||
with self._db.transaction() as cursor:
|
||||
try:
|
||||
board_id = uuid_string()
|
||||
cursor.execute(
|
||||
"""--sql
|
||||
INSERT OR IGNORE INTO boards (board_id, board_name, user_id)
|
||||
VALUES (?, ?, ?);
|
||||
""",
|
||||
(board_id, board_name, user_id),
|
||||
)
|
||||
except sqlite3.Error as e:
|
||||
raise BoardRecordSaveException from e
|
||||
return self.get(board_id)
|
||||
|
||||
def get(
|
||||
self,
|
||||
board_id: str,
|
||||
) -> BoardRecord:
|
||||
with self._db.transaction() as cursor:
|
||||
try:
|
||||
cursor.execute(
|
||||
"""--sql
|
||||
SELECT *
|
||||
FROM boards
|
||||
WHERE board_id = ?;
|
||||
""",
|
||||
(board_id,),
|
||||
)
|
||||
|
||||
result = cast(Union[sqlite3.Row, None], cursor.fetchone())
|
||||
except sqlite3.Error as e:
|
||||
raise BoardRecordNotFoundException from e
|
||||
if result is None:
|
||||
raise BoardRecordNotFoundException
|
||||
return BoardRecord(**dict(result))
|
||||
|
||||
def update(
|
||||
self,
|
||||
board_id: str,
|
||||
changes: BoardChanges,
|
||||
) -> BoardRecord:
|
||||
with self._db.transaction() as cursor:
|
||||
try:
|
||||
# Change the name of a board
|
||||
if changes.board_name is not None:
|
||||
cursor.execute(
|
||||
"""--sql
|
||||
UPDATE boards
|
||||
SET board_name = ?
|
||||
WHERE board_id = ?;
|
||||
""",
|
||||
(changes.board_name, board_id),
|
||||
)
|
||||
|
||||
# Change the cover image of a board
|
||||
if changes.cover_image_name is not None:
|
||||
cursor.execute(
|
||||
"""--sql
|
||||
UPDATE boards
|
||||
SET cover_image_name = ?
|
||||
WHERE board_id = ?;
|
||||
""",
|
||||
(changes.cover_image_name, board_id),
|
||||
)
|
||||
|
||||
# Change the archived status of a board
|
||||
if changes.archived is not None:
|
||||
cursor.execute(
|
||||
"""--sql
|
||||
UPDATE boards
|
||||
SET archived = ?
|
||||
WHERE board_id = ?;
|
||||
""",
|
||||
(changes.archived, board_id),
|
||||
)
|
||||
|
||||
# Change the visibility of a board
|
||||
if changes.board_visibility is not None:
|
||||
cursor.execute(
|
||||
"""--sql
|
||||
UPDATE boards
|
||||
SET board_visibility = ?
|
||||
WHERE board_id = ?;
|
||||
""",
|
||||
(changes.board_visibility.value, board_id),
|
||||
)
|
||||
|
||||
except sqlite3.Error as e:
|
||||
raise BoardRecordSaveException from e
|
||||
return self.get(board_id)
|
||||
|
||||
def get_many(
|
||||
self,
|
||||
user_id: str,
|
||||
is_admin: bool,
|
||||
order_by: BoardRecordOrderBy,
|
||||
direction: SQLiteDirection,
|
||||
offset: int = 0,
|
||||
limit: int = 10,
|
||||
include_archived: bool = False,
|
||||
) -> OffsetPaginatedResults[BoardRecord]:
|
||||
with self._db.transaction() as cursor:
|
||||
# Build base query - admins see all boards, regular users see owned, shared, or public boards
|
||||
if is_admin:
|
||||
base_query = """
|
||||
SELECT DISTINCT boards.*
|
||||
FROM boards
|
||||
{archived_filter}
|
||||
ORDER BY {order_by} {direction}
|
||||
LIMIT ? OFFSET ?;
|
||||
"""
|
||||
|
||||
# Determine archived filter condition
|
||||
archived_filter = "WHERE 1=1" if include_archived else "WHERE boards.archived = 0"
|
||||
|
||||
final_query = base_query.format(
|
||||
archived_filter=archived_filter, order_by=order_by.value, direction=direction.value
|
||||
)
|
||||
|
||||
# Execute query to fetch boards
|
||||
cursor.execute(final_query, (limit, offset))
|
||||
else:
|
||||
base_query = """
|
||||
SELECT DISTINCT boards.*
|
||||
FROM boards
|
||||
LEFT JOIN shared_boards ON boards.board_id = shared_boards.board_id
|
||||
WHERE (boards.user_id = ? OR shared_boards.user_id = ? OR boards.board_visibility IN ('shared', 'public'))
|
||||
{archived_filter}
|
||||
ORDER BY {order_by} {direction}
|
||||
LIMIT ? OFFSET ?;
|
||||
"""
|
||||
|
||||
# Determine archived filter condition
|
||||
archived_filter = "" if include_archived else "AND boards.archived = 0"
|
||||
|
||||
final_query = base_query.format(
|
||||
archived_filter=archived_filter, order_by=order_by.value, direction=direction.value
|
||||
)
|
||||
|
||||
# Execute query to fetch boards
|
||||
cursor.execute(final_query, (user_id, user_id, limit, offset))
|
||||
|
||||
result = cast(list[sqlite3.Row], cursor.fetchall())
|
||||
boards = [deserialize_board_record(dict(r)) for r in result]
|
||||
|
||||
# Determine count query - admins count all boards, regular users count accessible boards
|
||||
if is_admin:
|
||||
if include_archived:
|
||||
count_query = """
|
||||
SELECT COUNT(DISTINCT boards.board_id)
|
||||
FROM boards;
|
||||
"""
|
||||
else:
|
||||
count_query = """
|
||||
SELECT COUNT(DISTINCT boards.board_id)
|
||||
FROM boards
|
||||
WHERE boards.archived = 0;
|
||||
"""
|
||||
cursor.execute(count_query)
|
||||
else:
|
||||
if include_archived:
|
||||
count_query = """
|
||||
SELECT COUNT(DISTINCT boards.board_id)
|
||||
FROM boards
|
||||
LEFT JOIN shared_boards ON boards.board_id = shared_boards.board_id
|
||||
WHERE (boards.user_id = ? OR shared_boards.user_id = ? OR boards.board_visibility IN ('shared', 'public'));
|
||||
"""
|
||||
else:
|
||||
count_query = """
|
||||
SELECT COUNT(DISTINCT boards.board_id)
|
||||
FROM boards
|
||||
LEFT JOIN shared_boards ON boards.board_id = shared_boards.board_id
|
||||
WHERE (boards.user_id = ? OR shared_boards.user_id = ? OR boards.board_visibility IN ('shared', 'public'))
|
||||
AND boards.archived = 0;
|
||||
"""
|
||||
|
||||
# Execute count query
|
||||
cursor.execute(count_query, (user_id, user_id))
|
||||
|
||||
count = cast(int, cursor.fetchone()[0])
|
||||
|
||||
return OffsetPaginatedResults[BoardRecord](items=boards, offset=offset, limit=limit, total=count)
|
||||
|
||||
def get_all(
|
||||
self,
|
||||
user_id: str,
|
||||
is_admin: bool,
|
||||
order_by: BoardRecordOrderBy,
|
||||
direction: SQLiteDirection,
|
||||
include_archived: bool = False,
|
||||
) -> list[BoardRecord]:
|
||||
with self._db.transaction() as cursor:
|
||||
# Build query - admins see all boards, regular users see owned, shared, or public boards
|
||||
if is_admin:
|
||||
if order_by == BoardRecordOrderBy.Name:
|
||||
base_query = """
|
||||
SELECT DISTINCT boards.*
|
||||
FROM boards
|
||||
{archived_filter}
|
||||
ORDER BY LOWER(boards.board_name) {direction}
|
||||
"""
|
||||
else:
|
||||
base_query = """
|
||||
SELECT DISTINCT boards.*
|
||||
FROM boards
|
||||
{archived_filter}
|
||||
ORDER BY {order_by} {direction}
|
||||
"""
|
||||
|
||||
archived_filter = "WHERE 1=1" if include_archived else "WHERE boards.archived = 0"
|
||||
|
||||
final_query = base_query.format(
|
||||
archived_filter=archived_filter, order_by=order_by.value, direction=direction.value
|
||||
)
|
||||
|
||||
cursor.execute(final_query)
|
||||
else:
|
||||
if order_by == BoardRecordOrderBy.Name:
|
||||
base_query = """
|
||||
SELECT DISTINCT boards.*
|
||||
FROM boards
|
||||
LEFT JOIN shared_boards ON boards.board_id = shared_boards.board_id
|
||||
WHERE (boards.user_id = ? OR shared_boards.user_id = ? OR boards.board_visibility IN ('shared', 'public'))
|
||||
{archived_filter}
|
||||
ORDER BY LOWER(boards.board_name) {direction}
|
||||
"""
|
||||
else:
|
||||
base_query = """
|
||||
SELECT DISTINCT boards.*
|
||||
FROM boards
|
||||
LEFT JOIN shared_boards ON boards.board_id = shared_boards.board_id
|
||||
WHERE (boards.user_id = ? OR shared_boards.user_id = ? OR boards.board_visibility IN ('shared', 'public'))
|
||||
{archived_filter}
|
||||
ORDER BY {order_by} {direction}
|
||||
"""
|
||||
|
||||
archived_filter = "" if include_archived else "AND boards.archived = 0"
|
||||
|
||||
final_query = base_query.format(
|
||||
archived_filter=archived_filter, order_by=order_by.value, direction=direction.value
|
||||
)
|
||||
|
||||
cursor.execute(final_query, (user_id, user_id))
|
||||
|
||||
result = cast(list[sqlite3.Row], cursor.fetchall())
|
||||
boards = [deserialize_board_record(dict(r)) for r in result]
|
||||
|
||||
return boards
|
||||
@@ -0,0 +1,70 @@
|
||||
from abc import ABC, abstractmethod
|
||||
|
||||
from invokeai.app.services.board_records.board_records_common import BoardChanges, BoardRecordOrderBy
|
||||
from invokeai.app.services.boards.boards_common import BoardDTO
|
||||
from invokeai.app.services.shared.pagination import OffsetPaginatedResults
|
||||
from invokeai.app.services.shared.sqlite.sqlite_common import SQLiteDirection
|
||||
|
||||
|
||||
class BoardServiceABC(ABC):
|
||||
"""High-level service for board management."""
|
||||
|
||||
@abstractmethod
|
||||
def create(
|
||||
self,
|
||||
board_name: str,
|
||||
user_id: str,
|
||||
) -> BoardDTO:
|
||||
"""Creates a board for a specific user."""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def get_dto(
|
||||
self,
|
||||
board_id: str,
|
||||
) -> BoardDTO:
|
||||
"""Gets a board."""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def update(
|
||||
self,
|
||||
board_id: str,
|
||||
changes: BoardChanges,
|
||||
) -> BoardDTO:
|
||||
"""Updates a board."""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def delete(
|
||||
self,
|
||||
board_id: str,
|
||||
) -> None:
|
||||
"""Deletes a board."""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def get_many(
|
||||
self,
|
||||
user_id: str,
|
||||
is_admin: bool,
|
||||
order_by: BoardRecordOrderBy,
|
||||
direction: SQLiteDirection,
|
||||
offset: int = 0,
|
||||
limit: int = 10,
|
||||
include_archived: bool = False,
|
||||
) -> OffsetPaginatedResults[BoardDTO]:
|
||||
"""Gets many boards for a specific user, including shared boards. Admin users see all boards."""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def get_all(
|
||||
self,
|
||||
user_id: str,
|
||||
is_admin: bool,
|
||||
order_by: BoardRecordOrderBy,
|
||||
direction: SQLiteDirection,
|
||||
include_archived: bool = False,
|
||||
) -> list[BoardDTO]:
|
||||
"""Gets all boards for a specific user, including shared boards. Admin users see all boards."""
|
||||
pass
|
||||
@@ -0,0 +1,35 @@
|
||||
from typing import Optional
|
||||
|
||||
from pydantic import Field
|
||||
|
||||
from invokeai.app.services.board_records.board_records_common import BoardRecord
|
||||
|
||||
|
||||
class BoardDTO(BoardRecord):
|
||||
"""Deserialized board record with cover image URL and image count."""
|
||||
|
||||
cover_image_name: Optional[str] = Field(description="The name of the board's cover image.")
|
||||
"""The URL of the thumbnail of the most recent image in the board."""
|
||||
image_count: int = Field(description="The number of images in the board.")
|
||||
"""The number of images in the board."""
|
||||
asset_count: int = Field(description="The number of assets in the board.")
|
||||
"""The number of assets in the board."""
|
||||
owner_username: Optional[str] = Field(default=None, description="The username of the board owner (for admin view).")
|
||||
"""The username of the board owner (for admin view)."""
|
||||
|
||||
|
||||
def board_record_to_dto(
|
||||
board_record: BoardRecord,
|
||||
cover_image_name: Optional[str],
|
||||
image_count: int,
|
||||
asset_count: int,
|
||||
owner_username: Optional[str] = None,
|
||||
) -> BoardDTO:
|
||||
"""Converts a board record to a board DTO."""
|
||||
return BoardDTO(
|
||||
**board_record.model_dump(exclude={"cover_image_name"}),
|
||||
cover_image_name=cover_image_name,
|
||||
image_count=image_count,
|
||||
asset_count=asset_count,
|
||||
owner_username=owner_username,
|
||||
)
|
||||
@@ -0,0 +1,119 @@
|
||||
from invokeai.app.services.board_records.board_records_common import BoardChanges, BoardRecordOrderBy
|
||||
from invokeai.app.services.boards.boards_base import BoardServiceABC
|
||||
from invokeai.app.services.boards.boards_common import BoardDTO, board_record_to_dto
|
||||
from invokeai.app.services.invoker import Invoker
|
||||
from invokeai.app.services.shared.pagination import OffsetPaginatedResults
|
||||
from invokeai.app.services.shared.sqlite.sqlite_common import SQLiteDirection
|
||||
|
||||
|
||||
class BoardService(BoardServiceABC):
|
||||
__invoker: Invoker
|
||||
|
||||
def start(self, invoker: Invoker) -> None:
|
||||
self.__invoker = invoker
|
||||
|
||||
def create(
|
||||
self,
|
||||
board_name: str,
|
||||
user_id: str,
|
||||
) -> BoardDTO:
|
||||
board_record = self.__invoker.services.board_records.save(board_name, user_id)
|
||||
return board_record_to_dto(board_record, None, 0, 0)
|
||||
|
||||
def get_dto(self, board_id: str) -> BoardDTO:
|
||||
board_record = self.__invoker.services.board_records.get(board_id)
|
||||
cover_image = self.__invoker.services.image_records.get_most_recent_image_for_board(board_record.board_id)
|
||||
if cover_image:
|
||||
cover_image_name = cover_image.image_name
|
||||
else:
|
||||
cover_image_name = None
|
||||
image_count = self.__invoker.services.board_image_records.get_image_count_for_board(board_id)
|
||||
asset_count = self.__invoker.services.board_image_records.get_asset_count_for_board(board_id)
|
||||
return board_record_to_dto(board_record, cover_image_name, image_count, asset_count)
|
||||
|
||||
def update(
|
||||
self,
|
||||
board_id: str,
|
||||
changes: BoardChanges,
|
||||
) -> BoardDTO:
|
||||
board_record = self.__invoker.services.board_records.update(board_id, changes)
|
||||
cover_image = self.__invoker.services.image_records.get_most_recent_image_for_board(board_record.board_id)
|
||||
if cover_image:
|
||||
cover_image_name = cover_image.image_name
|
||||
else:
|
||||
cover_image_name = None
|
||||
|
||||
image_count = self.__invoker.services.board_image_records.get_image_count_for_board(board_id)
|
||||
asset_count = self.__invoker.services.board_image_records.get_asset_count_for_board(board_id)
|
||||
return board_record_to_dto(board_record, cover_image_name, image_count, asset_count)
|
||||
|
||||
def delete(self, board_id: str) -> None:
|
||||
self.__invoker.services.board_records.delete(board_id)
|
||||
|
||||
def get_many(
|
||||
self,
|
||||
user_id: str,
|
||||
is_admin: bool,
|
||||
order_by: BoardRecordOrderBy,
|
||||
direction: SQLiteDirection,
|
||||
offset: int = 0,
|
||||
limit: int = 10,
|
||||
include_archived: bool = False,
|
||||
) -> OffsetPaginatedResults[BoardDTO]:
|
||||
board_records = self.__invoker.services.board_records.get_many(
|
||||
user_id, is_admin, order_by, direction, offset, limit, include_archived
|
||||
)
|
||||
board_dtos = []
|
||||
for r in board_records.items:
|
||||
cover_image = self.__invoker.services.image_records.get_most_recent_image_for_board(r.board_id)
|
||||
if cover_image:
|
||||
cover_image_name = cover_image.image_name
|
||||
else:
|
||||
cover_image_name = None
|
||||
|
||||
image_count = self.__invoker.services.board_image_records.get_image_count_for_board(r.board_id)
|
||||
asset_count = self.__invoker.services.board_image_records.get_asset_count_for_board(r.board_id)
|
||||
|
||||
# For admin users, include owner username
|
||||
owner_username = None
|
||||
if is_admin:
|
||||
owner = self.__invoker.services.users.get(r.user_id)
|
||||
if owner:
|
||||
owner_username = owner.display_name or owner.email
|
||||
|
||||
board_dtos.append(board_record_to_dto(r, cover_image_name, image_count, asset_count, owner_username))
|
||||
|
||||
return OffsetPaginatedResults[BoardDTO](items=board_dtos, offset=offset, limit=limit, total=len(board_dtos))
|
||||
|
||||
def get_all(
|
||||
self,
|
||||
user_id: str,
|
||||
is_admin: bool,
|
||||
order_by: BoardRecordOrderBy,
|
||||
direction: SQLiteDirection,
|
||||
include_archived: bool = False,
|
||||
) -> list[BoardDTO]:
|
||||
board_records = self.__invoker.services.board_records.get_all(
|
||||
user_id, is_admin, order_by, direction, include_archived
|
||||
)
|
||||
board_dtos = []
|
||||
for r in board_records:
|
||||
cover_image = self.__invoker.services.image_records.get_most_recent_image_for_board(r.board_id)
|
||||
if cover_image:
|
||||
cover_image_name = cover_image.image_name
|
||||
else:
|
||||
cover_image_name = None
|
||||
|
||||
image_count = self.__invoker.services.board_image_records.get_image_count_for_board(r.board_id)
|
||||
asset_count = self.__invoker.services.board_image_records.get_asset_count_for_board(r.board_id)
|
||||
|
||||
# For admin users, include owner username
|
||||
owner_username = None
|
||||
if is_admin:
|
||||
owner = self.__invoker.services.users.get(r.user_id)
|
||||
if owner:
|
||||
owner_username = owner.display_name or owner.email
|
||||
|
||||
board_dtos.append(board_record_to_dto(r, cover_image_name, image_count, asset_count, owner_username))
|
||||
|
||||
return board_dtos
|
||||
@@ -0,0 +1,58 @@
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import Optional
|
||||
|
||||
|
||||
class BulkDownloadBase(ABC):
|
||||
"""Responsible for creating a zip file containing the images specified by the given image names or board id."""
|
||||
|
||||
@abstractmethod
|
||||
def handler(
|
||||
self,
|
||||
image_names: Optional[list[str]],
|
||||
board_id: Optional[str],
|
||||
bulk_download_item_id: Optional[str],
|
||||
user_id: str = "system",
|
||||
) -> None:
|
||||
"""
|
||||
Create a zip file containing the images specified by the given image names or board id.
|
||||
|
||||
:param image_names: A list of image names to include in the zip file.
|
||||
:param board_id: The ID of the board. If provided, all images associated with the board will be included in the zip file.
|
||||
:param bulk_download_item_id: The bulk_download_item_id that will be used to retrieve the bulk download item when it is prepared, if none is provided a uuid will be generated.
|
||||
:param user_id: The ID of the user who initiated the download.
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
def get_path(self, bulk_download_item_name: str) -> str:
|
||||
"""
|
||||
Get the path to the bulk download file.
|
||||
|
||||
:param bulk_download_item_name: The name of the bulk download item.
|
||||
:return: The path to the bulk download file.
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
def generate_item_id(self, board_id: Optional[str]) -> str:
|
||||
"""
|
||||
Generate an item ID for a bulk download item.
|
||||
|
||||
:param board_id: The ID of the board whose name is to be included in the item id.
|
||||
:return: The generated item ID.
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
def delete(self, bulk_download_item_name: str) -> None:
|
||||
"""
|
||||
Delete the bulk download file.
|
||||
|
||||
:param bulk_download_item_name: The name of the bulk download item.
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
def get_owner(self, bulk_download_item_name: str) -> Optional[str]:
|
||||
"""
|
||||
Get the user_id of the user who initiated the download.
|
||||
|
||||
:param bulk_download_item_name: The name of the bulk download item.
|
||||
:return: The user_id of the owner, or None if not tracked.
|
||||
"""
|
||||
@@ -0,0 +1,25 @@
|
||||
DEFAULT_BULK_DOWNLOAD_ID = "default"
|
||||
|
||||
|
||||
class BulkDownloadException(Exception):
|
||||
"""Exception raised when a bulk download fails."""
|
||||
|
||||
def __init__(self, message="Bulk download failed"):
|
||||
super().__init__(message)
|
||||
self.message = message
|
||||
|
||||
|
||||
class BulkDownloadTargetException(BulkDownloadException):
|
||||
"""Exception raised when a bulk download target is not found."""
|
||||
|
||||
def __init__(self, message="The bulk download target was not found"):
|
||||
super().__init__(message)
|
||||
self.message = message
|
||||
|
||||
|
||||
class BulkDownloadParametersException(BulkDownloadException):
|
||||
"""Exception raised when a bulk download parameter is invalid."""
|
||||
|
||||
def __init__(self, message="No image names or board ID provided"):
|
||||
super().__init__(message)
|
||||
self.message = message
|
||||
@@ -0,0 +1,190 @@
|
||||
from pathlib import Path
|
||||
from tempfile import TemporaryDirectory
|
||||
from typing import Optional, Union
|
||||
from zipfile import ZipFile
|
||||
|
||||
from invokeai.app.services.board_records.board_records_common import BoardRecordNotFoundException
|
||||
from invokeai.app.services.bulk_download.bulk_download_base import BulkDownloadBase
|
||||
from invokeai.app.services.bulk_download.bulk_download_common import (
|
||||
DEFAULT_BULK_DOWNLOAD_ID,
|
||||
BulkDownloadException,
|
||||
BulkDownloadParametersException,
|
||||
BulkDownloadTargetException,
|
||||
)
|
||||
from invokeai.app.services.image_records.image_records_common import ImageRecordNotFoundException
|
||||
from invokeai.app.services.images.images_common import ImageDTO
|
||||
from invokeai.app.services.invoker import Invoker
|
||||
from invokeai.app.util.misc import uuid_string
|
||||
|
||||
|
||||
class BulkDownloadService(BulkDownloadBase):
|
||||
def start(self, invoker: Invoker) -> None:
|
||||
self._invoker = invoker
|
||||
|
||||
def __init__(self):
|
||||
self._temp_directory = TemporaryDirectory()
|
||||
self._bulk_downloads_folder = Path(self._temp_directory.name) / "bulk_downloads"
|
||||
self._bulk_downloads_folder.mkdir(parents=True, exist_ok=True)
|
||||
# Track which user owns each download so the fetch endpoint can enforce ownership
|
||||
self._download_owners: dict[str, str] = {}
|
||||
|
||||
def handler(
|
||||
self,
|
||||
image_names: Optional[list[str]],
|
||||
board_id: Optional[str],
|
||||
bulk_download_item_id: Optional[str],
|
||||
user_id: str = "system",
|
||||
) -> None:
|
||||
bulk_download_id: str = DEFAULT_BULK_DOWNLOAD_ID
|
||||
bulk_download_item_id = bulk_download_item_id or uuid_string()
|
||||
bulk_download_item_name = bulk_download_item_id + ".zip"
|
||||
|
||||
# Record ownership so the fetch endpoint can verify the caller
|
||||
self._download_owners[bulk_download_item_name] = user_id
|
||||
|
||||
self._signal_job_started(bulk_download_id, bulk_download_item_id, bulk_download_item_name, user_id)
|
||||
|
||||
try:
|
||||
image_dtos: list[ImageDTO] = []
|
||||
|
||||
if board_id:
|
||||
image_dtos = self._board_handler(board_id)
|
||||
elif image_names:
|
||||
image_dtos = self._image_handler(image_names)
|
||||
else:
|
||||
raise BulkDownloadParametersException()
|
||||
|
||||
bulk_download_item_name: str = self._create_zip_file(image_dtos, bulk_download_item_id)
|
||||
self._signal_job_completed(bulk_download_id, bulk_download_item_id, bulk_download_item_name, user_id)
|
||||
except (
|
||||
ImageRecordNotFoundException,
|
||||
BoardRecordNotFoundException,
|
||||
BulkDownloadException,
|
||||
BulkDownloadParametersException,
|
||||
) as e:
|
||||
self._signal_job_failed(bulk_download_id, bulk_download_item_id, bulk_download_item_name, e, user_id)
|
||||
except Exception as e:
|
||||
self._signal_job_failed(bulk_download_id, bulk_download_item_id, bulk_download_item_name, e, user_id)
|
||||
self._invoker.services.logger.error("Problem bulk downloading images.")
|
||||
raise e
|
||||
|
||||
def _image_handler(self, image_names: list[str]) -> list[ImageDTO]:
|
||||
return [self._invoker.services.images.get_dto(image_name) for image_name in image_names]
|
||||
|
||||
def _board_handler(self, board_id: str) -> list[ImageDTO]:
|
||||
image_names = self._invoker.services.board_image_records.get_all_board_image_names_for_board(
|
||||
board_id,
|
||||
categories=None,
|
||||
is_intermediate=None,
|
||||
)
|
||||
return self._image_handler(image_names)
|
||||
|
||||
def generate_item_id(self, board_id: Optional[str]) -> str:
|
||||
return uuid_string() if board_id is None else self._get_clean_board_name(board_id) + "_" + uuid_string()
|
||||
|
||||
def _get_clean_board_name(self, board_id: str) -> str:
|
||||
if board_id == "none":
|
||||
return "Uncategorized"
|
||||
|
||||
return self._clean_string_to_path_safe(self._invoker.services.board_records.get(board_id).board_name)
|
||||
|
||||
def _create_zip_file(self, image_dtos: list[ImageDTO], bulk_download_item_id: str) -> str:
|
||||
"""
|
||||
Create a zip file containing the images specified by the given image names or board id.
|
||||
If download with the same bulk_download_id already exists, it will be overwritten.
|
||||
|
||||
:return: The name of the zip file.
|
||||
"""
|
||||
zip_file_name = bulk_download_item_id + ".zip"
|
||||
zip_file_path = self._bulk_downloads_folder / (zip_file_name)
|
||||
|
||||
with ZipFile(zip_file_path, "w") as zip_file:
|
||||
for image_dto in image_dtos:
|
||||
image_zip_path = Path(image_dto.image_category.value) / image_dto.image_name
|
||||
image_disk_path = self._invoker.services.images.get_path(image_dto.image_name)
|
||||
zip_file.write(image_disk_path, arcname=image_zip_path)
|
||||
|
||||
return str(zip_file_name)
|
||||
|
||||
# from https://stackoverflow.com/questions/7406102/create-sane-safe-filename-from-any-unsafe-string
|
||||
def _clean_string_to_path_safe(self, s: str) -> str:
|
||||
"""Clean a string to be path safe."""
|
||||
return "".join([c for c in s if c.isalpha() or c.isdigit() or c == " " or c == "_" or c == "-"]).rstrip()
|
||||
|
||||
def _signal_job_started(
|
||||
self,
|
||||
bulk_download_id: str,
|
||||
bulk_download_item_id: str,
|
||||
bulk_download_item_name: str,
|
||||
user_id: str = "system",
|
||||
) -> None:
|
||||
"""Signal that a bulk download job has started."""
|
||||
if self._invoker:
|
||||
assert bulk_download_id is not None
|
||||
self._invoker.services.events.emit_bulk_download_started(
|
||||
bulk_download_id, bulk_download_item_id, bulk_download_item_name, user_id=user_id
|
||||
)
|
||||
|
||||
def _signal_job_completed(
|
||||
self,
|
||||
bulk_download_id: str,
|
||||
bulk_download_item_id: str,
|
||||
bulk_download_item_name: str,
|
||||
user_id: str = "system",
|
||||
) -> None:
|
||||
"""Signal that a bulk download job has completed."""
|
||||
if self._invoker:
|
||||
assert bulk_download_id is not None
|
||||
assert bulk_download_item_name is not None
|
||||
self._invoker.services.events.emit_bulk_download_complete(
|
||||
bulk_download_id, bulk_download_item_id, bulk_download_item_name, user_id=user_id
|
||||
)
|
||||
|
||||
def _signal_job_failed(
|
||||
self,
|
||||
bulk_download_id: str,
|
||||
bulk_download_item_id: str,
|
||||
bulk_download_item_name: str,
|
||||
exception: Exception,
|
||||
user_id: str = "system",
|
||||
) -> None:
|
||||
"""Signal that a bulk download job has failed."""
|
||||
if self._invoker:
|
||||
assert bulk_download_id is not None
|
||||
assert exception is not None
|
||||
self._invoker.services.events.emit_bulk_download_error(
|
||||
bulk_download_id, bulk_download_item_id, bulk_download_item_name, str(exception), user_id=user_id
|
||||
)
|
||||
|
||||
def stop(self, *args, **kwargs):
|
||||
self._temp_directory.cleanup()
|
||||
|
||||
def get_owner(self, bulk_download_item_name: str) -> Optional[str]:
|
||||
return self._download_owners.get(bulk_download_item_name)
|
||||
|
||||
def delete(self, bulk_download_item_name: str) -> None:
|
||||
path = self.get_path(bulk_download_item_name)
|
||||
Path(path).unlink()
|
||||
self._download_owners.pop(bulk_download_item_name, None)
|
||||
|
||||
def get_path(self, bulk_download_item_name: str) -> str:
|
||||
path = str(self._bulk_downloads_folder / bulk_download_item_name)
|
||||
if not self._is_valid_path(path):
|
||||
raise BulkDownloadTargetException()
|
||||
return path
|
||||
|
||||
def _is_valid_path(self, path: Union[str, Path]) -> bool:
|
||||
"""Validates the path given for a bulk download."""
|
||||
path = path if isinstance(path, Path) else Path(path)
|
||||
|
||||
# Resolve the path to handle any path traversal attempts (e.g., ../)
|
||||
resolved_path = path.resolve()
|
||||
|
||||
# The path may not traverse out of the bulk downloads folder or its subfolders
|
||||
does_not_traverse = resolved_path.parent == self._bulk_downloads_folder.resolve()
|
||||
|
||||
# The path must exist and be a .zip file
|
||||
does_exist = resolved_path.exists()
|
||||
is_zip_file = resolved_path.suffix == ".zip"
|
||||
|
||||
return does_exist and is_zip_file and does_not_traverse
|
||||
@@ -0,0 +1,72 @@
|
||||
from abc import ABC, abstractmethod
|
||||
|
||||
|
||||
class ClientStatePersistenceABC(ABC):
|
||||
"""
|
||||
Base class for client persistence implementations.
|
||||
This class defines the interface for persisting client data per user.
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
def set_by_key(self, user_id: str, key: str, value: str) -> str:
|
||||
"""
|
||||
Set a key-value pair for the client.
|
||||
|
||||
Args:
|
||||
user_id (str): The user ID to set state for.
|
||||
key (str): The key to set.
|
||||
value (str): The value to set for the key.
|
||||
|
||||
Returns:
|
||||
str: The value that was set.
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def get_by_key(self, user_id: str, key: str) -> str | None:
|
||||
"""
|
||||
Get the value for a specific key of the client.
|
||||
|
||||
Args:
|
||||
user_id (str): The user ID to get state for.
|
||||
key (str): The key to retrieve the value for.
|
||||
|
||||
Returns:
|
||||
str | None: The value associated with the key, or None if the key does not exist.
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def get_keys_by_prefix(self, user_id: str, prefix: str) -> list[str]:
|
||||
"""
|
||||
Get all keys matching a prefix for a user.
|
||||
|
||||
Args:
|
||||
user_id (str): The user ID to get keys for.
|
||||
prefix (str): The prefix to filter keys by.
|
||||
|
||||
Returns:
|
||||
list[str]: A list of keys matching the prefix.
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def delete_by_key(self, user_id: str, key: str) -> None:
|
||||
"""
|
||||
Delete a specific key-value pair for a user.
|
||||
|
||||
Args:
|
||||
user_id (str): The user ID to delete state for.
|
||||
key (str): The key to delete.
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def delete(self, user_id: str) -> None:
|
||||
"""
|
||||
Delete all client state for a user.
|
||||
|
||||
Args:
|
||||
user_id (str): The user ID to delete state for.
|
||||
"""
|
||||
pass
|
||||
@@ -0,0 +1,80 @@
|
||||
from invokeai.app.services.client_state_persistence.client_state_persistence_base import ClientStatePersistenceABC
|
||||
from invokeai.app.services.invoker import Invoker
|
||||
from invokeai.app.services.shared.sqlite.sqlite_database import SqliteDatabase
|
||||
|
||||
|
||||
class ClientStatePersistenceSqlite(ClientStatePersistenceABC):
|
||||
"""
|
||||
SQLite implementation for client state persistence.
|
||||
This class stores client state data per user to prevent data leakage between users.
|
||||
"""
|
||||
|
||||
def __init__(self, db: SqliteDatabase) -> None:
|
||||
super().__init__()
|
||||
self._db = db
|
||||
|
||||
def start(self, invoker: Invoker) -> None:
|
||||
self._invoker = invoker
|
||||
|
||||
def set_by_key(self, user_id: str, key: str, value: str) -> str:
|
||||
with self._db.transaction() as cursor:
|
||||
cursor.execute(
|
||||
"""
|
||||
INSERT INTO client_state (user_id, key, value)
|
||||
VALUES (?, ?, ?)
|
||||
ON CONFLICT(user_id, key) DO UPDATE
|
||||
SET value = excluded.value;
|
||||
""",
|
||||
(user_id, key, value),
|
||||
)
|
||||
|
||||
return value
|
||||
|
||||
def get_by_key(self, user_id: str, key: str) -> str | None:
|
||||
with self._db.transaction() as cursor:
|
||||
cursor.execute(
|
||||
"""
|
||||
SELECT value FROM client_state
|
||||
WHERE user_id = ? AND key = ?
|
||||
""",
|
||||
(user_id, key),
|
||||
)
|
||||
row = cursor.fetchone()
|
||||
if row is None:
|
||||
return None
|
||||
return row[0]
|
||||
|
||||
def get_keys_by_prefix(self, user_id: str, prefix: str) -> list[str]:
|
||||
# Escape LIKE wildcards (%, _) and the escape char itself so callers can pass
|
||||
# arbitrary strings as a literal prefix without accidental pattern matching.
|
||||
escaped_prefix = prefix.replace("\\", "\\\\").replace("%", "\\%").replace("_", "\\_")
|
||||
with self._db.transaction() as cursor:
|
||||
cursor.execute(
|
||||
"""
|
||||
SELECT key FROM client_state
|
||||
WHERE user_id = ? AND key LIKE ? ESCAPE '\\'
|
||||
ORDER BY updated_at DESC
|
||||
""",
|
||||
(user_id, f"{escaped_prefix}%"),
|
||||
)
|
||||
return [row[0] for row in cursor.fetchall()]
|
||||
|
||||
def delete_by_key(self, user_id: str, key: str) -> None:
|
||||
with self._db.transaction() as cursor:
|
||||
cursor.execute(
|
||||
"""
|
||||
DELETE FROM client_state
|
||||
WHERE user_id = ? AND key = ?
|
||||
""",
|
||||
(user_id, key),
|
||||
)
|
||||
|
||||
def delete(self, user_id: str) -> None:
|
||||
with self._db.transaction() as cursor:
|
||||
cursor.execute(
|
||||
"""
|
||||
DELETE FROM client_state
|
||||
WHERE user_id = ?
|
||||
""",
|
||||
(user_id,),
|
||||
)
|
||||
@@ -0,0 +1,6 @@
|
||||
"""Init file for InvokeAI configure package."""
|
||||
|
||||
from invokeai.app.services.config.config_common import PagingArgumentParser
|
||||
from invokeai.app.services.config.config_default import InvokeAIAppConfig, get_config
|
||||
|
||||
__all__ = ["InvokeAIAppConfig", "get_config", "PagingArgumentParser"]
|
||||
@@ -0,0 +1,25 @@
|
||||
# Copyright (c) 2023 Lincoln Stein (https://github.com/lstein) and the InvokeAI Development Team
|
||||
|
||||
"""
|
||||
Base class for the InvokeAI configuration system.
|
||||
It defines a type of pydantic BaseSettings object that
|
||||
is able to read and write from an omegaconf-based config file,
|
||||
with overriding of settings from environment variables and/or
|
||||
the command line.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import pydoc
|
||||
|
||||
|
||||
class PagingArgumentParser(argparse.ArgumentParser):
|
||||
"""
|
||||
A custom ArgumentParser that uses pydoc to page its output.
|
||||
It also supports reading defaults from an init file.
|
||||
"""
|
||||
|
||||
def print_help(self, file=None) -> None:
|
||||
text = self.format_help()
|
||||
pydoc.pager(text)
|
||||
@@ -0,0 +1,691 @@
|
||||
# TODO(psyche): pydantic-settings supports YAML settings sources. If we can figure out a way to integrate the YAML
|
||||
# migration logic, we could use that for simpler config loading.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import copy
|
||||
import filecmp
|
||||
import locale
|
||||
import os
|
||||
import re
|
||||
import shutil
|
||||
from functools import lru_cache
|
||||
from pathlib import Path
|
||||
from typing import Any, Literal, Optional
|
||||
|
||||
import yaml
|
||||
from pydantic import BaseModel, Field, PrivateAttr, field_validator
|
||||
from pydantic_settings import BaseSettings, PydanticBaseSettingsSource, SettingsConfigDict
|
||||
|
||||
import invokeai.configs as model_configs
|
||||
from invokeai.backend.model_hash.model_hash import HASHING_ALGORITHMS
|
||||
from invokeai.frontend.cli.arg_parser import InvokeAIArgs
|
||||
|
||||
INIT_FILE = Path("invokeai.yaml")
|
||||
API_KEYS_FILE = Path("api_keys.yaml")
|
||||
DB_FILE = Path("invokeai.db")
|
||||
LEGACY_INIT_FILE = Path("invokeai.init")
|
||||
PRECISION = Literal["auto", "float16", "bfloat16", "float32"]
|
||||
ATTENTION_TYPE = Literal["auto", "normal", "xformers", "sliced", "torch-sdp"]
|
||||
ATTENTION_SLICE_SIZE = Literal["auto", "balanced", "max", 1, 2, 3, 4, 5, 6, 7, 8]
|
||||
LOG_FORMAT = Literal["plain", "color", "syslog", "legacy"]
|
||||
LOG_LEVEL = Literal["debug", "info", "warning", "error", "critical"]
|
||||
SESSION_QUEUE_MODE = Literal["FIFO", "round_robin"]
|
||||
IMAGE_SUBFOLDER_STRATEGY = Literal["flat", "date", "type", "hash"]
|
||||
CONFIG_SCHEMA_VERSION = "4.0.3"
|
||||
EXTERNAL_PROVIDER_CONFIG_FIELDS = (
|
||||
"external_alibabacloud_api_key",
|
||||
"external_alibabacloud_base_url",
|
||||
"external_gemini_api_key",
|
||||
"external_gemini_base_url",
|
||||
"external_openai_api_key",
|
||||
"external_openai_base_url",
|
||||
"external_seedream_api_key",
|
||||
"external_seedream_base_url",
|
||||
)
|
||||
|
||||
|
||||
class URLRegexTokenPair(BaseModel):
|
||||
url_regex: str = Field(description="Regular expression to match against the URL")
|
||||
token: str = Field(description="Token to use when the URL matches the regex")
|
||||
|
||||
@field_validator("url_regex")
|
||||
@classmethod
|
||||
def validate_url_regex(cls, v: str) -> str:
|
||||
"""Validate that the value is a valid regex."""
|
||||
try:
|
||||
re.compile(v)
|
||||
except re.error as e:
|
||||
raise ValueError(f"Invalid regex: {e}")
|
||||
return v
|
||||
|
||||
|
||||
class InvokeAIAppConfig(BaseSettings):
|
||||
"""Invoke's global app configuration.
|
||||
|
||||
Typically, you won't need to interact with this class directly. Instead, use the `get_config` function from `invokeai.app.services.config` to get a singleton config object.
|
||||
|
||||
Attributes:
|
||||
host: IP address to bind to. Use `0.0.0.0` to serve to your local network.
|
||||
port: Port to bind to.
|
||||
allow_origins: Allowed CORS origins.
|
||||
allow_credentials: Allow CORS credentials.
|
||||
allow_methods: Methods allowed for CORS.
|
||||
allow_headers: Headers allowed for CORS.
|
||||
ssl_certfile: SSL certificate file for HTTPS. See https://www.uvicorn.dev/settings/#https.
|
||||
ssl_keyfile: SSL key file for HTTPS. See https://www.uvicorn.dev/settings/#https.
|
||||
log_tokenization: Enable logging of parsed prompt tokens.
|
||||
patchmatch: Enable patchmatch inpaint code.
|
||||
models_dir: Path to the models directory.
|
||||
convert_cache_dir: Path to the converted models cache directory (DEPRECATED, but do not delete because it is needed for migration from previous versions).
|
||||
download_cache_dir: Path to the directory that contains dynamically downloaded models.
|
||||
legacy_conf_dir: Path to directory of legacy checkpoint config files.
|
||||
db_dir: Path to InvokeAI databases directory.
|
||||
outputs_dir: Path to directory for outputs.
|
||||
image_subfolder_strategy: Strategy for organizing images into subfolders. 'flat' stores all images in a single folder. 'date' organizes by YYYY/MM/DD. 'type' organizes by image category. 'hash' uses first 2 characters of UUID for filesystem performance.<br>Valid values: `flat`, `date`, `type`, `hash`
|
||||
custom_nodes_dir: Path to directory for custom nodes.
|
||||
style_presets_dir: Path to directory for style presets.
|
||||
workflow_thumbnails_dir: Path to directory for workflow thumbnails.
|
||||
log_handlers: Log handler. Valid options are "console", "file=<path>", "syslog=path|address:host:port", "http=<url>".
|
||||
log_format: Log format. Use "plain" for text-only, "color" for colorized output, "legacy" for 2.3-style logging and "syslog" for syslog-style.<br>Valid values: `plain`, `color`, `syslog`, `legacy`
|
||||
log_level: Emit logging messages at this level or higher.<br>Valid values: `debug`, `info`, `warning`, `error`, `critical`
|
||||
log_sql: Log SQL queries. `log_level` must be `debug` for this to do anything. Extremely verbose.
|
||||
log_level_network: Log level for network-related messages. 'info' and 'debug' are very verbose.<br>Valid values: `debug`, `info`, `warning`, `error`, `critical`
|
||||
use_memory_db: Use in-memory database. Useful for development.
|
||||
dev_reload: Automatically reload when Python sources are changed. Does not reload node definitions.
|
||||
profile_graphs: Enable graph profiling using `cProfile`.
|
||||
profile_prefix: An optional prefix for profile output files.
|
||||
profiles_dir: Path to profiles output directory.
|
||||
max_cache_ram_gb: The maximum amount of CPU RAM to use for model caching in GB. If unset, the limit will be configured based on the available RAM. In most cases, it is recommended to leave this unset.
|
||||
max_cache_vram_gb: The amount of VRAM to use for model caching in GB. If unset, the limit will be configured based on the available VRAM and the device_working_mem_gb. In most cases, it is recommended to leave this unset.
|
||||
log_memory_usage: If True, a memory snapshot will be captured before and after every model cache operation, and the result will be logged (at debug level). There is a time cost to capturing the memory snapshots, so it is recommended to only enable this feature if you are actively inspecting the model cache's behaviour.
|
||||
model_cache_keep_alive_min: How long to keep models in cache after last use, in minutes. A value of 0 (the default) means models are kept in cache indefinitely. If no model generations occur within the timeout period, the model cache is cleared using the same logic as the 'Clear Model Cache' button.
|
||||
device_working_mem_gb: The amount of working memory to keep available on the compute device (in GB). Has no effect if running on CPU. If you are experiencing OOM errors, try increasing this value.
|
||||
enable_partial_loading: Enable partial loading of models. This enables models to run with reduced VRAM requirements (at the cost of slower speed) by streaming the model from RAM to VRAM as its used. In some edge cases, partial loading can cause models to run more slowly if they were previously being fully loaded into VRAM.
|
||||
keep_ram_copy_of_weights: Whether to keep a full RAM copy of a model's weights when the model is loaded in VRAM. Keeping a RAM copy increases average RAM usage, but speeds up model switching and LoRA patching (assuming there is sufficient RAM). Set this to False if RAM pressure is consistently high.
|
||||
ram: DEPRECATED: This setting is no longer used. It has been replaced by `max_cache_ram_gb`, but most users will not need to use this config since automatic cache size limits should work well in most cases. This config setting will be removed once the new model cache behavior is stable.
|
||||
vram: DEPRECATED: This setting is no longer used. It has been replaced by `max_cache_vram_gb`, but most users will not need to use this config since automatic cache size limits should work well in most cases. This config setting will be removed once the new model cache behavior is stable.
|
||||
lazy_offload: DEPRECATED: This setting is no longer used. Lazy-offloading is enabled by default. This config setting will be removed once the new model cache behavior is stable.
|
||||
pytorch_cuda_alloc_conf: Configure the Torch CUDA memory allocator. This will impact peak reserved VRAM usage and performance. Setting to "backend:cudaMallocAsync" works well on many systems. The optimal configuration is highly dependent on the system configuration (device type, VRAM, CUDA driver version, etc.), so must be tuned experimentally.
|
||||
device: Preferred execution device. `auto` will choose the device depending on the hardware platform and the installed torch capabilities.<br>Valid values: `auto`, `cpu`, `cuda`, `mps`, `cuda:N` (where N is a device number)
|
||||
precision: Floating point precision. `float16` will consume half the memory of `float32` but produce slightly lower-quality images. The `auto` setting will guess the proper precision based on your video card and operating system.<br>Valid values: `auto`, `float16`, `bfloat16`, `float32`
|
||||
sequential_guidance: Whether to calculate guidance in serial instead of in parallel, lowering memory requirements.
|
||||
attention_type: Attention type.<br>Valid values: `auto`, `normal`, `xformers`, `sliced`, `torch-sdp`
|
||||
attention_slice_size: Slice size, valid when attention_type=="sliced".<br>Valid values: `auto`, `balanced`, `max`, `1`, `2`, `3`, `4`, `5`, `6`, `7`, `8`
|
||||
force_tiled_decode: Whether to enable tiled VAE decode (reduces memory consumption with some performance penalty).
|
||||
pil_compress_level: The compress_level setting of PIL.Image.save(), used for PNG encoding. All settings are lossless. 0 = no compression, 1 = fastest with slightly larger filesize, 9 = slowest with smallest filesize. 1 is typically the best setting.
|
||||
max_queue_size: Maximum number of items in the session queue.
|
||||
session_queue_mode: Session queue mode. Use 'FIFO' for traditional first-in-first-out, or 'round_robin' to serve each user's jobs in turn. In single-user mode, FIFO is always used regardless of this setting.<br>Valid values: `FIFO`, `round_robin`
|
||||
clear_queue_on_startup: Empties session queue on startup. If true, disables `max_queue_history`.
|
||||
max_queue_history: Keep the last N completed, failed, and canceled queue items. Older items are deleted on startup. Set to 0 to prune all terminal items. Ignored if `clear_queue_on_startup` is true.
|
||||
allow_nodes: List of nodes to allow. Omit to allow all.
|
||||
deny_nodes: List of nodes to deny. Omit to deny none.
|
||||
node_cache_size: How many cached nodes to keep in memory.
|
||||
hashing_algorithm: Model hashing algorthim for model installs. 'blake3_multi' is best for SSDs. 'blake3_single' is best for spinning disk HDDs. 'random' disables hashing, instead assigning a UUID to models. Useful when using a memory db to reduce model installation time, or if you don't care about storing stable hashes for models. Alternatively, any other hashlib algorithm is accepted, though these are not nearly as performant as blake3.<br>Valid values: `blake3_multi`, `blake3_single`, `random`, `md5`, `sha1`, `sha224`, `sha256`, `sha384`, `sha512`, `blake2b`, `blake2s`, `sha3_224`, `sha3_256`, `sha3_384`, `sha3_512`, `shake_128`, `shake_256`
|
||||
remote_api_tokens: List of regular expression and token pairs used when downloading models from URLs. The download URL is tested against the regex, and if it matches, the token is provided in as a Bearer token.
|
||||
scan_models_on_startup: Scan the models directory on startup, registering orphaned models. This is typically only used in conjunction with `use_memory_db` for testing purposes.
|
||||
unsafe_disable_picklescan: UNSAFE. Disable the picklescan security check during model installation. Recommended only for development and testing purposes. This will allow arbitrary code execution during model installation, so should never be used in production.
|
||||
allow_unknown_models: Allow installation of models that we are unable to identify. If enabled, models will be marked as `unknown` in the database, and will not have any metadata associated with them. If disabled, unknown models will be rejected during installation.
|
||||
multiuser: Enable multiuser support. When disabled, the application runs in single-user mode using a default system account with administrator privileges. When enabled, requires user authentication and authorization.
|
||||
strict_password_checking: Enforce strict password requirements. When True, passwords must contain uppercase, lowercase, and numbers. When False (default), any password is accepted but its strength (weak/moderate/strong) is reported to the user.
|
||||
external_alibabacloud_api_key: API key for Alibaba Cloud DashScope image generation.
|
||||
external_alibabacloud_base_url: Base URL override for Alibaba Cloud DashScope image generation.
|
||||
external_gemini_api_key: API key for Gemini image generation.
|
||||
external_openai_api_key: API key for OpenAI image generation.
|
||||
external_gemini_base_url: Base URL override for Gemini image generation.
|
||||
external_openai_base_url: Base URL override for OpenAI image generation.
|
||||
external_seedream_api_key: API key for Seedream image generation.
|
||||
external_seedream_base_url: Base URL override for Seedream image generation.
|
||||
"""
|
||||
|
||||
_root: Optional[Path] = PrivateAttr(default=None)
|
||||
_config_file: Optional[Path] = PrivateAttr(default=None)
|
||||
|
||||
# fmt: off
|
||||
|
||||
# INTERNAL
|
||||
schema_version: str = Field(default=CONFIG_SCHEMA_VERSION, description="Schema version of the config file. This is not a user-configurable setting.")
|
||||
# This is only used during v3 models.yaml migration
|
||||
legacy_models_yaml_path: Optional[Path] = Field(default=None, description="Path to the legacy models.yaml file. This is not a user-configurable setting.")
|
||||
|
||||
# WEB
|
||||
host: str = Field(default="127.0.0.1", description="IP address to bind to. Use `0.0.0.0` to serve to your local network.")
|
||||
port: int = Field(default=9090, description="Port to bind to.")
|
||||
allow_origins: list[str] = Field(default=[], description="Allowed CORS origins.")
|
||||
allow_credentials: bool = Field(default=True, description="Allow CORS credentials.")
|
||||
allow_methods: list[str] = Field(default=["*"], description="Methods allowed for CORS.")
|
||||
allow_headers: list[str] = Field(default=["*"], description="Headers allowed for CORS.")
|
||||
ssl_certfile: Optional[Path] = Field(default=None, description="SSL certificate file for HTTPS. See https://www.uvicorn.dev/settings/#https.")
|
||||
ssl_keyfile: Optional[Path] = Field(default=None, description="SSL key file for HTTPS. See https://www.uvicorn.dev/settings/#https.")
|
||||
|
||||
# MISC FEATURES
|
||||
log_tokenization: bool = Field(default=False, description="Enable logging of parsed prompt tokens.")
|
||||
patchmatch: bool = Field(default=True, description="Enable patchmatch inpaint code.")
|
||||
|
||||
# PATHS
|
||||
models_dir: Path = Field(default=Path("models"), description="Path to the models directory.")
|
||||
convert_cache_dir: Path = Field(default=Path("models/.convert_cache"), description="Path to the converted models cache directory (DEPRECATED, but do not delete because it is needed for migration from previous versions).")
|
||||
download_cache_dir: Path = Field(default=Path("models/.download_cache"), description="Path to the directory that contains dynamically downloaded models.")
|
||||
legacy_conf_dir: Path = Field(default=Path("configs"), description="Path to directory of legacy checkpoint config files.")
|
||||
db_dir: Path = Field(default=Path("databases"), description="Path to InvokeAI databases directory.")
|
||||
outputs_dir: Path = Field(default=Path("outputs"), description="Path to directory for outputs.")
|
||||
image_subfolder_strategy: IMAGE_SUBFOLDER_STRATEGY = Field(default="flat", description="Strategy for organizing images into subfolders. 'flat' stores all images in a single folder. 'date' organizes by YYYY/MM/DD. 'type' organizes by image category. 'hash' uses first 2 characters of UUID for filesystem performance.")
|
||||
custom_nodes_dir: Path = Field(default=Path("nodes"), description="Path to directory for custom nodes.")
|
||||
style_presets_dir: Path = Field(default=Path("style_presets"), description="Path to directory for style presets.")
|
||||
workflow_thumbnails_dir: Path = Field(default=Path("workflow_thumbnails"), description="Path to directory for workflow thumbnails.")
|
||||
|
||||
# LOGGING
|
||||
log_handlers: list[str] = Field(default=["console"], description='Log handler. Valid options are "console", "file=<path>", "syslog=path|address:host:port", "http=<url>".')
|
||||
# note - would be better to read the log_format values from logging.py, but this creates circular dependencies issues
|
||||
log_format: LOG_FORMAT = Field(default="color", description='Log format. Use "plain" for text-only, "color" for colorized output, "legacy" for 2.3-style logging and "syslog" for syslog-style.')
|
||||
log_level: LOG_LEVEL = Field(default="info", description="Emit logging messages at this level or higher.")
|
||||
log_sql: bool = Field(default=False, description="Log SQL queries. `log_level` must be `debug` for this to do anything. Extremely verbose.")
|
||||
log_level_network: LOG_LEVEL = Field(default='warning', description="Log level for network-related messages. 'info' and 'debug' are very verbose.")
|
||||
|
||||
# Development
|
||||
use_memory_db: bool = Field(default=False, description="Use in-memory database. Useful for development.")
|
||||
dev_reload: bool = Field(default=False, description="Automatically reload when Python sources are changed. Does not reload node definitions.")
|
||||
profile_graphs: bool = Field(default=False, description="Enable graph profiling using `cProfile`.")
|
||||
profile_prefix: Optional[str] = Field(default=None, description="An optional prefix for profile output files.")
|
||||
profiles_dir: Path = Field(default=Path("profiles"), description="Path to profiles output directory.")
|
||||
|
||||
# CACHE
|
||||
max_cache_ram_gb: Optional[float] = Field(default=None, gt=0, description="The maximum amount of CPU RAM to use for model caching in GB. If unset, the limit will be configured based on the available RAM. In most cases, it is recommended to leave this unset.")
|
||||
max_cache_vram_gb: Optional[float] = Field(default=None, ge=0, description="The amount of VRAM to use for model caching in GB. If unset, the limit will be configured based on the available VRAM and the device_working_mem_gb. In most cases, it is recommended to leave this unset.")
|
||||
log_memory_usage: bool = Field(default=False, description="If True, a memory snapshot will be captured before and after every model cache operation, and the result will be logged (at debug level). There is a time cost to capturing the memory snapshots, so it is recommended to only enable this feature if you are actively inspecting the model cache's behaviour.")
|
||||
model_cache_keep_alive_min: float = Field(default=0, ge=0, description="How long to keep models in cache after last use, in minutes. A value of 0 (the default) means models are kept in cache indefinitely. If no model generations occur within the timeout period, the model cache is cleared using the same logic as the 'Clear Model Cache' button.")
|
||||
device_working_mem_gb: float = Field(default=3, description="The amount of working memory to keep available on the compute device (in GB). Has no effect if running on CPU. If you are experiencing OOM errors, try increasing this value.")
|
||||
enable_partial_loading: bool = Field(default=True, description="Enable partial loading of models. This enables models to run with reduced VRAM requirements (at the cost of slower speed) by streaming the model from RAM to VRAM as its used. In some edge cases, partial loading can cause models to run more slowly if they were previously being fully loaded into VRAM.")
|
||||
keep_ram_copy_of_weights: bool = Field(default=True, description="Whether to keep a full RAM copy of a model's weights when the model is loaded in VRAM. Keeping a RAM copy increases average RAM usage, but speeds up model switching and LoRA patching (assuming there is sufficient RAM). Set this to False if RAM pressure is consistently high.")
|
||||
# Deprecated CACHE configs
|
||||
ram: Optional[float] = Field(default=None, gt=0, description="DEPRECATED: This setting is no longer used. It has been replaced by `max_cache_ram_gb`, but most users will not need to use this config since automatic cache size limits should work well in most cases. This config setting will be removed once the new model cache behavior is stable.")
|
||||
vram: Optional[float] = Field(default=None, ge=0, description="DEPRECATED: This setting is no longer used. It has been replaced by `max_cache_vram_gb`, but most users will not need to use this config since automatic cache size limits should work well in most cases. This config setting will be removed once the new model cache behavior is stable.")
|
||||
lazy_offload: bool = Field(default=True, description="DEPRECATED: This setting is no longer used. Lazy-offloading is enabled by default. This config setting will be removed once the new model cache behavior is stable.")
|
||||
|
||||
# PyTorch Memory Allocator
|
||||
pytorch_cuda_alloc_conf: Optional[str] = Field(default=None, description="Configure the Torch CUDA memory allocator. This will impact peak reserved VRAM usage and performance. Setting to \"backend:cudaMallocAsync\" works well on many systems. The optimal configuration is highly dependent on the system configuration (device type, VRAM, CUDA driver version, etc.), so must be tuned experimentally.")
|
||||
|
||||
# DEVICE
|
||||
device: str = Field(default="auto", description="Preferred execution device. `auto` will choose the device depending on the hardware platform and the installed torch capabilities.<br>Valid values: `auto`, `cpu`, `cuda`, `mps`, `cuda:N` (where N is a device number)", pattern=r"^(auto|cpu|mps|cuda(:\d+)?)$")
|
||||
precision: PRECISION = Field(default="auto", description="Floating point precision. `float16` will consume half the memory of `float32` but produce slightly lower-quality images. The `auto` setting will guess the proper precision based on your video card and operating system.")
|
||||
|
||||
# GENERATION
|
||||
sequential_guidance: bool = Field(default=False, description="Whether to calculate guidance in serial instead of in parallel, lowering memory requirements.")
|
||||
attention_type: ATTENTION_TYPE = Field(default="auto", description="Attention type.")
|
||||
attention_slice_size: ATTENTION_SLICE_SIZE = Field(default="auto", description='Slice size, valid when attention_type=="sliced".')
|
||||
force_tiled_decode: bool = Field(default=False, description="Whether to enable tiled VAE decode (reduces memory consumption with some performance penalty).")
|
||||
pil_compress_level: int = Field(default=1, description="The compress_level setting of PIL.Image.save(), used for PNG encoding. All settings are lossless. 0 = no compression, 1 = fastest with slightly larger filesize, 9 = slowest with smallest filesize. 1 is typically the best setting.")
|
||||
max_queue_size: int = Field(default=10000, gt=0, description="Maximum number of items in the session queue.")
|
||||
session_queue_mode: SESSION_QUEUE_MODE = Field(default="round_robin", description="Session queue mode. Use 'FIFO' for traditional first-in-first-out, or 'round_robin' to serve each user's jobs in turn. In single-user mode, FIFO is always used regardless of this setting.")
|
||||
clear_queue_on_startup: bool = Field(default=False, description="Empties session queue on startup. If true, disables `max_queue_history`.")
|
||||
max_queue_history: Optional[int] = Field(default=None, ge=0, description="Keep the last N completed, failed, and canceled queue items. Older items are deleted on startup. Set to 0 to prune all terminal items. Ignored if `clear_queue_on_startup` is true.")
|
||||
|
||||
# NODES
|
||||
allow_nodes: Optional[list[str]] = Field(default=None, description="List of nodes to allow. Omit to allow all.")
|
||||
deny_nodes: Optional[list[str]] = Field(default=None, description="List of nodes to deny. Omit to deny none.")
|
||||
node_cache_size: int = Field(default=512, description="How many cached nodes to keep in memory.")
|
||||
|
||||
# MODEL INSTALL
|
||||
hashing_algorithm: HASHING_ALGORITHMS = Field(default="blake3_single", description="Model hashing algorthim for model installs. 'blake3_multi' is best for SSDs. 'blake3_single' is best for spinning disk HDDs. 'random' disables hashing, instead assigning a UUID to models. Useful when using a memory db to reduce model installation time, or if you don't care about storing stable hashes for models. Alternatively, any other hashlib algorithm is accepted, though these are not nearly as performant as blake3.")
|
||||
remote_api_tokens: Optional[list[URLRegexTokenPair]] = Field(default=None, description="List of regular expression and token pairs used when downloading models from URLs. The download URL is tested against the regex, and if it matches, the token is provided in as a Bearer token.")
|
||||
scan_models_on_startup: bool = Field(default=False, description="Scan the models directory on startup, registering orphaned models. This is typically only used in conjunction with `use_memory_db` for testing purposes.")
|
||||
unsafe_disable_picklescan: bool = Field(default=False, description="UNSAFE. Disable the picklescan security check during model installation. Recommended only for development and testing purposes. This will allow arbitrary code execution during model installation, so should never be used in production.")
|
||||
allow_unknown_models: bool = Field(default=True, description="Allow installation of models that we are unable to identify. If enabled, models will be marked as `unknown` in the database, and will not have any metadata associated with them. If disabled, unknown models will be rejected during installation.")
|
||||
|
||||
# MULTIUSER
|
||||
multiuser: bool = Field(default=False, description="Enable multiuser support. When disabled, the application runs in single-user mode using a default system account with administrator privileges. When enabled, requires user authentication and authorization.")
|
||||
strict_password_checking: bool = Field(default=False, description="Enforce strict password requirements. When True, passwords must contain uppercase, lowercase, and numbers. When False (default), any password is accepted but its strength (weak/moderate/strong) is reported to the user.")
|
||||
|
||||
# EXTERNAL PROVIDERS
|
||||
external_alibabacloud_api_key: Optional[str] = Field(default=None, description="API key for Alibaba Cloud DashScope image generation.")
|
||||
external_alibabacloud_base_url: Optional[str] = Field(
|
||||
default=None, description="Base URL override for Alibaba Cloud DashScope image generation."
|
||||
)
|
||||
external_gemini_api_key: Optional[str] = Field(default=None, description="API key for Gemini image generation.")
|
||||
external_openai_api_key: Optional[str] = Field(default=None, description="API key for OpenAI image generation.")
|
||||
external_gemini_base_url: Optional[str] = Field(
|
||||
default=None, description="Base URL override for Gemini image generation."
|
||||
)
|
||||
external_openai_base_url: Optional[str] = Field(
|
||||
default=None, description="Base URL override for OpenAI image generation."
|
||||
)
|
||||
external_seedream_api_key: Optional[str] = Field(
|
||||
default=None, description="API key for Seedream image generation."
|
||||
)
|
||||
external_seedream_base_url: Optional[str] = Field(
|
||||
default=None, description="Base URL override for Seedream image generation."
|
||||
)
|
||||
|
||||
# fmt: on
|
||||
|
||||
model_config = SettingsConfigDict(env_prefix="INVOKEAI_", env_ignore_empty=True)
|
||||
|
||||
def update_config(self, config: dict[str, Any] | InvokeAIAppConfig, clobber: bool = True) -> None:
|
||||
"""Updates the config, overwriting existing values.
|
||||
|
||||
Args:
|
||||
config: A dictionary of config settings, or instance of `InvokeAIAppConfig`. If an instance of \
|
||||
`InvokeAIAppConfig`, only the explicitly set fields will be merged into the singleton config.
|
||||
clobber: If `True`, overwrite existing values. If `False`, only update fields that are not already set.
|
||||
"""
|
||||
|
||||
if isinstance(config, dict):
|
||||
new_config = self.model_validate(config)
|
||||
else:
|
||||
new_config = config
|
||||
|
||||
for field_name in new_config.model_fields_set:
|
||||
new_value = getattr(new_config, field_name)
|
||||
current_value = getattr(self, field_name)
|
||||
|
||||
if field_name in self.model_fields_set and not clobber:
|
||||
continue
|
||||
|
||||
if new_value != current_value:
|
||||
setattr(self, field_name, new_value)
|
||||
|
||||
def write_file(self, dest_path: Path, as_example: bool = False) -> None:
|
||||
"""Write the current configuration to file. This will overwrite the existing file.
|
||||
|
||||
A `meta` stanza is added to the top of the file, containing metadata about the config file. This is not stored in the config object.
|
||||
|
||||
Args:
|
||||
dest_path: Path to write the config to.
|
||||
"""
|
||||
dest_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
with open(dest_path, "w") as file:
|
||||
# Meta fields should be written in a separate stanza - skip legacy_models_yaml_path
|
||||
meta_dict = self.model_dump(mode="json", include={"schema_version"})
|
||||
|
||||
# User settings
|
||||
config_dict = self.model_dump(
|
||||
mode="json",
|
||||
exclude_unset=False if as_example else True,
|
||||
exclude_defaults=False if as_example else True,
|
||||
exclude_none=True if as_example else False,
|
||||
exclude={"schema_version", "legacy_models_yaml_path"},
|
||||
)
|
||||
|
||||
if as_example:
|
||||
file.write("# This is an example file with default and example settings.\n")
|
||||
file.write("# You should not copy this whole file into your config.\n")
|
||||
file.write("# Only add the settings you need to change to your config file.\n\n")
|
||||
file.write("# Internal metadata - do not edit:\n")
|
||||
file.write(yaml.dump(meta_dict, sort_keys=False))
|
||||
file.write("\n")
|
||||
file.write("# Put user settings here - see https://invoke.ai/configuration/invokeai-yaml/:\n")
|
||||
if len(config_dict) > 0:
|
||||
file.write(yaml.dump(config_dict, sort_keys=False))
|
||||
|
||||
def _resolve(self, partial_path: Path) -> Path:
|
||||
return (self.root_path / partial_path).resolve()
|
||||
|
||||
@property
|
||||
def root_path(self) -> Path:
|
||||
"""Path to the runtime root directory, resolved to an absolute path."""
|
||||
if self._root:
|
||||
root = Path(self._root).expanduser().absolute()
|
||||
else:
|
||||
root = self.find_root().expanduser().absolute()
|
||||
self._root = root # insulate ourselves from relative paths that may change
|
||||
return root.resolve()
|
||||
|
||||
@property
|
||||
def config_file_path(self) -> Path:
|
||||
"""Path to invokeai.yaml, resolved to an absolute path.."""
|
||||
resolved_path = self._resolve(self._config_file or INIT_FILE)
|
||||
assert resolved_path is not None
|
||||
return resolved_path
|
||||
|
||||
@property
|
||||
def api_keys_file_path(self) -> Path:
|
||||
"""Path to api_keys.yaml, resolved to an absolute path.."""
|
||||
resolved_path = self._resolve(API_KEYS_FILE)
|
||||
assert resolved_path is not None
|
||||
return resolved_path
|
||||
|
||||
@property
|
||||
def outputs_path(self) -> Optional[Path]:
|
||||
"""Path to the outputs directory, resolved to an absolute path.."""
|
||||
return self._resolve(self.outputs_dir)
|
||||
|
||||
@property
|
||||
def db_path(self) -> Path:
|
||||
"""Path to the invokeai.db file, resolved to an absolute path.."""
|
||||
db_dir = self._resolve(self.db_dir)
|
||||
assert db_dir is not None
|
||||
return db_dir / DB_FILE
|
||||
|
||||
@property
|
||||
def legacy_conf_path(self) -> Path:
|
||||
"""Path to directory of legacy configuration files (e.g. v1-inference.yaml), resolved to an absolute path.."""
|
||||
return self._resolve(self.legacy_conf_dir)
|
||||
|
||||
@property
|
||||
def models_path(self) -> Path:
|
||||
"""Path to the models directory, resolved to an absolute path.."""
|
||||
return self._resolve(self.models_dir)
|
||||
|
||||
@property
|
||||
def style_presets_path(self) -> Path:
|
||||
"""Path to the style presets directory, resolved to an absolute path.."""
|
||||
return self._resolve(self.style_presets_dir)
|
||||
|
||||
@property
|
||||
def workflow_thumbnails_path(self) -> Path:
|
||||
"""Path to the workflow thumbnails directory, resolved to an absolute path.."""
|
||||
return self._resolve(self.workflow_thumbnails_dir)
|
||||
|
||||
@property
|
||||
def convert_cache_path(self) -> Path:
|
||||
"""Path to the converted cache models directory, resolved to an absolute path.."""
|
||||
return self._resolve(self.convert_cache_dir)
|
||||
|
||||
@property
|
||||
def download_cache_path(self) -> Path:
|
||||
"""Path to the downloaded models directory, resolved to an absolute path.."""
|
||||
return self._resolve(self.download_cache_dir)
|
||||
|
||||
@property
|
||||
def custom_nodes_path(self) -> Path:
|
||||
"""Path to the custom nodes directory, resolved to an absolute path.."""
|
||||
custom_nodes_path = self._resolve(self.custom_nodes_dir)
|
||||
assert custom_nodes_path is not None
|
||||
return custom_nodes_path
|
||||
|
||||
@property
|
||||
def profiles_path(self) -> Path:
|
||||
"""Path to the graph profiles directory, resolved to an absolute path.."""
|
||||
return self._resolve(self.profiles_dir)
|
||||
|
||||
@staticmethod
|
||||
def find_root() -> Path:
|
||||
"""Choose the runtime root directory when not specified on command line or init file."""
|
||||
if os.environ.get("INVOKEAI_ROOT"):
|
||||
root = Path(os.environ["INVOKEAI_ROOT"])
|
||||
elif venv := os.environ.get("VIRTUAL_ENV", None):
|
||||
root = Path(venv).parent.resolve()
|
||||
else:
|
||||
root = Path("~/invokeai").expanduser().resolve()
|
||||
return root
|
||||
|
||||
|
||||
class DefaultInvokeAIAppConfig(InvokeAIAppConfig):
|
||||
"""A version of `InvokeAIAppConfig` that does not automatically parse any settings from environment variables
|
||||
or any file.
|
||||
|
||||
This is useful for writing out a default config file.
|
||||
|
||||
Note that init settings are set if provided.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def settings_customise_sources(
|
||||
cls,
|
||||
settings_cls: type[BaseSettings],
|
||||
init_settings: PydanticBaseSettingsSource,
|
||||
env_settings: PydanticBaseSettingsSource,
|
||||
dotenv_settings: PydanticBaseSettingsSource,
|
||||
file_secret_settings: PydanticBaseSettingsSource,
|
||||
) -> tuple[PydanticBaseSettingsSource, ...]:
|
||||
return (init_settings,)
|
||||
|
||||
|
||||
def migrate_v3_config_dict(config_dict: dict[str, Any]) -> dict[str, Any]:
|
||||
"""Migrate a v3 config dictionary to a v4.0.0.
|
||||
|
||||
Args:
|
||||
config_dict: A dictionary of settings from a v3 config file.
|
||||
|
||||
Returns:
|
||||
An `InvokeAIAppConfig` config dict.
|
||||
|
||||
"""
|
||||
parsed_config_dict: dict[str, Any] = {}
|
||||
for _category_name, category_dict in config_dict["InvokeAI"].items():
|
||||
for k, v in category_dict.items():
|
||||
# `outdir` was renamed to `outputs_dir` in v4
|
||||
if k == "outdir":
|
||||
parsed_config_dict["outputs_dir"] = v
|
||||
# `max_cache_size` was renamed to `ram` some time in v3, but both names were used
|
||||
if k == "max_cache_size" and "ram" not in category_dict:
|
||||
parsed_config_dict["ram"] = v
|
||||
# `max_vram_cache_size` was renamed to `vram` some time in v3, but both names were used
|
||||
if k == "max_vram_cache_size" and "vram" not in category_dict:
|
||||
parsed_config_dict["vram"] = v
|
||||
# autocast was removed in v4.0.1
|
||||
if k == "precision" and v == "autocast":
|
||||
parsed_config_dict["precision"] = "auto"
|
||||
if k == "conf_path":
|
||||
parsed_config_dict["legacy_models_yaml_path"] = v
|
||||
if k == "legacy_conf_dir":
|
||||
# The old default for this was "configs/stable-diffusion" ("configs\stable-diffusion" on Windows).
|
||||
if v == "configs/stable-diffusion" or v == "configs\\stable-diffusion":
|
||||
# If if the incoming config has the default value, skip
|
||||
continue
|
||||
elif Path(v).name == "stable-diffusion":
|
||||
# Else if the path ends in "stable-diffusion", we assume the parent is the new correct path.
|
||||
parsed_config_dict["legacy_conf_dir"] = str(Path(v).parent)
|
||||
else:
|
||||
# Else we do not attempt to migrate this setting
|
||||
parsed_config_dict["legacy_conf_dir"] = v
|
||||
elif k in InvokeAIAppConfig.model_fields:
|
||||
# skip unknown fields
|
||||
parsed_config_dict[k] = v
|
||||
parsed_config_dict["schema_version"] = "4.0.0"
|
||||
return parsed_config_dict
|
||||
|
||||
|
||||
def migrate_v4_0_0_to_4_0_1_config_dict(config_dict: dict[str, Any]) -> dict[str, Any]:
|
||||
"""Migrate v4.0.0 config dictionary to a v4.0.1 config dictionary
|
||||
|
||||
Args:
|
||||
config_dict: A dictionary of settings from a v4.0.0 config file.
|
||||
|
||||
Returns:
|
||||
A config dict with the settings migrated to v4.0.1.
|
||||
"""
|
||||
parsed_config_dict: dict[str, Any] = copy.deepcopy(config_dict)
|
||||
# precision "autocast" was replaced by "auto" in v4.0.1
|
||||
if parsed_config_dict.get("precision") == "autocast":
|
||||
parsed_config_dict["precision"] = "auto"
|
||||
parsed_config_dict["schema_version"] = "4.0.1"
|
||||
return parsed_config_dict
|
||||
|
||||
|
||||
def migrate_v4_0_1_to_4_0_2_config_dict(config_dict: dict[str, Any]) -> dict[str, Any]:
|
||||
"""Migrate v4.0.1 config dictionary to a v4.0.2 config dictionary.
|
||||
|
||||
Args:
|
||||
config_dict: A dictionary of settings from a v4.0.1 config file.
|
||||
|
||||
Returns:
|
||||
An config dict with the settings migrated to v4.0.2.
|
||||
"""
|
||||
parsed_config_dict: dict[str, Any] = copy.deepcopy(config_dict)
|
||||
# convert_cache was removed in 4.0.2
|
||||
parsed_config_dict.pop("convert_cache", None)
|
||||
parsed_config_dict["schema_version"] = "4.0.2"
|
||||
return parsed_config_dict
|
||||
|
||||
|
||||
def migrate_v4_0_2_to_4_0_3_config_dict(config_dict: dict[str, Any]) -> dict[str, Any]:
|
||||
"""Migrate v4.0.2 config dictionary to a v4.0.3 config dictionary.
|
||||
|
||||
Args:
|
||||
config_dict: A dictionary of settings from a v4.0.2 config file.
|
||||
|
||||
Returns:
|
||||
A config dict with the settings migrated to v4.0.3.
|
||||
"""
|
||||
parsed_config_dict: dict[str, Any] = copy.deepcopy(config_dict)
|
||||
parsed_config_dict["schema_version"] = "4.0.3"
|
||||
return parsed_config_dict
|
||||
|
||||
|
||||
def load_and_migrate_config(config_path: Path) -> InvokeAIAppConfig:
|
||||
"""Load and migrate a config file to the latest version.
|
||||
|
||||
Args:
|
||||
config_path: Path to the config file.
|
||||
|
||||
Returns:
|
||||
An instance of `InvokeAIAppConfig` with the loaded and migrated settings.
|
||||
"""
|
||||
assert config_path.suffix == ".yaml"
|
||||
with open(config_path, "rt", encoding=locale.getpreferredencoding()) as file:
|
||||
loaded_config_dict: dict[str, Any] = yaml.safe_load(file)
|
||||
|
||||
assert isinstance(loaded_config_dict, dict)
|
||||
|
||||
migrated = False
|
||||
if "InvokeAI" in loaded_config_dict:
|
||||
migrated = True
|
||||
loaded_config_dict = migrate_v3_config_dict(loaded_config_dict) # pyright: ignore [reportUnknownArgumentType]
|
||||
if loaded_config_dict["schema_version"] == "4.0.0":
|
||||
migrated = True
|
||||
loaded_config_dict = migrate_v4_0_0_to_4_0_1_config_dict(loaded_config_dict)
|
||||
if loaded_config_dict["schema_version"] == "4.0.1":
|
||||
migrated = True
|
||||
loaded_config_dict = migrate_v4_0_1_to_4_0_2_config_dict(loaded_config_dict)
|
||||
if loaded_config_dict["schema_version"] == "4.0.2":
|
||||
migrated = True
|
||||
loaded_config_dict = migrate_v4_0_2_to_4_0_3_config_dict(loaded_config_dict)
|
||||
|
||||
if migrated:
|
||||
shutil.copy(config_path, config_path.with_suffix(".yaml.bak"))
|
||||
try:
|
||||
# load and write without environment variables
|
||||
migrated_config = DefaultInvokeAIAppConfig.model_validate(loaded_config_dict)
|
||||
migrated_config.write_file(config_path)
|
||||
except Exception as e:
|
||||
shutil.copy(config_path.with_suffix(".yaml.bak"), config_path)
|
||||
raise RuntimeError(f"Failed to load and migrate v3 config file {config_path}: {e}") from e
|
||||
|
||||
try:
|
||||
# Meta is not included in the model fields, so we need to validate it separately
|
||||
config = InvokeAIAppConfig.model_validate(loaded_config_dict)
|
||||
assert config.schema_version == CONFIG_SCHEMA_VERSION, (
|
||||
f"Invalid schema version, expected {CONFIG_SCHEMA_VERSION}: {config.schema_version}"
|
||||
)
|
||||
return config
|
||||
except Exception as e:
|
||||
raise RuntimeError(f"Failed to load config file {config_path}: {e}") from e
|
||||
|
||||
|
||||
def load_external_api_keys(api_keys_file_path: Path) -> dict[str, str]:
|
||||
"""Load external provider config (API keys and base URLs) from a dedicated YAML file."""
|
||||
if not api_keys_file_path.exists():
|
||||
return {}
|
||||
|
||||
with open(api_keys_file_path, "rt", encoding=locale.getpreferredencoding()) as file:
|
||||
loaded_api_keys: Any = yaml.safe_load(file)
|
||||
|
||||
if loaded_api_keys is None:
|
||||
return {}
|
||||
|
||||
if not isinstance(loaded_api_keys, dict):
|
||||
raise RuntimeError(f"Failed to load api keys file {api_keys_file_path}: expected a mapping")
|
||||
|
||||
parsed_api_keys: dict[str, str] = {}
|
||||
for field_name in EXTERNAL_PROVIDER_CONFIG_FIELDS:
|
||||
value = loaded_api_keys.get(field_name)
|
||||
if value is None:
|
||||
continue
|
||||
if not isinstance(value, str):
|
||||
raise RuntimeError(
|
||||
f"Failed to load api keys file {api_keys_file_path}: value for '{field_name}' must be a string"
|
||||
)
|
||||
stripped_value = value.strip()
|
||||
if stripped_value:
|
||||
parsed_api_keys[field_name] = stripped_value
|
||||
|
||||
return parsed_api_keys
|
||||
|
||||
|
||||
@lru_cache(maxsize=1)
|
||||
def get_config() -> InvokeAIAppConfig:
|
||||
"""Get the global singleton app config.
|
||||
|
||||
When first called, this function:
|
||||
- Creates a config object. `pydantic-settings` handles merging of settings from environment variables, but not the init file.
|
||||
- Retrieves any provided CLI args from the InvokeAIArgs class. It does not _parse_ the CLI args; that is done in the main entrypoint.
|
||||
- Sets the root dir, if provided via CLI args.
|
||||
- Logs in to HF if there is no valid token already.
|
||||
- Copies all legacy configs to the legacy conf dir (needed for conversion from ckpt to diffusers).
|
||||
- Reads and merges in settings from the config file if it exists, else writes out a default config file.
|
||||
|
||||
On subsequent calls, the object is returned from the cache.
|
||||
"""
|
||||
# This object includes environment variables, as parsed by pydantic-settings
|
||||
config = InvokeAIAppConfig()
|
||||
env_fields_set = set(config.model_fields_set)
|
||||
|
||||
args = InvokeAIArgs.args
|
||||
|
||||
# This flag serves as a proxy for whether the config was retrieved in the context of the full application or not.
|
||||
# If it is False, we should just return a default config and not set the root, log in to HF, etc.
|
||||
if not InvokeAIArgs.did_parse:
|
||||
return config
|
||||
|
||||
# Set CLI args
|
||||
if root := getattr(args, "root", None):
|
||||
config._root = Path(root)
|
||||
if config_file := getattr(args, "config_file", None):
|
||||
config._config_file = Path(config_file)
|
||||
|
||||
# Create the example config file, with some extra example values provided
|
||||
example_config = DefaultInvokeAIAppConfig()
|
||||
example_config.remote_api_tokens = [
|
||||
URLRegexTokenPair(url_regex="cool-models.com", token="my_secret_token"),
|
||||
URLRegexTokenPair(url_regex="nifty-models.com", token="some_other_token"),
|
||||
]
|
||||
example_config.write_file(config.config_file_path.with_suffix(".example.yaml"), as_example=True)
|
||||
|
||||
# Copy all legacy configs only if needed
|
||||
# We know `__path__[0]` is correct here
|
||||
configs_src = Path(model_configs.__path__[0]) # pyright: ignore [reportUnknownMemberType, reportUnknownArgumentType, reportAttributeAccessIssue]
|
||||
dest_path = config.legacy_conf_path
|
||||
|
||||
# Create destination (we don't need to check for existence)
|
||||
dest_path.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# Compare directories recursively
|
||||
comparison = filecmp.dircmp(configs_src, dest_path)
|
||||
need_copy = any(
|
||||
[
|
||||
comparison.left_only, # Files exist only in source
|
||||
comparison.diff_files, # Files that differ
|
||||
comparison.common_funny, # Files that couldn't be compared
|
||||
]
|
||||
)
|
||||
|
||||
if need_copy:
|
||||
# Get permissions from destination directory
|
||||
dest_mode = dest_path.stat().st_mode
|
||||
|
||||
# Copy directory tree
|
||||
shutil.copytree(configs_src, dest_path, dirs_exist_ok=True)
|
||||
|
||||
# Set permissions on copied files to match destination directory
|
||||
dest_path.chmod(dest_mode)
|
||||
for p in dest_path.glob("**/*"):
|
||||
p.chmod(dest_mode)
|
||||
|
||||
if config.config_file_path.exists():
|
||||
config_from_file = load_and_migrate_config(config.config_file_path)
|
||||
# Clobbering here will overwrite any settings that were set via environment variables
|
||||
config.update_config(config_from_file, clobber=False)
|
||||
else:
|
||||
# We should never write env vars to the config file
|
||||
default_config = DefaultInvokeAIAppConfig()
|
||||
default_config.write_file(config.config_file_path, as_example=False)
|
||||
|
||||
api_keys_from_file = load_external_api_keys(config.api_keys_file_path)
|
||||
if api_keys_from_file:
|
||||
# API keys file should take precedence over invokeai.yaml, but not over environment variables.
|
||||
api_keys_to_apply = {key: value for key, value in api_keys_from_file.items() if key not in env_fields_set}
|
||||
if api_keys_to_apply:
|
||||
config.update_config(api_keys_to_apply, clobber=True)
|
||||
|
||||
return config
|
||||
@@ -0,0 +1,20 @@
|
||||
"""Init file for download queue."""
|
||||
|
||||
from invokeai.app.services.download.download_base import (
|
||||
DownloadJob,
|
||||
DownloadJobStatus,
|
||||
DownloadQueueServiceBase,
|
||||
MultiFileDownloadJob,
|
||||
UnknownJobIDException,
|
||||
)
|
||||
from invokeai.app.services.download.download_default import DownloadQueueService, TqdmProgress
|
||||
|
||||
__all__ = [
|
||||
"DownloadJob",
|
||||
"MultiFileDownloadJob",
|
||||
"DownloadQueueServiceBase",
|
||||
"DownloadQueueService",
|
||||
"TqdmProgress",
|
||||
"DownloadJobStatus",
|
||||
"UnknownJobIDException",
|
||||
]
|
||||
@@ -0,0 +1,368 @@
|
||||
# Copyright (c) 2023 Lincoln D. Stein and the InvokeAI Development Team
|
||||
"""Model download service."""
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
from enum import Enum
|
||||
from functools import total_ordering
|
||||
from pathlib import Path
|
||||
from typing import Any, Callable, List, Optional, Set, Union
|
||||
|
||||
from pydantic import BaseModel, Field, PrivateAttr
|
||||
from pydantic.networks import AnyHttpUrl
|
||||
|
||||
from invokeai.backend.model_manager.metadata import RemoteModelFile
|
||||
|
||||
|
||||
class DownloadJobStatus(str, Enum):
|
||||
"""State of a download job."""
|
||||
|
||||
WAITING = "waiting" # not enqueued, will not run
|
||||
RUNNING = "running" # actively downloading
|
||||
PAUSED = "paused" # paused, can be resumed
|
||||
COMPLETED = "completed" # finished running
|
||||
CANCELLED = "cancelled" # user cancelled
|
||||
ERROR = "error" # terminated with an error message
|
||||
|
||||
|
||||
class DownloadJobCancelledException(Exception):
|
||||
"""This exception is raised when a download job is cancelled."""
|
||||
|
||||
|
||||
class UnknownJobIDException(Exception):
|
||||
"""This exception is raised when an invalid job id is referened."""
|
||||
|
||||
|
||||
class ServiceInactiveException(Exception):
|
||||
"""This exception is raised when user attempts to initiate a download before the service is started."""
|
||||
|
||||
|
||||
SingleFileDownloadEventHandler = Callable[["DownloadJob"], None]
|
||||
SingleFileDownloadExceptionHandler = Callable[["DownloadJob", Optional[Exception]], None]
|
||||
MultiFileDownloadEventHandler = Callable[["MultiFileDownloadJob"], None]
|
||||
MultiFileDownloadExceptionHandler = Callable[["MultiFileDownloadJob", Optional[Exception]], None]
|
||||
DownloadEventHandler = Union[SingleFileDownloadEventHandler, MultiFileDownloadEventHandler]
|
||||
DownloadExceptionHandler = Union[SingleFileDownloadExceptionHandler, MultiFileDownloadExceptionHandler]
|
||||
|
||||
|
||||
class DownloadJobBase(BaseModel):
|
||||
"""Base of classes to monitor and control downloads."""
|
||||
|
||||
# automatically assigned on creation
|
||||
id: int = Field(description="Numeric ID of this job", default=-1) # default id is a sentinel
|
||||
|
||||
dest: Path = Field(description="Initial destination of downloaded model on local disk; a directory or file path")
|
||||
download_path: Optional[Path] = Field(default=None, description="Final location of downloaded file or directory")
|
||||
status: DownloadJobStatus = Field(default=DownloadJobStatus.WAITING, description="Status of the download")
|
||||
bytes: int = Field(default=0, description="Bytes downloaded so far")
|
||||
total_bytes: int = Field(default=0, description="Total file size (bytes)")
|
||||
|
||||
# set when an error occurs
|
||||
error_type: Optional[str] = Field(default=None, description="Name of exception that caused an error")
|
||||
error: Optional[str] = Field(default=None, description="Traceback of the exception that caused an error")
|
||||
|
||||
# internal flag
|
||||
_cancelled: bool = PrivateAttr(default=False)
|
||||
_paused: bool = PrivateAttr(default=False)
|
||||
|
||||
# optional event handlers passed in on creation
|
||||
_on_start: Optional[DownloadEventHandler] = PrivateAttr(default=None)
|
||||
_on_progress: Optional[DownloadEventHandler] = PrivateAttr(default=None)
|
||||
_on_complete: Optional[DownloadEventHandler] = PrivateAttr(default=None)
|
||||
_on_cancelled: Optional[DownloadEventHandler] = PrivateAttr(default=None)
|
||||
_on_error: Optional[DownloadExceptionHandler] = PrivateAttr(default=None)
|
||||
|
||||
def cancel(self) -> None:
|
||||
"""Call to cancel the job."""
|
||||
self._cancelled = True
|
||||
self._paused = False
|
||||
|
||||
def pause(self) -> None:
|
||||
"""Pause the job, preserving partial downloads."""
|
||||
self._paused = True
|
||||
self._cancelled = True
|
||||
|
||||
# cancelled and the callbacks are private attributes in order to prevent
|
||||
# them from being serialized and/or used in the Json Schema
|
||||
@property
|
||||
def cancelled(self) -> bool:
|
||||
"""Call to cancel the job."""
|
||||
return self._cancelled
|
||||
|
||||
@property
|
||||
def paused(self) -> bool:
|
||||
"""Return true if job is paused."""
|
||||
return self._paused
|
||||
|
||||
@property
|
||||
def complete(self) -> bool:
|
||||
"""Return true if job completed without errors."""
|
||||
return self.status == DownloadJobStatus.COMPLETED
|
||||
|
||||
@property
|
||||
def waiting(self) -> bool:
|
||||
"""Return true if the job is waiting to run."""
|
||||
return self.status == DownloadJobStatus.WAITING
|
||||
|
||||
@property
|
||||
def running(self) -> bool:
|
||||
"""Return true if the job is running."""
|
||||
return self.status == DownloadJobStatus.RUNNING
|
||||
|
||||
@property
|
||||
def errored(self) -> bool:
|
||||
"""Return true if the job is errored."""
|
||||
return self.status == DownloadJobStatus.ERROR
|
||||
|
||||
@property
|
||||
def in_terminal_state(self) -> bool:
|
||||
"""Return true if job has finished, one way or another."""
|
||||
return self.status not in [DownloadJobStatus.WAITING, DownloadJobStatus.RUNNING]
|
||||
|
||||
@property
|
||||
def on_start(self) -> Optional[DownloadEventHandler]:
|
||||
"""Return the on_start event handler."""
|
||||
return self._on_start
|
||||
|
||||
@property
|
||||
def on_progress(self) -> Optional[DownloadEventHandler]:
|
||||
"""Return the on_progress event handler."""
|
||||
return self._on_progress
|
||||
|
||||
@property
|
||||
def on_complete(self) -> Optional[DownloadEventHandler]:
|
||||
"""Return the on_complete event handler."""
|
||||
return self._on_complete
|
||||
|
||||
@property
|
||||
def on_error(self) -> Optional[DownloadExceptionHandler]:
|
||||
"""Return the on_error event handler."""
|
||||
return self._on_error
|
||||
|
||||
@property
|
||||
def on_cancelled(self) -> Optional[DownloadEventHandler]:
|
||||
"""Return the on_cancelled event handler."""
|
||||
return self._on_cancelled
|
||||
|
||||
def set_callbacks(
|
||||
self,
|
||||
on_start: Optional[DownloadEventHandler] = None,
|
||||
on_progress: Optional[DownloadEventHandler] = None,
|
||||
on_complete: Optional[DownloadEventHandler] = None,
|
||||
on_cancelled: Optional[DownloadEventHandler] = None,
|
||||
on_error: Optional[DownloadExceptionHandler] = None,
|
||||
) -> None:
|
||||
"""Set the callbacks for download events."""
|
||||
self._on_start = on_start
|
||||
self._on_progress = on_progress
|
||||
self._on_complete = on_complete
|
||||
self._on_error = on_error
|
||||
self._on_cancelled = on_cancelled
|
||||
|
||||
|
||||
@total_ordering
|
||||
class DownloadJob(DownloadJobBase):
|
||||
"""Class to monitor and control a model download request."""
|
||||
|
||||
# required variables to be passed in on creation
|
||||
source: AnyHttpUrl = Field(description="Where to download from. Specific types specified in child classes.")
|
||||
access_token: Optional[str] = Field(default=None, description="authorization token for protected resources")
|
||||
priority: int = Field(default=10, description="Queue priority; lower values are higher priority")
|
||||
|
||||
# set internally during download process
|
||||
job_started: Optional[str] = Field(default=None, description="Timestamp for when the download job started")
|
||||
job_ended: Optional[str] = Field(
|
||||
default=None, description="Timestamp for when the download job ende1d (completed or errored)"
|
||||
)
|
||||
content_type: Optional[str] = Field(default=None, description="Content type of downloaded file")
|
||||
canonical_url: Optional[str] = Field(default=None, description="Canonical URL to request on resume")
|
||||
etag: Optional[str] = Field(default=None, description="ETag from the remote server, if available")
|
||||
last_modified: Optional[str] = Field(default=None, description="Last-Modified from the remote server, if available")
|
||||
final_url: Optional[str] = Field(default=None, description="Final resolved URL after redirects, if available")
|
||||
expected_total_bytes: Optional[int] = Field(default=None, description="Expected total size of the download")
|
||||
resume_required: bool = Field(default=False, description="True if server refused resume; restart required")
|
||||
resume_message: Optional[str] = Field(default=None, description="Message explaining why resume is required")
|
||||
resume_from_scratch: bool = Field(
|
||||
default=False,
|
||||
description="True if resume metadata existed but the partial file was missing and the download restarted from the beginning",
|
||||
)
|
||||
|
||||
def __hash__(self) -> int:
|
||||
"""Return hash of the string representation of this object, for indexing."""
|
||||
return hash(str(self))
|
||||
|
||||
def __le__(self, other: "DownloadJob") -> bool:
|
||||
"""Return True if this job's priority is less than another's."""
|
||||
return self.priority <= other.priority
|
||||
|
||||
|
||||
class MultiFileDownloadJob(DownloadJobBase):
|
||||
"""Class to monitor and control multifile downloads."""
|
||||
|
||||
download_parts: Set[DownloadJob] = Field(default_factory=set, description="List of download parts.")
|
||||
|
||||
|
||||
class DownloadQueueServiceBase(ABC):
|
||||
"""Multithreaded queue for downloading models via URL."""
|
||||
|
||||
@abstractmethod
|
||||
def start(self, *args: Any, **kwargs: Any) -> None:
|
||||
"""Start the download worker threads."""
|
||||
|
||||
@abstractmethod
|
||||
def stop(self, *args: Any, **kwargs: Any) -> None:
|
||||
"""Stop the download worker threads."""
|
||||
|
||||
@abstractmethod
|
||||
def download(
|
||||
self,
|
||||
source: AnyHttpUrl,
|
||||
dest: Path,
|
||||
priority: int = 10,
|
||||
access_token: Optional[str] = None,
|
||||
on_start: Optional[DownloadEventHandler] = None,
|
||||
on_progress: Optional[DownloadEventHandler] = None,
|
||||
on_complete: Optional[DownloadEventHandler] = None,
|
||||
on_cancelled: Optional[DownloadEventHandler] = None,
|
||||
on_error: Optional[DownloadExceptionHandler] = None,
|
||||
) -> DownloadJob:
|
||||
"""
|
||||
Create and enqueue download job.
|
||||
|
||||
:param source: Source of the download as a URL.
|
||||
:param dest: Path to download to. See below.
|
||||
:param on_start, on_progress, on_complete, on_error: Callbacks for the indicated
|
||||
events.
|
||||
:returns: A DownloadJob object for monitoring the state of the download.
|
||||
|
||||
The `dest` argument is a Path object. Its behavior is:
|
||||
|
||||
1. If the path exists and is a directory, then the URL contents will be downloaded
|
||||
into that directory using the filename indicated in the response's `Content-Disposition` field.
|
||||
If no content-disposition is present, then the last component of the URL will be used (similar to
|
||||
wget's behavior).
|
||||
2. If the path does not exist, then it is taken as the name of a new file to create with the downloaded
|
||||
content.
|
||||
3. If the path exists and is an existing file, then the downloader will try to resume the download from
|
||||
the end of the existing file.
|
||||
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def multifile_download(
|
||||
self,
|
||||
parts: List[RemoteModelFile],
|
||||
dest: Path,
|
||||
access_token: Optional[str] = None,
|
||||
submit_job: bool = True,
|
||||
on_start: Optional[DownloadEventHandler] = None,
|
||||
on_progress: Optional[DownloadEventHandler] = None,
|
||||
on_complete: Optional[DownloadEventHandler] = None,
|
||||
on_cancelled: Optional[DownloadEventHandler] = None,
|
||||
on_error: Optional[DownloadExceptionHandler] = None,
|
||||
) -> MultiFileDownloadJob:
|
||||
"""
|
||||
Create and enqueue a multifile download job.
|
||||
|
||||
:param parts: Set of URL / filename pairs
|
||||
:param dest: Path to download to. See below.
|
||||
:param access_token: Access token to download the indicated files. If not provided,
|
||||
each file's URL may be matched to an access token using the config file matching
|
||||
system.
|
||||
:param submit_job: If true [default] then submit the job for execution. Otherwise,
|
||||
you will need to pass the job to submit_multifile_download().
|
||||
:param on_start, on_progress, on_complete, on_error: Callbacks for the indicated
|
||||
events.
|
||||
:returns: A MultiFileDownloadJob object for monitoring the state of the download.
|
||||
|
||||
The `dest` argument is a Path object pointing to a directory. All downloads
|
||||
with be placed inside this directory. The callbacks will receive the
|
||||
MultiFileDownloadJob.
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def submit_multifile_download(self, job: MultiFileDownloadJob) -> None:
|
||||
"""
|
||||
Enqueue a previously-created multi-file download job.
|
||||
|
||||
:param job: A MultiFileDownloadJob created with multifile_download()
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def submit_download_job(
|
||||
self,
|
||||
job: DownloadJob,
|
||||
on_start: Optional[DownloadEventHandler] = None,
|
||||
on_progress: Optional[DownloadEventHandler] = None,
|
||||
on_complete: Optional[DownloadEventHandler] = None,
|
||||
on_cancelled: Optional[DownloadEventHandler] = None,
|
||||
on_error: Optional[DownloadExceptionHandler] = None,
|
||||
) -> None:
|
||||
"""
|
||||
Enqueue a download job.
|
||||
|
||||
:param job: The DownloadJob
|
||||
:param on_start, on_progress, on_complete, on_error: Callbacks for the indicated
|
||||
events.
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def list_jobs(self) -> List[DownloadJob]:
|
||||
"""
|
||||
List active download jobs.
|
||||
|
||||
:returns List[DownloadJob]: List of download jobs whose state is not "completed."
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def id_to_job(self, id: int) -> DownloadJob:
|
||||
"""
|
||||
Return the DownloadJob corresponding to the integer ID.
|
||||
|
||||
:param id: ID of the DownloadJob.
|
||||
|
||||
Exceptions:
|
||||
* UnknownJobIDException
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def cancel_all_jobs(self) -> None:
|
||||
"""Cancel all active and enquedjobs."""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def prune_jobs(self) -> None:
|
||||
"""Prune completed and errored queue items from the job list."""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def cancel_job(self, job: DownloadJobBase) -> None:
|
||||
"""Cancel the job, clearing partial downloads and putting it into ERROR state."""
|
||||
pass
|
||||
|
||||
def pause_job(self, job: DownloadJobBase) -> None: # noqa D401
|
||||
"""Pause the job, preserving partial downloads."""
|
||||
raise NotImplementedError
|
||||
|
||||
@abstractmethod
|
||||
def join(self) -> None:
|
||||
"""Wait until all jobs are off the queue."""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def wait_for_job(self, job: DownloadJobBase, timeout: int = 0) -> DownloadJobBase:
|
||||
"""Wait until the indicated download job has reached a terminal state.
|
||||
|
||||
This will block until the indicated install job has completed,
|
||||
been cancelled, or errored out.
|
||||
|
||||
:param job: The job to wait on.
|
||||
:param timeout: Wait up to indicated number of seconds. Raise a TimeoutError if
|
||||
the job hasn't completed within the indicated time.
|
||||
"""
|
||||
pass
|
||||
@@ -0,0 +1,831 @@
|
||||
# Copyright (c) 2023,2026 Lincoln D. Stein
|
||||
"""Implementation of multithreaded download queue for invokeai."""
|
||||
|
||||
import os
|
||||
import re
|
||||
import threading
|
||||
import time
|
||||
import traceback
|
||||
from pathlib import Path
|
||||
from queue import Empty, PriorityQueue
|
||||
from shutil import disk_usage
|
||||
from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Set
|
||||
from urllib.parse import urlparse
|
||||
|
||||
import requests
|
||||
from pydantic.networks import AnyHttpUrl
|
||||
from requests import HTTPError
|
||||
from tqdm import tqdm
|
||||
|
||||
from invokeai.app.services.config import InvokeAIAppConfig, get_config
|
||||
from invokeai.app.services.download.download_base import (
|
||||
DownloadEventHandler,
|
||||
DownloadExceptionHandler,
|
||||
DownloadJob,
|
||||
DownloadJobBase,
|
||||
DownloadJobCancelledException,
|
||||
DownloadJobStatus,
|
||||
DownloadQueueServiceBase,
|
||||
MultiFileDownloadJob,
|
||||
ServiceInactiveException,
|
||||
UnknownJobIDException,
|
||||
)
|
||||
from invokeai.app.util.misc import get_iso_timestamp
|
||||
from invokeai.backend.model_manager.metadata import RemoteModelFile
|
||||
from invokeai.backend.util.logging import InvokeAILogger
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from invokeai.app.services.events.events_base import EventServiceBase
|
||||
|
||||
# Maximum number of bytes to download during each call to requests.iter_content()
|
||||
DOWNLOAD_CHUNK_SIZE = 100000
|
||||
|
||||
|
||||
class DownloadQueueService(DownloadQueueServiceBase):
|
||||
"""Class for queued download of models."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
max_parallel_dl: int = 5,
|
||||
app_config: Optional[InvokeAIAppConfig] = None,
|
||||
event_bus: Optional["EventServiceBase"] = None,
|
||||
requests_session: Optional[requests.sessions.Session] = None,
|
||||
):
|
||||
"""
|
||||
Initialize DownloadQueue.
|
||||
|
||||
:param app_config: InvokeAIAppConfig object
|
||||
:param max_parallel_dl: Number of simultaneous downloads allowed [5].
|
||||
:param requests_session: Optional requests.sessions.Session object, for unit tests.
|
||||
"""
|
||||
self._app_config = app_config or get_config()
|
||||
self._jobs: Dict[int, DownloadJob] = {}
|
||||
self._download_part2parent: Dict[int, MultiFileDownloadJob] = {}
|
||||
self._mfd_pending: Dict[int, list[DownloadJob]] = {}
|
||||
self._mfd_active: Dict[int, DownloadJob] = {}
|
||||
self._next_job_id = 0
|
||||
self._queue: PriorityQueue[DownloadJob] = PriorityQueue()
|
||||
self._stop_event = threading.Event()
|
||||
self._job_terminated_event = threading.Event()
|
||||
self._worker_pool: Set[threading.Thread] = set()
|
||||
self._lock = threading.Lock()
|
||||
self._logger = InvokeAILogger.get_logger("DownloadQueueService")
|
||||
self._event_bus = event_bus
|
||||
self._requests = requests_session or requests.Session()
|
||||
self._accept_download_requests = False
|
||||
self._max_parallel_dl = max_parallel_dl
|
||||
|
||||
def start(self, *args: Any, **kwargs: Any) -> None:
|
||||
"""Start the download worker threads."""
|
||||
with self._lock:
|
||||
if self._worker_pool:
|
||||
raise Exception("Attempt to start the download service twice")
|
||||
self._stop_event.clear()
|
||||
self._start_workers(self._max_parallel_dl)
|
||||
self._accept_download_requests = True
|
||||
|
||||
def stop(self, *args: Any, **kwargs: Any) -> None:
|
||||
"""Stop the download worker threads."""
|
||||
with self._lock:
|
||||
if not self._worker_pool:
|
||||
return
|
||||
self._accept_download_requests = False # reject attempts to add new jobs to queue
|
||||
queued_jobs = [x for x in self.list_jobs() if x.status == DownloadJobStatus.WAITING]
|
||||
active_jobs = [x for x in self.list_jobs() if x.status == DownloadJobStatus.RUNNING]
|
||||
if queued_jobs:
|
||||
self._logger.warning(f"Cancelling {len(queued_jobs)} queued downloads")
|
||||
if active_jobs:
|
||||
self._logger.info(f"Waiting for {len(active_jobs)} active download jobs to complete")
|
||||
with self._queue.mutex:
|
||||
self._queue.queue.clear()
|
||||
self.cancel_all_jobs()
|
||||
self._stop_event.set()
|
||||
for thread in self._worker_pool:
|
||||
thread.join()
|
||||
self._worker_pool.clear()
|
||||
|
||||
def submit_download_job(
|
||||
self,
|
||||
job: DownloadJob,
|
||||
on_start: Optional[DownloadEventHandler] = None,
|
||||
on_progress: Optional[DownloadEventHandler] = None,
|
||||
on_complete: Optional[DownloadEventHandler] = None,
|
||||
on_cancelled: Optional[DownloadEventHandler] = None,
|
||||
on_error: Optional[DownloadExceptionHandler] = None,
|
||||
) -> None:
|
||||
"""Enqueue a download job."""
|
||||
if not self._accept_download_requests:
|
||||
raise ServiceInactiveException(
|
||||
"The download service is not currently accepting requests. Please call start() to initialize the service."
|
||||
)
|
||||
if job.id == -1:
|
||||
job.id = self._next_id()
|
||||
job.set_callbacks(
|
||||
on_start=on_start,
|
||||
on_progress=on_progress,
|
||||
on_complete=on_complete,
|
||||
on_cancelled=on_cancelled,
|
||||
on_error=on_error,
|
||||
)
|
||||
self._jobs[job.id] = job
|
||||
self._queue.put(job)
|
||||
|
||||
def pause_job(self, job: DownloadJobBase) -> None:
|
||||
"""Pause the indicated job, preserving partial downloads."""
|
||||
if job.status in [DownloadJobStatus.WAITING, DownloadJobStatus.RUNNING]:
|
||||
job.pause()
|
||||
|
||||
def download(
|
||||
self,
|
||||
source: AnyHttpUrl,
|
||||
dest: Path,
|
||||
priority: int = 10,
|
||||
access_token: Optional[str] = None,
|
||||
on_start: Optional[DownloadEventHandler] = None,
|
||||
on_progress: Optional[DownloadEventHandler] = None,
|
||||
on_complete: Optional[DownloadEventHandler] = None,
|
||||
on_cancelled: Optional[DownloadEventHandler] = None,
|
||||
on_error: Optional[DownloadExceptionHandler] = None,
|
||||
) -> DownloadJob:
|
||||
"""Create and enqueue a download job and return it."""
|
||||
if not self._accept_download_requests:
|
||||
raise ServiceInactiveException(
|
||||
"The download service is not currently accepting requests. Please call start() to initialize the service."
|
||||
)
|
||||
job = DownloadJob(
|
||||
source=source,
|
||||
dest=dest,
|
||||
priority=priority,
|
||||
access_token=access_token or self._lookup_access_token(source),
|
||||
)
|
||||
self.submit_download_job(
|
||||
job,
|
||||
on_start=on_start,
|
||||
on_progress=on_progress,
|
||||
on_complete=on_complete,
|
||||
on_cancelled=on_cancelled,
|
||||
on_error=on_error,
|
||||
)
|
||||
return job
|
||||
|
||||
def multifile_download(
|
||||
self,
|
||||
parts: List[RemoteModelFile],
|
||||
dest: Path,
|
||||
access_token: Optional[str] = None,
|
||||
submit_job: bool = True,
|
||||
on_start: Optional[DownloadEventHandler] = None,
|
||||
on_progress: Optional[DownloadEventHandler] = None,
|
||||
on_complete: Optional[DownloadEventHandler] = None,
|
||||
on_cancelled: Optional[DownloadEventHandler] = None,
|
||||
on_error: Optional[DownloadExceptionHandler] = None,
|
||||
) -> MultiFileDownloadJob:
|
||||
mfdj = MultiFileDownloadJob(dest=dest, id=self._next_id())
|
||||
mfdj.set_callbacks(
|
||||
on_start=on_start,
|
||||
on_progress=on_progress,
|
||||
on_complete=on_complete,
|
||||
on_cancelled=on_cancelled,
|
||||
on_error=on_error,
|
||||
)
|
||||
|
||||
for part in parts:
|
||||
url = part.url
|
||||
path = dest / part.path
|
||||
assert path.is_relative_to(dest), "only relative download paths accepted"
|
||||
job = DownloadJob(
|
||||
source=url,
|
||||
dest=path,
|
||||
access_token=access_token or self._lookup_access_token(url),
|
||||
)
|
||||
job.id = self._next_id() # pre-assign ID so _download_part2parent can be keyed by ID
|
||||
if part.size and part.size > 0:
|
||||
job.total_bytes = part.size
|
||||
job.expected_total_bytes = part.size
|
||||
job.canonical_url = str(url)
|
||||
mfdj.download_parts.add(job)
|
||||
self._download_part2parent[job.id] = mfdj
|
||||
if submit_job:
|
||||
self.submit_multifile_download(mfdj)
|
||||
return mfdj
|
||||
|
||||
def submit_multifile_download(self, job: MultiFileDownloadJob) -> None:
|
||||
pending = sorted(job.download_parts, key=lambda j: str(j.source))
|
||||
self._mfd_pending[job.id] = list(pending)
|
||||
self._mfd_active.pop(job.id, None)
|
||||
self._submit_next_mfd_part(job)
|
||||
|
||||
def _submit_next_mfd_part(self, job: MultiFileDownloadJob) -> None:
|
||||
pending = self._mfd_pending.get(job.id, [])
|
||||
if not pending:
|
||||
return
|
||||
if self._mfd_active.get(job.id) is not None:
|
||||
return
|
||||
download_job = pending.pop(0)
|
||||
self._mfd_active[job.id] = download_job
|
||||
self.submit_download_job(
|
||||
download_job,
|
||||
on_start=self._mfd_started,
|
||||
on_progress=self._mfd_progress,
|
||||
on_complete=self._mfd_complete,
|
||||
on_cancelled=self._mfd_cancelled,
|
||||
on_error=self._mfd_error,
|
||||
)
|
||||
|
||||
def join(self) -> None:
|
||||
"""Wait for all jobs to complete."""
|
||||
self._queue.join()
|
||||
|
||||
def _next_id(self) -> int:
|
||||
with self._lock:
|
||||
id = self._next_job_id
|
||||
self._next_job_id += 1
|
||||
return id
|
||||
|
||||
def list_jobs(self) -> List[DownloadJob]:
|
||||
"""List all the jobs."""
|
||||
return list(self._jobs.values())
|
||||
|
||||
def prune_jobs(self) -> None:
|
||||
"""Prune completed and errored queue items from the job list."""
|
||||
with self._lock:
|
||||
to_delete = set()
|
||||
for job_id, job in self._jobs.items():
|
||||
if job.in_terminal_state:
|
||||
to_delete.add(job_id)
|
||||
for job_id in to_delete:
|
||||
del self._jobs[job_id]
|
||||
|
||||
def id_to_job(self, id: int) -> DownloadJob:
|
||||
"""Translate a job ID into a DownloadJob object."""
|
||||
try:
|
||||
return self._jobs[id]
|
||||
except KeyError as excp:
|
||||
raise UnknownJobIDException("Unrecognized job") from excp
|
||||
|
||||
def cancel_job(self, job: DownloadJobBase) -> None:
|
||||
"""
|
||||
Cancel the indicated job.
|
||||
|
||||
If it is running it will be stopped.
|
||||
job.status will be set to DownloadJobStatus.CANCELLED
|
||||
"""
|
||||
if job.status in [DownloadJobStatus.WAITING, DownloadJobStatus.RUNNING]:
|
||||
job.cancel()
|
||||
|
||||
def cancel_all_jobs(self) -> None:
|
||||
"""Cancel all jobs (those not in enqueued, running or paused state)."""
|
||||
for job in self._jobs.values():
|
||||
if not job.in_terminal_state:
|
||||
self.cancel_job(job)
|
||||
|
||||
def wait_for_job(self, job: DownloadJobBase, timeout: int = 0) -> DownloadJobBase:
|
||||
"""Block until the indicated job has reached terminal state, or when timeout limit reached."""
|
||||
start = time.time()
|
||||
while not job.in_terminal_state:
|
||||
if self._job_terminated_event.wait(timeout=0.25): # in case we miss an event
|
||||
self._job_terminated_event.clear()
|
||||
if timeout > 0 and time.time() - start > timeout:
|
||||
raise TimeoutError("Timeout exceeded")
|
||||
return job
|
||||
|
||||
def _start_workers(self, max_workers: int) -> None:
|
||||
"""Start the requested number of worker threads."""
|
||||
self._stop_event.clear()
|
||||
for i in range(0, max_workers): # noqa B007
|
||||
worker = threading.Thread(target=self._download_next_item, daemon=True)
|
||||
self._logger.debug(f"Download queue worker thread {worker.name} starting.")
|
||||
worker.start()
|
||||
self._worker_pool.add(worker)
|
||||
|
||||
def _download_next_item(self) -> None:
|
||||
"""Worker thread gets next job on priority queue."""
|
||||
done = False
|
||||
while not done:
|
||||
if self._stop_event.is_set():
|
||||
done = True
|
||||
continue
|
||||
try:
|
||||
job = self._queue.get(timeout=1)
|
||||
except Empty:
|
||||
continue
|
||||
try:
|
||||
if job.cancelled:
|
||||
raise DownloadJobCancelledException("Job was cancelled before start")
|
||||
job.job_started = get_iso_timestamp()
|
||||
self._do_download(job)
|
||||
if job.status != DownloadJobStatus.COMPLETED:
|
||||
self._signal_job_complete(job)
|
||||
except DownloadJobCancelledException:
|
||||
if job.paused:
|
||||
self._signal_job_paused(job)
|
||||
else:
|
||||
self._signal_job_cancelled(job)
|
||||
self._cleanup_cancelled_job(job)
|
||||
except Exception as excp:
|
||||
job.error_type = excp.__class__.__name__ + f"({str(excp)})"
|
||||
job.error = traceback.format_exc()
|
||||
self._signal_job_error(job, excp)
|
||||
finally:
|
||||
job.job_ended = get_iso_timestamp()
|
||||
self._job_terminated_event.set() # signal a change to terminal state
|
||||
self._download_part2parent.pop(job.id, None) # if this is a subpart of a multipart job, remove it
|
||||
self._queue.task_done()
|
||||
|
||||
self._logger.debug(f"Download queue worker thread {threading.current_thread().name} exiting.")
|
||||
|
||||
def _do_download(self, job: DownloadJob) -> None:
|
||||
"""Do the actual download."""
|
||||
|
||||
url = job.canonical_url or str(job.source)
|
||||
header = {"Authorization": f"Bearer {job.access_token}"} if job.access_token else {}
|
||||
had_resume_metadata = bool(job.etag or job.last_modified)
|
||||
open_mode = "wb"
|
||||
resume_from = 0
|
||||
|
||||
if not job.dest.is_dir():
|
||||
job.download_path = job.dest
|
||||
in_progress_path = self._in_progress_path(job.download_path)
|
||||
if in_progress_path.exists():
|
||||
resume_from = in_progress_path.stat().st_size
|
||||
job.bytes = resume_from
|
||||
self._logger.debug(
|
||||
f"Resume check: in-progress file found at {in_progress_path} size={resume_from} bytes"
|
||||
)
|
||||
if resume_from > 0:
|
||||
if job.etag:
|
||||
header["If-Range"] = job.etag
|
||||
elif job.last_modified:
|
||||
header["If-Range"] = job.last_modified
|
||||
header["Range"] = f"bytes={resume_from}-"
|
||||
open_mode = "ab"
|
||||
else:
|
||||
self._logger.debug(f"Resume check: no in-progress file at {in_progress_path}")
|
||||
elif job.download_path:
|
||||
# Resume for directory downloads when we already know the filename.
|
||||
in_progress_path = self._in_progress_path(job.download_path)
|
||||
if in_progress_path.exists():
|
||||
resume_from = in_progress_path.stat().st_size
|
||||
job.bytes = resume_from
|
||||
self._logger.debug(
|
||||
f"Resume check (dir): in-progress file found at {in_progress_path} size={resume_from} bytes"
|
||||
)
|
||||
if resume_from > 0:
|
||||
if job.etag:
|
||||
header["If-Range"] = job.etag
|
||||
elif job.last_modified:
|
||||
header["If-Range"] = job.last_modified
|
||||
header["Range"] = f"bytes={resume_from}-"
|
||||
open_mode = "ab"
|
||||
else:
|
||||
self._logger.debug(f"Resume check (dir): no in-progress file at {in_progress_path}")
|
||||
elif job.dest.is_dir():
|
||||
# Attempt to infer a single in-progress file from disk for directory downloads.
|
||||
try:
|
||||
candidates = sorted(job.dest.glob("*.downloading"))
|
||||
except OSError:
|
||||
candidates = []
|
||||
if len(candidates) == 1:
|
||||
inferred = candidates[0].with_name(candidates[0].name.removesuffix(".downloading"))
|
||||
job.download_path = inferred
|
||||
try:
|
||||
resume_from = candidates[0].stat().st_size
|
||||
except FileNotFoundError:
|
||||
# The .downloading file was renamed/deleted between glob and stat (race condition); skip resume.
|
||||
job.download_path = None
|
||||
else:
|
||||
job.bytes = resume_from
|
||||
self._logger.debug(
|
||||
f"Resume check (dir): inferred in-progress file path={candidates[0]} size={resume_from} bytes"
|
||||
)
|
||||
if resume_from > 0:
|
||||
if job.etag:
|
||||
header["If-Range"] = job.etag
|
||||
elif job.last_modified:
|
||||
header["If-Range"] = job.last_modified
|
||||
header["Range"] = f"bytes={resume_from}-"
|
||||
open_mode = "ab"
|
||||
else:
|
||||
self._logger.debug(
|
||||
"Resume check (dir): no prior download_path available; cannot resume from disk "
|
||||
f"(candidates={len(candidates)})"
|
||||
)
|
||||
|
||||
if resume_from == 0:
|
||||
job.bytes = 0
|
||||
if had_resume_metadata:
|
||||
job.resume_from_scratch = True
|
||||
job.resume_message = "Partial file missing. Restarted download from the beginning."
|
||||
|
||||
# Make a streaming request. This will retrieve headers including
|
||||
# content-length and content-disposition, but not fetch any content itself
|
||||
resp = self._requests.get(str(url), headers=header, stream=True)
|
||||
job.final_url = str(resp.url) if resp.url else None
|
||||
self._logger.debug(
|
||||
"Resume response: "
|
||||
f"status={resp.status_code} "
|
||||
f"content_length={resp.headers.get('content-length')} "
|
||||
f"content_range={resp.headers.get('Content-Range')} "
|
||||
f"etag={resp.headers.get('ETag')} "
|
||||
f"last_modified={resp.headers.get('Last-Modified')}"
|
||||
)
|
||||
if resp.status_code == 416 and resume_from > 0:
|
||||
# Range not satisfiable - local partial is already complete
|
||||
expected = job.expected_total_bytes or job.total_bytes or resume_from
|
||||
if resume_from == expected:
|
||||
job.total_bytes = expected
|
||||
job.bytes = resume_from
|
||||
job.download_path = job.download_path or job.dest
|
||||
self._signal_job_started(job)
|
||||
self._signal_job_complete(job)
|
||||
return
|
||||
job.resume_required = True
|
||||
job.resume_message = "Resume refused by server. Restart required."
|
||||
job.pause()
|
||||
raise DownloadJobCancelledException("Resume refused by server. Restart required.")
|
||||
if not resp.ok:
|
||||
host = urlparse(str(resp.url or url)).netloc
|
||||
status = resp.status_code
|
||||
reason = resp.reason
|
||||
if status >= 500:
|
||||
self._logger.error(f"Remote server error from {host}: HTTP {status} {reason}")
|
||||
raise HTTPError(reason)
|
||||
self._logger.error(f"Download failed from {host}: HTTP {status} {reason}")
|
||||
raise HTTPError(reason)
|
||||
|
||||
job.content_type = resp.headers.get("Content-Type")
|
||||
job.etag = resp.headers.get("ETag") or job.etag
|
||||
job.last_modified = resp.headers.get("Last-Modified") or job.last_modified
|
||||
content_length = int(resp.headers.get("content-length", 0))
|
||||
|
||||
if job.dest.is_dir():
|
||||
parsed_url = urlparse(str(url))
|
||||
file_name = os.path.basename(parsed_url.path) # default is to use the last bit of the URL
|
||||
|
||||
if match := re.search('filename="(.+)"', resp.headers.get("Content-Disposition", "")):
|
||||
remote_name = match.group(1)
|
||||
if self._validate_filename(job.dest.as_posix(), remote_name):
|
||||
file_name = remote_name
|
||||
|
||||
job.download_path = job.dest / file_name
|
||||
|
||||
else:
|
||||
job.dest.parent.mkdir(parents=True, exist_ok=True)
|
||||
job.download_path = job.dest
|
||||
|
||||
assert job.download_path
|
||||
|
||||
in_progress_path = self._in_progress_path(job.download_path)
|
||||
|
||||
if resume_from > 0 and resp.status_code == 200:
|
||||
# Server ignored Range. Restart download from scratch.
|
||||
job.resume_required = True
|
||||
job.resume_message = "Resume refused by server. Restart required."
|
||||
job.pause()
|
||||
raise DownloadJobCancelledException("Resume refused by server. Restart required.")
|
||||
|
||||
if resume_from > 0 and resp.status_code == 206:
|
||||
content_range = resp.headers.get("Content-Range", "")
|
||||
total_from_range = None
|
||||
if match := re.match(r"bytes\s+\d+-\d+/(\d+)", content_range):
|
||||
total_from_range = int(match.group(1))
|
||||
if total_from_range is not None:
|
||||
job.total_bytes = total_from_range
|
||||
else:
|
||||
job.total_bytes = resume_from + content_length
|
||||
job.bytes = resume_from
|
||||
job.expected_total_bytes = job.total_bytes
|
||||
else:
|
||||
job.total_bytes = content_length
|
||||
job.expected_total_bytes = content_length
|
||||
|
||||
if job.download_path.exists() and resume_from == 0:
|
||||
existing_size = job.download_path.stat().st_size
|
||||
if job.total_bytes > 0 and existing_size == job.total_bytes:
|
||||
job.bytes = existing_size
|
||||
self._signal_job_started(job)
|
||||
self._signal_job_complete(job)
|
||||
return
|
||||
# Existing file does not match expected size; treat as corrupt and restart.
|
||||
self._logger.debug(
|
||||
"Resume check: existing file size mismatch; deleting and restarting "
|
||||
f"path={job.download_path} existing_size={existing_size} expected={job.total_bytes}"
|
||||
)
|
||||
job.download_path.unlink()
|
||||
|
||||
free_space = disk_usage(job.download_path.parent).free
|
||||
GB = 2**30
|
||||
remaining_bytes = max(job.total_bytes - job.bytes, 0)
|
||||
if free_space < remaining_bytes:
|
||||
raise RuntimeError(
|
||||
f"Free disk space {free_space / GB:.2f} GB is not enough for download of {remaining_bytes / GB:.2f} GB."
|
||||
)
|
||||
|
||||
# Don't clobber an existing file. See commit 82c2c85202f88c6d24ff84710f297cfc6ae174af
|
||||
# for code that instead resumes an interrupted download.
|
||||
if job.download_path.exists() and resume_from == 0:
|
||||
raise OSError(f"[Errno 17] File {job.download_path} exists")
|
||||
|
||||
# append ".downloading" to the path
|
||||
# signal caller that the download is starting. At this point, key fields such as
|
||||
# download_path and total_bytes will be populated. We call it here because the might
|
||||
# discover that the local file is already complete and generate a COMPLETED status.
|
||||
self._signal_job_started(job)
|
||||
|
||||
expected_total = job.total_bytes or job.expected_total_bytes or content_length
|
||||
# "range not satisfiable" - local file is at least as large as the remote file
|
||||
if resp.status_code == 416 or (expected_total > 0 and job.bytes >= expected_total):
|
||||
self._logger.info(f"{job.download_path}: complete file found. Skipping.")
|
||||
return
|
||||
|
||||
# "partial content" - local file is smaller than remote file
|
||||
elif resp.status_code == 206 or job.bytes > 0:
|
||||
self._logger.info(f"{job.download_path}: partial file found. Resuming")
|
||||
|
||||
# some other error
|
||||
elif resp.status_code != 200:
|
||||
host = urlparse(str(resp.url or url)).netloc
|
||||
status = resp.status_code
|
||||
reason = resp.reason
|
||||
if status >= 500:
|
||||
self._logger.error(f"Remote server error from {host}: HTTP {status} {reason}")
|
||||
raise HTTPError(reason)
|
||||
self._logger.error(f"Download failed from {host}: HTTP {status} {reason}")
|
||||
raise HTTPError(reason)
|
||||
|
||||
self._logger.debug(f"{job.source}: Downloading {job.download_path}")
|
||||
report_delta = job.total_bytes / 100 # report every 1% change
|
||||
last_report_bytes = 0
|
||||
|
||||
# DOWNLOAD LOOP
|
||||
with open(in_progress_path, open_mode) as file:
|
||||
for data in resp.iter_content(chunk_size=DOWNLOAD_CHUNK_SIZE):
|
||||
if job.cancelled:
|
||||
raise DownloadJobCancelledException("Job was cancelled at caller's request")
|
||||
|
||||
job.bytes += file.write(data)
|
||||
if (job.bytes - last_report_bytes >= report_delta) or (job.bytes >= job.total_bytes):
|
||||
last_report_bytes = job.bytes
|
||||
self._signal_job_progress(job)
|
||||
|
||||
if job.total_bytes > 0 and job.bytes < job.total_bytes:
|
||||
job.resume_required = True
|
||||
job.resume_message = "Download interrupted. Resume required."
|
||||
job.pause()
|
||||
raise DownloadJobCancelledException("Download interrupted. Resume required.")
|
||||
|
||||
# if we get here we are done and can rename the file to the original dest
|
||||
self._logger.debug(f"{job.source}: saved to {job.download_path} (bytes={job.bytes})")
|
||||
in_progress_path.rename(job.download_path)
|
||||
|
||||
def _validate_filename(self, directory: str, filename: str) -> bool:
|
||||
pc_name_max = get_pc_name_max(directory)
|
||||
pc_path_max = get_pc_path_max(directory)
|
||||
if "/" in filename:
|
||||
return False
|
||||
if filename.startswith(".."):
|
||||
return False
|
||||
if len(filename) > pc_name_max:
|
||||
return False
|
||||
if len(os.path.join(directory, filename)) > pc_path_max:
|
||||
return False
|
||||
return True
|
||||
|
||||
def _in_progress_path(self, path: Path) -> Path:
|
||||
return path.with_name(path.name + ".downloading")
|
||||
|
||||
def _lookup_access_token(self, source: AnyHttpUrl) -> Optional[str]:
|
||||
# Pull the token from config if it exists and matches the URL
|
||||
token = None
|
||||
for pair in self._app_config.remote_api_tokens or []:
|
||||
if re.search(pair.url_regex, str(source)):
|
||||
token = pair.token
|
||||
break
|
||||
return token
|
||||
|
||||
def _signal_job_started(self, job: DownloadJob) -> None:
|
||||
job.status = DownloadJobStatus.RUNNING
|
||||
self._execute_cb(job, "on_start")
|
||||
if self._event_bus:
|
||||
self._event_bus.emit_download_started(job)
|
||||
|
||||
def _signal_job_progress(self, job: DownloadJob) -> None:
|
||||
self._execute_cb(job, "on_progress")
|
||||
if self._event_bus:
|
||||
self._event_bus.emit_download_progress(job)
|
||||
|
||||
def _signal_job_complete(self, job: DownloadJob) -> None:
|
||||
job.status = DownloadJobStatus.COMPLETED
|
||||
self._execute_cb(job, "on_complete")
|
||||
if self._event_bus:
|
||||
self._event_bus.emit_download_complete(job)
|
||||
|
||||
def _signal_job_cancelled(self, job: DownloadJob) -> None:
|
||||
if job.status not in [DownloadJobStatus.RUNNING, DownloadJobStatus.WAITING]:
|
||||
return
|
||||
job.status = DownloadJobStatus.CANCELLED
|
||||
self._execute_cb(job, "on_cancelled")
|
||||
if self._event_bus:
|
||||
self._event_bus.emit_download_cancelled(job)
|
||||
|
||||
# if multifile download, then signal the parent
|
||||
if parent_job := self._download_part2parent.get(job.id, None):
|
||||
if not parent_job.in_terminal_state:
|
||||
parent_job.status = DownloadJobStatus.CANCELLED
|
||||
self._execute_cb(parent_job, "on_cancelled")
|
||||
|
||||
def _signal_job_paused(self, job: DownloadJob) -> None:
|
||||
if job.status not in [DownloadJobStatus.RUNNING, DownloadJobStatus.WAITING]:
|
||||
return
|
||||
if job.download_path:
|
||||
in_progress_path = self._in_progress_path(job.download_path)
|
||||
if in_progress_path.exists():
|
||||
job.bytes = in_progress_path.stat().st_size
|
||||
job.status = DownloadJobStatus.PAUSED
|
||||
self._execute_cb(job, "on_cancelled")
|
||||
if self._event_bus:
|
||||
self._event_bus.emit_download_paused(job)
|
||||
|
||||
if parent_job := self._download_part2parent.get(job.id, None):
|
||||
if not parent_job.in_terminal_state:
|
||||
parent_job.status = DownloadJobStatus.PAUSED
|
||||
self._execute_cb(parent_job, "on_cancelled")
|
||||
|
||||
def _signal_job_error(self, job: DownloadJob, excp: Optional[Exception] = None) -> None:
|
||||
job.status = DownloadJobStatus.ERROR
|
||||
self._logger.error(f"{str(job.source)}: {traceback.format_exception(excp)}")
|
||||
self._execute_cb(job, "on_error", excp)
|
||||
|
||||
if self._event_bus:
|
||||
self._event_bus.emit_download_error(job)
|
||||
|
||||
def _cleanup_cancelled_job(self, job: DownloadJob) -> None:
|
||||
if job.paused:
|
||||
return
|
||||
self._logger.debug(f"Cleaning up leftover files from cancelled download job {job.download_path}")
|
||||
try:
|
||||
if job.download_path:
|
||||
partial_file = self._in_progress_path(job.download_path)
|
||||
partial_file.unlink()
|
||||
except OSError as excp:
|
||||
self._logger.warning(excp)
|
||||
|
||||
########################################
|
||||
# callbacks used for multifile downloads
|
||||
########################################
|
||||
def _mfd_started(self, download_job: DownloadJob) -> None:
|
||||
self._logger.info(f"File download started: {download_job.source}")
|
||||
with self._lock:
|
||||
mf_job = self._download_part2parent[download_job.id]
|
||||
if mf_job.waiting:
|
||||
mf_job.total_bytes = sum(x.total_bytes for x in mf_job.download_parts)
|
||||
mf_job.status = DownloadJobStatus.RUNNING
|
||||
assert download_job.download_path is not None
|
||||
path_relative_to_destdir = download_job.download_path.relative_to(mf_job.dest)
|
||||
mf_job.download_path = (
|
||||
mf_job.dest / path_relative_to_destdir.parts[0]
|
||||
) # keep just the first component of the path
|
||||
self._execute_cb(mf_job, "on_start")
|
||||
|
||||
def _mfd_progress(self, download_job: DownloadJob) -> None:
|
||||
with self._lock:
|
||||
mf_job = self._download_part2parent[download_job.id]
|
||||
if mf_job.cancelled:
|
||||
for part in mf_job.download_parts:
|
||||
self.cancel_job(part)
|
||||
elif mf_job.running:
|
||||
mf_job.total_bytes = sum(x.total_bytes for x in mf_job.download_parts)
|
||||
mf_job.bytes = sum(x.bytes for x in mf_job.download_parts)
|
||||
self._execute_cb(mf_job, "on_progress")
|
||||
|
||||
def _mfd_complete(self, download_job: DownloadJob) -> None:
|
||||
self._logger.info(f"Download complete: {download_job.source}")
|
||||
submit_next = False
|
||||
mf_job: Optional[MultiFileDownloadJob] = None
|
||||
with self._lock:
|
||||
mf_job = self._download_part2parent[download_job.id]
|
||||
self._mfd_active.pop(mf_job.id, None)
|
||||
mf_job.total_bytes = sum(x.total_bytes for x in mf_job.download_parts)
|
||||
mf_job.bytes = sum(x.bytes for x in mf_job.download_parts)
|
||||
|
||||
# are there any more active jobs left in this task?
|
||||
if all(x.complete for x in mf_job.download_parts):
|
||||
mf_job.status = DownloadJobStatus.COMPLETED
|
||||
self._execute_cb(mf_job, "on_complete")
|
||||
elif not mf_job.in_terminal_state and not mf_job.paused:
|
||||
submit_next = True
|
||||
|
||||
# we're done with this sub-job
|
||||
self._job_terminated_event.set()
|
||||
if submit_next and mf_job is not None:
|
||||
self._submit_next_mfd_part(mf_job)
|
||||
|
||||
def _mfd_cancelled(self, download_job: DownloadJob) -> None:
|
||||
with self._lock:
|
||||
mf_job = self._download_part2parent[download_job.id]
|
||||
assert mf_job is not None
|
||||
self._mfd_active.pop(mf_job.id, None)
|
||||
|
||||
if not mf_job.in_terminal_state:
|
||||
if download_job.paused:
|
||||
self._logger.warning(f"Download paused: {download_job.source}")
|
||||
mf_job.pause()
|
||||
else:
|
||||
self._logger.warning(f"Download cancelled: {download_job.source}")
|
||||
mf_job.cancel()
|
||||
|
||||
if download_job.paused:
|
||||
return
|
||||
for s in mf_job.download_parts:
|
||||
self.cancel_job(s)
|
||||
self._mfd_pending.pop(mf_job.id, None)
|
||||
|
||||
def _mfd_error(self, download_job: DownloadJob, excp: Optional[Exception] = None) -> None:
|
||||
with self._lock:
|
||||
mf_job = self._download_part2parent[download_job.id]
|
||||
assert mf_job is not None
|
||||
self._mfd_active.pop(mf_job.id, None)
|
||||
if not mf_job.in_terminal_state:
|
||||
mf_job.status = download_job.status
|
||||
mf_job.error = download_job.error
|
||||
mf_job.error_type = download_job.error_type
|
||||
self._execute_cb(mf_job, "on_error", excp)
|
||||
self._logger.error(
|
||||
f"Cancelling {mf_job.dest} due to an error while downloading {download_job.source}: {str(excp)}"
|
||||
)
|
||||
for s in [x for x in mf_job.download_parts if x.running]:
|
||||
self.cancel_job(s)
|
||||
self._mfd_pending.pop(mf_job.id, None)
|
||||
self._job_terminated_event.set()
|
||||
|
||||
def _execute_cb(
|
||||
self,
|
||||
job: DownloadJob | MultiFileDownloadJob,
|
||||
callback_name: Literal[
|
||||
"on_start",
|
||||
"on_progress",
|
||||
"on_complete",
|
||||
"on_cancelled",
|
||||
"on_error",
|
||||
],
|
||||
excp: Optional[Exception] = None,
|
||||
) -> None:
|
||||
if callback := getattr(job, callback_name, None):
|
||||
args = [job, excp] if excp else [job]
|
||||
try:
|
||||
callback(*args)
|
||||
except Exception as e:
|
||||
self._logger.error(
|
||||
f"An error occurred while processing the {callback_name} callback: {traceback.format_exception(e)}"
|
||||
)
|
||||
|
||||
|
||||
def get_pc_name_max(directory: str) -> int:
|
||||
if hasattr(os, "pathconf"):
|
||||
try:
|
||||
return os.pathconf(directory, "PC_NAME_MAX")
|
||||
except OSError:
|
||||
# macOS w/ external drives raise OSError
|
||||
pass
|
||||
return 260 # hardcoded for windows
|
||||
|
||||
|
||||
def get_pc_path_max(directory: str) -> int:
|
||||
if hasattr(os, "pathconf"):
|
||||
try:
|
||||
return os.pathconf(directory, "PC_PATH_MAX")
|
||||
except OSError:
|
||||
# some platforms may not have this value
|
||||
pass
|
||||
return 32767 # hardcoded for windows with long names enabled
|
||||
|
||||
|
||||
# Example on_progress event handler to display a TQDM status bar
|
||||
# Activate with:
|
||||
# download_service.download(DownloadJob('http://foo.bar/baz', '/tmp', on_progress=TqdmProgress().update))
|
||||
class TqdmProgress(object):
|
||||
"""TQDM-based progress bar object to use in on_progress handlers."""
|
||||
|
||||
_bars: Dict[int, tqdm] # type: ignore
|
||||
_last: Dict[int, int] # last bytes downloaded
|
||||
|
||||
def __init__(self) -> None: # noqa D107
|
||||
self._bars = {}
|
||||
self._last = {}
|
||||
|
||||
def update(self, job: DownloadJob) -> None: # noqa D102
|
||||
job_id = job.id
|
||||
# new job
|
||||
if job_id not in self._bars:
|
||||
assert job.download_path
|
||||
dest = Path(job.download_path).name
|
||||
self._bars[job_id] = tqdm(
|
||||
desc=dest,
|
||||
initial=0,
|
||||
total=job.total_bytes,
|
||||
unit="iB",
|
||||
unit_scale=True,
|
||||
)
|
||||
self._last[job_id] = 0
|
||||
self._bars[job_id].update(job.bytes - self._last[job_id])
|
||||
self._last[job_id] = job.bytes
|
||||
@@ -0,0 +1,263 @@
|
||||
# Copyright (c) 2022 Kyle Schouviller (https://github.com/kyle0654)
|
||||
|
||||
|
||||
from typing import TYPE_CHECKING, Optional
|
||||
|
||||
from invokeai.app.services.events.events_common import (
|
||||
BatchEnqueuedEvent,
|
||||
BulkDownloadCompleteEvent,
|
||||
BulkDownloadErrorEvent,
|
||||
BulkDownloadStartedEvent,
|
||||
DownloadCancelledEvent,
|
||||
DownloadCompleteEvent,
|
||||
DownloadErrorEvent,
|
||||
DownloadPausedEvent,
|
||||
DownloadProgressEvent,
|
||||
DownloadStartedEvent,
|
||||
EventBase,
|
||||
InvocationCompleteEvent,
|
||||
InvocationErrorEvent,
|
||||
InvocationProgressEvent,
|
||||
InvocationStartedEvent,
|
||||
ModelInstallCancelledEvent,
|
||||
ModelInstallCompleteEvent,
|
||||
ModelInstallDownloadProgressEvent,
|
||||
ModelInstallDownloadsCompleteEvent,
|
||||
ModelInstallDownloadStartedEvent,
|
||||
ModelInstallErrorEvent,
|
||||
ModelInstallStartedEvent,
|
||||
ModelLoadCompleteEvent,
|
||||
ModelLoadStartedEvent,
|
||||
QueueClearedEvent,
|
||||
QueueItemsRetriedEvent,
|
||||
QueueItemStatusChangedEvent,
|
||||
RecallParametersUpdatedEvent,
|
||||
WorkflowCreatedEvent,
|
||||
WorkflowDeletedEvent,
|
||||
WorkflowUpdatedEvent,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from invokeai.app.invocations.baseinvocation import BaseInvocation, BaseInvocationOutput
|
||||
from invokeai.app.services.download.download_base import DownloadJob
|
||||
from invokeai.app.services.model_install.model_install_common import ModelInstallJob
|
||||
from invokeai.app.services.session_processor.session_processor_common import ProgressImage
|
||||
from invokeai.app.services.session_queue.session_queue_common import (
|
||||
BatchStatus,
|
||||
EnqueueBatchResult,
|
||||
RetryItemsResult,
|
||||
SessionQueueItem,
|
||||
SessionQueueStatus,
|
||||
)
|
||||
from invokeai.backend.model_manager.configs.factory import AnyModelConfig
|
||||
from invokeai.backend.model_manager.taxonomy import SubModelType
|
||||
|
||||
|
||||
class EventServiceBase:
|
||||
"""Basic event bus, to have an empty stand-in when not needed"""
|
||||
|
||||
def dispatch(self, event: "EventBase") -> None:
|
||||
pass
|
||||
|
||||
# region: Invocation
|
||||
|
||||
def emit_invocation_started(self, queue_item: "SessionQueueItem", invocation: "BaseInvocation") -> None:
|
||||
"""Emitted when an invocation is started"""
|
||||
self.dispatch(InvocationStartedEvent.build(queue_item, invocation))
|
||||
|
||||
def emit_invocation_progress(
|
||||
self,
|
||||
queue_item: "SessionQueueItem",
|
||||
invocation: "BaseInvocation",
|
||||
message: str,
|
||||
percentage: float | None = None,
|
||||
image: "ProgressImage | None" = None,
|
||||
) -> None:
|
||||
"""Emitted at periodically during an invocation"""
|
||||
self.dispatch(InvocationProgressEvent.build(queue_item, invocation, message, percentage, image))
|
||||
|
||||
def emit_invocation_complete(
|
||||
self, queue_item: "SessionQueueItem", invocation: "BaseInvocation", output: "BaseInvocationOutput"
|
||||
) -> None:
|
||||
"""Emitted when an invocation is complete"""
|
||||
self.dispatch(InvocationCompleteEvent.build(queue_item, invocation, output))
|
||||
|
||||
def emit_invocation_error(
|
||||
self,
|
||||
queue_item: "SessionQueueItem",
|
||||
invocation: "BaseInvocation",
|
||||
error_type: str,
|
||||
error_message: str,
|
||||
error_traceback: str,
|
||||
) -> None:
|
||||
"""Emitted when an invocation encounters an error"""
|
||||
self.dispatch(InvocationErrorEvent.build(queue_item, invocation, error_type, error_message, error_traceback))
|
||||
|
||||
# endregion
|
||||
|
||||
# region Queue
|
||||
|
||||
def emit_queue_item_status_changed(
|
||||
self, queue_item: "SessionQueueItem", batch_status: "BatchStatus", queue_status: "SessionQueueStatus"
|
||||
) -> None:
|
||||
"""Emitted when a queue item's status changes"""
|
||||
self.dispatch(QueueItemStatusChangedEvent.build(queue_item, batch_status, queue_status))
|
||||
|
||||
def emit_batch_enqueued(self, enqueue_result: "EnqueueBatchResult", user_id: str = "system") -> None:
|
||||
"""Emitted when a batch is enqueued"""
|
||||
self.dispatch(BatchEnqueuedEvent.build(enqueue_result, user_id))
|
||||
|
||||
def emit_queue_items_retried(
|
||||
self, retry_result: "RetryItemsResult", user_ids: list[str], retried_item_ids_by_user: dict[str, list[int]]
|
||||
) -> None:
|
||||
"""Emitted when a list of queue items are retried"""
|
||||
self.dispatch(QueueItemsRetriedEvent.build(retry_result, user_ids, retried_item_ids_by_user))
|
||||
|
||||
def emit_queue_cleared(self, queue_id: str, user_id: str | None = None) -> None:
|
||||
"""Emitted when a queue is cleared. `user_id` scopes the clear to one user's items; None means all."""
|
||||
self.dispatch(QueueClearedEvent.build(queue_id, user_id))
|
||||
|
||||
def emit_recall_parameters_updated(self, queue_id: str, user_id: str, parameters: dict) -> None:
|
||||
"""Emitted when recall parameters are updated"""
|
||||
self.dispatch(RecallParametersUpdatedEvent.build(queue_id, user_id, parameters))
|
||||
|
||||
# endregion
|
||||
|
||||
# region Workflow library
|
||||
|
||||
def emit_workflow_created(self, workflow_id: str, user_id: str, is_public: bool) -> None:
|
||||
"""Emitted when a workflow is created."""
|
||||
self.dispatch(WorkflowCreatedEvent.build(workflow_id=workflow_id, user_id=user_id, is_public=is_public))
|
||||
|
||||
def emit_workflow_updated(self, workflow_id: str, user_id: str, old_is_public: bool, new_is_public: bool) -> None:
|
||||
"""Emitted when a workflow is updated."""
|
||||
self.dispatch(
|
||||
WorkflowUpdatedEvent.build(
|
||||
workflow_id=workflow_id,
|
||||
user_id=user_id,
|
||||
old_is_public=old_is_public,
|
||||
new_is_public=new_is_public,
|
||||
)
|
||||
)
|
||||
|
||||
def emit_workflow_deleted(self, workflow_id: str, user_id: str, is_public: bool) -> None:
|
||||
"""Emitted when a workflow is deleted."""
|
||||
self.dispatch(WorkflowDeletedEvent.build(workflow_id=workflow_id, user_id=user_id, is_public=is_public))
|
||||
|
||||
# endregion
|
||||
|
||||
# region Download
|
||||
|
||||
def emit_download_started(self, job: "DownloadJob") -> None:
|
||||
"""Emitted when a download is started"""
|
||||
self.dispatch(DownloadStartedEvent.build(job))
|
||||
|
||||
def emit_download_progress(self, job: "DownloadJob") -> None:
|
||||
"""Emitted at intervals during a download"""
|
||||
self.dispatch(DownloadProgressEvent.build(job))
|
||||
|
||||
def emit_download_complete(self, job: "DownloadJob") -> None:
|
||||
"""Emitted when a download is completed"""
|
||||
self.dispatch(DownloadCompleteEvent.build(job))
|
||||
|
||||
def emit_download_cancelled(self, job: "DownloadJob") -> None:
|
||||
"""Emitted when a download is cancelled"""
|
||||
self.dispatch(DownloadCancelledEvent.build(job))
|
||||
|
||||
def emit_download_paused(self, job: "DownloadJob") -> None:
|
||||
"""Emitted when a download is paused"""
|
||||
self.dispatch(DownloadPausedEvent.build(job))
|
||||
|
||||
def emit_download_error(self, job: "DownloadJob") -> None:
|
||||
"""Emitted when a download encounters an error"""
|
||||
self.dispatch(DownloadErrorEvent.build(job))
|
||||
|
||||
# endregion
|
||||
|
||||
# region Model loading
|
||||
|
||||
def emit_model_load_started(self, config: "AnyModelConfig", submodel_type: Optional["SubModelType"] = None) -> None:
|
||||
"""Emitted when a model load is started."""
|
||||
self.dispatch(ModelLoadStartedEvent.build(config, submodel_type))
|
||||
|
||||
def emit_model_load_complete(
|
||||
self, config: "AnyModelConfig", submodel_type: Optional["SubModelType"] = None
|
||||
) -> None:
|
||||
"""Emitted when a model load is complete."""
|
||||
self.dispatch(ModelLoadCompleteEvent.build(config, submodel_type))
|
||||
|
||||
# endregion
|
||||
|
||||
# region Model install
|
||||
|
||||
def emit_model_install_download_started(self, job: "ModelInstallJob") -> None:
|
||||
"""Emitted at intervals while the install job is started (remote models only)."""
|
||||
self.dispatch(ModelInstallDownloadStartedEvent.build(job))
|
||||
|
||||
def emit_model_install_download_progress(self, job: "ModelInstallJob") -> None:
|
||||
"""Emitted at intervals while the install job is in progress (remote models only)."""
|
||||
self.dispatch(ModelInstallDownloadProgressEvent.build(job))
|
||||
|
||||
def emit_model_install_downloads_complete(self, job: "ModelInstallJob") -> None:
|
||||
self.dispatch(ModelInstallDownloadsCompleteEvent.build(job))
|
||||
|
||||
def emit_model_install_started(self, job: "ModelInstallJob") -> None:
|
||||
"""Emitted once when an install job is started (after any download)."""
|
||||
self.dispatch(ModelInstallStartedEvent.build(job))
|
||||
|
||||
def emit_model_install_complete(self, job: "ModelInstallJob") -> None:
|
||||
"""Emitted when an install job is completed successfully."""
|
||||
self.dispatch(ModelInstallCompleteEvent.build(job))
|
||||
|
||||
def emit_model_install_cancelled(self, job: "ModelInstallJob") -> None:
|
||||
"""Emitted when an install job is cancelled."""
|
||||
self.dispatch(ModelInstallCancelledEvent.build(job))
|
||||
|
||||
def emit_model_install_error(self, job: "ModelInstallJob") -> None:
|
||||
"""Emitted when an install job encounters an exception."""
|
||||
self.dispatch(ModelInstallErrorEvent.build(job))
|
||||
|
||||
# endregion
|
||||
|
||||
# region Bulk image download
|
||||
|
||||
def emit_bulk_download_started(
|
||||
self,
|
||||
bulk_download_id: str,
|
||||
bulk_download_item_id: str,
|
||||
bulk_download_item_name: str,
|
||||
user_id: str = "system",
|
||||
) -> None:
|
||||
"""Emitted when a bulk image download is started"""
|
||||
self.dispatch(
|
||||
BulkDownloadStartedEvent.build(bulk_download_id, bulk_download_item_id, bulk_download_item_name, user_id)
|
||||
)
|
||||
|
||||
def emit_bulk_download_complete(
|
||||
self,
|
||||
bulk_download_id: str,
|
||||
bulk_download_item_id: str,
|
||||
bulk_download_item_name: str,
|
||||
user_id: str = "system",
|
||||
) -> None:
|
||||
"""Emitted when a bulk image download is complete"""
|
||||
self.dispatch(
|
||||
BulkDownloadCompleteEvent.build(bulk_download_id, bulk_download_item_id, bulk_download_item_name, user_id)
|
||||
)
|
||||
|
||||
def emit_bulk_download_error(
|
||||
self,
|
||||
bulk_download_id: str,
|
||||
bulk_download_item_id: str,
|
||||
bulk_download_item_name: str,
|
||||
error: str,
|
||||
user_id: str = "system",
|
||||
) -> None:
|
||||
"""Emitted when a bulk image download has an error"""
|
||||
self.dispatch(
|
||||
BulkDownloadErrorEvent.build(
|
||||
bulk_download_id, bulk_download_item_id, bulk_download_item_name, error, user_id
|
||||
)
|
||||
)
|
||||
|
||||
# endregion
|
||||
@@ -0,0 +1,779 @@
|
||||
from typing import TYPE_CHECKING, Any, ClassVar, Coroutine, Generic, Optional, Protocol, TypeAlias, TypeVar
|
||||
|
||||
from fastapi_events.handlers.local import local_handler
|
||||
from fastapi_events.registry.payload_schema import registry as payload_schema
|
||||
from pydantic import BaseModel, ConfigDict, Field
|
||||
|
||||
from invokeai.app.services.model_install.model_install_common import ModelInstallJob, ModelSource
|
||||
from invokeai.app.services.session_processor.session_processor_common import ProgressImage
|
||||
from invokeai.app.services.session_queue.session_queue_common import (
|
||||
QUEUE_ITEM_STATUS,
|
||||
BatchStatus,
|
||||
EnqueueBatchResult,
|
||||
RetryItemsResult,
|
||||
SessionQueueItem,
|
||||
SessionQueueStatus,
|
||||
)
|
||||
from invokeai.app.services.shared.graph import AnyInvocation, AnyInvocationOutput
|
||||
from invokeai.app.util.misc import get_timestamp
|
||||
from invokeai.backend.model_manager.configs.factory import AnyModelConfig
|
||||
from invokeai.backend.model_manager.taxonomy import SubModelType
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from invokeai.app.services.download.download_base import DownloadJob
|
||||
from invokeai.app.services.model_install.model_install_common import ModelInstallJob, ModelSource
|
||||
|
||||
|
||||
class EventBase(BaseModel):
|
||||
"""Base class for all events. All events must inherit from this class.
|
||||
|
||||
Events must define a class attribute `__event_name__` to identify the event.
|
||||
|
||||
All other attributes should be defined as normal for a pydantic model.
|
||||
|
||||
A timestamp is automatically added to the event when it is created.
|
||||
"""
|
||||
|
||||
__event_name__: ClassVar[str]
|
||||
timestamp: int = Field(description="The timestamp of the event", default_factory=get_timestamp)
|
||||
|
||||
model_config = ConfigDict(json_schema_serialization_defaults_required=True)
|
||||
|
||||
@classmethod
|
||||
def get_events(cls) -> set[type["EventBase"]]:
|
||||
"""Get a set of all event models."""
|
||||
|
||||
event_subclasses: set[type["EventBase"]] = set()
|
||||
for subclass in cls.__subclasses__():
|
||||
# We only want to include subclasses that are event models, not intermediary classes
|
||||
if hasattr(subclass, "__event_name__"):
|
||||
event_subclasses.add(subclass)
|
||||
event_subclasses.update(subclass.get_events())
|
||||
|
||||
return event_subclasses
|
||||
|
||||
|
||||
TEvent = TypeVar("TEvent", bound=EventBase, contravariant=True)
|
||||
|
||||
FastAPIEvent: TypeAlias = tuple[str, TEvent]
|
||||
"""
|
||||
A tuple representing a `fastapi-events` event, with the event name and payload.
|
||||
Provide a generic type to `TEvent` to specify the payload type.
|
||||
"""
|
||||
|
||||
|
||||
class FastAPIEventFunc(Protocol, Generic[TEvent]):
|
||||
def __call__(self, event: FastAPIEvent[TEvent]) -> Optional[Coroutine[Any, Any, None]]: ...
|
||||
|
||||
|
||||
def register_events(events: set[type[TEvent]] | type[TEvent], func: FastAPIEventFunc[TEvent]) -> None:
|
||||
"""Register a function to handle specific events.
|
||||
|
||||
:param events: An event or set of events to handle
|
||||
:param func: The function to handle the events
|
||||
"""
|
||||
events = events if isinstance(events, set) else {events}
|
||||
for event in events:
|
||||
assert hasattr(event, "__event_name__")
|
||||
local_handler.register(event_name=event.__event_name__, _func=func) # pyright: ignore [reportUnknownMemberType, reportUnknownArgumentType, reportAttributeAccessIssue]
|
||||
|
||||
|
||||
class QueueEventBase(EventBase):
|
||||
"""Base class for queue events"""
|
||||
|
||||
queue_id: str = Field(description="The ID of the queue")
|
||||
|
||||
|
||||
class QueueItemEventBase(QueueEventBase):
|
||||
"""Base class for queue item events"""
|
||||
|
||||
item_id: int = Field(description="The ID of the queue item")
|
||||
batch_id: str = Field(description="The ID of the queue batch")
|
||||
origin: str | None = Field(default=None, description="The origin of the queue item")
|
||||
destination: str | None = Field(default=None, description="The destination of the queue item")
|
||||
user_id: str = Field(default="system", description="The ID of the user who created the queue item")
|
||||
|
||||
|
||||
class InvocationEventBase(QueueItemEventBase):
|
||||
"""Base class for invocation events"""
|
||||
|
||||
session_id: str = Field(description="The ID of the session (aka graph execution state)")
|
||||
queue_id: str = Field(description="The ID of the queue")
|
||||
session_id: str = Field(description="The ID of the session (aka graph execution state)")
|
||||
invocation: AnyInvocation = Field(description="The ID of the invocation")
|
||||
invocation_source_id: str = Field(description="The ID of the prepared invocation's source node")
|
||||
|
||||
|
||||
@payload_schema.register
|
||||
class InvocationStartedEvent(InvocationEventBase):
|
||||
"""Event model for invocation_started"""
|
||||
|
||||
__event_name__ = "invocation_started"
|
||||
|
||||
@classmethod
|
||||
def build(cls, queue_item: SessionQueueItem, invocation: AnyInvocation) -> "InvocationStartedEvent":
|
||||
return cls(
|
||||
queue_id=queue_item.queue_id,
|
||||
item_id=queue_item.item_id,
|
||||
batch_id=queue_item.batch_id,
|
||||
origin=queue_item.origin,
|
||||
destination=queue_item.destination,
|
||||
user_id=queue_item.user_id,
|
||||
session_id=queue_item.session_id,
|
||||
invocation=invocation,
|
||||
invocation_source_id=queue_item.session.prepared_source_mapping[invocation.id],
|
||||
)
|
||||
|
||||
|
||||
@payload_schema.register
|
||||
class InvocationProgressEvent(InvocationEventBase):
|
||||
"""Event model for invocation_progress"""
|
||||
|
||||
__event_name__ = "invocation_progress"
|
||||
|
||||
message: str = Field(description="A message to display")
|
||||
percentage: float | None = Field(
|
||||
default=None, ge=0, le=1, description="The percentage of the progress (omit to indicate indeterminate progress)"
|
||||
)
|
||||
image: ProgressImage | None = Field(
|
||||
default=None, description="An image representing the current state of the progress"
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def build(
|
||||
cls,
|
||||
queue_item: SessionQueueItem,
|
||||
invocation: AnyInvocation,
|
||||
message: str,
|
||||
percentage: float | None = None,
|
||||
image: ProgressImage | None = None,
|
||||
) -> "InvocationProgressEvent":
|
||||
return cls(
|
||||
queue_id=queue_item.queue_id,
|
||||
item_id=queue_item.item_id,
|
||||
batch_id=queue_item.batch_id,
|
||||
origin=queue_item.origin,
|
||||
destination=queue_item.destination,
|
||||
user_id=queue_item.user_id,
|
||||
session_id=queue_item.session_id,
|
||||
invocation=invocation,
|
||||
invocation_source_id=queue_item.session.prepared_source_mapping[invocation.id],
|
||||
percentage=percentage,
|
||||
image=image,
|
||||
message=message,
|
||||
)
|
||||
|
||||
|
||||
@payload_schema.register
|
||||
class InvocationCompleteEvent(InvocationEventBase):
|
||||
"""Event model for invocation_complete"""
|
||||
|
||||
__event_name__ = "invocation_complete"
|
||||
|
||||
result: AnyInvocationOutput = Field(description="The result of the invocation")
|
||||
|
||||
@classmethod
|
||||
def build(
|
||||
cls, queue_item: SessionQueueItem, invocation: AnyInvocation, result: AnyInvocationOutput
|
||||
) -> "InvocationCompleteEvent":
|
||||
return cls(
|
||||
queue_id=queue_item.queue_id,
|
||||
item_id=queue_item.item_id,
|
||||
batch_id=queue_item.batch_id,
|
||||
origin=queue_item.origin,
|
||||
destination=queue_item.destination,
|
||||
user_id=queue_item.user_id,
|
||||
session_id=queue_item.session_id,
|
||||
invocation=invocation,
|
||||
invocation_source_id=queue_item.session.prepared_source_mapping[invocation.id],
|
||||
result=result,
|
||||
)
|
||||
|
||||
|
||||
@payload_schema.register
|
||||
class InvocationErrorEvent(InvocationEventBase):
|
||||
"""Event model for invocation_error"""
|
||||
|
||||
__event_name__ = "invocation_error"
|
||||
|
||||
error_type: str = Field(description="The error type")
|
||||
error_message: str = Field(description="The error message")
|
||||
error_traceback: str = Field(description="The error traceback")
|
||||
|
||||
@classmethod
|
||||
def build(
|
||||
cls,
|
||||
queue_item: SessionQueueItem,
|
||||
invocation: AnyInvocation,
|
||||
error_type: str,
|
||||
error_message: str,
|
||||
error_traceback: str,
|
||||
) -> "InvocationErrorEvent":
|
||||
return cls(
|
||||
queue_id=queue_item.queue_id,
|
||||
item_id=queue_item.item_id,
|
||||
batch_id=queue_item.batch_id,
|
||||
origin=queue_item.origin,
|
||||
destination=queue_item.destination,
|
||||
user_id=queue_item.user_id,
|
||||
session_id=queue_item.session_id,
|
||||
invocation=invocation,
|
||||
invocation_source_id=queue_item.session.prepared_source_mapping[invocation.id],
|
||||
error_type=error_type,
|
||||
error_message=error_message,
|
||||
error_traceback=error_traceback,
|
||||
)
|
||||
|
||||
|
||||
@payload_schema.register
|
||||
class QueueItemStatusChangedEvent(QueueItemEventBase):
|
||||
"""Event model for queue_item_status_changed"""
|
||||
|
||||
__event_name__ = "queue_item_status_changed"
|
||||
|
||||
status: QUEUE_ITEM_STATUS = Field(description="The new status of the queue item")
|
||||
status_sequence: int | None = Field(
|
||||
default=None,
|
||||
description="A monotonically increasing version for this queue item's visible status lifecycle",
|
||||
)
|
||||
error_type: Optional[str] = Field(default=None, description="The error type, if any")
|
||||
error_message: Optional[str] = Field(default=None, description="The error message, if any")
|
||||
error_traceback: Optional[str] = Field(default=None, description="The error traceback, if any")
|
||||
created_at: str = Field(description="The timestamp when the queue item was created")
|
||||
updated_at: str = Field(description="The timestamp when the queue item was last updated")
|
||||
started_at: Optional[str] = Field(default=None, description="The timestamp when the queue item was started")
|
||||
completed_at: Optional[str] = Field(default=None, description="The timestamp when the queue item was completed")
|
||||
batch_status: BatchStatus = Field(description="The status of the batch")
|
||||
queue_status: SessionQueueStatus = Field(description="The status of the queue")
|
||||
session_id: str = Field(description="The ID of the session (aka graph execution state)")
|
||||
|
||||
@classmethod
|
||||
def build(
|
||||
cls, queue_item: SessionQueueItem, batch_status: BatchStatus, queue_status: SessionQueueStatus
|
||||
) -> "QueueItemStatusChangedEvent":
|
||||
return cls(
|
||||
queue_id=queue_item.queue_id,
|
||||
item_id=queue_item.item_id,
|
||||
batch_id=queue_item.batch_id,
|
||||
origin=queue_item.origin,
|
||||
destination=queue_item.destination,
|
||||
user_id=queue_item.user_id,
|
||||
session_id=queue_item.session_id,
|
||||
status=queue_item.status,
|
||||
status_sequence=queue_item.status_sequence,
|
||||
error_type=queue_item.error_type,
|
||||
error_message=queue_item.error_message,
|
||||
error_traceback=queue_item.error_traceback,
|
||||
created_at=str(queue_item.created_at),
|
||||
updated_at=str(queue_item.updated_at),
|
||||
started_at=str(queue_item.started_at) if queue_item.started_at else None,
|
||||
completed_at=str(queue_item.completed_at) if queue_item.completed_at else None,
|
||||
batch_status=batch_status,
|
||||
queue_status=queue_status,
|
||||
)
|
||||
|
||||
|
||||
@payload_schema.register
|
||||
class BatchEnqueuedEvent(QueueEventBase):
|
||||
"""Event model for batch_enqueued"""
|
||||
|
||||
__event_name__ = "batch_enqueued"
|
||||
|
||||
batch_id: str = Field(description="The ID of the batch")
|
||||
enqueued: int = Field(description="The number of invocations enqueued")
|
||||
requested: int = Field(
|
||||
description="The number of invocations initially requested to be enqueued (may be less than enqueued if queue was full)"
|
||||
)
|
||||
priority: int = Field(description="The priority of the batch")
|
||||
origin: str | None = Field(default=None, description="The origin of the batch")
|
||||
user_id: str = Field(default="system", description="The ID of the user who enqueued the batch")
|
||||
|
||||
@classmethod
|
||||
def build(cls, enqueue_result: EnqueueBatchResult, user_id: str = "system") -> "BatchEnqueuedEvent":
|
||||
return cls(
|
||||
queue_id=enqueue_result.queue_id,
|
||||
batch_id=enqueue_result.batch.batch_id,
|
||||
origin=enqueue_result.batch.origin,
|
||||
enqueued=enqueue_result.enqueued,
|
||||
requested=enqueue_result.requested,
|
||||
priority=enqueue_result.priority,
|
||||
user_id=user_id,
|
||||
)
|
||||
|
||||
|
||||
@payload_schema.register
|
||||
class QueueItemsRetriedEvent(QueueEventBase):
|
||||
"""Event model for queue_items_retried"""
|
||||
|
||||
__event_name__ = "queue_items_retried"
|
||||
|
||||
retried_item_ids: list[int] = Field(description="The IDs of the queue items that were retried")
|
||||
user_ids: list[str] = Field(description="The IDs of the users who own the retried root queue items")
|
||||
retried_item_ids_by_user: dict[str, list[int]] = Field(
|
||||
description="The retried root queue item IDs keyed by owner user ID."
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def build(
|
||||
cls, retry_result: RetryItemsResult, user_ids: list[str], retried_item_ids_by_user: dict[str, list[int]]
|
||||
) -> "QueueItemsRetriedEvent":
|
||||
return cls(
|
||||
queue_id=retry_result.queue_id,
|
||||
retried_item_ids=retry_result.retried_item_ids,
|
||||
user_ids=user_ids,
|
||||
retried_item_ids_by_user=retried_item_ids_by_user,
|
||||
)
|
||||
|
||||
|
||||
@payload_schema.register
|
||||
class QueueClearedEvent(QueueEventBase):
|
||||
"""Event model for queue_cleared"""
|
||||
|
||||
__event_name__ = "queue_cleared"
|
||||
|
||||
user_id: str | None = Field(
|
||||
default=None,
|
||||
description="The ID of the user whose queue items were cleared, or None if all users' items were cleared",
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def build(cls, queue_id: str, user_id: str | None = None) -> "QueueClearedEvent":
|
||||
return cls(queue_id=queue_id, user_id=user_id)
|
||||
|
||||
|
||||
class WorkflowEventBase(EventBase):
|
||||
"""Base class for workflow library CRUD events."""
|
||||
|
||||
workflow_id: str = Field(description="The ID of the workflow")
|
||||
user_id: str = Field(description="The owner of the workflow")
|
||||
|
||||
|
||||
@payload_schema.register
|
||||
class WorkflowCreatedEvent(WorkflowEventBase):
|
||||
"""Event model for workflow_created"""
|
||||
|
||||
__event_name__ = "workflow_created"
|
||||
|
||||
is_public: bool = Field(description="Whether the workflow is shared with all users")
|
||||
|
||||
@classmethod
|
||||
def build(cls, workflow_id: str, user_id: str, is_public: bool) -> "WorkflowCreatedEvent":
|
||||
return cls(workflow_id=workflow_id, user_id=user_id, is_public=is_public)
|
||||
|
||||
|
||||
@payload_schema.register
|
||||
class WorkflowUpdatedEvent(WorkflowEventBase):
|
||||
"""Event model for workflow_updated"""
|
||||
|
||||
__event_name__ = "workflow_updated"
|
||||
|
||||
old_is_public: bool = Field(description="Whether the workflow was shared before the update")
|
||||
new_is_public: bool = Field(description="Whether the workflow is shared after the update")
|
||||
|
||||
@classmethod
|
||||
def build(cls, workflow_id: str, user_id: str, old_is_public: bool, new_is_public: bool) -> "WorkflowUpdatedEvent":
|
||||
return cls(
|
||||
workflow_id=workflow_id,
|
||||
user_id=user_id,
|
||||
old_is_public=old_is_public,
|
||||
new_is_public=new_is_public,
|
||||
)
|
||||
|
||||
|
||||
@payload_schema.register
|
||||
class WorkflowDeletedEvent(WorkflowEventBase):
|
||||
"""Event model for workflow_deleted"""
|
||||
|
||||
__event_name__ = "workflow_deleted"
|
||||
|
||||
is_public: bool = Field(description="Whether the workflow was shared when it was deleted")
|
||||
|
||||
@classmethod
|
||||
def build(cls, workflow_id: str, user_id: str, is_public: bool) -> "WorkflowDeletedEvent":
|
||||
return cls(workflow_id=workflow_id, user_id=user_id, is_public=is_public)
|
||||
|
||||
|
||||
@payload_schema.register
|
||||
class WorkflowAccessRevokedEvent(WorkflowEventBase):
|
||||
"""Event model for workflow_access_revoked."""
|
||||
|
||||
__event_name__ = "workflow_access_revoked"
|
||||
|
||||
@classmethod
|
||||
def build(cls, workflow_id: str, user_id: str) -> "WorkflowAccessRevokedEvent":
|
||||
return cls(workflow_id=workflow_id, user_id=user_id)
|
||||
|
||||
|
||||
class DownloadEventBase(EventBase):
|
||||
"""Base class for events associated with a download"""
|
||||
|
||||
source: str = Field(description="The source of the download")
|
||||
|
||||
|
||||
@payload_schema.register
|
||||
class DownloadStartedEvent(DownloadEventBase):
|
||||
"""Event model for download_started"""
|
||||
|
||||
__event_name__ = "download_started"
|
||||
|
||||
download_path: str = Field(description="The local path where the download is saved")
|
||||
|
||||
@classmethod
|
||||
def build(cls, job: "DownloadJob") -> "DownloadStartedEvent":
|
||||
assert job.download_path
|
||||
return cls(source=str(job.source), download_path=job.download_path.as_posix())
|
||||
|
||||
|
||||
@payload_schema.register
|
||||
class DownloadProgressEvent(DownloadEventBase):
|
||||
"""Event model for download_progress"""
|
||||
|
||||
__event_name__ = "download_progress"
|
||||
|
||||
download_path: str = Field(description="The local path where the download is saved")
|
||||
current_bytes: int = Field(description="The number of bytes downloaded so far")
|
||||
total_bytes: int = Field(description="The total number of bytes to be downloaded")
|
||||
|
||||
@classmethod
|
||||
def build(cls, job: "DownloadJob") -> "DownloadProgressEvent":
|
||||
assert job.download_path
|
||||
return cls(
|
||||
source=str(job.source),
|
||||
download_path=job.download_path.as_posix(),
|
||||
current_bytes=job.bytes,
|
||||
total_bytes=job.total_bytes,
|
||||
)
|
||||
|
||||
|
||||
@payload_schema.register
|
||||
class DownloadCompleteEvent(DownloadEventBase):
|
||||
"""Event model for download_complete"""
|
||||
|
||||
__event_name__ = "download_complete"
|
||||
|
||||
download_path: str = Field(description="The local path where the download is saved")
|
||||
total_bytes: int = Field(description="The total number of bytes downloaded")
|
||||
|
||||
@classmethod
|
||||
def build(cls, job: "DownloadJob") -> "DownloadCompleteEvent":
|
||||
assert job.download_path
|
||||
return cls(source=str(job.source), download_path=job.download_path.as_posix(), total_bytes=job.total_bytes)
|
||||
|
||||
|
||||
@payload_schema.register
|
||||
class DownloadCancelledEvent(DownloadEventBase):
|
||||
"""Event model for download_cancelled"""
|
||||
|
||||
__event_name__ = "download_cancelled"
|
||||
|
||||
@classmethod
|
||||
def build(cls, job: "DownloadJob") -> "DownloadCancelledEvent":
|
||||
return cls(source=str(job.source))
|
||||
|
||||
|
||||
@payload_schema.register
|
||||
class DownloadPausedEvent(DownloadEventBase):
|
||||
"""Event model for download_paused"""
|
||||
|
||||
__event_name__ = "download_paused"
|
||||
|
||||
@classmethod
|
||||
def build(cls, job: "DownloadJob") -> "DownloadPausedEvent":
|
||||
return cls(source=str(job.source))
|
||||
|
||||
|
||||
@payload_schema.register
|
||||
class DownloadErrorEvent(DownloadEventBase):
|
||||
"""Event model for download_error"""
|
||||
|
||||
__event_name__ = "download_error"
|
||||
|
||||
error_type: str = Field(description="The type of error")
|
||||
error: str = Field(description="The error message")
|
||||
|
||||
@classmethod
|
||||
def build(cls, job: "DownloadJob") -> "DownloadErrorEvent":
|
||||
assert job.error_type
|
||||
assert job.error
|
||||
return cls(source=str(job.source), error_type=job.error_type, error=job.error)
|
||||
|
||||
|
||||
class ModelEventBase(EventBase):
|
||||
"""Base class for events associated with a model"""
|
||||
|
||||
|
||||
@payload_schema.register
|
||||
class ModelLoadStartedEvent(ModelEventBase):
|
||||
"""Event model for model_load_started"""
|
||||
|
||||
__event_name__ = "model_load_started"
|
||||
|
||||
config: AnyModelConfig = Field(description="The model's config")
|
||||
submodel_type: Optional[SubModelType] = Field(default=None, description="The submodel type, if any")
|
||||
|
||||
@classmethod
|
||||
def build(cls, config: AnyModelConfig, submodel_type: Optional[SubModelType] = None) -> "ModelLoadStartedEvent":
|
||||
return cls(config=config, submodel_type=submodel_type)
|
||||
|
||||
|
||||
@payload_schema.register
|
||||
class ModelLoadCompleteEvent(ModelEventBase):
|
||||
"""Event model for model_load_complete"""
|
||||
|
||||
__event_name__ = "model_load_complete"
|
||||
|
||||
config: AnyModelConfig = Field(description="The model's config")
|
||||
submodel_type: Optional[SubModelType] = Field(default=None, description="The submodel type, if any")
|
||||
|
||||
@classmethod
|
||||
def build(cls, config: AnyModelConfig, submodel_type: Optional[SubModelType] = None) -> "ModelLoadCompleteEvent":
|
||||
return cls(config=config, submodel_type=submodel_type)
|
||||
|
||||
|
||||
@payload_schema.register
|
||||
class ModelInstallDownloadStartedEvent(ModelEventBase):
|
||||
"""Event model for model_install_download_started"""
|
||||
|
||||
__event_name__ = "model_install_download_started"
|
||||
|
||||
id: int = Field(description="The ID of the install job")
|
||||
source: ModelSource = Field(description="Source of the model; local path, repo_id or url")
|
||||
local_path: str = Field(description="Where model is downloading to")
|
||||
bytes: int = Field(description="Number of bytes downloaded so far")
|
||||
total_bytes: int = Field(description="Total size of download, including all files")
|
||||
parts: list[dict[str, int | str]] = Field(
|
||||
description="Progress of downloading URLs that comprise the model, if any"
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def build(cls, job: "ModelInstallJob") -> "ModelInstallDownloadStartedEvent":
|
||||
parts: list[dict[str, str | int]] = [
|
||||
{
|
||||
"url": str(x.source),
|
||||
"local_path": str(x.download_path),
|
||||
"bytes": x.bytes,
|
||||
"total_bytes": x.total_bytes,
|
||||
}
|
||||
for x in job.download_parts
|
||||
]
|
||||
return cls(
|
||||
id=job.id,
|
||||
source=job.source,
|
||||
local_path=job.local_path.as_posix(),
|
||||
parts=parts,
|
||||
bytes=job.bytes,
|
||||
total_bytes=job.total_bytes,
|
||||
)
|
||||
|
||||
|
||||
@payload_schema.register
|
||||
class ModelInstallDownloadProgressEvent(ModelEventBase):
|
||||
"""Event model for model_install_download_progress"""
|
||||
|
||||
__event_name__ = "model_install_download_progress"
|
||||
|
||||
id: int = Field(description="The ID of the install job")
|
||||
source: ModelSource = Field(description="Source of the model; local path, repo_id or url")
|
||||
local_path: str = Field(description="Where model is downloading to")
|
||||
bytes: int = Field(description="Number of bytes downloaded so far")
|
||||
total_bytes: int = Field(description="Total size of download, including all files")
|
||||
parts: list[dict[str, int | str]] = Field(
|
||||
description="Progress of downloading URLs that comprise the model, if any"
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def build(cls, job: "ModelInstallJob") -> "ModelInstallDownloadProgressEvent":
|
||||
parts: list[dict[str, str | int]] = [
|
||||
{
|
||||
"url": str(x.source),
|
||||
"local_path": str(x.download_path),
|
||||
"bytes": x.bytes,
|
||||
"total_bytes": x.total_bytes,
|
||||
}
|
||||
for x in job.download_parts
|
||||
]
|
||||
return cls(
|
||||
id=job.id,
|
||||
source=job.source,
|
||||
local_path=job.local_path.as_posix(),
|
||||
parts=parts,
|
||||
bytes=job.bytes,
|
||||
total_bytes=job.total_bytes,
|
||||
)
|
||||
|
||||
|
||||
@payload_schema.register
|
||||
class ModelInstallDownloadsCompleteEvent(ModelEventBase):
|
||||
"""Emitted once when an install job becomes active."""
|
||||
|
||||
__event_name__ = "model_install_downloads_complete"
|
||||
|
||||
id: int = Field(description="The ID of the install job")
|
||||
source: ModelSource = Field(description="Source of the model; local path, repo_id or url")
|
||||
|
||||
@classmethod
|
||||
def build(cls, job: "ModelInstallJob") -> "ModelInstallDownloadsCompleteEvent":
|
||||
return cls(id=job.id, source=job.source)
|
||||
|
||||
|
||||
@payload_schema.register
|
||||
class ModelInstallStartedEvent(ModelEventBase):
|
||||
"""Event model for model_install_started"""
|
||||
|
||||
__event_name__ = "model_install_started"
|
||||
|
||||
id: int = Field(description="The ID of the install job")
|
||||
source: ModelSource = Field(description="Source of the model; local path, repo_id or url")
|
||||
|
||||
@classmethod
|
||||
def build(cls, job: "ModelInstallJob") -> "ModelInstallStartedEvent":
|
||||
return cls(id=job.id, source=job.source)
|
||||
|
||||
|
||||
@payload_schema.register
|
||||
class ModelInstallCompleteEvent(ModelEventBase):
|
||||
"""Event model for model_install_complete"""
|
||||
|
||||
__event_name__ = "model_install_complete"
|
||||
|
||||
id: int = Field(description="The ID of the install job")
|
||||
source: ModelSource = Field(description="Source of the model; local path, repo_id or url")
|
||||
key: str = Field(description="Model config record key")
|
||||
total_bytes: Optional[int] = Field(description="Size of the model (may be None for installation of a local path)")
|
||||
config: AnyModelConfig = Field(description="The installed model's config")
|
||||
|
||||
@classmethod
|
||||
def build(cls, job: "ModelInstallJob") -> "ModelInstallCompleteEvent":
|
||||
assert job.config_out is not None
|
||||
return cls(
|
||||
id=job.id,
|
||||
source=job.source,
|
||||
key=(job.config_out.key),
|
||||
total_bytes=job.total_bytes,
|
||||
config=job.config_out,
|
||||
)
|
||||
|
||||
|
||||
@payload_schema.register
|
||||
class ModelInstallCancelledEvent(ModelEventBase):
|
||||
"""Event model for model_install_cancelled"""
|
||||
|
||||
__event_name__ = "model_install_cancelled"
|
||||
|
||||
id: int = Field(description="The ID of the install job")
|
||||
source: ModelSource = Field(description="Source of the model; local path, repo_id or url")
|
||||
|
||||
@classmethod
|
||||
def build(cls, job: "ModelInstallJob") -> "ModelInstallCancelledEvent":
|
||||
return cls(id=job.id, source=job.source)
|
||||
|
||||
|
||||
@payload_schema.register
|
||||
class ModelInstallErrorEvent(ModelEventBase):
|
||||
"""Event model for model_install_error"""
|
||||
|
||||
__event_name__ = "model_install_error"
|
||||
|
||||
id: int = Field(description="The ID of the install job")
|
||||
source: ModelSource = Field(description="Source of the model; local path, repo_id or url")
|
||||
error_type: str = Field(description="The name of the exception")
|
||||
error: str = Field(description="A text description of the exception")
|
||||
|
||||
@classmethod
|
||||
def build(cls, job: "ModelInstallJob") -> "ModelInstallErrorEvent":
|
||||
assert job.error_type is not None
|
||||
assert job.error is not None
|
||||
return cls(id=job.id, source=job.source, error_type=job.error_type, error=job.error)
|
||||
|
||||
|
||||
class BulkDownloadEventBase(EventBase):
|
||||
"""Base class for events associated with a bulk image download"""
|
||||
|
||||
bulk_download_id: str = Field(description="The ID of the bulk image download")
|
||||
bulk_download_item_id: str = Field(description="The ID of the bulk image download item")
|
||||
bulk_download_item_name: str = Field(description="The name of the bulk image download item")
|
||||
user_id: str = Field(default="system", description="The ID of the user who initiated the download")
|
||||
|
||||
|
||||
@payload_schema.register
|
||||
class BulkDownloadStartedEvent(BulkDownloadEventBase):
|
||||
"""Event model for bulk_download_started"""
|
||||
|
||||
__event_name__ = "bulk_download_started"
|
||||
|
||||
@classmethod
|
||||
def build(
|
||||
cls,
|
||||
bulk_download_id: str,
|
||||
bulk_download_item_id: str,
|
||||
bulk_download_item_name: str,
|
||||
user_id: str = "system",
|
||||
) -> "BulkDownloadStartedEvent":
|
||||
return cls(
|
||||
bulk_download_id=bulk_download_id,
|
||||
bulk_download_item_id=bulk_download_item_id,
|
||||
bulk_download_item_name=bulk_download_item_name,
|
||||
user_id=user_id,
|
||||
)
|
||||
|
||||
|
||||
@payload_schema.register
|
||||
class BulkDownloadCompleteEvent(BulkDownloadEventBase):
|
||||
"""Event model for bulk_download_complete"""
|
||||
|
||||
__event_name__ = "bulk_download_complete"
|
||||
|
||||
@classmethod
|
||||
def build(
|
||||
cls,
|
||||
bulk_download_id: str,
|
||||
bulk_download_item_id: str,
|
||||
bulk_download_item_name: str,
|
||||
user_id: str = "system",
|
||||
) -> "BulkDownloadCompleteEvent":
|
||||
return cls(
|
||||
bulk_download_id=bulk_download_id,
|
||||
bulk_download_item_id=bulk_download_item_id,
|
||||
bulk_download_item_name=bulk_download_item_name,
|
||||
user_id=user_id,
|
||||
)
|
||||
|
||||
|
||||
@payload_schema.register
|
||||
class BulkDownloadErrorEvent(BulkDownloadEventBase):
|
||||
"""Event model for bulk_download_error"""
|
||||
|
||||
__event_name__ = "bulk_download_error"
|
||||
|
||||
error: str = Field(description="The error message")
|
||||
|
||||
@classmethod
|
||||
def build(
|
||||
cls,
|
||||
bulk_download_id: str,
|
||||
bulk_download_item_id: str,
|
||||
bulk_download_item_name: str,
|
||||
error: str,
|
||||
user_id: str = "system",
|
||||
) -> "BulkDownloadErrorEvent":
|
||||
return cls(
|
||||
bulk_download_id=bulk_download_id,
|
||||
bulk_download_item_id=bulk_download_item_id,
|
||||
bulk_download_item_name=bulk_download_item_name,
|
||||
error=error,
|
||||
user_id=user_id,
|
||||
)
|
||||
|
||||
|
||||
@payload_schema.register
|
||||
class RecallParametersUpdatedEvent(QueueEventBase):
|
||||
"""Event model for recall_parameters_updated"""
|
||||
|
||||
__event_name__ = "recall_parameters_updated"
|
||||
|
||||
user_id: str = Field(description="The ID of the user whose recall parameters were updated")
|
||||
parameters: dict[str, Any] = Field(description="The recall parameters that were updated")
|
||||
|
||||
@classmethod
|
||||
def build(cls, queue_id: str, user_id: str, parameters: dict[str, Any]) -> "RecallParametersUpdatedEvent":
|
||||
return cls(queue_id=queue_id, user_id=user_id, parameters=parameters)
|
||||
@@ -0,0 +1,54 @@
|
||||
import asyncio
|
||||
import threading
|
||||
|
||||
from fastapi_events.dispatcher import dispatch
|
||||
|
||||
from invokeai.app.services.events.events_base import EventServiceBase
|
||||
from invokeai.app.services.events.events_common import EventBase
|
||||
|
||||
|
||||
class FastAPIEventService(EventServiceBase):
|
||||
def __init__(self, event_handler_id: int, loop: asyncio.AbstractEventLoop) -> None:
|
||||
self.event_handler_id = event_handler_id
|
||||
self._queue = asyncio.Queue[EventBase | None]()
|
||||
self._stop_event = threading.Event()
|
||||
self._loop = loop
|
||||
|
||||
# We need to store a reference to the task so it doesn't get GC'd
|
||||
# See: https://docs.python.org/3/library/asyncio-task.html#creating-tasks
|
||||
self._background_tasks: set[asyncio.Task[None]] = set()
|
||||
task = self._loop.create_task(self._dispatch_from_queue(stop_event=self._stop_event))
|
||||
self._background_tasks.add(task)
|
||||
task.add_done_callback(self._background_tasks.remove)
|
||||
|
||||
super().__init__()
|
||||
|
||||
def stop(self, *args, **kwargs):
|
||||
self._stop_event.set()
|
||||
self._loop.call_soon_threadsafe(self._queue.put_nowait, None)
|
||||
|
||||
def dispatch(self, event: EventBase) -> None:
|
||||
if self._loop.is_closed():
|
||||
# The event loop was closed during shutdown. Events can no longer be dispatched;
|
||||
# silently drop this one so the generation thread can wind down cleanly.
|
||||
return
|
||||
self._loop.call_soon_threadsafe(self._queue.put_nowait, event)
|
||||
|
||||
async def _dispatch_from_queue(self, stop_event: threading.Event):
|
||||
"""Get events on from the queue and dispatch them, from the correct thread"""
|
||||
while not stop_event.is_set():
|
||||
try:
|
||||
event = await self._queue.get()
|
||||
if not event: # Probably stopping
|
||||
continue
|
||||
# Leave the payloads as live pydantic models
|
||||
dispatch(event, middleware_id=self.event_handler_id, payload_schema_dump=False)
|
||||
|
||||
except asyncio.CancelledError as e:
|
||||
raise e # Raise a proper error
|
||||
except Exception:
|
||||
import logging
|
||||
|
||||
logging.getLogger("InvokeAI").error(
|
||||
f"Error dispatching event {getattr(event, '__event_name__', event)}", exc_info=True
|
||||
)
|
||||
@@ -0,0 +1,23 @@
|
||||
from invokeai.app.services.external_generation.external_generation_base import (
|
||||
ExternalGenerationServiceBase,
|
||||
ExternalProvider,
|
||||
)
|
||||
from invokeai.app.services.external_generation.external_generation_common import (
|
||||
ExternalGeneratedImage,
|
||||
ExternalGenerationRequest,
|
||||
ExternalGenerationResult,
|
||||
ExternalProviderStatus,
|
||||
ExternalReferenceImage,
|
||||
)
|
||||
from invokeai.app.services.external_generation.external_generation_default import ExternalGenerationService
|
||||
|
||||
__all__ = [
|
||||
"ExternalGenerationRequest",
|
||||
"ExternalGenerationResult",
|
||||
"ExternalGeneratedImage",
|
||||
"ExternalGenerationService",
|
||||
"ExternalGenerationServiceBase",
|
||||
"ExternalProvider",
|
||||
"ExternalProviderStatus",
|
||||
"ExternalReferenceImage",
|
||||
]
|
||||
@@ -0,0 +1,28 @@
|
||||
class ExternalGenerationError(Exception):
|
||||
"""Base error for external generation."""
|
||||
|
||||
|
||||
class ExternalProviderNotFoundError(ExternalGenerationError):
|
||||
"""Raised when no provider is registered for a model."""
|
||||
|
||||
|
||||
class ExternalProviderNotConfiguredError(ExternalGenerationError):
|
||||
"""Raised when a provider is missing required credentials."""
|
||||
|
||||
|
||||
class ExternalProviderCapabilityError(ExternalGenerationError):
|
||||
"""Raised when a request is not supported by provider capabilities."""
|
||||
|
||||
|
||||
class ExternalProviderRequestError(ExternalGenerationError):
|
||||
"""Raised when a provider rejects the request or returns an error."""
|
||||
|
||||
|
||||
class ExternalProviderRateLimitError(ExternalProviderRequestError):
|
||||
"""Raised when a provider returns HTTP 429 (rate limit exceeded)."""
|
||||
|
||||
retry_after: float | None
|
||||
|
||||
def __init__(self, message: str, retry_after: float | None = None) -> None:
|
||||
super().__init__(message)
|
||||
self.retry_after = retry_after
|
||||
@@ -0,0 +1,40 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
from logging import Logger
|
||||
|
||||
from invokeai.app.services.config import InvokeAIAppConfig
|
||||
from invokeai.app.services.external_generation.external_generation_common import (
|
||||
ExternalGenerationRequest,
|
||||
ExternalGenerationResult,
|
||||
ExternalProviderStatus,
|
||||
)
|
||||
|
||||
|
||||
class ExternalProvider(ABC):
|
||||
provider_id: str
|
||||
|
||||
def __init__(self, app_config: InvokeAIAppConfig, logger: Logger) -> None:
|
||||
self._app_config = app_config
|
||||
self._logger = logger
|
||||
|
||||
@abstractmethod
|
||||
def is_configured(self) -> bool:
|
||||
raise NotImplementedError
|
||||
|
||||
@abstractmethod
|
||||
def generate(self, request: ExternalGenerationRequest) -> ExternalGenerationResult:
|
||||
raise NotImplementedError
|
||||
|
||||
def get_status(self) -> ExternalProviderStatus:
|
||||
return ExternalProviderStatus(provider_id=self.provider_id, configured=self.is_configured())
|
||||
|
||||
|
||||
class ExternalGenerationServiceBase(ABC):
|
||||
@abstractmethod
|
||||
def generate(self, request: ExternalGenerationRequest) -> ExternalGenerationResult:
|
||||
raise NotImplementedError
|
||||
|
||||
@abstractmethod
|
||||
def get_provider_statuses(self) -> dict[str, ExternalProviderStatus]:
|
||||
raise NotImplementedError
|
||||
@@ -0,0 +1,52 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Any
|
||||
|
||||
from PIL.Image import Image as PILImageType
|
||||
|
||||
from invokeai.backend.model_manager.configs.external_api import ExternalApiModelConfig, ExternalGenerationMode
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ExternalReferenceImage:
|
||||
image: PILImageType
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ExternalGenerationRequest:
|
||||
model: ExternalApiModelConfig
|
||||
mode: ExternalGenerationMode
|
||||
prompt: str
|
||||
seed: int | None
|
||||
num_images: int
|
||||
width: int
|
||||
height: int
|
||||
image_size: str | None
|
||||
init_image: PILImageType | None
|
||||
mask_image: PILImageType | None
|
||||
reference_images: list[ExternalReferenceImage]
|
||||
metadata: dict[str, Any] | None
|
||||
provider_options: dict[str, Any] | None = None
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ExternalGeneratedImage:
|
||||
image: PILImageType
|
||||
seed: int | None = None
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ExternalGenerationResult:
|
||||
images: list[ExternalGeneratedImage]
|
||||
seed_used: int | None = None
|
||||
provider_request_id: str | None = None
|
||||
provider_metadata: dict[str, Any] | None = None
|
||||
content_filters: dict[str, str] | None = None
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ExternalProviderStatus:
|
||||
provider_id: str
|
||||
configured: bool
|
||||
message: str | None = None
|
||||
@@ -0,0 +1,369 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import dataclasses
|
||||
import time
|
||||
from logging import Logger
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from PIL import Image
|
||||
from PIL.Image import Image as PILImageType
|
||||
|
||||
from invokeai.app.services.external_generation.errors import (
|
||||
ExternalProviderCapabilityError,
|
||||
ExternalProviderNotConfiguredError,
|
||||
ExternalProviderNotFoundError,
|
||||
ExternalProviderRateLimitError,
|
||||
)
|
||||
from invokeai.app.services.external_generation.external_generation_base import (
|
||||
ExternalGenerationServiceBase,
|
||||
ExternalProvider,
|
||||
)
|
||||
from invokeai.app.services.external_generation.external_generation_common import (
|
||||
ExternalGeneratedImage,
|
||||
ExternalGenerationRequest,
|
||||
ExternalGenerationResult,
|
||||
ExternalProviderStatus,
|
||||
)
|
||||
from invokeai.backend.model_manager.configs.external_api import ExternalApiModelConfig, ExternalImageSize
|
||||
from invokeai.backend.model_manager.starter_models import STARTER_MODELS
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from invokeai.app.services.model_records import ModelRecordServiceBase
|
||||
|
||||
|
||||
class ExternalGenerationService(ExternalGenerationServiceBase):
|
||||
def __init__(
|
||||
self,
|
||||
providers: dict[str, ExternalProvider],
|
||||
logger: Logger,
|
||||
record_store: ModelRecordServiceBase | None = None,
|
||||
) -> None:
|
||||
self._providers = providers
|
||||
self._logger = logger
|
||||
self._record_store = record_store
|
||||
|
||||
def generate(self, request: ExternalGenerationRequest) -> ExternalGenerationResult:
|
||||
provider = self._providers.get(request.model.provider_id)
|
||||
if provider is None:
|
||||
raise ExternalProviderNotFoundError(f"No external provider registered for '{request.model.provider_id}'")
|
||||
|
||||
if not provider.is_configured():
|
||||
raise ExternalProviderNotConfiguredError(f"Provider '{request.model.provider_id}' is missing credentials")
|
||||
|
||||
request = self._refresh_model_capabilities(request)
|
||||
resize_to_original_inpaint_size = _get_resize_target_for_inpaint(request)
|
||||
request = self._bucket_request(request)
|
||||
request = self._drop_unsupported_capabilities(request)
|
||||
|
||||
self._validate_request(request)
|
||||
result = self._generate_with_retry(provider, request)
|
||||
|
||||
if resize_to_original_inpaint_size is None:
|
||||
return result
|
||||
|
||||
width, height = resize_to_original_inpaint_size
|
||||
return _resize_result_images(result, width, height)
|
||||
|
||||
_MAX_RETRIES = 3
|
||||
_DEFAULT_RETRY_DELAY = 10.0
|
||||
_MAX_RETRY_DELAY = 60.0
|
||||
|
||||
def _generate_with_retry(
|
||||
self, provider: ExternalProvider, request: ExternalGenerationRequest
|
||||
) -> ExternalGenerationResult:
|
||||
for attempt in range(self._MAX_RETRIES):
|
||||
try:
|
||||
return provider.generate(request)
|
||||
except ExternalProviderRateLimitError as exc:
|
||||
if attempt == self._MAX_RETRIES - 1:
|
||||
raise
|
||||
delay = min(exc.retry_after or self._DEFAULT_RETRY_DELAY, self._MAX_RETRY_DELAY)
|
||||
self._logger.warning(
|
||||
"Rate limited by %s (attempt %d/%d), retrying in %.0fs",
|
||||
request.model.provider_id,
|
||||
attempt + 1,
|
||||
self._MAX_RETRIES,
|
||||
delay,
|
||||
)
|
||||
time.sleep(delay)
|
||||
raise ExternalProviderRateLimitError("Rate limit exceeded after all retries")
|
||||
|
||||
def get_provider_statuses(self) -> dict[str, ExternalProviderStatus]:
|
||||
return {provider_id: provider.get_status() for provider_id, provider in self._providers.items()}
|
||||
|
||||
def _validate_request(self, request: ExternalGenerationRequest) -> None:
|
||||
capabilities = request.model.capabilities
|
||||
|
||||
self._logger.debug(
|
||||
"Validating external request provider=%s model=%s mode=%s supported=%s",
|
||||
request.model.provider_id,
|
||||
request.model.provider_model_id,
|
||||
request.mode,
|
||||
capabilities.modes,
|
||||
)
|
||||
|
||||
if request.mode not in capabilities.modes:
|
||||
raise ExternalProviderCapabilityError(f"Mode '{request.mode}' is not supported by {request.model.name}")
|
||||
|
||||
if request.reference_images and not capabilities.supports_reference_images:
|
||||
raise ExternalProviderCapabilityError(f"Reference images are not supported by {request.model.name}")
|
||||
|
||||
if capabilities.max_reference_images is not None:
|
||||
if len(request.reference_images) > capabilities.max_reference_images:
|
||||
raise ExternalProviderCapabilityError(
|
||||
f"{request.model.name} supports at most {capabilities.max_reference_images} reference images"
|
||||
)
|
||||
|
||||
if capabilities.max_images_per_request is not None and request.num_images > capabilities.max_images_per_request:
|
||||
raise ExternalProviderCapabilityError(
|
||||
f"{request.model.name} supports at most {capabilities.max_images_per_request} images per request"
|
||||
)
|
||||
|
||||
if capabilities.max_image_size is not None:
|
||||
if request.width > capabilities.max_image_size.width or request.height > capabilities.max_image_size.height:
|
||||
raise ExternalProviderCapabilityError(
|
||||
f"{request.model.name} supports a maximum size of {capabilities.max_image_size.width}x{capabilities.max_image_size.height}"
|
||||
)
|
||||
|
||||
if capabilities.allowed_aspect_ratios:
|
||||
aspect_ratio = _format_aspect_ratio(request.width, request.height)
|
||||
if aspect_ratio not in capabilities.allowed_aspect_ratios:
|
||||
size_ratio = None
|
||||
if capabilities.aspect_ratio_sizes:
|
||||
size_ratio = _ratio_for_size(request.width, request.height, capabilities.aspect_ratio_sizes)
|
||||
if size_ratio is None or size_ratio not in capabilities.allowed_aspect_ratios:
|
||||
ratio_label = size_ratio or aspect_ratio
|
||||
raise ExternalProviderCapabilityError(
|
||||
f"{request.model.name} does not support aspect ratio {ratio_label}"
|
||||
)
|
||||
|
||||
required_modes = capabilities.input_image_required_for or ["img2img", "inpaint"]
|
||||
if request.mode in required_modes and request.init_image is None:
|
||||
raise ExternalProviderCapabilityError(
|
||||
f"Mode '{request.mode}' requires an init image for {request.model.name}"
|
||||
)
|
||||
|
||||
if request.mode == "inpaint" and request.mask_image is None:
|
||||
raise ExternalProviderCapabilityError(
|
||||
f"Mode '{request.mode}' requires a mask image for {request.model.name}"
|
||||
)
|
||||
|
||||
def _drop_unsupported_capabilities(self, request: ExternalGenerationRequest) -> ExternalGenerationRequest:
|
||||
"""Silently drop request fields the selected model does not support so workflow-editor runs don't fail
|
||||
when users wire them in regardless."""
|
||||
capabilities = request.model.capabilities
|
||||
updates: dict[str, object] = {}
|
||||
|
||||
if request.seed is not None and not capabilities.supports_seed:
|
||||
self._logger.debug(
|
||||
"Dropping seed for %s: model does not support seed control",
|
||||
request.model.name,
|
||||
)
|
||||
updates["seed"] = None
|
||||
|
||||
if updates:
|
||||
return dataclasses.replace(request, **updates)
|
||||
return request
|
||||
|
||||
def _refresh_model_capabilities(self, request: ExternalGenerationRequest) -> ExternalGenerationRequest:
|
||||
if self._record_store is None:
|
||||
return request
|
||||
|
||||
try:
|
||||
record = self._record_store.get_model(request.model.key)
|
||||
except Exception:
|
||||
record = None
|
||||
|
||||
if not isinstance(record, ExternalApiModelConfig):
|
||||
return request
|
||||
|
||||
if record.key != request.model.key:
|
||||
return request
|
||||
|
||||
if record.provider_id != request.model.provider_id:
|
||||
return request
|
||||
|
||||
if record.provider_model_id != request.model.provider_model_id:
|
||||
return request
|
||||
|
||||
record = _apply_starter_overrides(record)
|
||||
|
||||
if record == request.model:
|
||||
return request
|
||||
|
||||
return ExternalGenerationRequest(
|
||||
model=record,
|
||||
mode=request.mode,
|
||||
prompt=request.prompt,
|
||||
seed=request.seed,
|
||||
num_images=request.num_images,
|
||||
width=request.width,
|
||||
height=request.height,
|
||||
image_size=request.image_size,
|
||||
init_image=request.init_image,
|
||||
mask_image=request.mask_image,
|
||||
reference_images=request.reference_images,
|
||||
metadata=request.metadata,
|
||||
provider_options=request.provider_options,
|
||||
)
|
||||
|
||||
def _bucket_request(self, request: ExternalGenerationRequest) -> ExternalGenerationRequest:
|
||||
capabilities = request.model.capabilities
|
||||
if not capabilities.allowed_aspect_ratios:
|
||||
return request
|
||||
|
||||
aspect_ratio = _format_aspect_ratio(request.width, request.height)
|
||||
size = None
|
||||
if capabilities.aspect_ratio_sizes:
|
||||
size = capabilities.aspect_ratio_sizes.get(aspect_ratio)
|
||||
|
||||
if size is not None:
|
||||
if request.width == size.width and request.height == size.height:
|
||||
return request
|
||||
return self._bucket_to_size(request, size.width, size.height, aspect_ratio)
|
||||
|
||||
if aspect_ratio in capabilities.allowed_aspect_ratios:
|
||||
return request
|
||||
|
||||
if not capabilities.aspect_ratio_sizes:
|
||||
return request
|
||||
|
||||
closest = _select_closest_ratio(
|
||||
request.width,
|
||||
request.height,
|
||||
capabilities.allowed_aspect_ratios,
|
||||
)
|
||||
if closest is None:
|
||||
return request
|
||||
|
||||
size = capabilities.aspect_ratio_sizes.get(closest)
|
||||
if size is None:
|
||||
return request
|
||||
|
||||
return self._bucket_to_size(request, size.width, size.height, closest)
|
||||
|
||||
def _bucket_to_size(
|
||||
self,
|
||||
request: ExternalGenerationRequest,
|
||||
width: int,
|
||||
height: int,
|
||||
ratio: str,
|
||||
) -> ExternalGenerationRequest:
|
||||
self._logger.info(
|
||||
"Bucketing external request provider=%s model=%s %sx%s -> %sx%s (ratio %s)",
|
||||
request.model.provider_id,
|
||||
request.model.provider_model_id,
|
||||
request.width,
|
||||
request.height,
|
||||
width,
|
||||
height,
|
||||
ratio,
|
||||
)
|
||||
|
||||
return ExternalGenerationRequest(
|
||||
model=request.model,
|
||||
mode=request.mode,
|
||||
prompt=request.prompt,
|
||||
seed=request.seed,
|
||||
num_images=request.num_images,
|
||||
width=width,
|
||||
height=height,
|
||||
image_size=request.image_size,
|
||||
init_image=_resize_image(request.init_image, width, height, "RGB"),
|
||||
mask_image=_resize_image(request.mask_image, width, height, "L"),
|
||||
reference_images=request.reference_images,
|
||||
metadata=request.metadata,
|
||||
provider_options=request.provider_options,
|
||||
)
|
||||
|
||||
|
||||
def _format_aspect_ratio(width: int, height: int) -> str:
|
||||
divisor = _gcd(width, height)
|
||||
return f"{width // divisor}:{height // divisor}"
|
||||
|
||||
|
||||
def _select_closest_ratio(width: int, height: int, ratios: list[str]) -> str | None:
|
||||
ratio = width / height
|
||||
parsed: list[tuple[str, float]] = []
|
||||
for value in ratios:
|
||||
parsed_ratio = _parse_ratio(value)
|
||||
if parsed_ratio is not None:
|
||||
parsed.append((value, parsed_ratio))
|
||||
if not parsed:
|
||||
return None
|
||||
return min(parsed, key=lambda item: abs(item[1] - ratio))[0]
|
||||
|
||||
|
||||
def _ratio_for_size(width: int, height: int, sizes: dict[str, ExternalImageSize]) -> str | None:
|
||||
for ratio, size in sizes.items():
|
||||
if size.width == width and size.height == height:
|
||||
return ratio
|
||||
return None
|
||||
|
||||
|
||||
def _parse_ratio(value: str) -> float | None:
|
||||
if ":" not in value:
|
||||
return None
|
||||
left, right = value.split(":", 1)
|
||||
try:
|
||||
numerator = float(left)
|
||||
denominator = float(right)
|
||||
except ValueError:
|
||||
return None
|
||||
if denominator == 0:
|
||||
return None
|
||||
return numerator / denominator
|
||||
|
||||
|
||||
def _gcd(a: int, b: int) -> int:
|
||||
while b:
|
||||
a, b = b, a % b
|
||||
return a
|
||||
|
||||
|
||||
def _resize_image(image: PILImageType | None, width: int, height: int, mode: str) -> PILImageType | None:
|
||||
if image is None:
|
||||
return None
|
||||
if image.width == width and image.height == height:
|
||||
return image
|
||||
return image.convert(mode).resize((width, height), Image.Resampling.LANCZOS)
|
||||
|
||||
|
||||
def _get_resize_target_for_inpaint(request: ExternalGenerationRequest) -> tuple[int, int] | None:
|
||||
if request.mode != "inpaint" or request.init_image is None:
|
||||
return None
|
||||
return request.init_image.width, request.init_image.height
|
||||
|
||||
|
||||
def _resize_result_images(result: ExternalGenerationResult, width: int, height: int) -> ExternalGenerationResult:
|
||||
resized_images = [
|
||||
ExternalGeneratedImage(
|
||||
image=generated.image
|
||||
if generated.image.width == width and generated.image.height == height
|
||||
else generated.image.resize((width, height), Image.Resampling.LANCZOS),
|
||||
seed=generated.seed,
|
||||
)
|
||||
for generated in result.images
|
||||
]
|
||||
return ExternalGenerationResult(
|
||||
images=resized_images,
|
||||
seed_used=result.seed_used,
|
||||
provider_request_id=result.provider_request_id,
|
||||
provider_metadata=result.provider_metadata,
|
||||
content_filters=result.content_filters,
|
||||
)
|
||||
|
||||
|
||||
def _apply_starter_overrides(model: ExternalApiModelConfig) -> ExternalApiModelConfig:
|
||||
source = model.source or f"external://{model.provider_id}/{model.provider_model_id}"
|
||||
starter_match = next((starter for starter in STARTER_MODELS if starter.source == source), None)
|
||||
if starter_match is None:
|
||||
return model
|
||||
updates: dict[str, object] = {}
|
||||
if starter_match.capabilities is not None:
|
||||
updates["capabilities"] = starter_match.capabilities
|
||||
if starter_match.default_settings is not None:
|
||||
updates["default_settings"] = starter_match.default_settings
|
||||
if not updates:
|
||||
return model
|
||||
return model.model_copy(update=updates)
|
||||
@@ -0,0 +1,19 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
import io
|
||||
|
||||
from PIL import Image
|
||||
from PIL.Image import Image as PILImageType
|
||||
|
||||
|
||||
def encode_image_base64(image: PILImageType, format: str = "PNG") -> str:
|
||||
buffer = io.BytesIO()
|
||||
image.save(buffer, format=format)
|
||||
return base64.b64encode(buffer.getvalue()).decode("ascii")
|
||||
|
||||
|
||||
def decode_image_base64(encoded: str) -> PILImageType:
|
||||
data = base64.b64decode(encoded)
|
||||
image = Image.open(io.BytesIO(data))
|
||||
return image.convert("RGB")
|
||||
@@ -0,0 +1,6 @@
|
||||
from invokeai.app.services.external_generation.providers.alibabacloud import AlibabaCloudProvider
|
||||
from invokeai.app.services.external_generation.providers.gemini import GeminiProvider
|
||||
from invokeai.app.services.external_generation.providers.openai import OpenAIProvider
|
||||
from invokeai.app.services.external_generation.providers.seedream import SeedreamProvider
|
||||
|
||||
__all__ = ["AlibabaCloudProvider", "GeminiProvider", "OpenAIProvider", "SeedreamProvider"]
|
||||
@@ -0,0 +1,410 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import io
|
||||
import time
|
||||
|
||||
import requests
|
||||
from PIL import Image
|
||||
from PIL.Image import Image as PILImageType
|
||||
|
||||
from invokeai.app.services.external_generation.errors import ExternalProviderRequestError
|
||||
from invokeai.app.services.external_generation.external_generation_base import ExternalProvider
|
||||
from invokeai.app.services.external_generation.external_generation_common import (
|
||||
ExternalGeneratedImage,
|
||||
ExternalGenerationRequest,
|
||||
ExternalGenerationResult,
|
||||
)
|
||||
from invokeai.app.services.external_generation.image_utils import decode_image_base64, encode_image_base64
|
||||
|
||||
# Models that support the synchronous multimodal-generation endpoint with messages format
|
||||
_SYNC_MODELS = {
|
||||
"qwen-image-2.0-pro",
|
||||
"qwen-image-2.0",
|
||||
"qwen-image-max",
|
||||
"wan2.6-t2i",
|
||||
"qwen-image-edit-max",
|
||||
}
|
||||
|
||||
# Models that use the async image-generation endpoint with flat prompt format.
|
||||
# Currently no shipped starter model uses this path, but it is retained because
|
||||
# users may install custom external models via `external://alibabacloud/<model_id>`.
|
||||
_ASYNC_MODELS: set[str] = set()
|
||||
|
||||
_TASK_POLL_INTERVAL = 5 # seconds
|
||||
_TASK_POLL_TIMEOUT = 300 # seconds
|
||||
_DOWNLOAD_TIMEOUT = 60 # seconds
|
||||
_DOWNLOAD_MAX_BYTES = 32 * 1024 * 1024 # 32 MiB safety cap on image downloads
|
||||
_RETRY_STATUS_CODES = {429, 500, 502, 503, 504}
|
||||
_MAX_RETRIES = 2 # total attempts = 1 + _MAX_RETRIES
|
||||
_RETRY_BACKOFF_BASE = 2.0 # seconds
|
||||
|
||||
|
||||
class AlibabaCloudProvider(ExternalProvider):
|
||||
provider_id = "alibabacloud"
|
||||
|
||||
def is_configured(self) -> bool:
|
||||
return bool(self._app_config.external_alibabacloud_api_key)
|
||||
|
||||
def generate(self, request: ExternalGenerationRequest) -> ExternalGenerationResult:
|
||||
api_key = self._app_config.external_alibabacloud_api_key
|
||||
if not api_key:
|
||||
raise ExternalProviderRequestError("Alibaba Cloud DashScope API key is not configured")
|
||||
|
||||
base_url = (self._app_config.external_alibabacloud_base_url or "https://dashscope-intl.aliyuncs.com").rstrip(
|
||||
"/"
|
||||
)
|
||||
model_id = request.model.provider_model_id
|
||||
headers = {
|
||||
"Content-Type": "application/json",
|
||||
"Authorization": f"Bearer {api_key}",
|
||||
}
|
||||
size = f"{request.width}*{request.height}"
|
||||
|
||||
if model_id in _SYNC_MODELS:
|
||||
return self._generate_sync(request, base_url, headers, model_id, size)
|
||||
if model_id in _ASYNC_MODELS:
|
||||
return self._generate_async(request, base_url, headers, model_id, size)
|
||||
raise ExternalProviderRequestError(
|
||||
f"Unknown DashScope model_id '{model_id}'. Add it to _SYNC_MODELS or _ASYNC_MODELS in alibabacloud.py."
|
||||
)
|
||||
|
||||
def _generate_sync(
|
||||
self,
|
||||
request: ExternalGenerationRequest,
|
||||
base_url: str,
|
||||
headers: dict[str, str],
|
||||
model_id: str,
|
||||
size: str,
|
||||
) -> ExternalGenerationResult:
|
||||
"""Use the synchronous multimodal-generation endpoint (messages format)."""
|
||||
endpoint = f"{base_url}/api/v1/services/aigc/multimodal-generation/generation"
|
||||
|
||||
content: list[dict[str, str]] = []
|
||||
|
||||
# Reference images: DashScope multimodal accepts up to 3 input images for the
|
||||
# qwen-image-edit family; we let the API surface its own limit if exceeded.
|
||||
for ref in request.reference_images:
|
||||
content.append({"image": f"data:image/png;base64,{encode_image_base64(ref.image)}"})
|
||||
|
||||
content.append({"text": request.prompt})
|
||||
|
||||
parameters: dict[str, object] = {
|
||||
"size": size,
|
||||
"n": request.num_images,
|
||||
"prompt_extend": False,
|
||||
"watermark": False,
|
||||
}
|
||||
if request.seed is not None:
|
||||
parameters["seed"] = request.seed
|
||||
|
||||
payload: dict[str, object] = {
|
||||
"model": model_id,
|
||||
"input": {
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": content,
|
||||
}
|
||||
]
|
||||
},
|
||||
"parameters": parameters,
|
||||
}
|
||||
|
||||
response = self._post_with_retry(endpoint, headers=headers, json=payload, timeout=120, label="DashScope sync")
|
||||
if not response.ok:
|
||||
raise ExternalProviderRequestError(
|
||||
f"DashScope request failed with status {response.status_code} for model '{model_id}': {response.text}"
|
||||
)
|
||||
|
||||
data = response.json()
|
||||
request_id = data.get("request_id")
|
||||
return self._parse_sync_response(data, request, request_id)
|
||||
|
||||
def _generate_async(
|
||||
self,
|
||||
request: ExternalGenerationRequest,
|
||||
base_url: str,
|
||||
headers: dict[str, str],
|
||||
model_id: str,
|
||||
size: str,
|
||||
) -> ExternalGenerationResult:
|
||||
"""Use the async image-generation endpoint (flat prompt format) with task polling."""
|
||||
endpoint = f"{base_url}/api/v1/services/aigc/image-generation/generation"
|
||||
async_headers = {**headers, "X-DashScope-Async": "enable"}
|
||||
|
||||
parameters: dict[str, object] = {
|
||||
"size": size,
|
||||
"n": request.num_images,
|
||||
"prompt_extend": False,
|
||||
"watermark": False,
|
||||
}
|
||||
if request.seed is not None:
|
||||
parameters["seed"] = request.seed
|
||||
|
||||
input_data: dict[str, object] = {"prompt": request.prompt}
|
||||
|
||||
payload: dict[str, object] = {
|
||||
"model": model_id,
|
||||
"input": input_data,
|
||||
"parameters": parameters,
|
||||
}
|
||||
|
||||
response = self._post_with_retry(
|
||||
endpoint, headers=async_headers, json=payload, timeout=60, label="DashScope async submit"
|
||||
)
|
||||
if not response.ok:
|
||||
raise ExternalProviderRequestError(
|
||||
f"DashScope async request failed with status {response.status_code} for model '{model_id}': {response.text}"
|
||||
)
|
||||
|
||||
data = response.json()
|
||||
request_id = data.get("request_id")
|
||||
output = data.get("output", {})
|
||||
task_id = output.get("task_id")
|
||||
|
||||
if not task_id:
|
||||
raise ExternalProviderRequestError(f"DashScope async response missing task_id: {data}")
|
||||
|
||||
return self._poll_task(base_url, headers, task_id, request, request_id)
|
||||
|
||||
def _poll_task(
|
||||
self,
|
||||
base_url: str,
|
||||
headers: dict[str, str],
|
||||
task_id: str,
|
||||
request: ExternalGenerationRequest,
|
||||
request_id: str | None,
|
||||
) -> ExternalGenerationResult:
|
||||
"""Poll an async task until completion."""
|
||||
task_url = f"{base_url}/api/v1/tasks/{task_id}"
|
||||
start_time = time.monotonic()
|
||||
poll_headers = {"Authorization": headers["Authorization"]}
|
||||
first_poll = True
|
||||
|
||||
while True:
|
||||
elapsed = time.monotonic() - start_time
|
||||
if elapsed > _TASK_POLL_TIMEOUT:
|
||||
raise ExternalProviderRequestError(f"DashScope task {task_id} timed out after {_TASK_POLL_TIMEOUT}s")
|
||||
|
||||
response = self._get_with_retry(task_url, headers=poll_headers, timeout=30, label="DashScope task poll")
|
||||
if not response.ok:
|
||||
raise ExternalProviderRequestError(
|
||||
f"DashScope task poll failed with status {response.status_code}: {response.text}"
|
||||
)
|
||||
|
||||
data = response.json()
|
||||
output = data.get("output", {})
|
||||
status = output.get("task_status")
|
||||
|
||||
if first_poll:
|
||||
self._logger.info("DashScope task %s submitted (status=%s)", task_id, status)
|
||||
first_poll = False
|
||||
|
||||
if status == "SUCCEEDED":
|
||||
return self._parse_async_response(output, request, request_id)
|
||||
if status in ("FAILED", "UNKNOWN"):
|
||||
message = output.get("message", "Unknown error")
|
||||
raise ExternalProviderRequestError(f"DashScope task {task_id} failed: {message}")
|
||||
|
||||
self._logger.debug("DashScope task %s status: %s (%.0fs elapsed)", task_id, status, elapsed)
|
||||
time.sleep(_TASK_POLL_INTERVAL)
|
||||
|
||||
def _parse_sync_response(
|
||||
self,
|
||||
data: dict[str, object],
|
||||
request: ExternalGenerationRequest,
|
||||
request_id: str | None,
|
||||
) -> ExternalGenerationResult:
|
||||
"""Parse the synchronous multimodal-generation response."""
|
||||
output = data.get("output")
|
||||
if not isinstance(output, dict):
|
||||
raise ExternalProviderRequestError(f"DashScope response missing output: {data}")
|
||||
|
||||
choices = output.get("choices")
|
||||
if not isinstance(choices, list):
|
||||
raise ExternalProviderRequestError(f"DashScope response missing choices: {data}")
|
||||
|
||||
images: list[ExternalGeneratedImage] = []
|
||||
for choice in choices:
|
||||
if not isinstance(choice, dict):
|
||||
continue
|
||||
message = choice.get("message")
|
||||
if not isinstance(message, dict):
|
||||
continue
|
||||
content = message.get("content")
|
||||
if not isinstance(content, list):
|
||||
continue
|
||||
for part in content:
|
||||
if not isinstance(part, dict):
|
||||
continue
|
||||
image_url = part.get("image")
|
||||
if isinstance(image_url, str) and image_url:
|
||||
pil_image = self._download_image(image_url)
|
||||
images.append(ExternalGeneratedImage(image=pil_image, seed=request.seed))
|
||||
|
||||
if not images:
|
||||
raise ExternalProviderRequestError(f"DashScope response contained no images: {data}")
|
||||
|
||||
return ExternalGenerationResult(
|
||||
images=images,
|
||||
seed_used=request.seed,
|
||||
provider_request_id=request_id,
|
||||
provider_metadata={"model": request.model.provider_model_id},
|
||||
)
|
||||
|
||||
def _parse_async_response(
|
||||
self,
|
||||
output: dict[str, object],
|
||||
request: ExternalGenerationRequest,
|
||||
request_id: str | None,
|
||||
) -> ExternalGenerationResult:
|
||||
"""Parse the async task completion response."""
|
||||
results = output.get("results")
|
||||
if not isinstance(results, list):
|
||||
raise ExternalProviderRequestError(f"DashScope async response missing results: {output}")
|
||||
|
||||
images: list[ExternalGeneratedImage] = []
|
||||
for result in results:
|
||||
if not isinstance(result, dict):
|
||||
continue
|
||||
url = result.get("url")
|
||||
if isinstance(url, str) and url:
|
||||
pil_image = self._download_image(url)
|
||||
images.append(ExternalGeneratedImage(image=pil_image, seed=request.seed))
|
||||
continue
|
||||
b64_image = result.get("b64_image")
|
||||
if isinstance(b64_image, str) and b64_image:
|
||||
pil_image = decode_image_base64(b64_image)
|
||||
images.append(ExternalGeneratedImage(image=pil_image, seed=request.seed))
|
||||
|
||||
if not images:
|
||||
raise ExternalProviderRequestError(f"DashScope async response contained no images: {output}")
|
||||
|
||||
return ExternalGenerationResult(
|
||||
images=images,
|
||||
seed_used=request.seed,
|
||||
provider_request_id=request_id,
|
||||
provider_metadata={"model": request.model.provider_model_id},
|
||||
)
|
||||
|
||||
def _download_image(self, url: str) -> PILImageType:
|
||||
"""Download an image from a URL and return it as a PIL Image, with a size cap."""
|
||||
try:
|
||||
response = requests.get(url, timeout=_DOWNLOAD_TIMEOUT, stream=True)
|
||||
except requests.RequestException as exc:
|
||||
raise ExternalProviderRequestError(f"Failed to download image from DashScope: {exc}") from exc
|
||||
|
||||
with response:
|
||||
if not response.ok:
|
||||
raise ExternalProviderRequestError(
|
||||
f"Failed to download image from DashScope (status {response.status_code})"
|
||||
)
|
||||
|
||||
content_length = response.headers.get("Content-Length")
|
||||
if content_length is not None:
|
||||
try:
|
||||
if int(content_length) > _DOWNLOAD_MAX_BYTES:
|
||||
raise ExternalProviderRequestError(
|
||||
f"DashScope image exceeds {_DOWNLOAD_MAX_BYTES} byte cap (Content-Length={content_length})"
|
||||
)
|
||||
except ValueError:
|
||||
pass
|
||||
|
||||
buffer = bytearray()
|
||||
for chunk in response.iter_content(chunk_size=64 * 1024):
|
||||
if not chunk:
|
||||
continue
|
||||
buffer.extend(chunk)
|
||||
if len(buffer) > _DOWNLOAD_MAX_BYTES:
|
||||
raise ExternalProviderRequestError(f"DashScope image exceeds {_DOWNLOAD_MAX_BYTES} byte cap")
|
||||
|
||||
return Image.open(io.BytesIO(bytes(buffer))).convert("RGB")
|
||||
|
||||
def _post_with_retry(
|
||||
self,
|
||||
url: str,
|
||||
*,
|
||||
headers: dict[str, str],
|
||||
json: dict,
|
||||
timeout: int,
|
||||
label: str,
|
||||
) -> requests.Response:
|
||||
return self._request_with_retry("POST", url, headers=headers, json=json, timeout=timeout, label=label)
|
||||
|
||||
def _get_with_retry(
|
||||
self,
|
||||
url: str,
|
||||
*,
|
||||
headers: dict[str, str],
|
||||
timeout: int,
|
||||
label: str,
|
||||
) -> requests.Response:
|
||||
return self._request_with_retry("GET", url, headers=headers, timeout=timeout, label=label)
|
||||
|
||||
def _request_with_retry(
|
||||
self,
|
||||
method: str,
|
||||
url: str,
|
||||
*,
|
||||
headers: dict[str, str],
|
||||
timeout: int,
|
||||
label: str,
|
||||
json: dict | None = None,
|
||||
) -> requests.Response:
|
||||
"""Issue a request with limited retries on transient failures (429/5xx, network errors).
|
||||
|
||||
Honors `Retry-After` for 429 responses when present. Non-retryable errors
|
||||
(4xx other than 429, parse failures) are returned to the caller, which is
|
||||
responsible for raising a meaningful ExternalProviderRequestError.
|
||||
"""
|
||||
last_exc: Exception | None = None
|
||||
for attempt in range(_MAX_RETRIES + 1):
|
||||
try:
|
||||
if method == "POST":
|
||||
response = requests.post(url, headers=headers, json=json, timeout=timeout)
|
||||
else:
|
||||
response = requests.get(url, headers=headers, timeout=timeout)
|
||||
except requests.RequestException as exc:
|
||||
last_exc = exc
|
||||
if attempt >= _MAX_RETRIES:
|
||||
raise ExternalProviderRequestError(f"{label} network error: {exc}") from exc
|
||||
delay = _RETRY_BACKOFF_BASE * (2**attempt)
|
||||
self._logger.warning(
|
||||
"%s network error on attempt %d/%d: %s — retrying in %.1fs",
|
||||
label,
|
||||
attempt + 1,
|
||||
_MAX_RETRIES + 1,
|
||||
exc,
|
||||
delay,
|
||||
)
|
||||
time.sleep(delay)
|
||||
continue
|
||||
|
||||
if response.status_code in _RETRY_STATUS_CODES and attempt < _MAX_RETRIES:
|
||||
delay = self._retry_delay(response, attempt)
|
||||
self._logger.warning(
|
||||
"%s got status %d on attempt %d/%d — retrying in %.1fs",
|
||||
label,
|
||||
response.status_code,
|
||||
attempt + 1,
|
||||
_MAX_RETRIES + 1,
|
||||
delay,
|
||||
)
|
||||
time.sleep(delay)
|
||||
continue
|
||||
|
||||
return response
|
||||
|
||||
# Unreachable: the loop either returns a response or raises.
|
||||
assert last_exc is not None
|
||||
raise ExternalProviderRequestError(f"{label} failed after retries: {last_exc}") from last_exc
|
||||
|
||||
@staticmethod
|
||||
def _retry_delay(response: requests.Response, attempt: int) -> float:
|
||||
retry_after = response.headers.get("Retry-After")
|
||||
if retry_after:
|
||||
try:
|
||||
return max(0.0, float(retry_after))
|
||||
except ValueError:
|
||||
pass
|
||||
return _RETRY_BACKOFF_BASE * (2**attempt)
|
||||
@@ -0,0 +1,248 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import requests
|
||||
|
||||
from invokeai.app.services.external_generation.errors import (
|
||||
ExternalProviderRateLimitError,
|
||||
ExternalProviderRequestError,
|
||||
)
|
||||
from invokeai.app.services.external_generation.external_generation_base import ExternalProvider
|
||||
from invokeai.app.services.external_generation.external_generation_common import (
|
||||
ExternalGeneratedImage,
|
||||
ExternalGenerationRequest,
|
||||
ExternalGenerationResult,
|
||||
)
|
||||
from invokeai.app.services.external_generation.image_utils import decode_image_base64, encode_image_base64
|
||||
|
||||
|
||||
class GeminiProvider(ExternalProvider):
|
||||
provider_id = "gemini"
|
||||
_SYSTEM_INSTRUCTION = (
|
||||
"You are an image generation model. Always respond with an image based on the user's prompt. "
|
||||
"Do not return text-only responses. If the user input is not an edit instruction, "
|
||||
"interpret it as a request to create a new image."
|
||||
)
|
||||
|
||||
def is_configured(self) -> bool:
|
||||
return bool(self._app_config.external_gemini_api_key)
|
||||
|
||||
def generate(self, request: ExternalGenerationRequest) -> ExternalGenerationResult:
|
||||
api_key = self._app_config.external_gemini_api_key
|
||||
if not api_key:
|
||||
raise ExternalProviderRequestError("Gemini API key is not configured")
|
||||
|
||||
base_url = (self._app_config.external_gemini_base_url or "https://generativelanguage.googleapis.com").rstrip(
|
||||
"/"
|
||||
)
|
||||
if not base_url.endswith("/v1") and not base_url.endswith("/v1beta"):
|
||||
base_url = f"{base_url}/v1beta"
|
||||
model_id = request.model.provider_model_id.removeprefix("models/")
|
||||
endpoint = f"{base_url}/models/{model_id}:generateContent"
|
||||
|
||||
request_parts: list[dict[str, object]] = []
|
||||
|
||||
if request.init_image is not None:
|
||||
request_parts.append(
|
||||
{
|
||||
"inlineData": {
|
||||
"mimeType": "image/png",
|
||||
"data": encode_image_base64(request.init_image),
|
||||
}
|
||||
}
|
||||
)
|
||||
|
||||
request_parts.append({"text": request.prompt})
|
||||
|
||||
for reference in request.reference_images:
|
||||
request_parts.append(
|
||||
{
|
||||
"inlineData": {
|
||||
"mimeType": "image/png",
|
||||
"data": encode_image_base64(reference.image),
|
||||
}
|
||||
}
|
||||
)
|
||||
|
||||
opts = request.provider_options or {}
|
||||
|
||||
generation_config: dict[str, object] = {
|
||||
"candidateCount": request.num_images,
|
||||
"responseModalities": ["IMAGE"],
|
||||
}
|
||||
if "temperature" in opts:
|
||||
generation_config["temperature"] = opts["temperature"]
|
||||
aspect_ratio = _select_aspect_ratio(
|
||||
request.width,
|
||||
request.height,
|
||||
request.model.capabilities.allowed_aspect_ratios,
|
||||
)
|
||||
uses_image_config = request.model.capabilities.resolution_presets is not None
|
||||
if uses_image_config:
|
||||
image_config: dict[str, str] = {}
|
||||
if aspect_ratio is not None:
|
||||
image_config["aspectRatio"] = aspect_ratio
|
||||
if request.image_size is not None:
|
||||
image_config["imageSize"] = request.image_size
|
||||
if image_config:
|
||||
generation_config["imageConfig"] = image_config
|
||||
system_instruction = self._SYSTEM_INSTRUCTION
|
||||
if request.init_image is not None:
|
||||
system_instruction = (
|
||||
f"{system_instruction} An input image is provided. "
|
||||
"Treat the prompt as an edit instruction and modify the image accordingly. "
|
||||
"Do not return the original image unchanged."
|
||||
)
|
||||
if not uses_image_config and aspect_ratio is not None:
|
||||
system_instruction = f"{system_instruction} Use an aspect ratio of {aspect_ratio}."
|
||||
|
||||
payload: dict[str, object] = {
|
||||
"systemInstruction": {"parts": [{"text": system_instruction}]},
|
||||
"contents": [{"role": "user", "parts": request_parts}],
|
||||
"generationConfig": generation_config,
|
||||
}
|
||||
if "thinking_level" in opts:
|
||||
payload["thinkingConfig"] = {"thinkingLevel": opts["thinking_level"].upper()}
|
||||
|
||||
response = requests.post(
|
||||
endpoint,
|
||||
params={"key": api_key},
|
||||
json=payload,
|
||||
timeout=120,
|
||||
)
|
||||
|
||||
if not response.ok:
|
||||
if response.status_code == 429:
|
||||
retry_after = _parse_retry_after(response.headers.get("retry-after"))
|
||||
raise ExternalProviderRateLimitError(
|
||||
f"Gemini rate limit exceeded. {f'Retry after {retry_after:.0f}s.' if retry_after else 'Please try again later.'}",
|
||||
retry_after=retry_after,
|
||||
)
|
||||
raise ExternalProviderRequestError(
|
||||
f"Gemini request failed with status {response.status_code} for model '{model_id}': {response.text}"
|
||||
)
|
||||
|
||||
data = response.json()
|
||||
if not isinstance(data, dict):
|
||||
raise ExternalProviderRequestError("Gemini response payload was not a JSON object")
|
||||
images: list[ExternalGeneratedImage] = []
|
||||
text_parts: list[str] = []
|
||||
finish_messages: list[str] = []
|
||||
candidates = data.get("candidates")
|
||||
if not isinstance(candidates, list):
|
||||
raise ExternalProviderRequestError("Gemini response payload missing candidates")
|
||||
for candidate in candidates:
|
||||
if not isinstance(candidate, dict):
|
||||
continue
|
||||
finish_message = candidate.get("finishMessage")
|
||||
finish_reason = candidate.get("finishReason")
|
||||
if isinstance(finish_message, str):
|
||||
finish_messages.append(finish_message)
|
||||
elif isinstance(finish_reason, str):
|
||||
finish_messages.append(f"Finish reason: {finish_reason}")
|
||||
for part in _iter_response_parts(candidate):
|
||||
inline_data = part.get("inline_data") or part.get("inlineData")
|
||||
if isinstance(inline_data, dict):
|
||||
encoded = inline_data.get("data")
|
||||
if encoded:
|
||||
image = decode_image_base64(encoded)
|
||||
images.append(ExternalGeneratedImage(image=image, seed=request.seed))
|
||||
continue
|
||||
file_data = part.get("fileData") or part.get("file_data")
|
||||
if isinstance(file_data, dict):
|
||||
file_uri = file_data.get("fileUri") or file_data.get("file_uri")
|
||||
if isinstance(file_uri, str) and file_uri:
|
||||
raise ExternalProviderRequestError(
|
||||
f"Gemini returned fileUri instead of inline image data: {file_uri}"
|
||||
)
|
||||
text = part.get("text")
|
||||
if isinstance(text, str):
|
||||
text_parts.append(text)
|
||||
|
||||
if not images:
|
||||
self._logger.error("Gemini response contained no images: %s", data)
|
||||
detail = ""
|
||||
if finish_messages:
|
||||
combined = " ".join(message.strip() for message in finish_messages if message.strip())
|
||||
if combined:
|
||||
detail = f" Response status: {combined[:500]}"
|
||||
elif text_parts:
|
||||
combined = " ".join(text_parts).strip()
|
||||
if combined:
|
||||
detail = f" Response text: {combined[:500]}"
|
||||
raise ExternalProviderRequestError(f"Gemini response contained no images.{detail}")
|
||||
|
||||
return ExternalGenerationResult(
|
||||
images=images,
|
||||
seed_used=request.seed,
|
||||
provider_metadata={"model": request.model.provider_model_id},
|
||||
)
|
||||
|
||||
|
||||
def _iter_response_parts(candidate: dict[str, object]) -> list[dict[str, object]]:
|
||||
content = candidate.get("content")
|
||||
if isinstance(content, dict):
|
||||
content_parts = content.get("parts")
|
||||
if isinstance(content_parts, list):
|
||||
return [part for part in content_parts if isinstance(part, dict)]
|
||||
contents = candidate.get("contents")
|
||||
if isinstance(contents, list):
|
||||
parts: list[dict[str, object]] = []
|
||||
for item in contents:
|
||||
if not isinstance(item, dict):
|
||||
continue
|
||||
item_parts = item.get("parts")
|
||||
if isinstance(item_parts, list):
|
||||
parts.extend([part for part in item_parts if isinstance(part, dict)])
|
||||
if parts:
|
||||
return parts
|
||||
return []
|
||||
|
||||
|
||||
def _select_aspect_ratio(width: int, height: int, allowed: list[str] | None) -> str | None:
|
||||
if width <= 0 or height <= 0:
|
||||
return None
|
||||
ratio = width / height
|
||||
default_ratio = _format_aspect_ratio(width, height)
|
||||
if not allowed:
|
||||
return default_ratio
|
||||
parsed = [(value, _parse_ratio(value)) for value in allowed]
|
||||
filtered = [(value, parsed_ratio) for value, parsed_ratio in parsed if parsed_ratio is not None]
|
||||
if not filtered:
|
||||
return default_ratio
|
||||
return min(filtered, key=lambda item: abs(item[1] - ratio))[0]
|
||||
|
||||
|
||||
def _format_aspect_ratio(width: int, height: int) -> str | None:
|
||||
if width <= 0 or height <= 0:
|
||||
return None
|
||||
divisor = _gcd(width, height)
|
||||
return f"{width // divisor}:{height // divisor}"
|
||||
|
||||
|
||||
def _parse_ratio(value: str) -> float | None:
|
||||
if ":" not in value:
|
||||
return None
|
||||
left, right = value.split(":", 1)
|
||||
try:
|
||||
numerator = float(left)
|
||||
denominator = float(right)
|
||||
except ValueError:
|
||||
return None
|
||||
if denominator == 0:
|
||||
return None
|
||||
return numerator / denominator
|
||||
|
||||
|
||||
def _parse_retry_after(value: str | None) -> float | None:
|
||||
if not value:
|
||||
return None
|
||||
try:
|
||||
return float(value)
|
||||
except ValueError:
|
||||
return None
|
||||
|
||||
|
||||
def _gcd(a: int, b: int) -> int:
|
||||
while b:
|
||||
a, b = b, a % b
|
||||
return a
|
||||
@@ -0,0 +1,162 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import io
|
||||
|
||||
import requests
|
||||
from PIL.Image import Image as PILImageType
|
||||
|
||||
from invokeai.app.services.external_generation.errors import (
|
||||
ExternalProviderRateLimitError,
|
||||
ExternalProviderRequestError,
|
||||
)
|
||||
from invokeai.app.services.external_generation.external_generation_base import ExternalProvider
|
||||
from invokeai.app.services.external_generation.external_generation_common import (
|
||||
ExternalGeneratedImage,
|
||||
ExternalGenerationRequest,
|
||||
ExternalGenerationResult,
|
||||
)
|
||||
from invokeai.app.services.external_generation.image_utils import decode_image_base64
|
||||
|
||||
|
||||
class OpenAIProvider(ExternalProvider):
|
||||
provider_id = "openai"
|
||||
|
||||
_GPT_IMAGE_MODELS = {"gpt-image-1", "gpt-image-1.5", "gpt-image-1-mini", "gpt-image-2"}
|
||||
_DEFAULT_TIMEOUT = 120
|
||||
_MODEL_TIMEOUTS: dict[str, int] = {"gpt-image-2": 300}
|
||||
|
||||
def is_configured(self) -> bool:
|
||||
return bool(self._app_config.external_openai_api_key)
|
||||
|
||||
def generate(self, request: ExternalGenerationRequest) -> ExternalGenerationResult:
|
||||
api_key = self._app_config.external_openai_api_key
|
||||
if not api_key:
|
||||
raise ExternalProviderRequestError("OpenAI API key is not configured")
|
||||
|
||||
model_id = request.model.provider_model_id
|
||||
is_gpt_image = model_id in self._GPT_IMAGE_MODELS
|
||||
timeout = self._MODEL_TIMEOUTS.get(model_id, self._DEFAULT_TIMEOUT)
|
||||
size = f"{request.width}x{request.height}"
|
||||
base_url = (self._app_config.external_openai_base_url or "https://api.openai.com").rstrip("/")
|
||||
headers = {"Authorization": f"Bearer {api_key}"}
|
||||
|
||||
use_edits_endpoint = request.mode != "txt2img" or bool(request.reference_images)
|
||||
|
||||
opts = request.provider_options or {}
|
||||
|
||||
if not use_edits_endpoint:
|
||||
payload: dict[str, object] = {
|
||||
"model": model_id,
|
||||
"prompt": request.prompt,
|
||||
"n": request.num_images,
|
||||
"size": size,
|
||||
}
|
||||
# GPT Image models use output_format; DALL-E uses response_format
|
||||
if is_gpt_image:
|
||||
payload["output_format"] = "png"
|
||||
else:
|
||||
payload["response_format"] = "b64_json"
|
||||
if is_gpt_image:
|
||||
if opts.get("quality") and opts["quality"] != "auto":
|
||||
payload["quality"] = opts["quality"]
|
||||
if opts.get("background") and opts["background"] != "auto":
|
||||
payload["background"] = opts["background"]
|
||||
response = requests.post(
|
||||
f"{base_url}/v1/images/generations",
|
||||
headers=headers,
|
||||
json=payload,
|
||||
timeout=timeout,
|
||||
)
|
||||
else:
|
||||
images: list[PILImageType] = []
|
||||
if request.init_image is not None:
|
||||
images.append(request.init_image)
|
||||
images.extend(reference.image for reference in request.reference_images)
|
||||
if not images:
|
||||
raise ExternalProviderRequestError(
|
||||
"OpenAI image edits require at least one image (init image or reference image)"
|
||||
)
|
||||
|
||||
files: list[tuple[str, tuple[str, io.BytesIO, str]]] = []
|
||||
image_field_name = "image" if len(images) == 1 else "image[]"
|
||||
for index, image in enumerate(images):
|
||||
image_buffer = io.BytesIO()
|
||||
image.save(image_buffer, format="PNG")
|
||||
image_buffer.seek(0)
|
||||
files.append((image_field_name, (f"image_{index}.png", image_buffer, "image/png")))
|
||||
|
||||
if request.mask_image is not None:
|
||||
mask_buffer = io.BytesIO()
|
||||
request.mask_image.save(mask_buffer, format="PNG")
|
||||
mask_buffer.seek(0)
|
||||
files.append(("mask", ("mask.png", mask_buffer, "image/png")))
|
||||
|
||||
data: dict[str, object] = {
|
||||
"model": model_id,
|
||||
"prompt": request.prompt,
|
||||
"n": request.num_images,
|
||||
"size": size,
|
||||
}
|
||||
if is_gpt_image:
|
||||
data["output_format"] = "png"
|
||||
else:
|
||||
data["response_format"] = "b64_json"
|
||||
if is_gpt_image:
|
||||
if opts.get("quality") and opts["quality"] != "auto":
|
||||
data["quality"] = opts["quality"]
|
||||
if opts.get("background") and opts["background"] != "auto":
|
||||
data["background"] = opts["background"]
|
||||
if opts.get("input_fidelity"):
|
||||
data["input_fidelity"] = opts["input_fidelity"]
|
||||
response = requests.post(
|
||||
f"{base_url}/v1/images/edits",
|
||||
headers=headers,
|
||||
data=data,
|
||||
files=files,
|
||||
timeout=timeout,
|
||||
)
|
||||
|
||||
if not response.ok:
|
||||
if response.status_code == 429:
|
||||
retry_after = _parse_retry_after(response.headers.get("retry-after"))
|
||||
raise ExternalProviderRateLimitError(
|
||||
f"OpenAI rate limit exceeded. {f'Retry after {retry_after:.0f}s.' if retry_after else 'Please try again later.'}",
|
||||
retry_after=retry_after,
|
||||
)
|
||||
raise ExternalProviderRequestError(
|
||||
f"OpenAI request failed with status {response.status_code}: {response.text}"
|
||||
)
|
||||
|
||||
response_payload = response.json()
|
||||
if not isinstance(response_payload, dict):
|
||||
raise ExternalProviderRequestError("OpenAI response payload was not a JSON object")
|
||||
images: list[ExternalGeneratedImage] = []
|
||||
data_items = response_payload.get("data")
|
||||
if not isinstance(data_items, list):
|
||||
raise ExternalProviderRequestError("OpenAI response payload missing image data")
|
||||
for item in data_items:
|
||||
if not isinstance(item, dict):
|
||||
continue
|
||||
encoded = item.get("b64_json")
|
||||
if not encoded:
|
||||
continue
|
||||
images.append(ExternalGeneratedImage(image=decode_image_base64(encoded), seed=request.seed))
|
||||
|
||||
if not images:
|
||||
raise ExternalProviderRequestError("OpenAI response contained no images")
|
||||
|
||||
return ExternalGenerationResult(
|
||||
images=images,
|
||||
seed_used=request.seed,
|
||||
provider_request_id=response.headers.get("x-request-id"),
|
||||
provider_metadata={"model": model_id},
|
||||
)
|
||||
|
||||
|
||||
def _parse_retry_after(value: str | None) -> float | None:
|
||||
if not value:
|
||||
return None
|
||||
try:
|
||||
return float(value)
|
||||
except ValueError:
|
||||
return None
|
||||
@@ -0,0 +1,171 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import requests
|
||||
|
||||
from invokeai.app.services.external_generation.errors import (
|
||||
ExternalProviderCapabilityError,
|
||||
ExternalProviderRateLimitError,
|
||||
ExternalProviderRequestError,
|
||||
)
|
||||
from invokeai.app.services.external_generation.external_generation_base import ExternalProvider
|
||||
from invokeai.app.services.external_generation.external_generation_common import (
|
||||
ExternalGeneratedImage,
|
||||
ExternalGenerationRequest,
|
||||
ExternalGenerationResult,
|
||||
)
|
||||
from invokeai.app.services.external_generation.image_utils import decode_image_base64, encode_image_base64
|
||||
|
||||
_SEEDREAM_BATCH_PREFIXES = (
|
||||
"seedream-5",
|
||||
"seedream-4.5",
|
||||
"seedream-4.0",
|
||||
"seedream-4-5",
|
||||
"seedream-4-0",
|
||||
"seedream-5-0",
|
||||
)
|
||||
|
||||
# Seedream batch endpoint accepts up to 15 total images counting both inputs (reference + init)
|
||||
# and outputs combined. Hitting this only after the API call wastes a request and produces a
|
||||
# confusing 400, so we enforce it locally for batch-capable models.
|
||||
_SEEDREAM_BATCH_MAX_TOTAL_IMAGES = 15
|
||||
|
||||
|
||||
class SeedreamProvider(ExternalProvider):
|
||||
provider_id = "seedream"
|
||||
|
||||
def is_configured(self) -> bool:
|
||||
return bool(self._app_config.external_seedream_api_key)
|
||||
|
||||
def generate(self, request: ExternalGenerationRequest) -> ExternalGenerationResult:
|
||||
api_key = self._app_config.external_seedream_api_key
|
||||
if not api_key:
|
||||
raise ExternalProviderRequestError("Seedream API key is not configured")
|
||||
|
||||
base_url = (self._app_config.external_seedream_base_url or "https://ark.ap-southeast.bytepluses.com").rstrip(
|
||||
"/"
|
||||
)
|
||||
endpoint = f"{base_url}/api/v3/images/generations"
|
||||
headers = {"Authorization": f"Bearer {api_key}"}
|
||||
|
||||
model_id = request.model.provider_model_id
|
||||
is_batch_model = any(model_id.startswith(prefix) for prefix in _SEEDREAM_BATCH_PREFIXES)
|
||||
|
||||
if is_batch_model:
|
||||
input_image_count = len(request.reference_images) + (1 if request.init_image is not None else 0)
|
||||
total_images = input_image_count + request.num_images
|
||||
if total_images > _SEEDREAM_BATCH_MAX_TOTAL_IMAGES:
|
||||
raise ExternalProviderCapabilityError(
|
||||
f"{request.model.name} supports at most {_SEEDREAM_BATCH_MAX_TOTAL_IMAGES} images total "
|
||||
f"(reference + init + output), got {total_images}"
|
||||
)
|
||||
|
||||
opts = request.provider_options or {}
|
||||
|
||||
payload: dict[str, object] = {
|
||||
"model": model_id,
|
||||
"prompt": request.prompt,
|
||||
"size": f"{request.width}x{request.height}",
|
||||
"response_format": "b64_json",
|
||||
"watermark": opts.get("watermark", False),
|
||||
}
|
||||
|
||||
if opts.get("optimize_prompt"):
|
||||
payload["optimize_prompt_options"] = {"optimize_prompt": True}
|
||||
|
||||
# Seed and guidance_scale are only supported on 3.0 models
|
||||
if not is_batch_model and request.seed is not None and request.seed >= 0:
|
||||
payload["seed"] = request.seed
|
||||
if not is_batch_model and opts.get("guidance_scale") is not None:
|
||||
payload["guidance_scale"] = opts["guidance_scale"]
|
||||
|
||||
# Batch generation for 4.x/5.x models
|
||||
if is_batch_model:
|
||||
if request.num_images > 1:
|
||||
payload["sequential_image_generation"] = "auto"
|
||||
payload["sequential_image_generation_options"] = {"max_images": request.num_images}
|
||||
else:
|
||||
payload["sequential_image_generation"] = "disabled"
|
||||
|
||||
# Image input: init_image for img2img, reference images for 4.x
|
||||
images_b64: list[str] = []
|
||||
if request.init_image is not None:
|
||||
images_b64.append(f"data:image/png;base64,{encode_image_base64(request.init_image)}")
|
||||
for reference in request.reference_images:
|
||||
images_b64.append(f"data:image/png;base64,{encode_image_base64(reference.image)}")
|
||||
|
||||
if images_b64:
|
||||
payload["image"] = images_b64 if len(images_b64) > 1 else images_b64[0]
|
||||
|
||||
response = requests.post(endpoint, headers=headers, json=payload, timeout=120)
|
||||
|
||||
if not response.ok:
|
||||
if response.status_code == 429:
|
||||
retry_after = _parse_retry_after(response.headers.get("retry-after"))
|
||||
raise ExternalProviderRateLimitError(
|
||||
f"Seedream rate limit exceeded. {f'Retry after {retry_after:.0f}s.' if retry_after else 'Please try again later.'}",
|
||||
retry_after=retry_after,
|
||||
)
|
||||
raise ExternalProviderRequestError(
|
||||
f"Seedream request failed with status {response.status_code}: {response.text}"
|
||||
)
|
||||
|
||||
body = response.json()
|
||||
if not isinstance(body, dict):
|
||||
raise ExternalProviderRequestError("Seedream response payload was not a JSON object")
|
||||
|
||||
generated_images: list[ExternalGeneratedImage] = []
|
||||
item_errors: list[dict[str, object]] = []
|
||||
data_items = body.get("data")
|
||||
if not isinstance(data_items, list):
|
||||
raise ExternalProviderRequestError("Seedream response payload missing image data")
|
||||
|
||||
for item in data_items:
|
||||
if not isinstance(item, dict):
|
||||
continue
|
||||
# Items may be error objects for failed images in batch — collect rather than discard
|
||||
# so partial-failure causes (e.g., content filter) are visible to the caller.
|
||||
if "error" in item:
|
||||
error_payload = item["error"]
|
||||
item_errors.append(
|
||||
error_payload if isinstance(error_payload, dict) else {"message": str(error_payload)}
|
||||
)
|
||||
continue
|
||||
encoded = item.get("b64_json")
|
||||
if not encoded:
|
||||
continue
|
||||
image = decode_image_base64(encoded)
|
||||
generated_images.append(ExternalGeneratedImage(image=image, seed=request.seed))
|
||||
|
||||
if not generated_images:
|
||||
if item_errors:
|
||||
first = item_errors[0]
|
||||
message = first.get("message") if isinstance(first, dict) else None
|
||||
raise ExternalProviderRequestError(
|
||||
f"Seedream returned no images. Provider reported: {message or item_errors}"
|
||||
)
|
||||
raise ExternalProviderRequestError("Seedream response contained no images")
|
||||
|
||||
provider_metadata: dict[str, object] = {"model": model_id}
|
||||
if item_errors:
|
||||
provider_metadata["partial_failures"] = item_errors
|
||||
self._logger.warning(
|
||||
"Seedream returned %d image(s) with %d partial failure(s): %s",
|
||||
len(generated_images),
|
||||
len(item_errors),
|
||||
item_errors,
|
||||
)
|
||||
|
||||
return ExternalGenerationResult(
|
||||
images=generated_images,
|
||||
seed_used=request.seed,
|
||||
provider_metadata=provider_metadata,
|
||||
)
|
||||
|
||||
|
||||
def _parse_retry_after(value: str | None) -> float | None:
|
||||
if not value:
|
||||
return None
|
||||
try:
|
||||
return float(value)
|
||||
except ValueError:
|
||||
return None
|
||||
@@ -0,0 +1,59 @@
|
||||
from logging import Logger
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from invokeai.app.services.model_records.model_records_base import ModelRecordChanges
|
||||
from invokeai.backend.model_manager.configs.external_api import ExternalApiModelConfig
|
||||
from invokeai.backend.model_manager.starter_models import STARTER_MODELS
|
||||
from invokeai.backend.model_manager.taxonomy import BaseModelType, ModelType
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from invokeai.app.services.model_manager.model_manager_base import ModelManagerServiceBase
|
||||
|
||||
|
||||
def sync_configured_external_starter_models(
|
||||
configured_provider_ids: set[str],
|
||||
model_manager: "ModelManagerServiceBase",
|
||||
logger: Logger,
|
||||
) -> list[str]:
|
||||
"""Queue missing external starter models for configured providers."""
|
||||
|
||||
if not configured_provider_ids:
|
||||
return []
|
||||
|
||||
installed_sources = {
|
||||
model.source
|
||||
for model in model_manager.store.search_by_attr(
|
||||
base_model=BaseModelType.External,
|
||||
model_type=ModelType.ExternalImageGenerator,
|
||||
)
|
||||
if isinstance(model, ExternalApiModelConfig) and model.source
|
||||
}
|
||||
|
||||
queued_sources: list[str] = []
|
||||
for starter_model in STARTER_MODELS:
|
||||
if not starter_model.source.startswith("external://"):
|
||||
continue
|
||||
|
||||
provider_id = starter_model.source.removeprefix("external://").split("/", 1)[0]
|
||||
if provider_id not in configured_provider_ids:
|
||||
continue
|
||||
|
||||
if starter_model.source in installed_sources:
|
||||
continue
|
||||
|
||||
model_manager.install.heuristic_import(
|
||||
starter_model.source,
|
||||
config=ModelRecordChanges(
|
||||
name=starter_model.name,
|
||||
base=starter_model.base,
|
||||
type=starter_model.type,
|
||||
description=starter_model.description,
|
||||
format=starter_model.format,
|
||||
capabilities=starter_model.capabilities,
|
||||
default_settings=starter_model.default_settings,
|
||||
),
|
||||
)
|
||||
queued_sources.append(starter_model.source)
|
||||
logger.info("Queued external starter model sync for %s", starter_model.source)
|
||||
|
||||
return queued_sources
|
||||
@@ -0,0 +1,72 @@
|
||||
from abc import ABC, abstractmethod
|
||||
from pathlib import Path
|
||||
from typing import Optional
|
||||
|
||||
from PIL.Image import Image as PILImageType
|
||||
|
||||
|
||||
class ImageFileStorageBase(ABC):
|
||||
"""Low-level service responsible for storing and retrieving image files."""
|
||||
|
||||
@abstractmethod
|
||||
def get(self, image_name: str, image_subfolder: str = "") -> PILImageType:
|
||||
"""Retrieves an image as PIL Image."""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def get_path(self, image_name: str, thumbnail: bool = False, image_subfolder: str = "") -> Path:
|
||||
"""Gets the internal path to an image or thumbnail."""
|
||||
pass
|
||||
|
||||
@property
|
||||
@abstractmethod
|
||||
def image_root(self) -> Path:
|
||||
"""Gets the root directory for full-size images."""
|
||||
pass
|
||||
|
||||
@property
|
||||
@abstractmethod
|
||||
def thumbnail_root(self) -> Path:
|
||||
"""Gets the root directory for thumbnails."""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def evict_cache_paths(self, paths: list[Path]) -> None:
|
||||
"""Evicts any cached image objects for the provided paths."""
|
||||
pass
|
||||
|
||||
# TODO: We need to validate paths before starlette makes the FileResponse, else we get a
|
||||
# 500 internal server error. I don't like having this method on the service.
|
||||
@abstractmethod
|
||||
def validate_path(self, path: str) -> bool:
|
||||
"""Validates the path given for an image or thumbnail."""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
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:
|
||||
"""Saves an image and a 256x256 WEBP thumbnail. Returns a tuple of the image name, thumbnail name, and created timestamp."""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def delete(self, image_name: str, image_subfolder: str = "") -> None:
|
||||
"""Deletes an image and its thumbnail (if one exists)."""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def get_workflow(self, image_name: str, image_subfolder: str = "") -> Optional[str]:
|
||||
"""Gets the workflow of an image."""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def get_graph(self, image_name: str, image_subfolder: str = "") -> Optional[str]:
|
||||
"""Gets the graph of an image."""
|
||||
pass
|
||||
@@ -0,0 +1,20 @@
|
||||
# TODO: Should these excpetions subclass existing python exceptions?
|
||||
class ImageFileNotFoundException(Exception):
|
||||
"""Raised when an image file is not found in storage."""
|
||||
|
||||
def __init__(self, message="Image file not found"):
|
||||
super().__init__(message)
|
||||
|
||||
|
||||
class ImageFileSaveException(Exception):
|
||||
"""Raised when an image cannot be saved."""
|
||||
|
||||
def __init__(self, message="Image file not saved"):
|
||||
super().__init__(message)
|
||||
|
||||
|
||||
class ImageFileDeleteException(Exception):
|
||||
"""Raised when an image cannot be deleted."""
|
||||
|
||||
def __init__(self, message="Image file not deleted"):
|
||||
super().__init__(message)
|
||||
@@ -0,0 +1,253 @@
|
||||
# Copyright (c) 2022 Kyle Schouviller (https://github.com/kyle0654) and the InvokeAI Team
|
||||
import io
|
||||
import zlib
|
||||
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
|
||||
|
||||
_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
|
||||
|
||||
|
||||
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
|
||||
|
||||
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
|
||||
|
||||
@property
|
||||
def image_root(self) -> Path:
|
||||
return self.__output_folder.resolve()
|
||||
|
||||
@property
|
||||
def thumbnail_root(self) -> Path:
|
||||
return self.__thumbnails_folder.resolve()
|
||||
|
||||
def evict_cache_paths(self, paths: list[Path]) -> None:
|
||||
for path in paths:
|
||||
self.__cache.pop(path.resolve(), 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 = self.__get_cache(image_path)
|
||||
if cache_item:
|
||||
return cache_item
|
||||
|
||||
image = Image.open(image_path)
|
||||
self.__set_cache(image_path, image)
|
||||
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:
|
||||
try:
|
||||
self.__validate_storage_folders()
|
||||
image_path = self.get_path(image_name, image_subfolder=image_subfolder)
|
||||
|
||||
# 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)
|
||||
|
||||
# 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_path = self.get_path(image_name, thumbnail=True, image_subfolder=image_subfolder)
|
||||
|
||||
# Ensure thumbnail subfolder directories exist
|
||||
thumbnail_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
thumbnail_image = make_thumbnail(image, thumbnail_size)
|
||||
thumbnail_image.save(thumbnail_path)
|
||||
|
||||
self.__set_cache(image_path, image)
|
||||
self.__set_cache(thumbnail_path, thumbnail_image)
|
||||
except Exception as e:
|
||||
raise ImageFileSaveException from e
|
||||
|
||||
def delete(self, image_name: str, image_subfolder: str = "") -> None:
|
||||
try:
|
||||
image_path = self.get_path(image_name, image_subfolder=image_subfolder)
|
||||
|
||||
if image_path.exists():
|
||||
image_path.unlink()
|
||||
if image_path in self.__cache:
|
||||
del self.__cache[image_path]
|
||||
|
||||
thumbnail_path = self.get_path(image_name, True, image_subfolder=image_subfolder)
|
||||
|
||||
if thumbnail_path.exists():
|
||||
thumbnail_path.unlink()
|
||||
if thumbnail_path in self.__cache:
|
||||
del self.__cache[thumbnail_path]
|
||||
except Exception as e:
|
||||
raise ImageFileDeleteException from e
|
||||
|
||||
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 __get_cache(self, image_name: Path) -> Optional[PILImageType]:
|
||||
return None if image_name not in self.__cache else self.__cache[image_name]
|
||||
|
||||
def __set_cache(self, image_name: Path, image: PILImageType):
|
||||
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]
|
||||
@@ -0,0 +1,58 @@
|
||||
from abc import ABC, abstractmethod
|
||||
from datetime import datetime
|
||||
|
||||
from invokeai.app.services.image_records.image_records_common import ImageCategory
|
||||
|
||||
|
||||
class ImageSubfolderStrategy(ABC):
|
||||
"""Base class for image subfolder strategies."""
|
||||
|
||||
@abstractmethod
|
||||
def get_subfolder(self, image_name: str, image_category: ImageCategory, is_intermediate: bool) -> str:
|
||||
"""Returns relative subfolder prefix (e.g. '2026/03/17', 'general'), or empty string for flat."""
|
||||
pass
|
||||
|
||||
|
||||
class FlatStrategy(ImageSubfolderStrategy):
|
||||
"""No subfolders - all images in one directory (default behavior)."""
|
||||
|
||||
def get_subfolder(self, image_name: str, image_category: ImageCategory, is_intermediate: bool) -> str:
|
||||
return ""
|
||||
|
||||
|
||||
class DateStrategy(ImageSubfolderStrategy):
|
||||
"""Organize images by date: YYYY/MM/DD."""
|
||||
|
||||
def get_subfolder(self, image_name: str, image_category: ImageCategory, is_intermediate: bool) -> str:
|
||||
now = datetime.now()
|
||||
return f"{now.year}/{now.month:02d}/{now.day:02d}"
|
||||
|
||||
|
||||
class TypeStrategy(ImageSubfolderStrategy):
|
||||
"""Organize images by category/type: general, intermediate, mask, control, etc."""
|
||||
|
||||
def get_subfolder(self, image_name: str, image_category: ImageCategory, is_intermediate: bool) -> str:
|
||||
if is_intermediate:
|
||||
return "intermediate"
|
||||
return image_category.value
|
||||
|
||||
|
||||
class HashStrategy(ImageSubfolderStrategy):
|
||||
"""Organize images by UUID prefix for filesystem performance (first 2 characters)."""
|
||||
|
||||
def get_subfolder(self, image_name: str, image_category: ImageCategory, is_intermediate: bool) -> str:
|
||||
return image_name[:2]
|
||||
|
||||
|
||||
def create_subfolder_strategy(strategy_name: str) -> ImageSubfolderStrategy:
|
||||
"""Factory function to create a subfolder strategy by name."""
|
||||
strategies: dict[str, type[ImageSubfolderStrategy]] = {
|
||||
"flat": FlatStrategy,
|
||||
"date": DateStrategy,
|
||||
"type": TypeStrategy,
|
||||
"hash": HashStrategy,
|
||||
}
|
||||
cls = strategies.get(strategy_name)
|
||||
if cls is None:
|
||||
raise ValueError(f"Unknown subfolder strategy: {strategy_name}. Valid options: {', '.join(strategies.keys())}")
|
||||
return cls()
|
||||
@@ -0,0 +1,3 @@
|
||||
from invokeai.app.services.image_moves.image_moves_default import ImageMoveService
|
||||
|
||||
__all__ = ["ImageMoveService"]
|
||||
@@ -0,0 +1,816 @@
|
||||
import os
|
||||
import tempfile
|
||||
import threading
|
||||
from concurrent.futures import Future, ThreadPoolExecutor
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
from typing import Literal, Sequence, cast
|
||||
|
||||
from PIL import Image
|
||||
|
||||
from invokeai.app.services.config import InvokeAIAppConfig
|
||||
from invokeai.app.services.image_files.image_files_base import ImageFileStorageBase
|
||||
from invokeai.app.services.image_records.image_records_common import ImageCategory
|
||||
from invokeai.app.services.session_queue.session_queue_common import DEFAULT_QUEUE_ID
|
||||
from invokeai.app.services.shared.sqlite.sqlite_database import SqliteDatabase
|
||||
from invokeai.app.util.thumbnails import make_thumbnail
|
||||
|
||||
MoveJobState = Literal["planned", "moving", "moved", "committed", "error"]
|
||||
ImageMoveBackgroundOperation = Literal["move_all", "recovery"]
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class PlannedImageMove:
|
||||
image_name: str
|
||||
old_subfolder: str
|
||||
new_subfolder: str
|
||||
is_intermediate: bool
|
||||
old_path: Path
|
||||
new_path: Path
|
||||
old_thumbnail_path: Path
|
||||
new_thumbnail_path: Path
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ImageMoveJob:
|
||||
id: int
|
||||
state: MoveJobState
|
||||
error_message: str | None
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ImageMoveResult:
|
||||
planned: int = 0
|
||||
committed: int = 0
|
||||
errors: int = 0
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ImageMoveBackgroundStatus:
|
||||
is_running: bool
|
||||
operation: ImageMoveBackgroundOperation | None
|
||||
active_job_id: int | None
|
||||
latest_job: ImageMoveJob | None
|
||||
last_error: str | None
|
||||
needs_move_count: int
|
||||
|
||||
|
||||
class ImageMoveJobAlreadyRunning(Exception):
|
||||
pass
|
||||
|
||||
|
||||
class ImageMoveQueueActive(Exception):
|
||||
pass
|
||||
|
||||
|
||||
class ImageMoveService:
|
||||
def __init__(
|
||||
self,
|
||||
db: SqliteDatabase,
|
||||
image_files: ImageFileStorageBase,
|
||||
config: InvokeAIAppConfig,
|
||||
logger,
|
||||
) -> None:
|
||||
self._db = db
|
||||
self.image_files = image_files
|
||||
self._config = config
|
||||
self._logger = logger
|
||||
self._executor = ThreadPoolExecutor(max_workers=1, thread_name_prefix="image-move")
|
||||
self._future_lock = threading.Lock()
|
||||
self._future: Future | None = None
|
||||
self._future_operation: ImageMoveBackgroundOperation | None = None
|
||||
self._last_background_error: str | None = None
|
||||
self._invoker = None
|
||||
self._session_queue = None
|
||||
|
||||
def start(self, invoker) -> None:
|
||||
self._invoker = invoker
|
||||
self._session_queue = getattr(invoker.services, "session_queue", None)
|
||||
result = self.startup_recovery()
|
||||
if result.committed > 0 or result.errors > 0:
|
||||
self._logger.info(
|
||||
"Image move startup recovery completed: committed=%s, errors=%s",
|
||||
result.committed,
|
||||
result.errors,
|
||||
)
|
||||
|
||||
def set_session_queue(self, session_queue) -> None:
|
||||
self._session_queue = session_queue
|
||||
|
||||
def stop(self, *args, **kwargs) -> None:
|
||||
self._executor.shutdown(wait=True, cancel_futures=False)
|
||||
|
||||
def start_background_move_all(self) -> ImageMoveBackgroundStatus:
|
||||
return self._start_background_operation("move_all", self.move_all_images, require_idle_queue=True)
|
||||
|
||||
def start_background_recovery(self) -> ImageMoveBackgroundStatus:
|
||||
return self._start_background_operation("recovery", self.startup_recovery)
|
||||
|
||||
def get_background_status(self) -> ImageMoveBackgroundStatus:
|
||||
with self._future_lock:
|
||||
self._refresh_finished_future_locked()
|
||||
return self._build_background_status_locked()
|
||||
|
||||
def is_maintenance_active(self) -> bool:
|
||||
with self._future_lock:
|
||||
self._refresh_finished_future_locked()
|
||||
is_running = self._future is not None and not self._future.done()
|
||||
operation_reserved = self._future_operation is not None
|
||||
return operation_reserved or is_running or self._get_active_job_id() is not None
|
||||
|
||||
def _assert_no_active_queue_work(self) -> None:
|
||||
session_queue = self._session_queue
|
||||
if session_queue is None and self._invoker is not None:
|
||||
session_queue = getattr(self._invoker.services, "session_queue", None)
|
||||
if session_queue is None:
|
||||
return
|
||||
queue_status = session_queue.get_queue_status(DEFAULT_QUEUE_ID)
|
||||
if queue_status.pending > 0 or queue_status.in_progress > 0:
|
||||
raise ImageMoveQueueActive("Cannot start image move while queue work is active")
|
||||
|
||||
def _start_background_operation(
|
||||
self,
|
||||
operation: ImageMoveBackgroundOperation,
|
||||
target,
|
||||
require_idle_queue: bool = False,
|
||||
) -> ImageMoveBackgroundStatus:
|
||||
with self._future_lock:
|
||||
self._refresh_finished_future_locked()
|
||||
if self._future_operation is not None or (self._future is not None and not self._future.done()):
|
||||
raise ImageMoveJobAlreadyRunning("An image move job is already running")
|
||||
active_job_id = self._get_active_job_id()
|
||||
if operation != "recovery" and active_job_id is not None:
|
||||
raise ImageMoveJobAlreadyRunning("An image move job is already active")
|
||||
self._last_background_error = None
|
||||
self._future_operation = operation
|
||||
|
||||
try:
|
||||
if require_idle_queue:
|
||||
self._assert_no_active_queue_work()
|
||||
future = self._executor.submit(self._run_background_operation, operation, target)
|
||||
except Exception:
|
||||
with self._future_lock:
|
||||
self._future_operation = None
|
||||
raise
|
||||
|
||||
with self._future_lock:
|
||||
self._future = future
|
||||
return self._build_background_status_locked()
|
||||
|
||||
def _run_background_operation(self, operation: ImageMoveBackgroundOperation, target) -> None:
|
||||
try:
|
||||
target()
|
||||
except Exception as e:
|
||||
self._record_background_error(str(e))
|
||||
self._logger.exception("Image move background operation failed: %s", operation)
|
||||
|
||||
def _record_background_error(self, message: str) -> None:
|
||||
with self._future_lock:
|
||||
self._last_background_error = message
|
||||
active_job_id = self._get_active_job_id()
|
||||
if active_job_id is not None:
|
||||
try:
|
||||
self.record_job_error_message(active_job_id, message)
|
||||
except Exception:
|
||||
self._logger.exception("Failed to record image move background error on active job")
|
||||
|
||||
def _refresh_finished_future_locked(self) -> None:
|
||||
if self._future is None or not self._future.done():
|
||||
return
|
||||
self._future = None
|
||||
self._future_operation = None
|
||||
|
||||
def _build_background_status_locked(self) -> ImageMoveBackgroundStatus:
|
||||
latest_job = self.get_latest_job()
|
||||
return ImageMoveBackgroundStatus(
|
||||
is_running=self._future is not None and not self._future.done(),
|
||||
operation=self._future_operation,
|
||||
active_job_id=self._get_active_job_id(),
|
||||
latest_job=latest_job,
|
||||
last_error=self._last_background_error,
|
||||
needs_move_count=self.count_images_needing_move(),
|
||||
)
|
||||
|
||||
def move_all_images(self) -> ImageMoveResult:
|
||||
recovered = self.startup_recovery()
|
||||
last_image_name = ""
|
||||
planned = 0
|
||||
committed = recovered.committed
|
||||
errors = recovered.errors
|
||||
|
||||
while True:
|
||||
moves, plan_errors = self._plan_batch(
|
||||
last_image_name=last_image_name, limit=100, record_missing_errors=True
|
||||
)
|
||||
errors += plan_errors
|
||||
if not moves:
|
||||
next_name = self._next_image_name(last_image_name)
|
||||
if next_name is None:
|
||||
break
|
||||
last_image_name = next_name
|
||||
continue
|
||||
|
||||
job_id = self.create_move_job(moves)
|
||||
planned += len(moves)
|
||||
try:
|
||||
self.perform_filesystem_moves(job_id)
|
||||
self.commit_database_updates(job_id)
|
||||
committed += len(moves)
|
||||
except Exception as e:
|
||||
errors += 1
|
||||
self.record_job_error_message(job_id, str(e))
|
||||
raise
|
||||
last_image_name = moves[-1].image_name
|
||||
|
||||
return ImageMoveResult(planned=planned, committed=committed, errors=errors)
|
||||
|
||||
def startup_recovery(self) -> ImageMoveResult:
|
||||
with self._db.transaction() as cursor:
|
||||
cursor.execute(
|
||||
"""--sql
|
||||
SELECT id FROM image_subfolder_move_jobs
|
||||
WHERE state IN ('planned', 'moving', 'moved')
|
||||
ORDER BY id;
|
||||
"""
|
||||
)
|
||||
job_ids = [cast(int, row[0]) for row in cursor.fetchall()]
|
||||
|
||||
committed = 0
|
||||
errors = 0
|
||||
for job_id in job_ids:
|
||||
try:
|
||||
self.complete_partial_filesystem_moves(job_id)
|
||||
self.cleanup_empty_source_dirs(job_id)
|
||||
self.commit_database_updates(job_id)
|
||||
committed += len(self._get_items(job_id))
|
||||
except Exception as e:
|
||||
errors += 1
|
||||
if self._is_unrecoverable_error(e):
|
||||
self.mark_job_unrecoverable(job_id, str(e))
|
||||
else:
|
||||
self.record_job_error_message(job_id, str(e))
|
||||
return ImageMoveResult(committed=committed, errors=errors)
|
||||
|
||||
def plan_batch(self, last_image_name: str, limit: int) -> list[PlannedImageMove]:
|
||||
moves, _errors = self._plan_batch(last_image_name=last_image_name, limit=limit, record_missing_errors=False)
|
||||
return moves
|
||||
|
||||
def count_images_needing_move(self) -> int:
|
||||
with self._db.transaction() as cursor:
|
||||
cursor.execute(
|
||||
"""--sql
|
||||
SELECT image_name, image_subfolder, image_category, is_intermediate, created_at
|
||||
FROM images
|
||||
WHERE deleted_at IS NULL;
|
||||
"""
|
||||
)
|
||||
rows = cursor.fetchall()
|
||||
|
||||
count = 0
|
||||
for row in rows:
|
||||
old_subfolder = cast(str, row["image_subfolder"] or "")
|
||||
new_subfolder = self._get_new_subfolder(
|
||||
image_name=cast(str, row["image_name"]),
|
||||
image_category=ImageCategory(row["image_category"]),
|
||||
is_intermediate=bool(row["is_intermediate"]),
|
||||
created_at=row["created_at"],
|
||||
)
|
||||
if new_subfolder != old_subfolder:
|
||||
count += 1
|
||||
return count
|
||||
|
||||
def _plan_batch(
|
||||
self, last_image_name: str, limit: int, record_missing_errors: bool
|
||||
) -> tuple[list[PlannedImageMove], int]:
|
||||
with self._db.transaction() as cursor:
|
||||
cursor.execute(
|
||||
"""--sql
|
||||
SELECT image_name, image_subfolder, image_category, is_intermediate, created_at
|
||||
FROM images
|
||||
WHERE image_name > ?
|
||||
AND deleted_at IS NULL
|
||||
ORDER BY image_name
|
||||
LIMIT ?;
|
||||
""",
|
||||
(last_image_name, limit),
|
||||
)
|
||||
rows = cursor.fetchall()
|
||||
|
||||
moves: list[PlannedImageMove] = []
|
||||
for row in rows:
|
||||
image_name = cast(str, row["image_name"])
|
||||
old_subfolder = cast(str, row["image_subfolder"] or "")
|
||||
new_subfolder = self._get_new_subfolder(
|
||||
image_name=image_name,
|
||||
image_category=ImageCategory(row["image_category"]),
|
||||
is_intermediate=bool(row["is_intermediate"]),
|
||||
created_at=row["created_at"],
|
||||
)
|
||||
if new_subfolder == old_subfolder:
|
||||
continue
|
||||
moves.append(
|
||||
PlannedImageMove(
|
||||
image_name=image_name,
|
||||
old_subfolder=old_subfolder,
|
||||
new_subfolder=new_subfolder,
|
||||
is_intermediate=bool(row["is_intermediate"]),
|
||||
old_path=self.image_files.get_path(image_name, image_subfolder=old_subfolder),
|
||||
new_path=self.image_files.get_path(image_name, image_subfolder=new_subfolder),
|
||||
old_thumbnail_path=self.image_files.get_path(
|
||||
image_name, thumbnail=True, image_subfolder=old_subfolder
|
||||
),
|
||||
new_thumbnail_path=self.image_files.get_path(
|
||||
image_name, thumbnail=True, image_subfolder=new_subfolder
|
||||
),
|
||||
)
|
||||
)
|
||||
errors = 0
|
||||
if record_missing_errors:
|
||||
moves, errors = self._record_missing_source_errors(moves)
|
||||
self.preflight_moves(moves)
|
||||
return moves, errors
|
||||
|
||||
def create_move_job(self, moves: Sequence[PlannedImageMove]) -> int:
|
||||
if not moves:
|
||||
raise ValueError("Cannot create an image move job with no items")
|
||||
with self._db.transaction() as cursor:
|
||||
cursor.execute(
|
||||
"""--sql
|
||||
SELECT 1
|
||||
FROM image_subfolder_move_jobs
|
||||
WHERE state NOT IN ('committed', 'error')
|
||||
LIMIT 1;
|
||||
"""
|
||||
)
|
||||
if cursor.fetchone() is not None:
|
||||
raise ValueError("Cannot create image move job while another active image move job exists")
|
||||
cursor.execute("INSERT INTO image_subfolder_move_jobs (state) VALUES ('planned');")
|
||||
job_id = cast(int, cursor.lastrowid)
|
||||
cursor.executemany(
|
||||
"""--sql
|
||||
INSERT INTO image_subfolder_move_items (
|
||||
job_id,
|
||||
image_name,
|
||||
old_subfolder,
|
||||
new_subfolder,
|
||||
is_intermediate,
|
||||
old_path,
|
||||
new_path,
|
||||
old_thumbnail_path,
|
||||
new_thumbnail_path,
|
||||
state
|
||||
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, 'planned');
|
||||
""",
|
||||
[
|
||||
(
|
||||
job_id,
|
||||
move.image_name,
|
||||
move.old_subfolder,
|
||||
move.new_subfolder,
|
||||
int(move.is_intermediate),
|
||||
str(move.old_path),
|
||||
str(move.new_path),
|
||||
str(move.old_thumbnail_path),
|
||||
str(move.new_thumbnail_path),
|
||||
)
|
||||
for move in moves
|
||||
],
|
||||
)
|
||||
return job_id
|
||||
|
||||
def create_error_move_job(self, move: PlannedImageMove, message: str) -> int:
|
||||
with self._db.transaction() as cursor:
|
||||
cursor.execute(
|
||||
"INSERT INTO image_subfolder_move_jobs (state, error_message) VALUES ('error', ?);",
|
||||
(message,),
|
||||
)
|
||||
job_id = cast(int, cursor.lastrowid)
|
||||
cursor.execute(
|
||||
"""--sql
|
||||
INSERT INTO image_subfolder_move_items (
|
||||
job_id,
|
||||
image_name,
|
||||
old_subfolder,
|
||||
new_subfolder,
|
||||
is_intermediate,
|
||||
old_path,
|
||||
new_path,
|
||||
old_thumbnail_path,
|
||||
new_thumbnail_path,
|
||||
state,
|
||||
error_message
|
||||
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, 'error', ?);
|
||||
""",
|
||||
(
|
||||
job_id,
|
||||
move.image_name,
|
||||
move.old_subfolder,
|
||||
move.new_subfolder,
|
||||
int(move.is_intermediate),
|
||||
str(move.old_path),
|
||||
str(move.new_path),
|
||||
str(move.old_thumbnail_path),
|
||||
str(move.new_thumbnail_path),
|
||||
message,
|
||||
),
|
||||
)
|
||||
return job_id
|
||||
|
||||
def preflight_moves(self, moves: Sequence[PlannedImageMove]) -> None:
|
||||
destinations: set[Path] = set()
|
||||
thumbnail_destinations: set[Path] = set()
|
||||
for move in moves:
|
||||
if not move.old_path.exists():
|
||||
if not move.is_intermediate:
|
||||
raise FileNotFoundError(f"Source image does not exist: {move.old_path}")
|
||||
continue
|
||||
if move.new_path.exists():
|
||||
raise FileExistsError(f"Destination image already exists: {move.new_path}")
|
||||
if move.old_path == move.new_path:
|
||||
raise ValueError(f"Old and new paths are identical for {move.image_name}")
|
||||
if move.new_path in destinations:
|
||||
raise ValueError(f"Duplicate destination path: {move.new_path}")
|
||||
destinations.add(move.new_path)
|
||||
if move.new_thumbnail_path in thumbnail_destinations:
|
||||
raise ValueError(f"Duplicate destination thumbnail path: {move.new_thumbnail_path}")
|
||||
thumbnail_destinations.add(move.new_thumbnail_path)
|
||||
if self._has_active_job_for_image(move.image_name):
|
||||
raise ValueError(f"Image {move.image_name} already has an active image move job")
|
||||
self._assert_same_filesystem(move.old_path, move.new_path)
|
||||
if move.old_thumbnail_path.exists():
|
||||
if move.new_thumbnail_path.exists():
|
||||
raise FileExistsError(f"Destination thumbnail already exists: {move.new_thumbnail_path}")
|
||||
self._assert_same_filesystem(move.old_thumbnail_path, move.new_thumbnail_path)
|
||||
|
||||
def _record_missing_source_errors(self, moves: Sequence[PlannedImageMove]) -> tuple[list[PlannedImageMove], int]:
|
||||
remaining_moves: list[PlannedImageMove] = []
|
||||
errors = 0
|
||||
for move in moves:
|
||||
if move.old_path.exists() or move.is_intermediate:
|
||||
remaining_moves.append(move)
|
||||
continue
|
||||
message = f"Source image does not exist: {move.old_path}"
|
||||
self.create_error_move_job(move, message)
|
||||
self._logger.error(message)
|
||||
errors += 1
|
||||
return remaining_moves, errors
|
||||
|
||||
def perform_filesystem_moves(self, job_id: int) -> None:
|
||||
self._set_job_state(job_id, "moving")
|
||||
self.complete_partial_filesystem_moves(job_id)
|
||||
self.cleanup_empty_source_dirs(job_id)
|
||||
self._set_job_state(job_id, "moved")
|
||||
|
||||
def complete_partial_filesystem_moves(self, job_id: int) -> None:
|
||||
items = self._get_items(job_id)
|
||||
if not items:
|
||||
raise ValueError(f"Image move job {job_id} has no items")
|
||||
for item in items:
|
||||
old_path = self.image_files.get_path(item.image_name, image_subfolder=item.old_subfolder)
|
||||
new_path = self.image_files.get_path(item.image_name, image_subfolder=item.new_subfolder)
|
||||
old_thumbnail_path = self.image_files.get_path(
|
||||
item.image_name, thumbnail=True, image_subfolder=item.old_subfolder
|
||||
)
|
||||
new_thumbnail_path = self.image_files.get_path(
|
||||
item.image_name, thumbnail=True, image_subfolder=item.new_subfolder
|
||||
)
|
||||
old_exists = old_path.exists()
|
||||
new_exists = new_path.exists()
|
||||
if old_exists and new_exists:
|
||||
raise RuntimeError(f"Both old and new image files exist for {item.image_name}")
|
||||
if not old_exists and not new_exists:
|
||||
if item.is_intermediate:
|
||||
self._mark_missing_intermediate_moved(
|
||||
job_id=job_id,
|
||||
image_name=item.image_name,
|
||||
old_path=old_path,
|
||||
new_path=new_path,
|
||||
old_thumbnail_path=old_thumbnail_path,
|
||||
new_thumbnail_path=new_thumbnail_path,
|
||||
)
|
||||
continue
|
||||
raise RuntimeError(f"Neither old nor new image file exists for {item.image_name}")
|
||||
if old_exists:
|
||||
new_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
os.replace(old_path, new_path)
|
||||
self._fsync_file(new_path)
|
||||
self._fsync_dir(new_path.parent)
|
||||
self._fsync_dir(old_path.parent)
|
||||
|
||||
old_thumbnail_exists = old_thumbnail_path.exists()
|
||||
new_thumbnail_exists = new_thumbnail_path.exists()
|
||||
if old_thumbnail_exists and new_thumbnail_exists:
|
||||
self._regenerate_thumbnail(new_path, new_thumbnail_path)
|
||||
old_thumbnail_path.unlink()
|
||||
self._fsync_dir(old_thumbnail_path.parent)
|
||||
elif old_thumbnail_exists and not new_thumbnail_exists:
|
||||
new_thumbnail_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
os.replace(old_thumbnail_path, new_thumbnail_path)
|
||||
self._fsync_file(new_thumbnail_path)
|
||||
self._fsync_dir(new_thumbnail_path.parent)
|
||||
self._fsync_dir(old_thumbnail_path.parent)
|
||||
elif not new_thumbnail_exists:
|
||||
self._regenerate_thumbnail(new_path, new_thumbnail_path)
|
||||
|
||||
self.image_files.evict_cache_paths([old_path, new_path, old_thumbnail_path, new_thumbnail_path])
|
||||
self.mark_item_moved(job_id, item.image_name)
|
||||
|
||||
def cleanup_empty_source_dirs(self, job_id: int) -> None:
|
||||
for item in self._get_items(job_id):
|
||||
self._remove_empty_parents(
|
||||
self.image_files.get_path(item.image_name, image_subfolder=item.old_subfolder).parent,
|
||||
self.image_files.image_root,
|
||||
)
|
||||
self._remove_empty_parents(
|
||||
self.image_files.get_path(item.image_name, thumbnail=True, image_subfolder=item.old_subfolder).parent,
|
||||
self.image_files.thumbnail_root,
|
||||
)
|
||||
|
||||
def commit_database_updates(self, job_id: int) -> None:
|
||||
with self._db.transaction() as cursor:
|
||||
cursor.execute(
|
||||
"""--sql
|
||||
UPDATE images
|
||||
SET image_subfolder = (
|
||||
SELECT item.new_subfolder
|
||||
FROM image_subfolder_move_items AS item
|
||||
WHERE item.job_id = ?
|
||||
AND item.image_name = images.image_name
|
||||
)
|
||||
WHERE image_name IN (
|
||||
SELECT image_name
|
||||
FROM image_subfolder_move_items
|
||||
WHERE job_id = ?
|
||||
AND state = 'moved'
|
||||
)
|
||||
AND image_subfolder = (
|
||||
SELECT item.old_subfolder
|
||||
FROM image_subfolder_move_items AS item
|
||||
WHERE item.job_id = ?
|
||||
AND item.image_name = images.image_name
|
||||
);
|
||||
""",
|
||||
(job_id, job_id, job_id),
|
||||
)
|
||||
cursor.execute(
|
||||
"""--sql
|
||||
SELECT COUNT(*)
|
||||
FROM image_subfolder_move_items AS item
|
||||
LEFT JOIN images ON images.image_name = item.image_name
|
||||
WHERE item.job_id = ?
|
||||
AND (
|
||||
images.image_name IS NULL
|
||||
OR images.deleted_at IS NOT NULL
|
||||
OR images.image_subfolder != item.new_subfolder
|
||||
);
|
||||
""",
|
||||
(job_id,),
|
||||
)
|
||||
invalid_count = cast(int, cursor.fetchone()[0])
|
||||
if invalid_count:
|
||||
raise RuntimeError(f"Image move job {job_id} failed commit validation")
|
||||
cursor.execute(
|
||||
"UPDATE image_subfolder_move_items SET state = 'committed' WHERE job_id = ?;",
|
||||
(job_id,),
|
||||
)
|
||||
cursor.execute(
|
||||
"UPDATE image_subfolder_move_jobs SET state = 'committed', error_message = NULL WHERE id = ?;",
|
||||
(job_id,),
|
||||
)
|
||||
|
||||
def mark_item_moved(self, job_id: int, image_name: str) -> None:
|
||||
with self._db.transaction() as cursor:
|
||||
cursor.execute(
|
||||
"UPDATE image_subfolder_move_items SET state = 'moved' WHERE job_id = ? AND image_name = ?;",
|
||||
(job_id, image_name),
|
||||
)
|
||||
|
||||
def record_job_error_message(self, job_id: int, message: str) -> None:
|
||||
with self._db.transaction() as cursor:
|
||||
cursor.execute(
|
||||
"UPDATE image_subfolder_move_jobs SET error_message = ? WHERE id = ?;",
|
||||
(message, job_id),
|
||||
)
|
||||
|
||||
def mark_job_unrecoverable(self, job_id: int, message: str) -> None:
|
||||
with self._db.transaction() as cursor:
|
||||
cursor.execute(
|
||||
"UPDATE image_subfolder_move_jobs SET state = 'error', error_message = ? WHERE id = ?;",
|
||||
(message, job_id),
|
||||
)
|
||||
cursor.execute(
|
||||
"UPDATE image_subfolder_move_items SET state = 'error', error_message = ? WHERE job_id = ?;",
|
||||
(message, job_id),
|
||||
)
|
||||
|
||||
def get_job(self, job_id: int) -> ImageMoveJob:
|
||||
with self._db.transaction() as cursor:
|
||||
cursor.execute(
|
||||
"SELECT id, state, error_message FROM image_subfolder_move_jobs WHERE id = ?;",
|
||||
(job_id,),
|
||||
)
|
||||
row = cursor.fetchone()
|
||||
if row is None:
|
||||
raise ValueError(f"Image move job not found: {job_id}")
|
||||
return ImageMoveJob(
|
||||
id=cast(int, row["id"]), state=cast(MoveJobState, row["state"]), error_message=row["error_message"]
|
||||
)
|
||||
|
||||
def get_latest_job(self) -> ImageMoveJob | None:
|
||||
with self._db.transaction() as cursor:
|
||||
cursor.execute(
|
||||
"""--sql
|
||||
SELECT id, state, error_message
|
||||
FROM image_subfolder_move_jobs
|
||||
ORDER BY id DESC
|
||||
LIMIT 1;
|
||||
"""
|
||||
)
|
||||
row = cursor.fetchone()
|
||||
if row is None:
|
||||
return None
|
||||
return ImageMoveJob(
|
||||
id=cast(int, row["id"]), state=cast(MoveJobState, row["state"]), error_message=row["error_message"]
|
||||
)
|
||||
|
||||
def _get_new_subfolder(
|
||||
self, image_name: str, image_category: ImageCategory, is_intermediate: bool, created_at: str | datetime
|
||||
) -> str:
|
||||
strategy = self._config.image_subfolder_strategy
|
||||
if strategy == "flat":
|
||||
return ""
|
||||
if strategy == "type":
|
||||
return "intermediate" if is_intermediate else image_category.value
|
||||
if strategy == "hash":
|
||||
return image_name[:2]
|
||||
if strategy == "date":
|
||||
timestamp = created_at if isinstance(created_at, datetime) else datetime.fromisoformat(created_at)
|
||||
return f"{timestamp.year}/{timestamp.month:02d}/{timestamp.day:02d}"
|
||||
raise ValueError(f"Unknown image subfolder strategy: {strategy}")
|
||||
|
||||
def _get_items(self, job_id: int) -> list[PlannedImageMove]:
|
||||
with self._db.transaction() as cursor:
|
||||
cursor.execute(
|
||||
"""--sql
|
||||
SELECT image_name, old_subfolder, new_subfolder, is_intermediate
|
||||
FROM image_subfolder_move_items
|
||||
WHERE job_id = ?
|
||||
ORDER BY image_name;
|
||||
""",
|
||||
(job_id,),
|
||||
)
|
||||
rows = cursor.fetchall()
|
||||
return [
|
||||
PlannedImageMove(
|
||||
image_name=row["image_name"],
|
||||
old_subfolder=row["old_subfolder"],
|
||||
new_subfolder=row["new_subfolder"],
|
||||
is_intermediate=bool(row["is_intermediate"]),
|
||||
old_path=self.image_files.get_path(row["image_name"], image_subfolder=row["old_subfolder"]),
|
||||
new_path=self.image_files.get_path(row["image_name"], image_subfolder=row["new_subfolder"]),
|
||||
old_thumbnail_path=self.image_files.get_path(
|
||||
row["image_name"], thumbnail=True, image_subfolder=row["old_subfolder"]
|
||||
),
|
||||
new_thumbnail_path=self.image_files.get_path(
|
||||
row["image_name"], thumbnail=True, image_subfolder=row["new_subfolder"]
|
||||
),
|
||||
)
|
||||
for row in rows
|
||||
]
|
||||
|
||||
def _set_job_state(self, job_id: int, state: MoveJobState) -> None:
|
||||
with self._db.transaction() as cursor:
|
||||
cursor.execute("UPDATE image_subfolder_move_jobs SET state = ? WHERE id = ?;", (state, job_id))
|
||||
|
||||
def _has_active_job_for_image(self, image_name: str) -> bool:
|
||||
with self._db.transaction() as cursor:
|
||||
cursor.execute(
|
||||
"""--sql
|
||||
SELECT 1
|
||||
FROM image_subfolder_move_items AS item
|
||||
JOIN image_subfolder_move_jobs AS job ON job.id = item.job_id
|
||||
WHERE item.image_name = ?
|
||||
AND job.state NOT IN ('committed', 'error')
|
||||
LIMIT 1;
|
||||
""",
|
||||
(image_name,),
|
||||
)
|
||||
return cursor.fetchone() is not None
|
||||
|
||||
def _get_active_job_id(self) -> int | None:
|
||||
with self._db.transaction() as cursor:
|
||||
cursor.execute(
|
||||
"""--sql
|
||||
SELECT id
|
||||
FROM image_subfolder_move_jobs
|
||||
WHERE state NOT IN ('committed', 'error')
|
||||
ORDER BY id
|
||||
LIMIT 1;
|
||||
"""
|
||||
)
|
||||
row = cursor.fetchone()
|
||||
return None if row is None else cast(int, row["id"])
|
||||
|
||||
def _next_image_name(self, last_image_name: str) -> str | None:
|
||||
with self._db.transaction() as cursor:
|
||||
cursor.execute(
|
||||
"""--sql
|
||||
SELECT image_name FROM images
|
||||
WHERE image_name > ?
|
||||
AND deleted_at IS NULL
|
||||
ORDER BY image_name
|
||||
LIMIT 1;
|
||||
""",
|
||||
(last_image_name,),
|
||||
)
|
||||
row = cursor.fetchone()
|
||||
return None if row is None else cast(str, row[0])
|
||||
|
||||
def _regenerate_thumbnail(self, image_path: Path, thumbnail_path: Path) -> None:
|
||||
thumbnail_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
with Image.open(image_path) as image:
|
||||
thumbnail = make_thumbnail(image)
|
||||
with tempfile.NamedTemporaryFile(
|
||||
dir=thumbnail_path.parent, prefix=f".{thumbnail_path.name}.", suffix=".tmp", delete=False
|
||||
) as temp_file:
|
||||
temp_path = Path(temp_file.name)
|
||||
try:
|
||||
thumbnail.save(temp_path, format="WEBP")
|
||||
self._fsync_file(temp_path)
|
||||
os.replace(temp_path, thumbnail_path)
|
||||
self._fsync_file(thumbnail_path)
|
||||
self._fsync_dir(thumbnail_path.parent)
|
||||
finally:
|
||||
temp_path.unlink(missing_ok=True)
|
||||
|
||||
def _mark_missing_intermediate_moved(
|
||||
self,
|
||||
job_id: int,
|
||||
image_name: str,
|
||||
old_path: Path,
|
||||
new_path: Path,
|
||||
old_thumbnail_path: Path,
|
||||
new_thumbnail_path: Path,
|
||||
) -> None:
|
||||
for path in (old_thumbnail_path, new_thumbnail_path):
|
||||
if path.exists():
|
||||
path.unlink()
|
||||
self._fsync_dir(path.parent)
|
||||
self.image_files.evict_cache_paths([old_path, new_path, old_thumbnail_path, new_thumbnail_path])
|
||||
self.mark_item_moved(job_id, image_name)
|
||||
|
||||
def _remove_empty_parents(self, start: Path, root: Path) -> None:
|
||||
root = root.resolve()
|
||||
current = start.resolve()
|
||||
while current != root and current.is_relative_to(root):
|
||||
try:
|
||||
current.rmdir()
|
||||
except OSError:
|
||||
return
|
||||
current = current.parent
|
||||
|
||||
def _assert_same_filesystem(self, source: Path, destination: Path) -> None:
|
||||
source_parent = source.parent
|
||||
destination_parent = self._nearest_existing_parent(destination.parent)
|
||||
if source_parent.stat().st_dev != destination_parent.stat().st_dev:
|
||||
raise ValueError(f"Cross-filesystem image move is not supported: {source} -> {destination}")
|
||||
|
||||
def _nearest_existing_parent(self, path: Path) -> Path:
|
||||
current = path
|
||||
while not current.exists():
|
||||
if current.parent == current:
|
||||
raise FileNotFoundError(f"No existing parent found for {path}")
|
||||
current = current.parent
|
||||
return current
|
||||
|
||||
def _fsync_file(self, path: Path) -> None:
|
||||
try:
|
||||
with path.open("rb") as file:
|
||||
os.fsync(file.fileno())
|
||||
except OSError as e:
|
||||
self._logger.debug("Unable to fsync file: %s: %s", path, e)
|
||||
|
||||
def _fsync_dir(self, path: Path) -> None:
|
||||
try:
|
||||
dir_fd = os.open(path, os.O_RDONLY)
|
||||
except OSError as e:
|
||||
self._logger.debug("Unable to open directory for fsync: %s: %s", path, e)
|
||||
return
|
||||
try:
|
||||
os.fsync(dir_fd)
|
||||
except OSError as e:
|
||||
self._logger.debug("Unable to fsync directory: %s: %s", path, e)
|
||||
finally:
|
||||
try:
|
||||
os.close(dir_fd)
|
||||
except OSError as e:
|
||||
self._logger.debug("Unable to close directory fsync handle: %s: %s", path, e)
|
||||
|
||||
def _is_unrecoverable_error(self, error: Exception) -> bool:
|
||||
return isinstance(error, RuntimeError) and (
|
||||
str(error).startswith("Both old and new image files exist")
|
||||
or str(error).startswith("Neither old nor new image file exists")
|
||||
or str(error).startswith("Image move job")
|
||||
and "has no items" in str(error)
|
||||
)
|
||||
@@ -0,0 +1,149 @@
|
||||
from abc import ABC, abstractmethod
|
||||
from datetime import datetime
|
||||
from typing import Optional
|
||||
|
||||
from invokeai.app.invocations.fields import MetadataField
|
||||
from invokeai.app.services.image_records.image_records_common import (
|
||||
ImageCategory,
|
||||
ImageNamesResult,
|
||||
ImageRecord,
|
||||
ImageRecordChanges,
|
||||
ResourceOrigin,
|
||||
)
|
||||
from invokeai.app.services.shared.pagination import OffsetPaginatedResults
|
||||
from invokeai.app.services.shared.sqlite.sqlite_common import SQLiteDirection
|
||||
from invokeai.app.services.virtual_boards.virtual_boards_common import VirtualSubBoardDTO
|
||||
|
||||
|
||||
class ImageRecordStorageBase(ABC):
|
||||
"""Low-level service responsible for interfacing with the image record store."""
|
||||
|
||||
# TODO: Implement an `update()` method
|
||||
|
||||
@abstractmethod
|
||||
def get(self, image_name: str) -> ImageRecord:
|
||||
"""Gets an image record."""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def get_metadata(self, image_name: str) -> Optional[MetadataField]:
|
||||
"""Gets an image's metadata'."""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def update(
|
||||
self,
|
||||
image_name: str,
|
||||
changes: ImageRecordChanges,
|
||||
) -> None:
|
||||
"""Updates an image record."""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def get_many(
|
||||
self,
|
||||
offset: int = 0,
|
||||
limit: int = 10,
|
||||
starred_first: bool = True,
|
||||
order_dir: SQLiteDirection = SQLiteDirection.Descending,
|
||||
image_origin: Optional[ResourceOrigin] = None,
|
||||
categories: Optional[list[ImageCategory]] = None,
|
||||
is_intermediate: Optional[bool] = None,
|
||||
board_id: Optional[str] = None,
|
||||
search_term: Optional[str] = None,
|
||||
user_id: Optional[str] = None,
|
||||
is_admin: bool = False,
|
||||
) -> OffsetPaginatedResults[ImageRecord]:
|
||||
"""Gets a page of image records. When board_id is 'none', filters by user_id for per-user uncategorized images unless is_admin is True."""
|
||||
pass
|
||||
|
||||
# TODO: The database has a nullable `deleted_at` column, currently unused.
|
||||
# Should we implement soft deletes? Would need coordination with ImageFileStorage.
|
||||
@abstractmethod
|
||||
def delete(self, image_name: str) -> None:
|
||||
"""Deletes an image record."""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def delete_many(self, image_names: list[str]) -> None:
|
||||
"""Deletes many image records."""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def delete_intermediates(self) -> list[tuple[str, str]]:
|
||||
"""Deletes all intermediate image records, returning a list of (image_name, image_subfolder) tuples."""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def get_intermediates_count(self, user_id: Optional[str] = None) -> int:
|
||||
"""Gets a count of intermediate images. If user_id is provided, only counts that user's intermediates."""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def save(
|
||||
self,
|
||||
image_name: str,
|
||||
image_origin: ResourceOrigin,
|
||||
image_category: ImageCategory,
|
||||
width: int,
|
||||
height: int,
|
||||
has_workflow: bool,
|
||||
is_intermediate: Optional[bool] = False,
|
||||
starred: Optional[bool] = False,
|
||||
session_id: Optional[str] = None,
|
||||
node_id: Optional[str] = None,
|
||||
metadata: Optional[str] = None,
|
||||
user_id: Optional[str] = None,
|
||||
image_subfolder: str = "",
|
||||
) -> datetime:
|
||||
"""Saves an image record."""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def get_user_id(self, image_name: str) -> Optional[str]:
|
||||
"""Gets the user_id of the image owner. Returns None if image not found."""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def get_most_recent_image_for_board(self, board_id: str) -> Optional[ImageRecord]:
|
||||
"""Gets the most recent image for a board."""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def get_image_names(
|
||||
self,
|
||||
starred_first: bool = True,
|
||||
order_dir: SQLiteDirection = SQLiteDirection.Descending,
|
||||
image_origin: Optional[ResourceOrigin] = None,
|
||||
categories: Optional[list[ImageCategory]] = None,
|
||||
is_intermediate: Optional[bool] = None,
|
||||
board_id: Optional[str] = None,
|
||||
search_term: Optional[str] = None,
|
||||
user_id: Optional[str] = None,
|
||||
is_admin: bool = False,
|
||||
) -> ImageNamesResult:
|
||||
"""Gets ordered list of image names with metadata for optimistic updates."""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def get_image_dates(
|
||||
self,
|
||||
user_id: Optional[str] = None,
|
||||
is_admin: bool = False,
|
||||
) -> list[VirtualSubBoardDTO]:
|
||||
"""Gets a list of dates with image counts, grouped by DATE(created_at)."""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def get_image_names_by_date(
|
||||
self,
|
||||
date: str,
|
||||
starred_first: bool = True,
|
||||
order_dir: SQLiteDirection = SQLiteDirection.Descending,
|
||||
categories: Optional[list[ImageCategory]] = None,
|
||||
search_term: Optional[str] = None,
|
||||
user_id: Optional[str] = None,
|
||||
is_admin: bool = False,
|
||||
) -> ImageNamesResult:
|
||||
"""Gets ordered list of image names for a specific date."""
|
||||
pass
|
||||
@@ -0,0 +1,235 @@
|
||||
# TODO: Should these excpetions subclass existing python exceptions?
|
||||
import datetime
|
||||
from enum import Enum
|
||||
from typing import Optional, Union
|
||||
|
||||
from pydantic import BaseModel, Field, StrictBool, StrictStr
|
||||
|
||||
from invokeai.app.util.metaenum import MetaEnum
|
||||
from invokeai.app.util.misc import get_iso_timestamp
|
||||
from invokeai.app.util.model_exclude_null import BaseModelExcludeNull
|
||||
|
||||
|
||||
class ResourceOrigin(str, Enum, metaclass=MetaEnum):
|
||||
"""The origin of a resource (eg image).
|
||||
|
||||
- INTERNAL: The resource was created by the application.
|
||||
- EXTERNAL: The resource was not created by the application.
|
||||
This may be a user-initiated upload, or an internal application upload (eg Canvas init image).
|
||||
"""
|
||||
|
||||
INTERNAL = "internal"
|
||||
"""The resource was created by the application."""
|
||||
EXTERNAL = "external"
|
||||
"""The resource was not created by the application.
|
||||
This may be a user-initiated upload, or an internal application upload (eg Canvas init image).
|
||||
"""
|
||||
|
||||
|
||||
class InvalidOriginException(ValueError):
|
||||
"""Raised when a provided value is not a valid ResourceOrigin.
|
||||
|
||||
Subclasses `ValueError`.
|
||||
"""
|
||||
|
||||
def __init__(self, message="Invalid resource origin."):
|
||||
super().__init__(message)
|
||||
|
||||
|
||||
class ImageCategory(str, Enum, metaclass=MetaEnum):
|
||||
"""The category of an image.
|
||||
|
||||
- GENERAL: The image is an output, init image, or otherwise an image without a specialized purpose.
|
||||
- MASK: The image is a mask image.
|
||||
- CONTROL: The image is a ControlNet control image.
|
||||
- USER: The image is a user-provide image.
|
||||
- OTHER: The image is some other type of image with a specialized purpose. To be used by external nodes.
|
||||
"""
|
||||
|
||||
GENERAL = "general"
|
||||
"""GENERAL: The image is an output, init image, or otherwise an image without a specialized purpose."""
|
||||
MASK = "mask"
|
||||
"""MASK: The image is a mask image."""
|
||||
CONTROL = "control"
|
||||
"""CONTROL: The image is a ControlNet control image."""
|
||||
USER = "user"
|
||||
"""USER: The image is a user-provide image."""
|
||||
OTHER = "other"
|
||||
"""OTHER: The image is some other type of image with a specialized purpose. To be used by external nodes."""
|
||||
|
||||
|
||||
IMAGE_CATEGORIES: list[ImageCategory] = [ImageCategory.GENERAL]
|
||||
ASSETS_CATEGORIES: list[ImageCategory] = [
|
||||
ImageCategory.CONTROL,
|
||||
ImageCategory.MASK,
|
||||
ImageCategory.USER,
|
||||
ImageCategory.OTHER,
|
||||
]
|
||||
|
||||
|
||||
class InvalidImageCategoryException(ValueError):
|
||||
"""Raised when a provided value is not a valid ImageCategory.
|
||||
|
||||
Subclasses `ValueError`.
|
||||
"""
|
||||
|
||||
def __init__(self, message="Invalid image category."):
|
||||
super().__init__(message)
|
||||
|
||||
|
||||
class ImageRecordNotFoundException(Exception):
|
||||
"""Raised when an image record is not found."""
|
||||
|
||||
def __init__(self, message="Image record not found"):
|
||||
super().__init__(message)
|
||||
|
||||
|
||||
class ImageRecordSaveException(Exception):
|
||||
"""Raised when an image record cannot be saved."""
|
||||
|
||||
def __init__(self, message="Image record not saved"):
|
||||
super().__init__(message)
|
||||
|
||||
|
||||
class ImageRecordDeleteException(Exception):
|
||||
"""Raised when an image record cannot be deleted."""
|
||||
|
||||
def __init__(self, message="Image record not deleted"):
|
||||
super().__init__(message)
|
||||
|
||||
|
||||
IMAGE_DTO_COLS = ", ".join(
|
||||
[
|
||||
"images." + c
|
||||
for c in [
|
||||
"image_name",
|
||||
"image_origin",
|
||||
"image_category",
|
||||
"width",
|
||||
"height",
|
||||
"session_id",
|
||||
"node_id",
|
||||
"has_workflow",
|
||||
"is_intermediate",
|
||||
"created_at",
|
||||
"updated_at",
|
||||
"deleted_at",
|
||||
"starred",
|
||||
"image_subfolder",
|
||||
]
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
class ImageRecord(BaseModelExcludeNull):
|
||||
"""Deserialized image record without metadata."""
|
||||
|
||||
image_name: str = Field(description="The unique name of the image.")
|
||||
"""The unique name of the image."""
|
||||
image_origin: ResourceOrigin = Field(description="The type of the image.")
|
||||
"""The origin of the image."""
|
||||
image_category: ImageCategory = Field(description="The category of the image.")
|
||||
"""The category of the image."""
|
||||
width: int = Field(description="The width of the image in px.")
|
||||
"""The actual width of the image in px. This may be different from the width in metadata."""
|
||||
height: int = Field(description="The height of the image in px.")
|
||||
"""The actual height of the image in px. This may be different from the height in metadata."""
|
||||
created_at: Union[datetime.datetime, str] = Field(description="The created timestamp of the image.")
|
||||
"""The created timestamp of the image."""
|
||||
updated_at: Union[datetime.datetime, str] = Field(description="The updated timestamp of the image.")
|
||||
"""The updated timestamp of the image."""
|
||||
deleted_at: Optional[Union[datetime.datetime, str]] = Field(
|
||||
default=None, description="The deleted timestamp of the image."
|
||||
)
|
||||
"""The deleted timestamp of the image."""
|
||||
is_intermediate: bool = Field(description="Whether this is an intermediate image.")
|
||||
"""Whether this is an intermediate image."""
|
||||
session_id: Optional[str] = Field(
|
||||
default=None,
|
||||
description="The session ID that generated this image, if it is a generated image.",
|
||||
)
|
||||
"""The session ID that generated this image, if it is a generated image."""
|
||||
node_id: Optional[str] = Field(
|
||||
default=None,
|
||||
description="The node ID that generated this image, if it is a generated image.",
|
||||
)
|
||||
"""The node ID that generated this image, if it is a generated image."""
|
||||
starred: bool = Field(description="Whether this image is starred.")
|
||||
"""Whether this image is starred."""
|
||||
has_workflow: bool = Field(description="Whether this image has a workflow.")
|
||||
image_subfolder: str = Field(default="", description="The subfolder where the image is stored on disk.")
|
||||
|
||||
|
||||
class ImageRecordChanges(BaseModelExcludeNull, extra="allow"):
|
||||
"""A set of changes to apply to an image record.
|
||||
|
||||
Only limited changes are valid:
|
||||
- `image_category`: change the category of an image
|
||||
- `session_id`: change the session associated with an image
|
||||
- `is_intermediate`: change the image's `is_intermediate` flag
|
||||
- `starred`: change whether the image is starred
|
||||
"""
|
||||
|
||||
image_category: Optional[ImageCategory] = Field(default=None, description="The image's new category.")
|
||||
"""The image's new category."""
|
||||
session_id: Optional[StrictStr] = Field(
|
||||
default=None,
|
||||
description="The image's new session ID.",
|
||||
)
|
||||
"""The image's new session ID."""
|
||||
is_intermediate: Optional[StrictBool] = Field(default=None, description="The image's new `is_intermediate` flag.")
|
||||
"""The image's new `is_intermediate` flag."""
|
||||
starred: Optional[StrictBool] = Field(default=None, description="The image's new `starred` state")
|
||||
"""The image's new `starred` state."""
|
||||
|
||||
|
||||
def deserialize_image_record(image_dict: dict) -> ImageRecord:
|
||||
"""Deserializes an image record."""
|
||||
|
||||
# Retrieve all the values, setting "reasonable" defaults if they are not present.
|
||||
|
||||
# TODO: do we really need to handle default values here? ideally the data is the correct shape...
|
||||
image_name = image_dict.get("image_name", "unknown")
|
||||
image_origin = ResourceOrigin(image_dict.get("image_origin", ResourceOrigin.INTERNAL.value))
|
||||
image_category = ImageCategory(image_dict.get("image_category", ImageCategory.GENERAL.value))
|
||||
width = image_dict.get("width", 0)
|
||||
height = image_dict.get("height", 0)
|
||||
session_id = image_dict.get("session_id", None)
|
||||
node_id = image_dict.get("node_id", None)
|
||||
created_at = image_dict.get("created_at", get_iso_timestamp())
|
||||
updated_at = image_dict.get("updated_at", get_iso_timestamp())
|
||||
deleted_at = image_dict.get("deleted_at", get_iso_timestamp())
|
||||
is_intermediate = image_dict.get("is_intermediate", False)
|
||||
starred = image_dict.get("starred", False)
|
||||
has_workflow = image_dict.get("has_workflow", False)
|
||||
image_subfolder = image_dict.get("image_subfolder", "")
|
||||
|
||||
return ImageRecord(
|
||||
image_name=image_name,
|
||||
image_origin=image_origin,
|
||||
image_category=image_category,
|
||||
width=width,
|
||||
height=height,
|
||||
session_id=session_id,
|
||||
node_id=node_id,
|
||||
created_at=created_at,
|
||||
updated_at=updated_at,
|
||||
deleted_at=deleted_at,
|
||||
is_intermediate=is_intermediate,
|
||||
starred=starred,
|
||||
has_workflow=has_workflow,
|
||||
image_subfolder=image_subfolder,
|
||||
)
|
||||
|
||||
|
||||
class ImageCollectionCounts(BaseModel):
|
||||
starred_count: int = Field(description="The number of starred images in the collection.")
|
||||
unstarred_count: int = Field(description="The number of unstarred images in the collection.")
|
||||
|
||||
|
||||
class ImageNamesResult(BaseModel):
|
||||
"""Response containing ordered image names with metadata for optimistic updates."""
|
||||
|
||||
image_names: list[str] = Field(description="Ordered list of image names")
|
||||
starred_count: int = Field(description="Number of starred images (when starred_first=True)")
|
||||
total_count: int = Field(description="Total number of images matching the query")
|
||||
@@ -0,0 +1,651 @@
|
||||
import sqlite3
|
||||
from datetime import datetime
|
||||
from typing import Optional, Union, cast
|
||||
|
||||
from invokeai.app.invocations.fields import MetadataField, MetadataFieldValidator
|
||||
from invokeai.app.services.image_records.image_records_base import ImageRecordStorageBase
|
||||
from invokeai.app.services.image_records.image_records_common import (
|
||||
IMAGE_DTO_COLS,
|
||||
ImageCategory,
|
||||
ImageNamesResult,
|
||||
ImageRecord,
|
||||
ImageRecordChanges,
|
||||
ImageRecordDeleteException,
|
||||
ImageRecordNotFoundException,
|
||||
ImageRecordSaveException,
|
||||
ResourceOrigin,
|
||||
deserialize_image_record,
|
||||
)
|
||||
from invokeai.app.services.shared.pagination import OffsetPaginatedResults
|
||||
from invokeai.app.services.shared.sqlite.sqlite_common import SQLiteDirection
|
||||
from invokeai.app.services.shared.sqlite.sqlite_database import SqliteDatabase
|
||||
from invokeai.app.services.virtual_boards.virtual_boards_common import VirtualSubBoardDTO
|
||||
|
||||
|
||||
class SqliteImageRecordStorage(ImageRecordStorageBase):
|
||||
def __init__(self, db: SqliteDatabase) -> None:
|
||||
super().__init__()
|
||||
self._db = db
|
||||
|
||||
def get(self, image_name: str) -> ImageRecord:
|
||||
with self._db.transaction() as cursor:
|
||||
try:
|
||||
cursor.execute(
|
||||
f"""--sql
|
||||
SELECT {IMAGE_DTO_COLS} FROM images
|
||||
WHERE image_name = ?;
|
||||
""",
|
||||
(image_name,),
|
||||
)
|
||||
|
||||
result = cast(Optional[sqlite3.Row], cursor.fetchone())
|
||||
except sqlite3.Error as e:
|
||||
raise ImageRecordNotFoundException from e
|
||||
|
||||
if not result:
|
||||
raise ImageRecordNotFoundException
|
||||
|
||||
return deserialize_image_record(dict(result))
|
||||
|
||||
def get_user_id(self, image_name: str) -> Optional[str]:
|
||||
with self._db.transaction() as cursor:
|
||||
cursor.execute(
|
||||
"""--sql
|
||||
SELECT user_id FROM images
|
||||
WHERE image_name = ?;
|
||||
""",
|
||||
(image_name,),
|
||||
)
|
||||
result = cast(Optional[sqlite3.Row], cursor.fetchone())
|
||||
if not result:
|
||||
return None
|
||||
return cast(Optional[str], dict(result).get("user_id"))
|
||||
|
||||
def get_metadata(self, image_name: str) -> Optional[MetadataField]:
|
||||
with self._db.transaction() as cursor:
|
||||
try:
|
||||
cursor.execute(
|
||||
"""--sql
|
||||
SELECT metadata FROM images
|
||||
WHERE image_name = ?;
|
||||
""",
|
||||
(image_name,),
|
||||
)
|
||||
|
||||
result = cast(Optional[sqlite3.Row], cursor.fetchone())
|
||||
|
||||
except sqlite3.Error as e:
|
||||
raise ImageRecordNotFoundException from e
|
||||
|
||||
if not result:
|
||||
raise ImageRecordNotFoundException
|
||||
|
||||
as_dict = dict(result)
|
||||
metadata_raw = cast(Optional[str], as_dict.get("metadata", None))
|
||||
return MetadataFieldValidator.validate_json(metadata_raw) if metadata_raw is not None else None
|
||||
|
||||
def update(
|
||||
self,
|
||||
image_name: str,
|
||||
changes: ImageRecordChanges,
|
||||
) -> None:
|
||||
with self._db.transaction() as cursor:
|
||||
try:
|
||||
# Change the category of the image
|
||||
if changes.image_category is not None:
|
||||
cursor.execute(
|
||||
"""--sql
|
||||
UPDATE images
|
||||
SET image_category = ?
|
||||
WHERE image_name = ?;
|
||||
""",
|
||||
(changes.image_category, image_name),
|
||||
)
|
||||
|
||||
# Change the session associated with the image
|
||||
if changes.session_id is not None:
|
||||
cursor.execute(
|
||||
"""--sql
|
||||
UPDATE images
|
||||
SET session_id = ?
|
||||
WHERE image_name = ?;
|
||||
""",
|
||||
(changes.session_id, image_name),
|
||||
)
|
||||
|
||||
# Change the image's `is_intermediate`` flag
|
||||
if changes.is_intermediate is not None:
|
||||
cursor.execute(
|
||||
"""--sql
|
||||
UPDATE images
|
||||
SET is_intermediate = ?
|
||||
WHERE image_name = ?;
|
||||
""",
|
||||
(changes.is_intermediate, image_name),
|
||||
)
|
||||
|
||||
# Change the image's `starred`` state
|
||||
if changes.starred is not None:
|
||||
cursor.execute(
|
||||
"""--sql
|
||||
UPDATE images
|
||||
SET starred = ?
|
||||
WHERE image_name = ?;
|
||||
""",
|
||||
(changes.starred, image_name),
|
||||
)
|
||||
|
||||
except sqlite3.Error as e:
|
||||
raise ImageRecordSaveException from e
|
||||
|
||||
def get_many(
|
||||
self,
|
||||
offset: int = 0,
|
||||
limit: int = 10,
|
||||
starred_first: bool = True,
|
||||
order_dir: SQLiteDirection = SQLiteDirection.Descending,
|
||||
image_origin: Optional[ResourceOrigin] = None,
|
||||
categories: Optional[list[ImageCategory]] = None,
|
||||
is_intermediate: Optional[bool] = None,
|
||||
board_id: Optional[str] = None,
|
||||
search_term: Optional[str] = None,
|
||||
user_id: Optional[str] = None,
|
||||
is_admin: bool = False,
|
||||
) -> OffsetPaginatedResults[ImageRecord]:
|
||||
with self._db.transaction() as cursor:
|
||||
# Manually build two queries - one for the count, one for the records
|
||||
count_query = """--sql
|
||||
SELECT COUNT(*)
|
||||
FROM images
|
||||
LEFT JOIN board_images ON board_images.image_name = images.image_name
|
||||
WHERE 1=1
|
||||
"""
|
||||
|
||||
images_query = f"""--sql
|
||||
SELECT {IMAGE_DTO_COLS}
|
||||
FROM images
|
||||
LEFT JOIN board_images ON board_images.image_name = images.image_name
|
||||
WHERE 1=1
|
||||
"""
|
||||
|
||||
query_conditions = ""
|
||||
query_params: list[Union[int, str, bool]] = []
|
||||
|
||||
if image_origin is not None:
|
||||
query_conditions += """--sql
|
||||
AND images.image_origin = ?
|
||||
"""
|
||||
query_params.append(image_origin.value)
|
||||
|
||||
if categories is not None:
|
||||
# Convert the enum values to unique list of strings
|
||||
category_strings = [c.value for c in set(categories)]
|
||||
# Create the correct length of placeholders
|
||||
placeholders = ",".join("?" * len(category_strings))
|
||||
|
||||
query_conditions += f"""--sql
|
||||
AND images.image_category IN ( {placeholders} )
|
||||
"""
|
||||
|
||||
# Unpack the included categories into the query params
|
||||
for c in category_strings:
|
||||
query_params.append(c)
|
||||
|
||||
if is_intermediate is not None:
|
||||
query_conditions += """--sql
|
||||
AND images.is_intermediate = ?
|
||||
"""
|
||||
|
||||
query_params.append(is_intermediate)
|
||||
|
||||
# board_id of "none" is reserved for images without a board
|
||||
if board_id == "none":
|
||||
query_conditions += """--sql
|
||||
AND board_images.board_id IS NULL
|
||||
"""
|
||||
# For uncategorized images, filter by user_id to ensure per-user isolation
|
||||
# Admin users can see all uncategorized images from all users
|
||||
if user_id is not None and not is_admin:
|
||||
query_conditions += """--sql
|
||||
AND images.user_id = ?
|
||||
"""
|
||||
query_params.append(user_id)
|
||||
elif board_id is not None:
|
||||
query_conditions += """--sql
|
||||
AND board_images.board_id = ?
|
||||
"""
|
||||
query_params.append(board_id)
|
||||
|
||||
# Search term condition
|
||||
if search_term:
|
||||
query_conditions += """--sql
|
||||
AND (
|
||||
images.metadata LIKE ?
|
||||
OR images.created_at LIKE ?
|
||||
)
|
||||
"""
|
||||
query_params.append(f"%{search_term.lower()}%")
|
||||
query_params.append(f"%{search_term.lower()}%")
|
||||
|
||||
if starred_first:
|
||||
query_pagination = f"""--sql
|
||||
ORDER BY images.starred DESC, images.created_at {order_dir.value} LIMIT ? OFFSET ?
|
||||
"""
|
||||
else:
|
||||
query_pagination = f"""--sql
|
||||
ORDER BY images.created_at {order_dir.value} LIMIT ? OFFSET ?
|
||||
"""
|
||||
|
||||
# Final images query with pagination
|
||||
images_query += query_conditions + query_pagination + ";"
|
||||
# Add all the parameters
|
||||
images_params = query_params.copy()
|
||||
# Add the pagination parameters
|
||||
images_params.extend([limit, offset])
|
||||
|
||||
# Build the list of images, deserializing each row
|
||||
cursor.execute(images_query, images_params)
|
||||
result = cast(list[sqlite3.Row], cursor.fetchall())
|
||||
|
||||
images = [deserialize_image_record(dict(r)) for r in result]
|
||||
|
||||
# Set up and execute the count query, without pagination
|
||||
count_query += query_conditions + ";"
|
||||
count_params = query_params.copy()
|
||||
cursor.execute(count_query, count_params)
|
||||
count = cast(int, cursor.fetchone()[0])
|
||||
|
||||
return OffsetPaginatedResults(items=images, offset=offset, limit=limit, total=count)
|
||||
|
||||
def delete(self, image_name: str) -> None:
|
||||
with self._db.transaction() as cursor:
|
||||
try:
|
||||
cursor.execute(
|
||||
"""--sql
|
||||
DELETE FROM images
|
||||
WHERE image_name = ?;
|
||||
""",
|
||||
(image_name,),
|
||||
)
|
||||
except sqlite3.Error as e:
|
||||
raise ImageRecordDeleteException from e
|
||||
|
||||
def delete_many(self, image_names: list[str]) -> None:
|
||||
with self._db.transaction() as cursor:
|
||||
try:
|
||||
placeholders = ",".join("?" for _ in image_names)
|
||||
|
||||
# Construct the SQLite query with the placeholders
|
||||
query = f"DELETE FROM images WHERE image_name IN ({placeholders})"
|
||||
|
||||
# Execute the query with the list of IDs as parameters
|
||||
cursor.execute(query, image_names)
|
||||
|
||||
except sqlite3.Error as e:
|
||||
raise ImageRecordDeleteException from e
|
||||
|
||||
def get_intermediates_count(self, user_id: Optional[str] = None) -> int:
|
||||
with self._db.transaction() as cursor:
|
||||
query = "SELECT COUNT(*) FROM images WHERE is_intermediate = TRUE"
|
||||
params: list[str] = []
|
||||
if user_id is not None:
|
||||
query += " AND user_id = ?"
|
||||
params.append(user_id)
|
||||
cursor.execute(query, params)
|
||||
count = cast(int, cursor.fetchone()[0])
|
||||
return count
|
||||
|
||||
def delete_intermediates(self) -> list[tuple[str, str]]:
|
||||
"""Deletes all intermediate image records.
|
||||
|
||||
Returns a list of (image_name, image_subfolder) tuples for file cleanup.
|
||||
"""
|
||||
with self._db.transaction() as cursor:
|
||||
try:
|
||||
cursor.execute(
|
||||
"""--sql
|
||||
SELECT image_name, image_subfolder FROM images
|
||||
WHERE is_intermediate = TRUE;
|
||||
"""
|
||||
)
|
||||
result = cast(list[sqlite3.Row], cursor.fetchall())
|
||||
image_name_subfolder_pairs = [(r[0], r[1]) for r in result]
|
||||
cursor.execute(
|
||||
"""--sql
|
||||
DELETE FROM images
|
||||
WHERE is_intermediate = TRUE;
|
||||
"""
|
||||
)
|
||||
except sqlite3.Error as e:
|
||||
raise ImageRecordDeleteException from e
|
||||
return image_name_subfolder_pairs
|
||||
|
||||
def save(
|
||||
self,
|
||||
image_name: str,
|
||||
image_origin: ResourceOrigin,
|
||||
image_category: ImageCategory,
|
||||
width: int,
|
||||
height: int,
|
||||
has_workflow: bool,
|
||||
is_intermediate: Optional[bool] = False,
|
||||
starred: Optional[bool] = False,
|
||||
session_id: Optional[str] = None,
|
||||
node_id: Optional[str] = None,
|
||||
metadata: Optional[str] = None,
|
||||
user_id: Optional[str] = None,
|
||||
image_subfolder: str = "",
|
||||
) -> datetime:
|
||||
with self._db.transaction() as cursor:
|
||||
try:
|
||||
cursor.execute(
|
||||
"""--sql
|
||||
INSERT OR IGNORE INTO images (
|
||||
image_name,
|
||||
image_origin,
|
||||
image_category,
|
||||
width,
|
||||
height,
|
||||
node_id,
|
||||
session_id,
|
||||
metadata,
|
||||
is_intermediate,
|
||||
starred,
|
||||
has_workflow,
|
||||
user_id,
|
||||
image_subfolder
|
||||
)
|
||||
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?);
|
||||
""",
|
||||
(
|
||||
image_name,
|
||||
image_origin.value,
|
||||
image_category.value,
|
||||
width,
|
||||
height,
|
||||
node_id,
|
||||
session_id,
|
||||
metadata,
|
||||
is_intermediate,
|
||||
starred,
|
||||
has_workflow,
|
||||
user_id or "system",
|
||||
image_subfolder,
|
||||
),
|
||||
)
|
||||
|
||||
cursor.execute(
|
||||
"""--sql
|
||||
SELECT created_at
|
||||
FROM images
|
||||
WHERE image_name = ?;
|
||||
""",
|
||||
(image_name,),
|
||||
)
|
||||
|
||||
created_at = datetime.fromisoformat(cursor.fetchone()[0])
|
||||
|
||||
except sqlite3.Error as e:
|
||||
raise ImageRecordSaveException from e
|
||||
return created_at
|
||||
|
||||
def get_most_recent_image_for_board(self, board_id: str) -> Optional[ImageRecord]:
|
||||
with self._db.transaction() as cursor:
|
||||
cursor.execute(
|
||||
"""--sql
|
||||
SELECT images.*
|
||||
FROM images
|
||||
JOIN board_images ON images.image_name = board_images.image_name
|
||||
WHERE board_images.board_id = ?
|
||||
AND images.is_intermediate = FALSE
|
||||
ORDER BY images.starred DESC, images.created_at DESC
|
||||
LIMIT 1;
|
||||
""",
|
||||
(board_id,),
|
||||
)
|
||||
|
||||
result = cast(Optional[sqlite3.Row], cursor.fetchone())
|
||||
|
||||
if result is None:
|
||||
return None
|
||||
|
||||
return deserialize_image_record(dict(result))
|
||||
|
||||
def get_image_names(
|
||||
self,
|
||||
starred_first: bool = True,
|
||||
order_dir: SQLiteDirection = SQLiteDirection.Descending,
|
||||
image_origin: Optional[ResourceOrigin] = None,
|
||||
categories: Optional[list[ImageCategory]] = None,
|
||||
is_intermediate: Optional[bool] = None,
|
||||
board_id: Optional[str] = None,
|
||||
search_term: Optional[str] = None,
|
||||
user_id: Optional[str] = None,
|
||||
is_admin: bool = False,
|
||||
) -> ImageNamesResult:
|
||||
with self._db.transaction() as cursor:
|
||||
# Build query conditions (reused for both starred count and image names queries)
|
||||
query_conditions = ""
|
||||
query_params: list[Union[int, str, bool]] = []
|
||||
|
||||
if image_origin is not None:
|
||||
query_conditions += """--sql
|
||||
AND images.image_origin = ?
|
||||
"""
|
||||
query_params.append(image_origin.value)
|
||||
|
||||
if categories is not None:
|
||||
category_strings = [c.value for c in set(categories)]
|
||||
placeholders = ",".join("?" * len(category_strings))
|
||||
query_conditions += f"""--sql
|
||||
AND images.image_category IN ( {placeholders} )
|
||||
"""
|
||||
for c in category_strings:
|
||||
query_params.append(c)
|
||||
|
||||
if is_intermediate is not None:
|
||||
query_conditions += """--sql
|
||||
AND images.is_intermediate = ?
|
||||
"""
|
||||
query_params.append(is_intermediate)
|
||||
|
||||
if board_id == "none":
|
||||
query_conditions += """--sql
|
||||
AND board_images.board_id IS NULL
|
||||
"""
|
||||
# For uncategorized images, filter by user_id to ensure per-user isolation
|
||||
# Admin users can see all uncategorized images from all users
|
||||
if user_id is not None and not is_admin:
|
||||
query_conditions += """--sql
|
||||
AND images.user_id = ?
|
||||
"""
|
||||
query_params.append(user_id)
|
||||
elif board_id is not None:
|
||||
query_conditions += """--sql
|
||||
AND board_images.board_id = ?
|
||||
"""
|
||||
query_params.append(board_id)
|
||||
|
||||
if search_term:
|
||||
query_conditions += """--sql
|
||||
AND (
|
||||
images.metadata LIKE ?
|
||||
OR images.created_at LIKE ?
|
||||
)
|
||||
"""
|
||||
query_params.append(f"%{search_term.lower()}%")
|
||||
query_params.append(f"%{search_term.lower()}%")
|
||||
|
||||
# Get starred count if starred_first is enabled
|
||||
starred_count = 0
|
||||
if starred_first:
|
||||
starred_count_query = f"""--sql
|
||||
SELECT COUNT(*)
|
||||
FROM images
|
||||
LEFT JOIN board_images ON board_images.image_name = images.image_name
|
||||
WHERE images.starred = TRUE AND (1=1{query_conditions})
|
||||
"""
|
||||
cursor.execute(starred_count_query, query_params)
|
||||
starred_count = cast(int, cursor.fetchone()[0])
|
||||
|
||||
# Get all image names with proper ordering
|
||||
if starred_first:
|
||||
names_query = f"""--sql
|
||||
SELECT images.image_name
|
||||
FROM images
|
||||
LEFT JOIN board_images ON board_images.image_name = images.image_name
|
||||
WHERE 1=1{query_conditions}
|
||||
ORDER BY images.starred DESC, images.created_at {order_dir.value}
|
||||
"""
|
||||
else:
|
||||
names_query = f"""--sql
|
||||
SELECT images.image_name
|
||||
FROM images
|
||||
LEFT JOIN board_images ON board_images.image_name = images.image_name
|
||||
WHERE 1=1{query_conditions}
|
||||
ORDER BY images.created_at {order_dir.value}
|
||||
"""
|
||||
|
||||
cursor.execute(names_query, query_params)
|
||||
result = cast(list[sqlite3.Row], cursor.fetchall())
|
||||
image_names = [row[0] for row in result]
|
||||
|
||||
return ImageNamesResult(image_names=image_names, starred_count=starred_count, total_count=len(image_names))
|
||||
|
||||
def get_image_dates(
|
||||
self,
|
||||
user_id: Optional[str] = None,
|
||||
is_admin: bool = False,
|
||||
) -> list[VirtualSubBoardDTO]:
|
||||
with self._db.transaction() as cursor:
|
||||
query_conditions = ""
|
||||
query_params: list[Union[int, str, bool]] = []
|
||||
|
||||
# Only non-intermediate images
|
||||
query_conditions += """--sql
|
||||
AND images.is_intermediate = 0
|
||||
"""
|
||||
|
||||
# User isolation for non-admin users
|
||||
if user_id is not None and not is_admin:
|
||||
query_conditions += """--sql
|
||||
AND images.user_id = ?
|
||||
"""
|
||||
query_params.append(user_id)
|
||||
|
||||
query = f"""--sql
|
||||
SELECT
|
||||
DATE(images.created_at) as date,
|
||||
SUM(CASE WHEN images.image_category = 'general' THEN 1 ELSE 0 END) as image_count,
|
||||
SUM(CASE WHEN images.image_category != 'general' THEN 1 ELSE 0 END) as asset_count,
|
||||
(
|
||||
SELECT i2.image_name FROM images i2
|
||||
WHERE DATE(i2.created_at) = DATE(images.created_at)
|
||||
AND i2.is_intermediate = 0
|
||||
ORDER BY i2.created_at DESC LIMIT 1
|
||||
) as cover_image_name
|
||||
FROM images
|
||||
WHERE 1=1
|
||||
{query_conditions}
|
||||
GROUP BY DATE(images.created_at)
|
||||
ORDER BY date DESC;
|
||||
"""
|
||||
|
||||
cursor.execute(query, query_params)
|
||||
result = cast(list[sqlite3.Row], cursor.fetchall())
|
||||
|
||||
return [
|
||||
VirtualSubBoardDTO(
|
||||
virtual_board_id=f"by_date:{dict(row)['date']}",
|
||||
board_name=dict(row)["date"],
|
||||
date=dict(row)["date"],
|
||||
image_count=dict(row)["image_count"],
|
||||
asset_count=dict(row)["asset_count"],
|
||||
cover_image_name=dict(row)["cover_image_name"],
|
||||
)
|
||||
for row in result
|
||||
]
|
||||
|
||||
def get_image_names_by_date(
|
||||
self,
|
||||
date: str,
|
||||
starred_first: bool = True,
|
||||
order_dir: SQLiteDirection = SQLiteDirection.Descending,
|
||||
categories: Optional[list[ImageCategory]] = None,
|
||||
search_term: Optional[str] = None,
|
||||
user_id: Optional[str] = None,
|
||||
is_admin: bool = False,
|
||||
) -> ImageNamesResult:
|
||||
with self._db.transaction() as cursor:
|
||||
query_conditions = ""
|
||||
query_params: list[Union[int, str, bool]] = []
|
||||
|
||||
# Filter by date
|
||||
query_conditions += """--sql
|
||||
AND DATE(images.created_at) = ?
|
||||
"""
|
||||
query_params.append(date)
|
||||
|
||||
# Only non-intermediate images
|
||||
query_conditions += """--sql
|
||||
AND images.is_intermediate = 0
|
||||
"""
|
||||
|
||||
if categories is not None:
|
||||
category_strings = [c.value for c in set(categories)]
|
||||
placeholders = ",".join("?" * len(category_strings))
|
||||
query_conditions += f"""--sql
|
||||
AND images.image_category IN ( {placeholders} )
|
||||
"""
|
||||
for c in category_strings:
|
||||
query_params.append(c)
|
||||
|
||||
# User isolation for non-admin users
|
||||
if user_id is not None and not is_admin:
|
||||
query_conditions += """--sql
|
||||
AND images.user_id = ?
|
||||
"""
|
||||
query_params.append(user_id)
|
||||
|
||||
if search_term:
|
||||
query_conditions += """--sql
|
||||
AND (
|
||||
images.metadata LIKE ?
|
||||
OR images.created_at LIKE ?
|
||||
)
|
||||
"""
|
||||
query_params.append(f"%{search_term.lower()}%")
|
||||
query_params.append(f"%{search_term.lower()}%")
|
||||
|
||||
# Get starred count if starred_first is enabled
|
||||
starred_count = 0
|
||||
if starred_first:
|
||||
starred_count_query = f"""--sql
|
||||
SELECT COUNT(*)
|
||||
FROM images
|
||||
WHERE images.starred = TRUE AND (1=1{query_conditions})
|
||||
"""
|
||||
cursor.execute(starred_count_query, query_params)
|
||||
starred_count = cast(int, cursor.fetchone()[0])
|
||||
|
||||
# Get all image names with proper ordering
|
||||
if starred_first:
|
||||
names_query = f"""--sql
|
||||
SELECT images.image_name
|
||||
FROM images
|
||||
WHERE 1=1{query_conditions}
|
||||
ORDER BY images.starred DESC, images.created_at {order_dir.value}
|
||||
"""
|
||||
else:
|
||||
names_query = f"""--sql
|
||||
SELECT images.image_name
|
||||
FROM images
|
||||
WHERE 1=1{query_conditions}
|
||||
ORDER BY images.created_at {order_dir.value}
|
||||
"""
|
||||
|
||||
cursor.execute(names_query, query_params)
|
||||
result = cast(list[sqlite3.Row], cursor.fetchall())
|
||||
image_names = [row[0] for row in result]
|
||||
|
||||
return ImageNamesResult(image_names=image_names, starred_count=starred_count, total_count=len(image_names))
|
||||
@@ -0,0 +1,169 @@
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import Callable, Optional
|
||||
|
||||
from PIL.Image import Image as PILImageType
|
||||
|
||||
from invokeai.app.invocations.fields import MetadataField
|
||||
from invokeai.app.services.image_records.image_records_common import (
|
||||
ImageCategory,
|
||||
ImageNamesResult,
|
||||
ImageRecord,
|
||||
ImageRecordChanges,
|
||||
ResourceOrigin,
|
||||
)
|
||||
from invokeai.app.services.images.images_common import ImageDTO
|
||||
from invokeai.app.services.shared.pagination import OffsetPaginatedResults
|
||||
from invokeai.app.services.shared.sqlite.sqlite_common import SQLiteDirection
|
||||
|
||||
|
||||
class ImageServiceABC(ABC):
|
||||
"""High-level service for image management."""
|
||||
|
||||
_on_changed_callbacks: list[Callable[[ImageDTO], None]]
|
||||
_on_deleted_callbacks: list[Callable[[str], None]]
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._on_changed_callbacks = []
|
||||
self._on_deleted_callbacks = []
|
||||
|
||||
def on_changed(self, on_changed: Callable[[ImageDTO], None]) -> None:
|
||||
"""Register a callback for when an image is changed"""
|
||||
self._on_changed_callbacks.append(on_changed)
|
||||
|
||||
def on_deleted(self, on_deleted: Callable[[str], None]) -> None:
|
||||
"""Register a callback for when an image is deleted"""
|
||||
self._on_deleted_callbacks.append(on_deleted)
|
||||
|
||||
def _on_changed(self, item: ImageDTO) -> None:
|
||||
for callback in self._on_changed_callbacks:
|
||||
callback(item)
|
||||
|
||||
def _on_deleted(self, item_id: str) -> None:
|
||||
for callback in self._on_deleted_callbacks:
|
||||
callback(item_id)
|
||||
|
||||
@abstractmethod
|
||||
def create(
|
||||
self,
|
||||
image: PILImageType,
|
||||
image_origin: ResourceOrigin,
|
||||
image_category: ImageCategory,
|
||||
node_id: Optional[str] = None,
|
||||
session_id: Optional[str] = None,
|
||||
board_id: Optional[str] = None,
|
||||
is_intermediate: Optional[bool] = False,
|
||||
metadata: Optional[str] = None,
|
||||
workflow: Optional[str] = None,
|
||||
graph: Optional[str] = None,
|
||||
user_id: Optional[str] = None,
|
||||
) -> ImageDTO:
|
||||
"""Creates an image, storing the file and its metadata."""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def update(
|
||||
self,
|
||||
image_name: str,
|
||||
changes: ImageRecordChanges,
|
||||
) -> ImageDTO:
|
||||
"""Updates an image."""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def get_pil_image(self, image_name: str) -> PILImageType:
|
||||
"""Gets an image as a PIL image."""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def get_record(self, image_name: str) -> ImageRecord:
|
||||
"""Gets an image record."""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def get_dto(self, image_name: str) -> ImageDTO:
|
||||
"""Gets an image DTO."""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def get_metadata(self, image_name: str) -> Optional[MetadataField]:
|
||||
"""Gets an image's metadata."""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def get_workflow(self, image_name: str) -> Optional[str]:
|
||||
"""Gets an image's workflow."""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def get_graph(self, image_name: str) -> Optional[str]:
|
||||
"""Gets an image's workflow."""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def get_path(self, image_name: str, thumbnail: bool = False) -> str:
|
||||
"""Gets an image's path."""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def validate_path(self, path: str) -> bool:
|
||||
"""Validates an image's path."""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def get_url(self, image_name: str, thumbnail: bool = False) -> str:
|
||||
"""Gets an image's or thumbnail's URL."""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def get_many(
|
||||
self,
|
||||
offset: int = 0,
|
||||
limit: int = 10,
|
||||
starred_first: bool = True,
|
||||
order_dir: SQLiteDirection = SQLiteDirection.Descending,
|
||||
image_origin: Optional[ResourceOrigin] = None,
|
||||
categories: Optional[list[ImageCategory]] = None,
|
||||
is_intermediate: Optional[bool] = None,
|
||||
board_id: Optional[str] = None,
|
||||
search_term: Optional[str] = None,
|
||||
user_id: Optional[str] = None,
|
||||
is_admin: bool = False,
|
||||
) -> OffsetPaginatedResults[ImageDTO]:
|
||||
"""Gets a paginated list of image DTOs with starred images first when starred_first=True."""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def delete(self, image_name: str):
|
||||
"""Deletes an image."""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def delete_intermediates(self) -> int:
|
||||
"""Deletes all intermediate images."""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def get_intermediates_count(self, user_id: Optional[str] = None) -> int:
|
||||
"""Gets the number of intermediate images. If user_id is provided, only counts that user's intermediates."""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def delete_images_on_board(self, board_id: str):
|
||||
"""Deletes all images on a board."""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def get_image_names(
|
||||
self,
|
||||
starred_first: bool = True,
|
||||
order_dir: SQLiteDirection = SQLiteDirection.Descending,
|
||||
image_origin: Optional[ResourceOrigin] = None,
|
||||
categories: Optional[list[ImageCategory]] = None,
|
||||
is_intermediate: Optional[bool] = None,
|
||||
board_id: Optional[str] = None,
|
||||
search_term: Optional[str] = None,
|
||||
user_id: Optional[str] = None,
|
||||
is_admin: bool = False,
|
||||
) -> ImageNamesResult:
|
||||
"""Gets ordered list of image names with metadata for optimistic updates."""
|
||||
pass
|
||||
@@ -0,0 +1,65 @@
|
||||
from typing import Optional
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
from invokeai.app.services.image_records.image_records_common import ImageRecord
|
||||
from invokeai.app.util.model_exclude_null import BaseModelExcludeNull
|
||||
|
||||
|
||||
class ImageUrlsDTO(BaseModelExcludeNull):
|
||||
"""The URLs for an image and its thumbnail."""
|
||||
|
||||
image_name: str = Field(description="The unique name of the image.")
|
||||
"""The unique name of the image."""
|
||||
image_url: str = Field(description="The URL of the image.")
|
||||
"""The URL of the image."""
|
||||
thumbnail_url: str = Field(description="The URL of the image's thumbnail.")
|
||||
"""The URL of the image's thumbnail."""
|
||||
|
||||
|
||||
class ImageDTO(ImageRecord, ImageUrlsDTO):
|
||||
"""Deserialized image record, enriched for the frontend."""
|
||||
|
||||
board_id: Optional[str] = Field(
|
||||
default=None, description="The id of the board the image belongs to, if one exists."
|
||||
)
|
||||
"""The id of the board the image belongs to, if one exists."""
|
||||
|
||||
|
||||
def image_record_to_dto(
|
||||
image_record: ImageRecord,
|
||||
image_url: str,
|
||||
thumbnail_url: str,
|
||||
board_id: Optional[str],
|
||||
) -> ImageDTO:
|
||||
"""Converts an image record to an image DTO."""
|
||||
return ImageDTO(
|
||||
**image_record.model_dump(),
|
||||
image_url=image_url,
|
||||
thumbnail_url=thumbnail_url,
|
||||
board_id=board_id,
|
||||
)
|
||||
|
||||
|
||||
class ResultWithAffectedBoards(BaseModel):
|
||||
affected_boards: list[str] = Field(description="The ids of boards affected by the delete operation")
|
||||
|
||||
|
||||
class DeleteImagesResult(ResultWithAffectedBoards):
|
||||
deleted_images: list[str] = Field(description="The names of the images that were deleted")
|
||||
|
||||
|
||||
class StarredImagesResult(ResultWithAffectedBoards):
|
||||
starred_images: list[str] = Field(description="The names of the images that were starred")
|
||||
|
||||
|
||||
class UnstarredImagesResult(ResultWithAffectedBoards):
|
||||
unstarred_images: list[str] = Field(description="The names of the images that were unstarred")
|
||||
|
||||
|
||||
class AddImagesToBoardResult(ResultWithAffectedBoards):
|
||||
added_images: list[str] = Field(description="The image names that were added to the board")
|
||||
|
||||
|
||||
class RemoveImagesFromBoardResult(ResultWithAffectedBoards):
|
||||
removed_images: list[str] = Field(description="The image names that were removed from their board")
|
||||
@@ -0,0 +1,371 @@
|
||||
from typing import Optional
|
||||
|
||||
from PIL.Image import Image as PILImageType
|
||||
|
||||
from invokeai.app.invocations.fields import MetadataField
|
||||
from invokeai.app.services.image_files.image_files_common import (
|
||||
ImageFileDeleteException,
|
||||
ImageFileNotFoundException,
|
||||
ImageFileSaveException,
|
||||
)
|
||||
from invokeai.app.services.image_files.image_subfolder_strategy import create_subfolder_strategy
|
||||
from invokeai.app.services.image_records.image_records_common import (
|
||||
ImageCategory,
|
||||
ImageNamesResult,
|
||||
ImageRecord,
|
||||
ImageRecordChanges,
|
||||
ImageRecordDeleteException,
|
||||
ImageRecordNotFoundException,
|
||||
ImageRecordSaveException,
|
||||
InvalidImageCategoryException,
|
||||
InvalidOriginException,
|
||||
ResourceOrigin,
|
||||
)
|
||||
from invokeai.app.services.images.images_base import ImageServiceABC
|
||||
from invokeai.app.services.images.images_common import ImageDTO, image_record_to_dto
|
||||
from invokeai.app.services.invoker import Invoker
|
||||
from invokeai.app.services.shared.pagination import OffsetPaginatedResults
|
||||
from invokeai.app.services.shared.sqlite.sqlite_common import SQLiteDirection
|
||||
|
||||
|
||||
class ImageService(ImageServiceABC):
|
||||
__invoker: Invoker
|
||||
|
||||
def start(self, invoker: Invoker) -> None:
|
||||
self.__invoker = invoker
|
||||
|
||||
def create(
|
||||
self,
|
||||
image: PILImageType,
|
||||
image_origin: ResourceOrigin,
|
||||
image_category: ImageCategory,
|
||||
node_id: Optional[str] = None,
|
||||
session_id: Optional[str] = None,
|
||||
board_id: Optional[str] = None,
|
||||
is_intermediate: Optional[bool] = False,
|
||||
metadata: Optional[str] = None,
|
||||
workflow: Optional[str] = None,
|
||||
graph: Optional[str] = None,
|
||||
user_id: Optional[str] = None,
|
||||
) -> ImageDTO:
|
||||
if image_origin not in ResourceOrigin:
|
||||
raise InvalidOriginException
|
||||
|
||||
if image_category not in ImageCategory:
|
||||
raise InvalidImageCategoryException
|
||||
|
||||
image_name = self.__invoker.services.names.create_image_name()
|
||||
|
||||
# Compute subfolder based on configured strategy
|
||||
strategy_name = self.__invoker.services.configuration.image_subfolder_strategy
|
||||
strategy = create_subfolder_strategy(strategy_name)
|
||||
image_subfolder = strategy.get_subfolder(image_name, image_category, is_intermediate or False)
|
||||
|
||||
(width, height) = image.size
|
||||
|
||||
try:
|
||||
# TODO: Consider using a transaction here to ensure consistency between storage and database
|
||||
self.__invoker.services.image_records.save(
|
||||
# Non-nullable fields
|
||||
image_name=image_name,
|
||||
image_origin=image_origin,
|
||||
image_category=image_category,
|
||||
width=width,
|
||||
height=height,
|
||||
has_workflow=workflow is not None or graph is not None,
|
||||
# Meta fields
|
||||
is_intermediate=is_intermediate,
|
||||
# Nullable fields
|
||||
node_id=node_id,
|
||||
metadata=metadata,
|
||||
session_id=session_id,
|
||||
user_id=user_id,
|
||||
image_subfolder=image_subfolder,
|
||||
)
|
||||
if board_id is not None:
|
||||
try:
|
||||
self.__invoker.services.board_image_records.add_image_to_board(
|
||||
board_id=board_id, image_name=image_name
|
||||
)
|
||||
except Exception as e:
|
||||
self.__invoker.services.logger.warning(f"Failed to add image to board {board_id}: {str(e)}")
|
||||
self.__invoker.services.image_files.save(
|
||||
image_name=image_name,
|
||||
image=image,
|
||||
metadata=metadata,
|
||||
workflow=workflow,
|
||||
graph=graph,
|
||||
image_subfolder=image_subfolder,
|
||||
)
|
||||
image_dto = self.get_dto(image_name)
|
||||
|
||||
self._on_changed(image_dto)
|
||||
return image_dto
|
||||
except ImageRecordSaveException:
|
||||
self.__invoker.services.logger.error("Failed to save image record")
|
||||
raise
|
||||
except ImageFileSaveException:
|
||||
self.__invoker.services.logger.error("Failed to save image file")
|
||||
raise
|
||||
except Exception as e:
|
||||
self.__invoker.services.logger.error(f"Problem saving image record and file: {str(e)}")
|
||||
raise e
|
||||
|
||||
def update(
|
||||
self,
|
||||
image_name: str,
|
||||
changes: ImageRecordChanges,
|
||||
) -> ImageDTO:
|
||||
try:
|
||||
self.__invoker.services.image_records.update(image_name, changes)
|
||||
image_dto = self.get_dto(image_name)
|
||||
self._on_changed(image_dto)
|
||||
return image_dto
|
||||
except ImageRecordSaveException:
|
||||
self.__invoker.services.logger.error("Failed to update image record")
|
||||
raise
|
||||
except Exception as e:
|
||||
self.__invoker.services.logger.error("Problem updating image record")
|
||||
raise e
|
||||
|
||||
def get_pil_image(self, image_name: str) -> PILImageType:
|
||||
try:
|
||||
record = self.__invoker.services.image_records.get(image_name)
|
||||
return self.__invoker.services.image_files.get(image_name, image_subfolder=record.image_subfolder)
|
||||
except ImageFileNotFoundException:
|
||||
self.__invoker.services.logger.error("Failed to get image file")
|
||||
raise
|
||||
except Exception as e:
|
||||
self.__invoker.services.logger.error("Problem getting image file")
|
||||
raise e
|
||||
|
||||
def get_record(self, image_name: str) -> ImageRecord:
|
||||
try:
|
||||
return self.__invoker.services.image_records.get(image_name)
|
||||
except ImageRecordNotFoundException:
|
||||
self.__invoker.services.logger.error("Image record not found")
|
||||
raise
|
||||
except Exception as e:
|
||||
self.__invoker.services.logger.error("Problem getting image record")
|
||||
raise e
|
||||
|
||||
def get_dto(self, image_name: str) -> ImageDTO:
|
||||
try:
|
||||
image_record = self.__invoker.services.image_records.get(image_name)
|
||||
|
||||
image_dto = image_record_to_dto(
|
||||
image_record=image_record,
|
||||
image_url=self.__invoker.services.urls.get_image_url(image_name),
|
||||
thumbnail_url=self.__invoker.services.urls.get_image_url(image_name, True),
|
||||
board_id=self.__invoker.services.board_image_records.get_board_for_image(image_name),
|
||||
)
|
||||
|
||||
return image_dto
|
||||
except ImageRecordNotFoundException:
|
||||
self.__invoker.services.logger.error("Image record not found")
|
||||
raise
|
||||
except Exception as e:
|
||||
self.__invoker.services.logger.error("Problem getting image DTO")
|
||||
raise e
|
||||
|
||||
def get_metadata(self, image_name: str) -> Optional[MetadataField]:
|
||||
try:
|
||||
return self.__invoker.services.image_records.get_metadata(image_name)
|
||||
except ImageRecordNotFoundException:
|
||||
self.__invoker.services.logger.error("Image record not found")
|
||||
raise
|
||||
except Exception as e:
|
||||
self.__invoker.services.logger.error("Problem getting image metadata")
|
||||
raise e
|
||||
|
||||
def get_workflow(self, image_name: str) -> Optional[str]:
|
||||
try:
|
||||
record = self.__invoker.services.image_records.get(image_name)
|
||||
return self.__invoker.services.image_files.get_workflow(image_name, image_subfolder=record.image_subfolder)
|
||||
except ImageFileNotFoundException:
|
||||
self.__invoker.services.logger.error("Image file not found")
|
||||
raise
|
||||
except Exception:
|
||||
self.__invoker.services.logger.error("Problem getting image workflow")
|
||||
raise
|
||||
|
||||
def get_graph(self, image_name: str) -> Optional[str]:
|
||||
try:
|
||||
record = self.__invoker.services.image_records.get(image_name)
|
||||
return self.__invoker.services.image_files.get_graph(image_name, image_subfolder=record.image_subfolder)
|
||||
except ImageFileNotFoundException:
|
||||
self.__invoker.services.logger.error("Image file not found")
|
||||
raise
|
||||
except Exception:
|
||||
self.__invoker.services.logger.error("Problem getting image graph")
|
||||
raise
|
||||
|
||||
def get_path(self, image_name: str, thumbnail: bool = False) -> str:
|
||||
try:
|
||||
record = self.__invoker.services.image_records.get(image_name)
|
||||
return str(
|
||||
self.__invoker.services.image_files.get_path(
|
||||
image_name, thumbnail, image_subfolder=record.image_subfolder
|
||||
)
|
||||
)
|
||||
except Exception as e:
|
||||
self.__invoker.services.logger.error("Problem getting image path")
|
||||
raise e
|
||||
|
||||
def validate_path(self, path: str) -> bool:
|
||||
try:
|
||||
return self.__invoker.services.image_files.validate_path(path)
|
||||
except Exception as e:
|
||||
self.__invoker.services.logger.error("Problem validating image path")
|
||||
raise e
|
||||
|
||||
def get_url(self, image_name: str, thumbnail: bool = False) -> str:
|
||||
try:
|
||||
return self.__invoker.services.urls.get_image_url(image_name, thumbnail)
|
||||
except Exception as e:
|
||||
self.__invoker.services.logger.error("Problem getting image path")
|
||||
raise e
|
||||
|
||||
def get_many(
|
||||
self,
|
||||
offset: int = 0,
|
||||
limit: int = 10,
|
||||
starred_first: bool = True,
|
||||
order_dir: SQLiteDirection = SQLiteDirection.Descending,
|
||||
image_origin: Optional[ResourceOrigin] = None,
|
||||
categories: Optional[list[ImageCategory]] = None,
|
||||
is_intermediate: Optional[bool] = None,
|
||||
board_id: Optional[str] = None,
|
||||
search_term: Optional[str] = None,
|
||||
user_id: Optional[str] = None,
|
||||
is_admin: bool = False,
|
||||
) -> OffsetPaginatedResults[ImageDTO]:
|
||||
try:
|
||||
results = self.__invoker.services.image_records.get_many(
|
||||
offset,
|
||||
limit,
|
||||
starred_first,
|
||||
order_dir,
|
||||
image_origin,
|
||||
categories,
|
||||
is_intermediate,
|
||||
board_id,
|
||||
search_term,
|
||||
user_id,
|
||||
is_admin,
|
||||
)
|
||||
|
||||
image_dtos = [
|
||||
image_record_to_dto(
|
||||
image_record=r,
|
||||
image_url=self.__invoker.services.urls.get_image_url(r.image_name),
|
||||
thumbnail_url=self.__invoker.services.urls.get_image_url(r.image_name, True),
|
||||
board_id=self.__invoker.services.board_image_records.get_board_for_image(r.image_name),
|
||||
)
|
||||
for r in results.items
|
||||
]
|
||||
|
||||
return OffsetPaginatedResults[ImageDTO](
|
||||
items=image_dtos,
|
||||
offset=results.offset,
|
||||
limit=results.limit,
|
||||
total=results.total,
|
||||
)
|
||||
except Exception as e:
|
||||
self.__invoker.services.logger.error("Problem getting paginated image DTOs")
|
||||
raise e
|
||||
|
||||
def delete(self, image_name: str):
|
||||
try:
|
||||
record = self.__invoker.services.image_records.get(image_name)
|
||||
self.__invoker.services.image_files.delete(image_name, image_subfolder=record.image_subfolder)
|
||||
self.__invoker.services.image_records.delete(image_name)
|
||||
self._on_deleted(image_name)
|
||||
except ImageRecordDeleteException:
|
||||
self.__invoker.services.logger.error("Failed to delete image record")
|
||||
raise
|
||||
except ImageFileDeleteException:
|
||||
self.__invoker.services.logger.error("Failed to delete image file")
|
||||
raise
|
||||
except Exception as e:
|
||||
self.__invoker.services.logger.error("Problem deleting image record and file")
|
||||
raise e
|
||||
|
||||
def delete_images_on_board(self, board_id: str):
|
||||
try:
|
||||
image_names = self.__invoker.services.board_image_records.get_all_board_image_names_for_board(
|
||||
board_id,
|
||||
categories=None,
|
||||
is_intermediate=None,
|
||||
)
|
||||
for image_name in image_names:
|
||||
try:
|
||||
record = self.__invoker.services.image_records.get(image_name)
|
||||
self.__invoker.services.image_files.delete(image_name, image_subfolder=record.image_subfolder)
|
||||
except Exception:
|
||||
pass
|
||||
self.__invoker.services.image_records.delete_many(image_names)
|
||||
for image_name in image_names:
|
||||
self._on_deleted(image_name)
|
||||
except ImageRecordDeleteException:
|
||||
self.__invoker.services.logger.error("Failed to delete image records")
|
||||
raise
|
||||
except ImageFileDeleteException:
|
||||
self.__invoker.services.logger.error("Failed to delete image files")
|
||||
raise
|
||||
except Exception as e:
|
||||
self.__invoker.services.logger.error(f"Problem deleting image records and files: {str(e)}")
|
||||
raise e
|
||||
|
||||
def delete_intermediates(self) -> int:
|
||||
try:
|
||||
image_name_subfolder_pairs = self.__invoker.services.image_records.delete_intermediates()
|
||||
count = len(image_name_subfolder_pairs)
|
||||
for image_name, image_subfolder in image_name_subfolder_pairs:
|
||||
self.__invoker.services.image_files.delete(image_name, image_subfolder=image_subfolder)
|
||||
self._on_deleted(image_name)
|
||||
return count
|
||||
except ImageRecordDeleteException:
|
||||
self.__invoker.services.logger.error("Failed to delete image records")
|
||||
raise
|
||||
except ImageFileDeleteException:
|
||||
self.__invoker.services.logger.error("Failed to delete image files")
|
||||
raise
|
||||
except Exception as e:
|
||||
self.__invoker.services.logger.error("Problem deleting image records and files")
|
||||
raise e
|
||||
|
||||
def get_intermediates_count(self, user_id: Optional[str] = None) -> int:
|
||||
try:
|
||||
return self.__invoker.services.image_records.get_intermediates_count(user_id=user_id)
|
||||
except Exception as e:
|
||||
self.__invoker.services.logger.error("Problem getting intermediates count")
|
||||
raise e
|
||||
|
||||
def get_image_names(
|
||||
self,
|
||||
starred_first: bool = True,
|
||||
order_dir: SQLiteDirection = SQLiteDirection.Descending,
|
||||
image_origin: Optional[ResourceOrigin] = None,
|
||||
categories: Optional[list[ImageCategory]] = None,
|
||||
is_intermediate: Optional[bool] = None,
|
||||
board_id: Optional[str] = None,
|
||||
search_term: Optional[str] = None,
|
||||
user_id: Optional[str] = None,
|
||||
is_admin: bool = False,
|
||||
) -> ImageNamesResult:
|
||||
try:
|
||||
return self.__invoker.services.image_records.get_image_names(
|
||||
starred_first=starred_first,
|
||||
order_dir=order_dir,
|
||||
image_origin=image_origin,
|
||||
categories=categories,
|
||||
is_intermediate=is_intermediate,
|
||||
board_id=board_id,
|
||||
search_term=search_term,
|
||||
user_id=user_id,
|
||||
is_admin=is_admin,
|
||||
)
|
||||
except Exception as e:
|
||||
self.__invoker.services.logger.error("Problem getting image names")
|
||||
raise e
|
||||
@@ -0,0 +1,63 @@
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import Optional, Union
|
||||
|
||||
from invokeai.app.invocations.baseinvocation import BaseInvocation, BaseInvocationOutput
|
||||
from invokeai.app.services.invocation_cache.invocation_cache_common import InvocationCacheStatus
|
||||
|
||||
|
||||
class InvocationCacheBase(ABC):
|
||||
"""
|
||||
Base class for invocation caches.
|
||||
When an invocation is executed, it is hashed and its output stored in the cache.
|
||||
When new invocations are executed, if they are flagged with `use_cache`, they
|
||||
will attempt to pull their value from the cache before executing.
|
||||
|
||||
Implementations should register for the `on_deleted` event of the `images` and `latents`
|
||||
services, and delete any cached outputs that reference the deleted image or latent.
|
||||
|
||||
See the memory implementation for an example.
|
||||
|
||||
Implementations should respect the `node_cache_size` configuration value, and skip all
|
||||
cache logic if the value is set to 0.
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
def get(self, key: Union[int, str]) -> Optional[BaseInvocationOutput]:
|
||||
"""Retrieves an invocation output from the cache"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def save(self, key: Union[int, str], invocation_output: BaseInvocationOutput) -> None:
|
||||
"""Stores an invocation output in the cache"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def delete(self, key: Union[int, str]) -> None:
|
||||
"""Deletes an invocation output from the cache"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def clear(self) -> None:
|
||||
"""Clears the cache"""
|
||||
pass
|
||||
|
||||
@staticmethod
|
||||
@abstractmethod
|
||||
def create_key(invocation: BaseInvocation) -> int:
|
||||
"""Gets the key for the invocation's cache item"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def disable(self) -> None:
|
||||
"""Disables the cache, overriding the max cache size"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def enable(self) -> None:
|
||||
"""Enables the cache, letting the the max cache size take effect"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def get_status(self) -> InvocationCacheStatus:
|
||||
"""Returns the status of the cache"""
|
||||
pass
|
||||
@@ -0,0 +1,9 @@
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
|
||||
class InvocationCacheStatus(BaseModel):
|
||||
size: int = Field(description="The current size of the invocation cache")
|
||||
hits: int = Field(description="The number of cache hits")
|
||||
misses: int = Field(description="The number of cache misses")
|
||||
enabled: bool = Field(description="Whether the invocation cache is enabled")
|
||||
max_size: int = Field(description="The maximum size of the invocation cache")
|
||||
@@ -0,0 +1,130 @@
|
||||
from collections import OrderedDict
|
||||
from dataclasses import dataclass, field
|
||||
from threading import Lock
|
||||
from typing import Optional, Union
|
||||
|
||||
from invokeai.app.invocations.baseinvocation import BaseInvocation, BaseInvocationOutput
|
||||
from invokeai.app.services.invocation_cache.invocation_cache_base import InvocationCacheBase
|
||||
from invokeai.app.services.invocation_cache.invocation_cache_common import InvocationCacheStatus
|
||||
from invokeai.app.services.invoker import Invoker
|
||||
|
||||
|
||||
@dataclass(order=True)
|
||||
class CachedItem:
|
||||
invocation_output: BaseInvocationOutput = field(compare=False)
|
||||
invocation_output_json: str = field(compare=False)
|
||||
|
||||
|
||||
class MemoryInvocationCache(InvocationCacheBase):
|
||||
_cache: OrderedDict[Union[int, str], CachedItem]
|
||||
_max_cache_size: int
|
||||
_disabled: bool
|
||||
_hits: int
|
||||
_misses: int
|
||||
_invoker: Invoker
|
||||
_lock: Lock
|
||||
|
||||
def __init__(self, max_cache_size: int = 0) -> None:
|
||||
self._cache = OrderedDict()
|
||||
self._max_cache_size = max_cache_size
|
||||
self._disabled = False
|
||||
self._hits = 0
|
||||
self._misses = 0
|
||||
self._lock = Lock()
|
||||
|
||||
def start(self, invoker: Invoker) -> None:
|
||||
self._invoker = invoker
|
||||
if self._max_cache_size == 0:
|
||||
return
|
||||
self._invoker.services.images.on_deleted(self._delete_by_match)
|
||||
self._invoker.services.tensors.on_deleted(self._delete_by_match)
|
||||
self._invoker.services.conditioning.on_deleted(self._delete_by_match)
|
||||
|
||||
def get(self, key: Union[int, str]) -> Optional[BaseInvocationOutput]:
|
||||
with self._lock:
|
||||
if self._max_cache_size == 0 or self._disabled:
|
||||
return None
|
||||
item = self._cache.get(key, None)
|
||||
if item is not None:
|
||||
self._hits += 1
|
||||
self._cache.move_to_end(key)
|
||||
return item.invocation_output
|
||||
self._misses += 1
|
||||
return None
|
||||
|
||||
def save(self, key: Union[int, str], invocation_output: BaseInvocationOutput) -> None:
|
||||
with self._lock:
|
||||
if self._max_cache_size == 0 or self._disabled or key in self._cache:
|
||||
return
|
||||
# If the cache is full, we need to remove the least used
|
||||
number_to_delete = len(self._cache) + 1 - self._max_cache_size
|
||||
self._delete_oldest_access(number_to_delete)
|
||||
self._cache[key] = CachedItem(
|
||||
invocation_output,
|
||||
invocation_output.model_dump_json(warnings=False, exclude_defaults=True, exclude_unset=True),
|
||||
)
|
||||
|
||||
def _delete_oldest_access(self, number_to_delete: int) -> None:
|
||||
number_to_delete = min(number_to_delete, len(self._cache))
|
||||
for _ in range(number_to_delete):
|
||||
self._cache.popitem(last=False)
|
||||
|
||||
def _delete(self, key: Union[int, str]) -> None:
|
||||
if self._max_cache_size == 0:
|
||||
return
|
||||
if key in self._cache:
|
||||
del self._cache[key]
|
||||
|
||||
def delete(self, key: Union[int, str]) -> None:
|
||||
with self._lock:
|
||||
return self._delete(key)
|
||||
|
||||
def clear(self) -> None:
|
||||
with self._lock:
|
||||
if self._max_cache_size == 0:
|
||||
return
|
||||
self._cache.clear()
|
||||
self._misses = 0
|
||||
self._hits = 0
|
||||
|
||||
@staticmethod
|
||||
def create_key(invocation: BaseInvocation) -> int:
|
||||
return hash(invocation.model_dump_json(exclude={"id"}, warnings=False))
|
||||
|
||||
def disable(self) -> None:
|
||||
with self._lock:
|
||||
if self._max_cache_size == 0:
|
||||
return
|
||||
self._disabled = True
|
||||
|
||||
def enable(self) -> None:
|
||||
with self._lock:
|
||||
if self._max_cache_size == 0:
|
||||
return
|
||||
self._disabled = False
|
||||
|
||||
def get_status(self) -> InvocationCacheStatus:
|
||||
with self._lock:
|
||||
return InvocationCacheStatus(
|
||||
hits=self._hits,
|
||||
misses=self._misses,
|
||||
enabled=not self._disabled and self._max_cache_size > 0,
|
||||
size=len(self._cache),
|
||||
max_size=self._max_cache_size,
|
||||
)
|
||||
|
||||
def _delete_by_match(self, to_match: str) -> None:
|
||||
with self._lock:
|
||||
if self._max_cache_size == 0:
|
||||
return
|
||||
keys_to_delete = set()
|
||||
for key, cached_item in self._cache.items():
|
||||
if to_match in cached_item.invocation_output_json:
|
||||
keys_to_delete.add(key)
|
||||
if not keys_to_delete:
|
||||
return
|
||||
for key in keys_to_delete:
|
||||
self._delete(key)
|
||||
self._invoker.services.logger.debug(
|
||||
f"Deleted {len(keys_to_delete)} cached invocation outputs for {to_match}"
|
||||
)
|
||||
@@ -0,0 +1,116 @@
|
||||
# Copyright (c) 2022 Kyle Schouviller (https://github.com/kyle0654) and the InvokeAI Team
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from invokeai.app.services.object_serializer.object_serializer_base import ObjectSerializerBase
|
||||
from invokeai.app.services.style_preset_images.style_preset_images_base import StylePresetImageFileStorageBase
|
||||
from invokeai.app.services.style_preset_records.style_preset_records_base import StylePresetRecordsStorageBase
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from logging import Logger
|
||||
|
||||
import torch
|
||||
|
||||
from invokeai.app.services.board_image_records.board_image_records_base import BoardImageRecordStorageBase
|
||||
from invokeai.app.services.board_images.board_images_base import BoardImagesServiceABC
|
||||
from invokeai.app.services.board_records.board_records_base import BoardRecordStorageBase
|
||||
from invokeai.app.services.boards.boards_base import BoardServiceABC
|
||||
from invokeai.app.services.bulk_download.bulk_download_base import BulkDownloadBase
|
||||
from invokeai.app.services.client_state_persistence.client_state_persistence_base import ClientStatePersistenceABC
|
||||
from invokeai.app.services.config import InvokeAIAppConfig
|
||||
from invokeai.app.services.download import DownloadQueueServiceBase
|
||||
from invokeai.app.services.events.events_base import EventServiceBase
|
||||
from invokeai.app.services.external_generation.external_generation_base import ExternalGenerationServiceBase
|
||||
from invokeai.app.services.image_files.image_files_base import ImageFileStorageBase
|
||||
from invokeai.app.services.image_moves.image_moves_default import ImageMoveService
|
||||
from invokeai.app.services.image_records.image_records_base import ImageRecordStorageBase
|
||||
from invokeai.app.services.images.images_base import ImageServiceABC
|
||||
from invokeai.app.services.invocation_cache.invocation_cache_base import InvocationCacheBase
|
||||
from invokeai.app.services.invocation_stats.invocation_stats_base import InvocationStatsServiceBase
|
||||
from invokeai.app.services.model_images.model_images_base import ModelImageFileStorageBase
|
||||
from invokeai.app.services.model_manager.model_manager_base import ModelManagerServiceBase
|
||||
from invokeai.app.services.model_relationship_records.model_relationship_records_base import (
|
||||
ModelRelationshipRecordStorageBase,
|
||||
)
|
||||
from invokeai.app.services.model_relationships.model_relationships_base import ModelRelationshipsServiceABC
|
||||
from invokeai.app.services.names.names_base import NameServiceBase
|
||||
from invokeai.app.services.session_processor.session_processor_base import SessionProcessorBase
|
||||
from invokeai.app.services.session_queue.session_queue_base import SessionQueueBase
|
||||
from invokeai.app.services.urls.urls_base import UrlServiceBase
|
||||
from invokeai.app.services.users.users_base import UserServiceBase
|
||||
from invokeai.app.services.workflow_records.workflow_records_base import WorkflowRecordsStorageBase
|
||||
from invokeai.app.services.workflow_thumbnails.workflow_thumbnails_base import WorkflowThumbnailServiceBase
|
||||
from invokeai.backend.stable_diffusion.diffusion.conditioning_data import ConditioningFieldData
|
||||
|
||||
|
||||
class InvocationServices:
|
||||
"""Services that can be used by invocations"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
board_images: "BoardImagesServiceABC",
|
||||
board_image_records: "BoardImageRecordStorageBase",
|
||||
boards: "BoardServiceABC",
|
||||
board_records: "BoardRecordStorageBase",
|
||||
bulk_download: "BulkDownloadBase",
|
||||
configuration: "InvokeAIAppConfig",
|
||||
events: "EventServiceBase",
|
||||
images: "ImageServiceABC",
|
||||
image_files: "ImageFileStorageBase",
|
||||
image_records: "ImageRecordStorageBase",
|
||||
logger: "Logger",
|
||||
model_images: "ModelImageFileStorageBase",
|
||||
model_manager: "ModelManagerServiceBase",
|
||||
model_relationships: "ModelRelationshipsServiceABC",
|
||||
model_relationship_records: "ModelRelationshipRecordStorageBase",
|
||||
download_queue: "DownloadQueueServiceBase",
|
||||
external_generation: "ExternalGenerationServiceBase",
|
||||
performance_statistics: "InvocationStatsServiceBase",
|
||||
session_queue: "SessionQueueBase",
|
||||
session_processor: "SessionProcessorBase",
|
||||
invocation_cache: "InvocationCacheBase",
|
||||
names: "NameServiceBase",
|
||||
urls: "UrlServiceBase",
|
||||
workflow_records: "WorkflowRecordsStorageBase",
|
||||
tensors: "ObjectSerializerBase[torch.Tensor]",
|
||||
conditioning: "ObjectSerializerBase[ConditioningFieldData]",
|
||||
style_preset_records: "StylePresetRecordsStorageBase",
|
||||
style_preset_image_files: "StylePresetImageFileStorageBase",
|
||||
workflow_thumbnails: "WorkflowThumbnailServiceBase",
|
||||
client_state_persistence: "ClientStatePersistenceABC",
|
||||
users: "UserServiceBase",
|
||||
image_moves: "ImageMoveService | None" = None,
|
||||
):
|
||||
self.board_images = board_images
|
||||
self.board_image_records = board_image_records
|
||||
self.boards = boards
|
||||
self.board_records = board_records
|
||||
self.bulk_download = bulk_download
|
||||
self.configuration = configuration
|
||||
self.events = events
|
||||
self.images = images
|
||||
self.image_files = image_files
|
||||
self.image_records = image_records
|
||||
self.logger = logger
|
||||
self.model_images = model_images
|
||||
self.model_manager = model_manager
|
||||
self.model_relationships = model_relationships
|
||||
self.model_relationship_records = model_relationship_records
|
||||
self.download_queue = download_queue
|
||||
self.external_generation = external_generation
|
||||
self.performance_statistics = performance_statistics
|
||||
self.session_queue = session_queue
|
||||
self.image_moves = image_moves
|
||||
self.session_processor = session_processor
|
||||
self.invocation_cache = invocation_cache
|
||||
self.names = names
|
||||
self.urls = urls
|
||||
self.workflow_records = workflow_records
|
||||
self.tensors = tensors
|
||||
self.conditioning = conditioning
|
||||
self.style_preset_records = style_preset_records
|
||||
self.style_preset_image_files = style_preset_image_files
|
||||
self.workflow_thumbnails = workflow_thumbnails
|
||||
self.client_state_persistence = client_state_persistence
|
||||
self.users = users
|
||||
@@ -0,0 +1,93 @@
|
||||
# Copyright 2023 Lincoln D. Stein <lincoln.stein@gmail.com>
|
||||
"""Utility to collect execution time and GPU usage stats on invocations in flight
|
||||
|
||||
Usage:
|
||||
|
||||
statistics = InvocationStatsService()
|
||||
with statistics.collect_stats(invocation, graph_execution_state.id):
|
||||
... execute graphs...
|
||||
statistics.log_stats()
|
||||
|
||||
Typical output:
|
||||
[2023-08-02 18:03:04,507]::[InvokeAI]::INFO --> Graph stats: c7764585-9c68-4d9d-a199-55e8186790f3
|
||||
[2023-08-02 18:03:04,507]::[InvokeAI]::INFO --> Node Calls Seconds VRAM Used
|
||||
[2023-08-02 18:03:04,507]::[InvokeAI]::INFO --> main_model_loader 1 0.005s 0.01G
|
||||
[2023-08-02 18:03:04,508]::[InvokeAI]::INFO --> clip_skip 1 0.004s 0.01G
|
||||
[2023-08-02 18:03:04,508]::[InvokeAI]::INFO --> compel 2 0.512s 0.26G
|
||||
[2023-08-02 18:03:04,508]::[InvokeAI]::INFO --> rand_int 1 0.001s 0.01G
|
||||
[2023-08-02 18:03:04,508]::[InvokeAI]::INFO --> range_of_size 1 0.001s 0.01G
|
||||
[2023-08-02 18:03:04,508]::[InvokeAI]::INFO --> iterate 1 0.001s 0.01G
|
||||
[2023-08-02 18:03:04,508]::[InvokeAI]::INFO --> metadata_accumulator 1 0.002s 0.01G
|
||||
[2023-08-02 18:03:04,508]::[InvokeAI]::INFO --> noise 1 0.002s 0.01G
|
||||
[2023-08-02 18:03:04,508]::[InvokeAI]::INFO --> t2l 1 3.541s 1.93G
|
||||
[2023-08-02 18:03:04,508]::[InvokeAI]::INFO --> l2i 1 0.679s 0.58G
|
||||
[2023-08-02 18:03:04,508]::[InvokeAI]::INFO --> TOTAL GRAPH EXECUTION TIME: 4.749s
|
||||
[2023-08-02 18:03:04,508]::[InvokeAI]::INFO --> Current VRAM utilization 0.01G
|
||||
|
||||
The abstract base class for this class is InvocationStatsServiceBase. An implementing class which
|
||||
writes to the system log is stored in InvocationServices.performance_statistics.
|
||||
"""
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
from pathlib import Path
|
||||
from typing import ContextManager
|
||||
|
||||
from invokeai.app.invocations.baseinvocation import BaseInvocation
|
||||
from invokeai.app.services.invocation_stats.invocation_stats_common import InvocationStatsSummary
|
||||
|
||||
|
||||
class InvocationStatsServiceBase(ABC):
|
||||
"Abstract base class for recording node memory/time performance statistics"
|
||||
|
||||
@abstractmethod
|
||||
def __init__(self) -> None:
|
||||
"""
|
||||
Initialize the InvocationStatsService and reset counters to zero
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
def collect_stats(
|
||||
self,
|
||||
invocation: BaseInvocation,
|
||||
graph_execution_state_id: str,
|
||||
) -> ContextManager[None]:
|
||||
"""
|
||||
Return a context object that will capture the statistics on the execution
|
||||
of invocaation. Use with: to place around the part of the code that executes the invocation.
|
||||
:param invocation: BaseInvocation object from the current graph.
|
||||
:param graph_execution_state_id: The id of the current session.
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def reset_stats(self, graph_execution_state_id: str) -> None:
|
||||
"""Reset all stored statistics."""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def log_stats(self, graph_execution_state_id: str) -> None:
|
||||
"""
|
||||
Write out the accumulated statistics to the log or somewhere else.
|
||||
:param graph_execution_state_id: The id of the session whose stats to log.
|
||||
:raises GESStatsNotFoundError: if the graph isn't tracked in the stats.
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def get_stats(self, graph_execution_state_id: str) -> InvocationStatsSummary:
|
||||
"""
|
||||
Gets the accumulated statistics for the indicated graph.
|
||||
:param graph_execution_state_id: The id of the session whose stats to get.
|
||||
:raises GESStatsNotFoundError: if the graph isn't tracked in the stats.
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def dump_stats(self, graph_execution_state_id: str, output_path: Path) -> None:
|
||||
"""
|
||||
Write out the accumulated statistics to the indicated path as JSON.
|
||||
:param graph_execution_state_id: The id of the session whose stats to dump.
|
||||
:param output_path: The file to write the stats to.
|
||||
:raises GESStatsNotFoundError: if the graph isn't tracked in the stats.
|
||||
"""
|
||||
pass
|
||||
@@ -0,0 +1,183 @@
|
||||
from collections import defaultdict
|
||||
from dataclasses import asdict, dataclass
|
||||
from typing import Any, Optional
|
||||
|
||||
|
||||
class GESStatsNotFoundError(Exception):
|
||||
"""Raised when execution stats are not found for a given Graph Execution State."""
|
||||
|
||||
|
||||
@dataclass
|
||||
class NodeExecutionStatsSummary:
|
||||
"""The stats for a specific type of node."""
|
||||
|
||||
node_type: str
|
||||
num_calls: int
|
||||
time_used_seconds: float
|
||||
delta_vram_gb: float
|
||||
|
||||
|
||||
@dataclass
|
||||
class ModelCacheStatsSummary:
|
||||
"""The stats for the model cache."""
|
||||
|
||||
high_water_mark_gb: float
|
||||
cache_size_gb: float
|
||||
total_usage_gb: float
|
||||
cache_hits: int
|
||||
cache_misses: int
|
||||
models_cached: int
|
||||
models_cleared: int
|
||||
|
||||
|
||||
@dataclass
|
||||
class GraphExecutionStatsSummary:
|
||||
"""The stats for the graph execution state."""
|
||||
|
||||
graph_execution_state_id: str
|
||||
execution_time_seconds: float
|
||||
# `wall_time_seconds`, `ram_usage_gb` and `ram_change_gb` are derived from the node execution stats.
|
||||
# In some situations, there are no node stats, so these values are optional.
|
||||
wall_time_seconds: Optional[float]
|
||||
ram_usage_gb: Optional[float]
|
||||
ram_change_gb: Optional[float]
|
||||
|
||||
|
||||
@dataclass
|
||||
class InvocationStatsSummary:
|
||||
"""
|
||||
The accumulated stats for a graph execution.
|
||||
Its `__str__` method returns a human-readable stats summary.
|
||||
"""
|
||||
|
||||
vram_usage_gb: Optional[float]
|
||||
graph_stats: GraphExecutionStatsSummary
|
||||
model_cache_stats: ModelCacheStatsSummary
|
||||
node_stats: list[NodeExecutionStatsSummary]
|
||||
|
||||
def __str__(self) -> str:
|
||||
_str = ""
|
||||
_str = f"Graph stats: {self.graph_stats.graph_execution_state_id}\n"
|
||||
_str += f"{'Node':>30} {'Calls':>7} {'Seconds':>9} {'VRAM Change':+>10}\n"
|
||||
|
||||
for summary in self.node_stats:
|
||||
_str += f"{summary.node_type:>30} {summary.num_calls:>7} {summary.time_used_seconds:>8.3f}s {summary.delta_vram_gb:+10.3f}G\n"
|
||||
|
||||
_str += f"TOTAL GRAPH EXECUTION TIME: {self.graph_stats.execution_time_seconds:7.3f}s\n"
|
||||
|
||||
if self.graph_stats.wall_time_seconds is not None:
|
||||
_str += f"TOTAL GRAPH WALL TIME: {self.graph_stats.wall_time_seconds:7.3f}s\n"
|
||||
|
||||
if self.graph_stats.ram_usage_gb is not None and self.graph_stats.ram_change_gb is not None:
|
||||
_str += f"RAM used by InvokeAI process: {self.graph_stats.ram_usage_gb:4.2f}G ({self.graph_stats.ram_change_gb:+5.3f}G)\n"
|
||||
|
||||
_str += f"RAM used to load models: {self.model_cache_stats.total_usage_gb:4.2f}G\n"
|
||||
if self.vram_usage_gb:
|
||||
_str += f"VRAM in use: {self.vram_usage_gb:4.3f}G\n"
|
||||
_str += "RAM cache statistics:\n"
|
||||
_str += f" Model cache hits: {self.model_cache_stats.cache_hits}\n"
|
||||
_str += f" Model cache misses: {self.model_cache_stats.cache_misses}\n"
|
||||
_str += f" Models cached: {self.model_cache_stats.models_cached}\n"
|
||||
_str += f" Models cleared from cache: {self.model_cache_stats.models_cleared}\n"
|
||||
_str += f" Cache high water mark: {self.model_cache_stats.high_water_mark_gb:4.2f}/{self.model_cache_stats.cache_size_gb:4.2f}G\n"
|
||||
|
||||
return _str
|
||||
|
||||
def as_dict(self) -> dict[str, Any]:
|
||||
"""Returns the stats as a dictionary."""
|
||||
return asdict(self)
|
||||
|
||||
|
||||
@dataclass
|
||||
class NodeExecutionStats:
|
||||
"""Class for tracking execution stats of an invocation node."""
|
||||
|
||||
invocation_type: str
|
||||
|
||||
start_time: float # Seconds since the epoch.
|
||||
end_time: float # Seconds since the epoch.
|
||||
|
||||
start_ram_gb: float # GB
|
||||
end_ram_gb: float # GB
|
||||
|
||||
delta_vram_gb: float # GB
|
||||
|
||||
def total_time(self) -> float:
|
||||
return self.end_time - self.start_time
|
||||
|
||||
|
||||
class GraphExecutionStats:
|
||||
"""Class for tracking execution stats of a graph."""
|
||||
|
||||
def __init__(self):
|
||||
self._node_stats_list: list[NodeExecutionStats] = []
|
||||
|
||||
def add_node_execution_stats(self, node_stats: NodeExecutionStats):
|
||||
self._node_stats_list.append(node_stats)
|
||||
|
||||
def get_total_run_time(self) -> float:
|
||||
"""Get the total time spent executing nodes in the graph."""
|
||||
total = 0.0
|
||||
for node_stats in self._node_stats_list:
|
||||
total += node_stats.total_time()
|
||||
return total
|
||||
|
||||
def get_first_node_stats(self) -> NodeExecutionStats | None:
|
||||
"""Get the stats of the first node in the graph (by start_time)."""
|
||||
first_node = None
|
||||
for node_stats in self._node_stats_list:
|
||||
if first_node is None or node_stats.start_time < first_node.start_time:
|
||||
first_node = node_stats
|
||||
|
||||
assert first_node is not None
|
||||
return first_node
|
||||
|
||||
def get_last_node_stats(self) -> NodeExecutionStats | None:
|
||||
"""Get the stats of the last node in the graph (by end_time)."""
|
||||
last_node = None
|
||||
for node_stats in self._node_stats_list:
|
||||
if last_node is None or node_stats.end_time > last_node.end_time:
|
||||
last_node = node_stats
|
||||
|
||||
return last_node
|
||||
|
||||
def get_graph_stats_summary(self, graph_execution_state_id: str) -> GraphExecutionStatsSummary:
|
||||
"""Get a summary of the graph stats."""
|
||||
first_node = self.get_first_node_stats()
|
||||
last_node = self.get_last_node_stats()
|
||||
|
||||
wall_time_seconds: Optional[float] = None
|
||||
ram_usage_gb: Optional[float] = None
|
||||
ram_change_gb: Optional[float] = None
|
||||
|
||||
if last_node and first_node:
|
||||
wall_time_seconds = last_node.end_time - first_node.start_time
|
||||
ram_usage_gb = last_node.end_ram_gb
|
||||
ram_change_gb = last_node.end_ram_gb - first_node.start_ram_gb
|
||||
|
||||
return GraphExecutionStatsSummary(
|
||||
graph_execution_state_id=graph_execution_state_id,
|
||||
execution_time_seconds=self.get_total_run_time(),
|
||||
wall_time_seconds=wall_time_seconds,
|
||||
ram_usage_gb=ram_usage_gb,
|
||||
ram_change_gb=ram_change_gb,
|
||||
)
|
||||
|
||||
def get_node_stats_summaries(self) -> list[NodeExecutionStatsSummary]:
|
||||
"""Get a summary of the node stats."""
|
||||
summaries: list[NodeExecutionStatsSummary] = []
|
||||
node_stats_by_type: dict[str, list[NodeExecutionStats]] = defaultdict(list)
|
||||
|
||||
for node_stats in self._node_stats_list:
|
||||
node_stats_by_type[node_stats.invocation_type].append(node_stats)
|
||||
|
||||
for node_type, node_type_stats_list in node_stats_by_type.items():
|
||||
num_calls = len(node_type_stats_list)
|
||||
time_used = sum([n.total_time() for n in node_type_stats_list])
|
||||
delta_vram = max([n.delta_vram_gb for n in node_type_stats_list])
|
||||
summary = NodeExecutionStatsSummary(
|
||||
node_type=node_type, num_calls=num_calls, time_used_seconds=time_used, delta_vram_gb=delta_vram
|
||||
)
|
||||
summaries.append(summary)
|
||||
|
||||
return summaries
|
||||
@@ -0,0 +1,143 @@
|
||||
import json
|
||||
import time
|
||||
from contextlib import contextmanager
|
||||
from pathlib import Path
|
||||
from typing import Generator
|
||||
|
||||
import psutil
|
||||
import torch
|
||||
|
||||
import invokeai.backend.util.logging as logger
|
||||
from invokeai.app.invocations.baseinvocation import BaseInvocation
|
||||
from invokeai.app.services.invocation_stats.invocation_stats_base import InvocationStatsServiceBase
|
||||
from invokeai.app.services.invocation_stats.invocation_stats_common import (
|
||||
GESStatsNotFoundError,
|
||||
GraphExecutionStats,
|
||||
GraphExecutionStatsSummary,
|
||||
InvocationStatsSummary,
|
||||
ModelCacheStatsSummary,
|
||||
NodeExecutionStats,
|
||||
NodeExecutionStatsSummary,
|
||||
)
|
||||
from invokeai.app.services.invoker import Invoker
|
||||
from invokeai.backend.model_manager.load.model_cache.cache_stats import CacheStats
|
||||
|
||||
# Size of 1GB in bytes.
|
||||
GB = 2**30
|
||||
|
||||
|
||||
class InvocationStatsService(InvocationStatsServiceBase):
|
||||
"""Accumulate performance information about a running graph. Collects time spent in each node,
|
||||
as well as the maximum and current VRAM utilisation for CUDA systems"""
|
||||
|
||||
def __init__(self):
|
||||
# Maps graph_execution_state_id to GraphExecutionStats.
|
||||
self._stats: dict[str, GraphExecutionStats] = {}
|
||||
# Maps graph_execution_state_id to model manager CacheStats.
|
||||
self._cache_stats: dict[str, CacheStats] = {}
|
||||
|
||||
def start(self, invoker: Invoker) -> None:
|
||||
self._invoker = invoker
|
||||
|
||||
@contextmanager
|
||||
def collect_stats(self, invocation: BaseInvocation, graph_execution_state_id: str) -> Generator[None, None, None]:
|
||||
# This is to handle case of the model manager not being initialized, which happens
|
||||
# during some tests.
|
||||
services = self._invoker.services
|
||||
if not self._stats.get(graph_execution_state_id):
|
||||
# First time we're seeing this graph_execution_state_id.
|
||||
self._stats[graph_execution_state_id] = GraphExecutionStats()
|
||||
self._cache_stats[graph_execution_state_id] = CacheStats()
|
||||
|
||||
# Record state before the invocation.
|
||||
start_time = time.time()
|
||||
start_ram = psutil.Process().memory_info().rss
|
||||
|
||||
# Remember current VRAM usage
|
||||
vram_in_use = torch.cuda.memory_allocated() if torch.cuda.is_available() else 0.0
|
||||
|
||||
assert services.model_manager.load is not None
|
||||
services.model_manager.load.ram_cache.stats = self._cache_stats[graph_execution_state_id]
|
||||
|
||||
try:
|
||||
# Let the invocation run.
|
||||
yield None
|
||||
finally:
|
||||
# Record delta VRAM
|
||||
delta_vram_gb = ((torch.cuda.memory_allocated() - vram_in_use) / GB) if torch.cuda.is_available() else 0.0
|
||||
|
||||
node_stats = NodeExecutionStats(
|
||||
invocation_type=invocation.get_type(),
|
||||
start_time=start_time,
|
||||
end_time=time.time(),
|
||||
start_ram_gb=start_ram / GB,
|
||||
end_ram_gb=psutil.Process().memory_info().rss / GB,
|
||||
delta_vram_gb=delta_vram_gb,
|
||||
)
|
||||
self._stats[graph_execution_state_id].add_node_execution_stats(node_stats)
|
||||
|
||||
def reset_stats(self, graph_execution_state_id: str) -> None:
|
||||
self._stats.pop(graph_execution_state_id, None)
|
||||
self._cache_stats.pop(graph_execution_state_id, None)
|
||||
|
||||
def get_stats(self, graph_execution_state_id: str) -> InvocationStatsSummary:
|
||||
graph_stats_summary = self._get_graph_summary(graph_execution_state_id)
|
||||
node_stats_summaries = self._get_node_summaries(graph_execution_state_id)
|
||||
model_cache_stats_summary = self._get_model_cache_summary(graph_execution_state_id)
|
||||
# Note: We use memory_allocated() here (not memory_reserved()) because we want to show
|
||||
# the current actively-used VRAM, not the total reserved memory including PyTorch's cache.
|
||||
vram_usage_gb = torch.cuda.memory_allocated() / GB if torch.cuda.is_available() else None
|
||||
|
||||
return InvocationStatsSummary(
|
||||
graph_stats=graph_stats_summary,
|
||||
model_cache_stats=model_cache_stats_summary,
|
||||
node_stats=node_stats_summaries,
|
||||
vram_usage_gb=vram_usage_gb,
|
||||
)
|
||||
|
||||
def log_stats(self, graph_execution_state_id: str) -> None:
|
||||
stats = self.get_stats(graph_execution_state_id)
|
||||
logger.info(str(stats))
|
||||
|
||||
def dump_stats(self, graph_execution_state_id: str, output_path: Path) -> None:
|
||||
stats = self.get_stats(graph_execution_state_id)
|
||||
with open(output_path, "w") as f:
|
||||
f.write(json.dumps(stats.as_dict(), indent=2))
|
||||
|
||||
def _get_model_cache_summary(self, graph_execution_state_id: str) -> ModelCacheStatsSummary:
|
||||
try:
|
||||
cache_stats = self._cache_stats[graph_execution_state_id]
|
||||
except KeyError as e:
|
||||
raise GESStatsNotFoundError(
|
||||
f"Attempted to get model cache statistics for unknown graph {graph_execution_state_id}: {e}."
|
||||
) from e
|
||||
|
||||
return ModelCacheStatsSummary(
|
||||
cache_hits=cache_stats.hits,
|
||||
cache_misses=cache_stats.misses,
|
||||
high_water_mark_gb=cache_stats.high_watermark / GB,
|
||||
cache_size_gb=cache_stats.cache_size / GB,
|
||||
total_usage_gb=sum(list(cache_stats.loaded_model_sizes.values())) / GB,
|
||||
models_cached=cache_stats.in_cache,
|
||||
models_cleared=cache_stats.cleared,
|
||||
)
|
||||
|
||||
def _get_graph_summary(self, graph_execution_state_id: str) -> GraphExecutionStatsSummary:
|
||||
try:
|
||||
graph_stats = self._stats[graph_execution_state_id]
|
||||
except KeyError as e:
|
||||
raise GESStatsNotFoundError(
|
||||
f"Attempted to get graph statistics for unknown graph {graph_execution_state_id}: {e}."
|
||||
) from e
|
||||
|
||||
return graph_stats.get_graph_stats_summary(graph_execution_state_id)
|
||||
|
||||
def _get_node_summaries(self, graph_execution_state_id: str) -> list[NodeExecutionStatsSummary]:
|
||||
try:
|
||||
graph_stats = self._stats[graph_execution_state_id]
|
||||
except KeyError as e:
|
||||
raise GESStatsNotFoundError(
|
||||
f"Attempted to get node statistics for unknown graph {graph_execution_state_id}: {e}."
|
||||
) from e
|
||||
|
||||
return graph_stats.get_node_stats_summaries()
|
||||
@@ -0,0 +1,37 @@
|
||||
# Copyright (c) 2022 Kyle Schouviller (https://github.com/kyle0654)
|
||||
|
||||
|
||||
from invokeai.app.services.invocation_services import InvocationServices
|
||||
|
||||
|
||||
class Invoker:
|
||||
"""The invoker, used to execute invocations"""
|
||||
|
||||
services: InvocationServices
|
||||
|
||||
def __init__(self, services: InvocationServices):
|
||||
self.services = services
|
||||
self._start()
|
||||
|
||||
def __start_service(self, service) -> None:
|
||||
# Call start() method on any services that have it
|
||||
start_op = getattr(service, "start", None)
|
||||
if callable(start_op):
|
||||
start_op(self)
|
||||
|
||||
def __stop_service(self, service) -> None:
|
||||
# Call stop() method on any services that have it
|
||||
stop_op = getattr(service, "stop", None)
|
||||
if callable(stop_op):
|
||||
stop_op(self)
|
||||
|
||||
def _start(self) -> None:
|
||||
"""Starts the invoker. This is called automatically when the invoker is created."""
|
||||
for service in vars(self.services):
|
||||
self.__start_service(getattr(self.services, service))
|
||||
|
||||
def stop(self) -> None:
|
||||
"""Stops the invoker. A new invoker will have to be created to execute further."""
|
||||
# First stop all services
|
||||
for service in vars(self.services):
|
||||
self.__stop_service(getattr(self.services, service))
|
||||
@@ -0,0 +1,59 @@
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import Callable, Generic, TypeVar
|
||||
|
||||
from pydantic import BaseModel
|
||||
|
||||
T = TypeVar("T", bound=BaseModel)
|
||||
|
||||
|
||||
class ItemStorageABC(ABC, Generic[T]):
|
||||
"""Provides storage for a single type of item. The type must be a Pydantic model."""
|
||||
|
||||
_on_changed_callbacks: list[Callable[[T], None]]
|
||||
_on_deleted_callbacks: list[Callable[[str], None]]
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._on_changed_callbacks = []
|
||||
self._on_deleted_callbacks = []
|
||||
|
||||
"""Base item storage class"""
|
||||
|
||||
@abstractmethod
|
||||
def get(self, item_id: str) -> T:
|
||||
"""
|
||||
Gets the item.
|
||||
:param item_id: the id of the item to get
|
||||
:raises ItemNotFoundError: if the item is not found
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def set(self, item: T) -> None:
|
||||
"""
|
||||
Sets the item.
|
||||
:param item: the item to set
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def delete(self, item_id: str) -> None:
|
||||
"""
|
||||
Deletes the item, if it exists.
|
||||
"""
|
||||
pass
|
||||
|
||||
def on_changed(self, on_changed: Callable[[T], None]) -> None:
|
||||
"""Register a callback for when an item is changed"""
|
||||
self._on_changed_callbacks.append(on_changed)
|
||||
|
||||
def on_deleted(self, on_deleted: Callable[[str], None]) -> None:
|
||||
"""Register a callback for when an item is deleted"""
|
||||
self._on_deleted_callbacks.append(on_deleted)
|
||||
|
||||
def _on_changed(self, item: T) -> None:
|
||||
for callback in self._on_changed_callbacks:
|
||||
callback(item)
|
||||
|
||||
def _on_deleted(self, item_id: str) -> None:
|
||||
for callback in self._on_deleted_callbacks:
|
||||
callback(item_id)
|
||||
@@ -0,0 +1,5 @@
|
||||
class ItemNotFoundError(KeyError):
|
||||
"""Raised when an item is not found in storage"""
|
||||
|
||||
def __init__(self, item_id: str) -> None:
|
||||
super().__init__(f"Item with id {item_id} not found")
|
||||
@@ -0,0 +1,52 @@
|
||||
from collections import OrderedDict
|
||||
from contextlib import suppress
|
||||
from typing import Generic, TypeVar
|
||||
|
||||
from pydantic import BaseModel
|
||||
|
||||
from invokeai.app.services.item_storage.item_storage_base import ItemStorageABC
|
||||
from invokeai.app.services.item_storage.item_storage_common import ItemNotFoundError
|
||||
|
||||
T = TypeVar("T", bound=BaseModel)
|
||||
|
||||
|
||||
class ItemStorageMemory(ItemStorageABC[T], Generic[T]):
|
||||
"""
|
||||
Provides a simple in-memory storage for items, with a maximum number of items to store.
|
||||
The storage uses the LRU strategy to evict items from storage when the max has been reached.
|
||||
"""
|
||||
|
||||
def __init__(self, id_field: str = "id", max_items: int = 10) -> None:
|
||||
super().__init__()
|
||||
if max_items < 1:
|
||||
raise ValueError("max_items must be at least 1")
|
||||
if not id_field:
|
||||
raise ValueError("id_field must not be empty")
|
||||
self._id_field = id_field
|
||||
self._items: OrderedDict[str, T] = OrderedDict()
|
||||
self._max_items = max_items
|
||||
|
||||
def get(self, item_id: str) -> T:
|
||||
# If the item exists, move it to the end of the OrderedDict.
|
||||
item = self._items.pop(item_id, None)
|
||||
if item is None:
|
||||
raise ItemNotFoundError(item_id)
|
||||
self._items[item_id] = item
|
||||
return item
|
||||
|
||||
def set(self, item: T) -> None:
|
||||
item_id = getattr(item, self._id_field)
|
||||
if item_id in self._items:
|
||||
# If item already exists, remove it and add it to the end
|
||||
self._items.pop(item_id)
|
||||
elif len(self._items) >= self._max_items:
|
||||
# If cache is full, evict the least recently used item
|
||||
self._items.popitem(last=False)
|
||||
self._items[item_id] = item
|
||||
self._on_changed(item)
|
||||
|
||||
def delete(self, item_id: str) -> None:
|
||||
# This is a no-op if the item doesn't exist.
|
||||
with suppress(KeyError):
|
||||
del self._items[item_id]
|
||||
self._on_deleted(item_id)
|
||||
@@ -0,0 +1,33 @@
|
||||
from abc import ABC, abstractmethod
|
||||
from pathlib import Path
|
||||
|
||||
from PIL.Image import Image as PILImageType
|
||||
|
||||
|
||||
class ModelImageFileStorageBase(ABC):
|
||||
"""Low-level service responsible for storing and retrieving image files."""
|
||||
|
||||
@abstractmethod
|
||||
def get(self, model_key: str) -> PILImageType:
|
||||
"""Retrieves a model image as PIL Image."""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def get_path(self, model_key: str) -> Path:
|
||||
"""Gets the internal path to a model image."""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def get_url(self, model_key: str) -> str | None:
|
||||
"""Gets the URL to fetch a model image."""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def save(self, image: PILImageType, model_key: str) -> None:
|
||||
"""Saves a model image."""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def delete(self, model_key: str) -> None:
|
||||
"""Deletes a model image."""
|
||||
pass
|
||||
@@ -0,0 +1,20 @@
|
||||
# TODO: Should these exceptions subclass existing python exceptions?
|
||||
class ModelImageFileNotFoundException(Exception):
|
||||
"""Raised when an image file is not found in storage."""
|
||||
|
||||
def __init__(self, message="Model image file not found"):
|
||||
super().__init__(message)
|
||||
|
||||
|
||||
class ModelImageFileSaveException(Exception):
|
||||
"""Raised when an image cannot be saved."""
|
||||
|
||||
def __init__(self, message="Model image file not saved"):
|
||||
super().__init__(message)
|
||||
|
||||
|
||||
class ModelImageFileDeleteException(Exception):
|
||||
"""Raised when an image cannot be deleted."""
|
||||
|
||||
def __init__(self, message="Model image file not deleted"):
|
||||
super().__init__(message)
|
||||
@@ -0,0 +1,83 @@
|
||||
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)
|
||||
@@ -0,0 +1,25 @@
|
||||
"""Initialization file for model install service package."""
|
||||
|
||||
from invokeai.app.services.model_install.model_install_base import ModelInstallServiceBase
|
||||
from invokeai.app.services.model_install.model_install_common import (
|
||||
HFModelSource,
|
||||
InstallStatus,
|
||||
LocalModelSource,
|
||||
ModelInstallJob,
|
||||
ModelSource,
|
||||
UnknownInstallJobException,
|
||||
URLModelSource,
|
||||
)
|
||||
from invokeai.app.services.model_install.model_install_default import ModelInstallService
|
||||
|
||||
__all__ = [
|
||||
"ModelInstallServiceBase",
|
||||
"ModelInstallService",
|
||||
"InstallStatus",
|
||||
"ModelInstallJob",
|
||||
"UnknownInstallJobException",
|
||||
"ModelSource",
|
||||
"LocalModelSource",
|
||||
"HFModelSource",
|
||||
"URLModelSource",
|
||||
]
|
||||
@@ -0,0 +1,265 @@
|
||||
# Copyright 2023 Lincoln D. Stein and the InvokeAI development team
|
||||
"""Baseclass definitions for the model installer."""
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING, List, Optional, Union
|
||||
|
||||
from pydantic.networks import AnyHttpUrl
|
||||
|
||||
from invokeai.app.services.config import InvokeAIAppConfig
|
||||
from invokeai.app.services.download import DownloadQueueServiceBase
|
||||
from invokeai.app.services.invoker import Invoker
|
||||
from invokeai.app.services.model_install.model_install_common import ModelInstallJob, ModelSource
|
||||
from invokeai.app.services.model_records import ModelRecordChanges, ModelRecordServiceBase
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from invokeai.app.services.events.events_base import EventServiceBase
|
||||
|
||||
|
||||
class ModelInstallServiceBase(ABC):
|
||||
"""Abstract base class for InvokeAI model installation."""
|
||||
|
||||
@abstractmethod
|
||||
def __init__(
|
||||
self,
|
||||
app_config: InvokeAIAppConfig,
|
||||
record_store: ModelRecordServiceBase,
|
||||
download_queue: DownloadQueueServiceBase,
|
||||
event_bus: Optional["EventServiceBase"] = None,
|
||||
):
|
||||
"""
|
||||
Create ModelInstallService object.
|
||||
|
||||
:param config: Systemwide InvokeAIAppConfig.
|
||||
:param store: Systemwide ModelConfigStore
|
||||
:param event_bus: InvokeAI event bus for reporting events to.
|
||||
"""
|
||||
|
||||
# make the invoker optional here because we don't need it and it
|
||||
# makes the installer harder to use outside the web app
|
||||
@abstractmethod
|
||||
def start(self, invoker: Optional[Invoker] = None) -> None:
|
||||
"""Start the installer service."""
|
||||
|
||||
@abstractmethod
|
||||
def stop(self, invoker: Optional[Invoker] = None) -> None:
|
||||
"""Stop the model install service. After this the objection can be safely deleted."""
|
||||
|
||||
@property
|
||||
@abstractmethod
|
||||
def app_config(self) -> InvokeAIAppConfig:
|
||||
"""Return the appConfig object associated with the installer."""
|
||||
|
||||
@property
|
||||
@abstractmethod
|
||||
def record_store(self) -> ModelRecordServiceBase:
|
||||
"""Return the ModelRecoreService object associated with the installer."""
|
||||
|
||||
@property
|
||||
@abstractmethod
|
||||
def event_bus(self) -> Optional["EventServiceBase"]:
|
||||
"""Return the event service base object associated with the installer."""
|
||||
|
||||
@abstractmethod
|
||||
def register_path(
|
||||
self,
|
||||
model_path: Union[Path, str],
|
||||
config: Optional[ModelRecordChanges] = None,
|
||||
) -> str:
|
||||
"""
|
||||
Probe and register the model at model_path.
|
||||
|
||||
This keeps the model in its current location.
|
||||
|
||||
:param model_path: Filesystem Path to the model.
|
||||
:param config: ModelRecordChanges object that will override autoassigned model record values.
|
||||
:returns id: The string ID of the registered model.
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
def unregister(self, key: str) -> None:
|
||||
"""Remove model with indicated key from the database."""
|
||||
|
||||
@abstractmethod
|
||||
def delete(self, key: str) -> None:
|
||||
"""Remove model with indicated key from the database. Delete its files only if they are within our models directory."""
|
||||
|
||||
@abstractmethod
|
||||
def unconditionally_delete(self, key: str) -> None:
|
||||
"""Remove model with indicated key from the database and unconditionally delete weight files from disk."""
|
||||
|
||||
@abstractmethod
|
||||
def install_path(
|
||||
self,
|
||||
model_path: Union[Path, str],
|
||||
config: Optional[ModelRecordChanges] = None,
|
||||
) -> str:
|
||||
"""
|
||||
Probe, register and install the model in the models directory.
|
||||
|
||||
This moves the model from its current location into
|
||||
the models directory handled by InvokeAI.
|
||||
|
||||
:param model_path: Filesystem Path to the model.
|
||||
:param config: ModelRecordChanges object that will override autoassigned model record values.
|
||||
:returns id: The string ID of the registered model.
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
def heuristic_import(
|
||||
self,
|
||||
source: str,
|
||||
config: Optional[ModelRecordChanges] = None,
|
||||
access_token: Optional[str] = None,
|
||||
inplace: Optional[bool] = False,
|
||||
) -> ModelInstallJob:
|
||||
r"""Install the indicated model using heuristics to interpret user intentions.
|
||||
|
||||
:param source: String source
|
||||
:param config: Optional ModelRecordChanges object. Any fields in this object
|
||||
will override corresponding autoassigned probe fields in the
|
||||
model's config record as described in `import_model()`.
|
||||
:param access_token: Optional access token for remote sources.
|
||||
|
||||
The source can be:
|
||||
1. A local file path in posix() format (`/foo/bar` or `C:\foo\bar`)
|
||||
2. An http or https URL (`https://foo.bar/foo`)
|
||||
3. A HuggingFace repo_id (`foo/bar`, `foo/bar:fp16`, `foo/bar:fp16:vae`)
|
||||
|
||||
We extend the HuggingFace repo_id syntax to include the variant and the
|
||||
subfolder or path. The following are acceptable alternatives:
|
||||
stabilityai/stable-diffusion-v4
|
||||
stabilityai/stable-diffusion-v4:fp16
|
||||
stabilityai/stable-diffusion-v4:fp16:vae
|
||||
stabilityai/stable-diffusion-v4::/checkpoints/sd4.safetensors
|
||||
stabilityai/stable-diffusion-v4:onnx:vae
|
||||
|
||||
Because a local file path can look like a huggingface repo_id, the logic
|
||||
first checks whether the path exists on disk, and if not, it is treated as
|
||||
a parseable huggingface repo.
|
||||
|
||||
The previous support for recursing into a local folder and loading all model-like files
|
||||
has been removed.
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def import_model(
|
||||
self,
|
||||
source: ModelSource,
|
||||
config: Optional[ModelRecordChanges] = None,
|
||||
) -> ModelInstallJob:
|
||||
"""Install the indicated model.
|
||||
|
||||
:param source: ModelSource object
|
||||
|
||||
:param config: Optional dict. Any fields in this dict
|
||||
will override corresponding autoassigned probe fields in the
|
||||
model's config record. Use it to override
|
||||
`name`, `description`, `base_type`, `model_type`, `format`,
|
||||
`prediction_type`, and/or `image_size`.
|
||||
|
||||
This will download the model located at `source`,
|
||||
probe it, and install it into the models directory.
|
||||
This call is executed asynchronously in a separate
|
||||
thread and will issue the following events on the event bus:
|
||||
|
||||
- model_install_started
|
||||
- model_install_error
|
||||
- model_install_completed
|
||||
|
||||
The `inplace` flag does not affect the behavior of downloaded
|
||||
models, which are always moved into the `models` directory.
|
||||
|
||||
The call returns a ModelInstallJob object which can be
|
||||
polled to learn the current status and/or error message.
|
||||
|
||||
Variants recognized by HuggingFace currently are:
|
||||
1. onnx
|
||||
2. openvino
|
||||
3. fp16
|
||||
4. None (usually returns fp32 model)
|
||||
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
def get_job_by_source(self, source: ModelSource) -> List[ModelInstallJob]:
|
||||
"""Return the ModelInstallJob(s) corresponding to the provided source."""
|
||||
|
||||
@abstractmethod
|
||||
def get_job_by_id(self, id: int) -> ModelInstallJob:
|
||||
"""Return the ModelInstallJob corresponding to the provided id. Raises ValueError if no job has that ID."""
|
||||
|
||||
@abstractmethod
|
||||
def list_jobs(self) -> List[ModelInstallJob]: # noqa D102
|
||||
"""
|
||||
List active and complete install jobs.
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
def prune_jobs(self) -> None:
|
||||
"""Prune all completed and errored jobs."""
|
||||
|
||||
@abstractmethod
|
||||
def cancel_job(self, job: ModelInstallJob) -> None:
|
||||
"""Cancel the indicated job."""
|
||||
|
||||
@abstractmethod
|
||||
def pause_job(self, job: ModelInstallJob) -> None:
|
||||
"""Pause the indicated job, preserving partial downloads."""
|
||||
|
||||
@abstractmethod
|
||||
def resume_job(self, job: ModelInstallJob) -> None:
|
||||
"""Resume a previously paused job."""
|
||||
|
||||
@abstractmethod
|
||||
def restart_failed(self, job: ModelInstallJob) -> None:
|
||||
"""Restart failed or non-resumable downloads for a job."""
|
||||
|
||||
@abstractmethod
|
||||
def restart_file(self, job: ModelInstallJob, file_source: str) -> None:
|
||||
"""Restart a specific file download for a job."""
|
||||
|
||||
@abstractmethod
|
||||
def wait_for_job(self, job: ModelInstallJob, timeout: int = 0) -> ModelInstallJob:
|
||||
"""Wait for the indicated job to reach a terminal state.
|
||||
|
||||
This will block until the indicated install job has completed,
|
||||
been cancelled, or errored out.
|
||||
|
||||
:param job: The job to wait on.
|
||||
:param timeout: Wait up to indicated number of seconds. Raise a TimeoutError if
|
||||
the job hasn't completed within the indicated time.
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
def wait_for_installs(self, timeout: int = 0) -> List[ModelInstallJob]:
|
||||
"""
|
||||
Wait for all pending installs to complete.
|
||||
|
||||
This will block until all pending installs have
|
||||
completed, been cancelled, or errored out.
|
||||
|
||||
:param timeout: Wait up to indicated number of seconds. Raise an Exception('timeout') if
|
||||
installs do not complete within the indicated time. A timeout of zero (the default)
|
||||
will block indefinitely until the installs complete.
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
def download_and_cache_model(self, source: str | AnyHttpUrl) -> Path:
|
||||
"""
|
||||
Download the model file located at source to the models cache and return its Path.
|
||||
|
||||
:param source: A string representing a URL or repo_id.
|
||||
|
||||
The model file will be downloaded into the system-wide model cache
|
||||
(`models/.cache`) if it isn't already there. Note that the model cache
|
||||
is periodically cleared of infrequently-used entries when the model
|
||||
converter runs.
|
||||
|
||||
Note that this doesn't automatically install or register the model, but is
|
||||
intended for use by nodes that need access to models that aren't directly
|
||||
supported by InvokeAI. The downloading process takes advantage of the download queue
|
||||
to avoid interrupting other operations.
|
||||
"""
|
||||
@@ -0,0 +1,270 @@
|
||||
import re
|
||||
import traceback
|
||||
from enum import Enum
|
||||
from pathlib import Path
|
||||
from typing import Literal, Optional, Set, Union
|
||||
|
||||
from pydantic import BaseModel, Field, PrivateAttr, field_validator
|
||||
from pydantic.networks import AnyHttpUrl
|
||||
from typing_extensions import Annotated
|
||||
|
||||
from invokeai.app.services.download import DownloadJob, MultiFileDownloadJob
|
||||
from invokeai.app.services.model_records import ModelRecordChanges
|
||||
from invokeai.backend.model_manager.configs.factory import AnyModelConfig
|
||||
from invokeai.backend.model_manager.metadata import AnyModelRepoMetadata
|
||||
from invokeai.backend.model_manager.taxonomy import ModelRepoVariant, ModelSourceType
|
||||
|
||||
|
||||
class InvalidModelConfigException(Exception):
|
||||
"""Raised when a model configuration is invalid."""
|
||||
|
||||
pass
|
||||
|
||||
|
||||
class InstallStatus(str, Enum):
|
||||
"""State of an install job running in the background."""
|
||||
|
||||
WAITING = "waiting" # waiting to be dequeued
|
||||
DOWNLOADING = "downloading" # downloading of model files in process
|
||||
DOWNLOADS_DONE = "downloads_done" # downloading done, waiting to run
|
||||
RUNNING = "running" # being processed
|
||||
PAUSED = "paused" # paused, can be resumed
|
||||
COMPLETED = "completed" # finished running
|
||||
ERROR = "error" # terminated with an error message
|
||||
CANCELLED = "cancelled" # terminated with an error message
|
||||
|
||||
|
||||
class UnknownInstallJobException(Exception):
|
||||
"""Raised when the status of an unknown job is requested."""
|
||||
|
||||
|
||||
class StringLikeSource(BaseModel):
|
||||
"""
|
||||
Base class for model sources, implements functions that lets the source be sorted and indexed.
|
||||
|
||||
These shenanigans let this stuff work:
|
||||
|
||||
source1 = LocalModelSource(path='C:/users/mort/foo.safetensors')
|
||||
mydict = {source1: 'model 1'}
|
||||
assert mydict['C:/users/mort/foo.safetensors'] == 'model 1'
|
||||
assert mydict[LocalModelSource(path='C:/users/mort/foo.safetensors')] == 'model 1'
|
||||
|
||||
source2 = LocalModelSource(path=Path('C:/users/mort/foo.safetensors'))
|
||||
assert source1 == source2
|
||||
assert source1 == 'C:/users/mort/foo.safetensors'
|
||||
"""
|
||||
|
||||
def __hash__(self) -> int:
|
||||
"""Return hash of the path field, for indexing."""
|
||||
return hash(str(self))
|
||||
|
||||
def __lt__(self, other: object) -> int:
|
||||
"""Return comparison of the stringified version, for sorting."""
|
||||
return str(self) < str(other)
|
||||
|
||||
def __eq__(self, other: object) -> bool:
|
||||
"""Return equality on the stringified version."""
|
||||
if isinstance(other, Path):
|
||||
return str(self) == other.as_posix()
|
||||
else:
|
||||
return str(self) == str(other)
|
||||
|
||||
|
||||
class LocalModelSource(StringLikeSource):
|
||||
"""A local file or directory path."""
|
||||
|
||||
path: str | Path
|
||||
inplace: Optional[bool] = False
|
||||
type: Literal["local"] = "local"
|
||||
|
||||
# these methods allow the source to be used in a string-like way,
|
||||
# for example as an index into a dict
|
||||
def __str__(self) -> str:
|
||||
"""Return string version of path when string rep needed."""
|
||||
return Path(self.path).as_posix()
|
||||
|
||||
|
||||
class HFModelSource(StringLikeSource):
|
||||
"""
|
||||
A HuggingFace repo_id with optional variant, sub-folder(s) and access token.
|
||||
Note that the variant option, if not provided to the constructor, will default to fp16, which is
|
||||
what people (almost) always want.
|
||||
|
||||
The subfolder can be a single path or multiple paths joined by '+' (e.g., "text_encoder+tokenizer").
|
||||
When multiple subfolders are specified, all of them will be downloaded and combined into the model directory.
|
||||
"""
|
||||
|
||||
repo_id: str
|
||||
variant: Optional[ModelRepoVariant] = ModelRepoVariant.FP16
|
||||
subfolder: Optional[Path] = None
|
||||
access_token: Optional[str] = None
|
||||
type: Literal["hf"] = "hf"
|
||||
|
||||
@field_validator("repo_id")
|
||||
@classmethod
|
||||
def proper_repo_id(cls, v: str) -> str: # noqa D102
|
||||
if not re.match(r"^([.\w-]+/[.\w-]+)$", v):
|
||||
raise ValueError(f"{v}: invalid repo_id format")
|
||||
return v
|
||||
|
||||
@property
|
||||
def subfolders(self) -> list[Path]:
|
||||
"""Return list of subfolders (supports '+' separated multiple subfolders)."""
|
||||
if self.subfolder is None:
|
||||
return []
|
||||
subfolder_str = self.subfolder.as_posix()
|
||||
if "+" in subfolder_str:
|
||||
return [Path(s.strip()) for s in subfolder_str.split("+")]
|
||||
return [self.subfolder]
|
||||
|
||||
def __str__(self) -> str:
|
||||
"""Return string version of repoid when string rep needed."""
|
||||
base: str = self.repo_id
|
||||
if self.variant:
|
||||
base += f":{self.variant or ''}"
|
||||
if self.subfolder:
|
||||
base += f"::{self.subfolder.as_posix()}"
|
||||
return base
|
||||
|
||||
|
||||
class URLModelSource(StringLikeSource):
|
||||
"""A generic URL point to a checkpoint file."""
|
||||
|
||||
url: AnyHttpUrl
|
||||
access_token: Optional[str] = None
|
||||
type: Literal["url"] = "url"
|
||||
|
||||
def __str__(self) -> str:
|
||||
"""Return string version of the url when string rep needed."""
|
||||
return str(self.url)
|
||||
|
||||
|
||||
class ExternalModelSource(StringLikeSource):
|
||||
"""An external provider model identifier."""
|
||||
|
||||
provider_id: str
|
||||
provider_model_id: str
|
||||
type: Literal["external"] = "external"
|
||||
|
||||
def __str__(self) -> str:
|
||||
return f"external://{self.provider_id}/{self.provider_model_id}"
|
||||
|
||||
|
||||
ModelSource = Annotated[
|
||||
Union[LocalModelSource, HFModelSource, URLModelSource, ExternalModelSource],
|
||||
Field(discriminator="type"),
|
||||
]
|
||||
|
||||
MODEL_SOURCE_TO_TYPE_MAP = {
|
||||
URLModelSource: ModelSourceType.Url,
|
||||
HFModelSource: ModelSourceType.HFRepoID,
|
||||
LocalModelSource: ModelSourceType.Path,
|
||||
ExternalModelSource: ModelSourceType.External,
|
||||
}
|
||||
|
||||
|
||||
class ModelInstallJob(BaseModel):
|
||||
"""Object that tracks the current status of an install request."""
|
||||
|
||||
id: int = Field(description="Unique ID for this job")
|
||||
status: InstallStatus = Field(default=InstallStatus.WAITING, description="Current status of install process")
|
||||
error_reason: Optional[str] = Field(default=None, description="Information about why the job failed")
|
||||
config_in: ModelRecordChanges = Field(
|
||||
default_factory=ModelRecordChanges,
|
||||
description="Configuration information (e.g. 'description') to apply to model.",
|
||||
)
|
||||
config_out: Optional[AnyModelConfig] = Field(
|
||||
default=None, description="After successful installation, this will hold the configuration object."
|
||||
)
|
||||
inplace: bool = Field(
|
||||
default=False, description="Leave model in its current location; otherwise install under models directory"
|
||||
)
|
||||
source: ModelSource = Field(description="Source (URL, repo_id, or local path) of model")
|
||||
local_path: Path = Field(description="Path to locally-downloaded model; may be the same as the source")
|
||||
bytes: int = Field(
|
||||
default=0, description="For a remote model, the number of bytes downloaded so far (may not be available)"
|
||||
)
|
||||
total_bytes: int = Field(default=0, description="Total size of the model to be installed")
|
||||
source_metadata: Optional[AnyModelRepoMetadata] = Field(
|
||||
default=None, description="Metadata provided by the model source"
|
||||
)
|
||||
download_parts: Set[DownloadJob] = Field(
|
||||
default_factory=set, description="Download jobs contributing to this install"
|
||||
)
|
||||
error: Optional[str] = Field(
|
||||
default=None, description="On an error condition, this field will contain the text of the exception"
|
||||
)
|
||||
error_traceback: Optional[str] = Field(
|
||||
default=None, description="On an error condition, this field will contain the exception traceback"
|
||||
)
|
||||
# internal flags and transitory settings
|
||||
_install_tmpdir: Optional[Path] = PrivateAttr(default=None)
|
||||
_multifile_job: Optional[MultiFileDownloadJob] = PrivateAttr(default=None)
|
||||
_exception: Optional[Exception] = PrivateAttr(default=None)
|
||||
_resume_metadata: Optional[dict] = PrivateAttr(default=None)
|
||||
|
||||
def set_error(self, e: Exception) -> None:
|
||||
"""Record the error and traceback from an exception."""
|
||||
self._exception = e
|
||||
self.error = str(e)
|
||||
self.error_traceback = self._format_error(e)
|
||||
self.status = InstallStatus.ERROR
|
||||
self.error_reason = self._exception.__class__.__name__ if self._exception else None
|
||||
|
||||
def cancel(self) -> None:
|
||||
"""Call to cancel the job."""
|
||||
self.status = InstallStatus.CANCELLED
|
||||
|
||||
@property
|
||||
def error_type(self) -> Optional[str]:
|
||||
"""Class name of the exception that led to status==ERROR."""
|
||||
return self._exception.__class__.__name__ if self._exception else None
|
||||
|
||||
def _format_error(self, exception: Exception) -> str:
|
||||
"""Error traceback."""
|
||||
return "".join(traceback.format_exception(exception))
|
||||
|
||||
@property
|
||||
def cancelled(self) -> bool:
|
||||
"""Set status to CANCELLED."""
|
||||
return self.status == InstallStatus.CANCELLED
|
||||
|
||||
@property
|
||||
def errored(self) -> bool:
|
||||
"""Return true if job has errored."""
|
||||
return self.status == InstallStatus.ERROR
|
||||
|
||||
@property
|
||||
def waiting(self) -> bool:
|
||||
"""Return true if job is waiting to run."""
|
||||
return self.status == InstallStatus.WAITING
|
||||
|
||||
@property
|
||||
def downloading(self) -> bool:
|
||||
"""Return true if job is downloading."""
|
||||
return self.status == InstallStatus.DOWNLOADING
|
||||
|
||||
@property
|
||||
def downloads_done(self) -> bool:
|
||||
"""Return true if job's downloads ae done."""
|
||||
return self.status == InstallStatus.DOWNLOADS_DONE
|
||||
|
||||
@property
|
||||
def paused(self) -> bool:
|
||||
"""Return true if job is paused."""
|
||||
return self.status == InstallStatus.PAUSED
|
||||
|
||||
@property
|
||||
def running(self) -> bool:
|
||||
"""Return true if job is running."""
|
||||
return self.status == InstallStatus.RUNNING
|
||||
|
||||
@property
|
||||
def complete(self) -> bool:
|
||||
"""Return true if job completed without errors."""
|
||||
return self.status == InstallStatus.COMPLETED
|
||||
|
||||
@property
|
||||
def in_terminal_state(self) -> bool:
|
||||
"""Return true if job is in a terminal state."""
|
||||
return self.status in [InstallStatus.COMPLETED, InstallStatus.ERROR, InstallStatus.CANCELLED]
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,6 @@
|
||||
"""Initialization file for model load service module."""
|
||||
|
||||
from invokeai.app.services.model_load.model_load_base import ModelLoadServiceBase
|
||||
from invokeai.app.services.model_load.model_load_default import ModelLoadService
|
||||
|
||||
__all__ = ["ModelLoadServiceBase", "ModelLoadService"]
|
||||
@@ -0,0 +1,52 @@
|
||||
# Copyright (c) 2024 Lincoln D. Stein and the InvokeAI Team
|
||||
"""Base class for model loader."""
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
from pathlib import Path
|
||||
from typing import Callable, Optional
|
||||
|
||||
from invokeai.backend.model_manager.configs.factory import AnyModelConfig
|
||||
from invokeai.backend.model_manager.load import LoadedModel, LoadedModelWithoutConfig
|
||||
from invokeai.backend.model_manager.load.model_cache.model_cache import ModelCache
|
||||
from invokeai.backend.model_manager.taxonomy import AnyModel, SubModelType
|
||||
|
||||
|
||||
class ModelLoadServiceBase(ABC):
|
||||
"""Wrapper around AnyModelLoader."""
|
||||
|
||||
@abstractmethod
|
||||
def load_model(self, model_config: AnyModelConfig, submodel_type: Optional[SubModelType] = None) -> LoadedModel:
|
||||
"""
|
||||
Given a model's configuration, load it and return the LoadedModel object.
|
||||
|
||||
:param model_config: Model configuration record (as returned by ModelRecordBase.get_model())
|
||||
:param submodel: For main (pipeline models), the submodel to fetch.
|
||||
"""
|
||||
|
||||
@property
|
||||
@abstractmethod
|
||||
def ram_cache(self) -> ModelCache:
|
||||
"""Return the RAM cache used by this loader."""
|
||||
|
||||
@abstractmethod
|
||||
def load_model_from_path(
|
||||
self, model_path: Path, loader: Optional[Callable[[Path], AnyModel]] = None
|
||||
) -> LoadedModelWithoutConfig:
|
||||
"""
|
||||
Load the model file or directory located at the indicated Path.
|
||||
|
||||
This will load an arbitrary model file into the RAM cache. If the optional loader
|
||||
argument is provided, the loader will be invoked to load the model into
|
||||
memory. Otherwise the method will call safetensors.torch.load_file() or
|
||||
torch.load() as appropriate to the file suffix.
|
||||
|
||||
Be aware that this returns a LoadedModelWithoutConfig object, which is the same as
|
||||
LoadedModel, but without the config attribute.
|
||||
|
||||
Args:
|
||||
model_path: A pathlib.Path to a checkpoint-style models file
|
||||
loader: A Callable that expects a Path and returns a Dict[str, Tensor]
|
||||
|
||||
Returns:
|
||||
A LoadedModel object.
|
||||
"""
|
||||
@@ -0,0 +1,128 @@
|
||||
# Copyright (c) 2024 Lincoln D. Stein and the InvokeAI Team
|
||||
"""Implementation of model loader service."""
|
||||
|
||||
from pathlib import Path
|
||||
from typing import Callable, Optional, Type
|
||||
|
||||
from picklescan.scanner import scan_file_path
|
||||
from safetensors.torch import load_file as safetensors_load_file
|
||||
from torch import load as torch_load
|
||||
|
||||
from invokeai.app.services.config import InvokeAIAppConfig
|
||||
from invokeai.app.services.invoker import Invoker
|
||||
from invokeai.app.services.model_load.model_load_base import ModelLoadServiceBase
|
||||
from invokeai.backend.model_manager.configs.factory import AnyModelConfig
|
||||
from invokeai.backend.model_manager.load import (
|
||||
LoadedModel,
|
||||
LoadedModelWithoutConfig,
|
||||
ModelLoaderRegistry,
|
||||
ModelLoaderRegistryBase,
|
||||
)
|
||||
from invokeai.backend.model_manager.load.model_cache.model_cache import ModelCache
|
||||
from invokeai.backend.model_manager.load.model_loaders.generic_diffusers import GenericDiffusersLoader
|
||||
from invokeai.backend.model_manager.taxonomy import AnyModel, SubModelType
|
||||
from invokeai.backend.util.devices import TorchDevice
|
||||
from invokeai.backend.util.logging import InvokeAILogger
|
||||
|
||||
|
||||
class ModelLoadService(ModelLoadServiceBase):
|
||||
"""Wrapper around ModelLoaderRegistry."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
app_config: InvokeAIAppConfig,
|
||||
ram_cache: ModelCache,
|
||||
registry: Optional[Type[ModelLoaderRegistryBase]] = ModelLoaderRegistry,
|
||||
):
|
||||
"""Initialize the model load service."""
|
||||
logger = InvokeAILogger.get_logger(self.__class__.__name__)
|
||||
logger.setLevel(app_config.log_level.upper())
|
||||
self._logger = logger
|
||||
self._app_config = app_config
|
||||
self._ram_cache = ram_cache
|
||||
self._registry = registry
|
||||
|
||||
def start(self, invoker: Invoker) -> None:
|
||||
self._invoker = invoker
|
||||
|
||||
@property
|
||||
def ram_cache(self) -> ModelCache:
|
||||
"""Return the RAM cache used by this loader."""
|
||||
return self._ram_cache
|
||||
|
||||
def load_model(self, model_config: AnyModelConfig, submodel_type: Optional[SubModelType] = None) -> LoadedModel:
|
||||
"""
|
||||
Given a model's configuration, load it and return the LoadedModel object.
|
||||
|
||||
:param model_config: Model configuration record (as returned by ModelRecordBase.get_model())
|
||||
:param submodel: For main (pipeline models), the submodel to fetch.
|
||||
"""
|
||||
|
||||
# We don't have an invoker during testing
|
||||
# TODO(psyche): Mock this method on the invoker in the tests
|
||||
if hasattr(self, "_invoker"):
|
||||
self._invoker.services.events.emit_model_load_started(model_config, submodel_type)
|
||||
|
||||
implementation, model_config, submodel_type = self._registry.get_implementation(model_config, submodel_type) # type: ignore
|
||||
loaded_model: LoadedModel = implementation(
|
||||
app_config=self._app_config,
|
||||
logger=self._logger,
|
||||
ram_cache=self._ram_cache,
|
||||
).load_model(model_config, submodel_type)
|
||||
|
||||
if hasattr(self, "_invoker"):
|
||||
self._invoker.services.events.emit_model_load_complete(model_config, submodel_type)
|
||||
|
||||
return loaded_model
|
||||
|
||||
def load_model_from_path(
|
||||
self, model_path: Path, loader: Optional[Callable[[Path], AnyModel]] = None
|
||||
) -> LoadedModelWithoutConfig:
|
||||
cache_key = str(model_path)
|
||||
try:
|
||||
return LoadedModelWithoutConfig(cache_record=self._ram_cache.get(key=cache_key), cache=self._ram_cache)
|
||||
except IndexError:
|
||||
pass
|
||||
|
||||
def torch_load_file(checkpoint: Path) -> AnyModel:
|
||||
scan_result = scan_file_path(checkpoint)
|
||||
if scan_result.infected_files != 0:
|
||||
if self._app_config.unsafe_disable_picklescan:
|
||||
self._logger.warning(
|
||||
f"Model at {checkpoint} is potentially infected by malware, but picklescan is disabled. "
|
||||
"Proceeding with caution."
|
||||
)
|
||||
else:
|
||||
raise Exception(f"The model at {checkpoint} is potentially infected by malware. Aborting load.")
|
||||
if scan_result.scan_err:
|
||||
if self._app_config.unsafe_disable_picklescan:
|
||||
self._logger.warning(
|
||||
f"Error scanning model at {checkpoint} for malware, but picklescan is disabled. "
|
||||
"Proceeding with caution."
|
||||
)
|
||||
else:
|
||||
raise Exception(f"Error scanning model at {checkpoint} for malware. Aborting load.")
|
||||
|
||||
result = torch_load(checkpoint, map_location="cpu")
|
||||
return result
|
||||
|
||||
def diffusers_load_directory(directory: Path) -> AnyModel:
|
||||
load_class = GenericDiffusersLoader(
|
||||
app_config=self._app_config,
|
||||
logger=self._logger,
|
||||
ram_cache=self._ram_cache,
|
||||
convert_cache=self.convert_cache,
|
||||
).get_hf_load_class(directory)
|
||||
return load_class.from_pretrained(model_path, torch_dtype=TorchDevice.choose_torch_dtype())
|
||||
|
||||
loader = loader or (
|
||||
diffusers_load_directory
|
||||
if model_path.is_dir()
|
||||
else torch_load_file
|
||||
if model_path.suffix.endswith((".ckpt", ".pt", ".pth", ".bin"))
|
||||
else lambda path: safetensors_load_file(path, device="cpu")
|
||||
)
|
||||
assert loader is not None
|
||||
raw_model = loader(model_path)
|
||||
self._ram_cache.put(key=cache_key, model=raw_model)
|
||||
return LoadedModelWithoutConfig(cache_record=self._ram_cache.get(key=cache_key), cache=self._ram_cache)
|
||||
@@ -0,0 +1,10 @@
|
||||
"""Initialization file for model manager service."""
|
||||
|
||||
from invokeai.app.services.model_manager.model_manager_default import ModelManagerService, ModelManagerServiceBase
|
||||
from invokeai.backend.model_manager.load import LoadedModel
|
||||
|
||||
__all__ = [
|
||||
"ModelManagerServiceBase",
|
||||
"ModelManagerService",
|
||||
"LoadedModel",
|
||||
]
|
||||
@@ -0,0 +1,67 @@
|
||||
# Copyright (c) 2023 Lincoln D. Stein and the InvokeAI Team
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
|
||||
import torch
|
||||
from typing_extensions import Self
|
||||
|
||||
from invokeai.app.services.config.config_default import InvokeAIAppConfig
|
||||
from invokeai.app.services.download.download_base import DownloadQueueServiceBase
|
||||
from invokeai.app.services.events.events_base import EventServiceBase
|
||||
from invokeai.app.services.invoker import Invoker
|
||||
from invokeai.app.services.model_install.model_install_base import ModelInstallServiceBase
|
||||
from invokeai.app.services.model_load.model_load_base import ModelLoadServiceBase
|
||||
from invokeai.app.services.model_records.model_records_base import ModelRecordServiceBase
|
||||
|
||||
|
||||
class ModelManagerServiceBase(ABC):
|
||||
"""Abstract base class for the model manager service."""
|
||||
|
||||
# attributes:
|
||||
# store: ModelRecordServiceBase = Field(description="An instance of the model record configuration service.")
|
||||
# install: ModelInstallServiceBase = Field(description="An instance of the model install service.")
|
||||
# load: ModelLoadServiceBase = Field(description="An instance of the model load service.")
|
||||
|
||||
@classmethod
|
||||
@abstractmethod
|
||||
def build_model_manager(
|
||||
cls,
|
||||
app_config: InvokeAIAppConfig,
|
||||
model_record_service: ModelRecordServiceBase,
|
||||
download_queue: DownloadQueueServiceBase,
|
||||
events: EventServiceBase,
|
||||
execution_device: torch.device,
|
||||
) -> Self:
|
||||
"""
|
||||
Construct the model manager service instance.
|
||||
|
||||
Use it rather than the __init__ constructor. This class
|
||||
method simplifies the construction considerably.
|
||||
"""
|
||||
pass
|
||||
|
||||
@property
|
||||
@abstractmethod
|
||||
def store(self) -> ModelRecordServiceBase:
|
||||
"""Return the ModelRecordServiceBase used to store and retrieve configuration records."""
|
||||
pass
|
||||
|
||||
@property
|
||||
@abstractmethod
|
||||
def load(self) -> ModelLoadServiceBase:
|
||||
"""Return the ModelLoadServiceBase used to load models from their configuration records."""
|
||||
pass
|
||||
|
||||
@property
|
||||
@abstractmethod
|
||||
def install(self) -> ModelInstallServiceBase:
|
||||
"""Return the ModelInstallServiceBase used to download and manipulate model files."""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def start(self, invoker: Invoker) -> None:
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def stop(self, invoker: Invoker) -> None:
|
||||
pass
|
||||
@@ -0,0 +1,111 @@
|
||||
# Copyright (c) 2023 Lincoln D. Stein and the InvokeAI Team
|
||||
"""Implementation of ModelManagerServiceBase."""
|
||||
|
||||
from typing import Optional
|
||||
|
||||
import torch
|
||||
from typing_extensions import Self
|
||||
|
||||
from invokeai.app.services.config.config_default import InvokeAIAppConfig
|
||||
from invokeai.app.services.download.download_base import DownloadQueueServiceBase
|
||||
from invokeai.app.services.events.events_base import EventServiceBase
|
||||
from invokeai.app.services.invoker import Invoker
|
||||
from invokeai.app.services.model_install.model_install_base import ModelInstallServiceBase
|
||||
from invokeai.app.services.model_install.model_install_default import ModelInstallService
|
||||
from invokeai.app.services.model_load.model_load_base import ModelLoadServiceBase
|
||||
from invokeai.app.services.model_load.model_load_default import ModelLoadService
|
||||
from invokeai.app.services.model_manager.model_manager_base import ModelManagerServiceBase
|
||||
from invokeai.app.services.model_records.model_records_base import ModelRecordServiceBase
|
||||
from invokeai.backend.model_manager.load.model_cache.model_cache import ModelCache
|
||||
from invokeai.backend.model_manager.load.model_loader_registry import ModelLoaderRegistry
|
||||
from invokeai.backend.util.devices import TorchDevice
|
||||
from invokeai.backend.util.logging import InvokeAILogger
|
||||
|
||||
|
||||
class ModelManagerService(ModelManagerServiceBase):
|
||||
"""
|
||||
The ModelManagerService handles various aspects of model installation, maintenance and loading.
|
||||
|
||||
It bundles three distinct services:
|
||||
model_manager.store -- Routines to manage the database of model configuration records.
|
||||
model_manager.install -- Routines to install, move and delete models.
|
||||
model_manager.load -- Routines to load models into memory.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
store: ModelRecordServiceBase,
|
||||
install: ModelInstallServiceBase,
|
||||
load: ModelLoadServiceBase,
|
||||
):
|
||||
self._store = store
|
||||
self._install = install
|
||||
self._load = load
|
||||
|
||||
@property
|
||||
def store(self) -> ModelRecordServiceBase:
|
||||
return self._store
|
||||
|
||||
@property
|
||||
def install(self) -> ModelInstallServiceBase:
|
||||
return self._install
|
||||
|
||||
@property
|
||||
def load(self) -> ModelLoadServiceBase:
|
||||
return self._load
|
||||
|
||||
def start(self, invoker: Invoker) -> None:
|
||||
for service in [self._store, self._install, self._load]:
|
||||
if hasattr(service, "start"):
|
||||
service.start(invoker)
|
||||
|
||||
def stop(self, invoker: Invoker) -> None:
|
||||
# Shutdown the model cache to cancel any pending timers
|
||||
if hasattr(self._load, "ram_cache"):
|
||||
self._load.ram_cache.shutdown()
|
||||
|
||||
for service in [self._store, self._install, self._load]:
|
||||
if hasattr(service, "stop"):
|
||||
service.stop(invoker)
|
||||
|
||||
@classmethod
|
||||
def build_model_manager(
|
||||
cls,
|
||||
app_config: InvokeAIAppConfig,
|
||||
model_record_service: ModelRecordServiceBase,
|
||||
download_queue: DownloadQueueServiceBase,
|
||||
events: EventServiceBase,
|
||||
execution_device: Optional[torch.device] = None,
|
||||
) -> Self:
|
||||
"""
|
||||
Construct the model manager service instance.
|
||||
|
||||
For simplicity, use this class method rather than the __init__ constructor.
|
||||
"""
|
||||
logger = InvokeAILogger.get_logger(cls.__name__)
|
||||
logger.setLevel(app_config.log_level.upper())
|
||||
|
||||
ram_cache = ModelCache(
|
||||
execution_device_working_mem_gb=app_config.device_working_mem_gb,
|
||||
enable_partial_loading=app_config.enable_partial_loading,
|
||||
keep_ram_copy_of_weights=app_config.keep_ram_copy_of_weights,
|
||||
max_ram_cache_size_gb=app_config.max_cache_ram_gb,
|
||||
max_vram_cache_size_gb=app_config.max_cache_vram_gb,
|
||||
execution_device=execution_device or TorchDevice.choose_torch_device(),
|
||||
storage_device="cpu",
|
||||
log_memory_usage=app_config.log_memory_usage,
|
||||
logger=logger,
|
||||
keep_alive_minutes=app_config.model_cache_keep_alive_min,
|
||||
)
|
||||
loader = ModelLoadService(
|
||||
app_config=app_config,
|
||||
ram_cache=ram_cache,
|
||||
registry=ModelLoaderRegistry,
|
||||
)
|
||||
installer = ModelInstallService(
|
||||
app_config=app_config,
|
||||
record_store=model_record_service,
|
||||
download_queue=download_queue,
|
||||
event_bus=events,
|
||||
)
|
||||
return cls(store=model_record_service, install=installer, load=loader)
|
||||
@@ -0,0 +1,23 @@
|
||||
"""Init file for model record services."""
|
||||
|
||||
from .model_records_base import ( # noqa F401
|
||||
DuplicateModelException,
|
||||
InvalidModelException,
|
||||
ModelRecordServiceBase,
|
||||
UnknownModelException,
|
||||
ModelSummary,
|
||||
ModelRecordChanges,
|
||||
ModelRecordOrderBy,
|
||||
)
|
||||
from .model_records_sql import ModelRecordServiceSQL # noqa F401
|
||||
|
||||
__all__ = [
|
||||
"ModelRecordServiceBase",
|
||||
"ModelRecordServiceSQL",
|
||||
"DuplicateModelException",
|
||||
"InvalidModelException",
|
||||
"UnknownModelException",
|
||||
"ModelSummary",
|
||||
"ModelRecordChanges",
|
||||
"ModelRecordOrderBy",
|
||||
]
|
||||
@@ -0,0 +1,298 @@
|
||||
# Copyright (c) 2023 Lincoln D. Stein and the InvokeAI Development Team
|
||||
"""
|
||||
Abstract base class for storing and retrieving model configuration records.
|
||||
"""
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
from enum import Enum
|
||||
from pathlib import Path
|
||||
from typing import Any, List, Optional, Set, Union
|
||||
|
||||
from pydantic import BaseModel, Field, field_validator
|
||||
|
||||
from invokeai.app.services.shared.pagination import PaginatedResults
|
||||
from invokeai.app.services.shared.sqlite.sqlite_common import SQLiteDirection
|
||||
from invokeai.app.util.model_exclude_null import BaseModelExcludeNull
|
||||
from invokeai.backend.model_manager.configs.controlnet import ControlAdapterDefaultSettings
|
||||
from invokeai.backend.model_manager.configs.external_api import (
|
||||
ExternalApiModelDefaultSettings,
|
||||
ExternalModelCapabilities,
|
||||
)
|
||||
from invokeai.backend.model_manager.configs.factory import AnyModelConfig
|
||||
from invokeai.backend.model_manager.configs.lora import LoraModelDefaultSettings
|
||||
from invokeai.backend.model_manager.configs.main import MainModelDefaultSettings
|
||||
from invokeai.backend.model_manager.taxonomy import (
|
||||
BaseModelType,
|
||||
ClipVariantType,
|
||||
Flux2VariantType,
|
||||
FluxVariantType,
|
||||
ModelFormat,
|
||||
ModelSourceType,
|
||||
ModelType,
|
||||
ModelVariantType,
|
||||
Qwen3VariantType,
|
||||
QwenImageVariantType,
|
||||
SchedulerPredictionType,
|
||||
ZImageVariantType,
|
||||
)
|
||||
|
||||
|
||||
class DuplicateModelException(Exception):
|
||||
"""Raised on an attempt to add a model with the same key twice."""
|
||||
|
||||
|
||||
class InvalidModelException(Exception):
|
||||
"""Raised when an invalid model is detected."""
|
||||
|
||||
|
||||
class UnknownModelException(Exception):
|
||||
"""Raised on an attempt to fetch or delete a model with a nonexistent key."""
|
||||
|
||||
|
||||
class ConfigFileVersionMismatchException(Exception):
|
||||
"""Raised on an attempt to open a config with an incompatible version."""
|
||||
|
||||
|
||||
class ModelRecordOrderBy(str, Enum):
|
||||
"""The order in which to return model summaries."""
|
||||
|
||||
Default = "default" # order by type, base, format and name
|
||||
Type = "type"
|
||||
Base = "base"
|
||||
Name = "name"
|
||||
Format = "format"
|
||||
Size = "size"
|
||||
DateAdded = "created_at"
|
||||
DateModified = "updated_at"
|
||||
Path = "path"
|
||||
|
||||
|
||||
class ModelSummary(BaseModel):
|
||||
"""A short summary of models for UI listing purposes."""
|
||||
|
||||
key: str = Field(description="model key")
|
||||
type: ModelType = Field(description="model type")
|
||||
base: BaseModelType = Field(description="base model")
|
||||
format: ModelFormat = Field(description="model format")
|
||||
name: str = Field(description="model name")
|
||||
description: str = Field(description="short description of model")
|
||||
tags: Set[str] = Field(description="tags associated with model")
|
||||
|
||||
|
||||
class ModelRecordChanges(BaseModelExcludeNull):
|
||||
"""A set of changes to apply to a model."""
|
||||
|
||||
# Changes applicable to all models
|
||||
source: Optional[str] = Field(description="original source of the model", default=None)
|
||||
source_type: Optional[ModelSourceType] = Field(description="type of model source", default=None)
|
||||
source_api_response: Optional[str] = Field(description="metadata from remote source", default=None)
|
||||
source_url: Optional[str] = Field(description="Optional URL for the model (e.g. download page)", default=None)
|
||||
|
||||
@field_validator("source_url", mode="before")
|
||||
@classmethod
|
||||
def validate_source_url(cls, v: Any) -> Optional[str]:
|
||||
if v is None or v == "":
|
||||
return None
|
||||
if not isinstance(v, str):
|
||||
raise ValueError("source_url must be a string")
|
||||
if not v.startswith(("https://", "http://")):
|
||||
raise ValueError("source_url must be an http or https URL")
|
||||
return v
|
||||
|
||||
name: Optional[str] = Field(description="Name of the model.", default=None)
|
||||
path: Optional[str] = Field(description="Path to the model.", default=None)
|
||||
description: Optional[str] = Field(description="Model description", default=None)
|
||||
base: Optional[BaseModelType] = Field(description="The base model.", default=None)
|
||||
type: Optional[ModelType] = Field(description="Type of model", default=None)
|
||||
key: Optional[str] = Field(description="Database ID for this model", default=None)
|
||||
hash: Optional[str] = Field(description="hash of model file", default=None)
|
||||
file_size: Optional[int] = Field(description="Size of model file", default=None)
|
||||
format: Optional[str] = Field(description="format of model file", default=None)
|
||||
trigger_phrases: Optional[set[str]] = Field(description="Set of trigger phrases for this model", default=None)
|
||||
default_settings: Optional[
|
||||
MainModelDefaultSettings
|
||||
| LoraModelDefaultSettings
|
||||
| ControlAdapterDefaultSettings
|
||||
| ExternalApiModelDefaultSettings
|
||||
] = Field(description="Default settings for this model", default=None)
|
||||
|
||||
# External API model changes
|
||||
provider_id: Optional[str] = Field(description="External provider identifier", default=None)
|
||||
provider_model_id: Optional[str] = Field(description="External provider model identifier", default=None)
|
||||
capabilities: Optional[ExternalModelCapabilities] = Field(
|
||||
description="External model capabilities",
|
||||
default=None,
|
||||
)
|
||||
cpu_only: Optional[bool] = Field(description="Whether this model should run on CPU only", default=None)
|
||||
|
||||
# Checkpoint-specific changes
|
||||
# TODO(MM2): Should we expose these? Feels footgun-y...
|
||||
variant: Optional[
|
||||
ModelVariantType
|
||||
| ClipVariantType
|
||||
| FluxVariantType
|
||||
| Flux2VariantType
|
||||
| ZImageVariantType
|
||||
| QwenImageVariantType
|
||||
| Qwen3VariantType
|
||||
] = Field(description="The variant of the model.", default=None)
|
||||
prediction_type: Optional[SchedulerPredictionType] = Field(
|
||||
description="The prediction type of the model.", default=None
|
||||
)
|
||||
upcast_attention: Optional[bool] = Field(description="Whether to upcast attention.", default=None)
|
||||
config_path: Optional[str] = Field(description="Path to config file for model", default=None)
|
||||
|
||||
|
||||
class ModelRecordServiceBase(ABC):
|
||||
"""Abstract base class for storage and retrieval of model configs."""
|
||||
|
||||
@abstractmethod
|
||||
def add_model(self, config: AnyModelConfig) -> AnyModelConfig:
|
||||
"""
|
||||
Add a model to the database.
|
||||
|
||||
:param key: Unique key for the model
|
||||
:param config: Model configuration record, either a dict with the
|
||||
required fields or a ModelConfigBase instance.
|
||||
|
||||
Can raise DuplicateModelException and InvalidModelConfigException exceptions.
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def del_model(self, key: str) -> None:
|
||||
"""
|
||||
Delete a model.
|
||||
|
||||
:param key: Unique key for the model to be deleted
|
||||
|
||||
Can raise an UnknownModelException
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def update_model(self, key: str, changes: ModelRecordChanges, allow_class_change: bool = False) -> AnyModelConfig:
|
||||
"""
|
||||
Update the model, returning the updated version.
|
||||
|
||||
:param key: Unique key for the model to be updated.
|
||||
:param changes: A set of changes to apply to this model. Changes are validated before being written.
|
||||
:param allow_class_change: If True, allows changes that would change the model config class. For example,
|
||||
changing a LoRA into a Main model. This does not disable validation, so the changes must still be valid.
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def replace_model(self, key: str, new_config: AnyModelConfig) -> AnyModelConfig:
|
||||
"""
|
||||
Replace the model record entirely, returning the new record.
|
||||
|
||||
This is used when we re-identify a model and have a new config object.
|
||||
|
||||
:param key: Unique key for the model to be updated.
|
||||
:param new_config: The new model config to write.
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def get_model(self, key: str) -> AnyModelConfig:
|
||||
"""
|
||||
Retrieve the configuration for the indicated model.
|
||||
|
||||
:param key: Key of model config to be fetched.
|
||||
|
||||
Exceptions: UnknownModelException
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def get_model_by_hash(self, hash: str) -> AnyModelConfig:
|
||||
"""
|
||||
Retrieve the configuration for the indicated model.
|
||||
|
||||
:param hash: Hash of model config to be fetched.
|
||||
|
||||
Exceptions: UnknownModelException
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def list_models(
|
||||
self,
|
||||
page: int = 0,
|
||||
per_page: int = 10,
|
||||
order_by: ModelRecordOrderBy = ModelRecordOrderBy.Default,
|
||||
direction: SQLiteDirection = SQLiteDirection.Ascending,
|
||||
) -> PaginatedResults[ModelSummary]:
|
||||
"""Return a paginated summary listing of each model in the database."""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def exists(self, key: str) -> bool:
|
||||
"""
|
||||
Return True if a model with the indicated key exists in the database.
|
||||
|
||||
:param key: Unique key for the model to be deleted
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def search_by_path(
|
||||
self,
|
||||
path: Union[str, Path],
|
||||
) -> List[AnyModelConfig]:
|
||||
"""Return the model(s) having the indicated path."""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def search_by_hash(
|
||||
self,
|
||||
hash: str,
|
||||
) -> List[AnyModelConfig]:
|
||||
"""Return the model(s) having the indicated original hash."""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def search_by_attr(
|
||||
self,
|
||||
model_name: Optional[str] = None,
|
||||
base_model: Optional[BaseModelType] = None,
|
||||
model_type: Optional[ModelType] = None,
|
||||
model_format: Optional[ModelFormat] = None,
|
||||
order_by: ModelRecordOrderBy = ModelRecordOrderBy.Default,
|
||||
direction: SQLiteDirection = SQLiteDirection.Ascending,
|
||||
) -> List[AnyModelConfig]:
|
||||
"""
|
||||
Return models matching name, base and/or type.
|
||||
|
||||
:param model_name: Filter by name of model (optional)
|
||||
:param base_model: Filter by base model (optional)
|
||||
:param model_type: Filter by type of model (optional)
|
||||
:param model_format: Filter by model format (e.g. "diffusers") (optional)
|
||||
|
||||
If none of the optional filters are passed, will return all
|
||||
models in the database.
|
||||
"""
|
||||
pass
|
||||
|
||||
def all_models(self) -> List[AnyModelConfig]:
|
||||
"""Return all the model configs in the database."""
|
||||
return self.search_by_attr()
|
||||
|
||||
def model_info_by_name(self, model_name: str, base_model: BaseModelType, model_type: ModelType) -> AnyModelConfig:
|
||||
"""
|
||||
Return information about a single model using its name, base type and model type.
|
||||
|
||||
If there are more than one model that match, raises a DuplicateModelException.
|
||||
If no model matches, raises an UnknownModelException
|
||||
"""
|
||||
model_configs = self.search_by_attr(model_name=model_name, base_model=base_model, model_type=model_type)
|
||||
if len(model_configs) > 1:
|
||||
raise DuplicateModelException(
|
||||
f"More than one model matched the search criteria: base_model='{base_model}', model_type='{model_type}', model_name='{model_name}'."
|
||||
)
|
||||
if len(model_configs) == 0:
|
||||
raise UnknownModelException(
|
||||
f"More than one model matched the search criteria: base_model='{base_model}', model_type='{model_type}', model_name='{model_name}'."
|
||||
)
|
||||
return model_configs[0]
|
||||
@@ -0,0 +1,458 @@
|
||||
# Copyright (c) 2023 Lincoln D. Stein and the InvokeAI Development Team
|
||||
"""
|
||||
SQL Implementation of the ModelRecordServiceBase API
|
||||
|
||||
Typical usage:
|
||||
|
||||
from invokeai.backend.model_manager import ModelConfigStoreSQL
|
||||
store = ModelConfigStoreSQL(sqlite_db)
|
||||
config = dict(
|
||||
path='/tmp/pokemon.bin',
|
||||
name='old name',
|
||||
base_model='sd-1',
|
||||
type='embedding',
|
||||
format='embedding_file',
|
||||
)
|
||||
|
||||
# adding - the key becomes the model's "key" field
|
||||
store.add_model('key1', config)
|
||||
|
||||
# updating
|
||||
config.name='new name'
|
||||
store.update_model('key1', config)
|
||||
|
||||
# checking for existence
|
||||
if store.exists('key1'):
|
||||
print("yes")
|
||||
|
||||
# fetching config
|
||||
new_config = store.get_model('key1')
|
||||
print(new_config.name, new_config.base)
|
||||
assert new_config.key == 'key1'
|
||||
|
||||
# deleting
|
||||
store.del_model('key1')
|
||||
|
||||
# searching
|
||||
configs = store.search_by_path(path='/tmp/pokemon.bin')
|
||||
configs = store.search_by_hash('750a499f35e43b7e1b4d15c207aa2f01')
|
||||
configs = store.search_by_attr(base_model='sd-2', model_type='main')
|
||||
"""
|
||||
|
||||
import json
|
||||
import logging
|
||||
import sqlite3
|
||||
from math import ceil
|
||||
from pathlib import Path
|
||||
from typing import List, Optional, Union
|
||||
|
||||
import pydantic
|
||||
from pydantic import ValidationError
|
||||
|
||||
from invokeai.app.services.model_records.model_records_base import (
|
||||
DuplicateModelException,
|
||||
ModelRecordChanges,
|
||||
ModelRecordOrderBy,
|
||||
ModelRecordServiceBase,
|
||||
ModelSummary,
|
||||
UnknownModelException,
|
||||
)
|
||||
from invokeai.app.services.shared.pagination import PaginatedResults
|
||||
from invokeai.app.services.shared.sqlite.sqlite_common import SQLiteDirection
|
||||
from invokeai.app.services.shared.sqlite.sqlite_database import SqliteDatabase
|
||||
from invokeai.backend.model_manager.configs.base import Config_Base
|
||||
from invokeai.backend.model_manager.configs.factory import AnyModelConfig, ModelConfigFactory
|
||||
from invokeai.backend.model_manager.taxonomy import BaseModelType, ModelFormat, ModelType
|
||||
|
||||
|
||||
def _construct_config_for_type(fields: dict, target_type: ModelType) -> AnyModelConfig:
|
||||
"""Try every config class whose `type` default matches `target_type` and return the first that validates.
|
||||
|
||||
Used when changing a model's type via the update endpoint: the existing record's `format`/`variant`
|
||||
fields belong to the old class and may not have a discriminator match in the new type space, so we
|
||||
fall back to constructing each candidate class directly with whatever fields it accepts.
|
||||
"""
|
||||
last_error: Exception | None = None
|
||||
for candidate_class in Config_Base.CONFIG_CLASSES:
|
||||
type_field = candidate_class.model_fields.get("type")
|
||||
if type_field is None or type_field.default != target_type:
|
||||
continue
|
||||
try:
|
||||
return candidate_class(**fields) # type: ignore[return-value]
|
||||
except ValidationError as e:
|
||||
last_error = e
|
||||
if last_error is not None:
|
||||
raise last_error
|
||||
raise ValidationError.from_exception_data(
|
||||
f"No model config class found for type={target_type!r}",
|
||||
line_errors=[],
|
||||
)
|
||||
|
||||
|
||||
class ModelRecordServiceSQL(ModelRecordServiceBase):
|
||||
"""Implementation of the ModelConfigStore ABC using a SQL database."""
|
||||
|
||||
def __init__(self, db: SqliteDatabase, logger: logging.Logger):
|
||||
"""
|
||||
Initialize a new object from preexisting sqlite3 connection and threading lock objects.
|
||||
|
||||
:param db: Sqlite connection object
|
||||
"""
|
||||
super().__init__()
|
||||
self._db = db
|
||||
self._logger = logger
|
||||
|
||||
def add_model(self, config: AnyModelConfig) -> AnyModelConfig:
|
||||
"""
|
||||
Add a model to the database.
|
||||
|
||||
:param key: Unique key for the model
|
||||
:param config: Model configuration record, either a dict with the
|
||||
required fields or a ModelConfigBase instance.
|
||||
|
||||
Can raise DuplicateModelException and InvalidModelConfigException exceptions.
|
||||
"""
|
||||
with self._db.transaction() as cursor:
|
||||
try:
|
||||
cursor.execute(
|
||||
"""--sql
|
||||
INSERT INTO models (
|
||||
id,
|
||||
config
|
||||
)
|
||||
VALUES (?,?);
|
||||
""",
|
||||
(
|
||||
config.key,
|
||||
config.model_dump_json(),
|
||||
),
|
||||
)
|
||||
|
||||
except sqlite3.IntegrityError as e:
|
||||
if "UNIQUE constraint failed" in str(e):
|
||||
if "models.path" in str(e):
|
||||
msg = f"A model with path '{config.path}' is already installed"
|
||||
elif "models.name" in str(e):
|
||||
msg = f"A model with name='{config.name}', type='{config.type}', base='{config.base}' is already installed"
|
||||
else:
|
||||
msg = f"A model with key '{config.key}' is already installed"
|
||||
raise DuplicateModelException(msg) from e
|
||||
else:
|
||||
raise e
|
||||
|
||||
return self.get_model(config.key)
|
||||
|
||||
def del_model(self, key: str) -> None:
|
||||
"""
|
||||
Delete a model.
|
||||
|
||||
:param key: Unique key for the model to be deleted
|
||||
|
||||
Can raise an UnknownModelException
|
||||
"""
|
||||
with self._db.transaction() as cursor:
|
||||
cursor.execute(
|
||||
"""--sql
|
||||
DELETE FROM models
|
||||
WHERE id=?;
|
||||
""",
|
||||
(key,),
|
||||
)
|
||||
if cursor.rowcount == 0:
|
||||
raise UnknownModelException("model not found")
|
||||
|
||||
def update_model(self, key: str, changes: ModelRecordChanges, allow_class_change: bool = False) -> AnyModelConfig:
|
||||
with self._db.transaction() as cursor:
|
||||
record = self.get_model(key)
|
||||
|
||||
if allow_class_change:
|
||||
# The changes may cause the model config class to change. To handle this, we need to construct the new
|
||||
# class from scratch rather than trying to modify the existing instance in place.
|
||||
#
|
||||
# 1. Convert the existing record to a dict
|
||||
# 2. Apply the changes to the dict
|
||||
# 3. Attempt to create a new model config from the updated dict
|
||||
|
||||
# 1. Convert the existing record to a dict
|
||||
record_as_dict = record.model_dump()
|
||||
|
||||
# 2. Apply the changes to the dict
|
||||
for field_name in changes.model_fields_set:
|
||||
record_as_dict[field_name] = getattr(changes, field_name)
|
||||
|
||||
# 3. Attempt to create a new model config from the updated dict.
|
||||
#
|
||||
# When the model type is being changed, the previous record's `format` and `variant` likely
|
||||
# belong to the old config class and won't validate against the new one (e.g. switching a
|
||||
# Qwen3 encoder to a Text LLM keeps format=qwen3_encoder, which has no matching discriminator
|
||||
# under text_llm). If the initial validation fails and the type changed, retry with stale
|
||||
# format/variant fields stripped so the new class can apply its own defaults.
|
||||
type_changed = "type" in changes.model_fields_set and changes.type != record.type
|
||||
try:
|
||||
record = ModelConfigFactory.from_dict(record_as_dict)
|
||||
except ValidationError:
|
||||
if not type_changed:
|
||||
raise
|
||||
fallback_dict = dict(record_as_dict)
|
||||
for stale_field in ("format", "variant"):
|
||||
if stale_field not in changes.model_fields_set:
|
||||
fallback_dict.pop(stale_field, None)
|
||||
record = _construct_config_for_type(fallback_dict, changes.type)
|
||||
|
||||
# If we get this far, the updated model config is valid, so we can save it to the database.
|
||||
json_serialized = record.model_dump_json()
|
||||
else:
|
||||
# We are not allowing the model config class to change, so we can just update the existing instance in
|
||||
# place. If the changes are invalid for the existing class, an exception will be raised by pydantic.
|
||||
for field_name in changes.model_fields_set:
|
||||
setattr(record, field_name, getattr(changes, field_name))
|
||||
json_serialized = record.model_dump_json()
|
||||
|
||||
cursor.execute(
|
||||
"""--sql
|
||||
UPDATE models
|
||||
SET
|
||||
config=?
|
||||
WHERE id=?;
|
||||
""",
|
||||
(json_serialized, key),
|
||||
)
|
||||
if cursor.rowcount == 0:
|
||||
raise UnknownModelException("model not found")
|
||||
|
||||
return self.get_model(key)
|
||||
|
||||
def replace_model(self, key: str, new_config: AnyModelConfig) -> AnyModelConfig:
|
||||
if key != new_config.key:
|
||||
raise ValueError("key does not match new_config.key")
|
||||
with self._db.transaction() as cursor:
|
||||
cursor.execute(
|
||||
"""--sql
|
||||
UPDATE models
|
||||
SET
|
||||
config=?
|
||||
WHERE id=?;
|
||||
""",
|
||||
(new_config.model_dump_json(), key),
|
||||
)
|
||||
if cursor.rowcount == 0:
|
||||
raise UnknownModelException("model not found")
|
||||
return self.get_model(key)
|
||||
|
||||
def get_model(self, key: str) -> AnyModelConfig:
|
||||
"""
|
||||
Retrieve the ModelConfigBase instance for the indicated model.
|
||||
|
||||
:param key: Key of model config to be fetched.
|
||||
|
||||
Exceptions: UnknownModelException
|
||||
"""
|
||||
with self._db.transaction() as cursor:
|
||||
cursor.execute(
|
||||
"""--sql
|
||||
SELECT config FROM models
|
||||
WHERE id=?;
|
||||
""",
|
||||
(key,),
|
||||
)
|
||||
rows = cursor.fetchone()
|
||||
if not rows:
|
||||
raise UnknownModelException("model not found")
|
||||
model = ModelConfigFactory.from_dict(json.loads(rows[0]))
|
||||
return model
|
||||
|
||||
def get_model_by_hash(self, hash: str) -> AnyModelConfig:
|
||||
with self._db.transaction() as cursor:
|
||||
cursor.execute(
|
||||
"""--sql
|
||||
SELECT config FROM models
|
||||
WHERE hash=?;
|
||||
""",
|
||||
(hash,),
|
||||
)
|
||||
rows = cursor.fetchone()
|
||||
if not rows:
|
||||
raise UnknownModelException("model not found")
|
||||
model = ModelConfigFactory.from_dict(json.loads(rows[0]))
|
||||
return model
|
||||
|
||||
def exists(self, key: str) -> bool:
|
||||
"""
|
||||
Return True if a model with the indicated key exists in the databse.
|
||||
|
||||
:param key: Unique key for the model to be deleted
|
||||
"""
|
||||
with self._db.transaction() as cursor:
|
||||
cursor.execute(
|
||||
"""--sql
|
||||
select count(*) FROM models
|
||||
WHERE id=?;
|
||||
""",
|
||||
(key,),
|
||||
)
|
||||
count = cursor.fetchone()[0]
|
||||
return count > 0
|
||||
|
||||
def search_by_attr(
|
||||
self,
|
||||
model_name: Optional[str] = None,
|
||||
base_model: Optional[BaseModelType] = None,
|
||||
model_type: Optional[ModelType] = None,
|
||||
model_format: Optional[ModelFormat] = None,
|
||||
order_by: ModelRecordOrderBy = ModelRecordOrderBy.Default,
|
||||
direction: SQLiteDirection = SQLiteDirection.Ascending,
|
||||
) -> List[AnyModelConfig]:
|
||||
"""
|
||||
Return models matching name, base and/or type.
|
||||
|
||||
:param model_name: Filter by name of model (optional)
|
||||
:param base_model: Filter by base model (optional)
|
||||
:param model_type: Filter by type of model (optional)
|
||||
:param model_format: Filter by model format (e.g. "diffusers") (optional)
|
||||
:param order_by: Result order
|
||||
:param direction: Result direction
|
||||
|
||||
If none of the optional filters are passed, will return all
|
||||
models in the database.
|
||||
"""
|
||||
with self._db.transaction() as cursor:
|
||||
assert isinstance(order_by, ModelRecordOrderBy)
|
||||
order_dir = "DESC" if direction == SQLiteDirection.Descending else "ASC"
|
||||
ordering = {
|
||||
ModelRecordOrderBy.Default: f"type {order_dir}, base COLLATE NOCASE {order_dir}, name COLLATE NOCASE {order_dir}, format",
|
||||
ModelRecordOrderBy.Type: "type",
|
||||
ModelRecordOrderBy.Base: "base COLLATE NOCASE",
|
||||
ModelRecordOrderBy.Name: "name COLLATE NOCASE",
|
||||
ModelRecordOrderBy.Format: "format",
|
||||
ModelRecordOrderBy.Size: "IFNULL(json_extract(config, '$.file_size'), 0)",
|
||||
ModelRecordOrderBy.DateAdded: "created_at",
|
||||
ModelRecordOrderBy.DateModified: "updated_at",
|
||||
ModelRecordOrderBy.Path: "path",
|
||||
}
|
||||
|
||||
where_clause: list[str] = []
|
||||
bindings: list[str] = []
|
||||
if model_name:
|
||||
where_clause.append("name=?")
|
||||
bindings.append(model_name)
|
||||
if base_model:
|
||||
where_clause.append("base=?")
|
||||
bindings.append(base_model)
|
||||
if model_type:
|
||||
where_clause.append("type=?")
|
||||
bindings.append(model_type)
|
||||
if model_format:
|
||||
where_clause.append("format=?")
|
||||
bindings.append(model_format)
|
||||
where = f"WHERE {' AND '.join(where_clause)}" if where_clause else ""
|
||||
|
||||
cursor.execute(
|
||||
f"""--sql
|
||||
SELECT config
|
||||
FROM models
|
||||
{where}
|
||||
ORDER BY {ordering[order_by]} {order_dir} -- using ? to bind doesn't work here for some reason;
|
||||
""",
|
||||
tuple(bindings),
|
||||
)
|
||||
result = cursor.fetchall()
|
||||
|
||||
# Parse the model configs.
|
||||
results: list[AnyModelConfig] = []
|
||||
for row in result:
|
||||
try:
|
||||
model_config = ModelConfigFactory.from_dict(json.loads(row[0]))
|
||||
except pydantic.ValidationError as e:
|
||||
# We catch this error so that the app can still run if there are invalid model configs in the database.
|
||||
# One reason that an invalid model config might be in the database is if someone had to rollback from a
|
||||
# newer version of the app that added a new model type.
|
||||
row_data = f"{row[0][:64]}..." if len(row[0]) > 64 else row[0]
|
||||
try:
|
||||
name = json.loads(row[0]).get("name", "<unknown>")
|
||||
except Exception:
|
||||
name = "<unknown>"
|
||||
self._logger.warning(
|
||||
f"Skipping invalid model config in the database with name {name}. Ignoring this model. ({row_data})"
|
||||
)
|
||||
self._logger.warning(f"Validation error: {e}")
|
||||
else:
|
||||
results.append(model_config)
|
||||
|
||||
return results
|
||||
|
||||
def search_by_path(self, path: Union[str, Path]) -> List[AnyModelConfig]:
|
||||
"""Return models with the indicated path."""
|
||||
with self._db.transaction() as cursor:
|
||||
cursor.execute(
|
||||
"""--sql
|
||||
SELECT config FROM models
|
||||
WHERE path=?;
|
||||
""",
|
||||
(str(path),),
|
||||
)
|
||||
results = [ModelConfigFactory.from_dict(json.loads(x[0])) for x in cursor.fetchall()]
|
||||
return results
|
||||
|
||||
def search_by_hash(self, hash: str) -> List[AnyModelConfig]:
|
||||
"""Return models with the indicated hash."""
|
||||
with self._db.transaction() as cursor:
|
||||
cursor.execute(
|
||||
"""--sql
|
||||
SELECT config FROM models
|
||||
WHERE hash=?;
|
||||
""",
|
||||
(hash,),
|
||||
)
|
||||
results = [ModelConfigFactory.from_dict(json.loads(x[0])) for x in cursor.fetchall()]
|
||||
return results
|
||||
|
||||
def list_models(
|
||||
self,
|
||||
page: int = 0,
|
||||
per_page: int = 10,
|
||||
order_by: ModelRecordOrderBy = ModelRecordOrderBy.Default,
|
||||
direction: SQLiteDirection = SQLiteDirection.Ascending,
|
||||
) -> PaginatedResults[ModelSummary]:
|
||||
"""Return a paginated summary listing of each model in the database."""
|
||||
with self._db.transaction() as cursor:
|
||||
assert isinstance(order_by, ModelRecordOrderBy)
|
||||
order_dir = "DESC" if direction == SQLiteDirection.Descending else "ASC"
|
||||
ordering = {
|
||||
ModelRecordOrderBy.Default: f"type {order_dir}, base COLLATE NOCASE {order_dir}, name COLLATE NOCASE {order_dir}, format",
|
||||
ModelRecordOrderBy.Type: "type",
|
||||
ModelRecordOrderBy.Base: "base COLLATE NOCASE",
|
||||
ModelRecordOrderBy.Name: "name COLLATE NOCASE",
|
||||
ModelRecordOrderBy.Format: "format",
|
||||
ModelRecordOrderBy.Size: "IFNULL(json_extract(config, '$.file_size'), 0)",
|
||||
ModelRecordOrderBy.DateAdded: "created_at",
|
||||
ModelRecordOrderBy.DateModified: "updated_at",
|
||||
ModelRecordOrderBy.Path: "path",
|
||||
}
|
||||
|
||||
# Lock so that the database isn't updated while we're doing the two queries.
|
||||
# query1: get the total number of model configs
|
||||
cursor.execute(
|
||||
"""--sql
|
||||
select count(*) from models;
|
||||
""",
|
||||
(),
|
||||
)
|
||||
total = int(cursor.fetchone()[0])
|
||||
|
||||
# query2: fetch key fields
|
||||
cursor.execute(
|
||||
f"""--sql
|
||||
SELECT config
|
||||
FROM models
|
||||
ORDER BY {ordering[order_by]} {order_dir} -- using ? to bind doesn't work here for some reason
|
||||
LIMIT ?
|
||||
OFFSET ?;
|
||||
""",
|
||||
(
|
||||
per_page,
|
||||
page * per_page,
|
||||
),
|
||||
)
|
||||
rows = cursor.fetchall()
|
||||
items = [ModelSummary.model_validate(dict(x)) for x in rows]
|
||||
return PaginatedResults(page=page, pages=ceil(total / per_page), per_page=per_page, total=total, items=items)
|
||||
@@ -0,0 +1,25 @@
|
||||
from abc import ABC, abstractmethod
|
||||
|
||||
|
||||
class ModelRelationshipRecordStorageBase(ABC):
|
||||
"""Abstract base class for model-to-model relationship record storage."""
|
||||
|
||||
@abstractmethod
|
||||
def add_model_relationship(self, model_key_1: str, model_key_2: str) -> None:
|
||||
"""Creates a relationship between two models by keys."""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def remove_model_relationship(self, model_key_1: str, model_key_2: str) -> None:
|
||||
"""Removes a relationship between two models by keys."""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def get_related_model_keys(self, model_key: str) -> list[str]:
|
||||
"""Gets all models keys related to a given model key."""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def get_related_model_keys_batch(self, model_keys: list[str]) -> list[str]:
|
||||
"""Get related model keys for multiple models given a list of keys."""
|
||||
pass
|
||||
@@ -0,0 +1,55 @@
|
||||
from invokeai.app.services.model_relationship_records.model_relationship_records_base import (
|
||||
ModelRelationshipRecordStorageBase,
|
||||
)
|
||||
from invokeai.app.services.shared.sqlite.sqlite_database import SqliteDatabase
|
||||
|
||||
|
||||
class SqliteModelRelationshipRecordStorage(ModelRelationshipRecordStorageBase):
|
||||
def __init__(self, db: SqliteDatabase) -> None:
|
||||
super().__init__()
|
||||
self._db = db
|
||||
|
||||
def add_model_relationship(self, model_key_1: str, model_key_2: str) -> None:
|
||||
with self._db.transaction() as cursor:
|
||||
if model_key_1 == model_key_2:
|
||||
raise ValueError("Cannot relate a model to itself.")
|
||||
a, b = sorted([model_key_1, model_key_2])
|
||||
cursor.execute(
|
||||
"INSERT OR IGNORE INTO model_relationships (model_key_1, model_key_2) VALUES (?, ?)",
|
||||
(a, b),
|
||||
)
|
||||
|
||||
def remove_model_relationship(self, model_key_1: str, model_key_2: str) -> None:
|
||||
with self._db.transaction() as cursor:
|
||||
a, b = sorted([model_key_1, model_key_2])
|
||||
cursor.execute(
|
||||
"DELETE FROM model_relationships WHERE model_key_1 = ? AND model_key_2 = ?",
|
||||
(a, b),
|
||||
)
|
||||
|
||||
def get_related_model_keys(self, model_key: str) -> list[str]:
|
||||
with self._db.transaction() as cursor:
|
||||
cursor.execute(
|
||||
"""
|
||||
SELECT model_key_2 FROM model_relationships WHERE model_key_1 = ?
|
||||
UNION
|
||||
SELECT model_key_1 FROM model_relationships WHERE model_key_2 = ?
|
||||
""",
|
||||
(model_key, model_key),
|
||||
)
|
||||
result = [row[0] for row in cursor.fetchall()]
|
||||
return result
|
||||
|
||||
def get_related_model_keys_batch(self, model_keys: list[str]) -> list[str]:
|
||||
with self._db.transaction() as cursor:
|
||||
key_list = ",".join("?" for _ in model_keys)
|
||||
cursor.execute(
|
||||
f"""
|
||||
SELECT model_key_2 FROM model_relationships WHERE model_key_1 IN ({key_list})
|
||||
UNION
|
||||
SELECT model_key_1 FROM model_relationships WHERE model_key_2 IN ({key_list})
|
||||
""",
|
||||
model_keys + model_keys,
|
||||
)
|
||||
result = [row[0] for row in cursor.fetchall()]
|
||||
return result
|
||||
@@ -0,0 +1,25 @@
|
||||
from abc import ABC, abstractmethod
|
||||
|
||||
|
||||
class ModelRelationshipsServiceABC(ABC):
|
||||
"""High-level service for managing model-to-model relationships."""
|
||||
|
||||
@abstractmethod
|
||||
def add_model_relationship(self, model_key_1: str, model_key_2: str) -> None:
|
||||
"""Creates a relationship between two models keys."""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def remove_model_relationship(self, model_key_1: str, model_key_2: str) -> None:
|
||||
"""Removes a relationship between two models keys."""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def get_related_model_keys(self, model_key: str) -> list[str]:
|
||||
"""Gets all models keys related to a given model key."""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def get_related_model_keys_batch(self, model_keys: list[str]) -> list[str]:
|
||||
"""Get related model keys for multiple models."""
|
||||
pass
|
||||
@@ -0,0 +1,9 @@
|
||||
from datetime import datetime
|
||||
|
||||
from invokeai.app.util.model_exclude_null import BaseModelExcludeNull
|
||||
|
||||
|
||||
class ModelRelationship(BaseModelExcludeNull):
|
||||
model_key_1: str
|
||||
model_key_2: str
|
||||
created_at: datetime
|
||||
@@ -0,0 +1,31 @@
|
||||
from invokeai.app.services.invoker import Invoker
|
||||
from invokeai.app.services.model_relationships.model_relationships_base import ModelRelationshipsServiceABC
|
||||
from invokeai.backend.model_manager.configs.factory import AnyModelConfig
|
||||
|
||||
|
||||
class ModelRelationshipsService(ModelRelationshipsServiceABC):
|
||||
__invoker: Invoker
|
||||
|
||||
def start(self, invoker: Invoker) -> None:
|
||||
self.__invoker = invoker
|
||||
|
||||
def add_model_relationship(self, model_key_1: str, model_key_2: str) -> None:
|
||||
self.__invoker.services.model_relationship_records.add_model_relationship(model_key_1, model_key_2)
|
||||
|
||||
def remove_model_relationship(self, model_key_1: str, model_key_2: str) -> None:
|
||||
self.__invoker.services.model_relationship_records.remove_model_relationship(model_key_1, model_key_2)
|
||||
|
||||
def get_related_model_keys(self, model_key: str) -> list[str]:
|
||||
return self.__invoker.services.model_relationship_records.get_related_model_keys(model_key)
|
||||
|
||||
def add_relationship_from_models(self, model_1: AnyModelConfig, model_2: AnyModelConfig) -> None:
|
||||
self.add_model_relationship(model_1.key, model_2.key)
|
||||
|
||||
def remove_relationship_from_models(self, model_1: AnyModelConfig, model_2: AnyModelConfig) -> None:
|
||||
self.remove_model_relationship(model_1.key, model_2.key)
|
||||
|
||||
def get_related_keys_from_model(self, model: AnyModelConfig) -> list[str]:
|
||||
return self.get_related_model_keys(model.key)
|
||||
|
||||
def get_related_model_keys_batch(self, model_keys: list[str]) -> list[str]:
|
||||
return self.__invoker.services.model_relationship_records.get_related_model_keys_batch(model_keys)
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user