|
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 |
|
|
|
|
|
hf_client = InferenceClient("stabilityai/stable-diffusion-3-medium") |
|
|
|
|
|
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"]) |
|
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: |
|
|
|
response = hf_client(inputs={"prompt": prompt}, parameters={"num_images": 1}, model_id="stabilityai/stable-diffusion-3-medium") |
|
image_data = response['images'][0] |
|
|
|
|
|
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')) |
|
|