From 0003ac29e980f5b2f3d003a0e7b2e19bb045f8bf Mon Sep 17 00:00:00 2001 From: ducoterra Date: Mon, 30 Mar 2026 13:06:31 -0400 Subject: [PATCH] allow multi-track drifting (multiple images to retcon) --- vibe_bot/llama_wrapper.py | 4 +--- vibe_bot/main.py | 11 +++++++---- 2 files changed, 8 insertions(+), 7 deletions(-) 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,