app.py

# img_bot.py
import discord, os, io, re, asyncio, logging, requests, replicate, subprocess
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"

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 값을 넣어주세요.")

# Replicate 라이브러리가 참조하도록 토큰 주입
os.environ["REPLICATE_API_TOKEN"] = REPL_TOKEN

# ── 모델 ────────────────────────────────────────────────────
MODEL = (
    "bytedance/sdxl-lightning-4step:"
    "6f7a773af6fc3e8de9d5a3c00be77c17308914bf67772726aff83496ba1e3bbe"
)

# ── 번역 파이프라인 (CPU) ───────────────────────────────────
translator_kwargs = {"device": -1}
if HF_TOKEN:
    translator_kwargs["token"] = HF_TOKEN
translator = transformers_pipeline(
    "translation",
    model="Helsinki-NLP/opus-mt-ko-en",
    **translator_kwargs
)

def ko2en(text: str) -> str:
    """프롬프트에 한글이 있으면 영어로 번역하여 반환."""
    if re.search(r"[가-힣]", text):
        try:
            return translator(text, max_length=512)[0]["translation_text"].strip()
        except Exception as e:
            logging.warning(f"번역 실패, 원문 사용: {e}")
    return text

# ── 로깅 ────────────────────────────────────────────────────
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s",
    handlers=[logging.StreamHandler()]
)

# ── 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)")

class ImageBot(discord.Client):
    async def on_ready(self):
        logging.info(f"Logged in as {self.user} (id={self.user.id})")
        # web.py 병렬 실행
        try:
            subprocess.Popen(["python", "web.py"])
            logging.info("web.py server has been started.")
        except Exception as e:
            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)

        await message.channel.typing()

        # ── Replicate 호출 ──────────────────────────────────
        def run_replicate():
            return list(replicate.run(MODEL, input={"prompt": prompt_en}))

        try:
            images = await asyncio.get_running_loop().run_in_executor(None, run_replicate)
        except Exception as e:
            logging.error(f"Replicate error: {e}")
            await message.reply("⚠️ 이미지 생성 실패!")
            return

        # ── 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"))
            except Exception as e:
                logging.warning(f"[IMG {idx}] 처리 실패: {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 인증
    ImageBot(intents=intents).run(TOKEN)