diff --git a/invokeai/app/api/dependencies.py b/invokeai/app/api/dependencies.py
index 83fc8d9e11..d64610b717 100644
--- a/invokeai/app/api/dependencies.py
+++ b/invokeai/app/api/dependencies.py
@@ -2,6 +2,7 @@
from typing import Optional
from logging import Logger
+import threading
import os
from invokeai.app.services.board_image_record_storage import (
SqliteBoardImageRecordStorage,
@@ -65,6 +66,8 @@ class ApiDependencies:
logger.info(f"Root directory = {str(config.root_path)}")
logger.debug(f"Internet connectivity is {config.internet_available}")
+ lock = threading.Lock()
+
events = FastAPIEventService(event_handler_id)
output_folder = config.output_path
@@ -74,17 +77,17 @@ class ApiDependencies:
db_location.parent.mkdir(parents=True, exist_ok=True)
graph_execution_manager = SqliteItemStorage[GraphExecutionState](
- filename=db_location, table_name="graph_executions"
+ filename=db_location, table_name="graph_executions", lock=lock
)
urls = LocalUrlService()
- image_record_storage = SqliteImageRecordStorage(db_location)
+ image_record_storage = SqliteImageRecordStorage(db_location, lock=lock)
image_file_storage = DiskImageFileStorage(f"{output_folder}/images")
names = SimpleNameService()
latents = ForwardCacheLatentsStorage(DiskLatentsStorage(f"{output_folder}/latents"))
- board_record_storage = SqliteBoardRecordStorage(db_location)
- board_image_record_storage = SqliteBoardImageRecordStorage(db_location)
+ board_record_storage = SqliteBoardRecordStorage(db_location, lock=lock)
+ board_image_record_storage = SqliteBoardImageRecordStorage(db_location, lock=lock)
boards = BoardService(
services=BoardServiceDependencies(
@@ -118,7 +121,7 @@ class ApiDependencies:
)
)
- batch_manager_storage = SqliteBatchProcessStorage(db_location)
+ batch_manager_storage = SqliteBatchProcessStorage(db_location, lock=lock)
batch_manager = BatchManager(batch_manager_storage)
services = InvocationServices(
@@ -130,7 +133,7 @@ class ApiDependencies:
boards=boards,
board_images=board_images,
queue=MemoryInvocationQueue(),
- graph_library=SqliteItemStorage[LibraryGraph](filename=db_location, table_name="graphs"),
+ graph_library=SqliteItemStorage[LibraryGraph](filename=db_location, table_name="graphs", lock=lock),
graph_execution_manager=graph_execution_manager,
processor=DefaultInvocationProcessor(),
configuration=config,
diff --git a/invokeai/app/cli_app.py b/invokeai/app/cli_app.py
index 39b22e9dfb..92c255eb45 100644
--- a/invokeai/app/cli_app.py
+++ b/invokeai/app/cli_app.py
@@ -38,6 +38,7 @@ from invokeai.app.services.images import ImageService, ImageServiceDependencies
from invokeai.app.services.resource_name import SimpleNameService
from invokeai.app.services.urls import LocalUrlService
from invokeai.app.services.batch_manager import BatchManager
+from invokeai.app.services.batch_manager_storage import SqliteBatchProcessStorage
from .services.default_graphs import default_text_to_image_graph_id, create_system_graphs
from .services.latent_storage import DiskLatentsStorage, ForwardCacheLatentsStorage
@@ -301,13 +302,16 @@ def invoke_cli():
)
)
+ batch_manager_storage = SqliteBatchProcessStorage(db_location)
+ batch_manager = BatchManager(batch_manager_storage)
+
services = InvocationServices(
model_manager=model_manager,
events=events,
latents=ForwardCacheLatentsStorage(DiskLatentsStorage(f"{output_folder}/latents")),
images=images,
boards=boards,
- batch_manager=BatchManager(),
+ batch_manager=batch_manager,
board_images=board_images,
queue=MemoryInvocationQueue(),
graph_library=SqliteItemStorage[LibraryGraph](filename=db_location, table_name="graphs"),
diff --git a/invokeai/app/services/batch_manager.py b/invokeai/app/services/batch_manager.py
index eba5a8d676..a5c92bb1fc 100644
--- a/invokeai/app/services/batch_manager.py
+++ b/invokeai/app/services/batch_manager.py
@@ -12,6 +12,7 @@ from invokeai.app.services.graph import Graph, GraphExecutionState
from invokeai.app.services.invoker import Invoker
from invokeai.app.services.batch_manager_storage import (
BatchProcessStorageBase,
+ BatchSessionNotFoundException,
Batch,
BatchProcess,
BatchSession,
@@ -98,7 +99,10 @@ class BatchManager(BatchManagerBase):
return GraphExecutionState(graph=graph)
def run_batch_process(self, batch_id: str):
- created_session = self.__batch_process_storage.get_created_session(batch_id)
+ try:
+ created_session = self.__batch_process_storage.get_created_session(batch_id)
+ except BatchSessionNotFoundException:
+ return
ges = self.__invoker.services.graph_execution_manager.get(created_session.session_id)
self.__invoker.invoke(ges, invoke_all=True)
diff --git a/invokeai/app/services/batch_manager_storage.py b/invokeai/app/services/batch_manager_storage.py
index 44fb5d928f..bc8804d21f 100644
--- a/invokeai/app/services/batch_manager_storage.py
+++ b/invokeai/app/services/batch_manager_storage.py
@@ -15,9 +15,8 @@ import json
from invokeai.app.invocations.baseinvocation import (
BaseInvocation,
)
-from invokeai.app.services.graph import Graph, GraphExecutionState
+from invokeai.app.services.graph import Graph
from invokeai.app.models.image import ImageField
-from invokeai.app.services.image_record_storage import OffsetPaginatedResults
from pydantic import BaseModel, Field, Extra, parse_raw_as
@@ -169,14 +168,14 @@ class SqliteBatchProcessStorage(BatchProcessStorageBase):
_cursor: sqlite3.Cursor
_lock: threading.Lock
- def __init__(self, filename: str) -> None:
+ def __init__(self, filename: str, lock: threading.Lock = threading.Lock()) -> None:
super().__init__()
self._filename = filename
self._conn = sqlite3.connect(filename, check_same_thread=False)
# Enable row factory to get rows as dictionaries (must be done before making the cursor!)
self._conn.row_factory = sqlite3.Row
self._cursor = self._conn.cursor()
- self._lock = threading.Lock()
+ self._lock = lock
try:
self._lock.acquire()
@@ -231,7 +230,7 @@ class SqliteBatchProcessStorage(BatchProcessStorageBase):
CREATE TABLE IF NOT EXISTS batch_session (
batch_id TEXT NOT NULL,
session_id TEXT NOT NULL,
- state TEXT NOT NULL DEFAULT('created'),
+ state TEXT NOT NULL,
created_at DATETIME NOT NULL DEFAULT(STRFTIME('%Y-%m-%d %H:%M:%f', 'NOW')),
-- updated via trigger
updated_at DATETIME NOT NULL DEFAULT(STRFTIME('%Y-%m-%d %H:%M:%f', 'NOW')),
@@ -267,7 +266,7 @@ class SqliteBatchProcessStorage(BatchProcessStorageBase):
ON batch_session FOR EACH ROW
BEGIN
UPDATE batch_session SET updated_at = STRFTIME('%Y-%m-%d %H:%M:%f', 'NOW')
- WHERE batch_id = old.batch_id AND image_name = old.image_name;
+ WHERE batch_id = old.batch_id AND session_id = old.session_id;
END;
"""
)
@@ -298,12 +297,13 @@ class SqliteBatchProcessStorage(BatchProcessStorageBase):
) -> BatchProcess:
try:
self._lock.acquire()
+ batches = [batch.json() for batch in batch_process.batches]
self._cursor.execute(
"""--sql
INSERT OR IGNORE INTO batch_process (batch_id, batches, graph)
VALUES (?, ?, ?);
""",
- (batch_process.batch_id, json.dumps([batch.json() for batch in batch_process.batches]), batch_process.graph.json()),
+ (batch_process.batch_id, json.dumps(batches), batch_process.graph.json()),
)
self._conn.commit()
except sqlite3.Error as e:
@@ -321,10 +321,11 @@ class SqliteBatchProcessStorage(BatchProcessStorageBase):
batch_id = session_dict.get("batch_id", "unknown")
batches_raw = session_dict.get("batches", "unknown")
graph_raw = session_dict.get("graph", "unknown")
-
+ batches = json.loads(batches_raw)
+ batches = [parse_raw_as(Batch, batch) for batch in batches]
return BatchProcess(
batch_id=batch_id,
- batches=[parse_raw_as(Batch, batch) for batch in json.loads(batches_raw)],
+ batches=batches,
graph=parse_raw_as(Graph, graph_raw),
)
@@ -398,7 +399,7 @@ class SqliteBatchProcessStorage(BatchProcessStorageBase):
self._lock.release()
if result is None:
raise BatchSessionNotFoundException
- return BatchSession(**dict(result))
+ return self._deserialize_batch_session(dict(result))
def _deserialize_batch_session(self, session_dict: dict) -> BatchSession:
"""Deserializes a batch session."""
diff --git a/invokeai/app/services/board_image_record_storage.py b/invokeai/app/services/board_image_record_storage.py
index f0007c8cef..c4c058422b 100644
--- a/invokeai/app/services/board_image_record_storage.py
+++ b/invokeai/app/services/board_image_record_storage.py
@@ -62,14 +62,14 @@ class SqliteBoardImageRecordStorage(BoardImageRecordStorageBase):
_cursor: sqlite3.Cursor
_lock: threading.Lock
- def __init__(self, filename: str) -> None:
+ def __init__(self, filename: str, lock: threading.Lock = threading.Lock()) -> None:
super().__init__()
self._filename = filename
self._conn = sqlite3.connect(filename, check_same_thread=False)
# Enable row factory to get rows as dictionaries (must be done before making the cursor!)
self._conn.row_factory = sqlite3.Row
self._cursor = self._conn.cursor()
- self._lock = threading.Lock()
+ self._lock = lock
try:
self._lock.acquire()
diff --git a/invokeai/app/services/board_record_storage.py b/invokeai/app/services/board_record_storage.py
index 2fad7b0ab3..2626799836 100644
--- a/invokeai/app/services/board_record_storage.py
+++ b/invokeai/app/services/board_record_storage.py
@@ -95,14 +95,14 @@ class SqliteBoardRecordStorage(BoardRecordStorageBase):
_cursor: sqlite3.Cursor
_lock: threading.Lock
- def __init__(self, filename: str) -> None:
+ def __init__(self, filename: str, lock: threading.Lock = threading.Lock()) -> None:
super().__init__()
self._filename = filename
self._conn = sqlite3.connect(filename, check_same_thread=False)
# Enable row factory to get rows as dictionaries (must be done before making the cursor!)
self._conn.row_factory = sqlite3.Row
self._cursor = self._conn.cursor()
- self._lock = threading.Lock()
+ self._lock = lock
try:
self._lock.acquire()
diff --git a/invokeai/app/services/image_record_storage.py b/invokeai/app/services/image_record_storage.py
index 8c274ab8f9..3dab08bcd7 100644
--- a/invokeai/app/services/image_record_storage.py
+++ b/invokeai/app/services/image_record_storage.py
@@ -155,14 +155,14 @@ class SqliteImageRecordStorage(ImageRecordStorageBase):
_cursor: sqlite3.Cursor
_lock: threading.Lock
- def __init__(self, filename: str) -> None:
+ def __init__(self, filename: str, lock: threading.Lock = threading.Lock()) -> None:
super().__init__()
self._filename = filename
self._conn = sqlite3.connect(filename, check_same_thread=False)
# Enable row factory to get rows as dictionaries (must be done before making the cursor!)
self._conn.row_factory = sqlite3.Row
self._cursor = self._conn.cursor()
- self._lock = threading.Lock()
+ self._lock = lock
try:
self._lock.acquire()
diff --git a/invokeai/app/services/sqlite.py b/invokeai/app/services/sqlite.py
index 4f63ffb368..f7a4cadbfe 100644
--- a/invokeai/app/services/sqlite.py
+++ b/invokeai/app/services/sqlite.py
@@ -20,16 +20,16 @@ class SqliteItemStorage(ItemStorageABC, Generic[T]):
_id_field: str
_lock: Lock
- def __init__(self, filename: str, table_name: str, id_field: str = "id"):
+ def __init__(self, filename: str, table_name: str, id_field: str = "id", lock: Lock = Lock()):
super().__init__()
self._filename = filename
self._table_name = table_name
self._id_field = id_field # TODO: validate that T has this field
- self._lock = Lock()
+ self._lock = lock
self._conn = sqlite3.connect(
self._filename, check_same_thread=False
) # TODO: figure out a better threading solution
- self._conn.set_trace_callback(print)
+ self._conn.execute('pragma journal_mode=wal')
self._cursor = self._conn.cursor()
self._create_table()
@@ -55,21 +55,12 @@ class SqliteItemStorage(ItemStorageABC, Generic[T]):
def set(self, item: T):
try:
self._lock.acquire()
- json = item.json()
- print('-----------------locking db-----------------')
- traceback.print_stack(limit=2)
self._cursor.execute(
f"""INSERT OR REPLACE INTO {self._table_name} (item) VALUES (?);""",
- (json,),
+ (item.json(),),
)
self._conn.commit()
- self._cursor.close()
- self._cursor = self._conn.cursor()
- print('-----------------unlocking db-----------------')
- except Exception as e:
- print("Exception!")
- print(e)
finally:
self._lock.release()
self._on_changed(item)
@@ -77,12 +68,8 @@ class SqliteItemStorage(ItemStorageABC, Generic[T]):
def get(self, id: str) -> Optional[T]:
try:
self._lock.acquire()
- print('-----------------locking db-----------------')
self._cursor.execute(f"""SELECT item FROM {self._table_name} WHERE id = ?;""", (str(id),))
result = self._cursor.fetchone()
- self._cursor.close()
- self._cursor = self._conn.cursor()
- print('-----------------unlocking db-----------------')
finally:
self._lock.release()