Update app.py

app.py

@@ -1,33 +1,32 @@
 # img_bot.py
 import discord, os, io, re, asyncio, logging, requests, replicate, subprocess
-from transformers import pipeline as transformers_pipeline
+from transformers import pipeline as transformers_pipeline   # 번역 파이프라인
 
 # ── 환경 변수 ────────────────────────────────────────────────
-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 토큰
-# message_content 인텐트 사용 여부(1 = ON, 0 = OFF)
-USE_MSG_INTENT = os.getenv("MESSAGE_CONTENT_INTENT", "1") != "0"
+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 값을 넣어주세요.")
+    raise RuntimeError("OPENAI_API_KEY 환경 변수에 Replicate Personal Access Token 값을 넣어주세요.")
 
 # Replicate 라이브러리가 참조하도록 토큰 주입
 os.environ["REPLICATE_API_TOKEN"] = REPL_TOKEN
 
-# ── 모델 ────────────────────────────────────────────────────
+# ── 이미지 생성 모델 (변경된 부분) ────────────────────────────
 MODEL = (
-    "bytedance/sdxl-lightning-4step:"
-    "6f7a773af6fc3e8de9d5a3c00be77c17308914bf67772726aff83496ba1e3bbe"
+    "nvidia/sana-sprint-1.6b:"
+    "6ed1ce77cdc8db65550e76d5ab82556d0cb31ac8ab3c4947b168a0bda7b962e4"
 )
 
 # ── 번역 파이프라인 (CPU) ───────────────────────────────────
 translator_kwargs = {"device": -1}
 if HF_TOKEN:
-    translator_kwargs["token"] = HF_TOKEN
+    translator_kwargs["token"] = HF_TOKEN        # 인증 토큰 전달
+
 translator = transformers_pipeline(
     "translation",
     model="Helsinki-NLP/opus-mt-ko-en",
@@ -35,7 +34,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()
@@ -52,12 +51,7 @@ logging.basicConfig(
 
 # ── Discord 설정 ────────────────────────────────────────────
 intents = discord.Intents.default()
-if USE_MSG_INTENT:
-    intents.message_content = True
-    logging.info("Message Content Intent = ON  "
-                 "(Discord Developer Portal → Bot → Privileged Intents 에서 활성화 필요)")
-else:
-    logging.info("Message Content Intent = OFF  (MESSAGE_CONTENT_INTENT=0)")
+intents.message_content = True  # 메시지 콘텐츠 읽기
 
 class ImageBot(discord.Client):
     async def on_ready(self):
@@ -70,23 +64,29 @@ class ImageBot(discord.Client):
             logging.warning(f"web.py 실행 실패: {e}")
 
     async def on_message(self, message: discord.Message):
-        # 인텐트 OFF 시 프롬프트를 읽을 수 없으므로 즉시 반환
-        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}))
+            # 모델에 따라 필요한 입력 파라미터 이름이 다를 수 있습니다.
+            # 일반적으로 "prompt" 키를 사용하지만 필요하면 수정하세요.
+            try:
+                result = replicate.run(MODEL, input={"prompt": prompt_en})
+            except replicate.exceptions.ReplicateError:
+                # 파라미터 이름이 다를 경우 빈 dict로 재시도
+                result = replicate.run(MODEL, input={})
+
+            # 결과가 리스트가 아닐 수 있으므로 리스트로 감싸 반환
+            return result if isinstance(result, list) else [result]
 
         try:
             images = await asyncio.get_running_loop().run_in_executor(None, run_replicate)
@@ -95,12 +95,12 @@ class ImageBot(discord.Client):
             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}.png"))
+                files.append(discord.File(io.BytesIO(data), filename=f"img_{idx}.jpg"))
             except Exception as e:
                 logging.warning(f"[IMG {idx}] 처리 실패: {e}")