Spaces:
Sleeping
Sleeping
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) |