app.py

import discord
from discord.ext import commands
from gradio_client import Client
import asyncio
import os
import logging
import subprocess
import sys

# 로깅 설정
logging.basicConfig(level=logging.DEBUG, format='%(asctime)s:%(levelname)s:%(name)s: %(message)s', handlers=[logging.StreamHandler()])

# 인텐트 설정
intents = discord.Intents.default()
intents.message_content = True

# Discord Bot 설정
bot = commands.Bot(command_prefix='!', intents=intents)

# Gradio API 클라이언트 설정
gradio_client = Client("http://211.233.58.202:7960/")

# Discord 채널 ID
CHANNEL_ID = 1269529561914413106

@bot.event
async def on_ready():
    logging.info(f'{bot.user.name} has connected to Discord!')
    try:
        # Gradio 연결 테스트
        response = gradio_client.predict("test", api_name="/infer_t2i")
        logging.info("Gradio API connection successful")
    except Exception as e:
        logging.error(f"Gradio API connection failed: {str(e)}")

@bot.event
async def on_message(message):
    if message.author == bot.user:
        return

    if message.channel.id != CHANNEL_ID and not isinstance(message.channel, discord.Thread):
        return

    try:
        result = await asyncio.wait_for(
            asyncio.to_thread(
                gradio_client.predict,
                prompt=message.content,
                seed=123,
                randomize_seed=False,
                width=1024,
                height=576,
                guidance_scale=5,
                num_inference_steps=30,
                api_name="/infer_t2i"
            ),
            timeout=60  # 60 seconds timeout
        )
        
        logging.debug(f"Predict result: {result}")
        if isinstance(result, tuple) and len(result) > 0:
            file_path = result[0]
            await message.channel.send(file=discord.File(file_path))
        else:
            logging.error("Unexpected result format from Gradio API")
    except asyncio.TimeoutError:
        await message.channel.send("Image generation timed out. Please try again.")
    except Exception as e:
        error_message = f"An error occurred: {str(e)}"
        await message.channel.send(error_message)
        logging.error(f"Error in on_message: {error_message}")

@bot.event
async def on_disconnect():
    logging.info("Bot disconnected. Attempting to reconnect...")

@bot.event
async def on_resume():
    logging.info("Bot connection resumed.")

def run_web():
    subprocess.run([sys.executable, "web.py"])

def run_discord_bot():
    discord_token = os.getenv('DISCORD_TOKEN')
    if discord_token:
        bot.run(discord_token)
    else:
        logging.error("Discord token is not set. Please check your environment variables.")
        sys.exit(1)  # Exit if token is not set

if __name__ == "__main__":
    web_process = subprocess.Popen([sys.executable, "web.py"], stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    logging.info("Web server process started.")

    run_discord_bot()

    web_process.wait()  # Ensure the web process completes