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

This commit is contained in:
wehub-resource-sync
2026-07-13 13:22:06 +08:00
commit cddb07a176
3370 changed files with 685519 additions and 0 deletions
View File
@@ -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
+1
View File
@@ -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"
+106
View File
@@ -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,),
)
+6
View File
@@ -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
+263
View File
@@ -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))
+169
View File
@@ -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()
+37
View File
@@ -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