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:
Khoa (Revenovich) Tran Gia
218
bot.py
Normal file
218
bot.py
Normal file
@@ -0,0 +1,218 @@
|
||||
"""
|
||||
bot.py
|
||||
======
|
||||
|
||||
Discord bot entry point. In WEB_ENABLED mode, also starts a FastAPI/Uvicorn
|
||||
web server in the same asyncio event loop via asyncio.gather.
|
||||
|
||||
Jobs are submitted directly to ComfyUI — no internal SerialJobQueue.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
from datetime import datetime, timezone
|
||||
from pathlib import Path
|
||||
|
||||
import discord
|
||||
from discord.ext import commands
|
||||
|
||||
from comfy_client import ComfyClient
|
||||
from config import BotConfig, COMMAND_PREFIX
|
||||
import generation_db
|
||||
from input_image_db import init_db, get_all_images
|
||||
from status_monitor import StatusMonitor
|
||||
from workflow_manager import WorkflowManager
|
||||
from workflow_state import WorkflowStateManager
|
||||
from commands import register_all_commands, CustomHelpCommand
|
||||
from commands.input_images import PersistentSetInputView
|
||||
from commands.server import autostart_comfy
|
||||
|
||||
try:
|
||||
from dotenv import load_dotenv
|
||||
load_dotenv()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
logging.basicConfig(level=logging.INFO)
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
_PROJECT_ROOT = Path(__file__).resolve().parent
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Bot setup
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def get_prefix(bot, message):
|
||||
"""Dynamic command prefix getter."""
|
||||
msg = message.content.lower()
|
||||
if msg.startswith(COMMAND_PREFIX):
|
||||
return COMMAND_PREFIX
|
||||
return COMMAND_PREFIX
|
||||
|
||||
|
||||
intents = discord.Intents.default()
|
||||
intents.message_content = True
|
||||
intents.guilds = True
|
||||
|
||||
bot = commands.Bot(
|
||||
command_prefix=get_prefix,
|
||||
intents=intents,
|
||||
help_command=CustomHelpCommand(),
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Event handlers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@bot.event
|
||||
async def on_ready() -> None:
|
||||
logger.info("Logged in as %s (ID: %s)", bot.user, bot.user.id)
|
||||
if not hasattr(bot, "start_time"):
|
||||
bot.start_time = datetime.now(timezone.utc)
|
||||
cfg = getattr(bot, "config", None)
|
||||
if cfg:
|
||||
for row in get_all_images():
|
||||
view = PersistentSetInputView(bot, cfg, row["id"])
|
||||
bot.add_view(view, message_id=row["bot_reply_id"])
|
||||
asyncio.create_task(autostart_comfy(cfg))
|
||||
|
||||
if not hasattr(bot, "status_monitor"):
|
||||
log_ch = getattr(getattr(bot, "config", None), "log_channel_id", None)
|
||||
if log_ch:
|
||||
bot.status_monitor = StatusMonitor(bot, log_ch)
|
||||
if hasattr(bot, "status_monitor"):
|
||||
await bot.status_monitor.start()
|
||||
|
||||
|
||||
@bot.event
|
||||
async def on_disconnect() -> None:
|
||||
logger.info("Discord connection closed")
|
||||
|
||||
|
||||
@bot.event
|
||||
async def on_resumed() -> None:
|
||||
logger.info("Discord session resumed")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Startup helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
async def create_comfy(config: BotConfig) -> ComfyClient:
|
||||
state_manager = WorkflowStateManager(state_file="current-workflow-changes.json")
|
||||
workflow_manager = WorkflowManager()
|
||||
return ComfyClient(
|
||||
server_address=config.comfy_server,
|
||||
workflow_manager=workflow_manager,
|
||||
state_manager=state_manager,
|
||||
logger=logger,
|
||||
history_limit=config.comfy_history_limit,
|
||||
output_path=config.comfy_output_path,
|
||||
)
|
||||
|
||||
|
||||
def _try_autoload_last_workflow(client: ComfyClient) -> None:
|
||||
"""Re-load the last used workflow from the workflows/ folder on startup."""
|
||||
last_wf = client.state_manager.get_last_workflow_file()
|
||||
if not last_wf:
|
||||
return
|
||||
wf_path = _PROJECT_ROOT / "workflows" / last_wf
|
||||
if not wf_path.exists():
|
||||
logger.warning("Last workflow file not found: %s", wf_path)
|
||||
return
|
||||
try:
|
||||
with open(wf_path, "r", encoding="utf-8") as f:
|
||||
workflow = json.load(f)
|
||||
# Restore template without clearing overrides on restart
|
||||
client.workflow_manager.set_workflow_template(workflow)
|
||||
logger.info("Auto-loaded last workflow: %s", last_wf)
|
||||
except Exception as exc:
|
||||
logger.error("Failed to auto-load workflow %s: %s", last_wf, exc)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Main entry point
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
async def main() -> None:
|
||||
try:
|
||||
config: BotConfig = BotConfig.from_env()
|
||||
logger.info("Configuration loaded")
|
||||
except RuntimeError as exc:
|
||||
logger.error("Configuration error: %s", exc)
|
||||
return
|
||||
|
||||
bot.comfy = await create_comfy(config)
|
||||
bot.config = config
|
||||
|
||||
# Auto-load last workflow (restores template without clearing overrides)
|
||||
_try_autoload_last_workflow(bot.comfy)
|
||||
|
||||
# Fallback: WORKFLOW_FILE env var
|
||||
if not bot.comfy.get_workflow_template() and config.workflow_file:
|
||||
try:
|
||||
bot.comfy.load_workflow_from_file(config.workflow_file)
|
||||
logger.info("Loaded workflow from %s", config.workflow_file)
|
||||
except Exception as exc:
|
||||
logger.error("Failed to load workflow %s: %s", config.workflow_file, exc)
|
||||
|
||||
from user_state_registry import UserStateRegistry
|
||||
bot.user_registry = UserStateRegistry(
|
||||
settings_dir=_PROJECT_ROOT / "user_settings",
|
||||
default_workflow=bot.comfy.get_workflow_template(),
|
||||
)
|
||||
|
||||
init_db()
|
||||
generation_db.init_db(_PROJECT_ROOT / "generation_history.db")
|
||||
register_all_commands(bot, config)
|
||||
logger.info("All commands registered")
|
||||
|
||||
coroutines = [bot.start(config.discord_bot_token)]
|
||||
|
||||
if config.web_enabled:
|
||||
try:
|
||||
import uvicorn
|
||||
from web.deps import set_bot
|
||||
from web.app import create_app
|
||||
|
||||
set_bot(bot)
|
||||
fastapi_app = create_app()
|
||||
|
||||
uvi_config = uvicorn.Config(
|
||||
fastapi_app,
|
||||
host=config.web_host,
|
||||
port=config.web_port,
|
||||
log_level="info",
|
||||
loop="none", # use existing event loop
|
||||
)
|
||||
uvi_server = uvicorn.Server(uvi_config)
|
||||
coroutines.append(uvi_server.serve())
|
||||
logger.info(
|
||||
"Web UI enabled at http://%s:%d", config.web_host, config.web_port
|
||||
)
|
||||
except ImportError:
|
||||
logger.warning(
|
||||
"uvicorn or fastapi not installed — web UI disabled. "
|
||||
"pip install fastapi uvicorn[standard]"
|
||||
)
|
||||
|
||||
try:
|
||||
await asyncio.gather(*coroutines)
|
||||
finally:
|
||||
if hasattr(bot, "status_monitor"):
|
||||
await bot.status_monitor.stop()
|
||||
if hasattr(bot, "comfy"):
|
||||
await bot.comfy.close()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
try:
|
||||
asyncio.run(main())
|
||||
except KeyboardInterrupt:
|
||||
logger.info("Received interrupt, shutting down…")
|
||||
Reference in New Issue
Block a user