Update app.py

app.py

@@ -1,23 +1,25 @@
 # img_bot.py
-import discord, os, io, re, asyncio, logging, requests, replicate, subprocess
-from transformers import pipeline as transformers_pipeline   # 번역 파이프라인
+import discord, os, io, re, random, asyncio, logging, requests, replicate, subprocess
+from transformers import pipeline as transformers_pipeline    # 번역 파이프라인
+from gradio_client import Client, handle_file                 # NEW ─ Gradio API
 
 # ── 환경 변수 ────────────────────────────────────────────────
-TOKEN       = os.getenv("DISCORD_TOKEN")                     # Discord 봇 토큰
-CHANNEL_ID  = int(os.getenv("DISCORD_CHANNEL_ID"))           # 감시할 채널 ID
-REPL_TOKEN  = (os.getenv("OPENAI_API_KEY") or "").strip()    # Replicate 토큰 재사용
-HF_TOKEN    = (os.getenv("HF_TOKEN") or "").strip()          # Hugging Face 토큰
-USE_MSG_INTENT = os.getenv("MESSAGE_CONTENT_INTENT", "1") != "0"  # 1=ON(기본), 0=OFF
+TOKEN       = os.getenv("DISCORD_TOKEN")                      # Discord 봇 토큰
+CHANNEL_ID  = int(os.getenv("DISCORD_CHANNEL_ID"))            # 감시할 채널 ID
+REPL_TOKEN  = (os.getenv("OPENAI_API_KEY") or "").strip()     # 그대로 유지(Replicate 토큰)
+HF_TOKEN    = (os.getenv("HF_TOKEN") or "").strip()           # Hugging Face Personal Access Token
 
 if not TOKEN or not CHANNEL_ID:
     raise RuntimeError("DISCORD_TOKEN 과 DISCORD_CHANNEL_ID 환경 변수를 모두 지정하세요.")
 if not REPL_TOKEN:
     raise RuntimeError("OPENAI_API_KEY 환경 변수에 Replicate Personal Access Token 값을 넣어주세요.")
 
-os.environ["REPLICATE_API_TOKEN"] = REPL_TOKEN       # Replicate 가 참조할 토큰
+# Replicate 라이브러리가 참조하도록 토큰 주입 (비록 사용 안 하지만 구조 유지)
+os.environ["REPLICATE_API_TOKEN"] = REPL_TOKEN
 
-# ── 이미지 생성 모델 (Flux Schnell) ─────────────────────────
-MODEL = "black-forest-labs/flux-schnell"
+# ── Gradio 서버 엔드포인트 ──────────────────────────────────
+GRADIO_URL = "http://211.233.58.201:7896"           # 서버 주소
+GRADIO_API = "/generate_image"                      # 사용할 함수 이름
 
 # ── 번역 파이프라인 (CPU) ───────────────────────────────────
 translator_kwargs = {"device": -1}
@@ -28,7 +30,7 @@ translator = transformers_pipeline(
 )
 
 def ko2en(text: str) -> str:
-    """프롬프트에 한글 포함 시 영어 번역."""
+    """프롬프트에 한글이 있으면 영어로 번역."""
     if re.search(r"[가-힣]", text):
         try:
             return translator(text, max_length=512)[0]["translation_text"].strip()
@@ -40,14 +42,9 @@ def ko2en(text: str) -> str:
 logging.basicConfig(level=logging.INFO,
                     format="%(asctime)s [%(levelname)s] %(message)s")
 
-# ── Discord 인텐트 설정 ────────────────────────────────────
+# ── Discord 인텐트 ──────────────────────────────────────────
 intents = discord.Intents.default()
-if USE_MSG_INTENT:
-    intents.message_content = True
-    logging.info("Message Content Intent = ON  "
-                 "(Developer Portal → Bot → Privileged Intents 에서 ON 필요)")
-else:
-    logging.info("Message Content Intent = OFF  (MESSAGE_CONTENT_INTENT=0)")
+intents.message_content = True   # 메시지 읽기 (Portal에서 ON 필요)
 
 class ImageBot(discord.Client):
     async def on_ready(self):
@@ -60,43 +57,57 @@ class ImageBot(discord.Client):
             logging.warning(f"web.py 실행 실패: {e}")
 
     async def on_message(self, message: discord.Message):
-        # 메시지 내용 읽을 권한이 없으면 바로 반환
-        if not USE_MSG_INTENT:
-            return
         if message.author.id == self.user.id or message.channel.id != CHANNEL_ID:
             return
-
         prompt_raw = message.content.strip()
         if not prompt_raw:
             return
 
-        prompt_en = ko2en(prompt_raw)
+        prompt_en = ko2en(prompt_raw)          # 한글 → 영어
         await message.channel.typing()
 
-        # ── Replicate 요청 ─────────────────────────────────
-        def run_replicate():
-            return list(replicate.run(MODEL, input={"prompt": prompt_en}))
+        # ── Gradio 호출 ────────────────────────────────────
+        def generate_image():
+            client = Client(GRADIO_URL)
+            seed = random.randint(0, 2**32 - 1)
+            result = client.predict(
+                prompt=prompt_en,
+                width=768,
+                height=768,
+                guidance=3.5,
+                inference_steps=30,
+                seed=seed,
+                do_img2img=False,
+                init_image=None,            # img2img 사용 안 함
+                image2image_strength=0.8,
+                resize_img=True,
+                api_name=GRADIO_API
+            )
+            return result[0]   # dict (path / url)
 
         try:
-            images = await asyncio.get_running_loop().run_in_executor(None, run_replicate)
+            img_info = await asyncio.get_running_loop().run_in_executor(None, generate_image)
         except Exception as e:
-            logging.error(f"Replicate error: {e}")
+            logging.error(f"Gradio API error: {e}")
             await message.reply("⚠️ 이미지 생성 실패!")
             return
 
-        # ── Discord 전송 ──────────────────────────────────
+        # ── 이미지 Discord 전송 ─────────────────────────────
         files = []
-        for idx, item in enumerate(images):
-            try:
-                data = item.read() if hasattr(item, "read") else requests.get(item).content
-                files.append(discord.File(io.BytesIO(data), filename=f"img_{idx}.webp"))
-            except Exception as e:
-                logging.warning(f"[IMG {idx}] 처리 실패: {e}")
+        try:
+            if img_info.get("path"):                 # 서버가 파일 경로로 반환
+                with open(img_info["path"], "rb") as f:
+                    files.append(discord.File(io.BytesIO(f.read()), filename="generated.webp"))
+            elif img_info.get("url"):                # URL 또는 base64
+                data = requests.get(img_info["url"]).content
+                files.append(discord.File(io.BytesIO(data), filename="generated.webp"))
+        except Exception as e:
+            logging.warning(f"이미지 처리 실패: {e}")
 
         await message.reply(files=files if files else None,
                             content=None if files else "⚠️ 이미지를 전송할 수 없습니다.")
 
 # ── 실행 ────────────────────────────────────────────────────
 if __name__ == "__main__":
-    replicate.Client(api_token=REPL_TOKEN)   # Replicate 인증
+    replicate.Client(api_token=REPL_TOKEN)    # 구조 유지: Replicate 인증(미사용)
     ImageBot(intents=intents).run(TOKEN)