Update app.py

app.py

@@ -5,6 +5,8 @@ from huggingface_hub import InferenceApi
 from gtts import gTTS
 import requests
 import tempfile
+import time
+import threading
 
 # —––––––– Page Config —–––––––
 st.set_page_config(page_title="Magic Story Generator", layout="centered")
@@ -14,10 +16,31 @@ st.title("📖✨ Turn Images into Children's Stories")
 @st.cache_resource
 def load_clients():
     hf_token = st.secrets["HF_TOKEN"]
-    return (
-        InferenceApi("Salesforce/blip-image-captioning-base", token=hf_token),
-        hf_token  # we'll use direct requests for story generation
-    )
+    caption_client = InferenceApi("Salesforce/blip-image-captioning-base", token=hf_token)
+
+    # Start background keep-alive for story model to avoid cold starts
+    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}
+    }
+    def keep_model_warm():
+        # Initial warm-up
+        try:
+            requests.post(api_url, headers=headers, json=keep_alive_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)
+            except:
+                pass
+    threading.Thread(target=keep_model_warm, daemon=True).start()
+
+    return caption_client, hf_token
 
 caption_client, hf_token = load_clients()
 
@@ -44,7 +67,7 @@ def process_image(uploaded_file):
         st.error(f"Image processing error: {type(e).__name__}: {e}")
         st.stop()
 
-# —––––––– Helper: Generate Story via HTTP —–––––––
+# —––––––– Helper: Generate Story —–––––––
 def generate_story(prompt: str) -> str:
     api_url = "https://api-inference.huggingface.co/models/deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B"
     headers = {"Authorization": f"Bearer {hf_token}"}
@@ -59,25 +82,39 @@ def generate_story(prompt: str) -> str:
             "no_repeat_ngram_size": 2
         }
     }
-    try:
-        resp = requests.post(api_url, headers=headers, json=payload, timeout=30)
-    except Exception as e:
-        st.error(f"🚨 Story magic failed: {type(e).__name__}: {e}")
-        st.stop()
-    if resp.status_code != 200:
+    max_retries = 5
+    retries = 0
+    while True:
+        try:
+            resp = requests.post(api_url, headers=headers, json=payload, timeout=30)
+        except Exception as e:
+            st.error(f"🚨 Story magic failed: {type(e).__name__}: {e}")
+            st.stop()
+
+        if resp.status_code == 200:
+            data = resp.json()
+            if isinstance(data, list) and data:
+                text = data[0].get("generated_text", "").strip()
+                story = text.split("Story:")[-1].strip()
+                if "." in story:
+                    story = story.rsplit(".", 1)[0] + "."
+                return story
+            else:
+                st.error("🚨 Story magic failed: invalid response format")
+                st.stop()
+
+        if resp.status_code == 503 and retries < max_retries:
+            try:
+                wait = int(resp.json().get("estimated_time", 5))
+            except:
+                wait = 5 * (2 ** retries)
+            st.info(f"Model loading; retrying in {wait}s (attempt {retries+1}/{max_retries})")
+            time.sleep(wait)
+            retries += 1
+            continue
+
         st.error(f"🚨 Story magic failed: HTTP {resp.status_code} - {resp.text}")
         st.stop()
-    data = resp.json()
-    # Expecting list of generations
-    if isinstance(data, list) and data:
-        text = data[0].get("generated_text", "").strip()
-        # extract after "Story:" if present
-        story = text.split("Story:")[-1].strip()
-        if "." in story:
-            story = story.rsplit(".", 1)[0] + "."
-        return story
-    st.error("🚨 Story magic failed: invalid response format")
-    st.stop()
 
 # —––––––– Main App Flow —–––––––
 uploaded = st.file_uploader("Upload an image:", type=["jpg", "png", "jpeg"])
@@ -85,7 +122,7 @@ if uploaded:
     img = process_image(uploaded)
     st.image(img, use_container_width=True)
 
-    # Caption
+    # Generate Caption
     with st.spinner("🔍 Discovering image secrets..."):
         caption = generate_caption(img)
         if not caption:
@@ -93,7 +130,7 @@ if uploaded:
             st.stop()
     st.success(f"**Caption:** {caption}")
 
-    # Prepare Prompt
+    # Prepare Story Prompt
     story_prompt = (
         f"Image description: {caption}\n\n"
         "Write a 50-100 word children's story that:\n"
@@ -104,11 +141,9 @@ if uploaded:
         "Story:\n"
     )
 
-    # Generate Story
+    # Generate and Display Story
     with st.spinner("📝 Writing magical story..."):
         story = generate_story(story_prompt)
-
-    # Display Story
     st.subheader("📚 Your Magical Story")
     st.write(story)
 
@@ -124,5 +159,3 @@ if uploaded:
 
 # Footer
 st.markdown("---\n*Made with ❤️ by your friendly story wizard*")
-
-