sij/khoj
1
0
Fork 0
mirror of https://github.com/khoj-ai/khoj.git synced 2025-02-17 08:04:21 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Use Django management cmd to update inline images in DB to/from WebP/PNG

Browse source 
This provides Khoj server admins more control on migrating their S3
images to WebP format from PNG
This commit is contained in:
Debanjum Singh Solanky 2024-04-15 20:12:12 +05:30
parent 7fab8d6586
commit 4e7812fe55
2 changed files with 63 additions and 73 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

67
src/khoj/database/management/commands/convert_images_png_to_webp.py
View file

@ -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"]:

69
src/khoj/database/migrations/0035_convert_png_to_webp.py
View file

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