Update app.py

app.py

@@ -1,11 +1,8 @@
 import discord
 import logging
 import os
+import requests
 import uuid
-import torch
-import subprocess
-from huggingface_hub import snapshot_download
-from diffusers import StableDiffusionPipeline, StableDiffusionImg2ImgPipeline
 from transformers import pipeline
 
 # 로깅 설정
@@ -13,91 +10,68 @@ logging.basicConfig(level=logging.DEBUG, format='%(asctime)s:%(levelname)s:%(nam
 
 # 인텐트 설정
 intents = discord.Intents.default()
-intents.message_content = True
-
-# Hugging Face 모델 다운로드
-huggingface_token = os.getenv("HF_TOKEN")
-model_local_dir = "Fluently-XL-Final"
-
-# 모델을 로컬 디렉토리에 다운로드
-model_path = snapshot_download(
-    repo_id="fluently/Fluently-XL-Final",
-    repo_type="model",
-    local_dir=model_local_dir,
-    token=huggingface_token,
-)
-
-# 모델 로드 함수
-def load_pipeline(pipeline_type):
-    logging.debug(f'Loading pipeline: {pipeline_type}')
-    if pipeline_type == "text2img":
-        return StableDiffusionPipeline.from_pretrained(model_local_dir, torch_dtype=torch.float32)
-    elif pipeline_type == "img2img":
-        return StableDiffusionImg2ImgPipeline.from_pretrained(model_local_dir, torch_dtype=torch.float32)
-
-# 디바이스 설정 (CPU 전용)
-device = torch.device("cpu")
+intents.messages = True
 
 # 번역 파이프라인 설정
 translator = pipeline("translation", model="Helsinki-NLP/opus-mt-ko-en")
 
-# 고정된 네거티브 프롬프트
-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"
+# Hugging Face API 토큰
+hf_token = os.getenv("HF_TOKEN")
+
+# 모델 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"
+}
 
 # 디스코드 봇 클래스
 class MyClient(discord.Client):
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
-        self.is_processing = False
-        self.text2img_pipeline = load_pipeline("text2img").to(device)
-        self.text2img_pipeline.enable_attention_slicing()  # 메모리 최적화
 
     async def on_ready(self):
-        logging.info(f'{self.user} logged in!')
-        subprocess.Popen(["python", "web.py"])
-        logging.info("web.py server has started.")
+        logging.info(f'{self.user} has logged in!')
 
     async def on_message(self, message):
         if message.author == self.user:
             return
         if message.content.startswith('!image '):
-            self.is_processing = True
-            try:
-                prompt = message.content[len('!image '):]
-                prompt_en = translate_prompt(prompt)
-                logging.debug(f'Translated prompt: {prompt_en}')
-                logging.debug(f'Fixed negative prompt: {negative_prompt}')
-
-                image_path = await self.generate_image(prompt_en, negative_prompt)
-                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')
-                )
-            finally:
-                self.is_processing = False
-
-    async def generate_image(self, prompt, negative_prompt):
-        try:
-            if negative_prompt is None:
-                negative_prompt = ""
-            logging.debug(f"Calling text2img_pipeline with prompt: {prompt} and negative_prompt: {negative_prompt}")
-
-            result = self.text2img_pipeline(
-                prompt=prompt,
-                negative_prompt=negative_prompt,
-                num_inference_steps=50
+            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')
             )
-            images = result.images
 
+    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'
-            images[0].save(image_path)
-            logging.debug(f"Image saved to {image_path}")
-
+            with open(image_path, 'wb') as f:
+                f.write(response.content)  # 이미지 데이터를 파일로 저장
             return image_path
-        except Exception as e:
-            logging.error(f"An error occurred during image generation: {e}")
-            raise
+        else:
+            raise Exception("Failed to generate image from the API")
 
 # 프롬프트 번역 함수
 def translate_prompt(prompt):