CTP-APP / app.py
cheemam0807's picture
Upload app.py
140212b verified
import streamlit as st
from PIL import Image
from transformers import pipeline
# Load models once
img_caption_model = pipeline("image-to-text", model="Salesforce/blip-image-captioning-base")
story_gen_model = pipeline(
"text-generation",
model="gpt2",
device=-1, # CPU
)
# Image-to-text function
def img2text(image):
caption = img_caption_model(image)[0]["generated_text"]
return caption
# Story generator function
def generator(topic):
prompt = f"Once upon a time, {topic}..."
response = story_gen_model(
prompt,
max_new_tokens=30, # Short story
num_return_sequences=1,
do_sample=True,
temperature=0.7,
)
story = response[0]["generated_text"].replace(prompt, "").strip()
return story
# Streamlit App UI
st.title("Image-to-Story Generator 📸📖")
st.write("Upload an image to get a caption and a short story!")
uploaded_image = st.file_uploader("Choose an image...", type=["png", "jpg", "jpeg"])
if uploaded_image is not None:
image = Image.open(uploaded_image)
st.image(image, caption="Uploaded Image", use_column_width=True)
# Process Image → Caption → Story
with st.spinner('Generating Caption...'):
caption = img2text(image)
st.write(f"**Image Caption:** {caption}")
with st.spinner('Generating Story...'):
story = generator(caption)
st.write("**Generated Story:**")
st.write(story)