Merge pull request 'dev' (#24) from dev into main

Reviewed-on: #24

.gitea/workflows/check-code_quality.yaml

@@ -15,12 +15,12 @@ jobs:
       GITHUB: ${{ secrets.GITHUB }}
     steps:
       - name: Checkout code
-        uses: actions/checkout@v4
+        uses: actions/checkout@v6
 
       - name: Set up Python
-        uses: actions/setup-python@v5
+        uses: actions/setup-python@v6
         with:
-          python-version: '3.10'
+          python-version: '3.14'
 
       - name: Install dependencies
         run: |
@@ -85,7 +85,7 @@ jobs:
       
       # Ergebnisse als Artefakte speichern
       - name: Upload results as artifacts
-        uses: actions/upload-artifact@v3
+        uses: actions/upload-artifact@v5
         with:
           name: quality-reports
           path: /tmp/results

.gitea/workflows/release.yaml

@@ -10,7 +10,7 @@ jobs:
     runs-on: ubuntu-latest
     steps:
       - name: Checkout Repository
-        uses: actions/checkout@v4
+        uses: actions/checkout@v6
 
       - name: Verify Docker installation
         run: docker --version
@@ -20,15 +20,15 @@ jobs:
 
       - name: Build Docker Image
         run: |
-          docker build -f src/Dockerfile -t git.0day.agency/${{ secrets.REGISTRY_USERNAME }}/novelai-discord-bot:${{ gitea.ref_name }} src/
+          docker build -f src/Dockerfile -t git.0day.agency/${{ secrets.REGISTRY_USERNAME }}/novelai-bot:${{ gitea.ref_name }} src/
           # latest-Tag setzen
-          docker tag git.0day.agency/${{ secrets.REGISTRY_USERNAME }}/novelai-discord-bot:${{ gitea.ref_name }} \
-                     git.0day.agency/${{ secrets.REGISTRY_USERNAME }}/novelai-discord-bot:latest
+          docker tag git.0day.agency/${{ secrets.REGISTRY_USERNAME }}/novelai-bot:${{ gitea.ref_name }} \
+                     git.0day.agency/${{ secrets.REGISTRY_USERNAME }}/novelai-bot:latest
 
       - name: Push Docker Image mit Versionstag
         run: |
-          docker push git.0day.agency/${{ secrets.REGISTRY_USERNAME }}/novelai-discord-bot:${{ gitea.ref_name }}
+          docker push git.0day.agency/${{ secrets.REGISTRY_USERNAME }}/novelai-bot:${{ gitea.ref_name }}
 
       - name: Push Docker Image als latest
         run: |
-          docker push git.0day.agency/${{ secrets.REGISTRY_USERNAME }}/novelai-discord-bot:latest
+          docker push git.0day.agency/${{ secrets.REGISTRY_USERNAME }}/novelai-bot:latest

.github/renovate.json

@@ -0,0 +1,20 @@
+{
+    "$schema": "https://docs.renovatebot.com/renovate-schema.json",
+    "extends": [
+      "config:recommended"
+    ],
+    "baseBranches": [
+      "dev"
+    ],
+    "packageRules": [
+      {
+        "matchUpdateTypes": [
+          "minor",
+          "patch",
+          "pin",
+          "digest"
+        ],
+        "automerge": false
+      }
+    ]
+  }

README.md

@@ -1,2 +1,52 @@
-# novelai-bot
+# NovelAI Discord Bot
 
+![Code Quality](https://git.0day.agency/wirehack7/novelai-bot/actions/workflows/check-code_quality.yaml/badge.svg)
+![Safety Check](https://git.0day.agency/wirehack7/novelai-bot/actions/workflows/check-safety.yaml/badge.svg)
+![Release](https://git.0day.agency/wirehack7/novelai-bot/actions/workflows/release.yaml/badge.svg)
+
+Simple bot which registers a command to generate images via the NovelAI API.
+
+## Usage
+
+Just enter `/generate` in the desired channel.
+
+```text
+/generate 
+    *prompt: string
+    undesired_prompt: string
+    orientation: landscape or portrait
+    seed: integer
+```
+
+`*: required`
+
+## Run
+
+You might create environvent vars:
+
+- **NOVELAI_API_TOKEN**: your NovelAI user API key
+- **DISCORD_BOT_TOKEN**: token of your Discord bot
+- **DISCORD_CHANNEL_ID**: ID of the channel to enable this bot
+
+Don't forget to install the requirements in [requirements.txt](src/requirements.txt)
+
+## Docker Compose
+
+Or simply run it via **docker-compose**:
+
+```yaml
+services:
+  bot:
+    restart: unless-stopped
+    image: git.0day.agency/wirehack7/novelai-bot:latest
+    environment:
+      - NOVELAI_API_TOKEN=
+      - DISCORD_BOT_TOKEN=
+      - DISCORD_CHANNEL_ID=
+networks: {}
+```
+
+## Notes
+
+To create a Discord bot go to [Discord Developers](https://discord.com/developers/).
+You should make the bot private.

src/Dockerfile

@@ -1,4 +1,4 @@
-FROM python:3.13.3-slim
+FROM python:3.14.2-slim
 
 WORKDIR /app
 

src/main.py

@@ -1,103 +1,115 @@
+"""Discord Bot zur Bildgenerierung mit NovelAI API."""
+
+# flake8: noqa E501
+
 import io
+import logging
 import os
 import random
+import sys
+import time
 import zipfile
-import logging
+
+import aiohttp
 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"))
 
-root = logging.getLogger()
-root.setLevel(logging.INFO)
-
-handler = logging.StreamHandler(sys.stdout)
-handler.setLevel(logging.DEBUG)
-formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
-handler.setFormatter(formatter)
-root.addHandler(handler)
+# 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 and DISCORD_CHANNEL_ID have to be set."
+    )
 
 intents = discord.Intents.default()
 bot = commands.Bot(command_prefix=commands.when_mentioned, intents=intents)
 
-async def get_available_antlas(api_token):
-    headers = {
-        "Authorization": f"Bearer {api_token}"
-    }
-    response = requests.get("https://api.novelai.net/user/info", headers=headers)
-    if response.status_code == 200:
-        data = response.json()
-        return data.get("imageAnlas", 0)
-    else:
-        logging.error(f"Fehler beim Abrufen der Antlas: {response.status_code}")
-        return None
 
 @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 synchronized.", len(synced))
+    except Exception as err:  # pylint: disable=W0718
+        logger.error("error while syncing the command: %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,too-many-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)
 
-    # Auflösung wählen
-    if orientation.lower() == "landscape":
-        width, height = 1216, 832
-    else:
-        width, height = 832, 1216
+    logger.info(
+        "User %s (%s, %s) requested image",
+        interaction.user.display_name,
+        interaction.user.name,
+        interaction.user.id,
+    )
+
+    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",
@@ -127,90 +139,88 @@ 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",
     }
 
-    response = requests.post(
-        "https://image.novelai.net/ai/generate-image",
-        json=payload,
-        headers=headers,
-        timeout=120
-    )
+    start_time = time.monotonic()
+    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}")
 
-    if response.status_code == 200:
-        antlas = get_available_antlas(NOVELAI_API_TOKEN)
-        logging.info("Generated image, current antlas: {antlas}")
-        response_bytes = response.content
+    duration = time.monotonic() - start_time
+    logger.info("Image creation took %.2f seconds", duration)
 
-        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```"""
+    if status_code == 200:
+        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```"
+        )
 
-        # 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_bytes.decode("utf-8")
+    except Exception:  # pylint: disable=W0718
+        error_data = "error message unreadable"
 
-        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 {status_code} at API-request.\n" f"Answer: {error_data}\n"
+        )
 
         await interaction.followup.send(error_message)
 

src/requirements.txt

@@ -1,2 +1,2 @@
 discord
-requests
+aiohttp