app.py

import discord
import logging
import os
import asyncio
from transformers import PaliGemmaForConditionalGeneration, PaliGemmaProcessor
import torch
import re
import requests
from PIL import Image
import io
import gradio as gr
import threading
from huggingface_hub import InferenceClient

# 로깅 설정
logging.basicConfig(level=logging.DEBUG, format='%(asctime)s:%(levelname)s:%(name)s:%(message)s', handlers=[logging.StreamHandler()])

# 디스코드 인텐트 설정
intents = discord.Intents.default()
intents.message_content = True
intents.messages = True
intents.guilds = True
intents.guild_messages = True

# 추론 API 클라이언트 설정
hf_client = InferenceClient("CohereForAI/aya-23-35B", token=os.getenv("HF_TOKEN"))

# PaliGemma 모델 설정 (CPU 모드)
model = PaliGemmaForConditionalGeneration.from_pretrained("gokaygokay/sd3-long-captioner").to("cpu").eval()
processor = PaliGemmaProcessor.from_pretrained("gokaygokay/sd3-long-captioner")

# 대화 히스토리를 저장할 전역 변수
conversation_history = []

def modify_caption(caption: str) -> str:
    prefix_substrings = [
        ('captured from ', ''),
        ('captured at ', '')
    ]
    pattern = '|'.join([re.escape(opening) for opening, _ in prefix_substrings])
    replacers = {opening: replacer for opening, replacer in prefix_substrings}
    
    def replace_fn(match):
        return replacers[match.group(0)]
    
    return re.sub(pattern, replace_fn, caption, count=1, flags=re.IGNORECASE)

async def create_captions_rich(image: Image.Image) -> str:
    prompt = "caption en"
    image_tensor = processor(images=image, return_tensors="pt").pixel_values.to("cpu")
    image_tensor = (image_tensor * 255).type(torch.uint8)
    model_inputs = processor(text=prompt, images=image_tensor, return_tensors="pt").to("cpu")
    input_len = model_inputs["input_ids"].shape[-1]

    loop = asyncio.get_event_loop()
    generation = await loop.run_in_executor(
        None,
        lambda: model.generate(**model_inputs, max_new_tokens=256, do_sample=False)
    )
    generation = generation[0][input_len:]
    decoded = processor.decode(generation, skip_special_tokens=True)

    modified_caption = modify_caption(decoded)
    return modified_caption

async def translate_to_korean(text: str) -> str:
    messages = [
        {"role": "system", "content": "Translate the following text from English to Korean."},
        {"role": "user", "content": text}
    ]

    loop = asyncio.get_event_loop()
    response = await loop.run_in_executor(
        None,
        lambda: hf_client.chat_completion(
            messages, max_tokens=1000, stream=True, temperature=0.7, top_p=0.85
        )
    )

    full_response = []
    for part in response:
        if part.choices and part.choices[0].delta and part.choices[0].delta.content:
            full_response.append(part.choices[0].delta.content)

    full_response_text = ''.join(full_response)
    return full_response_text.strip()

async def interact_with_model(user_input: str) -> str:
    global conversation_history
    conversation_history.append({"role": "user", "content": user_input})
    
    messages = [
        {"role": "system", "content": "Translate the following text from English to Korean and respond as if you are an assistant who provides detailed answers in Korean."},
    ] + conversation_history

    loop = asyncio.get_event_loop()
    response = await loop.run_in_executor(
        None,
        lambda: hf_client.chat_completion(
            messages, max_tokens=1000, stream=True, temperature=0.7, top_p=0.85
        )
    )

    full_response = []
    for part in response:
        if part.choices and part.choices[0].delta and part.choices[0].delta.content:
            full_response.append(part.choices[0].delta.content)

    full_response_text = ''.join(full_response)
    conversation_history.append({"role": "assistant", "content": full_response_text})
    return full_response_text.strip()

# Gradio 인터페이스 설정
def create_captions_rich_sync(image):
    caption = asyncio.run(create_captions_rich(image))
    translated_caption = asyncio.run(translate_to_korean(caption))
    return translated_caption

css = """
  #mkd {
    height: 500px; 
    overflow: auto; 
    border: 1px solid #ccc; 
  }
"""

with gr.Blocks(css=css) as demo:
    gr.HTML("<h1><center>PaliGemma Fine-tuned for Long Captioning<center><h1>")
    with gr.Tab(label="PaliGemma Long Captioner"):
        with gr.Row():
            with gr.Column():
                input_img = gr.Image(label="Input Picture")
                submit_btn = gr.Button(value="Submit")
            output = gr.Text(label="Caption")

        submit_btn.click(create_captions_rich_sync, [input_img], [output])

# Gradio 서버를 비동기적으로 실행
def run_gradio():
    demo.launch(
        server_name="0.0.0.0", 
        server_port=int(os.getenv("GRADIO_SERVER_PORT", 7861)), 
        inbrowser=True
    )

# 특정 채널 ID 설정
SPECIFIC_CHANNEL_ID = int(os.getenv("DISCORD_CHANNEL_ID", "123456789012345678"))

# 디스코드 봇 설정
class MyClient(discord.Client):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.is_processing = False

    async def on_ready(self):
        logging.info(f'{self.user}로 로그인되었습니다!')
        threading.Thread(target=run_gradio, daemon=True).start()
        logging.info("Gradio 서버가 시작되었습니다.")

    async def on_message(self, message):
        if message.author == self.user:
            return
        if not self.is_message_in_specific_channel(message):
            return
        if self.is_processing:
            return
        self.is_processing = True
        try:
            if message.attachments:
                image_url = message.attachments[0].url
                response = await process_image(image_url, message)
                await message.channel.send(response)
            else:
                response = await interact_with_model(message.content)
                await message.channel.send(response)
        finally:
            self.is_processing = False

    def is_message_in_specific_channel(self, message):
        return message.channel.id == SPECIFIC_CHANNEL_ID or (
            isinstance(message.channel, discord.Thread) and message.channel.parent_id == SPECIFIC_CHANNEL_ID
        )

async def process_image(image_url, message):
    image = await download_image(image_url)
    caption = await create_captions_rich(image)
    translated_caption = await translate_to_korean(caption)
    intro_message = f"{message.author.mention}, 인식된 이미지 설명: {translated_caption}\n\n질문이 있으면 물어보세요!"
    return intro_message

async def download_image(url):
    response = requests.get(url)
    image = Image.open(io.BytesIO(response.content)).convert("RGB")
    return image

if __name__ == "__main__":
    discord_client = MyClient(intents=intents)
    discord_client.run(os.getenv('DISCORD_TOKEN'))