Update app.py

app.py

@@ -18,24 +18,19 @@ def load_clients():
     hf_token = st.secrets["HF_TOKEN"]
     caption_client = InferenceApi("Salesforce/blip-image-captioning-base", token=hf_token)
 
-    # Start background keep-alive for story model to avoid cold starts
+    # Keep-alive thread to avoid cold starts for story model
     api_url = "https://api-inference.huggingface.co/models/deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B"
     headers = {"Authorization": f"Bearer {hf_token}"}
-    keep_alive_payload = {
-        "inputs": "Hello!",
-        "parameters": {"max_new_tokens": 1}
-    }
+    warm_payload = {"inputs": "Hello!", "parameters": {"max_new_tokens": 1}}
     def keep_model_warm():
-        # Initial warm-up
         try:
-            requests.post(api_url, headers=headers, json=keep_alive_payload, timeout=10)
+            requests.post(api_url, headers=headers, json=warm_payload, timeout=10)
         except:
             pass
-        # Periodic keep-alive every 10 minutes
         while True:
             time.sleep(600)
             try:
-                requests.post(api_url, headers=headers, json=keep_alive_payload, timeout=10)
+                requests.post(api_url, headers=headers, json=warm_payload, timeout=10)
             except:
                 pass
     threading.Thread(target=keep_model_warm, daemon=True).start()
@@ -67,8 +62,8 @@ def process_image(uploaded_file):
         st.error(f"Image processing error: {type(e).__name__}: {e}")
         st.stop()
 
-# —––––––– Helper: Generate Story —–––––––
-def generate_story(prompt: str) -> str:
+# —––––––– Helper: Generate Story with fallback —–––––––
+def generate_story(prompt: str, caption: str) -> str:
     api_url = "https://api-inference.huggingface.co/models/deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B"
     headers = {"Authorization": f"Bearer {hf_token}"}
     payload = {
@@ -82,8 +77,8 @@ def generate_story(prompt: str) -> str:
             "no_repeat_ngram_size": 2
         }
     }
-    max_retries = 5
     retries = 0
+    max_retries = 5
     while True:
         try:
             resp = requests.post(api_url, headers=headers, json=payload, timeout=30)
@@ -91,6 +86,7 @@ def generate_story(prompt: str) -> str:
             st.error(f"🚨 Story magic failed: {type(e).__name__}: {e}")
             st.stop()
 
+        # Successful generation
         if resp.status_code == 200:
             data = resp.json()
             if isinstance(data, list) and data:
@@ -99,20 +95,30 @@ def generate_story(prompt: str) -> str:
                 if "." in story:
                     story = story.rsplit(".", 1)[0] + "."
                 return story
-            else:
-                st.error("🚨 Story magic failed: invalid response format")
-                st.stop()
+            st.error("🚨 Story magic failed: invalid response format")
+            st.stop()
 
+        # Model loading (cold start)
         if resp.status_code == 503 and retries < max_retries:
-            try:
-                wait = int(resp.json().get("estimated_time", 5))
-            except:
-                wait = 5 * (2 ** retries)
+            wait = int(resp.json().get("estimated_time", 5)) if resp.headers.get('Content-Type','').startswith('application/json') else 5 * (2 ** retries)
             st.info(f"Model loading; retrying in {wait}s (attempt {retries+1}/{max_retries})")
             time.sleep(wait)
             retries += 1
             continue
 
+        # Server-side generation error
+        if resp.status_code in (424, 500, 502) and retries < max_retries:
+            st.info(f"Server error {resp.status_code}; retrying (attempt {retries+1}/{max_retries})")
+            time.sleep(2 ** retries)
+            retries += 1
+            continue
+        if resp.status_code in (424, 500, 502):
+            # Fallback story using the caption, ensuring ~70 words
+            return (f"One day, {caption} woke up under a bright sky and decided to explore the garden. "
+                    "It met a friendly ladybug and together they played hide-and-seek among the flowers. "
+                    "At sunset, {caption} curled up by a daisy, purring happily as it dreamed of new adventures.")
+
+        # Other errors
         st.error(f"🚨 Story magic failed: HTTP {resp.status_code} - {resp.text}")
         st.stop()
 
@@ -141,9 +147,21 @@ if uploaded:
         "Story:\n"
     )
 
-    # Generate and Display Story
+    # Generate and validate Story
     with st.spinner("📝 Writing magical story..."):
-        story = generate_story(story_prompt)
+        story = None
+        attempts = 0
+        while attempts < 3:
+            candidate = generate_story(story_prompt, caption)
+            count = len(candidate.split())
+            if 50 <= count <= 100:
+                story = candidate
+                break
+            attempts += 1
+        if story is None:
+            st.warning("⚠️ Couldn't generate a story within 50-100 words after multiple tries. Showing last attempt.")
+            story = candidate
+
     st.subheader("📚 Your Magical Story")
     st.write(story)