diff --git a/src/main.py b/src/main.py index ec7718d..c99f54a 100644 --- a/src/main.py +++ b/src/main.py @@ -58,7 +58,7 @@ async def on_ready(): orientation="portrait or landscape (Standard: portrait)", seed="Optional seed (integer). If empty, a random one will be generated", ) -# pylint: disable=too-many-locals,too-many-return-statements +# pylint: disable=too-many-locals,too-many-return-statements,too-many-statements async def generate( interaction: discord.Interaction, prompt: str, @@ -159,16 +159,20 @@ async def generate( } start_time = time.monotonic() - - async with aiohttp.ClientSession() as session: - async with session.post( - "https://image.novelai.net/ai/generate-image", - json=payload, - headers=headers, - timeout=aiohttp.ClientTimeout(total=120), - ) as response: - response_bytes = await response.read() - status_code = response.status + try: + async with aiohttp.ClientSession() as session: + async with session.post( + "https://image.novelai.net/ai/generate-image", + json=payload, + headers=headers, + timeout=aiohttp.ClientTimeout(total=120), + ) as response: + response_bytes = await response.read() + status_code = response.status + except aiohttp.TimeoutError: + await interaction.followup.send("Error: the request timed out.") + except aiohttp.ClientError as e: + await interaction.followup.send(f"Error: {e}") duration = time.monotonic() - start_time logger.info("Image creation took %.2f seconds", duration) @@ -214,9 +218,11 @@ async def generate( except Exception: # pylint: disable=W0718 error_data = "error message unreadable" - error_message = f"Error {status_code} at API-request.\n" f"Answer: {error_data}\n" + error_message = ( + f"Error {status_code} at API-request.\n" f"Answer: {error_data}\n" + ) - await interaction.followup.send(error_message) + await interaction.followup.send(error_message) bot.run(BOT_TOKEN)