Update app.py

app.py

@@ -1,9 +1,10 @@
 import discord
 import logging
 import os
-from huggingface_hub import InferenceClient
 import asyncio
 import subprocess
+from huggingface_hub import InferenceApi
+import base64
 from PIL import Image
 import io
 
@@ -18,14 +19,11 @@ intents.guilds = True
 intents.guild_messages = True
 
 # 추론 API 클라이언트 설정
-hf_client = InferenceClient("stabilityai/stable-diffusion-3-medium")
+hf_client = InferenceApi("stabilityai/stable-diffusion-3-medium", token=os.getenv("HF_TOKEN"))
 
 # 특정 채널 ID
 SPECIFIC_CHANNEL_ID = int(os.getenv("DISCORD_CHANNEL_ID"))
 
-# 대화 히스토리를 저장할 변수
-conversation_history = []
-
 class MyClient(discord.Client):
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
@@ -39,40 +37,34 @@ class MyClient(discord.Client):
     async def on_message(self, message):
         if message.author == self.user or not message.content:
             return
-        if message.channel.id != SPECIFIC_CHANNEL_ID:
+        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)
-            await send_image(message.channel, image_path)
+            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):
-    """
-    Generate an image using the Stable Diffusion model from Hugging Face.
-    """
     try:
-        # Using the Inference API correctly
-        response = hf_client(inputs={"prompt": prompt}, parameters={"num_images": 1}, model_id="stabilityai/stable-diffusion-3-medium")
-        image_data = response['images'][0]  # Assuming the response contains image data
-
-        # Save image data to a file
+        response = hf_client(payload={"inputs": prompt})
+        image_data = response.get('image')  # Assuming the response returns an image key
         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):
-    """Send an image to the specified Discord channel."""
     file = discord.File(image_path)
     await channel.send(file=file)