Update app.py

app.py

@@ -5,34 +5,34 @@ from PIL import Image
 from transformers import pipeline
 from gtts import gTTS
 import tempfile
+from llama_cpp import Llama
+
+# First install required package:
+# pip install llama-cpp-python
 
 # —––––––– Page Setup —–––––––
 st.set_page_config(page_title="Magic Story Generator", layout="centered")
 st.title("📖✨ Turn Images into Children's Stories")
 
-# —––––––– Load Pipelines (cached) —–––––––
+# —––––––– Load Models (cached) —–––––––
 @st.cache_resource(show_spinner=False)
-def load_pipelines():
-    # 1) Image captioning pipeline
+def load_models():
+    # 1) Image captioning model
     captioner = pipeline(
-        task="image-to-text",
-        model="Salesforce/blip-image-captioning-base",
-        device=-1
+        "image-to-text",
+        model="Salesforce/blip-image-captioning-base"
     )
     
-    # 2) Story generation pipeline using verified model
-    storyteller = pipeline(
-        task="text2text-generation",
-        model="laxya007/story-generator-t5-small",
-        tokenizer="t5-small",
-        device=-1,
-        max_length=200,
-        temperature=0.7,
-        do_sample=True
+    # 2) GGUF Story Model
+    storyteller = Llama(
+        model_path="DavidAU/L3-Grand-Story-Darkness-MOE-4X8-24.9B-e32-GGUF",
+        n_ctx=2048,
+        n_threads=4,
+        n_gpu_layers=0  # Set based on your GPU capacity
     )
     return captioner, storyteller
 
-captioner, storyteller = load_pipelines()
+captioner, storyteller = load_models()
 
 # —––––––– Main App —–––––––
 uploaded = st.file_uploader("Upload an image:", type=["jpg", "png", "jpeg"])
@@ -43,29 +43,31 @@ if uploaded:
     # Generate caption
     with st.spinner("🔍 Generating caption..."):
         cap = captioner(img)
-        caption = cap[0].get("generated_text", "").strip()
-    if not caption:
-        st.error("😢 Couldn't understand this image. Try another one!")
-        st.stop()
+        caption = cap[0]['generated_text']
     st.success(f"**Caption:** {caption}")
 
     # Generate story
-    prompt = f"generate story: {caption}"
-    with st.spinner("📝 Writing story..."):
+    prompt = f"""Below is an image description. Write a children's story based on it.
+    
+    Image Description: {caption}
+    Story:"""
+    
+    with st.spinner("📝 Crafting magical story..."):
         start = time.time()
-        story = storyteller(prompt)[0]['generated_text']
+        output = storyteller(
+            prompt=prompt,
+            max_tokens=500,
+            temperature=0.7,
+            top_p=0.9,
+            repeat_penalty=1.1
+        )
         gen_time = time.time() - start
+        story = output['choices'][0]['text'].strip()
         st.text(f"⏱ Generated in {gen_time:.1f}s")
 
-    # Format story output
-    story = story.replace("<pad>", "").replace("</s>", "").strip()
-    if story.startswith("generate story:"):
-        story = story[15:].strip()
+    # Post-process story
+    story = story.split("###")[0].strip()  # Remove any trailing artifacts
     
-    # Word limit enforcement
-    words = story.split()
-    story = " ".join(words[:100]) if len(words) > 100 else story
-
     # Display story
     st.subheader("📚 Your Magical Story")
     st.write(story)
@@ -82,4 +84,3 @@ if uploaded:
 
 # Footer
 st.markdown("---\n*Made with ❤️ by your friendly story wizard*")
-