Update app.py

app.py

@@ -4,6 +4,8 @@ import os
 from huggingface_hub import InferenceClient
 import asyncio
 import subprocess
+import torch
+from diffusers import StableDiffusionPipeline
 
 # 로깅 설정
 logging.basicConfig(level=logging.DEBUG, format='%(asctime)s:%(levelname)s:%(name)s: %(message)s', handlers=[logging.StreamHandler()])
@@ -24,6 +26,10 @@ SPECIFIC_CHANNEL_ID = int(os.getenv("DISCORD_CHANNEL_ID"))
 # 대화 히스토리를 저장할 전역 변수
 conversation_history = []
 
+# 이미지 생성 모델 로드
+if torch.cuda.is_available():
+    model = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4", revision="fp16", torch_dtype=torch.float16).to("cuda")
+
 class MyClient(discord.Client):
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
@@ -41,15 +47,31 @@ class MyClient(discord.Client):
             return
         if self.is_processing:
             return
-        self.is_processing = True
-        try:
-            response = await generate_response(message)
-            await message.channel.send(response)
-        finally:
-            self.is_processing = False
+        if message.content.startswith('!image '):
+            self.is_processing = True
+            try:
+                prompt = message.content[len('!image '):]  # 이미지 생성 프롬프트 파싱
+                image_path = await generate_image(prompt)
+                await message.channel.send(file=discord.File(image_path, 'generated_image.png'))
+            finally:
+                self.is_processing = False
+        else:
+            self.is_processing = True
+            try:
+                response = await generate_response(message)
+                await message.channel.send(response)
+            finally:
+                self.is_processing = False
+
+async def generate_image(prompt):
+    generator = torch.Generator(device="cuda").manual_seed(torch.seed())
+    image = model(prompt, num_inference_steps=50, generator=generator)["sample"][0]
+    image_path = '/tmp/generated_image.png'
+    image.save(image_path)
+    return image_path
 
 async def generate_response(message):
-    global conversation_history  # 전역 변수 사용을 명시
+    global conversation_history
     user_input = message.content
     user_mention = message.author.mention
     system_message = f"{user_mention}, DISCORD에서 사용자들의 질문에 답하는 어시스턴트입니다."