more fixes
This commit is contained in:
parent
1a76246f53
commit
8f1c68aa04
GPG Key ID:
C9F4CF85BA505ADC
168
src/main.py
168
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,11 +73,12 @@ 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:
|
||||
@ -67,17 +88,24 @@ async def generate(
|
||||
)
|
||||
return
|
||||
|
||||
MAX_PROMPT_LENGTH = 500
|
||||
max_prompt_length = 500
|
||||
negative_prompt = default_negative + undesired_prompt.strip()
|
||||
|
||||
if len(prompt) > MAX_PROMPT_LENGTH:
|
||||
await interaction.followup.send(f"Prompt too long! Maximum {MAX_PROMPT_LENGTH} chars allowed.")
|
||||
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)
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user