feat: add script to migrate image workflows

invokeai/backend/util/migrate_image_workflows.py

@@ -0,0 +1,84 @@
+import sqlite3
+from datetime import datetime
+from pathlib import Path
+
+from PIL import Image
+from tqdm import tqdm
+
+from invokeai.app.services.config.config_default import InvokeAIAppConfig
+
+
+def migrate_image_workflows(output_path: Path, database: Path, page_size=100):
+    """
+    In the v3.5.0 release, InvokeAI changed how it handles image workflows. The `images` table in
+    the database now has a `has_workflow` column, indicating if an image has a workflow embedded.
+
+    This script checks each image for the presence of an embedded workflow, then updates its entry
+    in the database accordingly.
+
+    1) Check if the database is updated to support image workflows. Aborts if it doesn't have the
+    `has_workflow` column yet.
+    2) Backs up the database.
+    3) Opens each image in the `images` table via PIL
+    4) Checks if the `"invokeai_workflow"` attribute its in the image's embedded metadata, indicating
+    that it has a workflow.
+    5) If it does, updates the `has_workflow` column for that image to `TRUE`.
+
+    If there are any problems, the script immediately aborts. Because the processing happens in chunks,
+    if there is a problem, it is suggested that you restore the database from the backup and try again.
+    """
+    output_path = output_path
+    database = database
+    conn = sqlite3.connect(database)
+    cursor = conn.cursor()
+
+    # We can only migrate if the `images` table has the `has_workflow` column
+    cursor.execute("PRAGMA table_info(images)")
+    columns = [column[1] for column in cursor.fetchall()]
+    if "has_workflow" not in columns:
+        raise Exception("Database needs to be updated to support image workflows")
+
+    # Back up the database before we start
+    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+    backup_path = database.parent / f"{database.stem}_migrate-image-workflows_{timestamp}.db"
+    print(f"Backing up database to {backup_path}")
+    backup_conn = sqlite3.connect(backup_path)
+    with backup_conn:
+        conn.backup(backup_conn)
+    backup_conn.close()
+
+    # Get the total number of images and chunk it into pages
+    cursor.execute("SELECT COUNT(*) FROM images")
+    total_images = cursor.fetchone()[0]
+    total_pages = (total_images + page_size - 1) // page_size
+    print(f"Processing {total_images} images in chunks of {page_size} images...")
+
+    # Migrate the images
+    migrated_count = 0
+    pbar = tqdm(range(total_pages))
+    for page in pbar:
+        pbar.set_description(f"Migrating page {page + 1}/{total_pages}")
+        offset = page * page_size
+        cursor.execute("SELECT image_name FROM images LIMIT ? OFFSET ?", (page_size, offset))
+        images = cursor.fetchall()
+        for image_name in images:
+            image_path = output_path / "images" / image_name[0]
+            with Image.open(image_path) as img:
+                if "invokeai_workflow" in img.info:
+                    cursor.execute("UPDATE images SET has_workflow = TRUE WHERE image_name = ?", (image_name[0],))
+                    migrated_count += 1
+        conn.commit()
+    conn.close()
+
+    print(f"Migrated workflows for {migrated_count} images.")
+
+
+if __name__ == "__main__":
+    config = InvokeAIAppConfig.get_config()
+    output_path = config.output_path
+    database = config.db_path
+
+    assert output_path is not None
+    assert output_path.exists()
+    assert database.exists()
+    migrate_image_workflows(output_path=output_path, database=database, page_size=100)