diff --git a/src/main.py b/src/main.py index 254cb39..12f8b91 100644 --- a/src/main.py +++ b/src/main.py @@ -34,6 +34,7 @@ if not all([NOVELAI_API_TOKEN, BOT_TOKEN, ALLOWED_CHANNEL_ID]): intents = discord.Intents.default() bot = commands.Bot(command_prefix=commands.when_mentioned, intents=intents) + @bot.event async def on_ready(): """Bot Startup Ereignis.""" @@ -41,34 +42,26 @@ async def on_ready(): try: synced = await bot.tree.sync() logger.info("%d Slash-Commands synchronisiert.", len(synced)) - except Exception as err: + except Exception as err: # pylint: disable=W0718 logger.error("Fehler beim Synchronisieren der Commands: %s", err) - activity = discord.Game( - name="generating juicy NovelAI images 🥵" - ) - await bot.change_presence( - status=discord.Status.online, - activity=activity - ) + activity = discord.Game(name="generating juicy NovelAI images 🥵") + await bot.change_presence(status=discord.Status.online, activity=activity) -@bot.tree.command( - name="generate", - description="Generate image with NovelAI v4 Full" -) +@bot.tree.command(name="generate", description="Generate image with NovelAI v4 Full") @app_commands.describe( prompt="What should be generated?", undesired_prompt="What should be avoided? (optional)", orientation="portrait or landscape (Standard: portrait)", - seed="Optional seed (integer). If empty, a random one will be generated" + seed="Optional seed (integer). If empty, a random one will be generated", ) async def generate( interaction: discord.Interaction, prompt: str, undesired_prompt: str = "", orientation: str = "portrait", - seed: int = None + seed: int = None, ): """Slash-Command zur Bildgenerierung.""" default_negative = ( @@ -80,8 +73,7 @@ async def generate( if interaction.channel.id != ALLOWED_CHANNEL_ID: await interaction.response.send_message( - "This command isn't allowed here.", - ephemeral=True + "This command isn't allowed here.", ephemeral=True ) return @@ -97,7 +89,7 @@ async def generate( if len(negative_prompt) > max_prompt_length: await interaction.response.send_message( f"Negative prompt too long! Max {max_prompt_length} characters.", - ephemeral=True + ephemeral=True, ) return @@ -107,7 +99,7 @@ async def generate( "User %s (%s, %s) requested image", interaction.user.display_name, interaction.user.name, - interaction.user.id + interaction.user.id, ) width, height = (1216, 832) if orientation.lower() == "landscape" else (832, 1216) @@ -144,29 +136,23 @@ async def generate( "seed": seed, "characterPrompts": [], "v4_prompt": { - "caption": { - "base_caption": prompt, - "char_captions": [] - }, + "caption": {"base_caption": prompt, "char_captions": []}, "use_coords": False, - "use_order": True + "use_order": True, }, "v4_negative_prompt": { - "caption": { - "base_caption": negative_prompt, - "char_captions": [] - }, - "legacy_uc": False + "caption": {"base_caption": negative_prompt, "char_captions": []}, + "legacy_uc": False, }, "negative_prompt": negative_prompt, "deliberate_euler_ancestral_bug": False, - "prefer_brownian": True - } + "prefer_brownian": True, + }, } headers = { "Authorization": f"Bearer {NOVELAI_API_TOKEN}", - "Content-Type": "application/json" + "Content-Type": "application/json", } start_time = time.monotonic() @@ -174,7 +160,7 @@ async def generate( "https://image.novelai.net/ai/generate-image", json=payload, headers=headers, - timeout=120 + timeout=120, ) duration = time.monotonic() - start_time logger.info("Bildgenerierung dauerte %.2f Sekunden", duration) @@ -192,7 +178,7 @@ async def generate( f"Model: nai-diffusion-4-full```" ) - if response_bytes[:2] == b'PK': + if response_bytes[:2] == b"PK": with zipfile.ZipFile(io.BytesIO(response_bytes)) as zip_file: namelist = zip_file.namelist() image_name = next( @@ -201,40 +187,24 @@ async def generate( if image_name: image_bytes = zip_file.read(image_name) filename = "novelai.png" - file = discord.File( - io.BytesIO(image_bytes), - filename=filename - ) - await interaction.followup.send( - content=param_text, - file=file - ) + file = discord.File(io.BytesIO(image_bytes), filename=filename) + await interaction.followup.send(content=param_text, file=file) return - await interaction.followup.send( - "Found no PNG inside archive" - ) + await interaction.followup.send("Found no PNG inside archive") return - if response_bytes[:4] == b'\x89PNG': + if response_bytes[:4] == b"\x89PNG": filename = "novelai.png" - file = discord.File( - io.BytesIO(response_bytes), - filename=filename - ) - await interaction.followup.send( - content=param_text, - file=file - ) + file = discord.File(io.BytesIO(response_bytes), filename=filename) + await interaction.followup.send(content=param_text, file=file) return - await interaction.followup.send( - "API didn't send any file" - ) + await interaction.followup.send("API didn't send any file") return try: error_data = response.json() - except Exception: + except Exception: # pylint: disable=W0718 error_data = {"error": response.text} error_message = f"Error {response.status_code} at API-request.\n"