diff --git a/src/main.py b/src/main.py index 6f3a290..09ed1f2 100644 --- a/src/main.py +++ b/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.") - except Exception as e: - logging.error(e) + logger.info("%d Slash-Commands synchronisiert.", len(synced)) + except Exception as err: + 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)", @@ -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, " - "very displeasing, chromatic aberration, multiple views, logo, too many watermarks, white blank page, " - "blank page, watermarks, watermark, text, " + "blurry, lowres, error, film grain, scan artifacts, worst quality, " + "bad quality, jpeg artifacts, very displeasing, chromatic aberration, " + "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: - await interaction.followup.send(f"Prompt too long! Maximum {MAX_PROMPT_LENGTH} chars allowed.") + max_prompt_length = 500 + 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: - await interaction.followup.send(f"Negative Prompt too long! Maximum {MAX_PROMPT_LENGTH} chars allowed.") + if len(negative_prompt) > max_prompt_length: + 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 - if orientation.lower() == "landscape": - width, height = 1216, 832 - else: - width, height = 832, 1216 + width, height = (1216, 832) if orientation.lower() == "landscape" else (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((name for name in namelist if name.endswith(".png")), None) - + image_name = next( + (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) - await interaction.followup.send(content=param_text, file=file) - else: - await interaction.followup.send( - "Found no PNG inside archive" + file = discord.File( + io.BytesIO(image_bytes), + filename=filename ) - - elif response_bytes[:4] == b'\x89PNG': - # Direkte PNG-Rückgabe - filename = "novelai.png" - file = discord.File(io.BytesIO(response_bytes), filename=filename) - await interaction.followup.send(content=param_text, file=file) - - else: + await interaction.followup.send( + content=param_text, + file=file + ) + return await interaction.followup.send( - "API didn't send any file" + "Found no PNG inside archive" ) + return - else: - try: - error_data = response.json() - except Exception: - error_data = {"error": response.text} + 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 + ) + return - 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( + "API didn't send any file" + ) + 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)