Update app.py

app.py

@@ -1,76 +1,108 @@
 import os
-import torch
-from transformers import GPT2Tokenizer, GPT2LMHeadModel
-from huggingface_hub import InferenceApi
+import time
+import streamlit as st
+from transformers import pipeline
 from PIL import Image
-from io import BytesIO
+from gtts import gTTS
+import tempfile
 
-def load_caption_client(token: str):
-    return InferenceApi(repo_id="Salesforce/blip-image-captioning-base", token=token)
+# —––––––– Page Config —–––––––
+st.set_page_config(page_title="Magic Story Generator", layout="centered")
+st.title("📖✨ Turn Images into Children's Stories")
 
-def generate_caption(image_path: str, caption_client) -> str:
-    img = Image.open(image_path).convert("RGB")
-    buf = BytesIO()
-    img.save(buf, format="JPEG")
-    resp = caption_client(data=buf.getvalue())
-    if isinstance(resp, list) and resp:
-        return resp[0].get("generated_text", "").strip()
-    return ""
+# —––––––– Load Pipelines (cached) —–––––––
+@st.cache_resource(show_spinner=False)
+def load_pipelines():
+    # Cache transformers models locally
+    os.environ.setdefault("TRANSFORMERS_CACHE", "./hf_cache")
 
-def load_gpt2(model_name="gpt2"):
-    tokenizer = GPT2Tokenizer.from_pretrained(model_name)
-    model = GPT2LMHeadModel.from_pretrained(model_name)
-    model.eval()
-    return tokenizer, model
-
-def generate_story(caption: str, tokenizer, model) -> str:
-    # Build a strong prompt
-    prompt = (
-        f"You are a creative children’s-story author.\n"
-        f"Image description: “{caption}”\n\n"
-        "Write a coherent, 50–100 word story:\n"
+    # 1) Image-to-text pipeline for captioning (BLIP)
+    captioner = pipeline(
+        task="image-to-text",
+        model="Salesforce/blip-image-captioning-base",
+        device=-1  # force CPU; use 0 for GPU
     )
-    # Tokenize and move to device
-    inputs = tokenizer(prompt, return_tensors="pt")
-    # Generate up to ~120 new tokens
-    outputs = model.generate(
-        **inputs,
-        max_new_tokens=120,
-        temperature=0.7,
-        top_p=0.9,
-        repetition_penalty=1.1,
-        no_repeat_ngram_size=3,
-        do_sample=True,
-        pad_token_id=tokenizer.eos_token_id
+
+    # 2) Text-generation pipeline for storytelling (GPT-2)
+    storyteller = pipeline(
+        task="text-generation",
+        model="gpt2",
+        tokenizer="gpt2",
+        device=-1  # CPU
     )
-    # Decode and strip the prompt echo
-    text = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    story = text[len(prompt):].strip()
-    # Truncate to 100 words if needed
-    words = story.split()
-    if len(words) > 100:
-        story = " ".join(words[:100])
-        if not story.endswith("."):
-            story += "."
-    return story
 
-if __name__ == "__main__":
-    # 1) Read your HF token
-    hf_token = os.environ.get("HF_TOKEN")
-    if not hf_token:
-        raise RuntimeError("Please set HF_TOKEN env var")
+    return captioner, storyteller
+
+captioner, storyteller = load_pipelines()
+
+# —––––––– Main App Flow —–––––––
+uploaded = st.file_uploader("Upload an image:", type=["jpg", "png", "jpeg"])
+if uploaded:
+    # Load image
+    img = Image.open(uploaded).convert("RGB")
+    if max(img.size) > 2048:
+        img.thumbnail((2048, 2048))
+    st.image(img, use_container_width=True)
 
-    # 2) Generate caption
-    caption_client = load_caption_client(hf_token)
-    image_path = "path/to/your/image.jpg"   # <-- change me
-    caption = generate_caption(image_path, caption_client)
-    print(f"Caption: {caption}\n")
+    # Generate caption
+    with st.spinner("🔍 Generating caption..."):
+        raw = captioner(img)
+        caption = raw[0].get("generated_text", "").strip()
+    if not caption:
+        st.error("😢 Couldn't understand this image. Try another one!")
+        st.stop()
+    st.success(f"**Caption:** {caption}")
+
+    # Build storytelling prompt
+    prompt = f"""
+You are a creative children’s-story author.
+Image description: “{caption}”
+
+Write a coherent, 50–100 word story that:
+1. Introduces the main character.
+2. Shows a simple problem or discovery.
+3. Has a happy resolution.
+4. Uses clear language for ages 3–8.
+5. Keeps sentences under 20 words.
+Story:
+"""
+    # Generate story
+    with st.spinner("📝 Writing story..."):
+        t0 = time.time()
+        outputs = storyteller(
+            prompt,
+            max_new_tokens=120,
+            temperature=0.7,
+            top_p=0.9,
+            repetition_penalty=1.1,
+            no_repeat_ngram_size=3,
+            do_sample=True,
+            pad_token_id=storyteller.tokenizer.eos_token_id
+        )
+        story_text = outputs[0]["generated_text"].strip()
+        load_time = time.time() - t0
+    st.text(f"⏱ Story generated in {load_time:.1f}s")
+
+    # Post-process: strip prompt echo and truncate
+    if story_text.startswith(prompt):
+        story_text = story_text[len(prompt):].strip()
+    words = story_text.split()
+    if len(words) > 100:
+        story_text = " ".join(words[:100]) + ("." if not story_text.endswith('.') else "")
 
-    # 3) Load GPT-2
-    tokenizer, model = load_gpt2("gpt2")
-    # (optionally move model to GPU: model.to("cuda"))
+    # Display story
+    st.subheader("📚 Your Magical Story")
+    st.write(story_text)
 
-    # 4) Generate & print story
-    story = generate_story(caption, tokenizer, model)
-    print("Story:\n", story)
+    # Convert to audio
+    with st.spinner("🔊 Converting to audio..."):
+        try:
+            tts = gTTS(text=story_text, lang="en", slow=False)
+            with tempfile.NamedTemporaryFile(delete=False, suffix=".mp3") as fp:
+                tts.save(fp.name)
+                st.audio(fp.name, format="audio/mp3")
+        except Exception as e:
+            st.warning(f"⚠️ TTS failed: {e}")
 
+# Footer
+st.markdown("---\n*Made with ❤️ by your friendly story wizard* ")