diff --git a/invokeai/backend/util/migrate_image_workflows.py b/invokeai/backend/util/migrate_image_workflows.py new file mode 100644 index 0000000000..24bd9a1538 --- /dev/null +++ b/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)