Update app.py

app.py

@@ -1,144 +1,64 @@
-# story_generator.py
-import re
-import time
-import tempfile
 import streamlit as st
+import torch
 from PIL import Image
 from gtts import gTTS
-from transformers import pipeline
+from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer
 
-# --- Initialize Streamlit Config ---
-st.set_page_config(
-    page_title="Magic Story Generator",
-    layout="centered",
-    page_icon="📖"
-)
+# Streamlit config must be first
+st.set_page_config(page_title="Magic Story Generator", layout="centered", page_icon="📖")
 
-# --- Model Loading (Cached) ---
-@st.cache_resource(show_spinner=False)
+# Model loading cached for performance
+@st.cache_resource
 def load_models():
-    # Image captioning model
-    captioner = pipeline(
-        "image-to-text",
-        model="Salesforce/blip-image-captioning-base",
-        device=-1  # Force CPU usage
-    )
-    
-    # Story generation model with updated parameters
-    storyteller = pipeline(
-        "text-generation",
-        model="Qwen/Qwen3-1.7B",
+    caption_model = pipeline("image-to-text", "Salesforce/blip-image-captioning-base")
+    story_model = AutoModelForCausalLM.from_pretrained(
+        "Qwen/Qwen3-1.7B",
         device_map="auto",
-        trust_remote_code=True,
-        torch_dtype="auto",
-        temperature=0.7,
-        top_p=0.9,
-        repetition_penalty=1.1,
-        pad_token_id=151645,
-        max_new_tokens=300
+        torch_dtype=torch.float16,
+        trust_remote_code=True
     )
-    
-    return captioner, storyteller
+    story_tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-1.7B", trust_remote_code=True)
+    return caption_model, story_model, story_tokenizer
 
-# --- 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")
+# Initialize models
+caption_pipe, story_model, story_tokenizer = load_models()
 
-uploaded_image = st.file_uploader(
-    "Upload a children's book style image:",
-    type=["jpg", "jpeg", "png"]
-)
+# Main app interface
+st.title("📖 Instant Story Generator")
+uploaded_image = st.file_uploader("Upload an image:", type=["jpg", "jpeg", "png"])
 
 if uploaded_image:
-    # Display uploaded image with modern parameter
-    image = Image.open(uploaded_image).convert("RGB")
-    st.image(image, use_container_width=True)  # Updated parameter
+    img = Image.open(uploaded_image).convert("RGB")
+    st.image(img, caption="Your Image", use_column_width=True)
     
-    # Load models only when needed
-    try:
-        caption_pipe, story_pipe = load_models()
-    except Exception as e:
-        st.error(f"❌ Model loading failed: {str(e)}")
-        st.stop()
+    # Generate caption
+    caption = caption_pipe(img)[0]['generated_text']
     
-    # 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("Empty caption generated")
-                
-            st.success(f"**Image Understanding:** {image_caption}")
-        except Exception as e:
-            st.error("❌ Image analysis failed. Please try another image.")
-            st.stop()
-    
-    # Create story prompt
-    story_prompt = (
-        f"<|im_start|>system\n"
-        f"You are a children's book author. Create a 150-word story based on: {image_caption}\n"
-    )
+    # Generate story
+    messages = [{
+        "role": "system",
+        "content": f"Create a 50 to 100 words children's story based on: {caption}."
+    }]
     
-    # Generate story text
-    with st.spinner("📝 Crafting magical story..."):
-        try:
-            story_result = story_pipe(
-                story_prompt,
-                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()
+    inputs = story_tokenizer.apply_chat_template(
+        messages,
+        return_tensors="pt"
+    ).to(story_model.device)
     
-    # Process and display story
-    final_story = clean_generated_text(raw_story)
+    outputs = story_model.generate(
+        inputs,
+        max_new_tokens=300,
+        temperature=0.7,
+        top_p=0.9
+    )
     
-    st.subheader("✨ Your Story")
-    st.write(final_story)
+    # Display results
+    story = story_tokenizer.decode(outputs[0][len(inputs[0]):], skip_special_tokens=True)
+    st.subheader("Generated Story")
+    st.write(story)
     
-    # Generate audio version
-    with st.spinner("🔊 Creating audio version..."):
-        try:
-            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.warning("⚠️ Audio conversion failed. Text version still available.")
-
-# Footer
-st.markdown("---")
-st.caption("Made with ♥ by The Story Wizard • [Report Issues](https://example.com)")
+    # Audio conversion
+    audio = gTTS(text=story, lang='en')
+    with tempfile.NamedTemporaryFile(delete=False) as fp:
+        audio.save(fp.name)
+        st.audio(fp.name, format='audio/mp3')