Update app.py

app.py

@@ -1,21 +1,18 @@
-# Must be FIRST import and FIRST Streamlit command
-import streamlit as st
-st.set_page_config(
-    page_title="Magic Story Generator",
-    layout="centered",
-    page_icon="📖"
-)
-
-# Other imports AFTER Streamlit config
+# story_generator.py
 import re
 import time
 import tempfile
+import streamlit as st
 from PIL import Image
 from gtts import gTTS
 from transformers import pipeline
 
-# --- Constants & Setup ---
-st.title("📖✨ Turn Images into Children's Stories")
+# --- Initialize Streamlit Config ---
+st.set_page_config(
+    page_title="Magic Story Generator",
+    layout="centered",
+    page_icon="📖"
+)
 
 # --- Model Loading (Cached) ---
 @st.cache_resource(show_spinner=False)
@@ -24,113 +21,134 @@ def load_models():
     captioner = pipeline(
         "image-to-text",
         model="Salesforce/blip-image-captioning-base",
-        device=-1  # Use -1 for CPU, 0 for GPU
+        device=-1
     )
     
-    # Story generation model (Qwen3-1.7B)
+    # Story generation model with optimized settings
     storyteller = pipeline(
         "text-generation",
         model="Qwen/Qwen3-1.7B",
         device_map="auto",
         trust_remote_code=True,
         torch_dtype="auto",
-        max_new_tokens=250,
-        temperature=0.7,
-        top_p=0.85,
-        repetition_penalty=1.15,
-        eos_token_id=151645
+        model_kwargs={
+            "revision": "main",
+            "temperature": 0.7,
+            "top_p": 0.9,
+            "repetition_penalty": 1.1,
+            "pad_token_id": 151645
+        }
     )
     
     return captioner, storyteller
 
-caption_pipe, story_pipe = load_models()
+# --- Text Processing Utilities ---
+def clean_generated_text(raw_text):
+    # Split at first assistant marker
+    clean_text = raw_text.split("<|im_start|>assistant\n", 1)[-1]
+    
+    # Remove any subsequent chat turns
+    clean_text = clean_text.split("<|im_start|>")[0]
+    
+    # Remove special tokens and whitespace
+    clean_text = clean_text.replace("<|im_end|>", "").strip()
+    
+    # Regex cleanup for remaining markers
+    clean_text = re.sub(
+        r'^(assistant[\s\-\:>]*)+',
+        '',
+        clean_text,
+        flags=re.IGNORECASE
+    ).strip()
+    
+    # Format punctuation and capitalization
+    sentences = []
+    for sent in re.split(r'(?<=[.!?]) +', clean_text):
+        sent = sent.strip()
+        if not sent:
+            continue
+        if sent[-1] not in {'.', '!', '?'}:
+            sent += '.'
+        sentences.append(sent[0].upper() + sent[1:])
+    
+    return ' '.join(sentences)
+
+# --- Main Application UI ---
+st.title("📖✨ Magic Story Generator")
 
-# --- Main Application Flow ---
 uploaded_image = st.file_uploader(
     "Upload a children's book style image:",
     type=["jpg", "jpeg", "png"]
 )
 
 if uploaded_image:
-    # Process image
+    # Display uploaded image
     image = Image.open(uploaded_image).convert("RGB")
-    st.image(image, use_container_width=True)
-
-    # Generate caption
-    with st.spinner("🔍 Analyzing image..."):
-        caption_result = caption_pipe(image)
-        image_caption = caption_result[0].get("generated_text", "").strip()
+    st.image(image, use_column_width=True)
     
-    if not image_caption:
-        st.error("❌ Couldn't understand this image. Please try another!")
-        st.stop()
+    # Load models only when needed
+    caption_pipe, story_pipe = load_models()
+    
+    # Generate image caption
+    with st.spinner("🔍 Analyzing image..."):
+        try:
+            caption_result = caption_pipe(image)
+            image_caption = caption_result[0].get("generated_text", "").strip()
+            
+            if not image_caption:
+                raise ValueError("Couldn't generate caption")
+                
+            st.success(f"**Image Understanding:** {image_caption}")
+        except Exception as e:
+            st.error("❌ Failed to analyze image. Please try another.")
+            st.stop()
     
-    st.success(f"**Image Understanding:** {image_caption}")
-
     # Create story prompt
     story_prompt = (
         f"<|im_start|>system\n"
-        f"You are a children's book author. Create a 100-150 word story based on: {image_caption}\n"
-        "Use simple language, friendly characters, and a positive lesson.<|im_end|>\n"
+        f"You are a children's book author. Create a 150-word story based on: {image_caption}\n"
+        "Include these elements:\n"
+        "- Friendly characters\n"
+        "- Simple vocabulary\n"
+        "- Positive lesson\n"
+        "- Clear story structure\n"
+        "<|im_end|>\n"
         f"<|im_start|>user\n"
-        f"Write a child-friendly story with a clear beginning, middle, and end.<|im_end|>\n"
+        f"Write an engaging story suitable for ages 6-8.<|im_end|>\n"
         f"<|im_start|>assistant\n"
     )
-
-    # Generate story
-    with st.spinner("📝 Crafting magical story..."):
-        start_time = time.time()
-        story_result = story_pipe(
-            story_prompt,
-            do_sample=True,
-            num_return_sequences=1,
-            pad_token_id=151645
-        )
-        generation_time = time.time() - start_time
-
-    # Process output
-    raw_story = story_result[0]['generated_text']
     
-    # Clean up story text
-    clean_story = raw_story.split("<|im_start|>assistant\n")[-1]
-    clean_story = clean_story.split("<|im_start|>")[0]  # Remove any new turns
-    clean_story = clean_story.replace("<|im_end|>", "").strip()
+    # Generate story text
+    with st.spinner("📝 Crafting magical story..."):
+        try:
+            story_result = story_pipe(
+                story_prompt,
+                max_new_tokens=300,
+                do_sample=True,
+                num_return_sequences=1
+            )
+            raw_story = story_result[0]['generated_text']
+        except Exception as e:
+            st.error("❌ Story generation failed. Please try again.")
+            st.stop()
     
-    # Remove assistant mentions using regex
-    clean_story = re.sub(
-        r'^(assistant[:>]?\s*)+', 
-        '', 
-        clean_story, 
-        flags=re.IGNORECASE
-    ).strip()
-
-    # Format story punctuation
-    final_story = []
-    for sentence in clean_story.split(". "):
-        sentence = sentence.strip()
-        if not sentence:
-            continue
-        if not sentence.endswith('.'):
-            sentence += '.'
-        final_story.append(sentence[0].upper() + sentence[1:])
+    # Process and display story
+    final_story = clean_generated_text(raw_story)
     
-    final_story = " ".join(final_story).replace("..", ".")[:800]
-
-    # Display story
-    st.subheader("✨ Your Magical Story")
+    st.subheader("✨ Your Story")
     st.write(final_story)
-
-    # Audio conversion
+    
+    # Generate audio version
     with st.spinner("🔊 Creating audio version..."):
         try:
-            audio = gTTS(text=final_story, lang="en", slow=False)
-            with tempfile.NamedTemporaryFile(delete=False, suffix=".mp3") as tmp_file:
-                audio.save(tmp_file.name)
-                st.audio(tmp_file.name, format="audio/mp3")
+            tts = gTTS(text=final_story, lang='en', slow=False)
+            with tempfile.NamedTemporaryFile(delete=False, suffix=".mp3") as fp:
+                tts.save(fp.name)
+                st.audio(fp.read(), format="audio/mp3")
         except Exception as e:
-            st.error(f"❌ Audio conversion failed: {str(e)}")
+            st.warning("⚠️ Audio conversion failed. Text version still available.")
 
 # Footer
 st.markdown("---")
-st.markdown("📚 Made with ♥ by The Story Wizard • [Report Issues](https://example.com)")
+st.caption("Made with ♥ by The Story Wizard • [Report Issues](https://example.com)")