src/main.py

@@ -1,45 +1,65 @@
+"""Discord Bot zur Bildgenerierung mit NovelAI API."""
+
 import io
 import os
 import sys
 import random
 import zipfile
 import logging
+import time
 import discord
 from discord import app_commands
 from discord.ext import commands
 import requests
 
+# ENV Variablen
 NOVELAI_API_TOKEN = os.getenv("NOVELAI_API_TOKEN")
 BOT_TOKEN = os.getenv("DISCORD_BOT_TOKEN")
-ALLOWED_CHANNEL_ID = int(os.getenv("DISCORD_CHANNEL_ID", "0"))  # falls nicht gesetzt → 0
+ALLOWED_CHANNEL_ID = int(os.getenv("DISCORD_CHANNEL_ID", "0"))
 
+# Logging Setup
 logging.basicConfig(
     stream=sys.stdout,
     level=logging.INFO,
     format="%(asctime)s - %(levelname)s - %(message)s"
 )
-
 logger = logging.getLogger(__name__)
 
 if not all([NOVELAI_API_TOKEN, BOT_TOKEN, ALLOWED_CHANNEL_ID]):
-    raise ValueError("Environment-Variablen NOVELAI_API_TOKEN, DISCORD_BOT_TOKEN und DISCORD_CHANNEL_ID müssen gesetzt sein.")
+    raise ValueError(
+        "NOVELAI_API_TOKEN, DISCORD_BOT_TOKEN und DISCORD_CHANNEL_ID müssen gesetzt sein."
+    )
 
 intents = discord.Intents.default()
-bot = commands.Bot(command_prefix=commands.when_mentioned, intents=intents)
+bot = commands.Bot(
+    command_prefix=commands.when_mentioned,
+    intents=intents
+)
+
 
 @bot.event
 async def on_ready():
-    print(f"Bot started as {bot.user}")
+    """Bot Startup Ereignis."""
+    print(f"Bot gestartet als {bot.user}")
     try:
         synced = await bot.tree.sync()
-        logging.info(f"{len(synced)} Slash-Commands synchronized.")
+        logger.info("%d Slash-Commands synchronisiert.", len(synced))
-    except Exception as e:
+    except Exception as err:
-        logging.error(e)
+        logger.error("Fehler beim Synchronisieren der Commands: %s", err)
 
-    activity = discord.Game(name="generating juicy NovelAI images 🥵")
+    activity = discord.Game(
-    await bot.change_presence(status=discord.Status.online, activity=activity)
+        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)",
@@ -53,31 +73,39 @@ async def generate(
     orientation: str = "portrait",
     seed: int = None
 ):
-    
+    """Slash-Command zur Bildgenerierung."""
     default_negative = (
-        "blurry, lowres, error, film grain, scan artifacts, worst quality, bad quality, jpeg artifacts, "
+        "blurry, lowres, error, film grain, scan artifacts, worst quality, "
-        "very displeasing, chromatic aberration, multiple views, logo, too many watermarks, white blank page, "
+        "bad quality, jpeg artifacts, very displeasing, chromatic aberration, "
-        "blank page, watermarks, watermark, text, "
+        "multiple views, logo, too many watermarks, white blank page, blank page, "
+        "watermarks, watermark, text, "
     )
-    
+
     if interaction.channel.id != ALLOWED_CHANNEL_ID:
         await interaction.response.send_message(
             "This command isn't allowed here.",
             ephemeral=True
         )
         return
-    
-    MAX_PROMPT_LENGTH = 500
 
-    if len(prompt) > MAX_PROMPT_LENGTH:
+    max_prompt_length = 500
-        await interaction.followup.send(f"Prompt too long! Maximum {MAX_PROMPT_LENGTH} chars allowed.")
+    negative_prompt = default_negative + undesired_prompt.strip()
+
+    if len(prompt) > max_prompt_length:
+        await interaction.response.send_message(
+            f"Prompt too long! Max {max_prompt_length} characters.", ephemeral=True
+        )
         return
 
-    if len(default_negative + undesired_prompt.strip()) > MAX_PROMPT_LENGTH:
+    if len(negative_prompt) > max_prompt_length:
-        await interaction.followup.send(f"Negative Prompt too long! Maximum {MAX_PROMPT_LENGTH} chars allowed.")
+        await interaction.response.send_message(
+            f"Negative prompt too long! Max {max_prompt_length} characters.",
+            ephemeral=True
+        )
         return
 
     await interaction.response.defer(thinking=True)
+
     logger.info(
         "User %s (%s, %s) requested image",
         interaction.user.display_name,
@@ -85,16 +113,11 @@ async def generate(
         interaction.user.id
     )
 
-    # Auflösung wählen
+    width, height = (1216, 832) if orientation.lower() == "landscape" else (832, 1216)
-    if orientation.lower() == "landscape":
-        width, height = 1216, 832
-    else:
-        width, height = 832, 1216
 
     if seed is None:
         seed = random.randint(0, 2**32 - 1)
 
-    # API-Payload zusammenbauen
     payload = {
         "input": prompt,
         "model": "nai-diffusion-4-full",
@@ -133,12 +156,12 @@ async def generate(
             },
             "v4_negative_prompt": {
                 "caption": {
-                    "base_caption": default_negative + undesired_prompt.strip(),
+                    "base_caption": negative_prompt,
                     "char_captions": []
                 },
                 "legacy_uc": False
             },
-            "negative_prompt": default_negative + undesired_prompt.strip(),
+            "negative_prompt": negative_prompt,
             "deliberate_euler_ancestral_bug": False,
             "prefer_brownian": True
         }
@@ -149,64 +172,79 @@ async def generate(
         "Content-Type": "application/json"
     }
 
+    start_time = time.monotonic()
     response = requests.post(
         "https://image.novelai.net/ai/generate-image",
         json=payload,
         headers=headers,
         timeout=120
     )
+    duration = time.monotonic() - start_time
+    logger.info("Bildgenerierung dauerte %.2f Sekunden", duration)
 
     if response.status_code == 200:
         response_bytes = response.content
+        param_text = (
+            f"```Prompt: {prompt}\n"
+            f"Undesired prompt: {negative_prompt}\n"
+            f"Seed: {seed}\n"
+            f"Resolution: {width}x{height}\n"
+            f"Sampler: k_euler_ancestral\n"
+            f"Steps: 28\n"
+            f"Scale: 3.7\n"
+            f"Model: nai-diffusion-4-full```"
+        )
 
-        param_text = f"""```Prompt: {prompt}
-Undesired prompt: {default_negative + undesired_prompt.strip()}
-Seed: {seed}
-Resolution: {width}x{height}
-Sampler: k_euler_ancestral
-Steps: 28
-Scale: 3.7
-Model: nai-diffusion-4-full```"""
-
-        # Prüfen ob ZIP (PK am Anfang)
         if response_bytes[:2] == b'PK':
-            # ZIP-Archiv entpacken
             with zipfile.ZipFile(io.BytesIO(response_bytes)) as zip_file:
                 namelist = zip_file.namelist()
-                # Erstes PNG suchen
+                image_name = next(
-                image_name = next((name for name in namelist if name.endswith(".png")), None)
+                    (name for name in namelist if name.endswith(".png")), None
-
+                )
                 if image_name:
                     image_bytes = zip_file.read(image_name)
                     filename = "novelai.png"
-                    file = discord.File(io.BytesIO(image_bytes), filename=filename)
+                    file = discord.File(
-                    await interaction.followup.send(content=param_text, file=file)
+                        io.BytesIO(image_bytes),
-                else:
+                        filename=filename
-                    await interaction.followup.send(
-                        "Found no PNG inside archive"
                     )
-
+                    await interaction.followup.send(
-        elif response_bytes[:4] == b'\x89PNG':
+                        content=param_text,
-            # Direkte PNG-Rückgabe
+                        file=file
-            filename = "novelai.png"
+                    )
-            file = discord.File(io.BytesIO(response_bytes), filename=filename)
+                    return
-            await interaction.followup.send(content=param_text, file=file)
-
-        else:
             await interaction.followup.send(
-                "API didn't send any file"
+                "Found no PNG inside archive"
             )
+            return
 
-    else:
+        if response_bytes[:4] == b'\x89PNG':
-        try:
+            filename = "novelai.png"
-            error_data = response.json()
+            file = discord.File(
-        except Exception:
+                io.BytesIO(response_bytes),
-            error_data = {"error": response.text}
+                filename=filename
+            )
+            await interaction.followup.send(
+                content=param_text,
+                file=file
+            )
+            return
 
-        error_message = f"Error {response.status_code} at API-request.\n"
+        await interaction.followup.send(
-        for key, value in error_data.items():
+            "API didn't send any file"
-            error_message += f"**{key}**: {value}\n"
+        )
+        return
+
+    try:
+        error_data = response.json()
+    except Exception:
+        error_data = {"error": response.text}
+
+    error_message = f"Error {response.status_code} at API-request.\n"
+    for key, value in error_data.items():
+        error_message += f"**{key}**: {value}\n"
+
+    await interaction.followup.send(error_message)
 
-        await interaction.followup.send(error_message)
 
 bot.run(BOT_TOKEN)