Update app.py

app.py

@@ -40,17 +40,22 @@ else:
 
     # 2) Generate caption
     with st.spinner("🔍 Generating caption..."):
-        buf = BytesIO()
-        img.save(buf, format="PNG")
-        cap_out = caption_client(data=buf.getvalue())
-
-        # Handle caption response
-        if isinstance(cap_out, list) and cap_out:
-            cap_text = cap_out[0].get("generated_text", "").strip()
-        elif isinstance(cap_out, dict):
-            cap_text = cap_out.get("generated_text", "").strip()
-        else:
-            cap_text = str(cap_out).strip()
+        try:
+            buf = BytesIO()
+            img.save(buf, format="PNG")
+            cap_out = caption_client(data=buf.getvalue())
+
+            # Handle caption response
+            if isinstance(cap_out, list) and cap_out:
+                cap_text = cap_out[0].get("generated_text", "").strip()
+            elif isinstance(cap_out, dict):
+                cap_text = cap_out.get("generated_text", "").strip()
+            else:
+                cap_text = str(cap_out).strip()
+                
+        except Exception as e:
+            st.error(f"🚨 Caption generation failed: {str(e)}")
+            st.stop()
 
     if not cap_text:
         st.error("😕 Couldn’t generate a caption. Try another image.")
@@ -68,36 +73,40 @@ else:
         "Story:"
     )
 
-   # 4) Generate story with proper parameter structure
-with st.spinner("✍️ Generating story..."):
-    try:
-        # Generation parameters must be nested under 'parameters'
-        story_out = story_client(
-            inputs=prompt,
-            params={
-                "max_new_tokens": 150,
-                "do_sample": True,
-                "temperature": 0.8,
-                "top_p": 0.9,
-                "top_k": 50,
-                "repetition_penalty": 1.1,
-                "no_repeat_ngram_size": 2
-            }
-        )
-
-        # Handle different response formats
-        if isinstance(story_out, list):
-            story_text = story_out[0].get("generated_text", "")
-        else:
-            story_text = story_out.get("generated_text", "")
-        
-        # Extract story content after the prompt
-        story = story_text.split("Story:")[-1].strip()
-
-    except Exception as e:
-        st.error(f"🚨 Story generation failed: {str(e)}")
+    # 4) Generate story
+    with st.spinner("✍️ Generating story..."):
+        try:
+            story_out = story_client(
+                inputs=prompt,
+                parameters={
+                    "max_new_tokens": 250,
+                    "temperature": 0.7,
+                    "top_p": 0.9,
+                    "repetition_penalty": 1.1,
+                    "do_sample": True
+                }
+            )
+            
+            # Parse response
+            if isinstance(story_out, list):
+                story_text = story_out[0].get("generated_text", "")
+            else:
+                story_text = story_out.get("generated_text", "")
+            
+            # Clean up output
+            story = story_text.split("Story:")[-1].strip()
+
+        except Exception as e:
+            st.error(f"🚨 Story generation failed: {str(e)}")
+            st.stop()
+
+    if not story:
+        st.error("😕 Couldn’t generate a story. Please try again!")
         st.stop()
 
+    st.markdown("**Story:**")
+    st.write(story)
+
     # 5) Text-to-Speech
     with st.spinner("🔊 Converting to speech..."):
         try: