LTH001's picture
Update app.py
cb3f72b verified
raw
history blame
1.76 kB
# import part
!pip install streamlit pyngrok
import streamlit as st
from transformers import pipeline
from PIL import Image
import io
# function part
def generate_image_caption(image):
"""Generates a caption for the given image using a pre-trained model."""
img2caption = pipeline("image-to-text", model="Salesforce/blip-image-captioning-base")
result = img2caption(image)
return result[0]['generated_text']
def text2story(text):
"""Generates a children's story from text input with genre adaptation"""
story_prompt = f"Create a funny 100-word story for 8-year-olds about: {text}. Include: "
story_prompt += "1) A silly character 2) Magical object 3) Sound effects 4) Happy ending"
pipe = pipeline("text-generation", model="pranavpsv/genre-story-generator-v2")
story_text = pipe(
story_prompt,
max_new_tokens=200,
temperature=0.9,
top_k=50
)[0]['generated_text']
return story_text.split("Happy ending")[-1].strip() # Clean output
def main():
st.title("πŸ“– Image Story Generator")
st.write("Upload an image and get a magical children's story!")
uploaded_image = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
if uploaded_image:
image = Image.open(uploaded_image).convert("RGB")
st.image(image, use_column_width=True)
with st.spinner("✨ Analyzing image..."):
caption = generate_image_caption(image)
st.subheader("Image Understanding")
st.write(caption)
with st.spinner("πŸ“– Writing story..."):
story = text2story(caption)
st.subheader("Magical Story")
st.write(story)
if __name__ == "__main__":
main()