Spaces:
Runtime error
Runtime error
import requests | |
import discord | |
import logging | |
import os | |
from transformers import pipeline | |
import subprocess | |
import torch | |
from diffusers import DiffusionPipeline | |
import io | |
from PIL import Image | |
# ๋ก๊น ์ค์ | |
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 | |
# ๋ฒ์ญ ํ์ดํ๋ผ์ธ ์ค์ | |
translator = pipeline("translation", model="Helsinki-NLP/opus-mt-ko-en") | |
# ํ๊ฒฝ ๋ณ์์์ ์ง์ ๋ ์ฑ๋ ID ๊ฐ์ ธ์ค๊ธฐ, ๊ธฐ๋ณธ๊ฐ ์ฌ์ฉ | |
TARGET_CHANNEL_ID = int(os.getenv('TARGET_CHANNEL_ID', '123456789012345678')) # ๊ธฐ๋ณธ๊ฐ ์ฌ์ฉ, ์ค์ ์ฑ๋ ID๋ก ๋์ฒด ํ์ | |
# ๊ณ ์ ๋ ๋ค๊ฑฐํฐ๋ธ ํ๋กฌํํธ | |
negative_prompt = "blur, low quality, bad composition, ugly, disfigured, weird colors, low quality, jpeg artifacts, lowres, grainy, deformed structures, blurry, opaque, low contrast, distorted details, details are low" | |
# ๋๋ฐ์ด์ค ์ค์ | |
device = "cuda" if torch.cuda.is_available() else "cpu" | |
# ์ด๋ฏธ์ง ์์ฑ ํ์ดํ๋ผ์ธ ์ค์ | |
pipeline = DiffusionPipeline.from_pretrained("fluently/Fluently-XL-Final", torch_dtype=torch.float32) | |
pipeline = pipeline.to(device) | |
# ํ๋กฌํํธ ๋ฒ์ญ ํจ์ | |
def translate_prompt(prompt): | |
logging.debug(f'ํ๋กฌํํธ ๋ฒ์ญ ์ค: {prompt}') | |
translation = translator(prompt, max_length=512) | |
translated_text = translation[0]['translation_text'] | |
logging.debug(f'๋ฒ์ญ๋ ํ ์คํธ: {translated_text}') | |
return translated_text | |
def generate_image(prompt, negative_prompt): | |
combined_prompt = f"{prompt}. {negative_prompt}" | |
result = pipeline(combined_prompt) | |
image = result.images[0] # ์ฒซ ๋ฒ์งธ ์ด๋ฏธ์ง ์ ํ | |
return image | |
class MyClient(discord.Client): | |
async def on_ready(self): | |
logging.info(f'{self.user}๋ก ๋ก๊ทธ์ธ๋์์ต๋๋ค!') | |
subprocess.Popen(["python", "web.py"]) # ๋ณ๋์ Python ์คํฌ๋ฆฝํธ ์คํ | |
logging.info("web.py ์๋ฒ๊ฐ ์์๋์์ต๋๋ค.") | |
async def on_message(self, message): | |
if message.author == self.user or message.channel.id != TARGET_CHANNEL_ID: | |
return | |
if message.content.startswith('!image '): | |
self.is_processing = True | |
try: | |
prompt = message.content[len('!image '):] | |
prompt_en = translate_prompt(prompt) | |
image = generate_image(prompt_en, negative_prompt) | |
user_id = message.author.id | |
if image: | |
# ์ด๋ฏธ์ง๋ฅผ Discord์ ์ง์ ์ ๋ก๋ | |
with io.BytesIO() as image_binary: | |
image.save(image_binary, 'PNG') | |
image_binary.seek(0) | |
await message.channel.send( | |
f"<@{user_id}> ๋์ด ์์ฒญํ์ ์ด๋ฏธ์ง์ ๋๋ค:", | |
file=discord.File(fp=image_binary, filename='image.png') | |
) | |
else: | |
await message.channel.send(f"<@{user_id}> ์ด๋ฏธ์ง ์์ฑ์ ์คํจํ์์ต๋๋ค.") | |
finally: | |
self.is_processing = False | |
else: | |
await message.channel.send('์ฌ๋ฐ๋ฅธ ๋ช ๋ น์ด๋ฅผ ์ ๋ ฅํด ์ฃผ์ธ์. ์) "!image ๊ท์ฌ์ด ๊ณ ์์ด๊ฐ ์ ์ ์๊ณ ์๋ค." ๋ฑ์ผ๋ก ์ ๋ ฅํ์๋ฉด ์ด๋ฏธ์ง๊ฐ ์์ฑ๋ฉ๋๋ค.') | |
# ๋ด ์คํ | |
if __name__ == "__main__": | |
discord_token = os.getenv('DISCORD_TOKEN') | |
discord_client = MyClient(intents=intents) | |
discord_client.run(discord_token) | |