From f0c70fe3f19e59d2be746e00a4e16029569562dd Mon Sep 17 00:00:00 2001 From: psychedelicious <4822129+psychedelicious@users.noreply.github.com> Date: Thu, 14 Dec 2023 07:47:08 +1100 Subject: [PATCH] fix(db): add error handling for workflow migration - Handle an image file not existing despite being in the database. - Add a simple pydantic model that tests only for the existence of a workflow's version. - Check against this new model when migrating workflows, skipping if the workflow fails validation. If it succeeds, the frontend should be able to handle the workflow. --- .../sqlite_migrator/migrations/migration_2.py | 18 ++++++++++++++++-- .../workflow_records_common.py | 12 ++++++++++++ 2 files changed, 28 insertions(+), 2 deletions(-) diff --git a/invokeai/app/services/shared/sqlite_migrator/migrations/migration_2.py b/invokeai/app/services/shared/sqlite_migrator/migrations/migration_2.py index 9efd2a1283..9b9dedcc58 100644 --- a/invokeai/app/services/shared/sqlite_migrator/migrations/migration_2.py +++ b/invokeai/app/services/shared/sqlite_migrator/migrations/migration_2.py @@ -1,10 +1,15 @@ import sqlite3 from logging import Logger +from pydantic import ValidationError from tqdm import tqdm from invokeai.app.services.image_files.image_files_base import ImageFileStorageBase +from invokeai.app.services.image_files.image_files_common import ImageFileNotFoundException from invokeai.app.services.shared.sqlite_migrator.sqlite_migrator_common import Migration +from invokeai.app.services.workflow_records.workflow_records_common import ( + UnsafeWorkflowWithVersionValidator, +) class Migration2Callback: @@ -134,7 +139,7 @@ class Migration2Callback: This migrate callback checks each image for the presence of an embedded workflow, then updates its entry in the database accordingly. """ - # Get the total number of images and chunk it into pages + # Get all image names cursor.execute("SELECT image_name FROM images") image_names: list[str] = [image[0] for image in cursor.fetchall()] total_image_names = len(image_names) @@ -149,8 +154,17 @@ class Migration2Callback: pbar = tqdm(image_names) for idx, image_name in enumerate(pbar): pbar.set_description(f"Checking image {idx + 1}/{total_image_names} for workflow") - pil_image = self._image_files.get(image_name) + try: + pil_image = self._image_files.get(image_name) + except ImageFileNotFoundException: + self._logger.warning(f"Image {image_name} not found, skipping") + continue if "invokeai_workflow" in pil_image.info: + try: + UnsafeWorkflowWithVersionValidator.validate_json(pil_image.info.get("invokeai_workflow", "")) + except ValidationError: + self._logger.warning(f"Image {image_name} has invalid embedded workflow, skipping") + continue to_migrate.append((True, image_name)) self._logger.info(f"Adding {len(to_migrate)} embedded workflows to database") diff --git a/invokeai/app/services/workflow_records/workflow_records_common.py b/invokeai/app/services/workflow_records/workflow_records_common.py index ffc98bb786..b7a1fba729 100644 --- a/invokeai/app/services/workflow_records/workflow_records_common.py +++ b/invokeai/app/services/workflow_records/workflow_records_common.py @@ -71,6 +71,18 @@ class WorkflowWithoutID(BaseModel): WorkflowWithoutIDValidator = TypeAdapter(WorkflowWithoutID) +class UnsafeWorkflowWithVersion(BaseModel): + """ + This utility model only requires a workflow to have a valid version string. + It is used to validate a workflow version without having to validate the entire workflow. + """ + + meta: WorkflowMeta = Field(description="The meta of the workflow.") + + +UnsafeWorkflowWithVersionValidator = TypeAdapter(UnsafeWorkflowWithVersion) + + class Workflow(WorkflowWithoutID): id: str = Field(description="The id of the workflow.")