seawolf2357's picture
Update app.py
57dc30c verified
raw
history blame
2.62 kB
import discord
import logging
import os
from huggingface_hub import InferenceClient
import asyncio
import subprocess
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("stabilityai/stable-diffusion-3-medium")
# ํŠน์ • ์ฑ„๋„ 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)
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:
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)
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
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)
if __name__ == "__main__":
discord_client = MyClient(intents=intents)
discord_client.run(os.getenv('DISCORD_TOKEN'))