Update app.py

app.py

@@ -2,12 +2,9 @@ import streamlit as st
 from PIL import Image
 from io import BytesIO
 from huggingface_hub import InferenceApi
+from transformers import pipeline
 from gtts import gTTS
-import requests
-from requests.exceptions import ReadTimeout
 import tempfile
-import time
-import threading
 
 # —––––––– Page Config —–––––––
 st.set_page_config(page_title="Magic Story Generator", layout="centered")
@@ -17,137 +14,69 @@ st.title("📖✨ Turn Images into Children's Stories")
 @st.cache_resource
 def load_clients():
     hf_token = st.secrets["HF_TOKEN"]
+    # image captioning client as before
     caption_client = InferenceApi("Salesforce/blip-image-captioning-base", token=hf_token)
+    # text-generation pipeline for story
+    story_generator = pipeline(
+        "text-generation",
+        model="deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B",
+        tokenizer="deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B",
+        use_auth_token=hf_token,
+        device=0  # or -1 for CPU
+    )
+    return caption_client, story_generator
 
-    # 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}"}
-    warm_payload = {"inputs": "Hello!", "parameters": {"max_new_tokens": 1}}
-    def keep_model_warm():
-        try:
-            requests.post(api_url, headers=headers, json=warm_payload, timeout=10)
-        except:
-            pass
-        while True:
-            time.sleep(600)
-            try:
-                requests.post(api_url, headers=headers, json=warm_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()
+caption_client, story_generator = load_clients()
 
 # —––––––– Helper: Generate Caption —–––––––
 def generate_caption(img):
-    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()
-    except Exception as e:
-        st.error(f"Caption generation error: {type(e).__name__}: {e}")
-    return ""
-
-# —––––––– Helper: Process Image —–––––––
-def process_image(uploaded_file):
+    buf = BytesIO()
+    img.save(buf, format="JPEG")
     try:
-        img = Image.open(uploaded_file).convert("RGB")
-        if max(img.size) > 2048:
-            img.thumbnail((2048, 2048))
-        return img
+        out = caption_client(data=buf.getvalue())
+        return out[0].get("generated_text", "").strip()
     except Exception as e:
-        st.error(f"Image processing error: {type(e).__name__}: {e}")
-        st.stop()
-
-# —––––––– Helper: Generate Story with improved retry and timeout —–––––––
-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 = {
-        "inputs": prompt,
-        "parameters": {
-            "max_new_tokens":       200,
-            "temperature":          0.8,
-            "top_p":                0.95,
-            "repetition_penalty":   1.15,
-            "do_sample":            True,
-            "no_repeat_ngram_size": 2
-        }
-    }
-    retries = 0
-    max_retries = 5
-    timeout = 60  # allow up to 60s for large model
-    while True:
-        try:
-            resp = requests.post(api_url, headers=headers, json=payload, timeout=timeout)
-        except ReadTimeout:
-            if retries < max_retries:
-                wait = 2 ** retries
-                st.info(f"Request timed out; retrying in {wait}s (attempt {retries+1}/{max_retries})")
-                time.sleep(wait)
-                retries += 1
-                continue
-            st.error("🚨 Story magic failed: request timed out after multiple attempts.")
-            st.stop()
-        except Exception as e:
-            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:
-                text = data[0].get("generated_text", "").strip()
-                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()
-
-        # Model loading (cold start)
-        if resp.status_code == 503 and retries < max_retries:
-            wait = int(resp.json().get("estimated_time", 5))
-            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:
-            wait = 2 ** retries
-            st.info(f"Server error {resp.status_code}; retrying in {wait}s (attempt {retries+1}/{max_retries})")
-            time.sleep(wait)
-            retries += 1
-            continue
-        if resp.status_code in (424, 500, 502):
-            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()
+        st.error(f"Caption error: {e}")
+        return ""
+
+# —––––––– Helper: Generate Story via pipeline —–––––––
+def generate_story(prompt: str) -> str:
+    # generate up to ~200 tokens to cover 100 words margin
+    outputs = story_generator(
+        prompt,
+        max_new_tokens=200,
+        temperature=0.8,
+        top_p=0.95,
+        repetition_penalty=1.15,
+        no_repeat_ngram_size=2,
+        do_sample=True
+    )
+    text = outputs[0]["generated_text"].strip()
+    # everything after "Story:" (if you kept that in your prompt)
+    if "Story:" in text:
+        text = text.split("Story:", 1)[1].strip()
+    # truncate to 100 words
+    words = text.split()
+    if len(words) > 100:
+        text = " ".join(words[:100])
+        if not text.endswith("."):
+            text += "."
+    return text
 
 # —––––––– Main App Flow —–––––––
 uploaded = st.file_uploader("Upload an image:", type=["jpg", "png", "jpeg"])
 if uploaded:
-    img = process_image(uploaded)
+    img = Image.open(uploaded).convert("RGB")
+    if max(img.size) > 2048:
+        img.thumbnail((2048, 2048))
     st.image(img, use_container_width=True)
 
-    # Generate Caption
-    with st.spinner("🔍 Discovering image secrets..."):
-        caption = generate_caption(img)
-        if not caption:
-            st.error("😢 Couldn't understand this image. Try another one!")
-            st.stop()
+    caption = generate_caption(img)
+    if not caption:
+        st.error("😢 Couldn't understand this image. Try another one!")
+        st.stop()
     st.success(f"**Caption:** {caption}")
 
-    # Prepare Story Prompt
     story_prompt = (
         f"Image description: {caption}\n\n"
         "Write a 50-100 word children's story that:\n"
@@ -158,21 +87,8 @@ if uploaded:
         "Story:\n"
     )
 
-    # Generate and validate Story
     with st.spinner("📝 Writing magical story..."):
-        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
-
+        story = generate_story(story_prompt)
     st.subheader("📚 Your Magical Story")
     st.write(story)
 
@@ -184,8 +100,9 @@ if uploaded:
                 tts.save(fp.name)
                 st.audio(fp.name, format="audio/mp3")
         except Exception as e:
-            st.warning(f"⚠️ Couldn't make audio version: {type(e).__name__}: {e}")
+            st.warning(f"⚠️ Couldn't make audio version: {e}")
 
 # Footer
 st.markdown("---\n*Made with ❤️ by your friendly story wizard*")
 
+