File size: 5,935 Bytes
64fd107 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 |
import streamlit as st
from transformers import pipeline
from PIL import Image
import io, textwrap, numpy as np, soundfile as sf
# ------------------ Streamlit Page Configuration ------------------
st.set_page_config(
page_title="Picture to Story Magic", # App title on browser tab
page_icon="🦄", # Fun unicorn icon
layout="centered"
)
# ------------------ Custom CSS for a Colorful Background ------------------
st.markdown(
"""
<style>
body {
background-color: #FDEBD0; /* A soft pastel color */
}
</style>
""",
unsafe_allow_html=True
)
# ------------------ Playful Header for Young Users ------------------
st.markdown(
"""
<h1 style='text-align: center; color: #ff66cc;'>Picture to Story Magic!</h1>
<p style='text-align: center; font-size: 24px;'>
Hi little artist! Upload your picture and let us create a fun story just for you! 🎉
</p>
""",
unsafe_allow_html=True
)
# ------------------ Lazy Model Loading ------------------
def load_models():
"""
Lazy-load the required pipelines and store them in session state.
Pipelines:
1. Captioner: Generates descriptive text from an image using a lighter model.
2. Storyer: Generates a humorous children's story using aspis/gpt2-genre-story-generation.
3. TTS: Converts text into audio.
"""
if "captioner" not in st.session_state:
# Use the "base" version for faster/cost-effective captioning.
st.session_state.captioner = pipeline(
"image-to-text",
model="Salesforce/blip-image-captioning-base"
)
if "storyer" not in st.session_state:
st.session_state.storyer = pipeline(
"text-generation",
model="aspis/gpt2-genre-story-generation"
)
if "tts" not in st.session_state:
st.session_state.tts = pipeline(
"text-to-speech",
model="facebook/mms-tts-eng"
)
# ------------------ Caching Functions ------------------
@st.cache_data(show_spinner=False)
def get_caption(image_bytes):
"""
Convert the image bytes into a smaller image to speed up captioning,
then return the generated caption.
"""
image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
# Resize the image (preserving aspect ratio) to only 256x256 for faster processing.
image.thumbnail((256, 256))
caption = st.session_state.captioner(image)[0]["generated_text"]
return caption
@st.cache_data(show_spinner=False)
def get_story(caption):
"""
Generate a humorous and engaging children's story using the caption.
The prompt instructs the model to produce a playful story (50-100 words).
We lower max_new_tokens to 80 so that it generates its text faster.
"""
prompt = (
f"Write a funny, warm, and imaginative children's story for ages 3-10, 50-100 words, "
f"in third-person narrative, as if the author is playfully describing the scene in the image: {caption}. "
"Explicitly mention the exact venue or location (such as a park, school, or home), describe specific characters "
"(for example, a little girl named Lily or a boy named Jack), and detail the humorous actions they perform. "
"Ensure the story is playful, engaging, and ends with a complete sentence."
)
raw_story = st.session_state.storyer(
prompt,
max_new_tokens=80, # Reduced token generation for faster response
do_sample=True,
temperature=0.7,
top_p=0.9,
return_full_text=False
)[0]["generated_text"].strip()
words = raw_story.split()
return " ".join(words[:100])
@st.cache_data(show_spinner=False)
def get_audio(story):
"""
Convert the generated story text into audio.
The text is split into 300-character chunks to reduce repeated TTS calls,
the audio chunks are concatenated, and then stored in an in-memory WAV buffer.
"""
chunks = textwrap.wrap(story, width=300)
audio_chunks = [st.session_state.tts(chunk)["audio"].squeeze() for chunk in chunks]
audio = np.concatenate(audio_chunks)
buffer = io.BytesIO()
sf.write(buffer, audio, st.session_state.tts.model.config.sampling_rate, format="WAV")
buffer.seek(0)
return buffer
# ------------------ Main App Logic ------------------
uploaded_file = st.file_uploader("Choose a Picture...", type=["jpg", "jpeg", "png"])
if uploaded_file is not None:
try:
load_models() # Ensure models are loaded once
image_bytes = uploaded_file.getvalue()
# Display the user-uploaded image
image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
st.image(image, caption="Your Amazing Picture!", use_column_width=True)
st.markdown("<h3 style='text-align: center;'>Ready for your story?</h3>", unsafe_allow_html=True)
if st.button("Story, Please!"):
with st.spinner("Generating caption..."):
caption = get_caption(image_bytes)
st.markdown("<h3 style='text-align: center;'>Caption:</h3>", unsafe_allow_html=True)
st.write(caption)
with st.spinner("Generating story..."):
story = get_story(caption)
st.markdown("<h3 style='text-align: center;'>Your Story:</h3>", unsafe_allow_html=True)
st.write(story)
with st.spinner("Generating audio..."):
audio_buffer = get_audio(story)
st.audio(audio_buffer, format="audio/wav", start_time=0)
st.markdown(
"<p style='text-align: center; font-weight: bold;'>Enjoy your magical story! 🎶</p>",
unsafe_allow_html=True
)
except Exception as e:
st.error("Oops! Something went wrong. Please try a different picture or check the file format!")
st.error(f"Error details: {e}")
|