add image editing

main.py

@@ -5,11 +5,13 @@ import os
 import json
 import base64
 from io import BytesIO
+from openai import OpenAI
 
 DISCORD_TOKEN = os.getenv("DISCORD_TOKEN", "placeholder")
 
 OPENAI_API_ENDPOINT = os.getenv("OPENAI_API_ENDPOINT")
 IMAGE_GEN_ENDPOINT = os.getenv("IMAGE_GEN_ENDPOINT")
+IMAGE_EDIT_ENDPOINT = os.getenv("IMAGE_EDIT_ENDPOINT")
 
 if not OPENAI_API_ENDPOINT:
     raise Exception("OPENAI_API_ENDPOINT required.")
@@ -115,14 +117,18 @@ async def doodlebob(ctx, *, message: str):
         ],
     }
 
+    # Wait for the generated image prompt
     image_prompt = await call_llm(ctx, image_prompt_payload)
 
+    # If the string is empty we had an error
     if image_prompt == "":
         print("No image prompt supplied. Check for errors.")
         return
 
+    # Alert the user we're generating the image
     await ctx.send(f"**Doodlebob calling drone strike on {image_prompt[:100]}...**")
 
+    # Create the image prompt payload
     image_payload = {
         "model": "default",
         "prompt": image_prompt,
@@ -130,6 +136,7 @@ async def doodlebob(ctx, *, message: str):
         "size": "1024x1024",
     }
 
+    # Call the image generation endpoint
     response = requests.post(
         f"{IMAGE_GEN_ENDPOINT}/images/generations",
         json=image_payload,
@@ -149,6 +156,32 @@ async def doodlebob(ctx, *, message: str):
         return None
 
 
+@bot.command(name="retcon")
+async def retcon(ctx, *, message: str):
+    image_url = ctx.message.attachments[0].url
+    image_data = requests.get(image_url).content
+    image_bytestream = BytesIO(image_data)
+
+    await ctx.send(f"**Rewriting history to match {message[:100]}...**")
+
+    client = OpenAI(base_url=IMAGE_EDIT_ENDPOINT, api_key=OPENAI_API_KEY)
+
+    result = client.images.edit(
+        model="placeholder",
+        image=[image_bytestream],
+        prompt=message,
+        size="1024x1024",
+    )
+
+    image_base64 = result.data[0].b64_json
+    image_bytes = base64.b64decode(image_base64)
+
+    # Save the image to a file
+    edited_image_data = BytesIO(image_bytes)
+    send_img = discord.File(edited_image_data, filename="image.png")
+    await ctx.send(file=send_img)
+
+
 async def handle_chat(ctx, *, message: str, payload: dict, response_prefix: str):
     # Check if API key is set
     if not OPENAI_API_KEY: