hskwon7/ISOM5240-Individual-Assignment

1# app.py2 3"""4Streamlit application for Image-to-Story demo with history sidebar.5Allows demo/upload image, generates a caption, a trimmed story,6and plays back as MP3 via gTTS. Keeps history of all runs.7"""8import streamlit as st9from PIL import Image10import warnings11from modules import (12    load_captioner, load_story_gen,13    generate_caption, generate_story_simple,14    generate_audio15)16import io17 18warnings.filterwarnings("ignore", category=DeprecationWarning)19 20# Reset state when switching image source21def reset_state():22    for key in ["caption", "story", "audio_bytes", "audio_mime", "selected_index"]:23        if key in st.session_state:24            del st.session_state[key]25 26def main():27    st.set_page_config(layout="wide")28    st.title("🎨 Magic Picture Story Time!")29    st.write("Pick or upload a picture, and watch it turn into a fun story with voice! Ready for a magical tale?")30 31    # --- Sidebar: History ---32    st.sidebar.header("History")33    if "history" not in st.session_state:34        st.session_state.history = []           # list of dicts35    if "selected_index" not in st.session_state:36        st.session_state.selected_index = None37 38    # Render thumbnails & select buttons39    for idx, entry in enumerate(st.session_state.history):40        with st.sidebar.container():41            st.sidebar.image(entry["image_bytes"], width=100)42            if st.sidebar.button(f"View #{idx+1}", key=f"view_{idx}"):43                st.session_state.selected_index = idx44 45    # Sidebar clear-all button46    if st.sidebar.button("Clear History"):47        st.session_state.history = []48        st.session_state.selected_index = None49 50    # --- Main panel: image selection ---51    source = st.radio("Image source:",52                      ("Upload my own image", "Use demo image"),53                      on_change=reset_state)54 55    # Load pipelines once56    if "models_loaded" not in st.session_state:57        with st.spinner("Loading models…"):58            st.session_state.captioner = load_captioner()59            st.session_state.story_gen = load_story_gen()60        st.session_state.models_loaded = True61 62    captioner = st.session_state.captioner63    story_gen = st.session_state.story_gen64 65    # If user clicked a history entry, load it66    sel = st.session_state.selected_index67    if sel is not None:68        entry = st.session_state.history[sel]69        img = Image.open(io.BytesIO(entry["image_bytes"])).convert("RGB")70        st.image(img, use_container_width=True)71        st.markdown(f"**Caption:** {entry['caption']}")72        st.markdown(f"**Story:** {entry['story']}")73        if st.button("🔊 Play Story Audio"):74            st.audio(data=entry["audio_bytes"], format=entry["audio_mime"])75        return76 77    # Otherwise, handle a fresh upload/demo78    if source == "Use demo image":79        img = Image.open("test_kids_playing.jpg").convert("RGB")80        # grab raw bytes for history81        buf = io.BytesIO()82        img.save(buf, format="JPEG")83        img_bytes = buf.getvalue()84    else:85        uploaded = st.file_uploader("Upload an image",86                                    type=["png", "jpg", "jpeg"])87        if not uploaded:88            return89        img = Image.open(uploaded).convert("RGB")90        img_bytes = uploaded.getvalue()91 92    st.image(img, use_container_width=True)93 94    # Step 1: Caption95    if "caption" not in st.session_state:96        with st.spinner("Captioning image…"):97            st.session_state.caption = generate_caption(captioner, img)98    st.markdown(f"**Caption:** {st.session_state.caption}")99 100    # Step 2: Story101    if "story" not in st.session_state:102        with st.spinner("Creating story…"):103            st.session_state.story = generate_story_simple(104                story_gen, st.session_state.caption, 50, 100105            )106    st.markdown(f"**Story:** {st.session_state.story}")107 108    # Step 3: Audio109    if "audio_bytes" not in st.session_state:110        with st.spinner("Generating audio…"):111            audio_bytes, mime = generate_audio(st.session_state.story)112            st.session_state.audio_bytes = audio_bytes113            st.session_state.audio_mime  = mime114 115    if st.button("🔊 Play Story Audio"):116        st.audio(data=st.session_state.audio_bytes,117                 format=st.session_state.audio_mime)118 119    # Step 4: Append to history (only once per new run)120    if not st.session_state.history or st.session_state.history[-1]["image_bytes"] != img_bytes:121        st.session_state.history.append({122            "image_bytes": img_bytes,123            "caption":     st.session_state.caption,124            "story":       st.session_state.story,125            "audio_bytes": st.session_state.audio_bytes,126            "audio_mime":  st.session_state.audio_mime127        })128 129if __name__ == "__main__":130    main()131