import discord import logging import os import asyncio import subprocess from huggingface_hub import InferenceClient import base64 from PIL import Image import io # 로깅 설정 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 intents.messages = True intents.guilds = True intents.guild_messages = True # 추론 API 클라이언트 설정 hf_client = InferenceClient(api_token=os.getenv("HF_TOKEN"), model_id="stabilityai/stable-diffusion-3-medium") # 특정 채널 ID SPECIFIC_CHANNEL_ID = int(os.getenv("DISCORD_CHANNEL_ID")) class MyClient(discord.Client): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) self.is_processing = False async def on_ready(self): logging.info(f'{self.user} has logged in!') subprocess.Popen(["python", "web.py"]) # Start the web.py server as a separate process logging.info("Web.py server has been started.") async def on_message(self, message): if message.author == self.user or not message.content: return if message.channel.id != SPECIFIC_CHANNEL_ID and not isinstance(message.channel, discord.Thread): return if self.is_processing: return self.is_processing = True try: image_path = await generate_image(message.content) if image_path: await send_image(message.channel, image_path) else: await message.channel.send("Failed to generate the image.") finally: self.is_processing = False async def generate_image(prompt): try: response = hf_client(inputs={"inputs": prompt}) image_data = response['images'][0] # Assuming the response contains image data image_bytes = io.BytesIO(base64.b64decode(image_data)) image = Image.open(image_bytes) image_path = "output.png" image.save(image_path) return image_path except Exception as e: logging.error(f"Failed to generate image: {str(e)}") return None async def send_image(channel, image_path): file = discord.File(image_path) await channel.send(file=file) if __name__ == "__main__": discord_client = MyClient(intents=intents) discord_client.run(os.getenv('DISCORD_TOKEN'))