Update app.py

app.py

@@ -1,125 +1,115 @@
-# app.py
-import streamlit as st
-from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
-from gtts import gTTS
 import os
 import time
-import torch
-from threading import Thread
+import streamlit as st
+from PIL import Image
+from transformers import pipeline
+from gtts import gTTS
+import tempfile
 
-# Initialize models
-@st.cache_resource
-def load_models():
-    model_name = "Qwen/Qwen3-1.7B"
-    
-    tokenizer = AutoTokenizer.from_pretrained(
-        model_name,
-        trust_remote_code=True
+# --- Requirements ---
+# Update requirements.txt to include:
+"""
+streamlit>=1.20
+pillow>=9.0
+torch>=2.0.0
+transformers>=4.40
+sentencepiece>=0.2.0
+gTTS>=2.3.1
+accelerate>=0.30
+"""
+
+# --- Page Setup ---
+st.set_page_config(page_title="Magic Story Generator", layout="centered")
+st.title("📖✨ Turn Images into Children's Stories")
+
+# --- Load Pipelines (cached) ---
+@st.cache_resource(show_spinner=False)
+def load_pipelines():
+    # 1) Image-captioning pipeline (BLIP)
+    captioner = pipeline(
+        task="image-to-text",
+        model="Salesforce/blip-image-captioning-base",
+        device=-1
     )
     
-    model = AutoModelForCausalLM.from_pretrained(
-        model_name,
-        torch_dtype="auto",
+    # 2) Modified story-generation pipeline using Qwen3-1.7B
+    storyteller = pipeline(
+        task="text-generation",
+        model="Qwen/Qwen3-1.7B",
         device_map="auto",
-        trust_remote_code=True
+        trust_remote_code=True,
+        torch_dtype="auto",
+        max_new_tokens=150,
+        temperature=0.7,
+        top_p=0.9,
+        repetition_penalty=1.2,
+        eos_token_id=151645  # Specific to Qwen3 tokenizer
     )
     
-    return model, tokenizer
+    return captioner, storyteller
+
+captioner, storyteller = load_pipelines()
 
-def parse_thinking_output(output_ids, tokenizer, thinking_token_id=151668):
-    try:
-        index = len(output_ids) - output_ids[::-1].index(thinking_token_id)
-    except ValueError:
-        index = 0
-        
-    thinking = tokenizer.decode(output_ids[:index], skip_special_tokens=True).strip("\n")
-    content = tokenizer.decode(output_ids[index:], skip_special_tokens=True).strip("\n")
-    return thinking, content
+# --- Main App ---
+uploaded = st.file_uploader("Upload an image:", type=["jpg", "png", "jpeg"])
+if uploaded:
+    # Load and display the image
+    img = Image.open(uploaded).convert("RGB")
+    st.image(img, use_container_width=True)
 
-def generate_response(prompt, model, tokenizer):
-    messages = [{"role": "user", "content": prompt}]
-    text = tokenizer.apply_chat_template(
-        messages,
-        tokenize=False,
-        add_generation_prompt=True,
-        enable_thinking=True
+    # Generate caption
+    with st.spinner("🔍 Generating caption..."):
+        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}")
+
+    # Build prompt and generate story
+    prompt = (
+        f"<|im_start|>system\n"
+        f"You are a children's story writer. Create a 50-100 word story based on this image description: {caption}\n"
+        f"<|im_end|>\n"
+        f"<|im_start|>user\n"
+        f"Write a coherent, child-friendly story that flows naturally with simple vocabulary.<|im_end|>\n"
+        f"<|im_start|>assistant\n"
     )
     
-    streamer = TextIteratorStreamer(tokenizer)
-    inputs = tokenizer([text], return_tensors="pt").to(model.device)
+    with st.spinner("📝 Writing story..."):
+        start = time.time()
+        out = storyteller(
+            prompt,
+            do_sample=True,
+            num_return_sequences=1
+        )
+        gen_time = time.time() - start
+        st.text(f"⏱ Generated in {gen_time:.1f}s")
     
-    generation_kwargs = dict(
-        **inputs,
-        streamer=streamer,
-        max_new_tokens=4096,
-        temperature=0.7,
-        do_sample=True
-    )
+    # Process output
+    story = out[0]['generated_text'].split("<|im_start|>assistant\n")[-1]
+    story = story.replace("<|im_end|>", "").strip()
     
-    thread = Thread(target=model.generate, kwargs=generation_kwargs)
-    thread.start()
-    
-    full_response = ""
-    thinking_content = ""
-    for new_text in streamer:
-        full_response += new_text
-        try:
-            current_ids = tokenizer.encode(full_response, return_tensors="pt")[0]
-            thinking, content = parse_thinking_output(current_ids, tokenizer)
-            yield thinking, content
-        except:
-            yield "", full_response
+    # Enforce ≤100 words and proper ending
+    words = story.split()
+    if len(words) > 100:
+        story = " ".join(words[:100])
+    if not story.endswith(('.', '!', '?')):
+        story += '.'
 
-def text_to_speech(text):
-    tts = gTTS(text=text, lang='en', slow=False)
-    audio_file = f"audio_{int(time.time())}.mp3"
-    tts.save(audio_file)
-    return audio_file
+    # Display story
+    st.subheader("📚 Your Magical Story")
+    st.write(story)
 
-# Streamlit UI
-def main():
-    st.title("🧠 Qwen3-1.7B Thinking Mode Demo")
-    
-    model, tokenizer = load_models()
-    
-    with st.sidebar:
-        st.header("Settings")
-        max_length = st.slider("Max Tokens", 100, 4096, 1024)
-        temperature = st.slider("Temperature", 0.1, 1.0, 0.7)
-    
-    prompt = st.text_area("Enter your prompt:", 
-                         "Explain quantum computing in simple terms")
-    
-    if st.button("Generate Response"):
-        with st.spinner("Generating response..."):
-            # Setup containers
-            thinking_container = st.container(border=True)
-            response_container = st.empty()
-            audio_container = st.empty()
-            
-            full_content = ""
-            current_thinking = ""
-            
-            for thinking, content in generate_response(prompt, model, tokenizer):
-                if thinking != current_thinking:
-                    thinking_container.markdown(f"**Thinking Process:**\n{thinking}")
-                    current_thinking = thinking
-                
-                if content != full_content:
-                    response_container.markdown(f"**Final Answer:**\n{content}")
-                    full_content = content
-            
-            # Add audio version
-            audio_file = text_to_speech(full_content)
-            audio_container.audio(audio_file, format='audio/mp3')
-            
-            # Add download button
-            st.download_button(
-                label="Download Response",
-                data=full_content,
-                file_name="qwen_response.txt",
-                mime="text/plain"
-            )
+    # Convert to audio
+    with st.spinner("🔊 Converting to audio..."):
+        try:
+            tts = gTTS(text=story, lang="en", slow=False)
+            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}")
 
-if __name__ == "__main__":
-    main()
+# Footer
+st.markdown("---\nMade with ❤️ by your friendly story wizard")