# app.py

import gradio as gr
import torch
from PIL import Image
from transformers import AutoProcessor, AutoModelForCausalLM

# 1. 장치 설정
device = "cuda" if torch.cuda.is_available() else "cpu"

# 2. Florence 모델 및 프로세서 로드
florence_model = AutoModelForCausalLM.from_pretrained(
    "microsoft/Florence-2-base",
    trust_remote_code=True
).to(device).eval()
florence_processor = AutoProcessor.from_pretrained(
    "microsoft/Florence-2-base",
    trust_remote_code=True
)

# 3. 이미지 설명 생성 함수
def generate_caption(image):
    if not isinstance(image, Image.Image):
        image = Image.fromarray(image)

    inputs = florence_processor(
        text="<MORE_DETAILED_CAPTION>",
        images=image,
        return_tensors="pt"
    ).to(device)
    generated_ids = florence_model.generate(
        input_ids=inputs["input_ids"],
        pixel_values=inputs["pixel_values"],
        max_new_tokens=1024,
        early_stopping=False,
        do_sample=False,
        num_beams=3,
    )
    generated_text = florence_processor.batch_decode(
        generated_ids,
        skip_special_tokens=False
    )[0]
    parsed_answer = florence_processor.post_process_generation(
        generated_text,
        task="<MORE_DETAILED_CAPTION>",
        image_size=(image.width, image.height)
    )
    prompt = parsed_answer["<MORE_DETAILED_CAPTION>"]
    print("Generation completed!:", prompt)
    return prompt

# 4. Gradio 인터페이스 (캐리커쳐 버튼 포함)
with gr.Blocks(theme="Yntec/HaleyCH_Theme_Orange") as demo:
    gr.Markdown("## 🖼️ 이미지 설명 생성기")
    gr.Markdown(
        "⚠ 현재 CPU 모드로 실행 중이므로 속도가 느릴 수 있습니다. 양해 부탁드립니다."
    )

    with gr.Row():
        with gr.Column():
            image_input = gr.Image(label="Input Image", type="pil")
        with gr.Column():
            caption_output = gr.Textbox(
                label="Output Prompt",
                lines=6,                # 이전보다 2배 높이
                show_copy_button=True
            )
            # 오른쪽 하단 '캐리커쳐 만들기' 버튼
            gr.HTML("""
            <div style="margin-top:10px; text-align:center;">
              <a href="https://huggingface.co/spaces/VIDraft/stable-diffusion-3.5-large-turboX" target="_blank">
                <button style="
                  padding:10px 20px;
                  background-color:#ff9900;
                  color:white;
                  border:none;
                  border-radius:10px;
                  font-size:16px;
                  box-shadow:2px 2px 8px rgba(0,0,0,0.3);
                  cursor:pointer;
                ">
                  🎨 캐리커쳐 만들기
                </button>
              </a>
            </div>
            """)

    image_input.upload(fn=generate_caption, inputs=image_input, outputs=caption_output)

if __name__ == "__main__":
    demo.launch(debug=True)