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

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"]:
+            for chat in conversation.conversation_log.get("chat", []):
-                if chat["by"] == "khoj" and chat["intent"]["type"] == ImageIntentType.TEXT_TO_IMAGE2.value:
+                if (
-                    if options["reverse"] and chat["message"].endswith(".webp"):
+                    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"]:

src/khoj/database/migrations/0035_convert_png_to_webp.py

@@ -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),
-    ]