Skip to content
This repository was archived by the owner on Dec 30, 2024. It is now read-only.

Commit 3336b4f

Browse files
author
FuseFairy
committed
refactor: enhance image creation function with environment variable handling
1 parent f422497 commit 3336b4f

File tree

1 file changed

+11
-8
lines changed

1 file changed

+11
-8
lines changed

src/image/image_create.py

+11-8
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,15 @@
11
import discord
22
import asyncio
33
import requests
4+
import os
45
from PIL import Image
56
from io import BytesIO
67
from re_edge_gpt import ImageGenAsync
78
from src.log import setup_logger
89
from src.image.button_view import ButtonView
10+
from dotenv import load_dotenv
11+
12+
load_dotenv()
913

1014
logger = setup_logger(__name__)
1115

@@ -17,13 +21,12 @@ async def create_image(interaction: discord.Interaction, users_chatbot: dict, pr
1721
user_id = interaction.user.id
1822
username = interaction.user
1923
channel = interaction.channel
20-
prompts = f"> **{prompt}** - <@{str(interaction.user.id)}> (***BingImageCreator***)\n\n"
24+
prompts = f"> **{prompt}** - <@{str(user_id)}> (***BingImageCreator***)\n\n"
2125

22-
logger.info(f"\x1b[31m{username}\x1b[0m'{prompt}' ({channel}) [BingImageCreator]")
26+
logger.info(f"\x1b[31m{username}\x1b[0m'{prompt}' ({channel}) [BingImageCreator]")
2327

24-
# Fetches image links
2528
async_gen = ImageGenAsync(auth_cookie=auth_cookie, quiet=True)
26-
images = await async_gen.get_images(prompt=prompt, timeout=300)
29+
images = await async_gen.get_images(prompt=prompt, timeout=int(os.getenv("IMAGE_TIMEOUT")), max_generate_time_sec=int(os.getenv("IMAGE_MAX_CREATE_SEC")))
2730
images = [file for file in images if not file.endswith('.svg')]
2831
new_image = await concatenate_images(images)
2932
image_data = BytesIO()
@@ -32,11 +35,11 @@ async def create_image(interaction: discord.Interaction, users_chatbot: dict, pr
3235

3336
await interaction.followup.send(prompts, file=discord.File(fp=image_data, filename='new_image.png'), view=ButtonView(interaction, prompt, images, users_chatbot, user_id))
3437
except asyncio.TimeoutError:
35-
await interaction.followup.send("> **Error: Request timed out.**")
36-
logger.error("Error while create image: Request timed out.")
38+
await interaction.followup.send("> **ErrorRequest timed out.**")
39+
logger.error("Error while create imageRequest timed out.")
3740
except Exception as e:
38-
await interaction.followup.send(f"> **Error: {e}**")
39-
logger.error(f"Error while create image: {e}")
41+
await interaction.followup.send(f"> **Error{e}**")
42+
logger.error(f"Error while create image{e}")
4043

4144
async def concatenate_images(image_urls):
4245
images = [Image.open(BytesIO(requests.get(url).content)) for url in image_urls]

0 commit comments

Comments
 (0)