diff --git a/src/khoj/database/migrations/0035_convert_png_to_webp.py b/src/khoj/database/migrations/0035_convert_png_to_webp.py new file mode 100644 index 00000000..7d28a07d --- /dev/null +++ b/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), + ]