app.py

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

# Replicate 라이브러리가 참조하도록 토큰 주입 (비록 사용 안 하지만 구조 유지)
os.environ["REPLICATE_API_TOKEN"] = REPL_TOKEN

# ── Gradio 서버 엔드포인트 ──────────────────────────────────
GRADIO_URL = "http://211.233.58.201:7896"           # 서버 주소
GRADIO_API = "/generate_image"                      # 사용할 함수 이름

# ── 번역 파이프라인 (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")

# ── Discord 인텐트 ──────────────────────────────────────────
intents = discord.Intents.default()
intents.message_content = True   # 메시지 읽기 (Portal에서 ON 필요)

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):
        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()

        # ── 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:
            img_info = 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 전송 ─────────────────────────────
        files = []
        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 인증(미사용)
    ImageBot(intents=intents).run(TOKEN)