File size: 5,680 Bytes
64fd107 95bff35 64fd107 95bff35 64fd107 95bff35 64fd107 95bff35 64fd107 95bff35 64fd107 95bff35 64fd107 95bff35 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 |
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:
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):
"""
Converts image bytes into a lower resolution image (256x256 maximum)
and generates a caption.
"""
image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
# Resize to speed up 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):
"""
Generates a humorous and engaging children's story based on the caption.
Uses a prompt to instruct the model and limits token generation to 80 tokens.
"""
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,
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):
"""
Converts the generated story text into audio.
Splits the text into 300-character chunks to reduce repeated TTS calls,
concatenates the resulting audio chunks, and returns 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() # Make sure models are loaded
image_bytes = uploaded_file.getvalue()
# Display the 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}")
|