app.py

import discord
import logging
import os
import asyncio
import aiohttp
from gradio_client import Client, handle_file
from aiohttp import web

# 로깅 설정
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

# API 클라이언트 설정
api_client = Client("http://hugpu.ai:7889")

# 디스코드 봇 클래스
class MyClient(discord.Client):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.is_processing = False
        self.web_app = None

    async def start_web_server(self):
        app = web.Application()
        app.router.add_get('/', self.handle_index)
        runner = web.AppRunner(app)
        await runner.setup()
        site = web.TCPSite(runner, 'localhost', 8080)
        await site.start()
        self.web_app = app
        logging.info("Web server started on http://localhost:8080")

    async def handle_index(self, request):
        return web.Response(text="Server is running")

    async def on_ready(self):
        logging.info(f'{self.user}로 로그인되었습니다!')
        channel = self.get_channel(int(os.getenv("DISCORD_CHANNEL_ID", "123456789012345678")))
        await channel.send("이미지와 프롬프트를 입력해 이미지 생성을 요청하세요. 사용 예: '!image [프롬프트]' (이미지 첨부)")
        
        # Web 서버 시작
        await self.start_web_server()

    async def on_message(self, message):
        if message.author == self.user:
            return
        if message.content.startswith('!image '):
            if self.is_processing:
                await message.channel.send("이미지 생성 중입니다. 잠시만 기다려 주세요.")
                return
            self.is_processing = True
            try:
                if len(message.attachments) > 0:
                    image_url = message.attachments[0].url
                    prompt = message.content[7:].strip()
                    generated_images = await self.generate_image(prompt, image_url)
                    user_id = message.author.id
                    if generated_images and len(generated_images) > 0:
                        image_path = generated_images[0]['image']
                        await message.channel.send(
                            f"<@{user_id}> 님이 요청하신 이미지입니다.\n"
                            f"프롬프트: {prompt}",
                            file=discord.File(image_path)
                        )
                    else:
                        await message.channel.send("이미지 생성에 실패했습니다.")
                else:
                    await message.channel.send("이미지를 첨부해 주세요.")
            except Exception as e:
                logging.error(f'이미지 생성 중 오류 발생: {e}')
            finally:
                self.is_processing = False

    async def generate_image(self, prompt, image_url):
        if not prompt or not image_url:
            raise ValueError("Prompt or image URL is empty")
        logging.debug(f"Sending request to API with prompt: {prompt} and image URL: {image_url}")
        try:
            result = api_client.predict(
                images=[handle_file(image_url)],
                prompt=prompt,
                negative_prompt="",
                preserve_face_structure=True,
                face_strength=1.3,
                likeness_strength=1,
                nfaa_negative_prompt="naked, bikini, skimpy, scanty, bare skin, lingerie, swimsuit, exposed, see-through",
                api_name="/generate_image"
            )
            logging.debug(f"API response received: {result}")
            return result
        except Exception as e:
            logging.error(f'API 요청 중 오류 발생: {e}', exc_info=True)
            raise RuntimeError(f"API 요청 중 오류 발생: {str(e)}")

async def main():
    discord_token = os.getenv('DISCORD_TOKEN')
    client = MyClient(intents=intents)
    try:
        await client.start(discord_token)
    except KeyboardInterrupt:
        await client.close()

if __name__ == "__main__":
    asyncio.run(main())