novelai-bot/src/main.py

"""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 aiohttp
import discord
import requests
from discord import app_commands
from discord.ext import commands

# 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"))

# 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(
        "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)


@bot.event
async def on_ready():
    """Bot Startup Ereignis."""
    print(f"Bot gestartet als {bot.user}")
    try:
        synced = await bot.tree.sync()
        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",
)
# 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,
):
    """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, "
    )

    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
    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(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,
    )

    width, height = (1216, 832) if orientation.lower() == "landscape" else (832, 1216)

    if seed is None:
        seed = random.randint(0, 2**32 - 1)

    payload = {
        "input": prompt,
        "model": "nai-diffusion-4-full",
        "action": "generate",
        "parameters": {
            "params_version": 3,
            "width": width,
            "height": height,
            "scale": 3.7,
            "sampler": "k_euler_ancestral",
            "steps": 28,
            "n_samples": 1,
            "ucPreset": 0,
            "qualityToggle": True,
            "autoSmea": False,
            "dynamic_thresholding": False,
            "controlnet_strength": 1,
            "legacy": False,
            "add_original_image": True,
            "cfg_rescale": 0,
            "noise_schedule": "karras",
            "legacy_v3_extend": False,
            "skip_cfg_above_sigma": None,
            "use_coords": False,
            "legacy_uc": False,
            "normalize_reference_strength_multiple": True,
            "seed": seed,
            "characterPrompts": [],
            "v4_prompt": {
                "caption": {"base_caption": prompt, "char_captions": []},
                "use_coords": False,
                "use_order": True,
            },
            "v4_negative_prompt": {
                "caption": {"base_caption": negative_prompt, "char_captions": []},
                "legacy_uc": False,
            },
            "negative_prompt": negative_prompt,
            "deliberate_euler_ancestral_bug": False,
            "prefer_brownian": True,
        },
    }

    headers = {
        "Authorization": f"Bearer {NOVELAI_API_TOKEN}",
        "Content-Type": "application/json",
    }

    start_time = time.monotonic()

    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

    duration = time.monotonic() - start_time
    logger.info("Image creation took %.2f seconds", duration)

    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```"
        )

        if response_bytes[:2] == b"PK":
            with zipfile.ZipFile(io.BytesIO(response_bytes)) as zip_file:
                namelist = zip_file.namelist()
                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)
                    return
            await interaction.followup.send("Found no PNG inside archive")
            return

        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

        await interaction.followup.send("API didn't send any file")
        return

    try:
        error_data = response_bytes.decode("utf-8")
    except Exception:  # pylint: disable=W0718
        error_data = "error message unreadable"

    error_message = f"Error {status_code} at API-request.\n" f"Answer: {error_data}\n"

    await interaction.followup.send(error_message)


bot.run(BOT_TOKEN)