From 4e7812fe5514965ed98a10d40f9bfdfd916819f0 Mon Sep 17 00:00:00 2001 From: Debanjum Singh Solanky Date: Mon, 15 Apr 2024 20:12:12 +0530 Subject: [PATCH] Use Django management cmd to update inline images in DB to/from WebP/PNG This provides Khoj server admins more control on migrating their S3 images to WebP format from PNG --- .../commands/convert_images_png_to_webp.py | 67 ++++++++++++++++-- .../migrations/0035_convert_png_to_webp.py | 69 ------------------- 2 files changed, 63 insertions(+), 73 deletions(-) delete mode 100644 src/khoj/database/migrations/0035_convert_png_to_webp.py diff --git a/src/khoj/database/management/commands/convert_images_png_to_webp.py b/src/khoj/database/management/commands/convert_images_png_to_webp.py index b1ad8615..d137b141 100644 --- a/src/khoj/database/management/commands/convert_images_png_to_webp.py +++ b/src/khoj/database/management/commands/convert_images_png_to_webp.py @@ -1,4 +1,8 @@ +import base64 +import io + from django.core.management.base import BaseCommand +from PIL import Image from khoj.database.models import Conversation from khoj.utils.helpers import ImageIntentType @@ -19,19 +23,74 @@ class Command(BaseCommand): updated_count = 0 for conversation in Conversation.objects.all(): conversation_updated = False - for chat in conversation.conversation_log["chat"]: - if chat["by"] == "khoj" and chat["intent"]["type"] == ImageIntentType.TEXT_TO_IMAGE2.value: - if options["reverse"] and chat["message"].endswith(".webp"): + for chat in conversation.conversation_log.get("chat", []): + if ( + chat.get("by", "") == "khoj" + and chat.get("intent", {}).get("type", "") == ImageIntentType.TEXT_TO_IMAGE.value + and not options["reverse"] + ): + # Decode the base64 encoded PNG image + print("Decode the base64 encoded PNG image") + decoded_image = base64.b64decode(chat["message"]) + + # Convert images from PNG to WebP format + print("Convert images from PNG to WebP format") + image_io = io.BytesIO(decoded_image) + with Image.open(image_io) as png_image: + webp_image_io = io.BytesIO() + png_image.save(webp_image_io, "WEBP") + + # Encode the WebP image back to base64 + webp_image_bytes = webp_image_io.getvalue() + chat["message"] = base64.b64encode(webp_image_bytes).decode() + chat["intent"]["type"] = ImageIntentType.TEXT_TO_IMAGE_V3.value + webp_image_io.close() + conversation_updated = True + updated_count += 1 + + elif ( + chat.get("by", "") == "khoj" + and chat.get("intent", {}).get("type", "") == ImageIntentType.TEXT_TO_IMAGE_V3.value + and options["reverse"] + ): + # Decode the base64 encoded WebP image + print("Decode the base64 encoded WebP image") + decoded_image = base64.b64decode(chat["message"]) + + # Convert images from WebP to PNG format + print("Convert images from WebP to PNG format") + image_io = io.BytesIO(decoded_image) + with Image.open(image_io) as png_image: + webp_image_io = io.BytesIO() + png_image.save(webp_image_io, "PNG") + + # Encode the WebP image back to base64 + webp_image_bytes = webp_image_io.getvalue() + chat["message"] = base64.b64encode(webp_image_bytes).decode() + chat["intent"]["type"] = ImageIntentType.TEXT_TO_IMAGE.value + webp_image_io.close() + conversation_updated = True + updated_count += 1 + + elif ( + chat.get("by", "") == "khoj" + and chat.get("intent", {}).get("type", "") == ImageIntentType.TEXT_TO_IMAGE2.value + ): + if options["reverse"] and chat.get("message", "").endswith(".webp"): # Convert WebP url to PNG url + print("Convert WebP url to PNG url") chat["message"] = chat["message"].replace(".webp", ".png") conversation_updated = True updated_count += 1 - elif chat["message"].endswith(".png"): + elif chat.get("message", "").endswith(".png"): # Convert PNG url to WebP url + print("Convert PNG url to WebP url") chat["message"] = chat["message"].replace(".png", ".webp") conversation_updated = True updated_count += 1 + if conversation_updated: + print("Save the updated conversation") conversation.save() if updated_count > 0 and options["reverse"]: diff --git a/src/khoj/database/migrations/0035_convert_png_to_webp.py b/src/khoj/database/migrations/0035_convert_png_to_webp.py deleted file mode 100644 index 35495629..00000000 --- a/src/khoj/database/migrations/0035_convert_png_to_webp.py +++ /dev/null @@ -1,69 +0,0 @@ -# Generated by Django 4.2.10 on 2024-04-13 17:54 - -import base64 -import io - -from django.db import migrations -from PIL import Image - -from khoj.utils.helpers import ImageIntentType - - -def convert_png_images_to_webp(apps, schema_editor): - # Get the model from the versioned app registry to ensure the correct version is used - Conversations = apps.get_model("database", "Conversation") - for conversation in Conversations.objects.all(): - for chat in conversation.conversation_log["chat"]: - if chat["by"] == "khoj" and chat["intent"]["type"] == ImageIntentType.TEXT_TO_IMAGE.value: - # Decode the base64 encoded PNG image - decoded_image = base64.b64decode(chat["message"]) - - # Convert images from PNG to WebP format - image_io = io.BytesIO(decoded_image) - with Image.open(image_io) as png_image: - webp_image_io = io.BytesIO() - png_image.save(webp_image_io, "WEBP") - - # Encode the WebP image back to base64 - webp_image_bytes = webp_image_io.getvalue() - chat["message"] = base64.b64encode(webp_image_bytes).decode() - chat["intent"]["type"] = ImageIntentType.TEXT_TO_IMAGE_V3.value - webp_image_io.close() - - # Save the updated conversation history - conversation.save() - - -def convert_webp_images_to_png(apps, schema_editor): - # Get the model from the versioned app registry to ensure the correct version is used - Conversations = apps.get_model("database", "Conversation") - for conversation in Conversations.objects.all(): - for chat in conversation.conversation_log["chat"]: - if chat["by"] == "khoj" and chat["intent"]["type"] == ImageIntentType.TEXT_TO_IMAGE_V3.value: - # Decode the base64 encoded PNG image - decoded_image = base64.b64decode(chat["message"]) - - # Convert images from PNG to WebP format - image_io = io.BytesIO(decoded_image) - with Image.open(image_io) as png_image: - webp_image_io = io.BytesIO() - png_image.save(webp_image_io, "PNG") - - # Encode the WebP image back to base64 - webp_image_bytes = webp_image_io.getvalue() - chat["message"] = base64.b64encode(webp_image_bytes).decode() - chat["intent"]["type"] = ImageIntentType.TEXT_TO_IMAGE.value - webp_image_io.close() - - # Save the updated conversation history - conversation.save() - - -class Migration(migrations.Migration): - dependencies = [ - ("database", "0034_alter_chatmodeloptions_chat_model"), - ] - - operations = [ - migrations.RunPython(convert_png_images_to_webp, reverse_code=convert_webp_images_to_png), - ]