Update app.py

app.py

@@ -1,4 +1,4 @@
-# Import Streamlit first
+# FIRST import and FIRST Streamlit command
 import streamlit as st
 st.set_page_config(
     page_title="Magic Story Generator",
@@ -13,114 +13,142 @@ import torch
 import tempfile
 from PIL import Image
 from gtts import gTTS
-from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer
+from transformers import pipeline, AutoTokenizer
 
-# --- Initialize Models First ---
+# --- Constants & Setup ---
+st.title("📖✨ Turn Images into Children's Stories")
+
+# --- Model Loading (Cached) ---
 @st.cache_resource(show_spinner=False)
 def load_models():
-    """Load and return both models at startup"""
-    try:
-        # 1. Image Captioning Model
-        caption_pipe = pipeline(
-            "image-to-text",
-            model="Salesforce/blip-image-captioning-base",
-            device=0 if torch.cuda.is_available() else -1
-        )
-
-        # 2. Story Generation Model
-        story_tokenizer = AutoTokenizer.from_pretrained(
-            "Qwen/Qwen3-0.6B",
-            trust_remote_code=True
-        )
-        
-        story_model = AutoModelForCausalLM.from_pretrained(
-            "Qwen/Qwen3-0.6B",
-            device_map="auto",
-            torch_dtype=torch.float16
-        )
-
-        story_pipe = pipeline(
-            "text-generation",
-            model=story_model,
-            tokenizer=story_tokenizer,
-            max_new_tokens=230,
-            temperature=0.9,
-            top_k=50,
-            top_p=0.9,
-            repetition_penalty=1.1,
-            eos_token_id=151645
-        )
-
-        return caption_pipe, story_pipe
-
-    except Exception as e:
-        st.error(f"🚨 Model loading failed: {str(e)}")
-        st.stop()
+    # Image captioning model
+    captioner = pipeline(
+        "image-to-text",
+        model="Salesforce/blip-image-captioning-base",
+        device=0 if torch.cuda.is_available() else -1
+    )
+    
+    # Optimized story generation model
+    tokenizer = AutoTokenizer.from_pretrained("Deepthoughtworks/gpt-neo-2.7B__low-cpu")
+    storyteller = pipeline(
+        "text-generation",
+        model="Deepthoughtworks/gpt-neo-2.7B__low-cpu",
+        tokenizer=tokenizer,
+        device_map="auto",
+        torch_dtype=torch.float32,  # Changed to float32 for better CPU compatibility
+        max_new_tokens=150,         # Reduced length for faster generation
+        temperature=0.85,
+        top_k=40,
+        top_p=0.92,
+        repetition_penalty=1.15,
+        pad_token_id=tokenizer.eos_token_id  # Added for padding control
+    )
+    
+    return captioner, storyteller
 
-# Initialize models immediately when app starts
 caption_pipe, story_pipe = load_models()
 
-# --- Rest of Application ---
-st.title("📖✨ Turn Images into Children's Stories")
-
-def clean_story_text(raw_text):
-    """Improved cleaning function"""
-    clean = re.sub(r'<\|.*?\|>', '', raw_text)  # Remove special tokens
-    clean = re.sub(r'Okay, I need.*?(?=\n|$)', '', clean, flags=re.DOTALL)  # Remove thinking chains
-    return clean.strip()
-
+# --- Main Application Flow ---
 uploaded_image = st.file_uploader(
     "Upload a children's book style image:",
     type=["jpg", "jpeg", "png"]
 )
 
 if uploaded_image:
+    # Process image
     image = Image.open(uploaded_image).convert("RGB")
-    # Updated parameter here ↓
-    st.image(image, use_container_width=True)  # Changed use_column_width to use_container_width
+    st.image(image, use_container_width=True)  # Fixed deprecated parameter
 
+    # Generate caption
     with st.spinner("🔍 Analyzing image..."):
         try:
             caption_result = caption_pipe(image)
-            image_caption = caption_result[0].get("generated_text", "")
-            st.success(f"**Image Understanding:** {image_caption}")
+            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!")
+        st.stop()
+    
+    st.success(f"**Image Understanding:** {image_caption}")
 
-    # Story generation prompt
-    story_prompt = f"""Write a children's story about: {image_caption}
-Rules:
-- Use simple words (Grade 2 level)
-- Exclude thinking processes
-- 3 paragraphs maximum
+    # Create story prompt
+    story_prompt = f"""Write a 50 to 100 words children's story based on: {image_caption}
+Requirements:
+- Exclude your thinking process
 Story:"""
 
+    # Generate story with progress
+    progress_bar = st.progress(0)
+    status_text = st.empty()
+    
     try:
         with st.spinner("📝 Crafting magical story..."):
+            start_time = time.time()
+            
+            def update_progress(step):
+                progress = min(step/5, 1.0)
+                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,
-                top_p=0.9,
-                repetition_penalty=1.2
+                num_return_sequences=1
             )
             
-            raw_story = story_result[0]['generated_text']
-            final_story = clean_story_text(raw_story.split("Story:")[-1])
-
-            st.subheader("✨ Your Magical Story")
-            st.write(final_story)
+            update_progress(4)
+            generation_time = time.time() - start_time
+            st.info(f"Story generated in {generation_time:.1f} seconds")
 
-            # Audio conversion
-            with st.spinner("🔊 Creating audio version..."):
-                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")
+            # Process output
+            raw_story = story_result[0]['generated_text']
+            clean_story = raw_story.split("Story:")[-1].strip()
+            clean_story = re.sub(r'\n+', '\n\n', clean_story)  # Improve paragraph spacing
+            
+            # Format story text
+            final_story = ""
+            for paragraph in clean_story.split('\n\n'):
+                paragraph = paragraph.strip()
+                if paragraph:
+                    sentences = []
+                    for sent in re.split(r'(?<=[.!?]) +', paragraph):
+                        sent = sent.strip()
+                        if sent:
+                            if len(sent) > 1 and not sent.endswith(('.','!','?')):
+                                sent += '.'
+                            sentences.append(sent[0].upper() + sent[1:])
+                    final_story += ' '.join(sentences) + '\n\n'
+
+            update_progress(5)
+            time.sleep(0.5)
 
     except Exception as e:
         st.error(f"❌ Story generation failed: {str(e)}")
+        st.stop()
+
+    finally:
+        progress_bar.empty()
+        status_text.empty()
+
+    # Display story
+    st.subheader("✨ Your Magical Story")
+    st.write(final_story.strip())
+
+    # 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)}")
 
 # Footer
 st.markdown("---")
 st.markdown("📚 Made with ♥ by The Story Wizard • [Report Issues](https://example.com)")
+