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)