Update app.py

app.py

@@ -4,42 +4,51 @@ from transformers import VisionEncoderDecoderModel, ViTImageProcessor, AutoToken
 import torch
 from TTS.api import TTS
 
-# Set page config
-st.set_page_config(page_title="Image Storytelling for Kids", layout="wide")
-
-st.title("🧒📖 AI Image Storytelling")
-st.write("Upload an image, and let AI generate a story with voice for children aged 3–10.")
+# Set page configuration
+st.set_page_config(page_title="Children's Image Storytelling", layout="wide")
 
 # Load models
 @st.cache_resource
 def load_models():
+    # Load image captioning model
     vision_model = VisionEncoderDecoderModel.from_pretrained("nlpconnect/vit-gpt2-image-captioning")
     processor = ViTImageProcessor.from_pretrained("nlpconnect/vit-gpt2-image-captioning")
     tokenizer = AutoTokenizer.from_pretrained("nlpconnect/vit-gpt2-image-captioning")
+    # Load text-to-speech model
     tts = TTS(model_name="tts_models/en/ljspeech/tacotron2-DDC", progress_bar=False, gpu=torch.cuda.is_available())
     return vision_model, processor, tokenizer, tts
 
-vision_model, processor, tokenizer, tts_model = load_models()
+# Main function
+def main():
+    # Display title
+    st.title("🧒📖 AI Image Storytelling")
+    st.write("Upload an image, and let AI generate a story for children aged 3–10 with voice narration.")
+
+    # Upload image
+    uploaded_file = st.file_uploader("Upload an image", type=["jpg", "jpeg", "png"])
+
+    if uploaded_file:
+        image = Image.open(uploaded_file).convert("RGB")
+        st.image(image, caption="Uploaded Image", use_column_width=True)
 
-# Upload image
-uploaded_file = st.file_uploader("Upload an image", type=["jpg", "jpeg", "png"])
+        if st.button("Generate Story"):
+            vision_model, processor, tokenizer, tts_model = load_models()
 
-if uploaded_file:
-    image = Image.open(uploaded_file).convert("RGB")
-    st.image(image, caption="Uploaded Image", use_column_width=True)
+            with st.spinner("Generating description..."):
+                pixel_values = processor(images=image, return_tensors="pt").pixel_values
+                output_ids = vision_model.generate(pixel_values, max_length=50, num_beams=4)
+                caption = tokenizer.decode(output_ids[0], skip_special_tokens=True)
+                st.success("Image Description: " + caption)
 
-    if st.button("Generate Story"):
-        with st.spinner("Generating description..."):
-            pixel_values = processor(images=image, return_tensors="pt").pixel_values
-            output_ids = vision_model.generate(pixel_values, max_length=50, num_beams=4)
-            caption = tokenizer.decode(output_ids[0], skip_special_tokens=True)
-            st.success("Image Description: " + caption)
+            with st.spinner("Generating story..."):
+                story_prompt = f"Based on the following description, tell me a short children's story: {caption}"
+                story = caption + " Once upon a time, " + caption.lower() + " entered a magical forest and met many new friends."
+                st.success("Story: " + story)
 
-        with st.spinner("Creating story..."):
-            story_prompt = f"Tell a short, friendly children's story based on: {caption}"
-            story = caption + " Once upon a time, " + caption.lower() + " went on an adventure and made new friends in a magical forest."
-            st.success("Story: " + story)
+            with st.spinner("Generating voice..."):
+                tts_model.tts_to_file(text=story, file_path="story.wav")
+                st.audio("story.wav", format="audio/wav")
 
-        with st.spinner("Generating voice..."):
-            tts_model.tts_to_file(text=story, file_path="story.wav")
-            st.audio("story.wav", format="audio/wav")
+# Run the main program
+if __name__ == "__main__":
+    main()