Update app.py

app.py

@@ -1,9 +1,14 @@
+import requests
 import discord
-import os
 import logging
+import os
 from transformers import pipeline
 import subprocess
+import torch
 from diffusers import DiffusionPipeline
+import io
+from PIL import Image
+import time
 
 # 로깅 설정
 logging.basicConfig(level=logging.DEBUG, format='%(asctime)s:%(levelname)s:%(name)s: %(message)s', handlers=[logging.StreamHandler()])
@@ -15,13 +20,33 @@ intents.message_content = True
 # 번역 파이프라인 설정
 translator = pipeline("translation", model="Helsinki-NLP/opus-mt-ko-en")
 
-# 이미지 생성 모델 로드
-model_id = "fluently/Fluently-XL-Final"
+# 환경 변수에서 지정된 채널 ID 가져오기
+TARGET_CHANNEL_ID = int(os.getenv('TARGET_CHANNEL_ID'))  # 채널 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"
-image_pipeline = DiffusionPipeline.from_pretrained(model_id).to(device)
 
-# 환경 변수에서 지정된 채널 ID 가져오기
-TARGET_CHANNEL_ID = os.getenv('TARGET_CHANNEL_ID')
+# 이미지 생성 파이프라인 설정
+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]  # 첫 번째 이미지 선택
+    # 이미지 URL 대신 이미지를 반환하도록 수정
+    return image
 
 class MyClient(discord.Client):
     async def on_ready(self):
@@ -33,22 +58,27 @@ class MyClient(discord.Client):
         if message.author == self.user or message.channel.id != TARGET_CHANNEL_ID:
             return
         if message.content.startswith('!image '):
-            prompt = message.content[len('!image '):]
-            translated_prompt = translate_prompt(prompt)
-            image_path = await generate_image(translated_prompt)
-            await message.channel.send(file=discord.File(image_path))
-
-# 프롬프트 번역 함수
-def translate_prompt(prompt):
-    translation = translator(prompt, max_length=512)
-    return translation[0]['translation_text']
-
-# 비동기 이미지 생성 함수
-async def generate_image(prompt):
-    image = image_pipeline(prompt)["sample"][0]  # 이미지 생성
-    image_path = f"./generated_images/{prompt.replace(' ', '_')}.png"
-    image.save(image_path)
-    return image_path
+            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__":