Initial commit — ComfyUI Discord bot + web UI
Full source for the-third-rev: Discord bot (discord.py), FastAPI web UI (React/TS/Vite/Tailwind), ComfyUI integration, generation history DB, preset manager, workflow inspector, and all supporting modules. Excluded from tracking: .env, invite_tokens.json, *.db (SQLite), current-workflow-changes.json, user_settings/, presets/, logs/, web-static/ (build output), frontend/node_modules/. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
156
backfill_image_data.py
Normal file
156
backfill_image_data.py
Normal file
@@ -0,0 +1,156 @@
|
||||
"""
|
||||
backfill_image_data.py
|
||||
======================
|
||||
|
||||
One-shot script to download image bytes from Discord and store them in
|
||||
input_images.db for rows that currently have image_data = NULL.
|
||||
|
||||
These rows were created before the BLOB-storage migration, so their bytes
|
||||
were never persisted. The script re-fetches each bot-reply message from
|
||||
Discord and writes the raw attachment bytes back into the DB.
|
||||
|
||||
Rows with bot_reply_id = 0 (web uploads that pre-date the migration) have
|
||||
no Discord source and are skipped — re-upload them via the web UI to
|
||||
backfill.
|
||||
|
||||
Usage
|
||||
-----
|
||||
python backfill_image_data.py
|
||||
|
||||
Requires:
|
||||
DISCORD_BOT_TOKEN in .env (same token the bot uses)
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
import sqlite3
|
||||
|
||||
import discord
|
||||
|
||||
try:
|
||||
from dotenv import load_dotenv
|
||||
load_dotenv()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
from config import BotConfig
|
||||
from input_image_db import DB_PATH
|
||||
|
||||
logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s")
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _load_null_rows() -> list[dict]:
|
||||
"""Return all rows that are missing image_data."""
|
||||
conn = sqlite3.connect(str(DB_PATH))
|
||||
conn.row_factory = sqlite3.Row
|
||||
rows = conn.execute(
|
||||
"SELECT id, bot_reply_id, channel_id, filename"
|
||||
" FROM input_images WHERE image_data IS NULL"
|
||||
).fetchall()
|
||||
conn.close()
|
||||
return [dict(r) for r in rows]
|
||||
|
||||
|
||||
def _save_image_data(row_id: int, data: bytes) -> None:
|
||||
conn = sqlite3.connect(str(DB_PATH))
|
||||
conn.execute("UPDATE input_images SET image_data = ? WHERE id = ?", (data, row_id))
|
||||
conn.commit()
|
||||
conn.close()
|
||||
|
||||
|
||||
async def _backfill(client: discord.Client) -> None:
|
||||
rows = _load_null_rows()
|
||||
|
||||
discord_rows = [r for r in rows if r["bot_reply_id"] != 0]
|
||||
web_rows = [r for r in rows if r["bot_reply_id"] == 0]
|
||||
|
||||
logger.info(
|
||||
"Rows missing image_data: %d total (%d from Discord, %d web-uploads skipped)",
|
||||
len(rows), len(discord_rows), len(web_rows),
|
||||
)
|
||||
|
||||
if web_rows:
|
||||
logger.info(
|
||||
"Skipped row IDs (no Discord source — re-upload via web UI): %s",
|
||||
[r["id"] for r in web_rows],
|
||||
)
|
||||
|
||||
if not discord_rows:
|
||||
logger.info("Nothing to fetch. Exiting.")
|
||||
return
|
||||
|
||||
ok = 0
|
||||
failed = 0
|
||||
|
||||
for row in discord_rows:
|
||||
row_id = row["id"]
|
||||
ch_id = row["channel_id"]
|
||||
msg_id = row["bot_reply_id"]
|
||||
filename = row["filename"]
|
||||
|
||||
try:
|
||||
channel = client.get_channel(ch_id) or await client.fetch_channel(ch_id)
|
||||
message = await channel.fetch_message(msg_id)
|
||||
|
||||
attachment = next(
|
||||
(a for a in message.attachments if a.filename == filename), None
|
||||
)
|
||||
if attachment is None:
|
||||
logger.warning(
|
||||
"Row %d: attachment '%s' not found on message %d — skipping",
|
||||
row_id, filename, msg_id,
|
||||
)
|
||||
failed += 1
|
||||
continue
|
||||
|
||||
data = await attachment.read()
|
||||
_save_image_data(row_id, data)
|
||||
logger.info("Row %d: saved '%s' (%d bytes)", row_id, filename, len(data))
|
||||
ok += 1
|
||||
|
||||
except discord.NotFound:
|
||||
logger.warning("Row %d: message %d not found (deleted?) — skipping", row_id, msg_id)
|
||||
failed += 1
|
||||
except discord.Forbidden:
|
||||
logger.warning("Row %d: no access to channel %d — skipping", row_id, ch_id)
|
||||
failed += 1
|
||||
except Exception as exc:
|
||||
logger.error("Row %d: unexpected error — %s", row_id, exc)
|
||||
failed += 1
|
||||
|
||||
logger.info(
|
||||
"Done. %d saved, %d failed/skipped, %d web-upload rows not touched.",
|
||||
ok, failed, len(web_rows),
|
||||
)
|
||||
|
||||
|
||||
async def _main(token: str) -> None:
|
||||
intents = discord.Intents.none() # no gateway events needed beyond connect
|
||||
client = discord.Client(intents=intents)
|
||||
|
||||
@client.event
|
||||
async def on_ready():
|
||||
logger.info("Logged in as %s", client.user)
|
||||
try:
|
||||
await _backfill(client)
|
||||
finally:
|
||||
await client.close()
|
||||
|
||||
await client.start(token)
|
||||
|
||||
|
||||
def main() -> None:
|
||||
try:
|
||||
config = BotConfig.from_env()
|
||||
except RuntimeError as exc:
|
||||
logger.error("Config error: %s", exc)
|
||||
return
|
||||
|
||||
asyncio.run(_main(config.discord_bot_token))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user