Text-to-Image / app.py
Futuretop's picture
Update app.py
39d15a3 verified
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)