diff --git a/.gitea/workflows/release.yaml b/.gitea/workflows/release.yaml index e1a0af9..2788e4f 100644 --- a/.gitea/workflows/release.yaml +++ b/.gitea/workflows/release.yaml @@ -31,10 +31,4 @@ jobs: - name: Push Docker Image als latest run: | - docker push git.0day.agency/${{ secrets.REGISTRY_USERNAME }}/novelai-bot:latest - - - name: Publish Container Package in Gitea - run: | - curl -X POST \ - -H "Authorization: token ${{ secrets.REGISTRY_TOKEN }}" \ - "https://git.0day.agency/api/packages/${{ secrets.REGISTRY_USERNAME }}/container/novelai-discord-bot/versions/${{ gitea.ref_name }}/publish" \ No newline at end of file + docker push git.0day.agency/${{ secrets.REGISTRY_USERNAME }}/novelai-bot:latest \ No newline at end of file diff --git a/src/main.py b/src/main.py index fe60f91..0e20154 100644 --- a/src/main.py +++ b/src/main.py @@ -1,97 +1,115 @@ +"""Discord Bot zur Bildgenerierung mit NovelAI API.""" + +# flake8: noqa E501 + import io -import os -import sys -import random -import zipfile import logging +import os +import random +import sys +import time +import zipfile + import discord +import requests 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" + 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.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: # 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) + @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", ) +# pylint: disable=too-many-locals,too-many-return-statements 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 = ( - "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 + "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, interaction.user.name, - interaction.user.id + 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", @@ -121,89 +139,82 @@ 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": default_negative + undesired_prompt.strip(), - "char_captions": [] - }, - "legacy_uc": False + "caption": {"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 - } + "prefer_brownian": True, + }, } headers = { "Authorization": f"Bearer {NOVELAI_API_TOKEN}", - "Content-Type": "application/json" + "Content-Type": "application/json", } + start_time = time.monotonic() response = requests.post( "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) 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 + if response_bytes[:2] == b"PK": 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" - ) + return + await interaction.followup.send("Found no PNG inside archive") + return - elif response_bytes[:4] == b'\x89PNG': - # Direkte PNG-Rückgabe + 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 - else: - await interaction.followup.send( - "API didn't send any file" - ) + await interaction.followup.send("API didn't send any file") + return - else: - try: - error_data = response.json() - except Exception: - error_data = {"error": response.text} + try: + error_data = response.json() + except Exception: # pylint: disable=W0718 + 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" + 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)