Jey813's picture
Upload app.py
7fb3a4f verified
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()