Update app.py

app.py

@@ -1,11 +1,7 @@
+import requests
 import discord
 import logging
 import os
-import uuid
-import torch
-import subprocess
-from huggingface_hub import snapshot_download
-from diffusers import StableDiffusionPipeline, StableDiffusionImg2ImgPipeline
 from transformers import pipeline
 
 # 로깅 설정
@@ -15,47 +11,45 @@ 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_path = snapshot_download(
-#    repo_id="SG161222/RealVisXL_V4.0",
-    repo_id="cagliostrolab/animagine-xl-3.1",
-
-    repo_type="model",
-#    local_dir="RealVisXL_V4.0",
-    local_dir="animagine-xl-3.1",    
-    token=huggingface_token,
-)
-
-# 모델 로드 함수
-def load_pipeline(pipeline_type):
-    logging.debug(f'파이프라인 로드 중: {pipeline_type}')
-    if pipeline_type == "text2img":
-        return StableDiffusionPipeline.from_pretrained(model_path, torch_dtype=torch.float16, use_fast=True)
-    elif pipeline_type == "img2img":
-        return StableDiffusionImg2ImgPipeline.from_pretrained(model_path, torch_dtype=torch.float16, use_fast=True)
-
-# 디바이스 설정
-device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
-
 # 번역 파이프라인 설정
 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"
 
+# 인퍼런스 API를 사용하기 위한 함수
+def generate_image(prompt, negative_prompt):
+    headers = {
+        "Authorization": f"Bearer {os.getenv('HF_TOKEN')}"
+    }
+    data = {
+        "inputs": {
+            "prompt": prompt,
+            "negative_prompt": negative_prompt,
+            "num_inference_steps": 50
+        }
+    }
+    api_url = "https://api-inference.huggingface.co/models/fluently/Fluently-XL-Final"
+    response = requests.post(api_url, headers=headers, json=data)
+    if response.status_code == 200:
+        image_url = response.json()[0]['url']  # 응답에서 이미지 URL 추출
+        return image_url
+    else:
+        logging.error("API 요청에 실패했습니다: " + response.text)
+        return None
+
+# 프롬프트 번역 함수
+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
+
 # 디스코드 봇 클래스
 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}로 로그인되었습니다!')
-        subprocess.Popen(["python", "web.py"])
-        logging.info("web.py 서버가 시작되었습니다.")
 
     async def on_message(self, message):
         if message.author == self.user:
@@ -68,29 +62,19 @@ class MyClient(discord.Client):
                 logging.debug(f'번역된 프롬프트: {prompt_en}')
                 logging.debug(f'고정된 네거티브 프롬프트: {negative_prompt}')
 
-                image_path = await self.generate_image(prompt_en, negative_prompt)
-                user_id = message.author.id  # 사용자의 ID 캡처
-                await message.channel.send(
-                    f"<@{user_id}> 님이 요청하신 이미지입니다.",
-                    file=discord.File(image_path, 'generated_image.png')
-                )
+                image_url = generate_image(prompt_en, negative_prompt)
+                user_id = message.author.id
+                if image_url:
+                    await message.channel.send(
+                        f"<@{user_id}> 님이 요청하신 이미지입니다: {image_url}"
+                    )
+                else:
+                    await message.channel.send(f"<@{user_id}> 이미지 생성에 실패하였습니다.")
             finally:
                 self.is_processing = False
-
-    async def generate_image(self, prompt, negative_prompt):
-        generator = torch.Generator(device=device).manual_seed(torch.seed())
-        images = self.text2img_pipeline(prompt, negative_prompt=negative_prompt, num_inference_steps=50, generator=generator)["images"]
-        image_path = f'/tmp/{uuid.uuid4()}.png'
-        images[0].save(image_path)
-        return image_path
-
-# 프롬프트 번역 함수
-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
+        else:
+            # "!image" 명령어가 아닐 경우 안내 메시지 전송
+            await message.channel.send('올바른 명령어를 입력해 주세요. 예) "!image 귀여운 고양이가 잠을 자고있다." 등으로 입력하시면 이미지가 생성됩니다.')
 
 # 디스코드 토큰 및 봇 실행
 if __name__ == "__main__":