chore: import upstream snapshot with attribution
docs / deploy (push) Has been cancelled
docs / changes (push) Has been cancelled
docs / check-and-build (push) Has been cancelled
build container image / cpu (push) Has been cancelled
build container image / cuda (push) Has been cancelled
build container image / rocm (push) Has been cancelled
frontend checks / frontend-checks (push) Has been cancelled
frontend tests / frontend-tests (push) Has been cancelled
lfs checks / lfs-check (push) Has been cancelled
python checks / python-checks (push) Has been cancelled
python tests / py3.12: macos-default (push) Has been cancelled
python tests / py3.11: windows-cpu (push) Has been cancelled
python tests / py3.12: windows-cpu (push) Has been cancelled
python tests / py3.11: linux-cpu (push) Has been cancelled
typegen checks / typegen-checks (push) Has been cancelled
uv lock checks / uv-lock-checks (push) Has been cancelled
openapi checks / openapi-checks (push) Has been cancelled
python tests / py3.11: macos-default (push) Has been cancelled
python tests / py3.12: linux-cpu (push) Has been cancelled
docs / deploy (push) Has been cancelled
docs / changes (push) Has been cancelled
docs / check-and-build (push) Has been cancelled
build container image / cpu (push) Has been cancelled
build container image / cuda (push) Has been cancelled
build container image / rocm (push) Has been cancelled
frontend checks / frontend-checks (push) Has been cancelled
frontend tests / frontend-tests (push) Has been cancelled
lfs checks / lfs-check (push) Has been cancelled
python checks / python-checks (push) Has been cancelled
python tests / py3.12: macos-default (push) Has been cancelled
python tests / py3.11: windows-cpu (push) Has been cancelled
python tests / py3.12: windows-cpu (push) Has been cancelled
python tests / py3.11: linux-cpu (push) Has been cancelled
typegen checks / typegen-checks (push) Has been cancelled
uv lock checks / uv-lock-checks (push) Has been cancelled
openapi checks / openapi-checks (push) Has been cancelled
python tests / py3.11: macos-default (push) Has been cancelled
python tests / py3.12: linux-cpu (push) Has been cancelled
This commit is contained in:
@@ -0,0 +1,166 @@
|
||||
"""FastAPI dependencies for authentication."""
|
||||
|
||||
from typing import Annotated
|
||||
|
||||
from fastapi import Depends, HTTPException, status
|
||||
from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer
|
||||
|
||||
from invokeai.app.api.dependencies import ApiDependencies
|
||||
from invokeai.app.services.auth.token_service import TokenData, verify_token
|
||||
from invokeai.backend.util.logging import logging
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# HTTP Bearer token security scheme
|
||||
security = HTTPBearer(auto_error=False)
|
||||
|
||||
|
||||
async def get_current_user(
|
||||
credentials: Annotated[HTTPAuthorizationCredentials | None, Depends(security)],
|
||||
) -> TokenData:
|
||||
"""Get current authenticated user from Bearer token.
|
||||
|
||||
Note: This function accesses ApiDependencies.invoker.services.users directly,
|
||||
which is the established pattern in this codebase. The ApiDependencies.invoker
|
||||
is initialized in the FastAPI lifespan context before any requests are handled.
|
||||
|
||||
Args:
|
||||
credentials: The HTTP authorization credentials containing the Bearer token
|
||||
|
||||
Returns:
|
||||
TokenData containing user information from the token
|
||||
|
||||
Raises:
|
||||
HTTPException: If token is missing, invalid, or expired (401 Unauthorized)
|
||||
"""
|
||||
if credentials is None:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_401_UNAUTHORIZED,
|
||||
detail="Missing authentication credentials",
|
||||
headers={"WWW-Authenticate": "Bearer"},
|
||||
)
|
||||
|
||||
token = credentials.credentials
|
||||
token_data = verify_token(token)
|
||||
|
||||
if token_data is None:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_401_UNAUTHORIZED,
|
||||
detail="Invalid or expired authentication token",
|
||||
headers={"WWW-Authenticate": "Bearer"},
|
||||
)
|
||||
|
||||
# Verify user still exists and is active
|
||||
user_service = ApiDependencies.invoker.services.users
|
||||
user = user_service.get(token_data.user_id)
|
||||
|
||||
if user is None or not user.is_active:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_401_UNAUTHORIZED,
|
||||
detail="User account is inactive or does not exist",
|
||||
headers={"WWW-Authenticate": "Bearer"},
|
||||
)
|
||||
|
||||
return token_data
|
||||
|
||||
|
||||
async def get_current_user_or_default(
|
||||
credentials: Annotated[HTTPAuthorizationCredentials | None, Depends(security)],
|
||||
) -> TokenData:
|
||||
"""Get current authenticated user from Bearer token, or return a default system user if not authenticated.
|
||||
|
||||
This dependency is useful for endpoints that should work in both single-user and multiuser modes.
|
||||
|
||||
When multiuser mode is disabled (default), this always returns a system user with admin privileges,
|
||||
allowing unrestricted access to all operations.
|
||||
|
||||
When multiuser mode is enabled, authentication is required and this function validates the token,
|
||||
returning authenticated user data or raising 401 Unauthorized if no valid credentials are provided.
|
||||
|
||||
Args:
|
||||
credentials: The HTTP authorization credentials containing the Bearer token
|
||||
|
||||
Returns:
|
||||
TokenData containing user information from the token, or system user in single-user mode
|
||||
|
||||
Raises:
|
||||
HTTPException: 401 Unauthorized if in multiuser mode and credentials are missing, invalid, or user is inactive
|
||||
"""
|
||||
# Get configuration to check if multiuser is enabled
|
||||
config = ApiDependencies.invoker.services.configuration
|
||||
|
||||
# In single-user mode (multiuser=False), always return system user with admin privileges
|
||||
if not config.multiuser:
|
||||
return TokenData(user_id="system", email="system@system.invokeai", is_admin=True)
|
||||
|
||||
# Multiuser mode is enabled - validate credentials
|
||||
if credentials is None:
|
||||
# In multiuser mode, authentication is required
|
||||
raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="Authentication required")
|
||||
|
||||
token = credentials.credentials
|
||||
token_data = verify_token(token)
|
||||
|
||||
if token_data is None:
|
||||
# Invalid token in multiuser mode - reject
|
||||
raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="Invalid or expired token")
|
||||
|
||||
# Verify user still exists and is active
|
||||
user_service = ApiDependencies.invoker.services.users
|
||||
user = user_service.get(token_data.user_id)
|
||||
|
||||
if user is None or not user.is_active:
|
||||
# User doesn't exist or is inactive in multiuser mode - reject
|
||||
raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="User not found or inactive")
|
||||
|
||||
return token_data
|
||||
|
||||
|
||||
async def require_admin(
|
||||
current_user: Annotated[TokenData, Depends(get_current_user)],
|
||||
) -> TokenData:
|
||||
"""Require admin role for the current user.
|
||||
|
||||
Args:
|
||||
current_user: The current authenticated user's token data
|
||||
|
||||
Returns:
|
||||
The token data if user is an admin
|
||||
|
||||
Raises:
|
||||
HTTPException: If user does not have admin privileges (403 Forbidden)
|
||||
"""
|
||||
if not current_user.is_admin:
|
||||
raise HTTPException(status_code=status.HTTP_403_FORBIDDEN, detail="Admin privileges required")
|
||||
return current_user
|
||||
|
||||
|
||||
async def require_admin_or_default(
|
||||
current_user: Annotated[TokenData, Depends(get_current_user_or_default)],
|
||||
) -> TokenData:
|
||||
"""Require admin role for the current user, or return default system admin in single-user mode.
|
||||
|
||||
This dependency is useful for admin-only endpoints that should work in both single-user and multiuser modes.
|
||||
|
||||
When multiuser mode is disabled (default), this always returns a system user with admin privileges.
|
||||
When multiuser mode is enabled, this validates that the authenticated user has admin privileges.
|
||||
|
||||
Args:
|
||||
current_user: The current authenticated user's token data (or default system user)
|
||||
|
||||
Returns:
|
||||
The token data if user is an admin (or system user in single-user mode)
|
||||
|
||||
Raises:
|
||||
HTTPException: If user does not have admin privileges (403 Forbidden) in multiuser mode
|
||||
"""
|
||||
if not current_user.is_admin:
|
||||
raise HTTPException(status_code=status.HTTP_403_FORBIDDEN, detail="Admin privileges required")
|
||||
return current_user
|
||||
|
||||
|
||||
# Type aliases for convenient use in route dependencies
|
||||
CurrentUser = Annotated[TokenData, Depends(get_current_user)]
|
||||
CurrentUserOrDefault = Annotated[TokenData, Depends(get_current_user_or_default)]
|
||||
AdminUser = Annotated[TokenData, Depends(require_admin)]
|
||||
AdminUserOrDefault = Annotated[TokenData, Depends(require_admin_or_default)]
|
||||
@@ -0,0 +1,245 @@
|
||||
# Copyright (c) 2022 Kyle Schouviller (https://github.com/kyle0654)
|
||||
|
||||
import asyncio
|
||||
from logging import Logger
|
||||
|
||||
import torch
|
||||
|
||||
from invokeai.app.services.app_settings import AppSettingsService
|
||||
from invokeai.app.services.auth.token_service import set_jwt_secret
|
||||
from invokeai.app.services.board_image_records.board_image_records_sqlite import SqliteBoardImageRecordStorage
|
||||
from invokeai.app.services.board_images.board_images_default import BoardImagesService
|
||||
from invokeai.app.services.board_records.board_records_sqlite import SqliteBoardRecordStorage
|
||||
from invokeai.app.services.boards.boards_default import BoardService
|
||||
from invokeai.app.services.bulk_download.bulk_download_default import BulkDownloadService
|
||||
from invokeai.app.services.client_state_persistence.client_state_persistence_sqlite import ClientStatePersistenceSqlite
|
||||
from invokeai.app.services.config.config_default import InvokeAIAppConfig
|
||||
from invokeai.app.services.download.download_default import DownloadQueueService
|
||||
from invokeai.app.services.events.events_fastapievents import FastAPIEventService
|
||||
from invokeai.app.services.external_generation.external_generation_default import ExternalGenerationService
|
||||
from invokeai.app.services.external_generation.providers import (
|
||||
AlibabaCloudProvider,
|
||||
GeminiProvider,
|
||||
OpenAIProvider,
|
||||
SeedreamProvider,
|
||||
)
|
||||
from invokeai.app.services.external_generation.startup import sync_configured_external_starter_models
|
||||
from invokeai.app.services.image_files.image_files_disk import DiskImageFileStorage
|
||||
from invokeai.app.services.image_moves.image_moves_default import ImageMoveService
|
||||
from invokeai.app.services.image_records.image_records_sqlite import SqliteImageRecordStorage
|
||||
from invokeai.app.services.images.images_default import ImageService
|
||||
from invokeai.app.services.invocation_cache.invocation_cache_memory import MemoryInvocationCache
|
||||
from invokeai.app.services.invocation_services import InvocationServices
|
||||
from invokeai.app.services.invocation_stats.invocation_stats_default import InvocationStatsService
|
||||
from invokeai.app.services.invoker import Invoker
|
||||
from invokeai.app.services.model_images.model_images_default import ModelImageFileStorageDisk
|
||||
from invokeai.app.services.model_manager.model_manager_default import ModelManagerService
|
||||
from invokeai.app.services.model_records.model_records_sql import ModelRecordServiceSQL
|
||||
from invokeai.app.services.model_relationship_records.model_relationship_records_sqlite import (
|
||||
SqliteModelRelationshipRecordStorage,
|
||||
)
|
||||
from invokeai.app.services.model_relationships.model_relationships_default import ModelRelationshipsService
|
||||
from invokeai.app.services.names.names_default import SimpleNameService
|
||||
from invokeai.app.services.object_serializer.object_serializer_disk import ObjectSerializerDisk
|
||||
from invokeai.app.services.object_serializer.object_serializer_forward_cache import ObjectSerializerForwardCache
|
||||
from invokeai.app.services.session_processor.session_processor_default import (
|
||||
DefaultSessionProcessor,
|
||||
DefaultSessionRunner,
|
||||
)
|
||||
from invokeai.app.services.session_queue.session_queue_sqlite import SqliteSessionQueue
|
||||
from invokeai.app.services.shared.sqlite.sqlite_util import init_db
|
||||
from invokeai.app.services.style_preset_images.style_preset_images_disk import StylePresetImageFileStorageDisk
|
||||
from invokeai.app.services.style_preset_records.style_preset_records_sqlite import SqliteStylePresetRecordsStorage
|
||||
from invokeai.app.services.urls.urls_default import LocalUrlService
|
||||
from invokeai.app.services.users.users_default import UserService
|
||||
from invokeai.app.services.workflow_records.workflow_records_sqlite import SqliteWorkflowRecordsStorage
|
||||
from invokeai.app.services.workflow_thumbnails.workflow_thumbnails_disk import WorkflowThumbnailFileStorageDisk
|
||||
from invokeai.backend.stable_diffusion.diffusion.conditioning_data import (
|
||||
AnimaConditioningInfo,
|
||||
BasicConditioningInfo,
|
||||
CogView4ConditioningInfo,
|
||||
ConditioningFieldData,
|
||||
FLUXConditioningInfo,
|
||||
QwenImageConditioningInfo,
|
||||
SD3ConditioningInfo,
|
||||
SDXLConditioningInfo,
|
||||
ZImageConditioningInfo,
|
||||
)
|
||||
from invokeai.backend.util.logging import InvokeAILogger
|
||||
from invokeai.version.invokeai_version import __version__
|
||||
|
||||
|
||||
# TODO: is there a better way to achieve this?
|
||||
def check_internet() -> bool:
|
||||
"""
|
||||
Return true if the internet is reachable.
|
||||
It does this by pinging huggingface.co.
|
||||
"""
|
||||
import urllib.request
|
||||
|
||||
host = "http://huggingface.co"
|
||||
try:
|
||||
urllib.request.urlopen(host, timeout=1)
|
||||
return True
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
|
||||
logger = InvokeAILogger.get_logger()
|
||||
|
||||
|
||||
class ApiDependencies:
|
||||
"""Contains and initializes all dependencies for the API"""
|
||||
|
||||
invoker: Invoker
|
||||
|
||||
@staticmethod
|
||||
def initialize(
|
||||
config: InvokeAIAppConfig,
|
||||
event_handler_id: int,
|
||||
loop: asyncio.AbstractEventLoop,
|
||||
logger: Logger = logger,
|
||||
) -> None:
|
||||
logger.info(f"InvokeAI version {__version__}")
|
||||
logger.info(f"Root directory = {str(config.root_path)}")
|
||||
|
||||
output_folder = config.outputs_path
|
||||
if output_folder is None:
|
||||
raise ValueError("Output folder is not set")
|
||||
|
||||
image_files = DiskImageFileStorage(f"{output_folder}/images")
|
||||
|
||||
model_images_folder = config.models_path
|
||||
style_presets_folder = config.style_presets_path
|
||||
workflow_thumbnails_folder = config.workflow_thumbnails_path
|
||||
|
||||
db = init_db(config=config, logger=logger, image_files=image_files)
|
||||
|
||||
# Initialize JWT secret from database
|
||||
app_settings = AppSettingsService(db=db)
|
||||
jwt_secret = app_settings.get_jwt_secret()
|
||||
set_jwt_secret(jwt_secret)
|
||||
logger.info("JWT secret loaded from database")
|
||||
|
||||
configuration = config
|
||||
logger = logger
|
||||
|
||||
board_image_records = SqliteBoardImageRecordStorage(db=db)
|
||||
board_images = BoardImagesService()
|
||||
board_records = SqliteBoardRecordStorage(db=db)
|
||||
boards = BoardService()
|
||||
events = FastAPIEventService(event_handler_id, loop=loop)
|
||||
bulk_download = BulkDownloadService()
|
||||
image_records = SqliteImageRecordStorage(db=db)
|
||||
image_moves = ImageMoveService(db=db, image_files=image_files, config=configuration, logger=logger)
|
||||
images = ImageService()
|
||||
invocation_cache = MemoryInvocationCache(max_cache_size=config.node_cache_size)
|
||||
tensors = ObjectSerializerForwardCache(
|
||||
ObjectSerializerDisk[torch.Tensor](
|
||||
output_folder / "tensors",
|
||||
safe_globals=[torch.Tensor],
|
||||
ephemeral=True,
|
||||
),
|
||||
)
|
||||
conditioning = ObjectSerializerForwardCache(
|
||||
ObjectSerializerDisk[ConditioningFieldData](
|
||||
output_folder / "conditioning",
|
||||
safe_globals=[
|
||||
ConditioningFieldData,
|
||||
BasicConditioningInfo,
|
||||
SDXLConditioningInfo,
|
||||
FLUXConditioningInfo,
|
||||
SD3ConditioningInfo,
|
||||
CogView4ConditioningInfo,
|
||||
ZImageConditioningInfo,
|
||||
QwenImageConditioningInfo,
|
||||
AnimaConditioningInfo,
|
||||
],
|
||||
ephemeral=True,
|
||||
),
|
||||
)
|
||||
download_queue_service = DownloadQueueService(app_config=configuration, event_bus=events)
|
||||
model_record_service = ModelRecordServiceSQL(db=db, logger=logger)
|
||||
model_manager = ModelManagerService.build_model_manager(
|
||||
app_config=configuration,
|
||||
model_record_service=model_record_service,
|
||||
download_queue=download_queue_service,
|
||||
events=events,
|
||||
)
|
||||
external_generation = ExternalGenerationService(
|
||||
providers={
|
||||
AlibabaCloudProvider.provider_id: AlibabaCloudProvider(app_config=configuration, logger=logger),
|
||||
GeminiProvider.provider_id: GeminiProvider(app_config=configuration, logger=logger),
|
||||
OpenAIProvider.provider_id: OpenAIProvider(app_config=configuration, logger=logger),
|
||||
SeedreamProvider.provider_id: SeedreamProvider(app_config=configuration, logger=logger),
|
||||
},
|
||||
logger=logger,
|
||||
record_store=model_record_service,
|
||||
)
|
||||
model_images_service = ModelImageFileStorageDisk(model_images_folder / "model_images")
|
||||
model_relationships = ModelRelationshipsService()
|
||||
model_relationship_records = SqliteModelRelationshipRecordStorage(db=db)
|
||||
names = SimpleNameService()
|
||||
performance_statistics = InvocationStatsService()
|
||||
session_processor = DefaultSessionProcessor(session_runner=DefaultSessionRunner())
|
||||
session_queue = SqliteSessionQueue(db=db)
|
||||
urls = LocalUrlService()
|
||||
workflow_records = SqliteWorkflowRecordsStorage(db=db)
|
||||
style_preset_records = SqliteStylePresetRecordsStorage(db=db)
|
||||
style_preset_image_files = StylePresetImageFileStorageDisk(style_presets_folder / "images")
|
||||
workflow_thumbnails = WorkflowThumbnailFileStorageDisk(workflow_thumbnails_folder)
|
||||
client_state_persistence = ClientStatePersistenceSqlite(db=db)
|
||||
users = UserService(db=db)
|
||||
|
||||
services = InvocationServices(
|
||||
board_image_records=board_image_records,
|
||||
board_images=board_images,
|
||||
board_records=board_records,
|
||||
boards=boards,
|
||||
bulk_download=bulk_download,
|
||||
configuration=configuration,
|
||||
events=events,
|
||||
image_files=image_files,
|
||||
image_moves=image_moves,
|
||||
image_records=image_records,
|
||||
images=images,
|
||||
invocation_cache=invocation_cache,
|
||||
logger=logger,
|
||||
model_images=model_images_service,
|
||||
model_manager=model_manager,
|
||||
model_relationships=model_relationships,
|
||||
model_relationship_records=model_relationship_records,
|
||||
download_queue=download_queue_service,
|
||||
external_generation=external_generation,
|
||||
names=names,
|
||||
performance_statistics=performance_statistics,
|
||||
session_processor=session_processor,
|
||||
session_queue=session_queue,
|
||||
urls=urls,
|
||||
workflow_records=workflow_records,
|
||||
tensors=tensors,
|
||||
conditioning=conditioning,
|
||||
style_preset_records=style_preset_records,
|
||||
style_preset_image_files=style_preset_image_files,
|
||||
workflow_thumbnails=workflow_thumbnails,
|
||||
client_state_persistence=client_state_persistence,
|
||||
users=users,
|
||||
)
|
||||
|
||||
ApiDependencies.invoker = Invoker(services)
|
||||
configured_external_providers = {
|
||||
provider_id
|
||||
for provider_id, status in external_generation.get_provider_statuses().items()
|
||||
if status.configured
|
||||
}
|
||||
sync_configured_external_starter_models(
|
||||
configured_provider_ids=configured_external_providers,
|
||||
model_manager=model_manager,
|
||||
logger=logger,
|
||||
)
|
||||
db.clean()
|
||||
|
||||
@staticmethod
|
||||
def shutdown() -> None:
|
||||
if ApiDependencies.invoker:
|
||||
ApiDependencies.invoker.stop()
|
||||
@@ -0,0 +1,134 @@
|
||||
import json
|
||||
import logging
|
||||
from dataclasses import dataclass
|
||||
|
||||
from PIL import Image
|
||||
|
||||
from invokeai.app.services.workflow_records.workflow_records_common import WorkflowWithoutIDValidator
|
||||
|
||||
|
||||
@dataclass
|
||||
class ExtractedMetadata:
|
||||
invokeai_metadata: str | None
|
||||
invokeai_workflow: str | None
|
||||
invokeai_graph: str | None
|
||||
|
||||
|
||||
def extract_metadata_from_image(
|
||||
pil_image: Image.Image,
|
||||
invokeai_metadata_override: str | None,
|
||||
invokeai_workflow_override: str | None,
|
||||
invokeai_graph_override: str | None,
|
||||
logger: logging.Logger,
|
||||
) -> ExtractedMetadata:
|
||||
"""
|
||||
Extracts the "invokeai_metadata", "invokeai_workflow", and "invokeai_graph" data embedded in the PIL Image.
|
||||
|
||||
These items are stored as stringified JSON in the image file's metadata, so we need to do some parsing to validate
|
||||
them. Once parsed, the values are returned as they came (as strings), or None if they are not present or invalid.
|
||||
|
||||
In some situations, we may prefer to override the values extracted from the image file with some other values.
|
||||
|
||||
For example, when uploading an image via API, the client can optionally provide the metadata directly in the request,
|
||||
as opposed to embedding it in the image file. In this case, the client-provided metadata will be used instead of the
|
||||
metadata embedded in the image file.
|
||||
|
||||
Args:
|
||||
pil_image: The PIL Image object.
|
||||
invokeai_metadata_override: The metadata override provided by the client.
|
||||
invokeai_workflow_override: The workflow override provided by the client.
|
||||
invokeai_graph_override: The graph override provided by the client.
|
||||
logger: The logger to use for debug logging.
|
||||
|
||||
Returns:
|
||||
ExtractedMetadata: The extracted metadata, workflow, and graph.
|
||||
"""
|
||||
|
||||
# The fallback value for metadata is None.
|
||||
stringified_metadata: str | None = None
|
||||
|
||||
# Use the metadata override if provided, else attempt to extract it from the image file.
|
||||
metadata_raw = (
|
||||
invokeai_metadata_override
|
||||
if invokeai_metadata_override is not None
|
||||
else pil_image.info.get("invokeai_metadata", None)
|
||||
)
|
||||
|
||||
# If the metadata is present in the image file, we will attempt to parse it as JSON. When we create images,
|
||||
# we always store metadata as a stringified JSON dict. So, we expect it to be a string here.
|
||||
if isinstance(metadata_raw, str):
|
||||
try:
|
||||
# Must be a JSON string
|
||||
metadata_parsed = json.loads(metadata_raw)
|
||||
# Must be a dict
|
||||
if isinstance(metadata_parsed, dict):
|
||||
# Looks good, overwrite the fallback value
|
||||
stringified_metadata = metadata_raw
|
||||
except Exception as e:
|
||||
logger.debug(f"Failed to parse metadata for uploaded image, {e}")
|
||||
pass
|
||||
|
||||
# We expect the workflow, if embedded in the image, to be a JSON-stringified WorkflowWithoutID. We will store it
|
||||
# as a string.
|
||||
workflow_raw: str | None = (
|
||||
invokeai_workflow_override
|
||||
if invokeai_workflow_override is not None
|
||||
else pil_image.info.get("invokeai_workflow", None)
|
||||
)
|
||||
|
||||
# The fallback value for workflow is None.
|
||||
stringified_workflow: str | None = None
|
||||
|
||||
# If the workflow is present in the image file, we will attempt to parse it as JSON. When we create images, we
|
||||
# always store workflows as a stringified JSON WorkflowWithoutID. So, we expect it to be a string here.
|
||||
if isinstance(workflow_raw, str):
|
||||
try:
|
||||
# Validate the workflow JSON before storing it
|
||||
WorkflowWithoutIDValidator.validate_json(workflow_raw)
|
||||
# Looks good, overwrite the fallback value
|
||||
stringified_workflow = workflow_raw
|
||||
except Exception:
|
||||
logger.debug("Failed to parse workflow for uploaded image")
|
||||
pass
|
||||
|
||||
# We expect the workflow, if embedded in the image, to be a JSON-stringified Graph. We will store it as a
|
||||
# string.
|
||||
graph_raw: str | None = (
|
||||
invokeai_graph_override if invokeai_graph_override is not None else pil_image.info.get("invokeai_graph", None)
|
||||
)
|
||||
|
||||
# The fallback value for graph is None.
|
||||
stringified_graph: str | None = None
|
||||
|
||||
# If the graph is present in the image file, we will attempt to parse it as JSON. When we create images, we
|
||||
# always store graphs as a stringified JSON Graph. So, we expect it to be a string here.
|
||||
if isinstance(graph_raw, str):
|
||||
try:
|
||||
# TODO(psyche): Due to pydantic's handling of None values, it is possible for the graph to fail validation,
|
||||
# even if it is a direct dump of a valid graph. Node fields in the graph are allowed to have be unset if
|
||||
# they have incoming connections, but something about the ser/de process cannot adequately handle this.
|
||||
#
|
||||
# In lieu of fixing the graph validation, we will just do a simple check here to see if the graph is dict
|
||||
# with the correct keys. This is not a perfect solution, but it should be good enough for now.
|
||||
|
||||
# FIX ME: Validate the graph JSON before storing it
|
||||
# Graph.model_validate_json(graph_raw)
|
||||
|
||||
# Crappy workaround to validate JSON
|
||||
graph_parsed = json.loads(graph_raw)
|
||||
if not isinstance(graph_parsed, dict):
|
||||
raise ValueError("Not a dict")
|
||||
if not isinstance(graph_parsed.get("nodes", None), dict):
|
||||
raise ValueError("'nodes' is not a dict")
|
||||
if not isinstance(graph_parsed.get("edges", None), list):
|
||||
raise ValueError("'edges' is not a list")
|
||||
|
||||
# Looks good, overwrite the fallback value
|
||||
stringified_graph = graph_raw
|
||||
except Exception as e:
|
||||
logger.debug(f"Failed to parse graph for uploaded image, {e}")
|
||||
pass
|
||||
|
||||
return ExtractedMetadata(
|
||||
invokeai_metadata=stringified_metadata, invokeai_workflow=stringified_workflow, invokeai_graph=stringified_graph
|
||||
)
|
||||
@@ -0,0 +1,50 @@
|
||||
from typing import Any
|
||||
|
||||
from starlette.exceptions import HTTPException
|
||||
from starlette.responses import Response
|
||||
from starlette.staticfiles import StaticFiles
|
||||
from starlette.types import Scope
|
||||
|
||||
|
||||
class NoCacheStaticFiles(StaticFiles):
|
||||
"""
|
||||
This class is used to override the default caching behavior of starlette for static files,
|
||||
ensuring we *never* cache static files. It modifies the file response headers to strictly
|
||||
never cache the files.
|
||||
|
||||
Static files include the javascript bundles, fonts, locales, and some images. Generated
|
||||
images are not included, as they are served by a router.
|
||||
|
||||
This class also implements proper SPA (Single Page Application) routing by serving index.html
|
||||
for any routes that don't match static files, enabling client-side routing to work correctly
|
||||
in production builds.
|
||||
"""
|
||||
|
||||
def __init__(self, *args: Any, **kwargs: Any):
|
||||
self.cachecontrol = "max-age=0, no-cache, no-store, , must-revalidate"
|
||||
self.pragma = "no-cache"
|
||||
self.expires = "0"
|
||||
super().__init__(*args, **kwargs)
|
||||
|
||||
def file_response(self, *args: Any, **kwargs: Any) -> Response:
|
||||
resp = super().file_response(*args, **kwargs)
|
||||
resp.headers.setdefault("Cache-Control", self.cachecontrol)
|
||||
resp.headers.setdefault("Pragma", self.pragma)
|
||||
resp.headers.setdefault("Expires", self.expires)
|
||||
return resp
|
||||
|
||||
async def get_response(self, path: str, scope: Scope) -> Response:
|
||||
"""
|
||||
Override get_response to implement SPA routing.
|
||||
|
||||
When a file is not found and html mode is enabled, serve index.html instead of raising a 404.
|
||||
This allows client-side routing to work correctly in SPAs.
|
||||
"""
|
||||
try:
|
||||
return await super().get_response(path, scope)
|
||||
except HTTPException as exc:
|
||||
# If the file is not found (404) and html mode is enabled, serve index.html
|
||||
# This allows client-side routing to handle the path
|
||||
if exc.status_code == 404 and self.html:
|
||||
return await super().get_response("index.html", scope)
|
||||
raise
|
||||
@@ -0,0 +1,92 @@
|
||||
"""Cross-router authorization helpers.
|
||||
|
||||
These helpers are imported by multiple router modules. Keep them free of router
|
||||
specifics so any route can call them after resolving `current_user`.
|
||||
"""
|
||||
|
||||
from fastapi import HTTPException
|
||||
|
||||
from invokeai.app.api.auth_dependencies import CurrentUserOrDefault
|
||||
from invokeai.app.api.dependencies import ApiDependencies
|
||||
from invokeai.app.services.board_records.board_records_common import BoardVisibility
|
||||
|
||||
|
||||
def assert_image_owner(image_name: str, current_user: CurrentUserOrDefault) -> None:
|
||||
"""Raise 403 if the current user does not own the image and is not an admin.
|
||||
|
||||
Ownership is satisfied when ANY of these hold:
|
||||
- The user is an admin.
|
||||
- The user is the image's direct owner (image_records.user_id).
|
||||
- The user owns the board the image sits on.
|
||||
- The image sits on a Public board (public boards grant mutation rights).
|
||||
"""
|
||||
if current_user.is_admin:
|
||||
return
|
||||
owner = ApiDependencies.invoker.services.image_records.get_user_id(image_name)
|
||||
if owner is not None and owner == current_user.user_id:
|
||||
return
|
||||
|
||||
board_id = ApiDependencies.invoker.services.board_image_records.get_board_for_image(image_name)
|
||||
if board_id is not None:
|
||||
try:
|
||||
board = ApiDependencies.invoker.services.boards.get_dto(board_id=board_id)
|
||||
if board.user_id == current_user.user_id:
|
||||
return
|
||||
if board.board_visibility == BoardVisibility.Public:
|
||||
return
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
raise HTTPException(status_code=403, detail="Not authorized to modify this image")
|
||||
|
||||
|
||||
def assert_image_read_access(image_name: str, current_user: CurrentUserOrDefault) -> None:
|
||||
"""Raise 403 if the current user may not view the image.
|
||||
|
||||
Access is granted when ANY of these hold:
|
||||
- The user is an admin.
|
||||
- The user owns the image.
|
||||
- The image sits on a shared or public board.
|
||||
"""
|
||||
if current_user.is_admin:
|
||||
return
|
||||
|
||||
owner = ApiDependencies.invoker.services.image_records.get_user_id(image_name)
|
||||
if owner is not None and owner == current_user.user_id:
|
||||
return
|
||||
|
||||
board_id = ApiDependencies.invoker.services.board_image_records.get_board_for_image(image_name)
|
||||
if board_id is not None:
|
||||
try:
|
||||
board = ApiDependencies.invoker.services.boards.get_dto(board_id=board_id)
|
||||
if board.board_visibility in (BoardVisibility.Shared, BoardVisibility.Public):
|
||||
return
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
raise HTTPException(status_code=403, detail="Not authorized to access this image")
|
||||
|
||||
|
||||
def assert_board_read_access(board_id: str, current_user: CurrentUserOrDefault) -> None:
|
||||
"""Raise 403 if the current user may not read images from this board.
|
||||
|
||||
Access is granted when ANY of these hold:
|
||||
- The user is an admin.
|
||||
- The user owns the board.
|
||||
- The board visibility is Shared or Public.
|
||||
"""
|
||||
if current_user.is_admin:
|
||||
return
|
||||
|
||||
try:
|
||||
board = ApiDependencies.invoker.services.boards.get_dto(board_id=board_id)
|
||||
except Exception:
|
||||
raise HTTPException(status_code=404, detail="Board not found")
|
||||
|
||||
if board.user_id == current_user.user_id:
|
||||
return
|
||||
|
||||
if board.board_visibility in (BoardVisibility.Shared, BoardVisibility.Public):
|
||||
return
|
||||
|
||||
raise HTTPException(status_code=403, detail="Not authorized to access this board")
|
||||
@@ -0,0 +1,399 @@
|
||||
import locale
|
||||
from enum import Enum
|
||||
from importlib.metadata import distributions
|
||||
from pathlib import Path as FilePath
|
||||
from threading import Lock
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
import yaml
|
||||
from fastapi import Body, HTTPException, Path
|
||||
from fastapi.routing import APIRouter
|
||||
from pydantic import BaseModel, Field, model_validator
|
||||
|
||||
from invokeai.app.api.auth_dependencies import AdminUserOrDefault
|
||||
from invokeai.app.api.dependencies import ApiDependencies
|
||||
from invokeai.app.services.config.config_default import (
|
||||
EXTERNAL_PROVIDER_CONFIG_FIELDS,
|
||||
IMAGE_SUBFOLDER_STRATEGY,
|
||||
DefaultInvokeAIAppConfig,
|
||||
InvokeAIAppConfig,
|
||||
get_config,
|
||||
load_and_migrate_config,
|
||||
load_external_api_keys,
|
||||
)
|
||||
from invokeai.app.services.external_generation.external_generation_common import ExternalProviderStatus
|
||||
from invokeai.app.services.invocation_cache.invocation_cache_common import InvocationCacheStatus
|
||||
from invokeai.app.services.model_records.model_records_base import UnknownModelException
|
||||
from invokeai.backend.image_util.infill_methods.patchmatch import PatchMatch
|
||||
from invokeai.backend.model_manager.taxonomy import BaseModelType, ModelType
|
||||
from invokeai.backend.util.logging import logging
|
||||
from invokeai.version import __version__
|
||||
|
||||
|
||||
class LogLevel(int, Enum):
|
||||
NotSet = logging.NOTSET
|
||||
Debug = logging.DEBUG
|
||||
Info = logging.INFO
|
||||
Warning = logging.WARNING
|
||||
Error = logging.ERROR
|
||||
Critical = logging.CRITICAL
|
||||
|
||||
|
||||
app_router = APIRouter(prefix="/v1/app", tags=["app"])
|
||||
|
||||
|
||||
class AppVersion(BaseModel):
|
||||
"""App Version Response"""
|
||||
|
||||
version: str = Field(description="App version")
|
||||
|
||||
|
||||
@app_router.get("/version", operation_id="app_version", status_code=200, response_model=AppVersion)
|
||||
async def get_version() -> AppVersion:
|
||||
return AppVersion(version=__version__)
|
||||
|
||||
|
||||
@app_router.get("/app_deps", operation_id="get_app_deps", status_code=200, response_model=dict[str, str])
|
||||
async def get_app_deps() -> dict[str, str]:
|
||||
deps: dict[str, str] = {dist.metadata["Name"]: dist.version for dist in distributions()}
|
||||
try:
|
||||
cuda = getattr(getattr(torch, "version", None), "cuda", None) or "N/A" # pyright: ignore[reportAttributeAccessIssue]
|
||||
except Exception:
|
||||
cuda = "N/A"
|
||||
|
||||
deps["CUDA"] = cuda
|
||||
|
||||
sorted_deps = dict(sorted(deps.items(), key=lambda item: item[0].lower()))
|
||||
|
||||
return sorted_deps
|
||||
|
||||
|
||||
@app_router.get("/patchmatch_status", operation_id="get_patchmatch_status", status_code=200, response_model=bool)
|
||||
async def get_patchmatch_status() -> bool:
|
||||
return PatchMatch.patchmatch_available()
|
||||
|
||||
|
||||
class InvokeAIAppConfigWithSetFields(BaseModel):
|
||||
"""InvokeAI App Config with model fields set"""
|
||||
|
||||
set_fields: set[str] = Field(description="The set fields")
|
||||
config: InvokeAIAppConfig = Field(description="The InvokeAI App Config")
|
||||
|
||||
|
||||
class ExternalProviderStatusModel(BaseModel):
|
||||
provider_id: str = Field(description="The external provider identifier")
|
||||
configured: bool = Field(description="Whether credentials are configured for the provider")
|
||||
message: str | None = Field(default=None, description="Optional provider status detail")
|
||||
|
||||
|
||||
class ExternalProviderConfigUpdate(BaseModel):
|
||||
api_key: str | None = Field(default=None, description="API key for the external provider")
|
||||
base_url: str | None = Field(default=None, description="Optional base URL override for the provider")
|
||||
|
||||
|
||||
class ExternalProviderConfigModel(BaseModel):
|
||||
provider_id: str = Field(description="The external provider identifier")
|
||||
api_key_configured: bool = Field(description="Whether an API key is configured")
|
||||
base_url: str | None = Field(default=None, description="Optional base URL override")
|
||||
|
||||
|
||||
EXTERNAL_PROVIDER_FIELDS: dict[str, tuple[str, str]] = {
|
||||
"alibabacloud": ("external_alibabacloud_api_key", "external_alibabacloud_base_url"),
|
||||
"gemini": ("external_gemini_api_key", "external_gemini_base_url"),
|
||||
"openai": ("external_openai_api_key", "external_openai_base_url"),
|
||||
"seedream": ("external_seedream_api_key", "external_seedream_base_url"),
|
||||
}
|
||||
_EXTERNAL_PROVIDER_CONFIG_LOCK = Lock()
|
||||
|
||||
|
||||
def _remove_nullable_default_from_schema(schema: dict[str, Any]) -> None:
|
||||
schema.pop("default", None)
|
||||
any_of = schema.pop("anyOf", None)
|
||||
if isinstance(any_of, list):
|
||||
non_null_schemas = [
|
||||
subschema for subschema in any_of if isinstance(subschema, dict) and subschema.get("type") != "null"
|
||||
]
|
||||
if len(non_null_schemas) == 1:
|
||||
schema.update(non_null_schemas[0])
|
||||
|
||||
|
||||
class UpdateAppGenerationSettingsRequest(BaseModel):
|
||||
"""Writable generation-related app settings."""
|
||||
|
||||
image_subfolder_strategy: IMAGE_SUBFOLDER_STRATEGY | None = Field(
|
||||
default=None,
|
||||
description="Strategy for organizing images into subfolders.",
|
||||
json_schema_extra=_remove_nullable_default_from_schema,
|
||||
)
|
||||
max_queue_history: int | None = Field(
|
||||
default=None,
|
||||
ge=0,
|
||||
description="Keep the last N completed, failed, and canceled queue items on startup. Set to 0 to prune all terminal items.",
|
||||
)
|
||||
|
||||
@model_validator(mode="after")
|
||||
def validate_explicit_nulls(self) -> "UpdateAppGenerationSettingsRequest":
|
||||
if "image_subfolder_strategy" in self.model_fields_set and self.image_subfolder_strategy is None:
|
||||
raise ValueError("image_subfolder_strategy may not be null")
|
||||
return self
|
||||
|
||||
|
||||
@app_router.get(
|
||||
"/runtime_config", operation_id="get_runtime_config", status_code=200, response_model=InvokeAIAppConfigWithSetFields
|
||||
)
|
||||
async def get_runtime_config() -> InvokeAIAppConfigWithSetFields:
|
||||
config = get_config()
|
||||
return InvokeAIAppConfigWithSetFields(set_fields=config.model_fields_set, config=config)
|
||||
|
||||
|
||||
@app_router.patch(
|
||||
"/runtime_config",
|
||||
operation_id="update_runtime_config",
|
||||
status_code=200,
|
||||
response_model=InvokeAIAppConfigWithSetFields,
|
||||
)
|
||||
async def update_runtime_config(
|
||||
_: AdminUserOrDefault,
|
||||
changes: UpdateAppGenerationSettingsRequest = Body(description="Writable runtime configuration changes"),
|
||||
) -> InvokeAIAppConfigWithSetFields:
|
||||
with _EXTERNAL_PROVIDER_CONFIG_LOCK:
|
||||
config = get_config()
|
||||
update_dict = changes.model_dump(exclude_unset=True)
|
||||
config.update_config(update_dict)
|
||||
|
||||
if config.config_file_path.exists():
|
||||
persisted_config = load_and_migrate_config(config.config_file_path)
|
||||
else:
|
||||
persisted_config = DefaultInvokeAIAppConfig()
|
||||
|
||||
persisted_config.update_config(update_dict)
|
||||
persisted_config.write_file(config.config_file_path)
|
||||
return InvokeAIAppConfigWithSetFields(set_fields=config.model_fields_set, config=config)
|
||||
|
||||
|
||||
@app_router.get(
|
||||
"/external_providers/status",
|
||||
operation_id="get_external_provider_statuses",
|
||||
status_code=200,
|
||||
response_model=list[ExternalProviderStatusModel],
|
||||
)
|
||||
async def get_external_provider_statuses() -> list[ExternalProviderStatusModel]:
|
||||
statuses = ApiDependencies.invoker.services.external_generation.get_provider_statuses()
|
||||
return [status_to_model(status) for status in statuses.values()]
|
||||
|
||||
|
||||
@app_router.get(
|
||||
"/external_providers/config",
|
||||
operation_id="get_external_provider_configs",
|
||||
status_code=200,
|
||||
response_model=list[ExternalProviderConfigModel],
|
||||
)
|
||||
async def get_external_provider_configs() -> list[ExternalProviderConfigModel]:
|
||||
config = get_config()
|
||||
return [_build_external_provider_config(provider_id, config) for provider_id in EXTERNAL_PROVIDER_FIELDS]
|
||||
|
||||
|
||||
@app_router.post(
|
||||
"/external_providers/config/{provider_id}",
|
||||
operation_id="set_external_provider_config",
|
||||
status_code=200,
|
||||
response_model=ExternalProviderConfigModel,
|
||||
)
|
||||
async def set_external_provider_config(
|
||||
_: AdminUserOrDefault,
|
||||
provider_id: str = Path(description="The external provider identifier"),
|
||||
update: ExternalProviderConfigUpdate = Body(description="External provider configuration settings"),
|
||||
) -> ExternalProviderConfigModel:
|
||||
api_key_field, base_url_field = _get_external_provider_fields(provider_id)
|
||||
updates: dict[str, str | None] = {}
|
||||
|
||||
if update.api_key is not None:
|
||||
api_key = update.api_key.strip()
|
||||
updates[api_key_field] = api_key or None
|
||||
if update.base_url is not None:
|
||||
base_url = update.base_url.strip()
|
||||
updates[base_url_field] = base_url or None
|
||||
|
||||
if not updates:
|
||||
raise HTTPException(status_code=400, detail="No external provider config fields provided")
|
||||
|
||||
api_key_removed = update.api_key is not None and updates.get(api_key_field) is None
|
||||
_apply_external_provider_update(updates)
|
||||
if api_key_removed:
|
||||
_remove_external_models_for_provider(provider_id)
|
||||
return _build_external_provider_config(provider_id, get_config())
|
||||
|
||||
|
||||
@app_router.delete(
|
||||
"/external_providers/config/{provider_id}",
|
||||
operation_id="reset_external_provider_config",
|
||||
status_code=200,
|
||||
response_model=ExternalProviderConfigModel,
|
||||
)
|
||||
async def reset_external_provider_config(
|
||||
_: AdminUserOrDefault,
|
||||
provider_id: str = Path(description="The external provider identifier"),
|
||||
) -> ExternalProviderConfigModel:
|
||||
api_key_field, base_url_field = _get_external_provider_fields(provider_id)
|
||||
_apply_external_provider_update({api_key_field: None, base_url_field: None})
|
||||
_remove_external_models_for_provider(provider_id)
|
||||
return _build_external_provider_config(provider_id, get_config())
|
||||
|
||||
|
||||
def status_to_model(status: ExternalProviderStatus) -> ExternalProviderStatusModel:
|
||||
return ExternalProviderStatusModel(
|
||||
provider_id=status.provider_id,
|
||||
configured=status.configured,
|
||||
message=status.message,
|
||||
)
|
||||
|
||||
|
||||
def _get_external_provider_fields(provider_id: str) -> tuple[str, str]:
|
||||
if provider_id not in EXTERNAL_PROVIDER_FIELDS:
|
||||
raise HTTPException(status_code=404, detail=f"Unknown external provider '{provider_id}'")
|
||||
return EXTERNAL_PROVIDER_FIELDS[provider_id]
|
||||
|
||||
|
||||
def _write_external_api_keys_file(api_keys_file_path: FilePath, api_keys: dict[str, str]) -> None:
|
||||
if not api_keys:
|
||||
if api_keys_file_path.exists():
|
||||
api_keys_file_path.unlink()
|
||||
return
|
||||
|
||||
api_keys_file_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
with open(api_keys_file_path, "w", encoding=locale.getpreferredencoding()) as api_keys_file:
|
||||
yaml.safe_dump(api_keys, api_keys_file, sort_keys=False)
|
||||
|
||||
|
||||
def _apply_external_provider_update(updates: dict[str, str | None]) -> None:
|
||||
with _EXTERNAL_PROVIDER_CONFIG_LOCK:
|
||||
runtime_config = get_config()
|
||||
config_path = runtime_config.config_file_path
|
||||
api_keys_file_path = runtime_config.api_keys_file_path
|
||||
if config_path.exists():
|
||||
file_config = load_and_migrate_config(config_path)
|
||||
else:
|
||||
file_config = DefaultInvokeAIAppConfig()
|
||||
|
||||
runtime_config.update_config(updates)
|
||||
provider_config_fields = set(EXTERNAL_PROVIDER_CONFIG_FIELDS)
|
||||
provider_updates = {field: value for field, value in updates.items() if field in provider_config_fields}
|
||||
non_provider_updates = {field: value for field, value in updates.items() if field not in provider_config_fields}
|
||||
|
||||
if non_provider_updates:
|
||||
file_config.update_config(non_provider_updates)
|
||||
|
||||
persisted_api_keys = load_external_api_keys(api_keys_file_path)
|
||||
for field_name in EXTERNAL_PROVIDER_CONFIG_FIELDS:
|
||||
file_value = getattr(file_config, field_name, None)
|
||||
if field_name not in persisted_api_keys and isinstance(file_value, str) and file_value.strip():
|
||||
persisted_api_keys[field_name] = file_value
|
||||
|
||||
for field_name, value in provider_updates.items():
|
||||
if value is None:
|
||||
persisted_api_keys.pop(field_name, None)
|
||||
else:
|
||||
persisted_api_keys[field_name] = value
|
||||
|
||||
_write_external_api_keys_file(api_keys_file_path, persisted_api_keys)
|
||||
|
||||
for field_name in EXTERNAL_PROVIDER_CONFIG_FIELDS:
|
||||
setattr(file_config, field_name, None)
|
||||
|
||||
file_config_to_write = type(file_config).model_validate(
|
||||
file_config.model_dump(exclude_unset=True, exclude_none=True)
|
||||
)
|
||||
file_config_to_write.write_file(config_path, as_example=False)
|
||||
|
||||
|
||||
def _build_external_provider_config(provider_id: str, config: InvokeAIAppConfig) -> ExternalProviderConfigModel:
|
||||
api_key_field, base_url_field = _get_external_provider_fields(provider_id)
|
||||
return ExternalProviderConfigModel(
|
||||
provider_id=provider_id,
|
||||
api_key_configured=bool(getattr(config, api_key_field)),
|
||||
base_url=getattr(config, base_url_field),
|
||||
)
|
||||
|
||||
|
||||
def _remove_external_models_for_provider(provider_id: str) -> None:
|
||||
model_manager = ApiDependencies.invoker.services.model_manager
|
||||
external_models = model_manager.store.search_by_attr(
|
||||
base_model=BaseModelType.External,
|
||||
model_type=ModelType.ExternalImageGenerator,
|
||||
)
|
||||
|
||||
for model in external_models:
|
||||
if getattr(model, "provider_id", None) != provider_id:
|
||||
continue
|
||||
try:
|
||||
model_manager.install.delete(model.key)
|
||||
except UnknownModelException:
|
||||
logging.warning(f"External model key '{model.key}' was already removed while resetting '{provider_id}'")
|
||||
except Exception as error:
|
||||
logging.warning(f"Failed removing external model key '{model.key}' for '{provider_id}': {error}")
|
||||
|
||||
|
||||
@app_router.get(
|
||||
"/logging",
|
||||
operation_id="get_log_level",
|
||||
responses={200: {"description": "The operation was successful"}},
|
||||
response_model=LogLevel,
|
||||
)
|
||||
async def get_log_level() -> LogLevel:
|
||||
"""Returns the log level"""
|
||||
return LogLevel(ApiDependencies.invoker.services.logger.level)
|
||||
|
||||
|
||||
@app_router.post(
|
||||
"/logging",
|
||||
operation_id="set_log_level",
|
||||
responses={200: {"description": "The operation was successful"}},
|
||||
response_model=LogLevel,
|
||||
)
|
||||
async def set_log_level(
|
||||
level: LogLevel = Body(description="New log verbosity level"),
|
||||
) -> LogLevel:
|
||||
"""Sets the log verbosity level"""
|
||||
ApiDependencies.invoker.services.logger.setLevel(level)
|
||||
return LogLevel(ApiDependencies.invoker.services.logger.level)
|
||||
|
||||
|
||||
@app_router.delete(
|
||||
"/invocation_cache",
|
||||
operation_id="clear_invocation_cache",
|
||||
responses={200: {"description": "The operation was successful"}},
|
||||
)
|
||||
async def clear_invocation_cache() -> None:
|
||||
"""Clears the invocation cache"""
|
||||
ApiDependencies.invoker.services.invocation_cache.clear()
|
||||
|
||||
|
||||
@app_router.put(
|
||||
"/invocation_cache/enable",
|
||||
operation_id="enable_invocation_cache",
|
||||
responses={200: {"description": "The operation was successful"}},
|
||||
)
|
||||
async def enable_invocation_cache() -> None:
|
||||
"""Clears the invocation cache"""
|
||||
ApiDependencies.invoker.services.invocation_cache.enable()
|
||||
|
||||
|
||||
@app_router.put(
|
||||
"/invocation_cache/disable",
|
||||
operation_id="disable_invocation_cache",
|
||||
responses={200: {"description": "The operation was successful"}},
|
||||
)
|
||||
async def disable_invocation_cache() -> None:
|
||||
"""Clears the invocation cache"""
|
||||
ApiDependencies.invoker.services.invocation_cache.disable()
|
||||
|
||||
|
||||
@app_router.get(
|
||||
"/invocation_cache/status",
|
||||
operation_id="get_invocation_cache_status",
|
||||
responses={200: {"model": InvocationCacheStatus}},
|
||||
)
|
||||
async def get_invocation_cache_status() -> InvocationCacheStatus:
|
||||
"""Clears the invocation cache"""
|
||||
return ApiDependencies.invoker.services.invocation_cache.get_status()
|
||||
@@ -0,0 +1,536 @@
|
||||
"""Authentication endpoints."""
|
||||
|
||||
import secrets
|
||||
import string
|
||||
from datetime import timedelta
|
||||
from typing import Annotated
|
||||
|
||||
from fastapi import APIRouter, Body, HTTPException, Path, status
|
||||
from pydantic import BaseModel, Field, field_validator
|
||||
|
||||
from invokeai.app.api.auth_dependencies import AdminUser, CurrentUser
|
||||
from invokeai.app.api.dependencies import ApiDependencies
|
||||
from invokeai.app.services.auth.token_service import TokenData, create_access_token
|
||||
from invokeai.app.services.users.users_common import (
|
||||
UserCreateRequest,
|
||||
UserDTO,
|
||||
UserUpdateRequest,
|
||||
validate_email_with_special_domains,
|
||||
)
|
||||
|
||||
auth_router = APIRouter(prefix="/v1/auth", tags=["authentication"])
|
||||
|
||||
# Token expiration constants (in days)
|
||||
TOKEN_EXPIRATION_NORMAL = 1 # 1 day for normal login
|
||||
TOKEN_EXPIRATION_REMEMBER_ME = 7 # 7 days for "remember me" login
|
||||
|
||||
|
||||
class LoginRequest(BaseModel):
|
||||
"""Request body for user login."""
|
||||
|
||||
email: str = Field(description="User email address")
|
||||
password: str = Field(description="User password")
|
||||
remember_me: bool = Field(default=False, description="Whether to extend session duration")
|
||||
|
||||
@field_validator("email")
|
||||
@classmethod
|
||||
def validate_email(cls, v: str) -> str:
|
||||
"""Validate email address, allowing special-use domains."""
|
||||
return validate_email_with_special_domains(v)
|
||||
|
||||
|
||||
class LoginResponse(BaseModel):
|
||||
"""Response from successful login."""
|
||||
|
||||
token: str = Field(description="JWT access token")
|
||||
user: UserDTO = Field(description="User information")
|
||||
expires_in: int = Field(description="Token expiration time in seconds")
|
||||
|
||||
|
||||
class SetupRequest(BaseModel):
|
||||
"""Request body for initial admin setup."""
|
||||
|
||||
email: str = Field(description="Admin email address")
|
||||
display_name: str | None = Field(default=None, description="Admin display name")
|
||||
password: str = Field(description="Admin password")
|
||||
|
||||
@field_validator("email")
|
||||
@classmethod
|
||||
def validate_email(cls, v: str) -> str:
|
||||
"""Validate email address, allowing special-use domains."""
|
||||
return validate_email_with_special_domains(v)
|
||||
|
||||
|
||||
class SetupResponse(BaseModel):
|
||||
"""Response from successful admin setup."""
|
||||
|
||||
success: bool = Field(description="Whether setup was successful")
|
||||
user: UserDTO = Field(description="Created admin user information")
|
||||
|
||||
|
||||
class LogoutResponse(BaseModel):
|
||||
"""Response from logout."""
|
||||
|
||||
success: bool = Field(description="Whether logout was successful")
|
||||
|
||||
|
||||
class SetupStatusResponse(BaseModel):
|
||||
"""Response for setup status check."""
|
||||
|
||||
setup_required: bool = Field(description="Whether initial setup is required")
|
||||
multiuser_enabled: bool = Field(description="Whether multiuser mode is enabled")
|
||||
strict_password_checking: bool = Field(description="Whether strict password requirements are enforced")
|
||||
admin_email: str | None = Field(default=None, description="Email of the first active admin user, if any")
|
||||
|
||||
|
||||
@auth_router.get("/status", response_model=SetupStatusResponse)
|
||||
async def get_setup_status() -> SetupStatusResponse:
|
||||
"""Check if initial administrator setup is required.
|
||||
|
||||
Returns:
|
||||
SetupStatusResponse indicating whether setup is needed and multiuser mode status
|
||||
"""
|
||||
config = ApiDependencies.invoker.services.configuration
|
||||
|
||||
# If multiuser is disabled, setup is never required
|
||||
if not config.multiuser:
|
||||
return SetupStatusResponse(
|
||||
setup_required=False,
|
||||
multiuser_enabled=False,
|
||||
strict_password_checking=config.strict_password_checking,
|
||||
admin_email=None,
|
||||
)
|
||||
|
||||
# In multiuser mode, check if an admin exists
|
||||
user_service = ApiDependencies.invoker.services.users
|
||||
setup_required = not user_service.has_admin()
|
||||
|
||||
# Only expose admin_email during initial setup to avoid leaking
|
||||
# administrator identity on public deployments.
|
||||
admin_email = user_service.get_admin_email() if setup_required else None
|
||||
|
||||
return SetupStatusResponse(
|
||||
setup_required=setup_required,
|
||||
multiuser_enabled=True,
|
||||
strict_password_checking=config.strict_password_checking,
|
||||
admin_email=admin_email,
|
||||
)
|
||||
|
||||
|
||||
@auth_router.post("/login", response_model=LoginResponse)
|
||||
async def login(
|
||||
request: Annotated[LoginRequest, Body(description="Login credentials")],
|
||||
) -> LoginResponse:
|
||||
"""Authenticate user and return access token.
|
||||
|
||||
Args:
|
||||
request: Login credentials (email and password)
|
||||
|
||||
Returns:
|
||||
LoginResponse containing JWT token and user information
|
||||
|
||||
Raises:
|
||||
HTTPException: 401 if credentials are invalid or user is inactive
|
||||
HTTPException: 403 if multiuser mode is disabled
|
||||
"""
|
||||
config = ApiDependencies.invoker.services.configuration
|
||||
|
||||
# Check if multiuser is enabled
|
||||
if not config.multiuser:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_403_FORBIDDEN,
|
||||
detail="Multiuser mode is disabled. Authentication is not required in single-user mode.",
|
||||
)
|
||||
|
||||
user_service = ApiDependencies.invoker.services.users
|
||||
user = user_service.authenticate(request.email, request.password)
|
||||
|
||||
if user is None:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_401_UNAUTHORIZED,
|
||||
detail="Incorrect email or password",
|
||||
headers={"WWW-Authenticate": "Bearer"},
|
||||
)
|
||||
|
||||
if not user.is_active:
|
||||
raise HTTPException(status_code=status.HTTP_403_FORBIDDEN, detail="User account is disabled")
|
||||
|
||||
# Create token with appropriate expiration
|
||||
expires_delta = timedelta(days=TOKEN_EXPIRATION_REMEMBER_ME if request.remember_me else TOKEN_EXPIRATION_NORMAL)
|
||||
token_data = TokenData(
|
||||
user_id=user.user_id,
|
||||
email=user.email,
|
||||
is_admin=user.is_admin,
|
||||
remember_me=request.remember_me,
|
||||
)
|
||||
token = create_access_token(token_data, expires_delta)
|
||||
|
||||
return LoginResponse(
|
||||
token=token,
|
||||
user=user,
|
||||
expires_in=int(expires_delta.total_seconds()),
|
||||
)
|
||||
|
||||
|
||||
@auth_router.post("/logout", response_model=LogoutResponse)
|
||||
async def logout(
|
||||
current_user: CurrentUser,
|
||||
) -> LogoutResponse:
|
||||
"""Logout current user.
|
||||
|
||||
Currently a no-op since we use stateless JWT tokens. For token invalidation in
|
||||
future implementations, consider:
|
||||
- Token blacklist: Store invalidated tokens in Redis/database with expiration
|
||||
- Token versioning: Add version field to user record, increment on logout
|
||||
- Short-lived tokens: Use refresh token pattern with token rotation
|
||||
- Session storage: Track active sessions server-side for revocation
|
||||
|
||||
Args:
|
||||
current_user: The authenticated user (validates token)
|
||||
|
||||
Returns:
|
||||
LogoutResponse indicating success
|
||||
"""
|
||||
# TODO: Implement token invalidation when server-side session management is added
|
||||
# For now, this is a no-op since we use stateless JWT tokens
|
||||
return LogoutResponse(success=True)
|
||||
|
||||
|
||||
@auth_router.get("/me", response_model=UserDTO)
|
||||
async def get_current_user_info(
|
||||
current_user: CurrentUser,
|
||||
) -> UserDTO:
|
||||
"""Get current authenticated user's information.
|
||||
|
||||
Args:
|
||||
current_user: The authenticated user's token data
|
||||
|
||||
Returns:
|
||||
UserDTO containing user information
|
||||
|
||||
Raises:
|
||||
HTTPException: 404 if user is not found (should not happen normally)
|
||||
"""
|
||||
user_service = ApiDependencies.invoker.services.users
|
||||
user = user_service.get(current_user.user_id)
|
||||
|
||||
if user is None:
|
||||
raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="User not found")
|
||||
|
||||
return user
|
||||
|
||||
|
||||
@auth_router.post("/setup", response_model=SetupResponse)
|
||||
async def setup_admin(
|
||||
request: Annotated[SetupRequest, Body(description="Admin account details")],
|
||||
) -> SetupResponse:
|
||||
"""Set up initial administrator account.
|
||||
|
||||
This endpoint can only be called once, when no admin user exists. It creates
|
||||
the first admin user for the system.
|
||||
|
||||
Args:
|
||||
request: Admin account details (email, display_name, password)
|
||||
|
||||
Returns:
|
||||
SetupResponse containing the created admin user
|
||||
|
||||
Raises:
|
||||
HTTPException: 400 if admin already exists or password is weak
|
||||
HTTPException: 403 if multiuser mode is disabled
|
||||
"""
|
||||
config = ApiDependencies.invoker.services.configuration
|
||||
|
||||
# Check if multiuser is enabled
|
||||
if not config.multiuser:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_403_FORBIDDEN,
|
||||
detail="Multiuser mode is disabled. Admin setup is not required in single-user mode.",
|
||||
)
|
||||
|
||||
user_service = ApiDependencies.invoker.services.users
|
||||
|
||||
# Check if any admin exists
|
||||
if user_service.has_admin():
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail="Administrator account already configured",
|
||||
)
|
||||
|
||||
# Create admin user - this will validate password strength
|
||||
try:
|
||||
user_data = UserCreateRequest(
|
||||
email=request.email,
|
||||
display_name=request.display_name,
|
||||
password=request.password,
|
||||
is_admin=True,
|
||||
)
|
||||
user = user_service.create_admin(user_data, strict_password_checking=config.strict_password_checking)
|
||||
except ValueError as e:
|
||||
raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail=str(e)) from e
|
||||
|
||||
return SetupResponse(success=True, user=user)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# User management models
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
_PASSWORD_ALPHABET = string.ascii_letters + string.digits + string.punctuation
|
||||
|
||||
|
||||
class AdminUserCreateRequest(BaseModel):
|
||||
"""Request body for admin to create a new user."""
|
||||
|
||||
email: str = Field(description="User email address")
|
||||
display_name: str | None = Field(default=None, description="Display name")
|
||||
password: str = Field(description="User password")
|
||||
is_admin: bool = Field(default=False, description="Whether user should have admin privileges")
|
||||
|
||||
@field_validator("email")
|
||||
@classmethod
|
||||
def validate_email(cls, v: str) -> str:
|
||||
"""Validate email address, allowing special-use domains."""
|
||||
return validate_email_with_special_domains(v)
|
||||
|
||||
|
||||
class AdminUserUpdateRequest(BaseModel):
|
||||
"""Request body for admin to update any user."""
|
||||
|
||||
display_name: str | None = Field(default=None, description="Display name")
|
||||
password: str | None = Field(default=None, description="New password")
|
||||
is_admin: bool | None = Field(default=None, description="Whether user should have admin privileges")
|
||||
is_active: bool | None = Field(default=None, description="Whether user account should be active")
|
||||
|
||||
|
||||
class UserProfileUpdateRequest(BaseModel):
|
||||
"""Request body for a user to update their own profile."""
|
||||
|
||||
display_name: str | None = Field(default=None, description="New display name")
|
||||
current_password: str | None = Field(default=None, description="Current password (required when changing password)")
|
||||
new_password: str | None = Field(default=None, description="New password")
|
||||
|
||||
|
||||
class GeneratePasswordResponse(BaseModel):
|
||||
"""Response containing a generated password."""
|
||||
|
||||
password: str = Field(description="Generated strong password")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# User management endpoints
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@auth_router.get("/generate-password", response_model=GeneratePasswordResponse)
|
||||
async def generate_password(
|
||||
current_user: CurrentUser,
|
||||
) -> GeneratePasswordResponse:
|
||||
"""Generate a strong random password.
|
||||
|
||||
Returns a cryptographically secure random password of 16 characters
|
||||
containing uppercase, lowercase, digits, and punctuation.
|
||||
"""
|
||||
# Ensure the generated password always meets strength requirements:
|
||||
# at least one uppercase, one lowercase, one digit, one special char.
|
||||
while True:
|
||||
password = "".join(secrets.choice(_PASSWORD_ALPHABET) for _ in range(16))
|
||||
if (
|
||||
any(c.isupper() for c in password)
|
||||
and any(c.islower() for c in password)
|
||||
and any(c.isdigit() for c in password)
|
||||
):
|
||||
return GeneratePasswordResponse(password=password)
|
||||
|
||||
|
||||
@auth_router.get("/users", response_model=list[UserDTO])
|
||||
async def list_users(
|
||||
current_user: AdminUser,
|
||||
) -> list[UserDTO]:
|
||||
"""List all users. Requires admin privileges.
|
||||
|
||||
The internal 'system' user (created for backward compatibility) is excluded
|
||||
from the results since it cannot be managed through this interface.
|
||||
|
||||
Returns:
|
||||
List of all real users (system user excluded)
|
||||
"""
|
||||
user_service = ApiDependencies.invoker.services.users
|
||||
return [u for u in user_service.list_users() if u.user_id != "system"]
|
||||
|
||||
|
||||
@auth_router.post("/users", response_model=UserDTO, status_code=status.HTTP_201_CREATED)
|
||||
async def create_user(
|
||||
request: Annotated[AdminUserCreateRequest, Body(description="New user details")],
|
||||
current_user: AdminUser,
|
||||
) -> UserDTO:
|
||||
"""Create a new user. Requires admin privileges.
|
||||
|
||||
Args:
|
||||
request: New user details
|
||||
|
||||
Returns:
|
||||
The created user
|
||||
|
||||
Raises:
|
||||
HTTPException: 400 if email already exists or password is weak
|
||||
"""
|
||||
user_service = ApiDependencies.invoker.services.users
|
||||
config = ApiDependencies.invoker.services.configuration
|
||||
try:
|
||||
user_data = UserCreateRequest(
|
||||
email=request.email,
|
||||
display_name=request.display_name,
|
||||
password=request.password,
|
||||
is_admin=request.is_admin,
|
||||
)
|
||||
return user_service.create(user_data, strict_password_checking=config.strict_password_checking)
|
||||
except ValueError as e:
|
||||
raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail=str(e)) from e
|
||||
|
||||
|
||||
@auth_router.get("/users/{user_id}", response_model=UserDTO)
|
||||
async def get_user(
|
||||
user_id: Annotated[str, Path(description="User ID")],
|
||||
current_user: AdminUser,
|
||||
) -> UserDTO:
|
||||
"""Get a user by ID. Requires admin privileges.
|
||||
|
||||
Args:
|
||||
user_id: The user ID
|
||||
|
||||
Returns:
|
||||
The user
|
||||
|
||||
Raises:
|
||||
HTTPException: 404 if user not found
|
||||
"""
|
||||
user_service = ApiDependencies.invoker.services.users
|
||||
user = user_service.get(user_id)
|
||||
if user is None:
|
||||
raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="User not found")
|
||||
return user
|
||||
|
||||
|
||||
@auth_router.patch("/users/{user_id}", response_model=UserDTO)
|
||||
async def update_user(
|
||||
user_id: Annotated[str, Path(description="User ID")],
|
||||
request: Annotated[AdminUserUpdateRequest, Body(description="User fields to update")],
|
||||
current_user: AdminUser,
|
||||
) -> UserDTO:
|
||||
"""Update a user. Requires admin privileges.
|
||||
|
||||
Args:
|
||||
user_id: The user ID
|
||||
request: Fields to update
|
||||
|
||||
Returns:
|
||||
The updated user
|
||||
|
||||
Raises:
|
||||
HTTPException: 400 if password is weak
|
||||
HTTPException: 404 if user not found
|
||||
"""
|
||||
user_service = ApiDependencies.invoker.services.users
|
||||
config = ApiDependencies.invoker.services.configuration
|
||||
try:
|
||||
changes = UserUpdateRequest(
|
||||
display_name=request.display_name,
|
||||
password=request.password,
|
||||
is_admin=request.is_admin,
|
||||
is_active=request.is_active,
|
||||
)
|
||||
return user_service.update(user_id, changes, strict_password_checking=config.strict_password_checking)
|
||||
except ValueError as e:
|
||||
raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail=str(e)) from e
|
||||
|
||||
|
||||
@auth_router.delete("/users/{user_id}", status_code=status.HTTP_204_NO_CONTENT)
|
||||
async def delete_user(
|
||||
user_id: Annotated[str, Path(description="User ID")],
|
||||
current_user: AdminUser,
|
||||
) -> None:
|
||||
"""Delete a user. Requires admin privileges.
|
||||
|
||||
Admins can delete any user including other admins, but cannot delete the last
|
||||
remaining admin.
|
||||
|
||||
Args:
|
||||
user_id: The user ID
|
||||
|
||||
Raises:
|
||||
HTTPException: 400 if attempting to delete the last admin
|
||||
HTTPException: 404 if user not found
|
||||
"""
|
||||
user_service = ApiDependencies.invoker.services.users
|
||||
user = user_service.get(user_id)
|
||||
if user is None:
|
||||
raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="User not found")
|
||||
|
||||
# Prevent deleting the last active admin
|
||||
if user.is_admin and user.is_active and user_service.count_admins() <= 1:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail="Cannot delete the last administrator",
|
||||
)
|
||||
|
||||
try:
|
||||
user_service.delete(user_id)
|
||||
except ValueError as e:
|
||||
raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail=str(e)) from e
|
||||
|
||||
|
||||
@auth_router.patch("/me", response_model=UserDTO)
|
||||
async def update_current_user(
|
||||
request: Annotated[UserProfileUpdateRequest, Body(description="Profile fields to update")],
|
||||
current_user: CurrentUser,
|
||||
) -> UserDTO:
|
||||
"""Update the current user's own profile.
|
||||
|
||||
To change the password, both ``current_password`` and ``new_password`` must
|
||||
be provided. The current password is verified before the change is applied.
|
||||
|
||||
Args:
|
||||
request: Profile fields to update
|
||||
current_user: The authenticated user
|
||||
|
||||
Returns:
|
||||
The updated user
|
||||
|
||||
Raises:
|
||||
HTTPException: 400 if current password is incorrect or new password is weak
|
||||
HTTPException: 404 if user not found
|
||||
"""
|
||||
user_service = ApiDependencies.invoker.services.users
|
||||
config = ApiDependencies.invoker.services.configuration
|
||||
|
||||
# Verify current password when attempting a password change
|
||||
if request.new_password is not None:
|
||||
if not request.current_password:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail="Current password is required to set a new password",
|
||||
)
|
||||
|
||||
# Re-authenticate to verify the current password
|
||||
user = user_service.get(current_user.user_id)
|
||||
if user is None:
|
||||
raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="User not found")
|
||||
|
||||
authenticated = user_service.authenticate(user.email, request.current_password)
|
||||
if authenticated is None:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail="Current password is incorrect",
|
||||
)
|
||||
|
||||
try:
|
||||
changes = UserUpdateRequest(
|
||||
display_name=request.display_name,
|
||||
password=request.new_password,
|
||||
)
|
||||
return user_service.update(
|
||||
current_user.user_id, changes, strict_password_checking=config.strict_password_checking
|
||||
)
|
||||
except ValueError as e:
|
||||
raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail=str(e)) from e
|
||||
@@ -0,0 +1,223 @@
|
||||
from fastapi import Body, HTTPException
|
||||
from fastapi.routing import APIRouter
|
||||
|
||||
from invokeai.app.api.auth_dependencies import CurrentUserOrDefault
|
||||
from invokeai.app.api.dependencies import ApiDependencies
|
||||
from invokeai.app.api.routers.image_move_maintenance import assert_image_move_maintenance_inactive
|
||||
from invokeai.app.services.images.images_common import AddImagesToBoardResult, RemoveImagesFromBoardResult
|
||||
|
||||
board_images_router = APIRouter(prefix="/v1/board_images", tags=["boards"])
|
||||
|
||||
|
||||
def _assert_board_write_access(board_id: str, current_user: CurrentUserOrDefault) -> None:
|
||||
"""Raise 403 if the current user may not mutate the given board.
|
||||
|
||||
Write access is granted when ANY of these hold:
|
||||
- The user is an admin.
|
||||
- The user owns the board.
|
||||
- The board visibility is Public (public boards accept contributions from any user).
|
||||
"""
|
||||
from invokeai.app.services.board_records.board_records_common import BoardVisibility
|
||||
|
||||
try:
|
||||
board = ApiDependencies.invoker.services.boards.get_dto(board_id=board_id)
|
||||
except Exception:
|
||||
raise HTTPException(status_code=404, detail="Board not found")
|
||||
if current_user.is_admin:
|
||||
return
|
||||
if board.user_id == current_user.user_id:
|
||||
return
|
||||
if board.board_visibility == BoardVisibility.Public:
|
||||
return
|
||||
raise HTTPException(status_code=403, detail="Not authorized to modify this board")
|
||||
|
||||
|
||||
def _assert_image_direct_owner(image_name: str, current_user: CurrentUserOrDefault) -> None:
|
||||
"""Raise 403 if the current user is not the direct owner of the image.
|
||||
|
||||
This is intentionally stricter than _assert_image_owner in images.py:
|
||||
board ownership is NOT sufficient here. Allowing a user to add someone
|
||||
else's image to their own board would grant them mutation rights via the
|
||||
board-ownership fallback in _assert_image_owner, escalating read access
|
||||
into write access.
|
||||
"""
|
||||
if current_user.is_admin:
|
||||
return
|
||||
owner = ApiDependencies.invoker.services.image_records.get_user_id(image_name)
|
||||
if owner is not None and owner == current_user.user_id:
|
||||
return
|
||||
raise HTTPException(status_code=403, detail="Not authorized to move this image")
|
||||
|
||||
|
||||
@board_images_router.post(
|
||||
"/",
|
||||
operation_id="add_image_to_board",
|
||||
responses={
|
||||
201: {"description": "The image was added to a board successfully"},
|
||||
},
|
||||
status_code=201,
|
||||
response_model=AddImagesToBoardResult,
|
||||
)
|
||||
async def add_image_to_board(
|
||||
current_user: CurrentUserOrDefault,
|
||||
board_id: str = Body(description="The id of the board to add to"),
|
||||
image_name: str = Body(description="The name of the image to add"),
|
||||
) -> AddImagesToBoardResult:
|
||||
"""Creates a board_image"""
|
||||
_assert_board_write_access(board_id, current_user)
|
||||
_assert_image_direct_owner(image_name, current_user)
|
||||
assert_image_move_maintenance_inactive()
|
||||
try:
|
||||
added_images: set[str] = set()
|
||||
affected_boards: set[str] = set()
|
||||
old_board_id = ApiDependencies.invoker.services.board_image_records.get_board_for_image(image_name) or "none"
|
||||
ApiDependencies.invoker.services.board_images.add_image_to_board(board_id=board_id, image_name=image_name)
|
||||
added_images.add(image_name)
|
||||
affected_boards.add(board_id)
|
||||
affected_boards.add(old_board_id)
|
||||
|
||||
return AddImagesToBoardResult(
|
||||
added_images=list(added_images),
|
||||
affected_boards=list(affected_boards),
|
||||
)
|
||||
except Exception:
|
||||
raise HTTPException(status_code=500, detail="Failed to add image to board")
|
||||
|
||||
|
||||
@board_images_router.delete(
|
||||
"/",
|
||||
operation_id="remove_image_from_board",
|
||||
responses={
|
||||
201: {"description": "The image was removed from the board successfully"},
|
||||
},
|
||||
status_code=201,
|
||||
response_model=RemoveImagesFromBoardResult,
|
||||
)
|
||||
async def remove_image_from_board(
|
||||
current_user: CurrentUserOrDefault,
|
||||
image_name: str = Body(description="The name of the image to remove", embed=True),
|
||||
) -> RemoveImagesFromBoardResult:
|
||||
"""Removes an image from its board, if it had one"""
|
||||
try:
|
||||
old_board_id = ApiDependencies.invoker.services.images.get_dto(image_name).board_id or "none"
|
||||
if old_board_id != "none":
|
||||
_assert_board_write_access(old_board_id, current_user)
|
||||
assert_image_move_maintenance_inactive()
|
||||
removed_images: set[str] = set()
|
||||
affected_boards: set[str] = set()
|
||||
ApiDependencies.invoker.services.board_images.remove_image_from_board(image_name=image_name)
|
||||
removed_images.add(image_name)
|
||||
affected_boards.add("none")
|
||||
affected_boards.add(old_board_id)
|
||||
return RemoveImagesFromBoardResult(
|
||||
removed_images=list(removed_images),
|
||||
affected_boards=list(affected_boards),
|
||||
)
|
||||
|
||||
except HTTPException:
|
||||
raise
|
||||
except Exception:
|
||||
raise HTTPException(status_code=500, detail="Failed to remove image from board")
|
||||
|
||||
|
||||
@board_images_router.post(
|
||||
"/batch",
|
||||
operation_id="add_images_to_board",
|
||||
responses={
|
||||
201: {"description": "Images were added to board successfully"},
|
||||
},
|
||||
status_code=201,
|
||||
response_model=AddImagesToBoardResult,
|
||||
)
|
||||
async def add_images_to_board(
|
||||
current_user: CurrentUserOrDefault,
|
||||
board_id: str = Body(description="The id of the board to add to"),
|
||||
image_names: list[str] = Body(description="The names of the images to add", embed=True),
|
||||
) -> AddImagesToBoardResult:
|
||||
"""Adds a list of images to a board"""
|
||||
_assert_board_write_access(board_id, current_user)
|
||||
try:
|
||||
assert_image_move_maintenance_inactive()
|
||||
except HTTPException:
|
||||
for image_name in image_names:
|
||||
_assert_image_direct_owner(image_name, current_user)
|
||||
raise
|
||||
|
||||
try:
|
||||
added_images: set[str] = set()
|
||||
affected_boards: set[str] = set()
|
||||
for image_name in image_names:
|
||||
try:
|
||||
_assert_image_direct_owner(image_name, current_user)
|
||||
old_board_id = (
|
||||
ApiDependencies.invoker.services.board_image_records.get_board_for_image(image_name) or "none"
|
||||
)
|
||||
ApiDependencies.invoker.services.board_images.add_image_to_board(
|
||||
board_id=board_id,
|
||||
image_name=image_name,
|
||||
)
|
||||
added_images.add(image_name)
|
||||
affected_boards.add(board_id)
|
||||
affected_boards.add(old_board_id)
|
||||
|
||||
except HTTPException:
|
||||
raise
|
||||
except Exception:
|
||||
pass
|
||||
return AddImagesToBoardResult(
|
||||
added_images=list(added_images),
|
||||
affected_boards=list(affected_boards),
|
||||
)
|
||||
except HTTPException:
|
||||
raise
|
||||
except Exception:
|
||||
raise HTTPException(status_code=500, detail="Failed to add images to board")
|
||||
|
||||
|
||||
@board_images_router.post(
|
||||
"/batch/delete",
|
||||
operation_id="remove_images_from_board",
|
||||
responses={
|
||||
201: {"description": "Images were removed from board successfully"},
|
||||
},
|
||||
status_code=201,
|
||||
response_model=RemoveImagesFromBoardResult,
|
||||
)
|
||||
async def remove_images_from_board(
|
||||
current_user: CurrentUserOrDefault,
|
||||
image_names: list[str] = Body(description="The names of the images to remove", embed=True),
|
||||
) -> RemoveImagesFromBoardResult:
|
||||
"""Removes a list of images from their board, if they had one"""
|
||||
try:
|
||||
assert_image_move_maintenance_inactive()
|
||||
except HTTPException:
|
||||
for image_name in image_names:
|
||||
old_board_id = ApiDependencies.invoker.services.images.get_dto(image_name).board_id or "none"
|
||||
if old_board_id != "none":
|
||||
_assert_board_write_access(old_board_id, current_user)
|
||||
raise
|
||||
|
||||
try:
|
||||
removed_images: set[str] = set()
|
||||
affected_boards: set[str] = set()
|
||||
for image_name in image_names:
|
||||
try:
|
||||
old_board_id = ApiDependencies.invoker.services.images.get_dto(image_name).board_id or "none"
|
||||
if old_board_id != "none":
|
||||
_assert_board_write_access(old_board_id, current_user)
|
||||
ApiDependencies.invoker.services.board_images.remove_image_from_board(image_name=image_name)
|
||||
removed_images.add(image_name)
|
||||
affected_boards.add("none")
|
||||
affected_boards.add(old_board_id)
|
||||
except HTTPException:
|
||||
raise
|
||||
except Exception:
|
||||
pass
|
||||
return RemoveImagesFromBoardResult(
|
||||
removed_images=list(removed_images),
|
||||
affected_boards=list(affected_boards),
|
||||
)
|
||||
except HTTPException:
|
||||
raise
|
||||
except Exception:
|
||||
raise HTTPException(status_code=500, detail="Failed to remove images from board")
|
||||
@@ -0,0 +1,225 @@
|
||||
from typing import Optional, Union
|
||||
|
||||
from fastapi import Body, HTTPException, Path, Query
|
||||
from fastapi.routing import APIRouter
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
from invokeai.app.api.auth_dependencies import CurrentUserOrDefault
|
||||
from invokeai.app.api.dependencies import ApiDependencies
|
||||
from invokeai.app.api.routers.image_move_maintenance import assert_image_move_maintenance_inactive
|
||||
from invokeai.app.services.board_records.board_records_common import BoardChanges, BoardRecordOrderBy, BoardVisibility
|
||||
from invokeai.app.services.boards.boards_common import BoardDTO
|
||||
from invokeai.app.services.image_records.image_records_common import ImageCategory
|
||||
from invokeai.app.services.shared.pagination import OffsetPaginatedResults
|
||||
from invokeai.app.services.shared.sqlite.sqlite_common import SQLiteDirection
|
||||
|
||||
boards_router = APIRouter(prefix="/v1/boards", tags=["boards"])
|
||||
|
||||
|
||||
class DeleteBoardResult(BaseModel):
|
||||
board_id: str = Field(description="The id of the board that was deleted.")
|
||||
deleted_board_images: list[str] = Field(
|
||||
description="The image names of the board-images relationships that were deleted."
|
||||
)
|
||||
deleted_images: list[str] = Field(description="The names of the images that were deleted.")
|
||||
|
||||
|
||||
@boards_router.post(
|
||||
"/",
|
||||
operation_id="create_board",
|
||||
responses={
|
||||
201: {"description": "The board was created successfully"},
|
||||
},
|
||||
status_code=201,
|
||||
response_model=BoardDTO,
|
||||
)
|
||||
async def create_board(
|
||||
current_user: CurrentUserOrDefault,
|
||||
board_name: str = Query(description="The name of the board to create", max_length=300),
|
||||
) -> BoardDTO:
|
||||
"""Creates a board for the current user"""
|
||||
try:
|
||||
result = ApiDependencies.invoker.services.boards.create(board_name=board_name, user_id=current_user.user_id)
|
||||
return result
|
||||
except Exception:
|
||||
raise HTTPException(status_code=500, detail="Failed to create board")
|
||||
|
||||
|
||||
@boards_router.get("/{board_id}", operation_id="get_board", response_model=BoardDTO)
|
||||
async def get_board(
|
||||
current_user: CurrentUserOrDefault,
|
||||
board_id: str = Path(description="The id of board to get"),
|
||||
) -> BoardDTO:
|
||||
"""Gets a board (user must have access to it)"""
|
||||
|
||||
try:
|
||||
result = ApiDependencies.invoker.services.boards.get_dto(board_id=board_id)
|
||||
except Exception:
|
||||
raise HTTPException(status_code=404, detail="Board not found")
|
||||
|
||||
# Admins can access any board.
|
||||
# Owners can access their own boards.
|
||||
# Shared and public boards are visible to all authenticated users.
|
||||
if (
|
||||
not current_user.is_admin
|
||||
and result.user_id != current_user.user_id
|
||||
and result.board_visibility == BoardVisibility.Private
|
||||
):
|
||||
raise HTTPException(status_code=403, detail="Not authorized to access this board")
|
||||
|
||||
return result
|
||||
|
||||
|
||||
@boards_router.patch(
|
||||
"/{board_id}",
|
||||
operation_id="update_board",
|
||||
responses={
|
||||
201: {
|
||||
"description": "The board was updated successfully",
|
||||
},
|
||||
},
|
||||
status_code=201,
|
||||
response_model=BoardDTO,
|
||||
)
|
||||
async def update_board(
|
||||
current_user: CurrentUserOrDefault,
|
||||
board_id: str = Path(description="The id of board to update"),
|
||||
changes: BoardChanges = Body(description="The changes to apply to the board"),
|
||||
) -> BoardDTO:
|
||||
"""Updates a board (user must have access to it)"""
|
||||
try:
|
||||
board = ApiDependencies.invoker.services.boards.get_dto(board_id=board_id)
|
||||
except Exception:
|
||||
raise HTTPException(status_code=404, detail="Board not found")
|
||||
|
||||
if not current_user.is_admin and board.user_id != current_user.user_id:
|
||||
raise HTTPException(status_code=403, detail="Not authorized to update this board")
|
||||
|
||||
try:
|
||||
result = ApiDependencies.invoker.services.boards.update(board_id=board_id, changes=changes)
|
||||
return result
|
||||
except Exception:
|
||||
raise HTTPException(status_code=500, detail="Failed to update board")
|
||||
|
||||
|
||||
@boards_router.delete("/{board_id}", operation_id="delete_board", response_model=DeleteBoardResult)
|
||||
async def delete_board(
|
||||
current_user: CurrentUserOrDefault,
|
||||
board_id: str = Path(description="The id of board to delete"),
|
||||
include_images: Optional[bool] = Query(description="Permanently delete all images on the board", default=False),
|
||||
) -> DeleteBoardResult:
|
||||
"""Deletes a board (user must have access to it)"""
|
||||
try:
|
||||
board = ApiDependencies.invoker.services.boards.get_dto(board_id=board_id)
|
||||
except Exception:
|
||||
raise HTTPException(status_code=404, detail="Board not found")
|
||||
|
||||
if not current_user.is_admin and board.user_id != current_user.user_id:
|
||||
raise HTTPException(status_code=403, detail="Not authorized to delete this board")
|
||||
|
||||
try:
|
||||
if include_images is True:
|
||||
assert_image_move_maintenance_inactive()
|
||||
deleted_images = ApiDependencies.invoker.services.board_images.get_all_board_image_names_for_board(
|
||||
board_id=board_id,
|
||||
categories=None,
|
||||
is_intermediate=None,
|
||||
)
|
||||
ApiDependencies.invoker.services.images.delete_images_on_board(board_id=board_id)
|
||||
ApiDependencies.invoker.services.boards.delete(board_id=board_id)
|
||||
return DeleteBoardResult(
|
||||
board_id=board_id,
|
||||
deleted_board_images=[],
|
||||
deleted_images=deleted_images,
|
||||
)
|
||||
else:
|
||||
deleted_board_images = ApiDependencies.invoker.services.board_images.get_all_board_image_names_for_board(
|
||||
board_id=board_id,
|
||||
categories=None,
|
||||
is_intermediate=None,
|
||||
)
|
||||
ApiDependencies.invoker.services.boards.delete(board_id=board_id)
|
||||
return DeleteBoardResult(
|
||||
board_id=board_id,
|
||||
deleted_board_images=deleted_board_images,
|
||||
deleted_images=[],
|
||||
)
|
||||
except HTTPException:
|
||||
raise
|
||||
except Exception:
|
||||
raise HTTPException(status_code=500, detail="Failed to delete board")
|
||||
|
||||
|
||||
@boards_router.get(
|
||||
"/",
|
||||
operation_id="list_boards",
|
||||
response_model=Union[OffsetPaginatedResults[BoardDTO], list[BoardDTO]],
|
||||
)
|
||||
async def list_boards(
|
||||
current_user: CurrentUserOrDefault,
|
||||
order_by: BoardRecordOrderBy = Query(default=BoardRecordOrderBy.CreatedAt, description="The attribute to order by"),
|
||||
direction: SQLiteDirection = Query(default=SQLiteDirection.Descending, description="The direction to order by"),
|
||||
all: Optional[bool] = Query(default=None, description="Whether to list all boards"),
|
||||
offset: Optional[int] = Query(default=None, description="The page offset"),
|
||||
limit: Optional[int] = Query(default=None, description="The number of boards per page"),
|
||||
include_archived: bool = Query(default=False, description="Whether or not to include archived boards in list"),
|
||||
) -> Union[OffsetPaginatedResults[BoardDTO], list[BoardDTO]]:
|
||||
"""Gets a list of boards for the current user, including shared boards. Admin users see all boards."""
|
||||
if all:
|
||||
return ApiDependencies.invoker.services.boards.get_all(
|
||||
current_user.user_id, current_user.is_admin, order_by, direction, include_archived
|
||||
)
|
||||
elif offset is not None and limit is not None:
|
||||
return ApiDependencies.invoker.services.boards.get_many(
|
||||
current_user.user_id, current_user.is_admin, order_by, direction, offset, limit, include_archived
|
||||
)
|
||||
else:
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail="Invalid request: Must provide either 'all' or both 'offset' and 'limit'",
|
||||
)
|
||||
|
||||
|
||||
@boards_router.get(
|
||||
"/{board_id}/image_names",
|
||||
operation_id="list_all_board_image_names",
|
||||
response_model=list[str],
|
||||
)
|
||||
async def list_all_board_image_names(
|
||||
current_user: CurrentUserOrDefault,
|
||||
board_id: str = Path(description="The id of the board or 'none' for uncategorized images"),
|
||||
categories: list[ImageCategory] | None = Query(default=None, description="The categories of image to include."),
|
||||
is_intermediate: bool | None = Query(default=None, description="Whether to list intermediate images."),
|
||||
) -> list[str]:
|
||||
"""Gets a list of images for a board"""
|
||||
|
||||
if board_id != "none":
|
||||
try:
|
||||
board = ApiDependencies.invoker.services.boards.get_dto(board_id=board_id)
|
||||
except Exception:
|
||||
raise HTTPException(status_code=404, detail="Board not found")
|
||||
|
||||
if (
|
||||
not current_user.is_admin
|
||||
and board.user_id != current_user.user_id
|
||||
and board.board_visibility == BoardVisibility.Private
|
||||
):
|
||||
raise HTTPException(status_code=403, detail="Not authorized to access this board")
|
||||
|
||||
image_names = ApiDependencies.invoker.services.board_images.get_all_board_image_names_for_board(
|
||||
board_id,
|
||||
categories,
|
||||
is_intermediate,
|
||||
)
|
||||
|
||||
# For uncategorized images (board_id="none"), filter to only the caller's
|
||||
# images so that one user cannot enumerate another's uncategorized images.
|
||||
# Admin users can see all uncategorized images.
|
||||
if board_id == "none" and not current_user.is_admin:
|
||||
image_names = [
|
||||
name
|
||||
for name in image_names
|
||||
if ApiDependencies.invoker.services.image_records.get_user_id(name) == current_user.user_id
|
||||
]
|
||||
|
||||
return image_names
|
||||
@@ -0,0 +1,100 @@
|
||||
from fastapi import Body, HTTPException, Path, Query
|
||||
from fastapi.routing import APIRouter
|
||||
|
||||
from invokeai.app.api.auth_dependencies import CurrentUserOrDefault
|
||||
from invokeai.app.api.dependencies import ApiDependencies
|
||||
from invokeai.backend.util.logging import logging
|
||||
|
||||
client_state_router = APIRouter(prefix="/v1/client_state", tags=["client_state"])
|
||||
|
||||
|
||||
@client_state_router.get(
|
||||
"/{queue_id}/get_by_key",
|
||||
operation_id="get_client_state_by_key",
|
||||
response_model=str | None,
|
||||
)
|
||||
async def get_client_state_by_key(
|
||||
current_user: CurrentUserOrDefault,
|
||||
queue_id: str = Path(description="The queue id (ignored, kept for backwards compatibility)"),
|
||||
key: str = Query(..., description="Key to get"),
|
||||
) -> str | None:
|
||||
"""Gets the client state for the current user (or system user if not authenticated)"""
|
||||
try:
|
||||
return ApiDependencies.invoker.services.client_state_persistence.get_by_key(current_user.user_id, key)
|
||||
except Exception as e:
|
||||
logging.error(f"Error getting client state: {e}")
|
||||
raise HTTPException(status_code=500, detail="Error getting client state")
|
||||
|
||||
|
||||
@client_state_router.post(
|
||||
"/{queue_id}/set_by_key",
|
||||
operation_id="set_client_state",
|
||||
response_model=str,
|
||||
)
|
||||
async def set_client_state(
|
||||
current_user: CurrentUserOrDefault,
|
||||
queue_id: str = Path(description="The queue id (ignored, kept for backwards compatibility)"),
|
||||
key: str = Query(..., description="Key to set"),
|
||||
value: str = Body(..., description="Stringified value to set"),
|
||||
) -> str:
|
||||
"""Sets the client state for the current user (or system user if not authenticated)"""
|
||||
try:
|
||||
return ApiDependencies.invoker.services.client_state_persistence.set_by_key(current_user.user_id, key, value)
|
||||
except Exception as e:
|
||||
logging.error(f"Error setting client state: {e}")
|
||||
raise HTTPException(status_code=500, detail="Error setting client state")
|
||||
|
||||
|
||||
@client_state_router.get(
|
||||
"/{queue_id}/get_keys_by_prefix",
|
||||
operation_id="get_client_state_keys_by_prefix",
|
||||
response_model=list[str],
|
||||
)
|
||||
async def get_client_state_keys_by_prefix(
|
||||
current_user: CurrentUserOrDefault,
|
||||
queue_id: str = Path(description="The queue id (ignored, kept for backwards compatibility)"),
|
||||
prefix: str = Query(..., description="Prefix to filter keys by"),
|
||||
) -> list[str]:
|
||||
"""Gets client state keys matching a prefix for the current user"""
|
||||
try:
|
||||
return ApiDependencies.invoker.services.client_state_persistence.get_keys_by_prefix(
|
||||
current_user.user_id, prefix
|
||||
)
|
||||
except Exception as e:
|
||||
logging.error(f"Error getting client state keys: {e}")
|
||||
raise HTTPException(status_code=500, detail="Error getting client state keys")
|
||||
|
||||
|
||||
@client_state_router.post(
|
||||
"/{queue_id}/delete_by_key",
|
||||
operation_id="delete_client_state_by_key",
|
||||
responses={204: {"description": "Client state key deleted"}},
|
||||
)
|
||||
async def delete_client_state_by_key(
|
||||
current_user: CurrentUserOrDefault,
|
||||
queue_id: str = Path(description="The queue id (ignored, kept for backwards compatibility)"),
|
||||
key: str = Query(..., description="Key to delete"),
|
||||
) -> None:
|
||||
"""Deletes a specific client state key for the current user"""
|
||||
try:
|
||||
ApiDependencies.invoker.services.client_state_persistence.delete_by_key(current_user.user_id, key)
|
||||
except Exception as e:
|
||||
logging.error(f"Error deleting client state key: {e}")
|
||||
raise HTTPException(status_code=500, detail="Error deleting client state key")
|
||||
|
||||
|
||||
@client_state_router.post(
|
||||
"/{queue_id}/delete",
|
||||
operation_id="delete_client_state",
|
||||
responses={204: {"description": "Client state deleted"}},
|
||||
)
|
||||
async def delete_client_state(
|
||||
current_user: CurrentUserOrDefault,
|
||||
queue_id: str = Path(description="The queue id (ignored, kept for backwards compatibility)"),
|
||||
) -> None:
|
||||
"""Deletes the client state for the current user (or system user if not authenticated)"""
|
||||
try:
|
||||
ApiDependencies.invoker.services.client_state_persistence.delete(current_user.user_id)
|
||||
except Exception as e:
|
||||
logging.error(f"Error deleting client state: {e}")
|
||||
raise HTTPException(status_code=500, detail="Error deleting client state")
|
||||
@@ -0,0 +1,504 @@
|
||||
"""FastAPI routes for custom node management."""
|
||||
|
||||
import json
|
||||
import shutil
|
||||
import subprocess
|
||||
import sys
|
||||
import traceback
|
||||
from importlib.util import module_from_spec, spec_from_file_location
|
||||
from pathlib import Path
|
||||
from typing import Optional
|
||||
|
||||
from fastapi import Body
|
||||
from fastapi.routing import APIRouter
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
from invokeai.app.api.auth_dependencies import AdminUserOrDefault
|
||||
from invokeai.app.api.dependencies import ApiDependencies
|
||||
from invokeai.app.invocations.baseinvocation import InvocationRegistry
|
||||
from invokeai.app.services.config.config_default import get_config
|
||||
from invokeai.app.services.workflow_records.workflow_records_common import WorkflowWithoutIDValidator
|
||||
from invokeai.backend.util.logging import InvokeAILogger
|
||||
|
||||
custom_nodes_router = APIRouter(prefix="/v2/custom_nodes", tags=["custom_nodes"])
|
||||
|
||||
logger = InvokeAILogger.get_logger()
|
||||
|
||||
# Name of the manifest file written inside a pack directory to track which workflows
|
||||
# were imported by that pack. Used on uninstall to delete only pack-imported workflows
|
||||
# — deleting by tag alone is unsafe because users can edit tags on their own workflows.
|
||||
PACK_MANIFEST_FILENAME = ".invokeai_pack_manifest.json"
|
||||
|
||||
|
||||
class NodePackInfo(BaseModel):
|
||||
"""Information about an installed node pack."""
|
||||
|
||||
name: str = Field(description="The name of the node pack.")
|
||||
path: str = Field(description="The path to the node pack directory.")
|
||||
node_count: int = Field(description="The number of nodes in the pack.")
|
||||
node_types: list[str] = Field(description="The invocation types provided by this node pack.")
|
||||
|
||||
|
||||
class NodePackListResponse(BaseModel):
|
||||
"""Response for listing installed node packs."""
|
||||
|
||||
node_packs: list[NodePackInfo] = Field(description="List of installed node packs.")
|
||||
custom_nodes_path: str = Field(description="The configured custom nodes directory path.")
|
||||
|
||||
|
||||
class InstallNodePackRequest(BaseModel):
|
||||
"""Request to install a node pack from a git URL."""
|
||||
|
||||
source: str = Field(description="Git URL of the node pack to install.")
|
||||
|
||||
|
||||
class InstallNodePackResponse(BaseModel):
|
||||
"""Response after installing a node pack."""
|
||||
|
||||
name: str = Field(description="The name of the installed node pack.")
|
||||
success: bool = Field(description="Whether the installation was successful.")
|
||||
message: str = Field(description="Status message.")
|
||||
workflows_imported: int = Field(default=0, description="Number of workflows imported from the pack.")
|
||||
requires_dependencies: bool = Field(
|
||||
default=False,
|
||||
description="Whether the pack ships a dependency manifest (requirements.txt or pyproject.toml) "
|
||||
"that the user must install manually following the pack's documentation.",
|
||||
)
|
||||
dependency_file: Optional[str] = Field(
|
||||
default=None,
|
||||
description="Name of the detected dependency manifest file, if any.",
|
||||
)
|
||||
|
||||
|
||||
class UninstallNodePackResponse(BaseModel):
|
||||
"""Response after uninstalling a node pack."""
|
||||
|
||||
name: str = Field(description="The name of the uninstalled node pack.")
|
||||
success: bool = Field(description="Whether the uninstall was successful.")
|
||||
message: str = Field(description="Status message.")
|
||||
|
||||
|
||||
def _get_custom_nodes_path() -> Path:
|
||||
"""Returns the configured custom nodes directory path."""
|
||||
config = get_config()
|
||||
return config.custom_nodes_path
|
||||
|
||||
|
||||
def _get_installed_packs() -> list[NodePackInfo]:
|
||||
"""Scans the custom nodes directory and returns info about installed packs."""
|
||||
custom_nodes_path = _get_custom_nodes_path()
|
||||
|
||||
if not custom_nodes_path.exists():
|
||||
return []
|
||||
|
||||
packs: list[NodePackInfo] = []
|
||||
|
||||
# Get all node types grouped by node_pack
|
||||
node_types_by_pack: dict[str, list[str]] = {}
|
||||
for inv_class in InvocationRegistry._invocation_classes:
|
||||
node_pack = inv_class.UIConfig.node_pack
|
||||
inv_type = inv_class.get_type()
|
||||
if node_pack not in node_types_by_pack:
|
||||
node_types_by_pack[node_pack] = []
|
||||
node_types_by_pack[node_pack].append(inv_type)
|
||||
|
||||
for d in sorted(custom_nodes_path.iterdir()):
|
||||
if not d.is_dir():
|
||||
continue
|
||||
if d.name.startswith("_") or d.name.startswith("."):
|
||||
continue
|
||||
init = d / "__init__.py"
|
||||
if not init.exists():
|
||||
continue
|
||||
|
||||
pack_name = d.name
|
||||
node_types = node_types_by_pack.get(pack_name, [])
|
||||
|
||||
packs.append(
|
||||
NodePackInfo(
|
||||
name=pack_name,
|
||||
path=str(d),
|
||||
node_count=len(node_types),
|
||||
node_types=node_types,
|
||||
)
|
||||
)
|
||||
|
||||
return packs
|
||||
|
||||
|
||||
@custom_nodes_router.get(
|
||||
"/",
|
||||
operation_id="list_custom_node_packs",
|
||||
response_model=NodePackListResponse,
|
||||
)
|
||||
async def list_custom_node_packs(current_admin: AdminUserOrDefault) -> NodePackListResponse:
|
||||
"""Lists all installed custom node packs.
|
||||
|
||||
Admin-only: the response includes absolute filesystem paths, and non-admins have no
|
||||
legitimate use for pack management data (install/uninstall/reload are also admin-only).
|
||||
"""
|
||||
packs = _get_installed_packs()
|
||||
return NodePackListResponse(node_packs=packs, custom_nodes_path=str(_get_custom_nodes_path()))
|
||||
|
||||
|
||||
@custom_nodes_router.post(
|
||||
"/install",
|
||||
operation_id="install_custom_node_pack",
|
||||
response_model=InstallNodePackResponse,
|
||||
)
|
||||
async def install_custom_node_pack(
|
||||
current_admin: AdminUserOrDefault,
|
||||
request: InstallNodePackRequest = Body(description="The source URL to install from."),
|
||||
) -> InstallNodePackResponse:
|
||||
"""Installs a custom node pack from a git URL by cloning it into the nodes directory."""
|
||||
custom_nodes_path = _get_custom_nodes_path()
|
||||
custom_nodes_path.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
source = request.source.strip()
|
||||
|
||||
# Extract pack name from URL
|
||||
pack_name = source.rstrip("/").split("/")[-1]
|
||||
if pack_name.endswith(".git"):
|
||||
pack_name = pack_name[:-4]
|
||||
|
||||
target_dir = custom_nodes_path / pack_name
|
||||
|
||||
if target_dir.exists():
|
||||
return InstallNodePackResponse(
|
||||
name=pack_name,
|
||||
success=False,
|
||||
message=f"Node pack '{pack_name}' already exists. Uninstall it first to reinstall.",
|
||||
)
|
||||
|
||||
try:
|
||||
# Clone the repository
|
||||
result = subprocess.run(
|
||||
["git", "clone", source, str(target_dir)],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=120,
|
||||
)
|
||||
|
||||
if result.returncode != 0:
|
||||
# Clean up on failure
|
||||
if target_dir.exists():
|
||||
shutil.rmtree(target_dir)
|
||||
return InstallNodePackResponse(
|
||||
name=pack_name,
|
||||
success=False,
|
||||
message=f"Git clone failed: {result.stderr.strip()}",
|
||||
)
|
||||
|
||||
# Detect dependency manifests but do NOT install them automatically.
|
||||
# The user is responsible for installing dependencies per the pack's documentation,
|
||||
# since arbitrary pip installs can break the InvokeAI environment.
|
||||
dependency_file: Optional[str] = None
|
||||
for candidate in ("requirements.txt", "pyproject.toml"):
|
||||
if (target_dir / candidate).exists():
|
||||
dependency_file = candidate
|
||||
logger.info(f"Node pack '{pack_name}' ships a {candidate}; user must install dependencies manually.")
|
||||
break
|
||||
|
||||
# Check for __init__.py
|
||||
init_file = target_dir / "__init__.py"
|
||||
if not init_file.exists():
|
||||
shutil.rmtree(target_dir)
|
||||
return InstallNodePackResponse(
|
||||
name=pack_name,
|
||||
success=False,
|
||||
message=f"Node pack '{pack_name}' does not contain an __init__.py file.",
|
||||
)
|
||||
|
||||
# Load the node pack at runtime
|
||||
_load_node_pack(pack_name, target_dir)
|
||||
|
||||
# Import any workflows found in the pack, owned by the installing admin and shared with all users
|
||||
imported_workflow_ids = _import_workflows_from_pack(target_dir, pack_name, owner_user_id=current_admin.user_id)
|
||||
_write_pack_manifest(target_dir, imported_workflow_ids)
|
||||
workflows_imported = len(imported_workflow_ids)
|
||||
workflow_msg = f" Imported {workflows_imported} workflow(s)." if workflows_imported > 0 else ""
|
||||
dependency_msg = (
|
||||
f" This pack includes a {dependency_file} — install its dependencies manually following the pack's documentation."
|
||||
if dependency_file
|
||||
else ""
|
||||
)
|
||||
|
||||
return InstallNodePackResponse(
|
||||
name=pack_name,
|
||||
success=True,
|
||||
message=f"Successfully installed node pack '{pack_name}'.{workflow_msg}{dependency_msg}",
|
||||
workflows_imported=workflows_imported,
|
||||
requires_dependencies=dependency_file is not None,
|
||||
dependency_file=dependency_file,
|
||||
)
|
||||
|
||||
except subprocess.TimeoutExpired:
|
||||
if target_dir.exists():
|
||||
shutil.rmtree(target_dir)
|
||||
return InstallNodePackResponse(
|
||||
name=pack_name,
|
||||
success=False,
|
||||
message="Installation timed out.",
|
||||
)
|
||||
except Exception:
|
||||
if target_dir.exists():
|
||||
shutil.rmtree(target_dir)
|
||||
error = traceback.format_exc()
|
||||
logger.error(f"Failed to install node pack {pack_name}: {error}")
|
||||
return InstallNodePackResponse(
|
||||
name=pack_name,
|
||||
success=False,
|
||||
message=f"Installation failed: {error}",
|
||||
)
|
||||
|
||||
|
||||
@custom_nodes_router.delete(
|
||||
"/{pack_name}",
|
||||
operation_id="uninstall_custom_node_pack",
|
||||
response_model=UninstallNodePackResponse,
|
||||
)
|
||||
async def uninstall_custom_node_pack(
|
||||
current_admin: AdminUserOrDefault,
|
||||
pack_name: str,
|
||||
) -> UninstallNodePackResponse:
|
||||
"""Uninstalls a custom node pack by removing its directory.
|
||||
|
||||
Note: A restart is required for the node removal to take full effect.
|
||||
Installed nodes from the pack will remain registered until restart.
|
||||
"""
|
||||
custom_nodes_path = _get_custom_nodes_path()
|
||||
target_dir = custom_nodes_path / pack_name
|
||||
|
||||
if not target_dir.exists():
|
||||
return UninstallNodePackResponse(
|
||||
name=pack_name,
|
||||
success=False,
|
||||
message=f"Node pack '{pack_name}' not found.",
|
||||
)
|
||||
|
||||
try:
|
||||
# Read the manifest BEFORE removing the directory — it records exactly which
|
||||
# workflow IDs this pack imported, so uninstall doesn't accidentally delete
|
||||
# user workflows that happen to share the pack tag.
|
||||
imported_workflow_ids = _read_pack_manifest(target_dir)
|
||||
|
||||
shutil.rmtree(target_dir)
|
||||
|
||||
# Unregister the nodes from the registry so they disappear immediately
|
||||
removed_types = InvocationRegistry.unregister_pack(pack_name)
|
||||
if removed_types:
|
||||
# Invalidate OpenAPI schema cache so frontend gets updated node definitions
|
||||
from invokeai.app.api_app import app
|
||||
|
||||
app.openapi_schema = None
|
||||
logger.info(
|
||||
f"Unregistered {len(removed_types)} node(s) from pack '{pack_name}': {', '.join(removed_types)}"
|
||||
)
|
||||
|
||||
# Remove the pack's module subtree from sys.modules. Only dropping the
|
||||
# root module would leave submodules cached; on reinstall the cached
|
||||
# submodules would be reused without re-running their @invocation
|
||||
# decorators, so the pack would show up with 0 nodes until restart.
|
||||
_purge_pack_modules(pack_name)
|
||||
|
||||
# Remove only workflows this pack imported, using the manifest-recorded IDs
|
||||
workflows_removed = _remove_workflows_by_ids(imported_workflow_ids, pack_name)
|
||||
workflow_msg = f" Removed {workflows_removed} workflow(s)." if workflows_removed > 0 else ""
|
||||
|
||||
return UninstallNodePackResponse(
|
||||
name=pack_name,
|
||||
success=True,
|
||||
message=f"Successfully uninstalled node pack '{pack_name}'.{workflow_msg}",
|
||||
)
|
||||
except Exception:
|
||||
error = traceback.format_exc()
|
||||
logger.error(f"Failed to uninstall node pack {pack_name}: {error}")
|
||||
return UninstallNodePackResponse(
|
||||
name=pack_name,
|
||||
success=False,
|
||||
message=f"Uninstall failed: {error}",
|
||||
)
|
||||
|
||||
|
||||
@custom_nodes_router.post(
|
||||
"/reload",
|
||||
operation_id="reload_custom_nodes",
|
||||
)
|
||||
async def reload_custom_nodes(current_admin: AdminUserOrDefault) -> dict[str, str]:
|
||||
"""Triggers a reload of all custom nodes.
|
||||
|
||||
This re-scans the nodes directory and loads any new node packs.
|
||||
Already loaded packs are skipped.
|
||||
"""
|
||||
config = get_config()
|
||||
custom_nodes_path = config.custom_nodes_path
|
||||
|
||||
if not custom_nodes_path.exists():
|
||||
return {"status": "No custom nodes directory found."}
|
||||
|
||||
from invokeai.app.invocations.load_custom_nodes import load_custom_nodes
|
||||
|
||||
load_custom_nodes(custom_nodes_path, logger)
|
||||
|
||||
# Invalidate the OpenAPI schema cache so the frontend gets updated node definitions
|
||||
from invokeai.app.api_app import app
|
||||
|
||||
app.openapi_schema = None
|
||||
|
||||
return {"status": "Custom nodes reloaded successfully."}
|
||||
|
||||
|
||||
def _purge_pack_modules(pack_name: str) -> list[str]:
|
||||
"""Removes the pack's root module and all of its submodules from sys.modules.
|
||||
|
||||
After uninstall, cached submodules (e.g. `pack_name.nodes`, `pack_name.foo.bar`)
|
||||
must be evicted as well — otherwise a subsequent reinstall reuses the cached
|
||||
objects, the @invocation decorators never re-run, and the pack ends up loaded
|
||||
with zero registered nodes until a full process restart.
|
||||
"""
|
||||
prefix = f"{pack_name}."
|
||||
to_remove = [name for name in sys.modules if name == pack_name or name.startswith(prefix)]
|
||||
for name in to_remove:
|
||||
del sys.modules[name]
|
||||
return to_remove
|
||||
|
||||
|
||||
def _load_node_pack(pack_name: str, pack_dir: Path) -> None:
|
||||
"""Loads a single node pack at runtime."""
|
||||
init = pack_dir / "__init__.py"
|
||||
if not init.exists():
|
||||
return
|
||||
|
||||
if pack_name in sys.modules:
|
||||
logger.info(f"Node pack {pack_name} already loaded, skipping.")
|
||||
return
|
||||
|
||||
spec = spec_from_file_location(pack_name, init.absolute())
|
||||
if spec is None or spec.loader is None:
|
||||
logger.warning(f"Could not load {init}")
|
||||
return
|
||||
|
||||
logger.info(f"Loading node pack {pack_name}")
|
||||
module = module_from_spec(spec)
|
||||
sys.modules[spec.name] = module
|
||||
spec.loader.exec_module(module)
|
||||
|
||||
# Invalidate OpenAPI schema cache
|
||||
from invokeai.app.api_app import app
|
||||
|
||||
app.openapi_schema = None
|
||||
|
||||
logger.info(f"Successfully loaded node pack {pack_name}")
|
||||
|
||||
|
||||
def _import_workflows_from_pack(pack_dir: Path, pack_name: str, owner_user_id: str) -> list[str]:
|
||||
"""Scans a node pack directory for workflow JSON files and imports them into the workflow library.
|
||||
|
||||
A JSON file is considered a workflow if it contains 'nodes' and 'edges' keys at the top level.
|
||||
Workflows are imported as user workflows owned by the installing admin and marked public so all
|
||||
users can see them — a pack is an admin-installed shared resource, not a private asset.
|
||||
|
||||
Returns the list of workflow IDs successfully created, in import order.
|
||||
"""
|
||||
imported_ids: list[str] = []
|
||||
|
||||
# Search for .json files recursively
|
||||
for json_file in pack_dir.rglob("*.json"):
|
||||
# Skip our own manifest file
|
||||
if json_file.name == PACK_MANIFEST_FILENAME:
|
||||
continue
|
||||
try:
|
||||
with open(json_file, "r", encoding="utf-8") as f:
|
||||
data = json.load(f)
|
||||
|
||||
# Check if this looks like a workflow (must have nodes and edges)
|
||||
if not isinstance(data, dict):
|
||||
continue
|
||||
if "nodes" not in data or "edges" not in data:
|
||||
continue
|
||||
|
||||
# Ensure the workflow has a meta section with category set to "user"
|
||||
if "meta" not in data:
|
||||
data["meta"] = {"version": "3.0.0", "category": "user"}
|
||||
else:
|
||||
data["meta"]["category"] = "user"
|
||||
|
||||
# Add the node pack name to tags for discoverability (display only — uninstall
|
||||
# does not rely on this tag, since users can edit tags on their own workflows).
|
||||
existing_tags = data.get("tags", "")
|
||||
pack_tag = f"node-pack:{pack_name}"
|
||||
if pack_tag not in existing_tags:
|
||||
data["tags"] = f"{existing_tags}, {pack_tag}".strip(", ") if existing_tags else pack_tag
|
||||
|
||||
# Remove the 'id' field if present — the system will assign a new one
|
||||
data.pop("id", None)
|
||||
|
||||
# Validate and import the workflow
|
||||
workflow = WorkflowWithoutIDValidator.validate_python(data)
|
||||
created = ApiDependencies.invoker.services.workflow_records.create(
|
||||
workflow=workflow, user_id=owner_user_id, is_public=True
|
||||
)
|
||||
imported_ids.append(created.workflow_id)
|
||||
logger.info(f"Imported workflow '{workflow.name}' from node pack '{pack_name}'")
|
||||
|
||||
except Exception:
|
||||
logger.warning(f"Skipped non-workflow or invalid JSON file: {json_file}")
|
||||
continue
|
||||
|
||||
if imported_ids:
|
||||
logger.info(f"Imported {len(imported_ids)} workflow(s) from node pack '{pack_name}'")
|
||||
|
||||
return imported_ids
|
||||
|
||||
|
||||
def _write_pack_manifest(pack_dir: Path, workflow_ids: list[str]) -> None:
|
||||
"""Writes the pack manifest recording which workflow IDs were imported from the pack."""
|
||||
manifest_path = pack_dir / PACK_MANIFEST_FILENAME
|
||||
try:
|
||||
with open(manifest_path, "w", encoding="utf-8") as f:
|
||||
json.dump({"workflow_ids": workflow_ids}, f)
|
||||
except Exception:
|
||||
logger.warning(f"Failed to write pack manifest at {manifest_path}")
|
||||
|
||||
|
||||
def _read_pack_manifest(pack_dir: Path) -> list[str]:
|
||||
"""Reads workflow IDs that this pack's install recorded in its manifest.
|
||||
|
||||
Returns an empty list if the manifest is missing or malformed. We deliberately do NOT
|
||||
fall back to tag-based lookup: workflow tags are user-editable and could collide with
|
||||
unrelated workflows, so we only delete what we recorded ourselves at install time.
|
||||
"""
|
||||
manifest_path = pack_dir / PACK_MANIFEST_FILENAME
|
||||
if not manifest_path.exists():
|
||||
return []
|
||||
try:
|
||||
with open(manifest_path, "r", encoding="utf-8") as f:
|
||||
data = json.load(f)
|
||||
ids = data.get("workflow_ids", [])
|
||||
if not isinstance(ids, list):
|
||||
return []
|
||||
return [str(x) for x in ids if isinstance(x, str)]
|
||||
except Exception:
|
||||
logger.warning(f"Failed to read pack manifest at {manifest_path}")
|
||||
return []
|
||||
|
||||
|
||||
def _remove_workflows_by_ids(workflow_ids: list[str], pack_name: str) -> int:
|
||||
"""Deletes the given workflow IDs. Used during uninstall to remove only the workflows
|
||||
this pack's install recorded in its manifest.
|
||||
"""
|
||||
if not workflow_ids:
|
||||
return 0
|
||||
|
||||
removed_count = 0
|
||||
for workflow_id in workflow_ids:
|
||||
try:
|
||||
ApiDependencies.invoker.services.workflow_records.delete(workflow_id)
|
||||
removed_count += 1
|
||||
except Exception:
|
||||
logger.warning(f"Failed to remove workflow '{workflow_id}' (from node pack '{pack_name}')")
|
||||
|
||||
if removed_count > 0:
|
||||
logger.info(f"Removed {removed_count} workflow(s) from node pack '{pack_name}'")
|
||||
|
||||
return removed_count
|
||||
@@ -0,0 +1,138 @@
|
||||
# Copyright (c) 2023 Lincoln D. Stein
|
||||
"""FastAPI route for the download queue."""
|
||||
|
||||
from pathlib import Path as FsPath
|
||||
from pathlib import PurePosixPath, PureWindowsPath
|
||||
from typing import List, Optional
|
||||
|
||||
from fastapi import Body, Path, Response
|
||||
from fastapi.routing import APIRouter
|
||||
from pydantic.networks import AnyHttpUrl
|
||||
from starlette.exceptions import HTTPException
|
||||
|
||||
from invokeai.app.api.auth_dependencies import AdminUserOrDefault, CurrentUserOrDefault
|
||||
from invokeai.app.api.dependencies import ApiDependencies
|
||||
from invokeai.app.services.download import (
|
||||
DownloadJob,
|
||||
UnknownJobIDException,
|
||||
)
|
||||
|
||||
download_queue_router = APIRouter(prefix="/v1/download_queue", tags=["download_queue"])
|
||||
|
||||
|
||||
def _validate_dest(dest: str) -> str:
|
||||
"""Reject absolute paths and parent-traversal segments.
|
||||
|
||||
Accepts a relative POSIX- or Windows-style path. Returns the original string
|
||||
for the caller to wrap in `Path(...)`. Raises 400 on suspicious input so the
|
||||
download service never sees it.
|
||||
"""
|
||||
if not dest or not dest.strip():
|
||||
raise HTTPException(status_code=400, detail="Download destination must not be empty.")
|
||||
|
||||
posix = PurePosixPath(dest)
|
||||
windows = PureWindowsPath(dest)
|
||||
if posix.is_absolute() or windows.is_absolute():
|
||||
raise HTTPException(status_code=400, detail="Download destination must be a relative path.")
|
||||
|
||||
if ".." in posix.parts or ".." in windows.parts:
|
||||
raise HTTPException(status_code=400, detail="Download destination must not contain '..' segments.")
|
||||
|
||||
return dest
|
||||
|
||||
|
||||
@download_queue_router.get(
|
||||
"/",
|
||||
operation_id="list_downloads",
|
||||
)
|
||||
async def list_downloads(current_user: CurrentUserOrDefault) -> List[DownloadJob]:
|
||||
"""Get a list of active and inactive jobs."""
|
||||
queue = ApiDependencies.invoker.services.download_queue
|
||||
return queue.list_jobs()
|
||||
|
||||
|
||||
@download_queue_router.patch(
|
||||
"/",
|
||||
operation_id="prune_downloads",
|
||||
responses={
|
||||
204: {"description": "All completed jobs have been pruned"},
|
||||
400: {"description": "Bad request"},
|
||||
},
|
||||
)
|
||||
async def prune_downloads(current_user: AdminUserOrDefault) -> Response:
|
||||
"""Prune completed and errored jobs."""
|
||||
queue = ApiDependencies.invoker.services.download_queue
|
||||
queue.prune_jobs()
|
||||
return Response(status_code=204)
|
||||
|
||||
|
||||
@download_queue_router.post(
|
||||
"/i/",
|
||||
operation_id="download",
|
||||
)
|
||||
async def download(
|
||||
current_user: CurrentUserOrDefault,
|
||||
source: AnyHttpUrl = Body(description="download source"),
|
||||
dest: str = Body(description="download destination"),
|
||||
priority: int = Body(default=10, description="queue priority"),
|
||||
access_token: Optional[str] = Body(default=None, description="token for authorization to download"),
|
||||
) -> DownloadJob:
|
||||
"""Download the source URL to the file or directory indicted in dest."""
|
||||
validated_dest = _validate_dest(dest)
|
||||
queue = ApiDependencies.invoker.services.download_queue
|
||||
return queue.download(source, FsPath(validated_dest), priority, access_token)
|
||||
|
||||
|
||||
@download_queue_router.get(
|
||||
"/i/{id}",
|
||||
operation_id="get_download_job",
|
||||
responses={
|
||||
200: {"description": "Success"},
|
||||
404: {"description": "The requested download JobID could not be found"},
|
||||
},
|
||||
)
|
||||
async def get_download_job(
|
||||
current_user: CurrentUserOrDefault,
|
||||
id: int = Path(description="ID of the download job to fetch."),
|
||||
) -> DownloadJob:
|
||||
"""Get a download job using its ID."""
|
||||
try:
|
||||
job = ApiDependencies.invoker.services.download_queue.id_to_job(id)
|
||||
return job
|
||||
except UnknownJobIDException as e:
|
||||
raise HTTPException(status_code=404, detail=str(e))
|
||||
|
||||
|
||||
@download_queue_router.delete(
|
||||
"/i/{id}",
|
||||
operation_id="cancel_download_job",
|
||||
responses={
|
||||
204: {"description": "Job has been cancelled"},
|
||||
404: {"description": "The requested download JobID could not be found"},
|
||||
},
|
||||
)
|
||||
async def cancel_download_job(
|
||||
current_user: CurrentUserOrDefault,
|
||||
id: int = Path(description="ID of the download job to cancel."),
|
||||
) -> Response:
|
||||
"""Cancel a download job using its ID."""
|
||||
try:
|
||||
queue = ApiDependencies.invoker.services.download_queue
|
||||
job = queue.id_to_job(id)
|
||||
queue.cancel_job(job)
|
||||
return Response(status_code=204)
|
||||
except UnknownJobIDException as e:
|
||||
raise HTTPException(status_code=404, detail=str(e))
|
||||
|
||||
|
||||
@download_queue_router.delete(
|
||||
"/i",
|
||||
operation_id="cancel_all_download_jobs",
|
||||
responses={
|
||||
204: {"description": "Download jobs have been cancelled"},
|
||||
},
|
||||
)
|
||||
async def cancel_all_download_jobs(current_user: AdminUserOrDefault) -> Response:
|
||||
"""Cancel all download jobs."""
|
||||
ApiDependencies.invoker.services.download_queue.cancel_all_jobs()
|
||||
return Response(status_code=204)
|
||||
@@ -0,0 +1,17 @@
|
||||
from fastapi import HTTPException, status
|
||||
|
||||
from invokeai.app.api.dependencies import ApiDependencies
|
||||
|
||||
IMAGE_MOVE_MAINTENANCE_ACTIVE_DETAIL = "Image storage maintenance is active"
|
||||
|
||||
|
||||
def assert_image_move_maintenance_inactive() -> None:
|
||||
invoker = getattr(ApiDependencies, "invoker", None)
|
||||
if invoker is None:
|
||||
return
|
||||
image_moves = getattr(invoker.services, "image_moves", None)
|
||||
if image_moves is not None and image_moves.is_maintenance_active():
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_409_CONFLICT,
|
||||
detail=IMAGE_MOVE_MAINTENANCE_ACTIVE_DETAIL,
|
||||
)
|
||||
@@ -0,0 +1,94 @@
|
||||
from fastapi import HTTPException, status
|
||||
from fastapi.routing import APIRouter
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
from invokeai.app.api.auth_dependencies import AdminUserOrDefault
|
||||
from invokeai.app.api.dependencies import ApiDependencies
|
||||
from invokeai.app.services.image_moves.image_moves_default import (
|
||||
ImageMoveBackgroundOperation,
|
||||
ImageMoveBackgroundStatus,
|
||||
ImageMoveJob,
|
||||
ImageMoveJobAlreadyRunning,
|
||||
ImageMoveQueueActive,
|
||||
MoveJobState,
|
||||
)
|
||||
|
||||
image_moves_router = APIRouter(prefix="/v1/image_moves", tags=["image_moves"])
|
||||
|
||||
|
||||
class ImageMoveJobResponse(BaseModel):
|
||||
id: int = Field(description="The image move job id.")
|
||||
state: MoveJobState = Field(description="The image move job state.")
|
||||
error_message: str | None = Field(default=None, description="The last error recorded for the job, if any.")
|
||||
|
||||
|
||||
class ImageMoveStatusResponse(BaseModel):
|
||||
is_running: bool = Field(description="Whether an image move background operation is currently running.")
|
||||
operation: ImageMoveBackgroundOperation | None = Field(
|
||||
default=None, description="The active background operation, if any."
|
||||
)
|
||||
active_job_id: int | None = Field(default=None, description="The active journal job id, if any.")
|
||||
latest_job: ImageMoveJobResponse | None = Field(default=None, description="The latest journal job, if any.")
|
||||
last_error: str | None = Field(default=None, description="The last background worker error, if any.")
|
||||
needs_move_count: int = Field(description="The number of images that do not match the current subfolder strategy.")
|
||||
|
||||
|
||||
def _get_image_move_service():
|
||||
image_moves = getattr(ApiDependencies.invoker.services, "image_moves", None)
|
||||
if image_moves is None:
|
||||
raise HTTPException(status_code=status.HTTP_503_SERVICE_UNAVAILABLE, detail="Image move service unavailable")
|
||||
return image_moves
|
||||
|
||||
|
||||
def _job_to_response(job: ImageMoveJob | None) -> ImageMoveJobResponse | None:
|
||||
if job is None:
|
||||
return None
|
||||
return ImageMoveJobResponse(id=job.id, state=job.state, error_message=job.error_message)
|
||||
|
||||
|
||||
def _status_to_response(service_status: ImageMoveBackgroundStatus | dict) -> ImageMoveStatusResponse:
|
||||
if isinstance(service_status, dict):
|
||||
return ImageMoveStatusResponse(**service_status)
|
||||
return ImageMoveStatusResponse(
|
||||
is_running=service_status.is_running,
|
||||
operation=service_status.operation,
|
||||
active_job_id=service_status.active_job_id,
|
||||
latest_job=_job_to_response(service_status.latest_job),
|
||||
last_error=service_status.last_error,
|
||||
needs_move_count=service_status.needs_move_count,
|
||||
)
|
||||
|
||||
|
||||
@image_moves_router.post(
|
||||
"/start",
|
||||
operation_id="start_image_move",
|
||||
response_model=ImageMoveStatusResponse,
|
||||
status_code=status.HTTP_202_ACCEPTED,
|
||||
)
|
||||
async def start_image_move(_: AdminUserOrDefault) -> ImageMoveStatusResponse:
|
||||
try:
|
||||
return _status_to_response(_get_image_move_service().start_background_move_all())
|
||||
except (ImageMoveJobAlreadyRunning, ImageMoveQueueActive) as e:
|
||||
raise HTTPException(status_code=status.HTTP_409_CONFLICT, detail=str(e)) from e
|
||||
|
||||
|
||||
@image_moves_router.post(
|
||||
"/recover",
|
||||
operation_id="start_image_move_recovery",
|
||||
response_model=ImageMoveStatusResponse,
|
||||
status_code=status.HTTP_202_ACCEPTED,
|
||||
)
|
||||
async def start_image_move_recovery(_: AdminUserOrDefault) -> ImageMoveStatusResponse:
|
||||
try:
|
||||
return _status_to_response(_get_image_move_service().start_background_recovery())
|
||||
except ImageMoveJobAlreadyRunning as e:
|
||||
raise HTTPException(status_code=status.HTTP_409_CONFLICT, detail=str(e)) from e
|
||||
|
||||
|
||||
@image_moves_router.get(
|
||||
"/status",
|
||||
operation_id="get_image_move_status",
|
||||
response_model=ImageMoveStatusResponse,
|
||||
)
|
||||
async def get_image_move_status(_: AdminUserOrDefault) -> ImageMoveStatusResponse:
|
||||
return _status_to_response(_get_image_move_service().get_background_status())
|
||||
@@ -0,0 +1,789 @@
|
||||
import io
|
||||
import json
|
||||
import traceback
|
||||
from typing import ClassVar, Optional
|
||||
|
||||
from fastapi import BackgroundTasks, Body, HTTPException, Path, Query, Request, Response, UploadFile
|
||||
from fastapi.responses import FileResponse
|
||||
from fastapi.routing import APIRouter
|
||||
from PIL import Image
|
||||
from pydantic import BaseModel, Field, model_validator
|
||||
|
||||
from invokeai.app.api.auth_dependencies import CurrentUserOrDefault
|
||||
from invokeai.app.api.dependencies import ApiDependencies
|
||||
from invokeai.app.api.extract_metadata_from_image import extract_metadata_from_image
|
||||
from invokeai.app.api.routers._access import (
|
||||
assert_board_read_access as _assert_board_read_access,
|
||||
)
|
||||
from invokeai.app.api.routers._access import (
|
||||
assert_image_owner as _assert_image_owner,
|
||||
)
|
||||
from invokeai.app.api.routers._access import (
|
||||
assert_image_read_access as _assert_image_read_access,
|
||||
)
|
||||
from invokeai.app.api.routers.image_move_maintenance import assert_image_move_maintenance_inactive
|
||||
from invokeai.app.invocations.fields import MetadataField
|
||||
from invokeai.app.services.image_records.image_records_common import (
|
||||
ImageCategory,
|
||||
ImageNamesResult,
|
||||
ImageRecordChanges,
|
||||
ResourceOrigin,
|
||||
)
|
||||
from invokeai.app.services.images.images_common import (
|
||||
DeleteImagesResult,
|
||||
ImageDTO,
|
||||
ImageUrlsDTO,
|
||||
StarredImagesResult,
|
||||
UnstarredImagesResult,
|
||||
)
|
||||
from invokeai.app.services.shared.pagination import OffsetPaginatedResults
|
||||
from invokeai.app.services.shared.sqlite.sqlite_common import SQLiteDirection
|
||||
from invokeai.app.util.controlnet_utils import heuristic_resize_fast
|
||||
from invokeai.backend.image_util.util import np_to_pil, pil_to_np
|
||||
|
||||
images_router = APIRouter(prefix="/v1/images", tags=["images"])
|
||||
|
||||
|
||||
# images are immutable; set a high max-age
|
||||
IMAGE_MAX_AGE = 31536000
|
||||
|
||||
|
||||
class ResizeToDimensions(BaseModel):
|
||||
width: int = Field(..., gt=0)
|
||||
height: int = Field(..., gt=0)
|
||||
|
||||
MAX_SIZE: ClassVar[int] = 4096 * 4096
|
||||
|
||||
@model_validator(mode="after")
|
||||
def validate_total_output_size(self):
|
||||
if self.width * self.height > self.MAX_SIZE:
|
||||
raise ValueError(f"Max total output size for resizing is {self.MAX_SIZE} pixels")
|
||||
return self
|
||||
|
||||
|
||||
@images_router.post(
|
||||
"/upload",
|
||||
operation_id="upload_image",
|
||||
responses={
|
||||
201: {"description": "The image was uploaded successfully"},
|
||||
415: {"description": "Image upload failed"},
|
||||
},
|
||||
status_code=201,
|
||||
response_model=ImageDTO,
|
||||
)
|
||||
async def upload_image(
|
||||
current_user: CurrentUserOrDefault,
|
||||
file: UploadFile,
|
||||
request: Request,
|
||||
response: Response,
|
||||
image_category: ImageCategory = Query(description="The category of the image"),
|
||||
is_intermediate: bool = Query(description="Whether this is an intermediate image"),
|
||||
board_id: Optional[str] = Query(default=None, description="The board to add this image to, if any"),
|
||||
session_id: Optional[str] = Query(default=None, description="The session ID associated with this upload, if any"),
|
||||
crop_visible: Optional[bool] = Query(default=False, description="Whether to crop the image"),
|
||||
resize_to: Optional[str] = Body(
|
||||
default=None,
|
||||
description=f"Dimensions to resize the image to, must be stringified tuple of 2 integers. Max total pixel count: {ResizeToDimensions.MAX_SIZE}",
|
||||
examples=['"[1024,1024]"'],
|
||||
),
|
||||
metadata: Optional[str] = Body(
|
||||
default=None,
|
||||
description="The metadata to associate with the image, must be a stringified JSON dict",
|
||||
embed=True,
|
||||
),
|
||||
) -> ImageDTO:
|
||||
"""Uploads an image for the current user"""
|
||||
# If uploading into a board, verify the user has write access.
|
||||
# Public boards allow uploads from any authenticated user.
|
||||
if board_id is not None:
|
||||
from invokeai.app.services.board_records.board_records_common import BoardVisibility
|
||||
|
||||
try:
|
||||
board = ApiDependencies.invoker.services.boards.get_dto(board_id=board_id)
|
||||
except Exception:
|
||||
raise HTTPException(status_code=404, detail="Board not found")
|
||||
if (
|
||||
not current_user.is_admin
|
||||
and board.user_id != current_user.user_id
|
||||
and board.board_visibility != BoardVisibility.Public
|
||||
):
|
||||
raise HTTPException(status_code=403, detail="Not authorized to upload to this board")
|
||||
|
||||
assert_image_move_maintenance_inactive()
|
||||
|
||||
if not file.content_type or not file.content_type.startswith("image"):
|
||||
raise HTTPException(status_code=415, detail="Not an image")
|
||||
|
||||
contents = await file.read()
|
||||
try:
|
||||
pil_image = Image.open(io.BytesIO(contents))
|
||||
except Exception:
|
||||
ApiDependencies.invoker.services.logger.error(traceback.format_exc())
|
||||
raise HTTPException(status_code=415, detail="Failed to read image")
|
||||
|
||||
if crop_visible:
|
||||
try:
|
||||
bbox = pil_image.getbbox()
|
||||
pil_image = pil_image.crop(bbox)
|
||||
except Exception:
|
||||
raise HTTPException(status_code=500, detail="Failed to crop image")
|
||||
|
||||
if resize_to:
|
||||
try:
|
||||
dims = json.loads(resize_to)
|
||||
resize_dims = ResizeToDimensions(**dims)
|
||||
except Exception:
|
||||
raise HTTPException(status_code=400, detail="Invalid resize_to format or size")
|
||||
|
||||
try:
|
||||
# heuristic_resize_fast expects an RGB or RGBA image
|
||||
pil_rgba = pil_image.convert("RGBA")
|
||||
np_image = pil_to_np(pil_rgba)
|
||||
np_image = heuristic_resize_fast(np_image, (resize_dims.width, resize_dims.height))
|
||||
pil_image = np_to_pil(np_image)
|
||||
except Exception:
|
||||
raise HTTPException(status_code=500, detail="Failed to resize image")
|
||||
|
||||
extracted_metadata = extract_metadata_from_image(
|
||||
pil_image=pil_image,
|
||||
invokeai_metadata_override=metadata,
|
||||
invokeai_workflow_override=None,
|
||||
invokeai_graph_override=None,
|
||||
logger=ApiDependencies.invoker.services.logger,
|
||||
)
|
||||
|
||||
try:
|
||||
image_dto = ApiDependencies.invoker.services.images.create(
|
||||
image=pil_image,
|
||||
image_origin=ResourceOrigin.EXTERNAL,
|
||||
image_category=image_category,
|
||||
session_id=session_id,
|
||||
board_id=board_id,
|
||||
metadata=extracted_metadata.invokeai_metadata,
|
||||
workflow=extracted_metadata.invokeai_workflow,
|
||||
graph=extracted_metadata.invokeai_graph,
|
||||
is_intermediate=is_intermediate,
|
||||
user_id=current_user.user_id,
|
||||
)
|
||||
|
||||
response.status_code = 201
|
||||
response.headers["Location"] = image_dto.image_url
|
||||
|
||||
return image_dto
|
||||
except Exception:
|
||||
ApiDependencies.invoker.services.logger.error(traceback.format_exc())
|
||||
raise HTTPException(status_code=500, detail="Failed to create image")
|
||||
|
||||
|
||||
class ImageUploadEntry(BaseModel):
|
||||
image_dto: ImageDTO = Body(description="The image DTO")
|
||||
presigned_url: str = Body(description="The URL to get the presigned URL for the image upload")
|
||||
|
||||
|
||||
@images_router.post("/", operation_id="create_image_upload_entry")
|
||||
async def create_image_upload_entry(
|
||||
width: int = Body(description="The width of the image"),
|
||||
height: int = Body(description="The height of the image"),
|
||||
board_id: Optional[str] = Body(default=None, description="The board to add this image to, if any"),
|
||||
) -> ImageUploadEntry:
|
||||
"""Uploads an image from a URL, not implemented"""
|
||||
|
||||
raise HTTPException(status_code=501, detail="Not implemented")
|
||||
|
||||
|
||||
@images_router.delete("/i/{image_name}", operation_id="delete_image", response_model=DeleteImagesResult)
|
||||
async def delete_image(
|
||||
current_user: CurrentUserOrDefault,
|
||||
image_name: str = Path(description="The name of the image to delete"),
|
||||
) -> DeleteImagesResult:
|
||||
"""Deletes an image"""
|
||||
_assert_image_owner(image_name, current_user)
|
||||
assert_image_move_maintenance_inactive()
|
||||
|
||||
deleted_images: set[str] = set()
|
||||
affected_boards: set[str] = set()
|
||||
|
||||
try:
|
||||
image_dto = ApiDependencies.invoker.services.images.get_dto(image_name)
|
||||
board_id = image_dto.board_id or "none"
|
||||
ApiDependencies.invoker.services.images.delete(image_name)
|
||||
deleted_images.add(image_name)
|
||||
affected_boards.add(board_id)
|
||||
except Exception:
|
||||
# TODO: Does this need any exception handling at all?
|
||||
pass
|
||||
|
||||
return DeleteImagesResult(
|
||||
deleted_images=list(deleted_images),
|
||||
affected_boards=list(affected_boards),
|
||||
)
|
||||
|
||||
|
||||
@images_router.delete("/intermediates", operation_id="clear_intermediates")
|
||||
async def clear_intermediates(
|
||||
current_user: CurrentUserOrDefault,
|
||||
) -> int:
|
||||
"""Clears all intermediates. Requires admin."""
|
||||
if not current_user.is_admin:
|
||||
raise HTTPException(status_code=403, detail="Only admins can clear all intermediates")
|
||||
assert_image_move_maintenance_inactive()
|
||||
|
||||
try:
|
||||
count_deleted = ApiDependencies.invoker.services.images.delete_intermediates()
|
||||
return count_deleted
|
||||
except Exception:
|
||||
raise HTTPException(status_code=500, detail="Failed to clear intermediates")
|
||||
|
||||
|
||||
@images_router.get("/intermediates", operation_id="get_intermediates_count")
|
||||
async def get_intermediates_count(
|
||||
current_user: CurrentUserOrDefault,
|
||||
) -> int:
|
||||
"""Gets the count of intermediate images. Non-admin users only see their own intermediates."""
|
||||
|
||||
try:
|
||||
user_id = None if current_user.is_admin else current_user.user_id
|
||||
return ApiDependencies.invoker.services.images.get_intermediates_count(user_id=user_id)
|
||||
except Exception:
|
||||
raise HTTPException(status_code=500, detail="Failed to get intermediates")
|
||||
|
||||
|
||||
@images_router.patch(
|
||||
"/i/{image_name}",
|
||||
operation_id="update_image",
|
||||
response_model=ImageDTO,
|
||||
)
|
||||
async def update_image(
|
||||
current_user: CurrentUserOrDefault,
|
||||
image_name: str = Path(description="The name of the image to update"),
|
||||
image_changes: ImageRecordChanges = Body(description="The changes to apply to the image"),
|
||||
) -> ImageDTO:
|
||||
"""Updates an image"""
|
||||
_assert_image_owner(image_name, current_user)
|
||||
assert_image_move_maintenance_inactive()
|
||||
|
||||
try:
|
||||
return ApiDependencies.invoker.services.images.update(image_name, image_changes)
|
||||
except Exception:
|
||||
raise HTTPException(status_code=400, detail="Failed to update image")
|
||||
|
||||
|
||||
@images_router.get(
|
||||
"/i/{image_name}",
|
||||
operation_id="get_image_dto",
|
||||
response_model=ImageDTO,
|
||||
)
|
||||
async def get_image_dto(
|
||||
current_user: CurrentUserOrDefault,
|
||||
image_name: str = Path(description="The name of image to get"),
|
||||
) -> ImageDTO:
|
||||
"""Gets an image's DTO"""
|
||||
_assert_image_read_access(image_name, current_user)
|
||||
|
||||
try:
|
||||
return ApiDependencies.invoker.services.images.get_dto(image_name)
|
||||
except Exception:
|
||||
raise HTTPException(status_code=404)
|
||||
|
||||
|
||||
@images_router.get(
|
||||
"/i/{image_name}/metadata",
|
||||
operation_id="get_image_metadata",
|
||||
response_model=Optional[MetadataField],
|
||||
)
|
||||
async def get_image_metadata(
|
||||
current_user: CurrentUserOrDefault,
|
||||
image_name: str = Path(description="The name of image to get"),
|
||||
) -> Optional[MetadataField]:
|
||||
"""Gets an image's metadata"""
|
||||
_assert_image_read_access(image_name, current_user)
|
||||
|
||||
try:
|
||||
return ApiDependencies.invoker.services.images.get_metadata(image_name)
|
||||
except Exception:
|
||||
raise HTTPException(status_code=404)
|
||||
|
||||
|
||||
class WorkflowAndGraphResponse(BaseModel):
|
||||
workflow: Optional[str] = Field(description="The workflow used to generate the image, as stringified JSON")
|
||||
graph: Optional[str] = Field(description="The graph used to generate the image, as stringified JSON")
|
||||
|
||||
|
||||
@images_router.get(
|
||||
"/i/{image_name}/workflow", operation_id="get_image_workflow", response_model=WorkflowAndGraphResponse
|
||||
)
|
||||
async def get_image_workflow(
|
||||
current_user: CurrentUserOrDefault,
|
||||
image_name: str = Path(description="The name of image whose workflow to get"),
|
||||
) -> WorkflowAndGraphResponse:
|
||||
_assert_image_read_access(image_name, current_user)
|
||||
assert_image_move_maintenance_inactive()
|
||||
|
||||
try:
|
||||
workflow = ApiDependencies.invoker.services.images.get_workflow(image_name)
|
||||
graph = ApiDependencies.invoker.services.images.get_graph(image_name)
|
||||
return WorkflowAndGraphResponse(workflow=workflow, graph=graph)
|
||||
except Exception:
|
||||
raise HTTPException(status_code=404)
|
||||
|
||||
|
||||
@images_router.get(
|
||||
"/i/{image_name}/full",
|
||||
operation_id="get_image_full",
|
||||
response_class=Response,
|
||||
responses={
|
||||
200: {
|
||||
"description": "Return the full-resolution image",
|
||||
"content": {"image/png": {}},
|
||||
},
|
||||
404: {"description": "Image not found"},
|
||||
},
|
||||
)
|
||||
@images_router.head(
|
||||
"/i/{image_name}/full",
|
||||
operation_id="get_image_full_head",
|
||||
response_class=Response,
|
||||
responses={
|
||||
200: {
|
||||
"description": "Return the full-resolution image",
|
||||
"content": {"image/png": {}},
|
||||
},
|
||||
404: {"description": "Image not found"},
|
||||
},
|
||||
)
|
||||
async def get_image_full(
|
||||
image_name: str = Path(description="The name of full-resolution image file to get"),
|
||||
) -> Response:
|
||||
"""Gets a full-resolution image file.
|
||||
|
||||
This endpoint is intentionally unauthenticated because browsers load images
|
||||
via <img src> tags which cannot send Bearer tokens. Image names are UUIDs,
|
||||
providing security through unguessability. Returns 409 while image storage
|
||||
maintenance is active.
|
||||
"""
|
||||
assert_image_move_maintenance_inactive()
|
||||
|
||||
try:
|
||||
path = ApiDependencies.invoker.services.images.get_path(image_name)
|
||||
with open(path, "rb") as f:
|
||||
content = f.read()
|
||||
response = Response(content, media_type="image/png")
|
||||
response.headers["Cache-Control"] = f"max-age={IMAGE_MAX_AGE}"
|
||||
response.headers["Content-Disposition"] = f'inline; filename="{image_name}"'
|
||||
return response
|
||||
except Exception:
|
||||
raise HTTPException(status_code=404)
|
||||
|
||||
|
||||
@images_router.get(
|
||||
"/i/{image_name}/thumbnail",
|
||||
operation_id="get_image_thumbnail",
|
||||
response_class=Response,
|
||||
responses={
|
||||
200: {
|
||||
"description": "Return the image thumbnail",
|
||||
"content": {"image/webp": {}},
|
||||
},
|
||||
404: {"description": "Image not found"},
|
||||
},
|
||||
)
|
||||
async def get_image_thumbnail(
|
||||
image_name: str = Path(description="The name of thumbnail image file to get"),
|
||||
) -> Response:
|
||||
"""Gets a thumbnail image file.
|
||||
|
||||
This endpoint is intentionally unauthenticated because browsers load images
|
||||
via <img src> tags which cannot send Bearer tokens. Image names are UUIDs,
|
||||
providing security through unguessability. Returns 409 while image storage
|
||||
maintenance is active.
|
||||
"""
|
||||
assert_image_move_maintenance_inactive()
|
||||
|
||||
try:
|
||||
path = ApiDependencies.invoker.services.images.get_path(image_name, thumbnail=True)
|
||||
with open(path, "rb") as f:
|
||||
content = f.read()
|
||||
response = Response(content, media_type="image/webp")
|
||||
response.headers["Cache-Control"] = f"max-age={IMAGE_MAX_AGE}"
|
||||
return response
|
||||
except Exception:
|
||||
raise HTTPException(status_code=404)
|
||||
|
||||
|
||||
@images_router.get(
|
||||
"/i/{image_name}/urls",
|
||||
operation_id="get_image_urls",
|
||||
response_model=ImageUrlsDTO,
|
||||
)
|
||||
async def get_image_urls(
|
||||
current_user: CurrentUserOrDefault,
|
||||
image_name: str = Path(description="The name of the image whose URL to get"),
|
||||
) -> ImageUrlsDTO:
|
||||
"""Gets an image and thumbnail URL"""
|
||||
_assert_image_read_access(image_name, current_user)
|
||||
|
||||
try:
|
||||
image_url = ApiDependencies.invoker.services.images.get_url(image_name)
|
||||
thumbnail_url = ApiDependencies.invoker.services.images.get_url(image_name, thumbnail=True)
|
||||
return ImageUrlsDTO(
|
||||
image_name=image_name,
|
||||
image_url=image_url,
|
||||
thumbnail_url=thumbnail_url,
|
||||
)
|
||||
except Exception:
|
||||
raise HTTPException(status_code=404)
|
||||
|
||||
|
||||
@images_router.get(
|
||||
"/",
|
||||
operation_id="list_image_dtos",
|
||||
response_model=OffsetPaginatedResults[ImageDTO],
|
||||
)
|
||||
async def list_image_dtos(
|
||||
current_user: CurrentUserOrDefault,
|
||||
image_origin: Optional[ResourceOrigin] = Query(default=None, description="The origin of images to list."),
|
||||
categories: Optional[list[ImageCategory]] = Query(default=None, description="The categories of image to include."),
|
||||
is_intermediate: Optional[bool] = Query(default=None, description="Whether to list intermediate images."),
|
||||
board_id: Optional[str] = Query(
|
||||
default=None,
|
||||
description="The board id to filter by. Use 'none' to find images without a board.",
|
||||
),
|
||||
offset: int = Query(default=0, description="The page offset"),
|
||||
limit: int = Query(default=10, description="The number of images per page"),
|
||||
order_dir: SQLiteDirection = Query(default=SQLiteDirection.Descending, description="The order of sort"),
|
||||
starred_first: bool = Query(default=True, description="Whether to sort by starred images first"),
|
||||
search_term: Optional[str] = Query(default=None, description="The term to search for"),
|
||||
) -> OffsetPaginatedResults[ImageDTO]:
|
||||
"""Gets a list of image DTOs for the current user"""
|
||||
|
||||
# Validate that the caller can read from this board before listing its images.
|
||||
# "none" is a sentinel for uncategorized images and is handled by the SQL layer.
|
||||
if board_id is not None and board_id != "none":
|
||||
_assert_board_read_access(board_id, current_user)
|
||||
|
||||
image_dtos = ApiDependencies.invoker.services.images.get_many(
|
||||
offset,
|
||||
limit,
|
||||
starred_first,
|
||||
order_dir,
|
||||
image_origin,
|
||||
categories,
|
||||
is_intermediate,
|
||||
board_id,
|
||||
search_term,
|
||||
current_user.user_id,
|
||||
)
|
||||
|
||||
return image_dtos
|
||||
|
||||
|
||||
@images_router.post("/delete", operation_id="delete_images_from_list", response_model=DeleteImagesResult)
|
||||
async def delete_images_from_list(
|
||||
current_user: CurrentUserOrDefault,
|
||||
image_names: list[str] = Body(description="The list of names of images to delete", embed=True),
|
||||
) -> DeleteImagesResult:
|
||||
try:
|
||||
assert_image_move_maintenance_inactive()
|
||||
except HTTPException:
|
||||
for image_name in image_names:
|
||||
_assert_image_owner(image_name, current_user)
|
||||
raise
|
||||
|
||||
try:
|
||||
deleted_images: set[str] = set()
|
||||
affected_boards: set[str] = set()
|
||||
for image_name in image_names:
|
||||
try:
|
||||
_assert_image_owner(image_name, current_user)
|
||||
image_dto = ApiDependencies.invoker.services.images.get_dto(image_name)
|
||||
board_id = image_dto.board_id or "none"
|
||||
ApiDependencies.invoker.services.images.delete(image_name)
|
||||
deleted_images.add(image_name)
|
||||
affected_boards.add(board_id)
|
||||
except HTTPException:
|
||||
raise
|
||||
except Exception:
|
||||
pass
|
||||
return DeleteImagesResult(
|
||||
deleted_images=list(deleted_images),
|
||||
affected_boards=list(affected_boards),
|
||||
)
|
||||
except HTTPException:
|
||||
raise
|
||||
except Exception:
|
||||
raise HTTPException(status_code=500, detail="Failed to delete images")
|
||||
|
||||
|
||||
@images_router.delete("/uncategorized", operation_id="delete_uncategorized_images", response_model=DeleteImagesResult)
|
||||
async def delete_uncategorized_images(
|
||||
current_user: CurrentUserOrDefault,
|
||||
) -> DeleteImagesResult:
|
||||
"""Deletes all uncategorized images owned by the current user (or all if admin)"""
|
||||
assert_image_move_maintenance_inactive()
|
||||
|
||||
image_names = ApiDependencies.invoker.services.board_images.get_all_board_image_names_for_board(
|
||||
board_id="none", categories=None, is_intermediate=None
|
||||
)
|
||||
|
||||
try:
|
||||
deleted_images: set[str] = set()
|
||||
affected_boards: set[str] = set()
|
||||
for image_name in image_names:
|
||||
try:
|
||||
_assert_image_owner(image_name, current_user)
|
||||
ApiDependencies.invoker.services.images.delete(image_name)
|
||||
deleted_images.add(image_name)
|
||||
affected_boards.add("none")
|
||||
except HTTPException:
|
||||
# Skip images not owned by the current user
|
||||
pass
|
||||
except Exception:
|
||||
pass
|
||||
return DeleteImagesResult(
|
||||
deleted_images=list(deleted_images),
|
||||
affected_boards=list(affected_boards),
|
||||
)
|
||||
except Exception:
|
||||
raise HTTPException(status_code=500, detail="Failed to delete images")
|
||||
|
||||
|
||||
class ImagesUpdatedFromListResult(BaseModel):
|
||||
updated_image_names: list[str] = Field(description="The image names that were updated")
|
||||
|
||||
|
||||
@images_router.post("/star", operation_id="star_images_in_list", response_model=StarredImagesResult)
|
||||
async def star_images_in_list(
|
||||
current_user: CurrentUserOrDefault,
|
||||
image_names: list[str] = Body(description="The list of names of images to star", embed=True),
|
||||
) -> StarredImagesResult:
|
||||
try:
|
||||
assert_image_move_maintenance_inactive()
|
||||
except HTTPException:
|
||||
for image_name in image_names:
|
||||
_assert_image_owner(image_name, current_user)
|
||||
raise
|
||||
|
||||
try:
|
||||
starred_images: set[str] = set()
|
||||
affected_boards: set[str] = set()
|
||||
for image_name in image_names:
|
||||
try:
|
||||
_assert_image_owner(image_name, current_user)
|
||||
updated_image_dto = ApiDependencies.invoker.services.images.update(
|
||||
image_name, changes=ImageRecordChanges(starred=True)
|
||||
)
|
||||
starred_images.add(image_name)
|
||||
affected_boards.add(updated_image_dto.board_id or "none")
|
||||
except HTTPException:
|
||||
raise
|
||||
except Exception:
|
||||
pass
|
||||
return StarredImagesResult(
|
||||
starred_images=list(starred_images),
|
||||
affected_boards=list(affected_boards),
|
||||
)
|
||||
except HTTPException:
|
||||
raise
|
||||
except Exception:
|
||||
raise HTTPException(status_code=500, detail="Failed to star images")
|
||||
|
||||
|
||||
@images_router.post("/unstar", operation_id="unstar_images_in_list", response_model=UnstarredImagesResult)
|
||||
async def unstar_images_in_list(
|
||||
current_user: CurrentUserOrDefault,
|
||||
image_names: list[str] = Body(description="The list of names of images to unstar", embed=True),
|
||||
) -> UnstarredImagesResult:
|
||||
try:
|
||||
assert_image_move_maintenance_inactive()
|
||||
except HTTPException:
|
||||
for image_name in image_names:
|
||||
_assert_image_owner(image_name, current_user)
|
||||
raise
|
||||
|
||||
try:
|
||||
unstarred_images: set[str] = set()
|
||||
affected_boards: set[str] = set()
|
||||
for image_name in image_names:
|
||||
try:
|
||||
_assert_image_owner(image_name, current_user)
|
||||
updated_image_dto = ApiDependencies.invoker.services.images.update(
|
||||
image_name, changes=ImageRecordChanges(starred=False)
|
||||
)
|
||||
unstarred_images.add(image_name)
|
||||
affected_boards.add(updated_image_dto.board_id or "none")
|
||||
except HTTPException:
|
||||
raise
|
||||
except Exception:
|
||||
pass
|
||||
return UnstarredImagesResult(
|
||||
unstarred_images=list(unstarred_images),
|
||||
affected_boards=list(affected_boards),
|
||||
)
|
||||
except HTTPException:
|
||||
raise
|
||||
except Exception:
|
||||
raise HTTPException(status_code=500, detail="Failed to unstar images")
|
||||
|
||||
|
||||
class ImagesDownloaded(BaseModel):
|
||||
response: Optional[str] = Field(
|
||||
default=None, description="The message to display to the user when images begin downloading"
|
||||
)
|
||||
bulk_download_item_name: Optional[str] = Field(
|
||||
default=None, description="The name of the bulk download item for which events will be emitted"
|
||||
)
|
||||
|
||||
|
||||
@images_router.post(
|
||||
"/download", operation_id="download_images_from_list", response_model=ImagesDownloaded, status_code=202
|
||||
)
|
||||
async def download_images_from_list(
|
||||
current_user: CurrentUserOrDefault,
|
||||
background_tasks: BackgroundTasks,
|
||||
image_names: Optional[list[str]] = Body(
|
||||
default=None, description="The list of names of images to download", embed=True
|
||||
),
|
||||
board_id: Optional[str] = Body(
|
||||
default=None, description="The board from which image should be downloaded", embed=True
|
||||
),
|
||||
) -> ImagesDownloaded:
|
||||
if (image_names is None or len(image_names) == 0) and board_id is None:
|
||||
raise HTTPException(status_code=400, detail="No images or board id specified.")
|
||||
|
||||
# Validate that the caller can read every image they are requesting.
|
||||
# For a board_id request, check board visibility; for explicit image names,
|
||||
# check each image individually.
|
||||
if board_id:
|
||||
_assert_board_read_access(board_id, current_user)
|
||||
if image_names:
|
||||
for name in image_names:
|
||||
_assert_image_read_access(name, current_user)
|
||||
|
||||
assert_image_move_maintenance_inactive()
|
||||
|
||||
bulk_download_item_id: str = ApiDependencies.invoker.services.bulk_download.generate_item_id(board_id)
|
||||
|
||||
background_tasks.add_task(
|
||||
ApiDependencies.invoker.services.bulk_download.handler,
|
||||
image_names,
|
||||
board_id,
|
||||
bulk_download_item_id,
|
||||
current_user.user_id,
|
||||
)
|
||||
return ImagesDownloaded(bulk_download_item_name=bulk_download_item_id + ".zip")
|
||||
|
||||
|
||||
@images_router.api_route(
|
||||
"/download/{bulk_download_item_name}",
|
||||
methods=["GET"],
|
||||
operation_id="get_bulk_download_item",
|
||||
response_class=Response,
|
||||
responses={
|
||||
200: {
|
||||
"description": "Return the complete bulk download item",
|
||||
"content": {"application/zip": {}},
|
||||
},
|
||||
404: {"description": "Image not found"},
|
||||
},
|
||||
)
|
||||
async def get_bulk_download_item(
|
||||
current_user: CurrentUserOrDefault,
|
||||
background_tasks: BackgroundTasks,
|
||||
bulk_download_item_name: str = Path(description="The bulk_download_item_name of the bulk download item to get"),
|
||||
) -> FileResponse:
|
||||
"""Gets a bulk download zip file.
|
||||
|
||||
Requires authentication. The caller must be the user who initiated the
|
||||
download (tracked by the bulk download service) or an admin.
|
||||
"""
|
||||
try:
|
||||
# Verify the caller owns this download (or is an admin)
|
||||
owner = ApiDependencies.invoker.services.bulk_download.get_owner(bulk_download_item_name)
|
||||
if owner is not None and owner != current_user.user_id and not current_user.is_admin:
|
||||
raise HTTPException(status_code=403, detail="Not authorized to access this download")
|
||||
|
||||
path = ApiDependencies.invoker.services.bulk_download.get_path(bulk_download_item_name)
|
||||
|
||||
response = FileResponse(
|
||||
path,
|
||||
media_type="application/zip",
|
||||
filename=bulk_download_item_name,
|
||||
content_disposition_type="inline",
|
||||
)
|
||||
response.headers["Cache-Control"] = f"max-age={IMAGE_MAX_AGE}"
|
||||
background_tasks.add_task(ApiDependencies.invoker.services.bulk_download.delete, bulk_download_item_name)
|
||||
return response
|
||||
except HTTPException:
|
||||
raise
|
||||
except Exception:
|
||||
raise HTTPException(status_code=404)
|
||||
|
||||
|
||||
@images_router.get("/names", operation_id="get_image_names")
|
||||
async def get_image_names(
|
||||
current_user: CurrentUserOrDefault,
|
||||
image_origin: Optional[ResourceOrigin] = Query(default=None, description="The origin of images to list."),
|
||||
categories: Optional[list[ImageCategory]] = Query(default=None, description="The categories of image to include."),
|
||||
is_intermediate: Optional[bool] = Query(default=None, description="Whether to list intermediate images."),
|
||||
board_id: Optional[str] = Query(
|
||||
default=None,
|
||||
description="The board id to filter by. Use 'none' to find images without a board.",
|
||||
),
|
||||
order_dir: SQLiteDirection = Query(default=SQLiteDirection.Descending, description="The order of sort"),
|
||||
starred_first: bool = Query(default=True, description="Whether to sort by starred images first"),
|
||||
search_term: Optional[str] = Query(default=None, description="The term to search for"),
|
||||
) -> ImageNamesResult:
|
||||
"""Gets ordered list of image names with metadata for optimistic updates"""
|
||||
|
||||
# Validate that the caller can read from this board before listing its images.
|
||||
if board_id is not None and board_id != "none":
|
||||
_assert_board_read_access(board_id, current_user)
|
||||
|
||||
try:
|
||||
result = ApiDependencies.invoker.services.images.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=current_user.user_id,
|
||||
is_admin=current_user.is_admin,
|
||||
)
|
||||
return result
|
||||
except Exception:
|
||||
raise HTTPException(status_code=500, detail="Failed to get image names")
|
||||
|
||||
|
||||
@images_router.post(
|
||||
"/images_by_names",
|
||||
operation_id="get_images_by_names",
|
||||
responses={200: {"model": list[ImageDTO]}},
|
||||
)
|
||||
async def get_images_by_names(
|
||||
current_user: CurrentUserOrDefault,
|
||||
image_names: list[str] = Body(embed=True, description="Object containing list of image names to fetch DTOs for"),
|
||||
) -> list[ImageDTO]:
|
||||
"""Gets image DTOs for the specified image names. Maintains order of input names."""
|
||||
|
||||
try:
|
||||
image_service = ApiDependencies.invoker.services.images
|
||||
|
||||
# Fetch DTOs preserving the order of requested names
|
||||
image_dtos: list[ImageDTO] = []
|
||||
for name in image_names:
|
||||
try:
|
||||
_assert_image_read_access(name, current_user)
|
||||
dto = image_service.get_dto(name)
|
||||
image_dtos.append(dto)
|
||||
except HTTPException:
|
||||
# Skip images the user is not authorized to view
|
||||
continue
|
||||
except Exception:
|
||||
# Skip missing images - they may have been deleted between name fetch and DTO fetch
|
||||
continue
|
||||
|
||||
return image_dtos
|
||||
except Exception:
|
||||
raise HTTPException(status_code=500, detail="Failed to get image DTOs")
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,210 @@
|
||||
"""FastAPI route for model relationship records."""
|
||||
|
||||
from typing import List
|
||||
|
||||
from fastapi import APIRouter, Body, HTTPException, Path, status
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
from invokeai.app.api.auth_dependencies import AdminUserOrDefault, CurrentUserOrDefault
|
||||
from invokeai.app.api.dependencies import ApiDependencies
|
||||
|
||||
model_relationships_router = APIRouter(prefix="/v1/model_relationships", tags=["model_relationships"])
|
||||
|
||||
# === Schemas ===
|
||||
|
||||
|
||||
class ModelRelationshipCreateRequest(BaseModel):
|
||||
model_key_1: str = Field(
|
||||
...,
|
||||
description="The key of the first model in the relationship",
|
||||
examples=[
|
||||
"aa3b247f-90c9-4416-bfcd-aeaa57a5339e",
|
||||
"ac32b914-10ab-496e-a24a-3068724b9c35",
|
||||
"d944abfd-c7c3-42e2-a4ff-da640b29b8b4",
|
||||
"b1c2d3e4-f5a6-7890-abcd-ef1234567890",
|
||||
"12345678-90ab-cdef-1234-567890abcdef",
|
||||
"fedcba98-7654-3210-fedc-ba9876543210",
|
||||
],
|
||||
)
|
||||
model_key_2: str = Field(
|
||||
...,
|
||||
description="The key of the second model in the relationship",
|
||||
examples=[
|
||||
"3bb7c0eb-b6c8-469c-ad8c-4d69c06075e4",
|
||||
"f0c3da4e-d9ff-42b5-a45c-23be75c887c9",
|
||||
"38170dd8-f1e5-431e-866c-2c81f1277fcc",
|
||||
"c57fea2d-7646-424c-b9ad-c0ba60fc68be",
|
||||
"10f7807b-ab54-46a9-ab03-600e88c630a1",
|
||||
"f6c1d267-cf87-4ee0-bee0-37e791eacab7",
|
||||
],
|
||||
)
|
||||
|
||||
|
||||
class ModelRelationshipBatchRequest(BaseModel):
|
||||
model_keys: List[str] = Field(
|
||||
...,
|
||||
description="List of model keys to fetch related models for",
|
||||
examples=[
|
||||
[
|
||||
"aa3b247f-90c9-4416-bfcd-aeaa57a5339e",
|
||||
"ac32b914-10ab-496e-a24a-3068724b9c35",
|
||||
],
|
||||
[
|
||||
"b1c2d3e4-f5a6-7890-abcd-ef1234567890",
|
||||
"12345678-90ab-cdef-1234-567890abcdef",
|
||||
"fedcba98-7654-3210-fedc-ba9876543210",
|
||||
],
|
||||
[
|
||||
"3bb7c0eb-b6c8-469c-ad8c-4d69c06075e4",
|
||||
],
|
||||
],
|
||||
)
|
||||
|
||||
|
||||
# === Routes ===
|
||||
|
||||
|
||||
@model_relationships_router.get(
|
||||
"/i/{model_key}",
|
||||
operation_id="get_related_models",
|
||||
response_model=list[str],
|
||||
responses={
|
||||
200: {
|
||||
"description": "A list of related model keys was retrieved successfully",
|
||||
"content": {
|
||||
"application/json": {
|
||||
"example": [
|
||||
"15e9eb28-8cfe-47c9-b610-37907a79fc3c",
|
||||
"71272e82-0e5f-46d5-bca9-9a61f4bd8a82",
|
||||
"a5d7cd49-1b98-4534-a475-aeee4ccf5fa2",
|
||||
]
|
||||
}
|
||||
},
|
||||
},
|
||||
404: {"description": "The specified model could not be found"},
|
||||
422: {"description": "Validation error"},
|
||||
},
|
||||
)
|
||||
async def get_related_models(
|
||||
current_user: CurrentUserOrDefault,
|
||||
model_key: str = Path(..., description="The key of the model to get relationships for"),
|
||||
) -> list[str]:
|
||||
"""
|
||||
Get a list of model keys related to a given model.
|
||||
"""
|
||||
return ApiDependencies.invoker.services.model_relationships.get_related_model_keys(model_key)
|
||||
|
||||
|
||||
@model_relationships_router.post(
|
||||
"/",
|
||||
status_code=status.HTTP_204_NO_CONTENT,
|
||||
responses={
|
||||
204: {"description": "The relationship was successfully created"},
|
||||
400: {"description": "Invalid model keys or self-referential relationship"},
|
||||
409: {"description": "The relationship already exists"},
|
||||
422: {"description": "Validation error"},
|
||||
500: {"description": "Internal server error"},
|
||||
},
|
||||
summary="Add Model Relationship",
|
||||
description="Creates a **bidirectional** relationship between two models, allowing each to reference the other as related.",
|
||||
)
|
||||
async def add_model_relationship(
|
||||
current_user: AdminUserOrDefault,
|
||||
req: ModelRelationshipCreateRequest = Body(..., description="The model keys to relate"),
|
||||
) -> None:
|
||||
"""
|
||||
Add a relationship between two models.
|
||||
|
||||
Relationships are bidirectional and will be accessible from both models.
|
||||
|
||||
- Raises 400 if keys are invalid or identical.
|
||||
- Raises 409 if the relationship already exists.
|
||||
"""
|
||||
if req.model_key_1 == req.model_key_2:
|
||||
raise HTTPException(status_code=400, detail="Cannot relate a model to itself.")
|
||||
|
||||
try:
|
||||
ApiDependencies.invoker.services.model_relationships.add_model_relationship(
|
||||
req.model_key_1,
|
||||
req.model_key_2,
|
||||
)
|
||||
except ValueError as e:
|
||||
raise HTTPException(status_code=409, detail=str(e))
|
||||
|
||||
|
||||
@model_relationships_router.delete(
|
||||
"/",
|
||||
status_code=status.HTTP_204_NO_CONTENT,
|
||||
responses={
|
||||
204: {"description": "The relationship was successfully removed"},
|
||||
400: {"description": "Invalid model keys or self-referential relationship"},
|
||||
404: {"description": "The relationship does not exist"},
|
||||
422: {"description": "Validation error"},
|
||||
500: {"description": "Internal server error"},
|
||||
},
|
||||
summary="Remove Model Relationship",
|
||||
description="Removes a **bidirectional** relationship between two models. The relationship must already exist.",
|
||||
)
|
||||
async def remove_model_relationship(
|
||||
current_user: AdminUserOrDefault,
|
||||
req: ModelRelationshipCreateRequest = Body(..., description="The model keys to disconnect"),
|
||||
) -> None:
|
||||
"""
|
||||
Removes a bidirectional relationship between two model keys.
|
||||
|
||||
- Raises 400 if attempting to unlink a model from itself.
|
||||
- Raises 404 if the relationship was not found.
|
||||
"""
|
||||
if req.model_key_1 == req.model_key_2:
|
||||
raise HTTPException(status_code=400, detail="Cannot unlink a model from itself.")
|
||||
|
||||
try:
|
||||
ApiDependencies.invoker.services.model_relationships.remove_model_relationship(
|
||||
req.model_key_1,
|
||||
req.model_key_2,
|
||||
)
|
||||
except ValueError as e:
|
||||
raise HTTPException(status_code=404, detail=str(e))
|
||||
|
||||
|
||||
@model_relationships_router.post(
|
||||
"/batch",
|
||||
operation_id="get_related_models_batch",
|
||||
response_model=List[str],
|
||||
responses={
|
||||
200: {
|
||||
"description": "Related model keys retrieved successfully",
|
||||
"content": {
|
||||
"application/json": {
|
||||
"example": [
|
||||
"ca562b14-995e-4a42-90c1-9528f1a5921d",
|
||||
"cc0c2b8a-c62e-41d6-878e-cc74dde5ca8f",
|
||||
"18ca7649-6a9e-47d5-bc17-41ab1e8cec81",
|
||||
"7c12d1b2-0ef9-4bec-ba55-797b2d8f2ee1",
|
||||
"c382eaa3-0e28-4ab0-9446-408667699aeb",
|
||||
"71272e82-0e5f-46d5-bca9-9a61f4bd8a82",
|
||||
"a5d7cd49-1b98-4534-a475-aeee4ccf5fa2",
|
||||
]
|
||||
}
|
||||
},
|
||||
},
|
||||
422: {"description": "Validation error"},
|
||||
500: {"description": "Internal server error"},
|
||||
},
|
||||
summary="Get Related Model Keys (Batch)",
|
||||
description="Retrieves all **unique related model keys** for a list of given models. This is useful for contextual suggestions or filtering.",
|
||||
)
|
||||
async def get_related_models_batch(
|
||||
current_user: CurrentUserOrDefault,
|
||||
req: ModelRelationshipBatchRequest = Body(..., description="Model keys to check for related connections"),
|
||||
) -> list[str]:
|
||||
"""
|
||||
Accepts multiple model keys and returns a flat list of all unique related keys.
|
||||
|
||||
Useful when working with multiple selections in the UI or cross-model comparisons.
|
||||
"""
|
||||
all_related: set[str] = set()
|
||||
for key in req.model_keys:
|
||||
related = ApiDependencies.invoker.services.model_relationships.get_related_model_keys(key)
|
||||
all_related.update(related)
|
||||
return list(all_related)
|
||||
@@ -0,0 +1,620 @@
|
||||
"""Router for updating recallable parameters on the frontend."""
|
||||
|
||||
import json
|
||||
from typing import Any, Literal, Optional
|
||||
|
||||
from fastapi import Body, HTTPException, Path, Query
|
||||
from fastapi.routing import APIRouter
|
||||
from pydantic import BaseModel, ConfigDict, Field
|
||||
|
||||
from invokeai.app.api.auth_dependencies import CurrentUserOrDefault
|
||||
from invokeai.app.api.dependencies import ApiDependencies
|
||||
from invokeai.app.api.routers.image_move_maintenance import assert_image_move_maintenance_inactive
|
||||
from invokeai.backend.image_util.controlnet_processor import process_controlnet_image
|
||||
from invokeai.backend.model_manager.taxonomy import ModelType
|
||||
|
||||
recall_parameters_router = APIRouter(prefix="/v1/recall", tags=["recall"])
|
||||
|
||||
|
||||
class LoRARecallParameter(BaseModel):
|
||||
"""LoRA configuration for recall"""
|
||||
|
||||
model_name: str = Field(description="The name of the LoRA model")
|
||||
weight: float = Field(default=0.75, ge=-10, le=10, description="The weight for the LoRA")
|
||||
is_enabled: bool = Field(default=True, description="Whether the LoRA is enabled")
|
||||
|
||||
|
||||
class ControlNetRecallParameter(BaseModel):
|
||||
"""ControlNet configuration for recall"""
|
||||
|
||||
model_name: str = Field(description="The name of the ControlNet/T2I Adapter/Control LoRA model")
|
||||
image_name: Optional[str] = Field(default=None, description="The filename of the control image in outputs/images")
|
||||
weight: float = Field(default=1.0, ge=-1, le=2, description="The weight for the control adapter")
|
||||
begin_step_percent: Optional[float] = Field(
|
||||
default=None, ge=0, le=1, description="When the control adapter is first applied (% of total steps)"
|
||||
)
|
||||
end_step_percent: Optional[float] = Field(
|
||||
default=None, ge=0, le=1, description="When the control adapter is last applied (% of total steps)"
|
||||
)
|
||||
control_mode: Optional[Literal["balanced", "more_prompt", "more_control"]] = Field(
|
||||
default=None, description="The control mode (ControlNet only)"
|
||||
)
|
||||
|
||||
|
||||
class IPAdapterRecallParameter(BaseModel):
|
||||
"""IP Adapter configuration for recall"""
|
||||
|
||||
model_name: str = Field(description="The name of the IP Adapter model")
|
||||
image_name: Optional[str] = Field(default=None, description="The filename of the reference image in outputs/images")
|
||||
weight: float = Field(default=1.0, ge=-1, le=2, description="The weight for the IP Adapter")
|
||||
begin_step_percent: Optional[float] = Field(
|
||||
default=None, ge=0, le=1, description="When the IP Adapter is first applied (% of total steps)"
|
||||
)
|
||||
end_step_percent: Optional[float] = Field(
|
||||
default=None, ge=0, le=1, description="When the IP Adapter is last applied (% of total steps)"
|
||||
)
|
||||
method: Optional[Literal["full", "style", "composition"]] = Field(default=None, description="The IP Adapter method")
|
||||
image_influence: Optional[Literal["lowest", "low", "medium", "high", "highest"]] = Field(
|
||||
default=None, description="FLUX Redux image influence (if model is flux_redux)"
|
||||
)
|
||||
|
||||
|
||||
class ReferenceImageRecallParameter(BaseModel):
|
||||
"""Global reference-image configuration for recall.
|
||||
|
||||
Used for reference images that feed directly into the main model rather
|
||||
than through a separate IP-Adapter / ControlNet model — for example
|
||||
FLUX.2 Klein, FLUX Kontext, and Qwen Image Edit. The receiving frontend
|
||||
picks the correct config type (``flux2_reference_image`` /
|
||||
``qwen_image_reference_image`` / ``flux_kontext_reference_image``) based
|
||||
on the currently-selected main model.
|
||||
"""
|
||||
|
||||
image_name: str = Field(description="The filename of the reference image in outputs/images")
|
||||
|
||||
|
||||
class RecallParameter(BaseModel):
|
||||
"""Request model for updating recallable parameters."""
|
||||
|
||||
model_config = ConfigDict(extra="forbid")
|
||||
|
||||
# Prompts
|
||||
positive_prompt: Optional[str] = Field(None, description="Positive prompt text")
|
||||
negative_prompt: Optional[str] = Field(None, description="Negative prompt text")
|
||||
|
||||
# Model configuration
|
||||
model: Optional[str] = Field(None, description="Main model name/identifier")
|
||||
refiner_model: Optional[str] = Field(None, description="Refiner model name/identifier")
|
||||
vae_model: Optional[str] = Field(None, description="VAE model name/identifier")
|
||||
scheduler: Optional[str] = Field(None, description="Scheduler name")
|
||||
|
||||
# Generation parameters
|
||||
steps: Optional[int] = Field(None, ge=1, description="Number of generation steps")
|
||||
refiner_steps: Optional[int] = Field(None, ge=0, description="Number of refiner steps")
|
||||
cfg_scale: Optional[float] = Field(None, description="CFG scale for guidance")
|
||||
cfg_rescale_multiplier: Optional[float] = Field(None, description="CFG rescale multiplier")
|
||||
refiner_cfg_scale: Optional[float] = Field(None, description="Refiner CFG scale")
|
||||
guidance: Optional[float] = Field(None, description="Guidance scale")
|
||||
|
||||
# Image parameters
|
||||
width: Optional[int] = Field(None, ge=64, description="Image width in pixels")
|
||||
height: Optional[int] = Field(None, ge=64, description="Image height in pixels")
|
||||
seed: Optional[int] = Field(None, ge=0, description="Random seed")
|
||||
|
||||
# Advanced parameters
|
||||
denoise_strength: Optional[float] = Field(None, ge=0, le=1, description="Denoising strength")
|
||||
refiner_denoise_start: Optional[float] = Field(None, ge=0, le=1, description="Refiner denoising start")
|
||||
clip_skip: Optional[int] = Field(None, ge=0, description="CLIP skip layers")
|
||||
seamless_x: Optional[bool] = Field(None, description="Enable seamless X tiling")
|
||||
seamless_y: Optional[bool] = Field(None, description="Enable seamless Y tiling")
|
||||
|
||||
# Refiner aesthetics
|
||||
refiner_positive_aesthetic_score: Optional[float] = Field(None, description="Refiner positive aesthetic score")
|
||||
refiner_negative_aesthetic_score: Optional[float] = Field(None, description="Refiner negative aesthetic score")
|
||||
|
||||
# LoRAs, ControlNets, and IP Adapters
|
||||
loras: Optional[list[LoRARecallParameter]] = Field(None, description="List of LoRAs with their weights")
|
||||
control_layers: Optional[list[ControlNetRecallParameter]] = Field(
|
||||
None, description="List of control adapters (ControlNet, T2I Adapter, Control LoRA) with their settings"
|
||||
)
|
||||
ip_adapters: Optional[list[IPAdapterRecallParameter]] = Field(
|
||||
None, description="List of IP Adapters with their settings"
|
||||
)
|
||||
reference_images: Optional[list[ReferenceImageRecallParameter]] = Field(
|
||||
None,
|
||||
description=(
|
||||
"List of model-free reference images for architectures that consume reference "
|
||||
"images directly (FLUX.2 Klein, FLUX Kontext, Qwen Image Edit). The frontend "
|
||||
"picks the correct config type based on the currently-selected main model."
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def resolve_model_name_to_key(model_name: str, model_type: ModelType = ModelType.Main) -> Optional[str]:
|
||||
"""
|
||||
Look up a model by name and return its key.
|
||||
|
||||
Args:
|
||||
model_name: The name of the model to look up
|
||||
model_type: The type of model to search for (default: Main)
|
||||
|
||||
Returns:
|
||||
The key of the first matching model, or None if not found.
|
||||
"""
|
||||
logger = ApiDependencies.invoker.services.logger
|
||||
try:
|
||||
models = ApiDependencies.invoker.services.model_manager.store.search_by_attr(
|
||||
model_name=model_name, model_type=model_type
|
||||
)
|
||||
|
||||
if models:
|
||||
logger.info(f"Resolved {model_type.value} model name '{model_name}' to key '{models[0].key}'")
|
||||
return models[0].key
|
||||
|
||||
logger.warning(f"Could not find {model_type.value} model with name '{model_name}'")
|
||||
return None
|
||||
except Exception as e:
|
||||
logger.error(f"Exception during {model_type.value} model lookup: {e}", exc_info=True)
|
||||
return None
|
||||
|
||||
|
||||
def load_image_file(image_name: str) -> Optional[dict[str, Any]]:
|
||||
"""
|
||||
Load an image from the outputs/images directory.
|
||||
|
||||
Args:
|
||||
image_name: The filename of the image in outputs/images
|
||||
|
||||
Returns:
|
||||
A dictionary with image_name, width, and height, or None if the image cannot be found
|
||||
"""
|
||||
logger = ApiDependencies.invoker.services.logger
|
||||
try:
|
||||
images_service = ApiDependencies.invoker.services.images
|
||||
# Use images service which handles subfolder resolution via DB record
|
||||
path = images_service.get_path(image_name)
|
||||
|
||||
if not images_service.validate_path(path):
|
||||
logger.warning(f"Image file not found: {image_name}")
|
||||
return None
|
||||
|
||||
pil_image = images_service.get_pil_image(image_name)
|
||||
width, height = pil_image.size
|
||||
logger.info(f"Found image file: {image_name} ({width}x{height})")
|
||||
return {"image_name": image_name, "width": width, "height": height}
|
||||
except Exception as e:
|
||||
logger.warning(f"Error loading image file {image_name}: {e}")
|
||||
return None
|
||||
|
||||
|
||||
def resolve_lora_models(loras: list[LoRARecallParameter]) -> list[dict[str, Any]]:
|
||||
"""
|
||||
Resolve LoRA model names to keys and build configuration list.
|
||||
|
||||
Args:
|
||||
loras: List of LoRA recall parameters
|
||||
|
||||
Returns:
|
||||
List of resolved LoRA configurations with model keys
|
||||
"""
|
||||
logger = ApiDependencies.invoker.services.logger
|
||||
resolved_loras = []
|
||||
|
||||
for lora in loras:
|
||||
model_key = resolve_model_name_to_key(lora.model_name, ModelType.LoRA)
|
||||
if model_key:
|
||||
resolved_loras.append({"model_key": model_key, "weight": lora.weight, "is_enabled": lora.is_enabled})
|
||||
else:
|
||||
logger.warning(f"Skipping LoRA '{lora.model_name}' - model not found")
|
||||
|
||||
return resolved_loras
|
||||
|
||||
|
||||
def resolve_control_models(control_layers: list[ControlNetRecallParameter]) -> list[dict[str, Any]]:
|
||||
"""
|
||||
Resolve control adapter model names to keys and build configuration list.
|
||||
|
||||
Tries to resolve as ControlNet, T2I Adapter, or Control LoRA in that order.
|
||||
|
||||
Args:
|
||||
control_layers: List of control adapter recall parameters
|
||||
|
||||
Returns:
|
||||
List of resolved control adapter configurations with model keys
|
||||
"""
|
||||
logger = ApiDependencies.invoker.services.logger
|
||||
services = ApiDependencies.invoker.services
|
||||
resolved_controls = []
|
||||
|
||||
for control in control_layers:
|
||||
model_key = None
|
||||
|
||||
# Try ControlNet first
|
||||
model_key = resolve_model_name_to_key(control.model_name, ModelType.ControlNet)
|
||||
if not model_key:
|
||||
# Try T2I Adapter
|
||||
model_key = resolve_model_name_to_key(control.model_name, ModelType.T2IAdapter)
|
||||
if not model_key:
|
||||
# Try Control LoRA (also uses LoRA type)
|
||||
model_key = resolve_model_name_to_key(control.model_name, ModelType.LoRA)
|
||||
|
||||
if model_key:
|
||||
config: dict[str, Any] = {"model_key": model_key, "weight": control.weight}
|
||||
if control.image_name is not None:
|
||||
image_data = load_image_file(control.image_name)
|
||||
if image_data:
|
||||
config["image"] = image_data
|
||||
|
||||
# Try to process the image using the model's default processor
|
||||
processed_image_data = process_controlnet_image(control.image_name, model_key, services)
|
||||
if processed_image_data:
|
||||
config["processed_image"] = processed_image_data
|
||||
logger.info(f"Added processed image for control adapter {control.model_name}")
|
||||
else:
|
||||
logger.warning(f"Could not load image for control adapter: {control.image_name}")
|
||||
if control.begin_step_percent is not None:
|
||||
config["begin_step_percent"] = control.begin_step_percent
|
||||
if control.end_step_percent is not None:
|
||||
config["end_step_percent"] = control.end_step_percent
|
||||
if control.control_mode is not None:
|
||||
config["control_mode"] = control.control_mode
|
||||
|
||||
resolved_controls.append(config)
|
||||
else:
|
||||
logger.warning(f"Skipping control adapter '{control.model_name}' - model not found")
|
||||
|
||||
return resolved_controls
|
||||
|
||||
|
||||
def resolve_ip_adapter_models(ip_adapters: list[IPAdapterRecallParameter]) -> list[dict[str, Any]]:
|
||||
"""
|
||||
Resolve IP Adapter model names to keys and build configuration list.
|
||||
|
||||
Args:
|
||||
ip_adapters: List of IP Adapter recall parameters
|
||||
|
||||
Returns:
|
||||
List of resolved IP Adapter configurations with model keys
|
||||
"""
|
||||
logger = ApiDependencies.invoker.services.logger
|
||||
resolved_adapters = []
|
||||
|
||||
for adapter in ip_adapters:
|
||||
# Try resolving as IP Adapter; if not found, try FLUX Redux
|
||||
model_key = resolve_model_name_to_key(adapter.model_name, ModelType.IPAdapter)
|
||||
if not model_key:
|
||||
model_key = resolve_model_name_to_key(adapter.model_name, ModelType.FluxRedux)
|
||||
if model_key:
|
||||
config: dict[str, Any] = {
|
||||
"model_key": model_key,
|
||||
# Always include weight; ignored by FLUX Redux on the frontend
|
||||
"weight": adapter.weight,
|
||||
}
|
||||
if adapter.image_name is not None:
|
||||
image_data = load_image_file(adapter.image_name)
|
||||
if image_data:
|
||||
config["image"] = image_data
|
||||
else:
|
||||
logger.warning(f"Could not load image for IP Adapter: {adapter.image_name}")
|
||||
if adapter.begin_step_percent is not None:
|
||||
config["begin_step_percent"] = adapter.begin_step_percent
|
||||
if adapter.end_step_percent is not None:
|
||||
config["end_step_percent"] = adapter.end_step_percent
|
||||
if adapter.method is not None:
|
||||
config["method"] = adapter.method
|
||||
# Include FLUX Redux image influence when provided
|
||||
if adapter.image_influence is not None:
|
||||
config["image_influence"] = adapter.image_influence
|
||||
|
||||
resolved_adapters.append(config)
|
||||
else:
|
||||
logger.warning(f"Skipping IP Adapter '{adapter.model_name}' - model not found")
|
||||
|
||||
return resolved_adapters
|
||||
|
||||
|
||||
def resolve_reference_images(
|
||||
reference_images: list[ReferenceImageRecallParameter],
|
||||
) -> list[dict[str, Any]]:
|
||||
"""
|
||||
Validate model-free reference images and build the configuration list.
|
||||
|
||||
Unlike IP Adapters and ControlNets, these reference images are consumed
|
||||
directly by the main model (FLUX.2 Klein, FLUX Kontext, Qwen Image Edit),
|
||||
so there is no adapter-model name to resolve. We simply verify that each
|
||||
referenced file exists in ``outputs/images`` and pass the image metadata
|
||||
through to the frontend.
|
||||
|
||||
Args:
|
||||
reference_images: List of reference-image recall parameters
|
||||
|
||||
Returns:
|
||||
List of reference-image configurations with resolved image metadata.
|
||||
Entries whose image file cannot be loaded are dropped with a warning.
|
||||
"""
|
||||
logger = ApiDependencies.invoker.services.logger
|
||||
resolved: list[dict[str, Any]] = []
|
||||
|
||||
for ref in reference_images:
|
||||
image_data = load_image_file(ref.image_name)
|
||||
if image_data is None:
|
||||
logger.warning(f"Skipping reference image '{ref.image_name}' - file not found")
|
||||
continue
|
||||
resolved.append({"image": image_data})
|
||||
|
||||
return resolved
|
||||
|
||||
|
||||
def _assert_recall_image_access(parameters: "RecallParameter", current_user: CurrentUserOrDefault) -> None:
|
||||
"""Validate that the caller can read every image referenced in the recall parameters.
|
||||
|
||||
Control layers, IP adapters, and reference images may reference image_name fields.
|
||||
Without this check an attacker who knows another user's image UUID could use the recall
|
||||
endpoint to extract image dimensions and — for ControlNet preprocessors — mint
|
||||
a derived processed image they can then fetch.
|
||||
"""
|
||||
from invokeai.app.services.board_records.board_records_common import BoardVisibility
|
||||
|
||||
image_names: list[str] = []
|
||||
if parameters.control_layers:
|
||||
for layer in parameters.control_layers:
|
||||
if layer.image_name is not None:
|
||||
image_names.append(layer.image_name)
|
||||
if parameters.ip_adapters:
|
||||
for adapter in parameters.ip_adapters:
|
||||
if adapter.image_name is not None:
|
||||
image_names.append(adapter.image_name)
|
||||
if parameters.reference_images:
|
||||
for ref in parameters.reference_images:
|
||||
if ref.image_name is not None:
|
||||
image_names.append(ref.image_name)
|
||||
|
||||
if not image_names:
|
||||
return
|
||||
|
||||
# Admin can access all images
|
||||
if current_user.is_admin:
|
||||
return
|
||||
|
||||
for image_name in image_names:
|
||||
owner = ApiDependencies.invoker.services.image_records.get_user_id(image_name)
|
||||
if owner is not None and owner == current_user.user_id:
|
||||
continue
|
||||
|
||||
# Check board visibility
|
||||
board_id = ApiDependencies.invoker.services.board_image_records.get_board_for_image(image_name)
|
||||
if board_id is not None:
|
||||
try:
|
||||
board = ApiDependencies.invoker.services.boards.get_dto(board_id=board_id)
|
||||
if board.board_visibility in (BoardVisibility.Shared, BoardVisibility.Public):
|
||||
continue
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
raise HTTPException(status_code=403, detail=f"Not authorized to access image {image_name}")
|
||||
|
||||
|
||||
@recall_parameters_router.post(
|
||||
"/{queue_id}",
|
||||
operation_id="update_recall_parameters",
|
||||
response_model=dict[str, Any],
|
||||
)
|
||||
async def update_recall_parameters(
|
||||
current_user: CurrentUserOrDefault,
|
||||
queue_id: str = Path(..., description="The queue id to perform this operation on"),
|
||||
parameters: RecallParameter = Body(..., description="Recall parameters to update"),
|
||||
strict: bool = Query(
|
||||
default=False,
|
||||
description="When true, parameters not included in the request are reset to their defaults (cleared).",
|
||||
),
|
||||
append: bool = Query(
|
||||
default=False,
|
||||
description=(
|
||||
"When true, recalled reference images (ip_adapters and reference_images) are "
|
||||
"appended to the frontend's existing reference-image list instead of replacing it. "
|
||||
"Mutually exclusive with strict."
|
||||
),
|
||||
),
|
||||
) -> dict[str, Any]:
|
||||
"""
|
||||
Update recallable parameters that can be recalled on the frontend.
|
||||
|
||||
This endpoint allows updating parameters such as prompt, model, steps, and other
|
||||
generation settings. These parameters are stored in client state and can be
|
||||
accessed by the frontend to populate UI elements.
|
||||
|
||||
Args:
|
||||
queue_id: The queue ID to associate these parameters with
|
||||
parameters: The RecallParameter object containing the parameters to update
|
||||
strict: When true, parameters not included in the request body are reset
|
||||
to their defaults (cleared on the frontend). Defaults to false,
|
||||
which preserves the existing behaviour of only updating the
|
||||
parameters that are explicitly provided.
|
||||
append: When true, recalled reference images (``ip_adapters`` and
|
||||
``reference_images``) are appended to whatever reference images the
|
||||
frontend already has, instead of replacing the whole list. Mutually
|
||||
exclusive with ``strict`` (which clears omitted parameters).
|
||||
|
||||
Returns:
|
||||
A dictionary containing the updated parameters and status
|
||||
|
||||
Example:
|
||||
POST /api/v1/recall/{queue_id}?strict=true
|
||||
{
|
||||
"positive_prompt": "a beautiful landscape",
|
||||
"model": "sd-1.5",
|
||||
"steps": 20
|
||||
}
|
||||
# In strict mode, all other parameters (reference_images, loras, etc.)
|
||||
# are cleared. In non-strict mode (default) they would be left as-is.
|
||||
"""
|
||||
logger = ApiDependencies.invoker.services.logger
|
||||
|
||||
if strict and append:
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail="The 'strict' and 'append' query parameters are mutually exclusive",
|
||||
)
|
||||
|
||||
# Validate image access before processing — prevents information leakage
|
||||
# (dimensions) and derived-image minting via ControlNet preprocessors.
|
||||
_assert_recall_image_access(parameters, current_user)
|
||||
assert_image_move_maintenance_inactive()
|
||||
|
||||
try:
|
||||
# In strict mode, include all parameters so the frontend clears anything
|
||||
# not explicitly provided. List-typed fields use [] instead of None so
|
||||
# the frontend sees an empty collection rather than a null it might skip.
|
||||
if strict:
|
||||
_list_fields = {
|
||||
name for name, field in RecallParameter.model_fields.items() if "list" in str(field.annotation).lower()
|
||||
}
|
||||
provided_params = {
|
||||
k: ([] if v is None and k in _list_fields else v) for k, v in parameters.model_dump().items()
|
||||
}
|
||||
else:
|
||||
provided_params = {k: v for k, v in parameters.model_dump().items() if v is not None}
|
||||
|
||||
if not provided_params:
|
||||
return {"status": "no_parameters_provided", "updated_count": 0}
|
||||
|
||||
# Store each parameter in client state scoped to the current user
|
||||
updated_count = 0
|
||||
for param_key, param_value in provided_params.items():
|
||||
# Convert parameter values to JSON strings for storage
|
||||
value_str = json.dumps(param_value)
|
||||
try:
|
||||
ApiDependencies.invoker.services.client_state_persistence.set_by_key(
|
||||
current_user.user_id, f"recall_{param_key}", value_str
|
||||
)
|
||||
updated_count += 1
|
||||
except Exception as e:
|
||||
logger.error(f"Error setting recall parameter {param_key}: {e}")
|
||||
raise HTTPException(
|
||||
status_code=500,
|
||||
detail=f"Error setting recall parameter {param_key}",
|
||||
)
|
||||
|
||||
logger.info(f"Updated {updated_count} recall parameters for queue {queue_id}")
|
||||
|
||||
# Resolve model name to key if a model was provided
|
||||
if "model" in provided_params and isinstance(provided_params["model"], str):
|
||||
model_name = provided_params["model"]
|
||||
model_key = resolve_model_name_to_key(model_name, ModelType.Main)
|
||||
|
||||
if model_key:
|
||||
logger.info(f"Resolved model name '{model_name}' to key '{model_key}'")
|
||||
provided_params["model"] = model_key
|
||||
else:
|
||||
logger.warning(f"Could not resolve model name '{model_name}' to a model key")
|
||||
# Remove model from parameters if we couldn't resolve it
|
||||
del provided_params["model"]
|
||||
|
||||
# Process LoRAs if provided
|
||||
if "loras" in provided_params:
|
||||
loras_param = parameters.loras
|
||||
if loras_param is not None:
|
||||
resolved_loras = resolve_lora_models(loras_param)
|
||||
provided_params["loras"] = resolved_loras
|
||||
logger.info(f"Resolved {len(resolved_loras)} LoRA(s)")
|
||||
|
||||
# Process control layers if provided
|
||||
if "control_layers" in provided_params:
|
||||
control_layers_param = parameters.control_layers
|
||||
if control_layers_param is not None:
|
||||
resolved_controls = resolve_control_models(control_layers_param)
|
||||
provided_params["control_layers"] = resolved_controls
|
||||
logger.info(f"Resolved {len(resolved_controls)} control layer(s)")
|
||||
|
||||
# Process IP adapters if provided
|
||||
if "ip_adapters" in provided_params:
|
||||
ip_adapters_param = parameters.ip_adapters
|
||||
if ip_adapters_param is not None:
|
||||
resolved_adapters = resolve_ip_adapter_models(ip_adapters_param)
|
||||
provided_params["ip_adapters"] = resolved_adapters
|
||||
logger.info(f"Resolved {len(resolved_adapters)} IP adapter(s)")
|
||||
|
||||
# Process model-free reference images if provided
|
||||
if "reference_images" in provided_params:
|
||||
reference_images_param = parameters.reference_images
|
||||
if reference_images_param is not None:
|
||||
resolved_refs = resolve_reference_images(reference_images_param)
|
||||
provided_params["reference_images"] = resolved_refs
|
||||
logger.info(f"Resolved {len(resolved_refs)} reference image(s)")
|
||||
|
||||
# Append mode rides along inside the event's parameters dict rather
|
||||
# than as a new event field so the generated client schema (which
|
||||
# types parameters as a free-form object) doesn't need regenerating.
|
||||
# Added after the persistence loop above, so the flag itself is never
|
||||
# stored as a recall parameter.
|
||||
if append:
|
||||
provided_params["append"] = True
|
||||
|
||||
# Emit event to notify frontend of parameter updates
|
||||
try:
|
||||
logger.info(
|
||||
f"Emitting recall_parameters_updated event for queue {queue_id} with {len(provided_params)} parameters"
|
||||
)
|
||||
ApiDependencies.invoker.services.events.emit_recall_parameters_updated(
|
||||
queue_id, current_user.user_id, provided_params
|
||||
)
|
||||
logger.info("Successfully emitted recall_parameters_updated event")
|
||||
except Exception as e:
|
||||
logger.error(f"Error emitting recall parameters event: {e}", exc_info=True)
|
||||
# Don't fail the request if event emission fails, just log it
|
||||
|
||||
return {
|
||||
"status": "success",
|
||||
"queue_id": queue_id,
|
||||
"updated_count": updated_count,
|
||||
"parameters": provided_params,
|
||||
}
|
||||
|
||||
except HTTPException:
|
||||
raise
|
||||
except Exception as e:
|
||||
logger.error(f"Error updating recall parameters: {e}")
|
||||
raise HTTPException(
|
||||
status_code=500,
|
||||
detail="Error updating recall parameters",
|
||||
)
|
||||
|
||||
|
||||
@recall_parameters_router.get(
|
||||
"/{queue_id}",
|
||||
operation_id="get_recall_parameters",
|
||||
response_model=dict[str, Any],
|
||||
)
|
||||
async def get_recall_parameters(
|
||||
current_user: CurrentUserOrDefault,
|
||||
queue_id: str = Path(..., description="The queue id to retrieve parameters for"),
|
||||
) -> dict[str, Any]:
|
||||
"""
|
||||
Retrieve all stored recall parameters for a given queue.
|
||||
|
||||
Returns a dictionary of all recall parameters that have been set for the queue.
|
||||
|
||||
Args:
|
||||
queue_id: The queue ID to retrieve parameters for
|
||||
|
||||
Returns:
|
||||
A dictionary containing all stored recall parameters
|
||||
"""
|
||||
logger = ApiDependencies.invoker.services.logger
|
||||
|
||||
try:
|
||||
# Retrieve all recall parameters by iterating through expected keys
|
||||
# Since client_state_persistence doesn't have a "get_all" method, we'll
|
||||
# return an informative response
|
||||
return {
|
||||
"status": "success",
|
||||
"queue_id": queue_id,
|
||||
"note": "Use the frontend to access stored recall parameters, or set specific parameters using POST",
|
||||
}
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error retrieving recall parameters: {e}")
|
||||
raise HTTPException(
|
||||
status_code=500,
|
||||
detail="Error retrieving recall parameters",
|
||||
)
|
||||
@@ -0,0 +1,634 @@
|
||||
from typing import Optional
|
||||
|
||||
from fastapi import Body, HTTPException, Path, Query
|
||||
from fastapi.routing import APIRouter
|
||||
from pydantic import BaseModel
|
||||
|
||||
from invokeai.app.api.auth_dependencies import AdminUserOrDefault, CurrentUserOrDefault
|
||||
from invokeai.app.api.dependencies import ApiDependencies
|
||||
from invokeai.app.api.routers.image_move_maintenance import assert_image_move_maintenance_inactive
|
||||
from invokeai.app.services.session_processor.session_processor_common import SessionProcessorStatus
|
||||
from invokeai.app.services.session_queue.session_queue_common import (
|
||||
Batch,
|
||||
BatchStatus,
|
||||
CancelAllExceptCurrentResult,
|
||||
CancelByBatchIDsResult,
|
||||
CancelByDestinationResult,
|
||||
ClearResult,
|
||||
DeleteAllExceptCurrentResult,
|
||||
DeleteByDestinationResult,
|
||||
EnqueueBatchResult,
|
||||
ItemIdsResult,
|
||||
PruneResult,
|
||||
RetryItemsResult,
|
||||
SessionQueueCountsByDestination,
|
||||
SessionQueueItem,
|
||||
SessionQueueItemNotFoundError,
|
||||
SessionQueueStatus,
|
||||
)
|
||||
from invokeai.app.services.shared.graph import Graph, GraphExecutionState
|
||||
from invokeai.app.services.shared.sqlite.sqlite_common import SQLiteDirection
|
||||
|
||||
session_queue_router = APIRouter(prefix="/v1/queue", tags=["queue"])
|
||||
|
||||
|
||||
class SessionQueueAndProcessorStatus(BaseModel):
|
||||
"""The overall status of session queue and processor"""
|
||||
|
||||
queue: SessionQueueStatus
|
||||
processor: SessionProcessorStatus
|
||||
|
||||
|
||||
def _get_workflow_call_root_queue_item(queue_item: SessionQueueItem) -> SessionQueueItem:
|
||||
if queue_item.root_item_id is None:
|
||||
return queue_item
|
||||
return ApiDependencies.invoker.services.session_queue.get_queue_item(queue_item.root_item_id)
|
||||
|
||||
|
||||
def sanitize_queue_item_for_user(
|
||||
queue_item: SessionQueueItem, current_user_id: str, is_admin: bool
|
||||
) -> SessionQueueItem:
|
||||
"""Sanitize queue item for non-admin users viewing other users' items.
|
||||
|
||||
For non-admin users viewing queue items belonging to other users,
|
||||
only timestamps, status, and error information are exposed. All other
|
||||
fields (user identity, generation parameters, graphs, workflows) are stripped.
|
||||
|
||||
Args:
|
||||
queue_item: The queue item to sanitize
|
||||
current_user_id: The ID of the current user viewing the item
|
||||
is_admin: Whether the current user is an admin
|
||||
|
||||
Returns:
|
||||
The sanitized queue item (sensitive fields cleared if necessary)
|
||||
"""
|
||||
# Admins and item owners can see everything
|
||||
if is_admin or queue_item.user_id == current_user_id:
|
||||
return queue_item
|
||||
|
||||
# For non-admins viewing other users' items, strip everything except
|
||||
# item_id, queue_id, status, and timestamps
|
||||
sanitized_item = queue_item.model_copy(deep=False)
|
||||
sanitized_item.user_id = "redacted"
|
||||
sanitized_item.user_display_name = None
|
||||
sanitized_item.user_email = None
|
||||
sanitized_item.batch_id = "redacted"
|
||||
sanitized_item.session_id = "redacted"
|
||||
sanitized_item.origin = None
|
||||
sanitized_item.destination = None
|
||||
sanitized_item.priority = 0
|
||||
sanitized_item.field_values = None
|
||||
sanitized_item.retried_from_item_id = None
|
||||
sanitized_item.workflow_call_id = None
|
||||
sanitized_item.parent_item_id = None
|
||||
sanitized_item.parent_session_id = None
|
||||
sanitized_item.root_item_id = None
|
||||
sanitized_item.workflow_call_depth = None
|
||||
sanitized_item.workflow = None
|
||||
sanitized_item.error_type = None
|
||||
sanitized_item.error_message = None
|
||||
sanitized_item.error_traceback = None
|
||||
sanitized_item.session = GraphExecutionState(
|
||||
id="redacted",
|
||||
graph=Graph(),
|
||||
)
|
||||
return sanitized_item
|
||||
|
||||
|
||||
@session_queue_router.post(
|
||||
"/{queue_id}/enqueue_batch",
|
||||
operation_id="enqueue_batch",
|
||||
responses={
|
||||
201: {"model": EnqueueBatchResult},
|
||||
},
|
||||
)
|
||||
async def enqueue_batch(
|
||||
current_user: CurrentUserOrDefault,
|
||||
queue_id: str = Path(description="The queue id to perform this operation on"),
|
||||
batch: Batch = Body(description="Batch to process"),
|
||||
prepend: bool = Body(default=False, description="Whether or not to prepend this batch in the queue"),
|
||||
) -> EnqueueBatchResult:
|
||||
"""Processes a batch and enqueues the output graphs for execution for the current user."""
|
||||
assert_image_move_maintenance_inactive()
|
||||
|
||||
try:
|
||||
return await ApiDependencies.invoker.services.session_queue.enqueue_batch(
|
||||
queue_id=queue_id, batch=batch, prepend=prepend, user_id=current_user.user_id
|
||||
)
|
||||
except Exception as e:
|
||||
raise HTTPException(status_code=500, detail=f"Unexpected error while enqueuing batch: {e}")
|
||||
|
||||
|
||||
@session_queue_router.get(
|
||||
"/{queue_id}/list_all",
|
||||
operation_id="list_all_queue_items",
|
||||
responses={
|
||||
200: {"model": list[SessionQueueItem]},
|
||||
},
|
||||
)
|
||||
async def list_all_queue_items(
|
||||
current_user: CurrentUserOrDefault,
|
||||
queue_id: str = Path(description="The queue id to perform this operation on"),
|
||||
destination: Optional[str] = Query(default=None, description="The destination of queue items to fetch"),
|
||||
) -> list[SessionQueueItem]:
|
||||
"""Gets all queue items"""
|
||||
try:
|
||||
items = ApiDependencies.invoker.services.session_queue.list_all_queue_items(
|
||||
queue_id=queue_id,
|
||||
destination=destination,
|
||||
)
|
||||
# Sanitize items for non-admin users
|
||||
return [sanitize_queue_item_for_user(item, current_user.user_id, current_user.is_admin) for item in items]
|
||||
except Exception as e:
|
||||
raise HTTPException(status_code=500, detail=f"Unexpected error while listing all queue items: {e}")
|
||||
|
||||
|
||||
@session_queue_router.get(
|
||||
"/{queue_id}/item_ids",
|
||||
operation_id="get_queue_item_ids",
|
||||
responses={
|
||||
200: {"model": ItemIdsResult},
|
||||
},
|
||||
)
|
||||
async def get_queue_item_ids(
|
||||
current_user: CurrentUserOrDefault,
|
||||
queue_id: str = Path(description="The queue id to perform this operation on"),
|
||||
order_dir: SQLiteDirection = Query(default=SQLiteDirection.Descending, description="The order of sort"),
|
||||
) -> ItemIdsResult:
|
||||
"""Gets all queue item ids that match the given parameters.
|
||||
|
||||
IDs for every user's items are returned (item ids carry no sensitive data on their own).
|
||||
When the corresponding items are hydrated via get_queue_items_by_item_ids, those belonging
|
||||
to other users are redacted by sanitize_queue_item_for_user. This lets a non-admin see
|
||||
partially-redacted entries for other users' jobs in the queue list, while still revealing
|
||||
only timestamps and status for items they do not own.
|
||||
|
||||
current_user is required so the endpoint stays behind authentication in multiuser mode.
|
||||
"""
|
||||
try:
|
||||
return ApiDependencies.invoker.services.session_queue.get_queue_item_ids(queue_id=queue_id, order_dir=order_dir)
|
||||
except Exception as e:
|
||||
raise HTTPException(status_code=500, detail=f"Unexpected error while listing all queue item ids: {e}")
|
||||
|
||||
|
||||
@session_queue_router.post(
|
||||
"/{queue_id}/items_by_ids",
|
||||
operation_id="get_queue_items_by_item_ids",
|
||||
responses={200: {"model": list[SessionQueueItem]}},
|
||||
)
|
||||
async def get_queue_items_by_item_ids(
|
||||
current_user: CurrentUserOrDefault,
|
||||
queue_id: str = Path(description="The queue id to perform this operation on"),
|
||||
item_ids: list[int] = Body(
|
||||
embed=True, description="Object containing list of queue item ids to fetch queue items for"
|
||||
),
|
||||
) -> list[SessionQueueItem]:
|
||||
"""Gets queue items for the specified queue item ids. Maintains order of item ids."""
|
||||
try:
|
||||
session_queue_service = ApiDependencies.invoker.services.session_queue
|
||||
|
||||
# Fetch queue items preserving the order of requested item ids
|
||||
queue_items: list[SessionQueueItem] = []
|
||||
for item_id in item_ids:
|
||||
try:
|
||||
queue_item = session_queue_service.get_queue_item(item_id=item_id)
|
||||
if queue_item.queue_id != queue_id: # Auth protection for items from other queues
|
||||
continue
|
||||
# Sanitize item for non-admin users
|
||||
sanitized_item = sanitize_queue_item_for_user(queue_item, current_user.user_id, current_user.is_admin)
|
||||
queue_items.append(sanitized_item)
|
||||
except Exception:
|
||||
# Skip missing queue items - they may have been deleted between item id fetch and queue item fetch
|
||||
continue
|
||||
|
||||
return queue_items
|
||||
except Exception:
|
||||
raise HTTPException(status_code=500, detail="Failed to get queue items")
|
||||
|
||||
|
||||
@session_queue_router.put(
|
||||
"/{queue_id}/processor/resume",
|
||||
operation_id="resume",
|
||||
responses={200: {"model": SessionProcessorStatus}},
|
||||
)
|
||||
async def resume(
|
||||
current_user: AdminUserOrDefault,
|
||||
queue_id: str = Path(description="The queue id to perform this operation on"),
|
||||
) -> SessionProcessorStatus:
|
||||
"""Resumes session processor. Admin only."""
|
||||
try:
|
||||
return ApiDependencies.invoker.services.session_processor.resume()
|
||||
except Exception as e:
|
||||
raise HTTPException(status_code=500, detail=f"Unexpected error while resuming queue: {e}")
|
||||
|
||||
|
||||
@session_queue_router.put(
|
||||
"/{queue_id}/processor/pause",
|
||||
operation_id="pause",
|
||||
responses={200: {"model": SessionProcessorStatus}},
|
||||
)
|
||||
async def pause(
|
||||
current_user: AdminUserOrDefault,
|
||||
queue_id: str = Path(description="The queue id to perform this operation on"),
|
||||
) -> SessionProcessorStatus:
|
||||
"""Pauses session processor. Admin only."""
|
||||
try:
|
||||
return ApiDependencies.invoker.services.session_processor.pause()
|
||||
except Exception as e:
|
||||
raise HTTPException(status_code=500, detail=f"Unexpected error while pausing queue: {e}")
|
||||
|
||||
|
||||
@session_queue_router.put(
|
||||
"/{queue_id}/cancel_all_except_current",
|
||||
operation_id="cancel_all_except_current",
|
||||
responses={200: {"model": CancelAllExceptCurrentResult}},
|
||||
)
|
||||
async def cancel_all_except_current(
|
||||
current_user: CurrentUserOrDefault,
|
||||
queue_id: str = Path(description="The queue id to perform this operation on"),
|
||||
) -> CancelAllExceptCurrentResult:
|
||||
"""Immediately cancels all queue items except in-processing items. Non-admin users can only cancel their own items."""
|
||||
try:
|
||||
# Admin users can cancel all items, non-admin users can only cancel their own
|
||||
user_id = None if current_user.is_admin else current_user.user_id
|
||||
return ApiDependencies.invoker.services.session_queue.cancel_all_except_current(
|
||||
queue_id=queue_id, user_id=user_id
|
||||
)
|
||||
except Exception as e:
|
||||
raise HTTPException(status_code=500, detail=f"Unexpected error while canceling all except current: {e}")
|
||||
|
||||
|
||||
@session_queue_router.put(
|
||||
"/{queue_id}/delete_all_except_current",
|
||||
operation_id="delete_all_except_current",
|
||||
responses={200: {"model": DeleteAllExceptCurrentResult}},
|
||||
)
|
||||
async def delete_all_except_current(
|
||||
current_user: CurrentUserOrDefault,
|
||||
queue_id: str = Path(description="The queue id to perform this operation on"),
|
||||
) -> DeleteAllExceptCurrentResult:
|
||||
"""Immediately deletes all queue items except in-processing items. Non-admin users can only delete their own items."""
|
||||
try:
|
||||
# Admin users can delete all items, non-admin users can only delete their own
|
||||
user_id = None if current_user.is_admin else current_user.user_id
|
||||
return ApiDependencies.invoker.services.session_queue.delete_all_except_current(
|
||||
queue_id=queue_id, user_id=user_id
|
||||
)
|
||||
except Exception as e:
|
||||
raise HTTPException(status_code=500, detail=f"Unexpected error while deleting all except current: {e}")
|
||||
|
||||
|
||||
@session_queue_router.put(
|
||||
"/{queue_id}/cancel_by_batch_ids",
|
||||
operation_id="cancel_by_batch_ids",
|
||||
responses={200: {"model": CancelByBatchIDsResult}},
|
||||
)
|
||||
async def cancel_by_batch_ids(
|
||||
current_user: CurrentUserOrDefault,
|
||||
queue_id: str = Path(description="The queue id to perform this operation on"),
|
||||
batch_ids: list[str] = Body(description="The list of batch_ids to cancel all queue items for", embed=True),
|
||||
) -> CancelByBatchIDsResult:
|
||||
"""Immediately cancels all queue items from the given batch ids. Non-admin users can only cancel their own items."""
|
||||
try:
|
||||
# Admin users can cancel all items, non-admin users can only cancel their own
|
||||
user_id = None if current_user.is_admin else current_user.user_id
|
||||
return ApiDependencies.invoker.services.session_queue.cancel_by_batch_ids(
|
||||
queue_id=queue_id, batch_ids=batch_ids, user_id=user_id
|
||||
)
|
||||
except Exception as e:
|
||||
raise HTTPException(status_code=500, detail=f"Unexpected error while canceling by batch id: {e}")
|
||||
|
||||
|
||||
@session_queue_router.put(
|
||||
"/{queue_id}/cancel_by_destination",
|
||||
operation_id="cancel_by_destination",
|
||||
responses={200: {"model": CancelByDestinationResult}},
|
||||
)
|
||||
async def cancel_by_destination(
|
||||
current_user: CurrentUserOrDefault,
|
||||
queue_id: str = Path(description="The queue id to perform this operation on"),
|
||||
destination: str = Query(description="The destination to cancel all queue items for"),
|
||||
) -> CancelByDestinationResult:
|
||||
"""Immediately cancels all queue items with the given destination. Non-admin users can only cancel their own items."""
|
||||
try:
|
||||
# Admin users can cancel all items, non-admin users can only cancel their own
|
||||
user_id = None if current_user.is_admin else current_user.user_id
|
||||
return ApiDependencies.invoker.services.session_queue.cancel_by_destination(
|
||||
queue_id=queue_id, destination=destination, user_id=user_id
|
||||
)
|
||||
except Exception as e:
|
||||
raise HTTPException(status_code=500, detail=f"Unexpected error while canceling by destination: {e}")
|
||||
|
||||
|
||||
@session_queue_router.put(
|
||||
"/{queue_id}/retry_items_by_id",
|
||||
operation_id="retry_items_by_id",
|
||||
responses={200: {"model": RetryItemsResult}},
|
||||
)
|
||||
async def retry_items_by_id(
|
||||
current_user: CurrentUserOrDefault,
|
||||
queue_id: str = Path(description="The queue id to perform this operation on"),
|
||||
item_ids: list[int] = Body(description="The queue item ids to retry"),
|
||||
) -> RetryItemsResult:
|
||||
"""Retries the given queue items. Users can only retry their own items unless they are an admin."""
|
||||
try:
|
||||
# Check queue membership for all items and ownership for non-admins.
|
||||
valid_item_ids: list[int] = []
|
||||
for item_id in item_ids:
|
||||
try:
|
||||
queue_item = ApiDependencies.invoker.services.session_queue.get_queue_item(item_id)
|
||||
if queue_item.queue_id != queue_id:
|
||||
raise HTTPException(
|
||||
status_code=404, detail=f"Queue item with id {item_id} not found in queue {queue_id}"
|
||||
)
|
||||
root_queue_item = _get_workflow_call_root_queue_item(queue_item)
|
||||
if root_queue_item.queue_id != queue_id:
|
||||
raise HTTPException(
|
||||
status_code=404, detail=f"Queue item with id {item_id} not found in queue {queue_id}"
|
||||
)
|
||||
if not current_user.is_admin and root_queue_item.user_id != current_user.user_id:
|
||||
raise HTTPException(
|
||||
status_code=403, detail=f"You do not have permission to retry queue item {item_id}"
|
||||
)
|
||||
valid_item_ids.append(item_id)
|
||||
except SessionQueueItemNotFoundError:
|
||||
# Skip items that don't exist - they will be handled by retry_items_by_id
|
||||
continue
|
||||
|
||||
return ApiDependencies.invoker.services.session_queue.retry_items_by_id(
|
||||
queue_id=queue_id, item_ids=valid_item_ids
|
||||
)
|
||||
except HTTPException:
|
||||
raise
|
||||
except Exception as e:
|
||||
raise HTTPException(status_code=500, detail=f"Unexpected error while retrying queue items: {e}")
|
||||
|
||||
|
||||
@session_queue_router.put(
|
||||
"/{queue_id}/clear",
|
||||
operation_id="clear",
|
||||
responses={
|
||||
200: {"model": ClearResult},
|
||||
},
|
||||
)
|
||||
async def clear(
|
||||
current_user: CurrentUserOrDefault,
|
||||
queue_id: str = Path(description="The queue id to perform this operation on"),
|
||||
) -> ClearResult:
|
||||
"""Clears the queue entirely. Admin users clear all items; non-admin users only clear their own items. If there's a currently-executing item, users can only cancel it if they own it or are an admin."""
|
||||
try:
|
||||
queue_item = ApiDependencies.invoker.services.session_queue.get_current(queue_id)
|
||||
if queue_item is not None:
|
||||
# Check authorization for canceling the current item
|
||||
if queue_item.user_id != current_user.user_id and not current_user.is_admin:
|
||||
raise HTTPException(
|
||||
status_code=403, detail="You do not have permission to cancel the currently executing queue item"
|
||||
)
|
||||
ApiDependencies.invoker.services.session_queue.cancel_queue_item(queue_item.item_id)
|
||||
# Admin users can clear all items, non-admin users can only clear their own
|
||||
user_id = None if current_user.is_admin else current_user.user_id
|
||||
clear_result = ApiDependencies.invoker.services.session_queue.clear(queue_id, user_id=user_id)
|
||||
return clear_result
|
||||
except HTTPException:
|
||||
raise
|
||||
except Exception as e:
|
||||
raise HTTPException(status_code=500, detail=f"Unexpected error while clearing queue: {e}")
|
||||
|
||||
|
||||
@session_queue_router.put(
|
||||
"/{queue_id}/prune",
|
||||
operation_id="prune",
|
||||
responses={
|
||||
200: {"model": PruneResult},
|
||||
},
|
||||
)
|
||||
async def prune(
|
||||
current_user: CurrentUserOrDefault,
|
||||
queue_id: str = Path(description="The queue id to perform this operation on"),
|
||||
) -> PruneResult:
|
||||
"""Prunes all completed or errored queue items. Non-admin users can only prune their own items."""
|
||||
try:
|
||||
# Admin users can prune all items, non-admin users can only prune their own
|
||||
user_id = None if current_user.is_admin else current_user.user_id
|
||||
return ApiDependencies.invoker.services.session_queue.prune(queue_id, user_id=user_id)
|
||||
except Exception as e:
|
||||
raise HTTPException(status_code=500, detail=f"Unexpected error while pruning queue: {e}")
|
||||
|
||||
|
||||
@session_queue_router.get(
|
||||
"/{queue_id}/current",
|
||||
operation_id="get_current_queue_item",
|
||||
responses={
|
||||
200: {"model": Optional[SessionQueueItem]},
|
||||
},
|
||||
)
|
||||
async def get_current_queue_item(
|
||||
current_user: CurrentUserOrDefault,
|
||||
queue_id: str = Path(description="The queue id to perform this operation on"),
|
||||
) -> Optional[SessionQueueItem]:
|
||||
"""Gets the currently execution queue item"""
|
||||
try:
|
||||
item = ApiDependencies.invoker.services.session_queue.get_current(queue_id)
|
||||
if item is not None:
|
||||
item = sanitize_queue_item_for_user(item, current_user.user_id, current_user.is_admin)
|
||||
return item
|
||||
except Exception as e:
|
||||
raise HTTPException(status_code=500, detail=f"Unexpected error while getting current queue item: {e}")
|
||||
|
||||
|
||||
@session_queue_router.get(
|
||||
"/{queue_id}/next",
|
||||
operation_id="get_next_queue_item",
|
||||
responses={
|
||||
200: {"model": Optional[SessionQueueItem]},
|
||||
},
|
||||
)
|
||||
async def get_next_queue_item(
|
||||
current_user: CurrentUserOrDefault,
|
||||
queue_id: str = Path(description="The queue id to perform this operation on"),
|
||||
) -> Optional[SessionQueueItem]:
|
||||
"""Gets the next queue item, without executing it"""
|
||||
try:
|
||||
item = ApiDependencies.invoker.services.session_queue.get_next(queue_id)
|
||||
if item is not None:
|
||||
item = sanitize_queue_item_for_user(item, current_user.user_id, current_user.is_admin)
|
||||
return item
|
||||
except Exception as e:
|
||||
raise HTTPException(status_code=500, detail=f"Unexpected error while getting next queue item: {e}")
|
||||
|
||||
|
||||
@session_queue_router.get(
|
||||
"/{queue_id}/status",
|
||||
operation_id="get_queue_status",
|
||||
responses={
|
||||
200: {"model": SessionQueueAndProcessorStatus},
|
||||
},
|
||||
)
|
||||
async def get_queue_status(
|
||||
current_user: CurrentUserOrDefault,
|
||||
queue_id: str = Path(description="The queue id to perform this operation on"),
|
||||
) -> SessionQueueAndProcessorStatus:
|
||||
"""Gets the status of the session queue. Returns global counts; non-admin users additionally
|
||||
get their own pending/in_progress counts (so the UI can show an X/Y badge) and cannot see the
|
||||
current item's identifiers unless they own it."""
|
||||
try:
|
||||
user_id = None if current_user.is_admin else current_user.user_id
|
||||
queue = ApiDependencies.invoker.services.session_queue.get_queue_status(queue_id, user_id=user_id)
|
||||
processor = ApiDependencies.invoker.services.session_processor.get_status()
|
||||
return SessionQueueAndProcessorStatus(queue=queue, processor=processor)
|
||||
except Exception as e:
|
||||
raise HTTPException(status_code=500, detail=f"Unexpected error while getting queue status: {e}")
|
||||
|
||||
|
||||
@session_queue_router.get(
|
||||
"/{queue_id}/b/{batch_id}/status",
|
||||
operation_id="get_batch_status",
|
||||
responses={
|
||||
200: {"model": BatchStatus},
|
||||
},
|
||||
)
|
||||
async def get_batch_status(
|
||||
current_user: CurrentUserOrDefault,
|
||||
queue_id: str = Path(description="The queue id to perform this operation on"),
|
||||
batch_id: str = Path(description="The batch to get the status of"),
|
||||
) -> BatchStatus:
|
||||
"""Gets the status of a batch. Non-admin users only see their own batches."""
|
||||
try:
|
||||
user_id = None if current_user.is_admin else current_user.user_id
|
||||
return ApiDependencies.invoker.services.session_queue.get_batch_status(
|
||||
queue_id=queue_id, batch_id=batch_id, user_id=user_id
|
||||
)
|
||||
except Exception as e:
|
||||
raise HTTPException(status_code=500, detail=f"Unexpected error while getting batch status: {e}")
|
||||
|
||||
|
||||
@session_queue_router.get(
|
||||
"/{queue_id}/i/{item_id}",
|
||||
operation_id="get_queue_item",
|
||||
responses={
|
||||
200: {"model": SessionQueueItem},
|
||||
},
|
||||
response_model_exclude_none=True,
|
||||
)
|
||||
async def get_queue_item(
|
||||
current_user: CurrentUserOrDefault,
|
||||
queue_id: str = Path(description="The queue id to perform this operation on"),
|
||||
item_id: int = Path(description="The queue item to get"),
|
||||
) -> SessionQueueItem:
|
||||
"""Gets a queue item"""
|
||||
try:
|
||||
queue_item = ApiDependencies.invoker.services.session_queue.get_queue_item(item_id=item_id)
|
||||
if queue_item.queue_id != queue_id:
|
||||
raise HTTPException(status_code=404, detail=f"Queue item with id {item_id} not found in queue {queue_id}")
|
||||
# Sanitize item for non-admin users
|
||||
return sanitize_queue_item_for_user(queue_item, current_user.user_id, current_user.is_admin)
|
||||
except SessionQueueItemNotFoundError:
|
||||
raise HTTPException(status_code=404, detail=f"Queue item with id {item_id} not found in queue {queue_id}")
|
||||
except Exception as e:
|
||||
raise HTTPException(status_code=500, detail=f"Unexpected error while fetching queue item: {e}")
|
||||
|
||||
|
||||
@session_queue_router.delete(
|
||||
"/{queue_id}/i/{item_id}",
|
||||
operation_id="delete_queue_item",
|
||||
)
|
||||
async def delete_queue_item(
|
||||
current_user: CurrentUserOrDefault,
|
||||
queue_id: str = Path(description="The queue id to perform this operation on"),
|
||||
item_id: int = Path(description="The queue item to delete"),
|
||||
) -> None:
|
||||
"""Deletes a queue item. Users can only delete their own items unless they are an admin."""
|
||||
try:
|
||||
# Get the queue item to check ownership
|
||||
queue_item = ApiDependencies.invoker.services.session_queue.get_queue_item(item_id)
|
||||
if queue_item.queue_id != queue_id:
|
||||
raise HTTPException(status_code=404, detail=f"Queue item with id {item_id} not found in queue {queue_id}")
|
||||
|
||||
root_queue_item = _get_workflow_call_root_queue_item(queue_item)
|
||||
if root_queue_item.queue_id != queue_id:
|
||||
raise HTTPException(status_code=404, detail=f"Queue item with id {item_id} not found in queue {queue_id}")
|
||||
|
||||
# The queue service deletes the entire chain, so authorization must use the root owner.
|
||||
if root_queue_item.user_id != current_user.user_id and not current_user.is_admin:
|
||||
raise HTTPException(status_code=403, detail="You do not have permission to delete this queue item")
|
||||
|
||||
ApiDependencies.invoker.services.session_queue.delete_queue_item(item_id)
|
||||
except SessionQueueItemNotFoundError:
|
||||
raise HTTPException(status_code=404, detail=f"Queue item with id {item_id} not found in queue {queue_id}")
|
||||
except HTTPException:
|
||||
raise
|
||||
except Exception as e:
|
||||
raise HTTPException(status_code=500, detail=f"Unexpected error while deleting queue item: {e}")
|
||||
|
||||
|
||||
@session_queue_router.put(
|
||||
"/{queue_id}/i/{item_id}/cancel",
|
||||
operation_id="cancel_queue_item",
|
||||
responses={
|
||||
200: {"model": SessionQueueItem},
|
||||
},
|
||||
)
|
||||
async def cancel_queue_item(
|
||||
current_user: CurrentUserOrDefault,
|
||||
queue_id: str = Path(description="The queue id to perform this operation on"),
|
||||
item_id: int = Path(description="The queue item to cancel"),
|
||||
) -> SessionQueueItem:
|
||||
"""Cancels a queue item. Users can only cancel their own items unless they are an admin."""
|
||||
try:
|
||||
# Get the queue item to check ownership
|
||||
queue_item = ApiDependencies.invoker.services.session_queue.get_queue_item(item_id)
|
||||
if queue_item.queue_id != queue_id:
|
||||
raise HTTPException(status_code=404, detail=f"Queue item with id {item_id} not found in queue {queue_id}")
|
||||
|
||||
# Check authorization: user must own the item or be an admin
|
||||
if queue_item.user_id != current_user.user_id and not current_user.is_admin:
|
||||
raise HTTPException(status_code=403, detail="You do not have permission to cancel this queue item")
|
||||
|
||||
return ApiDependencies.invoker.services.session_queue.cancel_queue_item(item_id)
|
||||
except SessionQueueItemNotFoundError:
|
||||
raise HTTPException(status_code=404, detail=f"Queue item with id {item_id} not found in queue {queue_id}")
|
||||
except HTTPException:
|
||||
raise
|
||||
except Exception as e:
|
||||
raise HTTPException(status_code=500, detail=f"Unexpected error while canceling queue item: {e}")
|
||||
|
||||
|
||||
@session_queue_router.get(
|
||||
"/{queue_id}/counts_by_destination",
|
||||
operation_id="counts_by_destination",
|
||||
responses={200: {"model": SessionQueueCountsByDestination}},
|
||||
)
|
||||
async def counts_by_destination(
|
||||
current_user: CurrentUserOrDefault,
|
||||
queue_id: str = Path(description="The queue id to query"),
|
||||
destination: str = Query(description="The destination to query"),
|
||||
) -> SessionQueueCountsByDestination:
|
||||
"""Gets the counts of queue items by destination. Non-admin users only see their own items."""
|
||||
try:
|
||||
user_id = None if current_user.is_admin else current_user.user_id
|
||||
return ApiDependencies.invoker.services.session_queue.get_counts_by_destination(
|
||||
queue_id=queue_id, destination=destination, user_id=user_id
|
||||
)
|
||||
except Exception as e:
|
||||
raise HTTPException(status_code=500, detail=f"Unexpected error while fetching counts by destination: {e}")
|
||||
|
||||
|
||||
@session_queue_router.delete(
|
||||
"/{queue_id}/d/{destination}",
|
||||
operation_id="delete_by_destination",
|
||||
responses={200: {"model": DeleteByDestinationResult}},
|
||||
)
|
||||
async def delete_by_destination(
|
||||
current_user: CurrentUserOrDefault,
|
||||
queue_id: str = Path(description="The queue id to query"),
|
||||
destination: str = Path(description="The destination to query"),
|
||||
) -> DeleteByDestinationResult:
|
||||
"""Deletes all items with the given destination. Non-admin users can only delete their own items."""
|
||||
try:
|
||||
# Admin users can delete all items, non-admin users can only delete their own
|
||||
user_id = None if current_user.is_admin else current_user.user_id
|
||||
return ApiDependencies.invoker.services.session_queue.delete_by_destination(
|
||||
queue_id=queue_id, destination=destination, user_id=user_id
|
||||
)
|
||||
except Exception as e:
|
||||
raise HTTPException(status_code=500, detail=f"Unexpected error while deleting by destination: {e}")
|
||||
@@ -0,0 +1,339 @@
|
||||
import csv
|
||||
import io
|
||||
import json
|
||||
import traceback
|
||||
from typing import Optional
|
||||
|
||||
import pydantic
|
||||
from fastapi import APIRouter, File, Form, HTTPException, Path, Response, UploadFile
|
||||
from fastapi.responses import FileResponse
|
||||
from PIL import Image
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
from invokeai.app.api.auth_dependencies import AdminUserOrDefault, CurrentUserOrDefault
|
||||
from invokeai.app.api.dependencies import ApiDependencies
|
||||
from invokeai.app.api.routers.model_manager import IMAGE_MAX_AGE
|
||||
from invokeai.app.services.auth.token_service import TokenData
|
||||
from invokeai.app.services.style_preset_images.style_preset_images_common import StylePresetImageFileNotFoundException
|
||||
from invokeai.app.services.style_preset_records.style_preset_records_common import (
|
||||
InvalidPresetImportDataError,
|
||||
PresetData,
|
||||
PresetType,
|
||||
StylePresetChanges,
|
||||
StylePresetNotFoundError,
|
||||
StylePresetRecordDTO,
|
||||
StylePresetRecordWithImage,
|
||||
StylePresetWithoutId,
|
||||
UnsupportedFileTypeError,
|
||||
parse_presets_from_file,
|
||||
)
|
||||
|
||||
|
||||
class StylePresetFormData(BaseModel):
|
||||
name: str = Field(description="Preset name")
|
||||
positive_prompt: str = Field(description="Positive prompt")
|
||||
negative_prompt: str = Field(description="Negative prompt")
|
||||
type: PresetType = Field(description="Preset type")
|
||||
is_public: bool = Field(default=False, description="Whether the preset is visible to other users")
|
||||
|
||||
|
||||
style_presets_router = APIRouter(prefix="/v1/style_presets", tags=["style_presets"])
|
||||
|
||||
|
||||
def _assert_preset_read(record: StylePresetRecordDTO, current_user: TokenData) -> None:
|
||||
"""Allow read access if admin, owner, default preset, or public preset."""
|
||||
if current_user.is_admin:
|
||||
return
|
||||
if record.type == PresetType.Default:
|
||||
return
|
||||
if record.is_public:
|
||||
return
|
||||
if record.user_id == current_user.user_id:
|
||||
return
|
||||
raise HTTPException(status_code=403, detail="Not authorized to access this style preset")
|
||||
|
||||
|
||||
def _assert_preset_write(record: StylePresetRecordDTO, current_user: TokenData) -> None:
|
||||
"""Allow write access only for admin or owner. Defaults are immutable for non-admins."""
|
||||
if current_user.is_admin:
|
||||
return
|
||||
if record.type == PresetType.Default:
|
||||
raise HTTPException(status_code=403, detail="Default style presets cannot be modified")
|
||||
if record.user_id == current_user.user_id:
|
||||
return
|
||||
raise HTTPException(status_code=403, detail="Not authorized to modify this style preset")
|
||||
|
||||
|
||||
def _load_record_or_404(style_preset_id: str) -> StylePresetRecordDTO:
|
||||
try:
|
||||
return ApiDependencies.invoker.services.style_preset_records.get(style_preset_id)
|
||||
except StylePresetNotFoundError:
|
||||
raise HTTPException(status_code=404, detail="Style preset not found")
|
||||
|
||||
|
||||
@style_presets_router.get(
|
||||
"/i/{style_preset_id}",
|
||||
operation_id="get_style_preset",
|
||||
responses={
|
||||
200: {"model": StylePresetRecordWithImage},
|
||||
},
|
||||
)
|
||||
async def get_style_preset(
|
||||
current_user: CurrentUserOrDefault,
|
||||
style_preset_id: str = Path(description="The style preset to get"),
|
||||
) -> StylePresetRecordWithImage:
|
||||
"""Gets a style preset"""
|
||||
record = _load_record_or_404(style_preset_id)
|
||||
_assert_preset_read(record, current_user)
|
||||
image = ApiDependencies.invoker.services.style_preset_image_files.get_url(style_preset_id)
|
||||
return StylePresetRecordWithImage(image=image, **record.model_dump())
|
||||
|
||||
|
||||
@style_presets_router.patch(
|
||||
"/i/{style_preset_id}",
|
||||
operation_id="update_style_preset",
|
||||
responses={
|
||||
200: {"model": StylePresetRecordWithImage},
|
||||
},
|
||||
)
|
||||
async def update_style_preset(
|
||||
current_user: CurrentUserOrDefault,
|
||||
image: Optional[UploadFile] = File(description="The image file to upload", default=None),
|
||||
style_preset_id: str = Path(description="The id of the style preset to update"),
|
||||
data: str = Form(description="The data of the style preset to update"),
|
||||
) -> StylePresetRecordWithImage:
|
||||
"""Updates a style preset"""
|
||||
# Validate the data payload BEFORE any image-state mutation so a malformed
|
||||
# request can't leave the preset image partially updated.
|
||||
try:
|
||||
parsed_data = json.loads(data)
|
||||
validated_data = StylePresetFormData(**parsed_data)
|
||||
|
||||
name = validated_data.name
|
||||
type = validated_data.type
|
||||
positive_prompt = validated_data.positive_prompt
|
||||
negative_prompt = validated_data.negative_prompt
|
||||
is_public = validated_data.is_public
|
||||
|
||||
except (json.JSONDecodeError, pydantic.ValidationError):
|
||||
raise HTTPException(status_code=400, detail="Invalid preset data")
|
||||
|
||||
record = _load_record_or_404(style_preset_id)
|
||||
_assert_preset_write(record, current_user)
|
||||
|
||||
if image is not None:
|
||||
if not image.content_type or not image.content_type.startswith("image"):
|
||||
raise HTTPException(status_code=415, detail="Not an image")
|
||||
|
||||
contents = await image.read()
|
||||
try:
|
||||
pil_image = Image.open(io.BytesIO(contents))
|
||||
|
||||
except Exception:
|
||||
ApiDependencies.invoker.services.logger.error(traceback.format_exc())
|
||||
raise HTTPException(status_code=415, detail="Failed to read image")
|
||||
|
||||
try:
|
||||
ApiDependencies.invoker.services.style_preset_image_files.save(style_preset_id, pil_image)
|
||||
except ValueError as e:
|
||||
raise HTTPException(status_code=409, detail=str(e))
|
||||
else:
|
||||
try:
|
||||
ApiDependencies.invoker.services.style_preset_image_files.delete(style_preset_id)
|
||||
except StylePresetImageFileNotFoundException:
|
||||
pass
|
||||
|
||||
preset_data = PresetData(positive_prompt=positive_prompt, negative_prompt=negative_prompt)
|
||||
changes = StylePresetChanges(name=name, preset_data=preset_data, type=type, is_public=is_public)
|
||||
|
||||
style_preset_image = ApiDependencies.invoker.services.style_preset_image_files.get_url(style_preset_id)
|
||||
style_preset = ApiDependencies.invoker.services.style_preset_records.update(
|
||||
style_preset_id=style_preset_id, changes=changes
|
||||
)
|
||||
return StylePresetRecordWithImage(image=style_preset_image, **style_preset.model_dump())
|
||||
|
||||
|
||||
@style_presets_router.delete(
|
||||
"/i/{style_preset_id}",
|
||||
operation_id="delete_style_preset",
|
||||
)
|
||||
async def delete_style_preset(
|
||||
current_user: CurrentUserOrDefault,
|
||||
style_preset_id: str = Path(description="The style preset to delete"),
|
||||
) -> None:
|
||||
"""Deletes a style preset"""
|
||||
record = _load_record_or_404(style_preset_id)
|
||||
_assert_preset_write(record, current_user)
|
||||
|
||||
try:
|
||||
ApiDependencies.invoker.services.style_preset_image_files.delete(style_preset_id)
|
||||
except StylePresetImageFileNotFoundException:
|
||||
pass
|
||||
|
||||
ApiDependencies.invoker.services.style_preset_records.delete(style_preset_id)
|
||||
|
||||
|
||||
@style_presets_router.post(
|
||||
"/",
|
||||
operation_id="create_style_preset",
|
||||
responses={
|
||||
200: {"model": StylePresetRecordWithImage},
|
||||
},
|
||||
)
|
||||
async def create_style_preset(
|
||||
current_user: CurrentUserOrDefault,
|
||||
image: Optional[UploadFile] = File(description="The image file to upload", default=None),
|
||||
data: str = Form(description="The data of the style preset to create"),
|
||||
) -> StylePresetRecordWithImage:
|
||||
"""Creates a style preset"""
|
||||
|
||||
try:
|
||||
parsed_data = json.loads(data)
|
||||
validated_data = StylePresetFormData(**parsed_data)
|
||||
|
||||
name = validated_data.name
|
||||
type = validated_data.type
|
||||
positive_prompt = validated_data.positive_prompt
|
||||
negative_prompt = validated_data.negative_prompt
|
||||
is_public = validated_data.is_public
|
||||
|
||||
except (json.JSONDecodeError, pydantic.ValidationError):
|
||||
raise HTTPException(status_code=400, detail="Invalid preset data")
|
||||
|
||||
# Only admins may create default-typed presets — they're the shipped catalog.
|
||||
if type == PresetType.Default and not current_user.is_admin:
|
||||
raise HTTPException(status_code=403, detail="Only admins can create default presets")
|
||||
|
||||
preset_data = PresetData(positive_prompt=positive_prompt, negative_prompt=negative_prompt)
|
||||
style_preset = StylePresetWithoutId(name=name, preset_data=preset_data, type=type, is_public=is_public)
|
||||
new_style_preset = ApiDependencies.invoker.services.style_preset_records.create(
|
||||
style_preset=style_preset, user_id=current_user.user_id
|
||||
)
|
||||
|
||||
if image is not None:
|
||||
if not image.content_type or not image.content_type.startswith("image"):
|
||||
raise HTTPException(status_code=415, detail="Not an image")
|
||||
|
||||
contents = await image.read()
|
||||
try:
|
||||
pil_image = Image.open(io.BytesIO(contents))
|
||||
|
||||
except Exception:
|
||||
ApiDependencies.invoker.services.logger.error(traceback.format_exc())
|
||||
raise HTTPException(status_code=415, detail="Failed to read image")
|
||||
|
||||
try:
|
||||
ApiDependencies.invoker.services.style_preset_image_files.save(new_style_preset.id, pil_image)
|
||||
except ValueError as e:
|
||||
raise HTTPException(status_code=409, detail=str(e))
|
||||
|
||||
preset_image = ApiDependencies.invoker.services.style_preset_image_files.get_url(new_style_preset.id)
|
||||
return StylePresetRecordWithImage(image=preset_image, **new_style_preset.model_dump())
|
||||
|
||||
|
||||
@style_presets_router.get(
|
||||
"/",
|
||||
operation_id="list_style_presets",
|
||||
responses={
|
||||
200: {"model": list[StylePresetRecordWithImage]},
|
||||
},
|
||||
)
|
||||
async def list_style_presets(current_user: CurrentUserOrDefault) -> list[StylePresetRecordWithImage]:
|
||||
"""Gets the style presets visible to the current user."""
|
||||
style_presets_with_image: list[StylePresetRecordWithImage] = []
|
||||
style_presets = ApiDependencies.invoker.services.style_preset_records.get_many(
|
||||
user_id=current_user.user_id,
|
||||
is_admin=current_user.is_admin,
|
||||
)
|
||||
for preset in style_presets:
|
||||
image = ApiDependencies.invoker.services.style_preset_image_files.get_url(preset.id)
|
||||
style_preset_with_image = StylePresetRecordWithImage(image=image, **preset.model_dump())
|
||||
style_presets_with_image.append(style_preset_with_image)
|
||||
|
||||
return style_presets_with_image
|
||||
|
||||
|
||||
@style_presets_router.get(
|
||||
"/i/{style_preset_id}/image",
|
||||
operation_id="get_style_preset_image",
|
||||
responses={
|
||||
200: {
|
||||
"description": "The style preset image was fetched successfully",
|
||||
},
|
||||
400: {"description": "Bad request"},
|
||||
404: {"description": "The style preset image could not be found"},
|
||||
},
|
||||
status_code=200,
|
||||
)
|
||||
async def get_style_preset_image(
|
||||
current_user: CurrentUserOrDefault,
|
||||
style_preset_id: str = Path(description="The id of the style preset image to get"),
|
||||
) -> FileResponse:
|
||||
"""Gets an image file that previews the model"""
|
||||
record = _load_record_or_404(style_preset_id)
|
||||
_assert_preset_read(record, current_user)
|
||||
|
||||
try:
|
||||
path = ApiDependencies.invoker.services.style_preset_image_files.get_path(style_preset_id)
|
||||
|
||||
response = FileResponse(
|
||||
path,
|
||||
media_type="image/png",
|
||||
filename=style_preset_id + ".png",
|
||||
content_disposition_type="inline",
|
||||
)
|
||||
response.headers["Cache-Control"] = f"max-age={IMAGE_MAX_AGE}"
|
||||
return response
|
||||
except Exception:
|
||||
raise HTTPException(status_code=404)
|
||||
|
||||
|
||||
@style_presets_router.get(
|
||||
"/export",
|
||||
operation_id="export_style_presets",
|
||||
responses={200: {"content": {"text/csv": {}}, "description": "A CSV file with the requested data."}},
|
||||
status_code=200,
|
||||
)
|
||||
async def export_style_presets(current_user: AdminUserOrDefault):
|
||||
# Admin-only export covers every user preset.
|
||||
output = io.StringIO()
|
||||
writer = csv.writer(output)
|
||||
|
||||
writer.writerow(["name", "prompt", "negative_prompt"])
|
||||
|
||||
style_presets = ApiDependencies.invoker.services.style_preset_records.get_many(
|
||||
type=PresetType.User,
|
||||
user_id=current_user.user_id,
|
||||
is_admin=True,
|
||||
)
|
||||
|
||||
for preset in style_presets:
|
||||
writer.writerow([preset.name, preset.preset_data.positive_prompt, preset.preset_data.negative_prompt])
|
||||
|
||||
csv_data = output.getvalue()
|
||||
output.close()
|
||||
|
||||
return Response(
|
||||
content=csv_data,
|
||||
media_type="text/csv",
|
||||
headers={"Content-Disposition": "attachment; filename=prompt_templates.csv"},
|
||||
)
|
||||
|
||||
|
||||
@style_presets_router.post(
|
||||
"/import",
|
||||
operation_id="import_style_presets",
|
||||
)
|
||||
async def import_style_presets(
|
||||
current_user: AdminUserOrDefault,
|
||||
file: UploadFile = File(description="The file to import"),
|
||||
):
|
||||
try:
|
||||
style_presets = await parse_presets_from_file(file)
|
||||
ApiDependencies.invoker.services.style_preset_records.create_many(style_presets, user_id=current_user.user_id)
|
||||
except InvalidPresetImportDataError as e:
|
||||
ApiDependencies.invoker.services.logger.error(traceback.format_exc())
|
||||
raise HTTPException(status_code=400, detail=str(e))
|
||||
except UnsupportedFileTypeError as e:
|
||||
ApiDependencies.invoker.services.logger.error(traceback.format_exc())
|
||||
raise HTTPException(status_code=415, detail=str(e))
|
||||
@@ -0,0 +1,242 @@
|
||||
import asyncio
|
||||
import logging
|
||||
import threading
|
||||
from pathlib import Path
|
||||
from typing import Optional, Union
|
||||
|
||||
import torch
|
||||
from dynamicprompts.generators import CombinatorialPromptGenerator, RandomPromptGenerator
|
||||
from fastapi import Body, HTTPException
|
||||
from fastapi.routing import APIRouter
|
||||
from pydantic import BaseModel, Field
|
||||
from pyparsing import ParseException
|
||||
from transformers import AutoProcessor, AutoTokenizer, LlavaOnevisionForConditionalGeneration, LlavaOnevisionProcessor
|
||||
|
||||
from invokeai.app.api.auth_dependencies import CurrentUserOrDefault
|
||||
from invokeai.app.api.dependencies import ApiDependencies
|
||||
from invokeai.app.api.routers._access import assert_image_read_access
|
||||
from invokeai.app.api.routers.image_move_maintenance import assert_image_move_maintenance_inactive
|
||||
from invokeai.app.services.image_files.image_files_common import ImageFileNotFoundException
|
||||
from invokeai.app.services.model_records.model_records_base import UnknownModelException
|
||||
from invokeai.app.util.dynamicprompts import find_missing_wildcards
|
||||
from invokeai.backend.llava_onevision_pipeline import LlavaOnevisionPipeline
|
||||
from invokeai.backend.model_manager.taxonomy import ModelType
|
||||
from invokeai.backend.text_llm_pipeline import DEFAULT_SYSTEM_PROMPT, TextLLMPipeline
|
||||
from invokeai.backend.util.devices import TorchDevice
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
utilities_router = APIRouter(prefix="/v1/utilities", tags=["utilities"])
|
||||
|
||||
# The underlying model loader is not thread-safe, so we serialize load_model calls.
|
||||
_model_load_lock = threading.Lock()
|
||||
|
||||
|
||||
class DynamicPromptsResponse(BaseModel):
|
||||
prompts: list[str]
|
||||
error: Optional[str] = None
|
||||
|
||||
|
||||
@utilities_router.post(
|
||||
"/dynamicprompts",
|
||||
operation_id="parse_dynamicprompts",
|
||||
responses={
|
||||
200: {"model": DynamicPromptsResponse},
|
||||
},
|
||||
)
|
||||
async def parse_dynamicprompts(
|
||||
current_user: CurrentUserOrDefault,
|
||||
prompt: str = Body(description="The prompt to parse with dynamicprompts"),
|
||||
max_prompts: int = Body(ge=1, le=10000, default=1000, description="The max number of prompts to generate"),
|
||||
combinatorial: bool = Body(default=True, description="Whether to use the combinatorial generator"),
|
||||
seed: int | None = Body(None, description="The seed to use for random generation. Only used if not combinatorial"),
|
||||
) -> DynamicPromptsResponse:
|
||||
"""Creates a batch process"""
|
||||
max_prompts = min(max_prompts, 10000)
|
||||
generator: Union[RandomPromptGenerator, CombinatorialPromptGenerator]
|
||||
error: Optional[str] = None
|
||||
|
||||
# An unknown wildcard used as a variant value sends the combinatorial generator into an infinite
|
||||
# loop, so bail out early with a clear message instead of hanging the request (and with it the UI
|
||||
# preview). The random generator handles unknown wildcards gracefully, so only the combinatorial
|
||||
# path is guarded.
|
||||
if combinatorial:
|
||||
missing_wildcards = find_missing_wildcards(prompt)
|
||||
if missing_wildcards:
|
||||
wildcards = ", ".join(missing_wildcards)
|
||||
return DynamicPromptsResponse(prompts=[prompt], error=f"No values found for wildcard(s): {wildcards}")
|
||||
|
||||
try:
|
||||
if combinatorial:
|
||||
generator = CombinatorialPromptGenerator()
|
||||
prompts = generator.generate(prompt, max_prompts=max_prompts)
|
||||
else:
|
||||
generator = RandomPromptGenerator(seed=seed)
|
||||
prompts = generator.generate(prompt, num_images=max_prompts)
|
||||
except ParseException as e:
|
||||
prompts = [prompt]
|
||||
error = str(e)
|
||||
return DynamicPromptsResponse(prompts=prompts if prompts else [""], error=error)
|
||||
|
||||
|
||||
# --- Expand Prompt ---
|
||||
|
||||
|
||||
class ExpandPromptRequest(BaseModel):
|
||||
prompt: str
|
||||
model_key: str
|
||||
max_tokens: int = Field(default=300, ge=1, le=2048)
|
||||
system_prompt: str | None = None
|
||||
|
||||
|
||||
class ExpandPromptResponse(BaseModel):
|
||||
expanded_prompt: str
|
||||
error: str | None = None
|
||||
|
||||
|
||||
def _resolve_model_path(model_config_path: str) -> Path:
|
||||
"""Resolve a model config path to an absolute path."""
|
||||
model_path = Path(model_config_path)
|
||||
if model_path.is_absolute():
|
||||
return model_path.resolve()
|
||||
base_models_path = ApiDependencies.invoker.services.configuration.models_path
|
||||
return (base_models_path / model_path).resolve()
|
||||
|
||||
|
||||
def _run_expand_prompt(prompt: str, model_key: str, max_tokens: int, system_prompt: str | None) -> str:
|
||||
"""Run text LLM inference synchronously (called from thread)."""
|
||||
model_manager = ApiDependencies.invoker.services.model_manager
|
||||
model_config = model_manager.store.get_model(model_key)
|
||||
|
||||
if model_config.type != ModelType.TextLLM:
|
||||
raise ValueError(f"Model '{model_key}' is not a TextLLM model (got {model_config.type})")
|
||||
|
||||
with _model_load_lock:
|
||||
loaded_model = model_manager.load.load_model(model_config)
|
||||
|
||||
with torch.no_grad(), loaded_model.model_on_device() as (_, model):
|
||||
model_abs_path = _resolve_model_path(model_config.path)
|
||||
tokenizer = AutoTokenizer.from_pretrained(model_abs_path, local_files_only=True)
|
||||
|
||||
pipeline = TextLLMPipeline(model, tokenizer)
|
||||
model_device = next(model.parameters()).device
|
||||
output = pipeline.run(
|
||||
prompt=prompt,
|
||||
system_prompt=system_prompt or DEFAULT_SYSTEM_PROMPT,
|
||||
max_new_tokens=max_tokens,
|
||||
device=model_device,
|
||||
dtype=TorchDevice.choose_torch_dtype(),
|
||||
)
|
||||
|
||||
return output
|
||||
|
||||
|
||||
@utilities_router.post(
|
||||
"/expand-prompt",
|
||||
operation_id="expand_prompt",
|
||||
responses={
|
||||
200: {"model": ExpandPromptResponse},
|
||||
},
|
||||
)
|
||||
async def expand_prompt(current_user: CurrentUserOrDefault, body: ExpandPromptRequest) -> ExpandPromptResponse:
|
||||
"""Expand a brief prompt into a detailed image generation prompt using a text LLM."""
|
||||
try:
|
||||
expanded = await asyncio.to_thread(
|
||||
_run_expand_prompt,
|
||||
body.prompt,
|
||||
body.model_key,
|
||||
body.max_tokens,
|
||||
body.system_prompt,
|
||||
)
|
||||
return ExpandPromptResponse(expanded_prompt=expanded)
|
||||
except UnknownModelException:
|
||||
raise HTTPException(status_code=404, detail=f"Model '{body.model_key}' not found")
|
||||
except ValueError as e:
|
||||
raise HTTPException(status_code=422, detail=str(e))
|
||||
except Exception as e:
|
||||
logger.error(f"Error expanding prompt: {e}")
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
|
||||
# --- Image to Prompt ---
|
||||
|
||||
|
||||
class ImageToPromptRequest(BaseModel):
|
||||
image_name: str
|
||||
model_key: str
|
||||
instruction: str = "Describe this image in detail for use as an AI image generation prompt."
|
||||
|
||||
|
||||
class ImageToPromptResponse(BaseModel):
|
||||
prompt: str
|
||||
error: str | None = None
|
||||
|
||||
|
||||
def _run_image_to_prompt(image_name: str, model_key: str, instruction: str) -> str:
|
||||
"""Run LLaVA OneVision inference synchronously (called from thread)."""
|
||||
model_manager = ApiDependencies.invoker.services.model_manager
|
||||
model_config = model_manager.store.get_model(model_key)
|
||||
|
||||
if model_config.type != ModelType.LlavaOnevision:
|
||||
raise ValueError(f"Model '{model_key}' is not a LLaVA OneVision model (got {model_config.type})")
|
||||
|
||||
with _model_load_lock:
|
||||
loaded_model = model_manager.load.load_model(model_config)
|
||||
|
||||
# Load the image from InvokeAI's image store
|
||||
image = ApiDependencies.invoker.services.images.get_pil_image(image_name)
|
||||
image = image.convert("RGB")
|
||||
|
||||
with torch.no_grad(), loaded_model.model_on_device() as (_, model):
|
||||
if not isinstance(model, LlavaOnevisionForConditionalGeneration):
|
||||
raise TypeError(f"Expected LlavaOnevisionForConditionalGeneration, got {type(model).__name__}")
|
||||
|
||||
model_abs_path = _resolve_model_path(model_config.path)
|
||||
processor = AutoProcessor.from_pretrained(model_abs_path, local_files_only=True)
|
||||
if not isinstance(processor, LlavaOnevisionProcessor):
|
||||
raise TypeError(f"Expected LlavaOnevisionProcessor, got {type(processor).__name__}")
|
||||
|
||||
pipeline = LlavaOnevisionPipeline(model, processor)
|
||||
model_device = next(model.parameters()).device
|
||||
output = pipeline.run(
|
||||
prompt=instruction,
|
||||
images=[image],
|
||||
device=model_device,
|
||||
dtype=TorchDevice.choose_torch_dtype(),
|
||||
)
|
||||
|
||||
return output
|
||||
|
||||
|
||||
@utilities_router.post(
|
||||
"/image-to-prompt",
|
||||
operation_id="image_to_prompt",
|
||||
responses={
|
||||
200: {"model": ImageToPromptResponse},
|
||||
},
|
||||
)
|
||||
async def image_to_prompt(current_user: CurrentUserOrDefault, body: ImageToPromptRequest) -> ImageToPromptResponse:
|
||||
"""Generate a descriptive prompt from an image using a vision-language model."""
|
||||
assert_image_move_maintenance_inactive()
|
||||
|
||||
# Reuse the image-read access check so non-owners can't probe stored images
|
||||
# via this endpoint (mirrors the policy in routers/images.py).
|
||||
assert_image_read_access(body.image_name, current_user)
|
||||
|
||||
try:
|
||||
prompt = await asyncio.to_thread(
|
||||
_run_image_to_prompt,
|
||||
body.image_name,
|
||||
body.model_key,
|
||||
body.instruction,
|
||||
)
|
||||
return ImageToPromptResponse(prompt=prompt)
|
||||
except UnknownModelException:
|
||||
raise HTTPException(status_code=404, detail=f"Model '{body.model_key}' not found")
|
||||
except ImageFileNotFoundException:
|
||||
raise HTTPException(status_code=404, detail=f"Image '{body.image_name}' not found")
|
||||
except (ValueError, TypeError) as e:
|
||||
raise HTTPException(status_code=422, detail=str(e))
|
||||
except Exception as e:
|
||||
logger.error(f"Error generating prompt from image: {e}")
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
@@ -0,0 +1,56 @@
|
||||
from fastapi import HTTPException, Path, Query
|
||||
from fastapi.routing import APIRouter
|
||||
|
||||
from invokeai.app.api.auth_dependencies import CurrentUserOrDefault
|
||||
from invokeai.app.api.dependencies import ApiDependencies
|
||||
from invokeai.app.services.image_records.image_records_common import ImageCategory, ImageNamesResult
|
||||
from invokeai.app.services.shared.sqlite.sqlite_common import SQLiteDirection
|
||||
from invokeai.app.services.virtual_boards.virtual_boards_common import VirtualSubBoardDTO
|
||||
|
||||
virtual_boards_router = APIRouter(prefix="/v1/virtual_boards", tags=["virtual_boards"])
|
||||
|
||||
|
||||
@virtual_boards_router.get(
|
||||
"/by_date",
|
||||
operation_id="list_virtual_boards_by_date",
|
||||
response_model=list[VirtualSubBoardDTO],
|
||||
)
|
||||
async def list_virtual_boards_by_date(
|
||||
current_user: CurrentUserOrDefault,
|
||||
) -> list[VirtualSubBoardDTO]:
|
||||
"""Gets a list of virtual sub-boards grouped by date."""
|
||||
try:
|
||||
return ApiDependencies.invoker.services.image_records.get_image_dates(
|
||||
user_id=current_user.user_id,
|
||||
is_admin=current_user.is_admin,
|
||||
)
|
||||
except Exception:
|
||||
raise HTTPException(status_code=500, detail="Failed to get virtual boards by date")
|
||||
|
||||
|
||||
@virtual_boards_router.get(
|
||||
"/by_date/{date}/image_names",
|
||||
operation_id="list_virtual_board_image_names_by_date",
|
||||
response_model=ImageNamesResult,
|
||||
)
|
||||
async def list_virtual_board_image_names_by_date(
|
||||
current_user: CurrentUserOrDefault,
|
||||
date: str = Path(description="The ISO date string, e.g. '2026-03-18'"),
|
||||
starred_first: bool = Query(default=True, description="Whether to sort starred images first"),
|
||||
order_dir: SQLiteDirection = Query(default=SQLiteDirection.Descending, description="The sort direction"),
|
||||
categories: list[ImageCategory] | None = Query(default=None, description="The categories of images to include"),
|
||||
search_term: str | None = Query(default=None, description="Search term to filter images"),
|
||||
) -> ImageNamesResult:
|
||||
"""Gets ordered image names for a specific date."""
|
||||
try:
|
||||
return ApiDependencies.invoker.services.image_records.get_image_names_by_date(
|
||||
date=date,
|
||||
starred_first=starred_first,
|
||||
order_dir=order_dir,
|
||||
categories=categories,
|
||||
search_term=search_term,
|
||||
user_id=current_user.user_id,
|
||||
is_admin=current_user.is_admin,
|
||||
)
|
||||
except Exception:
|
||||
raise HTTPException(status_code=500, detail="Failed to get image names for date")
|
||||
@@ -0,0 +1,479 @@
|
||||
import io
|
||||
import math
|
||||
import traceback
|
||||
from typing import Optional
|
||||
|
||||
from fastapi import APIRouter, Body, File, HTTPException, Path, Query, UploadFile
|
||||
from fastapi.responses import FileResponse
|
||||
from PIL import Image
|
||||
|
||||
from invokeai.app.api.auth_dependencies import CurrentUserOrDefault
|
||||
from invokeai.app.api.dependencies import ApiDependencies
|
||||
from invokeai.app.services.shared.pagination import PaginatedResults
|
||||
from invokeai.app.services.shared.sqlite.sqlite_common import SQLiteDirection
|
||||
from invokeai.app.services.shared.workflow_call_compatibility import get_workflow_call_compatibility
|
||||
from invokeai.app.services.workflow_records.workflow_records_common import (
|
||||
Workflow,
|
||||
WorkflowCategory,
|
||||
WorkflowNotFoundError,
|
||||
WorkflowRecordDTO,
|
||||
WorkflowRecordListItemWithThumbnailDTO,
|
||||
WorkflowRecordOrderBy,
|
||||
WorkflowRecordWithThumbnailDTO,
|
||||
WorkflowWithoutID,
|
||||
)
|
||||
from invokeai.app.services.workflow_thumbnails.workflow_thumbnails_common import WorkflowThumbnailFileNotFoundException
|
||||
|
||||
IMAGE_MAX_AGE = 31536000
|
||||
workflows_router = APIRouter(prefix="/v1/workflows", tags=["workflows"])
|
||||
|
||||
|
||||
@workflows_router.get(
|
||||
"/i/{workflow_id}",
|
||||
operation_id="get_workflow",
|
||||
responses={
|
||||
200: {"model": WorkflowRecordWithThumbnailDTO},
|
||||
},
|
||||
)
|
||||
async def get_workflow(
|
||||
current_user: CurrentUserOrDefault,
|
||||
workflow_id: str = Path(description="The workflow to get"),
|
||||
) -> WorkflowRecordWithThumbnailDTO:
|
||||
"""Gets a workflow"""
|
||||
try:
|
||||
workflow = ApiDependencies.invoker.services.workflow_records.get(workflow_id)
|
||||
except WorkflowNotFoundError:
|
||||
raise HTTPException(status_code=404, detail="Workflow not found")
|
||||
|
||||
config = ApiDependencies.invoker.services.configuration
|
||||
if config.multiuser:
|
||||
is_default = workflow.workflow.meta.category is WorkflowCategory.Default
|
||||
is_owner = workflow.user_id == current_user.user_id
|
||||
if not (is_default or is_owner or workflow.is_public or current_user.is_admin):
|
||||
raise HTTPException(status_code=403, detail="Not authorized to access this workflow")
|
||||
|
||||
thumbnail_url = ApiDependencies.invoker.services.workflow_thumbnails.get_url(workflow_id)
|
||||
compatibility = get_workflow_call_compatibility(
|
||||
workflow=workflow.workflow.model_dump(),
|
||||
workflow_id=workflow.workflow_id,
|
||||
services=ApiDependencies.invoker.services,
|
||||
user_id=current_user.user_id,
|
||||
maximum_children=ApiDependencies.invoker.services.configuration.max_queue_size,
|
||||
)
|
||||
return WorkflowRecordWithThumbnailDTO(
|
||||
thumbnail_url=thumbnail_url,
|
||||
call_saved_workflow_compatibility=compatibility,
|
||||
**workflow.model_dump(),
|
||||
)
|
||||
|
||||
|
||||
@workflows_router.patch(
|
||||
"/i/{workflow_id}",
|
||||
operation_id="update_workflow",
|
||||
responses={
|
||||
200: {"model": WorkflowRecordDTO},
|
||||
},
|
||||
)
|
||||
async def update_workflow(
|
||||
current_user: CurrentUserOrDefault,
|
||||
workflow: Workflow = Body(description="The updated workflow", embed=True),
|
||||
) -> WorkflowRecordDTO:
|
||||
"""Updates a workflow"""
|
||||
try:
|
||||
existing = ApiDependencies.invoker.services.workflow_records.get(workflow.id)
|
||||
except WorkflowNotFoundError:
|
||||
raise HTTPException(status_code=404, detail="Workflow not found")
|
||||
|
||||
config = ApiDependencies.invoker.services.configuration
|
||||
if config.multiuser:
|
||||
if not current_user.is_admin and existing.user_id != current_user.user_id:
|
||||
raise HTTPException(status_code=403, detail="Not authorized to update this workflow")
|
||||
user_id = None if current_user.is_admin else current_user.user_id
|
||||
updated = ApiDependencies.invoker.services.workflow_records.update(workflow=workflow, user_id=user_id)
|
||||
ApiDependencies.invoker.services.events.emit_workflow_updated(
|
||||
workflow_id=updated.workflow_id,
|
||||
user_id=updated.user_id,
|
||||
old_is_public=existing.is_public,
|
||||
new_is_public=updated.is_public,
|
||||
)
|
||||
return updated
|
||||
|
||||
|
||||
@workflows_router.delete(
|
||||
"/i/{workflow_id}",
|
||||
operation_id="delete_workflow",
|
||||
)
|
||||
async def delete_workflow(
|
||||
current_user: CurrentUserOrDefault,
|
||||
workflow_id: str = Path(description="The workflow to delete"),
|
||||
) -> None:
|
||||
"""Deletes a workflow"""
|
||||
try:
|
||||
existing = ApiDependencies.invoker.services.workflow_records.get(workflow_id)
|
||||
except WorkflowNotFoundError:
|
||||
raise HTTPException(status_code=404, detail="Workflow not found")
|
||||
|
||||
config = ApiDependencies.invoker.services.configuration
|
||||
if config.multiuser:
|
||||
if not current_user.is_admin and existing.user_id != current_user.user_id:
|
||||
raise HTTPException(status_code=403, detail="Not authorized to delete this workflow")
|
||||
try:
|
||||
ApiDependencies.invoker.services.workflow_thumbnails.delete(workflow_id)
|
||||
except WorkflowThumbnailFileNotFoundException:
|
||||
# It's OK if the workflow has no thumbnail file. We can still delete the workflow.
|
||||
pass
|
||||
user_id = None if current_user.is_admin else current_user.user_id
|
||||
ApiDependencies.invoker.services.workflow_records.delete(workflow_id, user_id=user_id)
|
||||
ApiDependencies.invoker.services.events.emit_workflow_deleted(
|
||||
workflow_id=existing.workflow_id,
|
||||
user_id=existing.user_id,
|
||||
is_public=existing.is_public,
|
||||
)
|
||||
|
||||
|
||||
@workflows_router.post(
|
||||
"/",
|
||||
operation_id="create_workflow",
|
||||
responses={
|
||||
200: {"model": WorkflowRecordDTO},
|
||||
},
|
||||
)
|
||||
async def create_workflow(
|
||||
current_user: CurrentUserOrDefault,
|
||||
workflow: WorkflowWithoutID = Body(description="The workflow to create", embed=True),
|
||||
) -> WorkflowRecordDTO:
|
||||
"""Creates a workflow"""
|
||||
# In single-user mode, workflows are owned by 'system' and shared by default so all legacy/single-user
|
||||
# workflows remain visible. In multiuser mode, workflows are private to the creator by default.
|
||||
config = ApiDependencies.invoker.services.configuration
|
||||
is_public = not config.multiuser
|
||||
created = ApiDependencies.invoker.services.workflow_records.create(
|
||||
workflow=workflow, user_id=current_user.user_id, is_public=is_public
|
||||
)
|
||||
ApiDependencies.invoker.services.events.emit_workflow_created(
|
||||
workflow_id=created.workflow_id,
|
||||
user_id=created.user_id,
|
||||
is_public=created.is_public,
|
||||
)
|
||||
return created
|
||||
|
||||
|
||||
@workflows_router.get(
|
||||
"/",
|
||||
operation_id="list_workflows",
|
||||
responses={
|
||||
200: {"model": PaginatedResults[WorkflowRecordListItemWithThumbnailDTO]},
|
||||
},
|
||||
)
|
||||
async def list_workflows(
|
||||
current_user: CurrentUserOrDefault,
|
||||
page: int = Query(default=0, description="The page to get"),
|
||||
per_page: Optional[int] = Query(default=None, description="The number of workflows per page"),
|
||||
order_by: WorkflowRecordOrderBy = Query(
|
||||
default=WorkflowRecordOrderBy.Name, description="The attribute to order by"
|
||||
),
|
||||
direction: SQLiteDirection = Query(default=SQLiteDirection.Ascending, description="The direction to order by"),
|
||||
categories: Optional[list[WorkflowCategory]] = Query(default=None, description="The categories of workflow to get"),
|
||||
tags: Optional[list[str]] = Query(default=None, description="The tags of workflow to get"),
|
||||
query: Optional[str] = Query(default=None, description="The text to query by (matches name and description)"),
|
||||
has_been_opened: Optional[bool] = Query(default=None, description="Whether to include/exclude recent workflows"),
|
||||
is_public: Optional[bool] = Query(default=None, description="Filter by public/shared status"),
|
||||
is_callable: Optional[bool] = Query(
|
||||
default=None,
|
||||
alias="callable",
|
||||
description="Filter by whether workflows are callable by call_saved_workflow",
|
||||
),
|
||||
) -> PaginatedResults[WorkflowRecordListItemWithThumbnailDTO]:
|
||||
"""Gets a page of workflows"""
|
||||
config = ApiDependencies.invoker.services.configuration
|
||||
|
||||
# In multiuser mode, scope user-category workflows to the current user unless fetching shared workflows.
|
||||
# Admins skip the user_id filter so they can see and manage all workflows including system-owned ones.
|
||||
user_id_filter: Optional[str] = None
|
||||
if config.multiuser and not current_user.is_admin:
|
||||
has_user_category = not categories or WorkflowCategory.User in categories
|
||||
if has_user_category and is_public is not True:
|
||||
user_id_filter = current_user.user_id
|
||||
|
||||
workflows_with_thumbnails: list[WorkflowRecordListItemWithThumbnailDTO] = []
|
||||
workflows = ApiDependencies.invoker.services.workflow_records.get_many(
|
||||
order_by=order_by,
|
||||
direction=direction,
|
||||
page=page,
|
||||
per_page=None if is_callable is not None else per_page,
|
||||
query=query,
|
||||
categories=categories,
|
||||
tags=tags,
|
||||
has_been_opened=has_been_opened,
|
||||
user_id=user_id_filter,
|
||||
is_public=is_public,
|
||||
)
|
||||
skipped_missing_workflows = 0
|
||||
for workflow in workflows.items:
|
||||
try:
|
||||
full_workflow = ApiDependencies.invoker.services.workflow_records.get(workflow.workflow_id)
|
||||
except WorkflowNotFoundError:
|
||||
skipped_missing_workflows += 1
|
||||
continue
|
||||
compatibility = get_workflow_call_compatibility(
|
||||
workflow=full_workflow.workflow.model_dump(),
|
||||
workflow_id=full_workflow.workflow_id,
|
||||
services=ApiDependencies.invoker.services,
|
||||
user_id=current_user.user_id,
|
||||
maximum_children=ApiDependencies.invoker.services.configuration.max_queue_size,
|
||||
resolve_generator_items=False,
|
||||
)
|
||||
if is_callable is not None and compatibility.is_callable != is_callable:
|
||||
continue
|
||||
workflows_with_thumbnails.append(
|
||||
WorkflowRecordListItemWithThumbnailDTO(
|
||||
thumbnail_url=ApiDependencies.invoker.services.workflow_thumbnails.get_url(workflow.workflow_id),
|
||||
call_saved_workflow_compatibility=compatibility,
|
||||
**workflow.model_dump(),
|
||||
)
|
||||
)
|
||||
|
||||
if is_callable is not None:
|
||||
total = len(workflows_with_thumbnails)
|
||||
if per_page:
|
||||
start = page * per_page
|
||||
end = start + per_page
|
||||
page_items = workflows_with_thumbnails[start:end]
|
||||
pages = math.ceil(total / per_page)
|
||||
else:
|
||||
page_items = workflows_with_thumbnails
|
||||
pages = 1
|
||||
return PaginatedResults[WorkflowRecordListItemWithThumbnailDTO](
|
||||
items=page_items,
|
||||
total=total,
|
||||
page=page,
|
||||
pages=pages,
|
||||
per_page=per_page if per_page else total,
|
||||
)
|
||||
|
||||
return PaginatedResults[WorkflowRecordListItemWithThumbnailDTO](
|
||||
items=workflows_with_thumbnails,
|
||||
total=max(len(workflows_with_thumbnails), workflows.total - skipped_missing_workflows),
|
||||
page=workflows.page,
|
||||
pages=workflows.pages,
|
||||
per_page=workflows.per_page,
|
||||
)
|
||||
|
||||
|
||||
@workflows_router.put(
|
||||
"/i/{workflow_id}/thumbnail",
|
||||
operation_id="set_workflow_thumbnail",
|
||||
responses={
|
||||
200: {"model": WorkflowRecordDTO},
|
||||
},
|
||||
)
|
||||
async def set_workflow_thumbnail(
|
||||
current_user: CurrentUserOrDefault,
|
||||
workflow_id: str = Path(description="The workflow to update"),
|
||||
image: UploadFile = File(description="The image file to upload"),
|
||||
):
|
||||
"""Sets a workflow's thumbnail image"""
|
||||
try:
|
||||
existing = ApiDependencies.invoker.services.workflow_records.get(workflow_id)
|
||||
except WorkflowNotFoundError:
|
||||
raise HTTPException(status_code=404, detail="Workflow not found")
|
||||
|
||||
config = ApiDependencies.invoker.services.configuration
|
||||
if config.multiuser and not current_user.is_admin and existing.user_id != current_user.user_id:
|
||||
raise HTTPException(status_code=403, detail="Not authorized to update this workflow")
|
||||
|
||||
if not image.content_type or not image.content_type.startswith("image"):
|
||||
raise HTTPException(status_code=415, detail="Not an image")
|
||||
|
||||
contents = await image.read()
|
||||
try:
|
||||
pil_image = Image.open(io.BytesIO(contents))
|
||||
|
||||
except Exception:
|
||||
ApiDependencies.invoker.services.logger.error(traceback.format_exc())
|
||||
raise HTTPException(status_code=415, detail="Failed to read image")
|
||||
|
||||
try:
|
||||
ApiDependencies.invoker.services.workflow_thumbnails.save(workflow_id, pil_image)
|
||||
except Exception as e:
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
|
||||
@workflows_router.delete(
|
||||
"/i/{workflow_id}/thumbnail",
|
||||
operation_id="delete_workflow_thumbnail",
|
||||
responses={
|
||||
200: {"model": WorkflowRecordDTO},
|
||||
},
|
||||
)
|
||||
async def delete_workflow_thumbnail(
|
||||
current_user: CurrentUserOrDefault,
|
||||
workflow_id: str = Path(description="The workflow to update"),
|
||||
):
|
||||
"""Removes a workflow's thumbnail image"""
|
||||
try:
|
||||
existing = ApiDependencies.invoker.services.workflow_records.get(workflow_id)
|
||||
except WorkflowNotFoundError:
|
||||
raise HTTPException(status_code=404, detail="Workflow not found")
|
||||
|
||||
config = ApiDependencies.invoker.services.configuration
|
||||
if config.multiuser and not current_user.is_admin and existing.user_id != current_user.user_id:
|
||||
raise HTTPException(status_code=403, detail="Not authorized to update this workflow")
|
||||
|
||||
try:
|
||||
ApiDependencies.invoker.services.workflow_thumbnails.delete(workflow_id)
|
||||
except ValueError as e:
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
|
||||
@workflows_router.get(
|
||||
"/i/{workflow_id}/thumbnail",
|
||||
operation_id="get_workflow_thumbnail",
|
||||
responses={
|
||||
200: {
|
||||
"description": "The workflow thumbnail was fetched successfully",
|
||||
},
|
||||
400: {"description": "Bad request"},
|
||||
404: {"description": "The workflow thumbnail could not be found"},
|
||||
},
|
||||
status_code=200,
|
||||
)
|
||||
async def get_workflow_thumbnail(
|
||||
workflow_id: str = Path(description="The id of the workflow thumbnail to get"),
|
||||
) -> FileResponse:
|
||||
"""Gets a workflow's thumbnail image.
|
||||
|
||||
This endpoint is intentionally unauthenticated because browsers load images
|
||||
via <img src> tags which cannot send Bearer tokens. Workflow IDs are UUIDs,
|
||||
providing security through unguessability.
|
||||
"""
|
||||
try:
|
||||
path = ApiDependencies.invoker.services.workflow_thumbnails.get_path(workflow_id)
|
||||
|
||||
response = FileResponse(
|
||||
path,
|
||||
media_type="image/png",
|
||||
filename=workflow_id + ".png",
|
||||
content_disposition_type="inline",
|
||||
)
|
||||
response.headers["Cache-Control"] = f"max-age={IMAGE_MAX_AGE}"
|
||||
return response
|
||||
except Exception:
|
||||
raise HTTPException(status_code=404)
|
||||
|
||||
|
||||
@workflows_router.patch(
|
||||
"/i/{workflow_id}/is_public",
|
||||
operation_id="update_workflow_is_public",
|
||||
responses={
|
||||
200: {"model": WorkflowRecordDTO},
|
||||
},
|
||||
)
|
||||
async def update_workflow_is_public(
|
||||
current_user: CurrentUserOrDefault,
|
||||
workflow_id: str = Path(description="The workflow to update"),
|
||||
is_public: bool = Body(description="Whether the workflow should be shared publicly", embed=True),
|
||||
) -> WorkflowRecordDTO:
|
||||
"""Updates whether a workflow is shared publicly"""
|
||||
try:
|
||||
existing = ApiDependencies.invoker.services.workflow_records.get(workflow_id)
|
||||
except WorkflowNotFoundError:
|
||||
raise HTTPException(status_code=404, detail="Workflow not found")
|
||||
|
||||
config = ApiDependencies.invoker.services.configuration
|
||||
if config.multiuser and not current_user.is_admin and existing.user_id != current_user.user_id:
|
||||
raise HTTPException(status_code=403, detail="Not authorized to update this workflow")
|
||||
|
||||
user_id = None if current_user.is_admin else current_user.user_id
|
||||
updated = ApiDependencies.invoker.services.workflow_records.update_is_public(
|
||||
workflow_id=workflow_id, is_public=is_public, user_id=user_id
|
||||
)
|
||||
ApiDependencies.invoker.services.events.emit_workflow_updated(
|
||||
workflow_id=updated.workflow_id,
|
||||
user_id=updated.user_id,
|
||||
old_is_public=existing.is_public,
|
||||
new_is_public=updated.is_public,
|
||||
)
|
||||
return updated
|
||||
|
||||
|
||||
@workflows_router.get("/tags", operation_id="get_all_tags")
|
||||
async def get_all_tags(
|
||||
current_user: CurrentUserOrDefault,
|
||||
categories: Optional[list[WorkflowCategory]] = Query(default=None, description="The categories to include"),
|
||||
is_public: Optional[bool] = Query(default=None, description="Filter by public/shared status"),
|
||||
) -> list[str]:
|
||||
"""Gets all unique tags from workflows"""
|
||||
config = ApiDependencies.invoker.services.configuration
|
||||
user_id_filter: Optional[str] = None
|
||||
if config.multiuser and not current_user.is_admin:
|
||||
has_user_category = not categories or WorkflowCategory.User in categories
|
||||
if has_user_category and is_public is not True:
|
||||
user_id_filter = current_user.user_id
|
||||
|
||||
return ApiDependencies.invoker.services.workflow_records.get_all_tags(
|
||||
categories=categories, user_id=user_id_filter, is_public=is_public
|
||||
)
|
||||
|
||||
|
||||
@workflows_router.get("/counts_by_tag", operation_id="get_counts_by_tag")
|
||||
async def get_counts_by_tag(
|
||||
current_user: CurrentUserOrDefault,
|
||||
tags: list[str] = Query(description="The tags to get counts for"),
|
||||
categories: Optional[list[WorkflowCategory]] = Query(default=None, description="The categories to include"),
|
||||
has_been_opened: Optional[bool] = Query(default=None, description="Whether to include/exclude recent workflows"),
|
||||
is_public: Optional[bool] = Query(default=None, description="Filter by public/shared status"),
|
||||
) -> dict[str, int]:
|
||||
"""Counts workflows by tag"""
|
||||
config = ApiDependencies.invoker.services.configuration
|
||||
user_id_filter: Optional[str] = None
|
||||
if config.multiuser and not current_user.is_admin:
|
||||
has_user_category = not categories or WorkflowCategory.User in categories
|
||||
if has_user_category and is_public is not True:
|
||||
user_id_filter = current_user.user_id
|
||||
|
||||
return ApiDependencies.invoker.services.workflow_records.counts_by_tag(
|
||||
tags=tags, categories=categories, has_been_opened=has_been_opened, user_id=user_id_filter, is_public=is_public
|
||||
)
|
||||
|
||||
|
||||
@workflows_router.get("/counts_by_category", operation_id="counts_by_category")
|
||||
async def counts_by_category(
|
||||
current_user: CurrentUserOrDefault,
|
||||
categories: list[WorkflowCategory] = Query(description="The categories to include"),
|
||||
has_been_opened: Optional[bool] = Query(default=None, description="Whether to include/exclude recent workflows"),
|
||||
is_public: Optional[bool] = Query(default=None, description="Filter by public/shared status"),
|
||||
) -> dict[str, int]:
|
||||
"""Counts workflows by category"""
|
||||
config = ApiDependencies.invoker.services.configuration
|
||||
user_id_filter: Optional[str] = None
|
||||
if config.multiuser and not current_user.is_admin:
|
||||
has_user_category = WorkflowCategory.User in categories
|
||||
if has_user_category and is_public is not True:
|
||||
user_id_filter = current_user.user_id
|
||||
|
||||
return ApiDependencies.invoker.services.workflow_records.counts_by_category(
|
||||
categories=categories, has_been_opened=has_been_opened, user_id=user_id_filter, is_public=is_public
|
||||
)
|
||||
|
||||
|
||||
@workflows_router.put(
|
||||
"/i/{workflow_id}/opened_at",
|
||||
operation_id="update_opened_at",
|
||||
)
|
||||
async def update_opened_at(
|
||||
current_user: CurrentUserOrDefault,
|
||||
workflow_id: str = Path(description="The workflow to update"),
|
||||
) -> None:
|
||||
"""Updates the opened_at field of a workflow"""
|
||||
try:
|
||||
existing = ApiDependencies.invoker.services.workflow_records.get(workflow_id)
|
||||
except WorkflowNotFoundError:
|
||||
raise HTTPException(status_code=404, detail="Workflow not found")
|
||||
|
||||
config = ApiDependencies.invoker.services.configuration
|
||||
if config.multiuser and not current_user.is_admin and existing.user_id != current_user.user_id:
|
||||
raise HTTPException(status_code=403, detail="Not authorized to update this workflow")
|
||||
|
||||
user_id = None if current_user.is_admin else current_user.user_id
|
||||
ApiDependencies.invoker.services.workflow_records.update_opened_at(workflow_id, user_id=user_id)
|
||||
@@ -0,0 +1,543 @@
|
||||
# Copyright (c) 2022 Kyle Schouviller (https://github.com/kyle0654)
|
||||
|
||||
from typing import Any
|
||||
|
||||
from fastapi import FastAPI
|
||||
from pydantic import BaseModel
|
||||
from socketio import ASGIApp, AsyncServer
|
||||
|
||||
from invokeai.app.services.auth.token_service import verify_token
|
||||
from invokeai.app.services.events.events_common import (
|
||||
BatchEnqueuedEvent,
|
||||
BulkDownloadCompleteEvent,
|
||||
BulkDownloadErrorEvent,
|
||||
BulkDownloadEventBase,
|
||||
BulkDownloadStartedEvent,
|
||||
DownloadCancelledEvent,
|
||||
DownloadCompleteEvent,
|
||||
DownloadErrorEvent,
|
||||
DownloadEventBase,
|
||||
DownloadProgressEvent,
|
||||
DownloadStartedEvent,
|
||||
FastAPIEvent,
|
||||
InvocationCompleteEvent,
|
||||
InvocationErrorEvent,
|
||||
InvocationProgressEvent,
|
||||
InvocationStartedEvent,
|
||||
ModelEventBase,
|
||||
ModelInstallCancelledEvent,
|
||||
ModelInstallCompleteEvent,
|
||||
ModelInstallDownloadProgressEvent,
|
||||
ModelInstallDownloadsCompleteEvent,
|
||||
ModelInstallErrorEvent,
|
||||
ModelInstallStartedEvent,
|
||||
ModelLoadCompleteEvent,
|
||||
ModelLoadStartedEvent,
|
||||
QueueClearedEvent,
|
||||
QueueEventBase,
|
||||
QueueItemsRetriedEvent,
|
||||
QueueItemStatusChangedEvent,
|
||||
RecallParametersUpdatedEvent,
|
||||
WorkflowAccessRevokedEvent,
|
||||
WorkflowCreatedEvent,
|
||||
WorkflowDeletedEvent,
|
||||
WorkflowEventBase,
|
||||
WorkflowUpdatedEvent,
|
||||
register_events,
|
||||
)
|
||||
from invokeai.backend.util.logging import InvokeAILogger
|
||||
|
||||
logger = InvokeAILogger.get_logger()
|
||||
|
||||
|
||||
class QueueSubscriptionEvent(BaseModel):
|
||||
"""Event data for subscribing to the socket.io queue room.
|
||||
This is a pydantic model to ensure the data is in the correct format."""
|
||||
|
||||
queue_id: str
|
||||
|
||||
|
||||
class BulkDownloadSubscriptionEvent(BaseModel):
|
||||
"""Event data for subscribing to the socket.io bulk downloads room.
|
||||
This is a pydantic model to ensure the data is in the correct format."""
|
||||
|
||||
bulk_download_id: str
|
||||
|
||||
|
||||
QUEUE_EVENTS = {
|
||||
InvocationStartedEvent,
|
||||
InvocationProgressEvent,
|
||||
InvocationCompleteEvent,
|
||||
InvocationErrorEvent,
|
||||
QueueItemStatusChangedEvent,
|
||||
BatchEnqueuedEvent,
|
||||
QueueItemsRetriedEvent,
|
||||
QueueClearedEvent,
|
||||
RecallParametersUpdatedEvent,
|
||||
}
|
||||
|
||||
MODEL_EVENTS = {
|
||||
DownloadCancelledEvent,
|
||||
DownloadCompleteEvent,
|
||||
DownloadErrorEvent,
|
||||
DownloadProgressEvent,
|
||||
DownloadStartedEvent,
|
||||
ModelLoadStartedEvent,
|
||||
ModelLoadCompleteEvent,
|
||||
ModelInstallDownloadProgressEvent,
|
||||
ModelInstallDownloadsCompleteEvent,
|
||||
ModelInstallStartedEvent,
|
||||
ModelInstallCompleteEvent,
|
||||
ModelInstallCancelledEvent,
|
||||
ModelInstallErrorEvent,
|
||||
}
|
||||
|
||||
BULK_DOWNLOAD_EVENTS = {BulkDownloadStartedEvent, BulkDownloadCompleteEvent, BulkDownloadErrorEvent}
|
||||
WORKFLOW_EVENTS = {WorkflowCreatedEvent, WorkflowUpdatedEvent, WorkflowDeletedEvent}
|
||||
|
||||
|
||||
class SocketIO:
|
||||
_sub_queue = "subscribe_queue"
|
||||
_unsub_queue = "unsubscribe_queue"
|
||||
|
||||
_sub_bulk_download = "subscribe_bulk_download"
|
||||
_unsub_bulk_download = "unsubscribe_bulk_download"
|
||||
|
||||
def __init__(self, app: FastAPI):
|
||||
self._sio = AsyncServer(async_mode="asgi", cors_allowed_origins="*")
|
||||
self._app = ASGIApp(socketio_server=self._sio, socketio_path="/ws/socket.io")
|
||||
app.mount("/ws", self._app)
|
||||
|
||||
# Track user information for each socket connection
|
||||
self._socket_users: dict[str, dict[str, Any]] = {}
|
||||
|
||||
# Set up authentication middleware
|
||||
self._sio.on("connect", handler=self._handle_connect)
|
||||
self._sio.on("disconnect", handler=self._handle_disconnect)
|
||||
|
||||
self._sio.on(self._sub_queue, handler=self._handle_sub_queue)
|
||||
self._sio.on(self._unsub_queue, handler=self._handle_unsub_queue)
|
||||
self._sio.on(self._sub_bulk_download, handler=self._handle_sub_bulk_download)
|
||||
self._sio.on(self._unsub_bulk_download, handler=self._handle_unsub_bulk_download)
|
||||
|
||||
register_events(QUEUE_EVENTS, self._handle_queue_event)
|
||||
register_events(MODEL_EVENTS, self._handle_model_event)
|
||||
register_events(BULK_DOWNLOAD_EVENTS, self._handle_bulk_image_download_event)
|
||||
register_events(WORKFLOW_EVENTS, self._handle_workflow_event)
|
||||
|
||||
async def _handle_connect(self, sid: str, environ: dict, auth: dict | None) -> bool:
|
||||
"""Handle socket connection and authenticate the user.
|
||||
|
||||
Returns True to accept the connection, False to reject it.
|
||||
Stores user_id in the internal socket users dict for later use.
|
||||
|
||||
In multiuser mode, connections without a valid token are rejected outright
|
||||
so that anonymous clients cannot subscribe to queue rooms and observe
|
||||
queue activity belonging to other users. In single-user mode, unauthenticated
|
||||
connections are accepted as the system admin user.
|
||||
"""
|
||||
# Extract token from auth data or headers
|
||||
token = None
|
||||
if auth and isinstance(auth, dict):
|
||||
token = auth.get("token")
|
||||
|
||||
if not token and environ:
|
||||
# Try to get token from headers
|
||||
headers = environ.get("HTTP_AUTHORIZATION", "")
|
||||
if headers.startswith("Bearer "):
|
||||
token = headers[7:]
|
||||
|
||||
# Verify the token
|
||||
if token:
|
||||
token_data = verify_token(token)
|
||||
if token_data:
|
||||
# In multiuser mode, also verify the backing user record still
|
||||
# exists and is active — mirrors the REST auth check in
|
||||
# auth_dependencies.py. A deleted or deactivated user whose
|
||||
# JWT has not yet expired must not be allowed to open a socket.
|
||||
if self._is_multiuser_enabled():
|
||||
try:
|
||||
from invokeai.app.api.dependencies import ApiDependencies
|
||||
|
||||
user = ApiDependencies.invoker.services.users.get(token_data.user_id)
|
||||
if user is None or not user.is_active:
|
||||
logger.warning(f"Rejecting socket {sid}: user {token_data.user_id} not found or inactive")
|
||||
return False
|
||||
except Exception:
|
||||
# If user service is unavailable, fail closed
|
||||
logger.warning(f"Rejecting socket {sid}: unable to verify user record")
|
||||
return False
|
||||
|
||||
# Store user_id and is_admin in socket users dict
|
||||
self._socket_users[sid] = {
|
||||
"user_id": token_data.user_id,
|
||||
"is_admin": token_data.is_admin,
|
||||
}
|
||||
logger.info(
|
||||
f"Socket {sid} connected with user_id: {token_data.user_id}, is_admin: {token_data.is_admin}"
|
||||
)
|
||||
await self._sio.enter_room(sid, f"user:{token_data.user_id}")
|
||||
await self._sio.enter_room(sid, "workflows:shared")
|
||||
if token_data.is_admin:
|
||||
await self._sio.enter_room(sid, "admin")
|
||||
return True
|
||||
|
||||
# No valid token provided. In multiuser mode this is not allowed — reject
|
||||
# the connection so anonymous clients cannot subscribe to queue rooms.
|
||||
# In single-user mode, fall through and accept the socket as system admin.
|
||||
if self._is_multiuser_enabled():
|
||||
logger.warning(
|
||||
f"Rejecting socket {sid} connection: multiuser mode is enabled and no valid auth token was provided"
|
||||
)
|
||||
return False
|
||||
|
||||
self._socket_users[sid] = {
|
||||
"user_id": "system",
|
||||
"is_admin": True,
|
||||
}
|
||||
logger.debug(f"Socket {sid} connected as system admin (single-user mode)")
|
||||
await self._sio.enter_room(sid, "user:system")
|
||||
await self._sio.enter_room(sid, "workflows:shared")
|
||||
await self._sio.enter_room(sid, "admin")
|
||||
return True
|
||||
|
||||
@staticmethod
|
||||
def _is_multiuser_enabled() -> bool:
|
||||
"""Check whether multiuser mode is enabled. Fails closed if configuration
|
||||
is not yet initialized, which should not happen in practice but prevents
|
||||
accidentally opening the socket during startup races."""
|
||||
try:
|
||||
# Imported here to avoid a circular import at module load time.
|
||||
from invokeai.app.api.dependencies import ApiDependencies
|
||||
|
||||
return bool(ApiDependencies.invoker.services.configuration.multiuser)
|
||||
except Exception:
|
||||
# If dependencies are not initialized, fail closed (treat as multiuser)
|
||||
# so we never accidentally admit an anonymous socket.
|
||||
return True
|
||||
|
||||
async def _handle_disconnect(self, sid: str) -> None:
|
||||
"""Handle socket disconnection and cleanup user info."""
|
||||
if sid in self._socket_users:
|
||||
del self._socket_users[sid]
|
||||
logger.debug(f"Socket {sid} disconnected and cleaned up")
|
||||
|
||||
async def _handle_sub_queue(self, sid: str, data: Any) -> None:
|
||||
"""Handle queue subscription and add socket to both queue and user-specific rooms."""
|
||||
queue_id = QueueSubscriptionEvent(**data).queue_id
|
||||
|
||||
# Check if we have user info for this socket. In multiuser mode _handle_connect
|
||||
# will have already rejected any socket without a valid token, so missing user
|
||||
# info here is a bug — refuse the subscription rather than silently falling back
|
||||
# to an anonymous system user who could then receive queue item events.
|
||||
if sid not in self._socket_users:
|
||||
if self._is_multiuser_enabled():
|
||||
logger.warning(
|
||||
f"Refusing queue subscription for socket {sid}: no user info (socket not authenticated via connect event)"
|
||||
)
|
||||
return
|
||||
# Single-user mode: safe to fall back to the system admin user.
|
||||
self._socket_users[sid] = {
|
||||
"user_id": "system",
|
||||
"is_admin": True,
|
||||
}
|
||||
|
||||
user_id = self._socket_users[sid]["user_id"]
|
||||
is_admin = self._socket_users[sid]["is_admin"]
|
||||
|
||||
# Add socket to the queue room
|
||||
await self._sio.enter_room(sid, queue_id)
|
||||
|
||||
# Also add socket to a user-specific room for event filtering
|
||||
user_room = f"user:{user_id}"
|
||||
await self._sio.enter_room(sid, user_room)
|
||||
|
||||
# If admin, also add to admin room to receive all events
|
||||
if is_admin:
|
||||
await self._sio.enter_room(sid, "admin")
|
||||
|
||||
logger.debug(
|
||||
f"Socket {sid} (user_id: {user_id}, is_admin: {is_admin}) subscribed to queue {queue_id} and user room {user_room}"
|
||||
)
|
||||
|
||||
async def _handle_unsub_queue(self, sid: str, data: Any) -> None:
|
||||
await self._sio.leave_room(sid, QueueSubscriptionEvent(**data).queue_id)
|
||||
|
||||
async def _handle_sub_bulk_download(self, sid: str, data: Any) -> None:
|
||||
# In multiuser mode, only allow authenticated sockets to subscribe.
|
||||
# Bulk download events are routed to user-specific rooms, so the
|
||||
# bulk_download_id room subscription is only kept for single-user
|
||||
# backward compatibility.
|
||||
if self._is_multiuser_enabled() and sid not in self._socket_users:
|
||||
logger.warning(f"Refusing bulk download subscription for unknown socket {sid} in multiuser mode")
|
||||
return
|
||||
await self._sio.enter_room(sid, BulkDownloadSubscriptionEvent(**data).bulk_download_id)
|
||||
|
||||
async def _handle_unsub_bulk_download(self, sid: str, data: Any) -> None:
|
||||
await self._sio.leave_room(sid, BulkDownloadSubscriptionEvent(**data).bulk_download_id)
|
||||
|
||||
def _owner_and_admin_sids(self, owner_user_id: str) -> list[str]:
|
||||
"""Sids belonging to the event's owner or to any admin.
|
||||
|
||||
Used as `skip_sid` when broadcasting a sanitized companion event to the queue room,
|
||||
so the owner and admins (who already received the full event) don't get a second
|
||||
copy that would clobber their cache with redacted values.
|
||||
"""
|
||||
return [
|
||||
sid
|
||||
for sid, info in self._socket_users.items()
|
||||
if info.get("user_id") == owner_user_id or info.get("is_admin")
|
||||
]
|
||||
|
||||
async def _handle_queue_event(self, event: FastAPIEvent[QueueEventBase]):
|
||||
"""Handle queue events with user isolation.
|
||||
|
||||
Queue events split into two routing paths:
|
||||
|
||||
1. The owner and admins receive the full unsanitized event in their `user:{id}` /
|
||||
`admin` rooms. The full payload may include batch_id, session_id, origin,
|
||||
destination, error metadata, etc.
|
||||
|
||||
2. For events that other authenticated users need to know about so their queue list
|
||||
and badge counts stay in sync (QueueItemStatusChangedEvent and BatchEnqueuedEvent),
|
||||
a sanitized companion event is also emitted to the full queue room with the
|
||||
owner's and admins' sids in `skip_sid`. The companion uses `user_id="redacted"`
|
||||
as a sentinel so the frontend handler knows to do tag invalidation only and skip
|
||||
per-session side effects.
|
||||
|
||||
InvocationEventBase events stay private (owner + admins only). RecallParametersUpdatedEvent
|
||||
is also private. QueueClearedEvent is broadcast to the queue room when unscoped (an admin or
|
||||
single-user clear that deleted every user's items); a user-scoped clear goes full to
|
||||
owner + admins with a sanitized companion to the rest of the queue room, so other users
|
||||
refresh their queue lists without treating the clear as their own.
|
||||
|
||||
IMPORTANT: Check InvocationEventBase BEFORE QueueItemEventBase since InvocationEventBase
|
||||
inherits from QueueItemEventBase. The order of isinstance checks matters!
|
||||
"""
|
||||
try:
|
||||
event_name, event_data = event
|
||||
|
||||
# Import here to avoid circular dependency
|
||||
from invokeai.app.services.events.events_common import InvocationEventBase, QueueItemEventBase
|
||||
|
||||
# Check InvocationEventBase FIRST (before QueueItemEventBase) since it's a subclass
|
||||
# Invocation events (progress, started, complete, error) are private to owner + admins
|
||||
if isinstance(event_data, InvocationEventBase) and hasattr(event_data, "user_id"):
|
||||
user_room = f"user:{event_data.user_id}"
|
||||
|
||||
# Emit to the user's room
|
||||
await self._sio.emit(event=event_name, data=event_data.model_dump(mode="json"), room=user_room)
|
||||
|
||||
# Also emit to admin room so admins can see all events, but strip image preview data
|
||||
# from InvocationProgressEvent to prevent admins from seeing other users' image content
|
||||
if isinstance(event_data, InvocationProgressEvent):
|
||||
admin_event_data = event_data.model_copy(update={"image": None})
|
||||
await self._sio.emit(event=event_name, data=admin_event_data.model_dump(mode="json"), room="admin")
|
||||
else:
|
||||
await self._sio.emit(event=event_name, data=event_data.model_dump(mode="json"), room="admin")
|
||||
|
||||
logger.debug(f"Emitted private invocation event {event_name} to user room {user_room} and admin room")
|
||||
|
||||
# QueueItemStatusChangedEvent: full to owner+admin, sanitized to everyone else in
|
||||
# the queue room so their queue list, badge, and item caches refresh.
|
||||
elif isinstance(event_data, QueueItemStatusChangedEvent):
|
||||
user_room = f"user:{event_data.user_id}"
|
||||
await self._sio.emit(event=event_name, data=event_data.model_dump(mode="json"), room=user_room)
|
||||
await self._sio.emit(event=event_name, data=event_data.model_dump(mode="json"), room="admin")
|
||||
|
||||
sanitized = event_data.model_copy(
|
||||
update={
|
||||
"user_id": "redacted",
|
||||
"batch_id": "redacted",
|
||||
"session_id": "redacted",
|
||||
"origin": None,
|
||||
"destination": None,
|
||||
"error_type": None,
|
||||
"error_message": None,
|
||||
"error_traceback": None,
|
||||
}
|
||||
)
|
||||
# Strip identifying fields out of the embedded batch_status / queue_status too.
|
||||
sanitized.batch_status = sanitized.batch_status.model_copy(
|
||||
update={"batch_id": "redacted", "origin": None, "destination": None}
|
||||
)
|
||||
sanitized.queue_status = sanitized.queue_status.model_copy(
|
||||
update={
|
||||
"item_id": None,
|
||||
"session_id": None,
|
||||
"batch_id": None,
|
||||
"user_pending": None,
|
||||
"user_in_progress": None,
|
||||
}
|
||||
)
|
||||
await self._sio.emit(
|
||||
event=event_name,
|
||||
data=sanitized.model_dump(mode="json"),
|
||||
room=event_data.queue_id,
|
||||
skip_sid=self._owner_and_admin_sids(event_data.user_id),
|
||||
)
|
||||
|
||||
logger.debug(
|
||||
f"Emitted queue_item_status_changed: full to {user_room}+admin, sanitized to queue {event_data.queue_id}"
|
||||
)
|
||||
|
||||
# Other queue item events (currently none beyond QueueItemStatusChangedEvent that
|
||||
# carry user_id) stay private to owner + admins.
|
||||
elif isinstance(event_data, QueueItemEventBase) and hasattr(event_data, "user_id"):
|
||||
user_room = f"user:{event_data.user_id}"
|
||||
await self._sio.emit(event=event_name, data=event_data.model_dump(mode="json"), room=user_room)
|
||||
await self._sio.emit(event=event_name, data=event_data.model_dump(mode="json"), room="admin")
|
||||
|
||||
logger.debug(f"Emitted private queue item event {event_name} to user room {user_room} and admin room")
|
||||
|
||||
# RecallParametersUpdatedEvent is private - only emit to owner + admins.
|
||||
#
|
||||
# Emit to the union of the owner room and the admin room in a SINGLE
|
||||
# call. python-socketio deduplicates recipients across a room list,
|
||||
# so a socket that belongs to BOTH rooms — e.g. the "system" user in
|
||||
# single-user mode, which is also an admin — receives the event
|
||||
# exactly once. Two separate emits would deliver it twice: harmless
|
||||
# for the idempotent scalar recall fields (the frontend just re-sets
|
||||
# them), but the append-mode reference-image recall *pushes* rather
|
||||
# than replaces, so a double delivery adds the same reference image
|
||||
# twice.
|
||||
elif isinstance(event_data, RecallParametersUpdatedEvent):
|
||||
user_room = f"user:{event_data.user_id}"
|
||||
await self._sio.emit(
|
||||
event=event_name, data=event_data.model_dump(mode="json"), room=[user_room, "admin"]
|
||||
)
|
||||
logger.debug(f"Emitted private recall_parameters_updated event to user room {user_room} and admin room")
|
||||
|
||||
# BatchEnqueuedEvent: full to owner+admin, sanitized to everyone else in the queue
|
||||
# room so their badge total and queue list pick up the new items.
|
||||
elif isinstance(event_data, BatchEnqueuedEvent):
|
||||
user_room = f"user:{event_data.user_id}"
|
||||
await self._sio.emit(event=event_name, data=event_data.model_dump(mode="json"), room=user_room)
|
||||
await self._sio.emit(event=event_name, data=event_data.model_dump(mode="json"), room="admin")
|
||||
|
||||
sanitized = event_data.model_copy(
|
||||
update={"user_id": "redacted", "batch_id": "redacted", "origin": None}
|
||||
)
|
||||
await self._sio.emit(
|
||||
event=event_name,
|
||||
data=sanitized.model_dump(mode="json"),
|
||||
room=event_data.queue_id,
|
||||
skip_sid=self._owner_and_admin_sids(event_data.user_id),
|
||||
)
|
||||
logger.debug(
|
||||
f"Emitted batch_enqueued: full to {user_room}+admin, sanitized to queue {event_data.queue_id}"
|
||||
)
|
||||
|
||||
# QueueItemsRetriedEvent carries queue item ids that should only be visible
|
||||
# to the affected owners + admins.
|
||||
elif isinstance(event_data, QueueItemsRetriedEvent):
|
||||
for user_id in event_data.user_ids:
|
||||
user_room = f"user:{user_id}"
|
||||
owner_event_data = event_data.model_copy(
|
||||
update={
|
||||
"retried_item_ids": event_data.retried_item_ids_by_user.get(user_id, []),
|
||||
"user_ids": [user_id],
|
||||
"retried_item_ids_by_user": {user_id: event_data.retried_item_ids_by_user.get(user_id, [])},
|
||||
}
|
||||
)
|
||||
await self._sio.emit(
|
||||
event=event_name, data=owner_event_data.model_dump(mode="json"), room=user_room
|
||||
)
|
||||
await self._sio.emit(event=event_name, data=event_data.model_dump(mode="json"), room="admin")
|
||||
logger.debug(
|
||||
f"Emitted private queue_items_retried event to user rooms {event_data.user_ids} and admin room"
|
||||
)
|
||||
|
||||
# QueueClearedEvent: an unscoped clear (user_id=None — admin or single-user mode)
|
||||
# deleted every user's items, so everyone gets the full event. A user-scoped clear
|
||||
# only deleted that user's rows: full event to owner+admin (single emit to a room
|
||||
# list so a socket in both rooms receives it once), sanitized companion to the rest
|
||||
# of the queue room so their queue lists and badge counts refetch without treating
|
||||
# the clear as their own.
|
||||
elif isinstance(event_data, QueueClearedEvent):
|
||||
if event_data.user_id is None:
|
||||
await self._sio.emit(
|
||||
event=event_name, data=event_data.model_dump(mode="json"), room=event_data.queue_id
|
||||
)
|
||||
logger.debug(f"Emitted unscoped queue_cleared to all subscribers in queue {event_data.queue_id}")
|
||||
else:
|
||||
user_room = f"user:{event_data.user_id}"
|
||||
await self._sio.emit(
|
||||
event=event_name, data=event_data.model_dump(mode="json"), room=[user_room, "admin"]
|
||||
)
|
||||
sanitized = event_data.model_copy(update={"user_id": "redacted"})
|
||||
await self._sio.emit(
|
||||
event=event_name,
|
||||
data=sanitized.model_dump(mode="json"),
|
||||
room=event_data.queue_id,
|
||||
skip_sid=self._owner_and_admin_sids(event_data.user_id),
|
||||
)
|
||||
logger.debug(
|
||||
f"Emitted queue_cleared: full to {user_room}+admin, sanitized to queue {event_data.queue_id}"
|
||||
)
|
||||
|
||||
else:
|
||||
# For remaining queue events that do not carry user identity,
|
||||
# emit to all subscribers in the queue room.
|
||||
await self._sio.emit(
|
||||
event=event_name, data=event_data.model_dump(mode="json"), room=event_data.queue_id
|
||||
)
|
||||
logger.debug(
|
||||
f"Emitted general queue event {event_name} to all subscribers in queue {event_data.queue_id}"
|
||||
)
|
||||
except Exception as e:
|
||||
# Log any unhandled exceptions in event handling to prevent silent failures
|
||||
logger.error(f"Error handling queue event {event[0]}: {e}", exc_info=True)
|
||||
|
||||
async def _handle_model_event(self, event: FastAPIEvent[ModelEventBase | DownloadEventBase]) -> None:
|
||||
await self._sio.emit(event=event[0], data=event[1].model_dump(mode="json"))
|
||||
|
||||
async def _handle_bulk_image_download_event(self, event: FastAPIEvent[BulkDownloadEventBase]) -> None:
|
||||
event_name, event_data = event
|
||||
# Route to user-specific + admin rooms so that other authenticated
|
||||
# users cannot learn the bulk_download_item_name (the capability token
|
||||
# needed to fetch the zip from the unauthenticated GET endpoint).
|
||||
# In single-user mode (user_id="system"), fall back to the shared
|
||||
# bulk_download_id room for backward compatibility.
|
||||
if hasattr(event_data, "user_id") and event_data.user_id != "system":
|
||||
user_room = f"user:{event_data.user_id}"
|
||||
await self._sio.emit(event=event_name, data=event_data.model_dump(mode="json"), room=user_room)
|
||||
await self._sio.emit(event=event_name, data=event_data.model_dump(mode="json"), room="admin")
|
||||
else:
|
||||
await self._sio.emit(
|
||||
event=event_name, data=event_data.model_dump(mode="json"), room=event_data.bulk_download_id
|
||||
)
|
||||
|
||||
async def _handle_workflow_event(self, event: FastAPIEvent[WorkflowEventBase]) -> None:
|
||||
event_name, event_data = event
|
||||
payload = event_data.model_dump(mode="json")
|
||||
|
||||
if not self._is_multiuser_enabled():
|
||||
await self._sio.emit(event=event_name, data=payload, room="admin")
|
||||
return
|
||||
|
||||
await self._sio.emit(event=event_name, data=payload, room=f"user:{event_data.user_id}")
|
||||
await self._sio.emit(event=event_name, data=payload, room="admin")
|
||||
|
||||
if event_name == "workflow_created":
|
||||
if getattr(event_data, "is_public", False):
|
||||
await self._sio.emit(event=event_name, data=payload, room="workflows:shared")
|
||||
return
|
||||
|
||||
if event_name == "workflow_deleted":
|
||||
if getattr(event_data, "is_public", False):
|
||||
await self._sio.emit(event=event_name, data=payload, room="workflows:shared")
|
||||
return
|
||||
|
||||
if event_name == "workflow_updated":
|
||||
if getattr(event_data, "new_is_public", False):
|
||||
await self._sio.emit(event=event_name, data=payload, room="workflows:shared")
|
||||
elif getattr(event_data, "old_is_public", False):
|
||||
access_revoked = WorkflowAccessRevokedEvent.build(
|
||||
workflow_id=event_data.workflow_id, user_id=event_data.user_id
|
||||
)
|
||||
await self._sio.emit(
|
||||
event=access_revoked.__event_name__,
|
||||
data=access_revoked.model_dump(mode="json"),
|
||||
room="workflows:shared",
|
||||
)
|
||||
Reference in New Issue
Block a user