Update app.py

app.py

@@ -1,4 +1,4 @@
-# Must be FIRST import and FIRST Streamlit command
+# FIRST import and FIRST Streamlit command
 import streamlit as st
 st.set_page_config(
     page_title="Magic Story Generator",
@@ -6,9 +6,10 @@ st.set_page_config(
     page_icon="📖"
 )
 
-# Other imports AFTER Streamlit config
+# Other imports
 import re
 import time
+import torch
 import tempfile
 from PIL import Image
 from gtts import gTTS
@@ -24,21 +25,22 @@ def load_models():
     captioner = pipeline(
         "image-to-text",
         model="Salesforce/blip-image-captioning-base",
-        device=-1  # Use -1 for CPU, 0 for GPU
+        device=0 if torch.cuda.is_available() else -1
     )
     
-    # Story generation model (Qwen3-1.7B)
+    # Optimized story generation model
     storyteller = pipeline(
         "text-generation",
-        model="Qwen/Qwen3-0.6B",
+        model="Qwen/Qwen3-0.5B",
         device_map="auto",
         trust_remote_code=True,
-        torch_dtype="auto",
-        max_new_tokens=230,
-        temperature=0.8,
+        model_kwargs={"load_in_8bit": True},
+        torch_dtype=torch.float16,
+        max_new_tokens=200,
+        temperature=0.9,
         top_k=50,
-        top_p=0.85,
-        repetition_penalty=1.15,
+        top_p=0.9,
+        repetition_penalty=1.1,
         eos_token_id=151645
     )
     
@@ -55,12 +57,16 @@ uploaded_image = st.file_uploader(
 if uploaded_image:
     # Process image
     image = Image.open(uploaded_image).convert("RGB")
-    st.image(image, use_container_width=True)
+    st.image(image, use_column_width=True)
 
     # Generate caption
     with st.spinner("🔍 Analyzing image..."):
-        caption_result = caption_pipe(image)
-        image_caption = caption_result[0].get("generated_text", "").strip()
+        try:
+            caption_result = caption_pipe(image)
+            image_caption = caption_result[0].get("generated_text", "").strip()
+        except Exception as e:
+            st.error(f"❌ Image analysis failed: {str(e)}")
+            st.stop()
     
     if not image_caption:
         st.error("❌ Couldn't understand this image. Please try another!")
@@ -71,62 +77,38 @@ if uploaded_image:
     # 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"
+        f"You're a children's author. Create a short story (100-150 words) based on: {image_caption}\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()
-    
-    # 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:])
+    # Generate story with progress
+    progress_bar = st.progress(0)
+    status_text = st.empty()
     
-    final_story = " ".join(final_story).replace("..", ".")[:800]
-
-    # Display story
-    st.subheader("✨ Your Magical Story")
-    st.write(final_story)
-
-    # Audio conversion
-    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")
-        except Exception as e:
-            st.error(f"❌ Audio conversion failed: {str(e)}")
+    try:
+        with st.spinner("📝 Crafting magical story..."):
+            start_time = time.time()
+            
+            def update_progress(step):
+                progress = min(step/5, 1.0)  # Simulate progress steps
+                progress_bar.progress(progress)
+                status_text.text(f"Step {int(step)}/5: {'📖'*int(step)}")
+            
+            update_progress(1)
+            story_result = story_pipe(
+                story_prompt,
+                do_sample=True,
+                num_return_sequences=1
+            )
+            
+            update_progress(4)
+            generation_time = time.time() - start_time
+            st.info(f"Story generated in {generation_time:.1f} seconds")
 
-# Footer
-st.markdown("---")
-st.markdown("📚 Made with ♥ by The Story Wizard • [Report Issues](https://example.com)")
+            # Process output
+            raw_story = story_result[0]['generated_text']
+            clean_story = raw_story.split("<|im_start|>assistant\n")[-1]
+            clean_story = re.sub(r'<\|.*?\|>', '', clean_story).strip()
+            
+            # Format story text
+            sentences = []
+            for sent in re