Update app.py

app.py

@@ -229,13 +229,13 @@ def load_models():
     models = {}
     
     try:
-        # Use CPU-friendly model and float32 precision
+        # Use a lightweight model that works well on CPU
         models['text_to_image'] = DiffusionPipeline.from_pretrained(
             "runwayml/stable-diffusion-v1-5", 
-            torch_dtype=torch.float32,  # Use float32 instead of float16
-            safety_checker=None  # Disable safety checker for simplicity
+            torch_dtype=torch.float32,  # Use float32 for CPU compatibility
+            safety_checker=None  # Disable safety checker for faster performance
         )
-        # Move to GPU if available, otherwise keep on CPU
+        # Use GPU if available, otherwise CPU
         device = "cuda" if torch.cuda.is_available() else "cpu"
         models['text_to_image'].to(device)
     except Exception as e:
@@ -291,13 +291,13 @@ def create_story_scene(story, concept, models):
         style = "cartoon style, bright colors, children's book illustration"
         prompt = f"{story}. {CONCEPTS[concept]['name']} concept. {style}"
         
-        # Generate image with fixed settings
+        # Generate image with optimized settings
         image = models['text_to_image'](
             prompt,
-            num_inference_steps=25,  # Fewer steps for faster generation
+            num_inference_steps=20,  # Reduced steps for faster generation
             guidance_scale=7.5,
-            height=512,
-            width=512
+            height=384,  # Smaller image for faster generation
+            width=384
         ).images[0]
         
         # Convert to bytes
@@ -473,9 +473,6 @@ def main():
     st.title("🧙‍♂️ StoryCoder - Learn Python Through Stories!")
     st.subheader("Turn your story into an animation and discover coding secrets!")
     
-    # Load AI models
-    ai_models = load_models()
-    
     # Initialize session state
     if 'story' not in st.session_state:
         st.session_state.story = ""
@@ -548,7 +545,11 @@ def main():
                     # Get the main concept
                     main_concept = st.session_state.concepts[0] if st.session_state.concepts else "variable"
                     
-                    with st.spinner("🎨 Generating AI story scene..."):
+                    # Load models only when needed
+                    with st.spinner("🚀 Loading AI models..."):
+                        ai_models = load_models()
+                    
+                    with st.spinner("🎨 Generating AI story scene (this may take 10-20 seconds)..."):
                         st.session_state.story_scene = create_story_scene(
                             story, main_concept, ai_models
                         )