Add migration script to convert PNG to WebP references in database

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

@@ -0,0 +1,77 @@
+# 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
+
+
+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"] == "text-to-image":
+                # 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"] = "text-to-image-v3"
+                    webp_image_io.close()
+
+            if chat["by"] == "khoj" and chat["intent"]["type"] == "text-to-image2":
+                print("❗️ Please MANUALLY update PNG images created by Khoj in your AWS S3 bucket to WebP format.")
+                # Convert PNG url to WebP url
+                chat["message"] = chat["message"].replace(".png", ".webp")
+
+        # 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"] == "text-to-image":
+                # 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"] = "text-to-image"
+                    webp_image_io.close()
+
+            if chat["by"] == "khoj" and chat["intent"]["type"] == "text-to-image2":
+                # Convert WebP url to PNG url
+                print("❗️ Please MANUALLY update WebP images created by Khoj in your AWS S3 bucket to PNG format.")
+                chat["message"] = chat["message"].replace(".webp", ".png")
+
+        # 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),
+    ]