Update app.py

app.py

@@ -1,87 +1,56 @@
 import discord
-import logging
 import os
-import requests
-import uuid
+import logging
 from transformers import pipeline
+import subprocess
+from diffusers import DiffusionPipeline
 
 # 로깅 설정
 logging.basicConfig(level=logging.DEBUG, format='%(asctime)s:%(levelname)s:%(name)s: %(message)s', handlers=[logging.StreamHandler()])
 
 # 인텐트 설정
 intents = discord.Intents.default()
-intents.messages = True
+intents.message_content = True
 
 # 번역 파이프라인 설정
 translator = pipeline("translation", model="Helsinki-NLP/opus-mt-ko-en")
 
-# Hugging Face API 토큰
-hf_token = os.getenv("HF_TOKEN")
+# 이미지 생성 모델 로드
+model_id = "fluently/Fluently-XL-Final"
+device = "cuda" if torch.cuda.is_available() else "cpu"
+image_pipeline = DiffusionPipeline.from_pretrained(model_id).to(device)
 
-# 모델 URL 설정
-model_urls = {
-    "epiCPhotoGasm": "https://api-inference.huggingface.co/models/Yntec/epiCPhotoGasm",
-    "AM-mix1": "https://api-inference.huggingface.co/models/digiplay/AM-mix1",
-    "AbsoluteReality": "https://api-inference.huggingface.co/models/Yntec/AbsoluteReality",
-    "DreamPhotoGASM": "https://api-inference.huggingface.co/models/Yntec/DreamPhotoGASM",
-    "insaneRealistic_v1": "https://api-inference.huggingface.co/models/digiplay/insaneRealistic_v1",
-    "photoMovieXFinal": "https://api-inference.huggingface.co/models/Yntec/photoMovieXFinal",
-    "Memento": "https://api-inference.huggingface.co/models/Yntec/Memento",
-    "photoMovieRealistic": "https://api-inference.huggingface.co/models/Yntec/photoMovieRealistic",
-    "Timeless": "https://api-inference.huggingface.co/models/Yntec/Timeless",
-    "Fabulous": "https://api-inference.huggingface.co/models/Yntec/Fabulous",
-    "IncredibleLife": "https://api-inference.huggingface.co/models/Yntec/IncredibleLife",
-    "RealLife": "https://api-inference.huggingface.co/models/Yntec/RealLife",
-    "HyperRealism": "https://api-inference.huggingface.co/models/Yntec/HyperRealism"
-}
+# 환경 변수에서 지정된 채널 ID 가져오기
+TARGET_CHANNEL_ID = os.getenv('TARGET_CHANNEL_ID')
 
-# 디스코드 봇 클래스
 class MyClient(discord.Client):
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-
     async def on_ready(self):
-        logging.info(f'{self.user} has logged in!')
+        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:
+        if message.author == self.user or message.channel.id != TARGET_CHANNEL_ID:
             return
         if message.content.startswith('!image '):
-            parts = message.content.split(' ', 2)
-            if len(parts) < 3 or parts[1] not in model_urls:
-                await message.channel.send("Invalid model name or usage.")
-                return
-            
-            model_name = parts[1]
-            prompt = parts[2]
-            prompt_en = translate_prompt(prompt)  # 번역 기능 추가
-            image_path = await self.generate_image(model_urls[model_name], prompt_en)
-            user_id = message.author.id
-            await message.channel.send(
-                f"<@{user_id}> Here is your requested image:",
-                file=discord.File(image_path, 'generated_image.png')
-            )
-
-    async def generate_image(self, model_url, prompt):
-        headers = {"Authorization": f"Bearer {hf_token}"}
-        response = requests.post(model_url, headers=headers, json={"inputs": prompt})
-        if response.status_code == 200:
-            image_path = f'/tmp/{uuid.uuid4()}.png'
-            with open(image_path, 'wb') as f:
-                f.write(response.content)  # 이미지 데이터를 파일로 저장
-            return image_path
-        else:
-            raise Exception("Failed to generate image from the API")
+            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):
-    logging.debug(f'Translating prompt: {prompt}')
     translation = translator(prompt, max_length=512)
-    translated_text = translation[0]['translation_text']
-    logging.debug(f'Translated text: {translated_text}')
-    return translated_text
+    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
 
-# 디스코드 토큰 및 봇 실행
+# 봇 실행
 if __name__ == "__main__":
     discord_token = os.getenv('DISCORD_TOKEN')
     discord_client = MyClient(intents=intents)