diff --git a/vibe_bot/llama_wrapper.py b/vibe_bot/llama_wrapper.py index fab13a2..0805dce 100644 --- a/vibe_bot/llama_wrapper.py +++ b/vibe_bot/llama_wrapper.py @@ -5,8 +5,6 @@ import openai from typing import Iterable from openai.types.chat import ChatCompletionMessageParam -from openai._types import FileTypes, SequenceNotStr -from typing import Union from io import BufferedReader, BytesIO @@ -138,7 +136,7 @@ def image_generation(prompt: str, openai_url: str, openai_api_key: str, n=1) -> def image_edit( - image: BufferedReader | BytesIO, + image: BufferedReader | BytesIO | list[BufferedReader] | list[BytesIO], prompt: str, openai_url: str, openai_api_key: str, diff --git a/vibe_bot/main.py b/vibe_bot/main.py index 2f5a721..61de370 100644 --- a/vibe_bot/main.py +++ b/vibe_bot/main.py @@ -278,14 +278,17 @@ async def doodlebob(ctx, *, message: str): @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) + image_data_list = [] + for discord_image in ctx.message.attachments: + image_url = discord_image.url + image_data = requests.get(image_url).content + image_bytestream = BytesIO(image_data) + image_data_list.append(image_bytestream) await ctx.send(f"**Rewriting history to match {message[:100]}...**") image_b64 = llama_wrapper.image_edit( - image=image_bytestream, + image=image_data_list, prompt=message, openai_url=IMAGE_EDIT_ENDPOINT, openai_api_key=IMAGE_EDIT_ENDPOINT_KEY,