till
cfe5d36c70
- app/ai_edit.py: OpenRouter chat/completions (image in/out), Presets + Prompt, Retry - bot: AI-Preset-Buttons + Bildunterschrift = freier AI-Prompt; _process refaktoriert - Freistellen/Upscaling bleiben lokal (originaltreu); AI-Edit fuer generative Edits - httpx pin, config OPENROUTER_* Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01XNQ8ghPfzAfsyVYd6HgFb6
387 lines
14 KiB
Python
387 lines
14 KiB
Python
"""Klarbildbot – Telegram-Bot fuer AI-Bildaufbereitung (Freistellen + Upscaling)."""
|
||
from __future__ import annotations
|
||
|
||
import asyncio
|
||
import logging
|
||
import posixpath
|
||
from concurrent.futures import ThreadPoolExecutor
|
||
from pathlib import Path
|
||
|
||
from telegram import (
|
||
InlineKeyboardButton,
|
||
InlineKeyboardMarkup,
|
||
InputFile,
|
||
Update,
|
||
)
|
||
from telegram.constants import ChatAction
|
||
from telegram.ext import (
|
||
Application,
|
||
CallbackQueryHandler,
|
||
CommandHandler,
|
||
ContextTypes,
|
||
MessageHandler,
|
||
filters,
|
||
)
|
||
|
||
from . import ai_edit, config, enhance, picdrop
|
||
|
||
logging.basicConfig(
|
||
format="%(asctime)s %(levelname)s %(name)s | %(message)s",
|
||
level=logging.INFO,
|
||
)
|
||
log = logging.getLogger("klarbildbot")
|
||
|
||
# Schwere CPU-Jobs strikt serialisieren (Speicherschutz auf kleinen Boxen).
|
||
EXECUTOR = ThreadPoolExecutor(max_workers=1)
|
||
|
||
# Kurzzeit-Speicher: pending-Bild pro User (file_id + Herkunft).
|
||
PENDING: dict[int, dict] = {}
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Zugriffsschutz
|
||
# ---------------------------------------------------------------------------
|
||
def _allowed(update: Update) -> bool:
|
||
if not config.ALLOWED_USER_IDS:
|
||
return True
|
||
user = update.effective_user
|
||
return bool(user and user.id in config.ALLOWED_USER_IDS)
|
||
|
||
|
||
async def _deny(update: Update) -> None:
|
||
uid = update.effective_user.id if update.effective_user else "?"
|
||
await update.effective_message.reply_text(
|
||
f"⛔️ Kein Zugriff. Deine Telegram-ID: {uid}\n"
|
||
"Bitte vom Betreiber in ALLOWED_USER_IDS freischalten lassen."
|
||
)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Commands
|
||
# ---------------------------------------------------------------------------
|
||
async def cmd_start(update: Update, ctx: ContextTypes.DEFAULT_TYPE) -> None:
|
||
uid = update.effective_user.id if update.effective_user else "?"
|
||
if not _allowed(update):
|
||
return await _deny(update)
|
||
ai_block = ""
|
||
if config.AI_EDIT_ENABLED:
|
||
ai_block = (
|
||
"• 🎨 *Studio-Hintergrund* / 🪄 *AI-Retusche* / 🖼 *Restaurierung*\n"
|
||
"• 💬 *Eigener AI-Prompt*: Bild mit Bildunterschrift senden "
|
||
"(z. B. 'Hintergrund zu Strand', 'mach es schwarz-weiß')\n"
|
||
)
|
||
await update.effective_message.reply_text(
|
||
"👋 *Klarbildbot* – AI-Bildaufbereitung\n\n"
|
||
"Schick mir ein Bild, dann kannst du wählen:\n"
|
||
"• 🔍 *Klarbild ×2 / ×4* – Upscaling + Schärfen (lokal, originaltreu)\n"
|
||
"• ✂️ *Freistellen* – Hintergrund entfernen, echtes transparentes PNG (lokal)\n"
|
||
"• ✨ *Freistellen + Klarbild* – beides\n"
|
||
+ ai_block +
|
||
"\n💡 Für beste Qualität das Bild als *Datei* senden "
|
||
"(Büroklammer → Datei), nicht als komprimiertes Foto.\n\n"
|
||
f"Deine Telegram-ID: `{uid}`",
|
||
parse_mode="Markdown",
|
||
)
|
||
|
||
|
||
async def cmd_help(update: Update, ctx: ContextTypes.DEFAULT_TYPE) -> None:
|
||
if not _allowed(update):
|
||
return await _deny(update)
|
||
txt = (
|
||
"*So geht's:*\n"
|
||
"1. Bild senden (am besten als Datei)\n"
|
||
"2. Aktion wählen\n"
|
||
"3. Ergebnis kommt als Datei zurück\n\n"
|
||
"*Befehle:*\n"
|
||
"/start – Übersicht\n"
|
||
"/help – diese Hilfe\n"
|
||
)
|
||
if config.PICDROP_ENABLED:
|
||
txt += "/picdrop – Picdrop-Galerie stapelweise verarbeiten\n"
|
||
await update.effective_message.reply_text(txt, parse_mode="Markdown")
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Bild empfangen -> Aktionsmenue
|
||
# ---------------------------------------------------------------------------
|
||
def _menu() -> InlineKeyboardMarkup:
|
||
rows = [
|
||
[InlineKeyboardButton("🔍 Klarbild ×2", callback_data="up:2"),
|
||
InlineKeyboardButton("🔍 Klarbild ×4", callback_data="up:4")],
|
||
[InlineKeyboardButton("✂️ Freistellen", callback_data="cut:0")],
|
||
[InlineKeyboardButton("✨ Freistellen + Klarbild ×2", callback_data="cutup:2")],
|
||
]
|
||
if config.AI_EDIT_ENABLED:
|
||
rows += [
|
||
[InlineKeyboardButton("🎨 Studio-Hintergrund", callback_data="ai:studio"),
|
||
InlineKeyboardButton("🪄 AI-Retusche", callback_data="ai:retouch")],
|
||
[InlineKeyboardButton("🖼 Foto-Restaurierung", callback_data="ai:restore")],
|
||
]
|
||
return InlineKeyboardMarkup(rows)
|
||
|
||
|
||
def _label(action: str, scale: int = 0, prompt: str = "") -> str:
|
||
return {
|
||
"up": f"Klarbild ×{scale}",
|
||
"cut": "Freistellen",
|
||
"cutup": f"Freistellen + Klarbild ×{scale}",
|
||
"ai": ai_edit.PRESET_LABELS.get(prompt, prompt),
|
||
"aiprompt": f"AI-Edit: {prompt[:40]}",
|
||
}.get(action, action)
|
||
|
||
|
||
def _do_job(action: str, scale: int, prompt: str, image_bytes: bytes) -> enhance.Result:
|
||
if action == "up":
|
||
return enhance.upscale(image_bytes, scale)
|
||
if action == "cut":
|
||
return enhance.cutout(image_bytes)
|
||
if action == "cutup":
|
||
return enhance.cutout_then_upscale(image_bytes, scale)
|
||
if action == "ai":
|
||
return ai_edit.edit(image_bytes, ai_edit.PRESETS[prompt])
|
||
if action == "aiprompt":
|
||
return ai_edit.edit(image_bytes, prompt)
|
||
raise ValueError(action)
|
||
|
||
|
||
async def _process(ctx: ContextTypes.DEFAULT_TYPE, chat_id: int, status,
|
||
file_id: str, action: str, scale: int, prompt: str) -> None:
|
||
"""Bild laden, Job im Executor ausfuehren, Ergebnis als Dokument senden.
|
||
`status` ist eine Telegram-Message, deren Text als Fortschritt editiert wird."""
|
||
label = _label(action, scale, prompt)
|
||
await status.edit_text(f"⏳ {label} läuft … (kann bei großen Bildern dauern)")
|
||
try:
|
||
tg_file = await ctx.bot.get_file(file_id)
|
||
image_bytes = bytes(await tg_file.download_as_bytearray())
|
||
except Exception as exc:
|
||
log.exception("Download fehlgeschlagen")
|
||
return await status.edit_text(
|
||
f"❌ Konnte das Bild nicht laden: {exc}\n"
|
||
"(Telegram-Bots können Dateien bis 20 MB laden.)")
|
||
|
||
await ctx.bot.send_chat_action(chat_id, ChatAction.UPLOAD_DOCUMENT)
|
||
loop = asyncio.get_running_loop()
|
||
try:
|
||
result = await loop.run_in_executor(
|
||
EXECUTOR, _do_job, action, scale, prompt, image_bytes)
|
||
except Exception as exc:
|
||
log.exception("Verarbeitung fehlgeschlagen")
|
||
return await status.edit_text(f"❌ Fehler bei der Verarbeitung: {exc}")
|
||
|
||
caption = (f"✅ {label}\n{result.width}×{result.height}px · {result.seconds:.1f}s"
|
||
+ (f"\n{result.note}" if result.note else ""))
|
||
try:
|
||
await ctx.bot.send_document(
|
||
chat_id=chat_id,
|
||
document=InputFile(result.data, filename=f"klarbild.{result.fmt}"),
|
||
caption=caption,
|
||
)
|
||
await status.edit_text(f"✅ Fertig: {label}")
|
||
except Exception as exc:
|
||
log.exception("Senden fehlgeschlagen")
|
||
await status.edit_text(f"❌ Konnte Ergebnis nicht senden: {exc}")
|
||
|
||
|
||
async def on_image(update: Update, ctx: ContextTypes.DEFAULT_TYPE) -> None:
|
||
if not _allowed(update):
|
||
return await _deny(update)
|
||
msg = update.effective_message
|
||
|
||
if msg.photo:
|
||
file_id = msg.photo[-1].file_id
|
||
source = "photo"
|
||
elif msg.document and (msg.document.mime_type or "").startswith("image/"):
|
||
file_id = msg.document.file_id
|
||
source = "document"
|
||
else:
|
||
return await msg.reply_text("Bitte ein Bild senden.")
|
||
|
||
# Bild mit Bildunterschrift + AI aktiv => Unterschrift als freier AI-Prompt
|
||
caption = (msg.caption or "").strip()
|
||
if caption and config.AI_EDIT_ENABLED:
|
||
status = await msg.reply_text("🪄 AI-Edit vorbereiten …")
|
||
return await _process(ctx, msg.chat_id, status, file_id, "aiprompt", 0, caption)
|
||
|
||
PENDING[update.effective_user.id] = {"file_id": file_id, "source": source}
|
||
hint = "" if source == "document" else "\n_(Tipp: als Datei senden = volle Auflösung)_"
|
||
if config.AI_EDIT_ENABLED:
|
||
hint += "\n💬 _Oder Bild mit Bildunterschrift senden = eigener AI-Prompt._"
|
||
await msg.reply_text(
|
||
"Was soll ich damit machen?" + hint,
|
||
reply_markup=_menu(),
|
||
parse_mode="Markdown",
|
||
)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Button gedrueckt -> verarbeiten
|
||
# ---------------------------------------------------------------------------
|
||
async def on_action(update: Update, ctx: ContextTypes.DEFAULT_TYPE) -> None:
|
||
query = update.callback_query
|
||
await query.answer()
|
||
if not _allowed(update):
|
||
return await query.edit_message_text("⛔️ Kein Zugriff.")
|
||
|
||
uid = update.effective_user.id
|
||
pend = PENDING.get(uid)
|
||
if not pend:
|
||
return await query.edit_message_text(
|
||
"Kein Bild gemerkt – bitte schick es nochmal.")
|
||
|
||
action, _, arg = query.data.partition(":")
|
||
scale = int(arg) if arg.isdigit() else 0
|
||
prompt = arg if action == "ai" else "" # bei AI-Presets ist arg der Preset-Key
|
||
|
||
try:
|
||
await _process(ctx, update.effective_chat.id, query.message,
|
||
pend["file_id"], action, scale, prompt)
|
||
finally:
|
||
PENDING.pop(uid, None)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Picdrop-Batch
|
||
# ---------------------------------------------------------------------------
|
||
async def cmd_picdrop(update: Update, ctx: ContextTypes.DEFAULT_TYPE) -> None:
|
||
if not _allowed(update):
|
||
return await _deny(update)
|
||
if not config.PICDROP_ENABLED:
|
||
return await update.effective_message.reply_text(
|
||
"Picdrop ist nicht konfiguriert (PICDROP_USER / PICDROP_PASSWORD).")
|
||
|
||
args = ctx.args or []
|
||
loop = asyncio.get_running_loop()
|
||
|
||
if not args:
|
||
try:
|
||
entries = await loop.run_in_executor(EXECUTOR, picdrop.list_dir, "/")
|
||
except Exception as exc:
|
||
return await update.effective_message.reply_text(f"❌ SFTP-Fehler: {exc}")
|
||
folders = [e["name"] for e in entries if e["is_dir"]][:40]
|
||
listing = "\n".join(f"• {f}" for f in folders) or "(keine Ordner)"
|
||
return await update.effective_message.reply_text(
|
||
"*Picdrop – Ordner im Wurzelverzeichnis:*\n" + listing +
|
||
"\n\nNutzung: `/picdrop <Ordnername-Teil> up2|up4|cut`",
|
||
parse_mode="Markdown")
|
||
|
||
needle = args[0]
|
||
op = (args[1] if len(args) > 1 else "up2").lower()
|
||
|
||
await update.effective_message.reply_text(f"🔎 Suche Galerie „{needle}“ …")
|
||
try:
|
||
matches = await loop.run_in_executor(EXECUTOR, picdrop.find_gallery, needle)
|
||
except Exception as exc:
|
||
return await update.effective_message.reply_text(f"❌ SFTP-Fehler: {exc}")
|
||
if not matches:
|
||
return await update.effective_message.reply_text("Keine passende Galerie gefunden.")
|
||
|
||
gallery = matches[0]
|
||
images = await loop.run_in_executor(EXECUTOR, picdrop.list_images, gallery)
|
||
if not images:
|
||
return await update.effective_message.reply_text(
|
||
f"Galerie `{gallery}` enthält keine Bilder.", parse_mode="Markdown")
|
||
|
||
await update.effective_message.reply_text(
|
||
f"📁 `{gallery}` · {len(images)} Bilder · Aktion `{op}`\n"
|
||
"Verarbeite und lade in Unterordner `klarbild/` hoch …",
|
||
parse_mode="Markdown")
|
||
|
||
def process_one(remote_path: str) -> str:
|
||
data = picdrop.download(remote_path)
|
||
if op == "cut":
|
||
res = enhance.cutout(data)
|
||
elif op == "up4":
|
||
res = enhance.upscale(data, 4)
|
||
else:
|
||
res = enhance.upscale(data, 2)
|
||
base = posixpath.splitext(posixpath.basename(remote_path))[0]
|
||
out_name = f"{base}_klar.{res.fmt}"
|
||
out_path = posixpath.join(gallery, "klarbild", out_name)
|
||
picdrop.upload(out_path, res.data)
|
||
return out_name
|
||
|
||
done = 0
|
||
for img in images:
|
||
try:
|
||
name = await loop.run_in_executor(EXECUTOR, process_one, img["path"])
|
||
done += 1
|
||
if done % 5 == 0 or done == len(images):
|
||
await update.effective_message.reply_text(
|
||
f"… {done}/{len(images)} fertig (zuletzt: {name})")
|
||
except Exception as exc:
|
||
log.exception("Picdrop-Bild fehlgeschlagen")
|
||
await update.effective_message.reply_text(
|
||
f"⚠️ {img['name']}: {exc}")
|
||
|
||
await update.effective_message.reply_text(
|
||
f"✅ Fertig: {done}/{len(images)} Bilder → `{gallery}/klarbild/`",
|
||
parse_mode="Markdown")
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Health-Server (fuer Coolify) + App-Setup
|
||
# ---------------------------------------------------------------------------
|
||
async def _start_health(app: Application) -> None:
|
||
from aiohttp import web
|
||
|
||
async def health(_req):
|
||
return web.json_response({"status": "ok", "service": "klarbildbot"})
|
||
|
||
server = web.Application()
|
||
server.router.add_get("/", health)
|
||
server.router.add_get("/health", health)
|
||
runner = web.AppRunner(server)
|
||
await runner.setup()
|
||
site = web.TCPSite(runner, "0.0.0.0", config.HEALTH_PORT)
|
||
await site.start()
|
||
app.bot_data["health_runner"] = runner
|
||
log.info("Health-Server auf :%s", config.HEALTH_PORT)
|
||
|
||
|
||
async def _post_init(app: Application) -> None:
|
||
await _start_health(app)
|
||
# Bot-Kommandos in Telegram-UI registrieren
|
||
from telegram import BotCommand
|
||
|
||
cmds = [BotCommand("start", "Übersicht"), BotCommand("help", "Hilfe")]
|
||
if config.PICDROP_ENABLED:
|
||
cmds.append(BotCommand("picdrop", "Picdrop-Galerie verarbeiten"))
|
||
await app.bot.set_my_commands(cmds)
|
||
# Modelle vorwaermen (im Thread, blockiert den Loop nicht)
|
||
asyncio.get_running_loop().run_in_executor(EXECUTOR, enhance.warmup)
|
||
log.info("Klarbildbot bereit.")
|
||
|
||
|
||
async def _post_shutdown(app: Application) -> None:
|
||
runner = app.bot_data.get("health_runner")
|
||
if runner:
|
||
await runner.cleanup()
|
||
|
||
|
||
def build_app() -> Application:
|
||
config.validate()
|
||
app = (
|
||
Application.builder()
|
||
.token(config.BOT_TOKEN)
|
||
.post_init(_post_init)
|
||
.post_shutdown(_post_shutdown)
|
||
.build()
|
||
)
|
||
app.add_handler(CommandHandler("start", cmd_start))
|
||
app.add_handler(CommandHandler("help", cmd_help))
|
||
app.add_handler(CommandHandler("picdrop", cmd_picdrop))
|
||
app.add_handler(MessageHandler(
|
||
filters.PHOTO | filters.Document.IMAGE, on_image))
|
||
app.add_handler(CallbackQueryHandler(on_action))
|
||
return app
|
||
|
||
|
||
def main() -> None:
|
||
app = build_app()
|
||
log.info("Starte Polling …")
|
||
app.run_polling(allowed_updates=Update.ALL_TYPES, drop_pending_updates=True)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|