Update app.py

app.py

@@ -1,10 +1,15 @@
 # img_bot.py
-import discord, os, io, re, random, asyncio, logging, requests, replicate, subprocess, base64
+import os, io, re, random, asyncio, logging, subprocess, base64
 from urllib.parse import urljoin, quote_plus
+
+import discord
+import requests
+import replicate
 from transformers import pipeline as transformers_pipeline
 from gradio_client import Client
+from PIL import Image     # WEBP → PNG 변환용
 
-# ── 환경 변수 ────────────────────────────────────────────────
+# ────────────────── 환경 변수 ──────────────────
 TOKEN       = os.getenv("DISCORD_TOKEN")
 CHANNEL_ID  = int(os.getenv("DISCORD_CHANNEL_ID"))
 REPL_TOKEN  = (os.getenv("OPENAI_API_KEY") or "").strip()
@@ -14,14 +19,13 @@ 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 미사용이지만 구조 유지
 
-os.environ["REPLICATE_API_TOKEN"] = REPL_TOKEN            # 구조 유지 (Replicate 미사용)
-
-# ── Gradio 서버 ─────────────────────────────────────────────
-GRADIO_URL  = "http://211.233.58.201:7971"
-GRADIO_API  = "/process_and_save_image"
+# ────────────────── Gradio 서버 ─────────────────
+GRADIO_URL = "http://211.233.58.201:7971"
+GRADIO_API = "/process_and_save_image"
 
-# ── 번역 파이프라인 (CPU) ───────────────────────────────────
+# ────────────────── 번역 파이프라인 ─────────────
 translator = transformers_pipeline(
     "translation",
     model="Helsinki-NLP/opus-mt-ko-en",
@@ -30,7 +34,7 @@ translator = transformers_pipeline(
 )
 
 async def ko2en_async(text: str) -> str:
-    """한글 포함 시 비동기 영어 번역."""
+    """한글 포함 시 영어 번역(비동기)."""
     if not re.search(r"[가-힣]", text):
         return text
     loop = asyncio.get_running_loop()
@@ -43,12 +47,14 @@ async def ko2en_async(text: str) -> str:
         logging.warning(f"번역 실패, 원문 사용: {e}")
         return text
 
-# ── 로깅 ────────────────────────────────────────────────────
-logging.basicConfig(level=logging.INFO,
+# ────────────────── 로깅 ────────────────────────
+logging.basicConfig(
+    level=logging.INFO,
     format="%(asctime)s [%(levelname)s] %(message)s",
-    handlers=[logging.StreamHandler()])
+    handlers=[logging.StreamHandler()]
+)
 
-# ── Discord 인텐트 ──────────────────────────────────────────
+# ────────────────── Discord 클라이언트 ──────────
 intents = discord.Intents.default()
 intents.message_content = True
 
@@ -72,7 +78,7 @@ class ImageBot(discord.Client):
         prompt_en = await ko2en_async(prompt_raw)
         await message.channel.typing()
 
-        # ── Gradio 호출 ────────────────────────────────────
+        # ───────────── Gradio 호출 ──────────────
         def generate_image():
             client = Client(GRADIO_URL)
             return client.predict(
@@ -83,75 +89,69 @@ class ImageBot(discord.Client):
                 prompt=prompt_en,
                 seed=random.randint(0, 2**32 - 1),
                 api_name=GRADIO_API
-            )  # ❗ 리스트나 dict·str 그대로 반환
+            )
 
         try:
-            img_info = await asyncio.get_running_loop().run_in_executor(None, generate_image)
-            logging.info(f"Gradio result type={type(img_info)} → {img_info}")
+            result = await asyncio.get_running_loop().run_in_executor(None, generate_image)
         except Exception as e:
             logging.error(f"Gradio API error: {e}")
             await message.reply("⚠️ 이미지 생성 실패!")
             return
 
-        # ── Discord 파일 준비 ──────────────────────────────
-        data = None
-        filename = "generated.webp"
+        # Gradio 결과 → list / dict / str 중 첫 요소로 정규화
+        if isinstance(result, list):
+            result = result[0]
+
+        # ────────────��� URL 있으면 링크로 전송 ─────────────
+        url = None
+        if isinstance(result, dict):
+            url = result.get("url")
+            if url and url.startswith("/"):
+                url = urljoin(GRADIO_URL, url)
+        elif isinstance(result, str) and result.startswith(("http", "/")):
+            url = result if result.startswith("http") else urljoin(GRADIO_URL, result)
+
+        if url:
+            await message.reply(content=url)   # Discord가 embed 미리보기
+            return
 
-        def download(url: str) -> bytes | None:
-            try:
-                return requests.get(url, timeout=10).content
-            except Exception as err:
-                logging.warning(f"URL 다운로드 실패({url}): {err}")
-                return None
-
-        # ① `img_info` 가 리스트면 첫 요소 사용
-        if isinstance(img_info, list) and img_info:
-            img_info = img_info[0]
-
-        # ② dict
-        if isinstance(img_info, dict):
-            # url 우선
-            url = img_info.get("url")
-            if url:
-                if url.startswith("data:"):
-                    try:
-                        data = base64.b64decode(re.sub(r"^data:image/[^;]+;base64,", "", url))
-                    except Exception as e:
-                        logging.warning(f"base64 디코딩 실패: {e}")
-                else:
-                    if url.startswith("/"):
-                        url = urljoin(GRADIO_URL, url)
-                    data = download(url)
-            # path 보조
-            if data is None:
-                path = img_info.get("path")
-                if path:
-                    # 원격 Gradio 파일 end-point
-                    remote_url = urljoin(GRADIO_URL, f"/gradio_api/file={quote_plus(path)}")
-                    data = download(remote_url) or (open(path, "rb").read() if os.path.isfile(path) else None)
-
-        # ③ str (로컬 경로/URL/data URI)
-        elif isinstance(img_info, str):
-            s = img_info.strip()
-            if s.startswith("data:"):
+        # ───────────── 첨부 파일 준비 ─────────────
+        data = None
+        # dict → path
+        if isinstance(result, dict):
+            path = result.get("path")
+            if path:
+                # Gradio 파일 다운로드 우선
+                remote = urljoin(GRADIO_URL, f"/gradio_api/file={quote_plus(path)}")
                 try:
-                    data = base64.b64decode(re.sub(r"^data:image/[^;]+;base64,", "", s))
-                except Exception as e:
-                    logging.warning(f"base64 디코딩 실패: {e}")
-            elif s.startswith("http"):
-                data = download(s)
-            else:
-                # 원격 파일로 간주
-                remote_url = urljoin(GRADIO_URL, f"/gradio_api/file={quote_plus(s)}")
-                data = download(remote_url) or (open(s, "rb").read() if os.path.isfile(s) else None)
-
-        # ── Discord 전송 ──────────────────────────────────
-        if data:
-            await message.reply(files=[discord.File(io.BytesIO(data), filename=filename)])
-        else:
+                    data = requests.get(remote, timeout=10).content
+                except Exception:
+                    if os.path.isfile(path):
+                        with open(path, "rb") as f:
+                            data = f.read()
+        # str → data URI
+        elif isinstance(result, str) and result.startswith("data:"):
+            try:
+                b64 = re.sub(r"^data:image/[^;]+;base64,", "", result)
+                data = base64.b64decode(b64)
+            except Exception as e:
+                logging.warning(f"base64 디코딩 실패: {e}")
+
+        if not data:
             await message.reply("⚠️ 이미지를 전송할 수 없습니다.")
+            return
 
-# ── 실행 ────────────────────────────────────────────────────
+        # ───────────── WEBP → PNG 변환 (미리보기용) ─────────────
+        buf = io.BytesIO()
+        try:
+            Image.open(io.BytesIO(data)).convert("RGB").save(buf, format="PNG")
+            buf.seek(0)
+            await message.reply(files=[discord.File(buf, filename="image.png")])
+        except Exception:
+            # 변환 실패 시 원본 그대로
+            await message.reply(files=[discord.File(io.BytesIO(data), filename="image.webp")])
+
+# ────────────────── 실행 ────────────────────────
 if __name__ == "__main__":
     replicate.Client(api_token=REPL_TOKEN)  # 구조 유지
     ImageBot(intents=intents).run(TOKEN)