Update app.py

app.py

@@ -2,42 +2,38 @@ import os
 import time
 import streamlit as st
 from PIL import Image
-from io import BytesIO
 from transformers import pipeline
-from huggingface_hub import login
 from gtts import gTTS
 import tempfile
 
 # —––––––– Requirements —–––––––
-# pip install streamlit pillow gTTS transformers huggingface_hub
-
-# —––––––– Page Config —–––––––
-st.set_page_config(page_title="Magic Story Generator (Local Pipeline)", layout="centered")
+# streamlit>=1.20
+# pillow>=9.0
+# transformers>=4.30
+# torch>=2.0.0
+# sentencepiece>=0.1.97
+# gTTS>=2.3.1
+
+# —––––––– Page Setup —–––––––
+st.set_page_config(page_title="Magic Story Generator", layout="centered")
 st.title("📖✨ Turn Images into Children's Stories")
 
-# —––––––– Load Clients & Pipelines (cached) —–––––––
+# —––––––– Load Pipelines (cached) —–––––––
 @st.cache_resource(show_spinner=False)
-def load_clients():
-    # Authenticate to pull private or remote-code models if needed
-    hf_token = st.secrets.get("HF_TOKEN")
-    if hf_token:
-        os.environ["HUGGINGFACEHUB_API_TOKEN"] = hf_token
-        login(hf_token)
-
+def load_pipelines():
     # 1) Image-captioning pipeline (BLIP)
     captioner = pipeline(
         task="image-to-text",
         model="Salesforce/blip-image-captioning-base",
-        device=-1  # CPU; change to 0 for GPU
+        device=-1  # CPU; set to 0+ for GPU
     )
-
     # 2) Story-generation pipeline (DeepSeek-R1-Distill-Qwen)
     storyteller = pipeline(
-        task="text-generation",
+        task="text2text-generation",
         model="deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B",
         tokenizer="deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B",
         trust_remote_code=True,
-        device=-1,                # CPU; set 0+ for GPU
+        device=-1,                # CPU; set to 0+ for GPU
         temperature=0.6,
         top_p=0.9,
         repetition_penalty=1.1,
@@ -45,68 +41,54 @@ def load_clients():
         max_new_tokens=120,
         return_full_text=False
     )
-
     return captioner, storyteller
 
-captioner, storyteller = load_clients()
-
-# —––––––– Helpers —–––––––
-def generate_caption(img: Image.Image) -> str:
-    # Use the BLIP pipeline to generate a caption
-    result = captioner(img)
-    if isinstance(result, list) and result:
-        return result[0].get("generated_text", "").strip()
-    return ""
-
-
-def generate_story(caption: str) -> str:
-    # Build a simple prompt incorporating the caption
-    prompt = (
-        f"Image description: {caption}\n"
-        "Write a coherent 50-100 word children's story that flows naturally."
-    )
-
-    t0 = time.time()
-    outputs = storyteller(
-        prompt
-    )
-    gen_time = time.time() - t0
-    st.text(f"⏱ Generated in {gen_time:.1f}s")
-
-    story = outputs[0].get("generated_text", "").strip()
-    # Truncate to 100 words
-    words = story.split()
-    if len(words) > 100:
-        story = " ".join(words[:100]) + ('.' if not story.endswith('.') else '')
-    return story
+captioner, storyteller = load_pipelines()
 
 # —––––––– Main App —–––––––
 uploaded = st.file_uploader("Upload an image:", type=["jpg", "png", "jpeg"])
 if uploaded:
+    # Display uploaded image
     img = Image.open(uploaded).convert("RGB")
-    if max(img.size) > 2048:
-        img.thumbnail((2048, 2048))
+    img.thumbnail((2048, 2048), Image.ANTIALIAS)
     st.image(img, use_container_width=True)
 
+    # Generate caption
     with st.spinner("🔍 Generating caption..."):
-        caption = generate_caption(img)
+        cap = captioner(img)
+        caption = cap[0].get("generated_text", "").strip() if isinstance(cap, list) else ""
     if not caption:
         st.error("😢 Couldn't understand this image. Try another one!")
         st.stop()
     st.success(f"**Caption:** {caption}")
 
+    # Generate story
+    prompt = (
+        f"Image description: {caption}\n"
+        "Write a coherent, 50-100 word children’s story that flows naturally."
+    )
     with st.spinner("📝 Writing story..."):
-        story = generate_story(caption)
+        start = time.time()
+        out = storyteller(prompt)
+        gen_time = time.time() - start
+        st.text(f"⏱ Generated in {gen_time:.1f}s")
+    story = out[0].get("generated_text", "").strip()
+
+    # Enforce ≤100 words
+    words = story.split()
+    if len(words) > 100:
+        story = " ".join(words[:100]) + ("" if story.endswith('.') else ".")
 
     st.subheader("📚 Your Magical Story")
     st.write(story)
 
+    # Convert to audio
     with st.spinner("🔊 Converting to audio..."):
         try:
             tts = gTTS(text=story, lang="en", slow=False)
-            with tempfile.NamedTemporaryFile(delete=False, suffix=".mp3") as fp:
-                tts.save(fp.name)
-                st.audio(fp.name, format="audio/mp3")
+            tmp = tempfile.NamedTemporaryFile(delete=False, suffix=".mp3")
+            tts.save(tmp.name)
+            st.audio(tmp.name, format="audio/mp3")
         except Exception as e:
             st.warning(f"⚠️ TTS failed: {e}")