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