Update app.py

app.py

@@ -1,4 +1,3 @@
-# app.py
 import streamlit as st
 from PIL import Image
 from io import BytesIO
@@ -6,11 +5,10 @@ from huggingface_hub import InferenceApi
 from gtts import gTTS
 import tempfile
 
-# —––––––– Page Config
 st.set_page_config(page_title="Magic Story Generator", layout="centered")
 st.title("📖✨ Turn Images into Children's Stories")
 
-# —––––––– Clients (cached)
+# —––––––– Clients (cached) —–––––––
 @st.cache_resource
 def load_clients():
     hf_token = st.secrets["HF_TOKEN"]
@@ -21,52 +19,42 @@ def load_clients():
 
 caption_client, story_client = load_clients()
 
-# —––––––– Modified image processing (safe version) —–––––––
+def generate_caption(img):
+    """
+    Runs the BLIP caption model on a PIL.Image and returns the generated text.
+    """
+    img_bytes = BytesIO()
+    img.save(img_bytes, format="JPEG")
+    try:
+        result = caption_client(data=img_bytes.getvalue())
+        if isinstance(result, list) and result:
+            return result[0].get("generated_text", "").strip()
+        return ""
+    except Exception as e:
+        st.error(f"Caption generation error: {e}")
+        return ""
+
 def process_image(uploaded_file):
     try:
         img = Image.open(uploaded_file).convert("RGB")
-        
-        # Only resize if absolutely necessary
-        if img.size[0] > 2048 or img.size[1] > 2048:
+        if max(img.size) > 2048:
             img.thumbnail((2048, 2048))
-            
         return img
     except Exception as e:
-        st.error(f"Image processing error: {str(e)}")
+        st.error(f"Image processing error: {e}")
         st.stop()
-# —––––––– Main flow —–––––––
+
 uploaded = st.file_uploader("Upload an image:", type=["jpg", "png", "jpeg"])
 if uploaded:
     img = process_image(uploaded)
-    st.image(img, use_column_width=True)
-    
+    st.image(img, use_container_width=True)
+
     with st.spinner("Generating caption..."):
         caption = generate_caption(img)
-        
-        # Your original validation
         if not caption:
             st.error("😢 Couldn't understand this image. Try another one!")
             st.stop()
-            
     st.success(f"**Caption:** {caption}")
-    
-
-    # Generate Caption
-    with st.spinner("🔍 Discovering image secrets..."):
-        try:
-            img_bytes = BytesIO()
-            img.save(img_bytes, format="JPEG")
-            caption_response = caption_client(data=img_bytes.getvalue())
-            caption = caption_response[0]['generated_text'].strip() if isinstance(caption_response, list) else ""
-            
-            if not caption:
-                st.error("😢 Couldn't understand this image. Try another one!")
-                st.stop()
-        except Exception as e:
-            st.error(f"🚨 Oops! Problem making caption: {str(e)}")
-            st.stop()
-
-    st.success(f"**Caption Magic:** {caption}")
 
     # Story Generation Prompt
     story_prompt = (