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

# Device 설정
device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"🚀 Using device: {device}")

# 모델 로딩
print("📦 Loading BLIP model...")
processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base").to(device)
print("✅ Model loaded successfully!")

# 후킹 템플릿
def generate_hook_title(caption):
    templates = [
        f"You won't believe this: {caption}",
        f"This is what happens when {caption.lower()}",
        f"{caption}? Now that's a twist!",
        f"{caption} – but it’s not what you think!",
        f"When {caption.lower()}, something unexpected happens 😮",
    ]
    return random.choice(templates)

# 처리 함수
def process_image(image):
    inputs = processor(image, return_tensors="pt").to(device)
    outputs = model.generate(**inputs)
    caption = processor.decode(outputs[0], skip_special_tokens=True)
    hook_title = generate_hook_title(caption)
    return hook_title

# Gradio 인터페이스
demo = gr.Interface(
    fn=process_image,
    inputs=gr.Image(type="pil", label="📸 Upload your YouTube thumbnail"),
    outputs=gr.Textbox(label="🔥 Catchy English Title"),
    title="🎬 YouTube Thumbnail Hook Title Generator",
    description="Upload a thumbnail image and get a catchy, AI-generated English title!"
)

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