Update app.py

app.py

@@ -22,7 +22,7 @@ import requests
 from io import BytesIO
 import json
 import torch
-from diffusers import DiffusionPipeline, StableDiffusionPipeline
+from diffusers import DiffusionPipeline
 import torch
 from transformers import pipeline
 
@@ -229,13 +229,12 @@ def load_models():
     models = {}
     
     try:
-        # Use a lightweight model that works well on CPU
+        # Use a fast, optimized model
         models['text_to_image'] = DiffusionPipeline.from_pretrained(
-            "runwayml/stable-diffusion-v1-5", 
-            torch_dtype=torch.float32,  # Use float32 for CPU compatibility
-            safety_checker=None  # Disable safety checker for faster performance
+            "OFA-Sys/small-stable-diffusion-v0",  # Optimized for speed
+            torch_dtype=torch.float32,
+            safety_checker=None
         )
-        # 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:
@@ -294,10 +293,10 @@ def create_story_scene(story, concept, models):
         # Generate image with optimized settings
         image = models['text_to_image'](
             prompt,
-            num_inference_steps=20,  # Reduced steps for faster generation
+            num_inference_steps=15,  # Faster generation
             guidance_scale=7.5,
-            height=384,  # Smaller image for faster generation
-            width=384
+            height=256,  # Smaller image for faster generation
+            width=256
         ).images[0]
         
         # Convert to bytes
@@ -549,7 +548,7 @@ def main():
                     with st.spinner("🚀 Loading AI models..."):
                         ai_models = load_models()
                     
-                    with st.spinner("🎨 Generating AI story scene (this may take 10-20 seconds)..."):
+                    with st.spinner("🎨 Generating AI story scene (this may take 5-10 seconds)..."):
                         st.session_state.story_scene = create_story_scene(
                             story, main_concept, ai_models
                         )
@@ -580,19 +579,19 @@ def main():
                 st.caption("Loop Example")
                 st.code('"A dragon breathes fire 5 times at the castle"', language="text")
                 st.image(create_animation_preview("", "loop"), 
-                         use_column_width=True, 
+                         use_container_width=True, 
                          caption="Loop Animation Preview")
             with col2:
                 st.caption("Conditional Example")
                 st.code('"If it rains, the cat stays inside, else it goes out"', language="text")
                 st.image(create_animation_preview("", "conditional"), 
-                         use_column_width=True, 
+                         use_container_width=True, 
                          caption="Conditional Animation Preview")
             with col3:
                 st.caption("Function Example")
                 st.code('"A wizard casts a spell to make flowers grow"', language="text")
                 st.image(create_animation_preview("", "function"), 
-                         use_column_width=True, 
+                         use_container_width=True, 
                          caption="Function Animation Preview")
     
     # Animation tab
@@ -677,7 +676,7 @@ def main():
                     # Show animation preview for the concept
                     st.image(create_animation_preview("", concept), 
                              caption=f"{details['name']} Animation Example", 
-                             use_column_width=True)
+                             use_container_width=True)
         
         if st.button("See the Magic Code!", use_container_width=True):
             st.session_state.active_tab = "code"
@@ -717,7 +716,7 @@ def main():
             concept = st.session_state.concepts[0] if st.session_state.concepts else "loop"
             st.image(create_animation_preview("", concept), 
                     caption="This is similar to what your animation would look like", 
-                    use_column_width=True)
+                    use_container_width=True)
         else:
             st.warning("No code generated yet!")