Update app.py

app.py

@@ -1,135 +1,76 @@
 import os
-import time
-import streamlit as st
+import torch
+from transformers import GPT2Tokenizer, GPT2LMHeadModel
+from huggingface_hub import InferenceApi
 from PIL import Image
 from io import BytesIO
-from huggingface_hub import InferenceApi, login
-from transformers import pipeline
-from gtts import gTTS
-import tempfile
 
-# —––––––– Page Config —–––––––
-st.set_page_config(page_title="Magic Story Generator", layout="centered")
-st.title("📖✨ Turn Images into Children's Stories")
+def load_caption_client(token: str):
+    return InferenceApi(repo_id="Salesforce/blip-image-captioning-base", token=token)
 
-# —––––––– Clients (cached) —–––––––
-@st.cache_resource(show_spinner=False)
-def load_clients():
-    hf_token = st.secrets["HF_TOKEN"]
-
-    # Authenticate for both HF Hub and transformers
-    os.environ["HUGGINGFACEHUB_API_TOKEN"] = hf_token
-    login(hf_token)
-
-    # Pin transformers cache locally via env var
-    cache_dir = "./hf_cache"
-    os.makedirs(cache_dir, exist_ok=True)
-    os.environ["TRANSFORMERS_CACHE"] = cache_dir
-
-    # 1) BLIP image-captioning client
-    caption_client = InferenceApi(
-        repo_id="Salesforce/blip-image-captioning-base",
-        token=hf_token
-    )
-
-    # 2) Text-generation pipeline on CPU (no cache_dir arg here!)
-    t0 = time.time()
-    story_generator = pipeline(
-        task="text-generation",
-        model="deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B",
-        tokenizer="deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B",
-        device=-1  # force CPU
-    )
-    st.text(f"✅ Story model loaded in {time.time() - t0:.1f}s (cached thereafter)")
-
-    return caption_client, story_generator
-
-caption_client, story_generator = load_clients()
-
-
-# —––––––– Helper: Generate Caption —–––––––
-def generate_caption(img: Image.Image) -> str:
+def generate_caption(image_path: str, caption_client) -> str:
+    img = Image.open(image_path).convert("RGB")
     buf = BytesIO()
     img.save(buf, format="JPEG")
-    try:
-        resp = caption_client(data=buf.getvalue())
-        if isinstance(resp, list) and resp:
-            return resp[0].get("generated_text", "").strip()
-    except Exception as e:
-        st.error(f"Caption generation error: {e}")
+    resp = caption_client(data=buf.getvalue())
+    if isinstance(resp, list) and resp:
+        return resp[0].get("generated_text", "").strip()
     return ""
 
-
-# —––––––– Helper: Generate Story via pipeline —–––––––
-def generate_story(caption: str) -> str:
-    prompt = f"""
-You are a creative children’s-story author.
-Below is an 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:
-"""
-    t0 = time.time()
-    outputs = story_generator(
-        prompt,
+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"
+    )
+    # 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
+        do_sample=True,
+        pad_token_id=tokenizer.eos_token_id
     )
-    st.text(f"⏱ Generated in {time.time() - t0:.1f}s on CPU")
-
-    text = outputs[0]["generated_text"].strip()
-    # strip the prompt echo
-    if text.startswith(prompt):
-        text = text[len(prompt):].strip()
-    # enforce ≤100 words
-    words = text.split()
+    # 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:
-        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 = Image.open(uploaded).convert("RGB")
-    if max(img.size) > 2048:
-        img.thumbnail((2048, 2048))
-    st.image(img, use_container_width=True)
-
-    with st.spinner("🔍 Generating caption..."):
-        caption = generate_caption(img)
-    if not caption:
-        st.error("😢 Couldn't understand this image. Try another one!")
-        st.stop()
-    st.success(f"**Caption:** {caption}")
-
-    with st.spinner("📝 Writing magical story..."):
-        story = generate_story(caption)
-
-    st.subheader("📚 Your Magical Story")
-    st.write(story)
-
-    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")
-        except Exception as e:
-            st.warning(f"⚠️ TTS failed: {e}")
-
-# Footer
-st.markdown("---\n*Made with ❤️ by your friendly story wizard*")
+        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")
+
+    # 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")
+
+    # 3) Load GPT-2
+    tokenizer, model = load_gpt2("gpt2")
+    # (optionally move model to GPU: model.to("cuda"))
+
+    # 4) Generate & print story
+    story = generate_story(caption, tokenizer, model)
+    print("Story:\n", story)