import gradio as gr
import torch
from PIL import Image
from transformers import BlipProcessor, BlipForConditionalGeneration

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

# 2. 모델 및 프로세서 로드
processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base").to(device)

# 3. 이미지 설명 생성 함수
def generate_caption(image):
    if image is None:
        return "이미지를 업로드해주세요."
    
    # 고속 처리를 위한 리사이즈
    image = image.resize((384, 384))
    
    # 설명 생성
    inputs = processor(images=image, return_tensors="pt").to(device)
    output_ids = model.generate(**inputs, max_length=50)
    caption = processor.decode(output_ids[0], skip_special_tokens=True)
    print("✅ 생성된 설명:", caption)
    if "Asian" in caption:
        caption = caption.replace("Asian", "Korean")
    print("✅ 생성된 설명:", caption)
    return caption
    return caption

# 4. Gradio 인터페이스 구성
with gr.Blocks(title="이미지 설명 생성기") as demo:
    gr.Markdown("## 🖼️ 이미지를 업로드하면 설명이 자동 생성됩니다.")

    with gr.Row():
        with gr.Column():
            image_input = gr.Image(label="입력 이미지", type="pil")
        with gr.Column():
            caption_output = gr.Textbox(label="생성된 설명", lines=3, show_copy_button=True)
            # HTML로 버튼 생성
            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)

# 5. 앱 실행
demo.launch(debug=True)