Update app.py

app.py

@@ -28,12 +28,22 @@ def generate_comic_strip(story):
     # Split the story into sentences to identify key scenes
     scenes = split_into_sentences(story)
     
+    # Limit the number of scenes to prevent excessive image generation
+    max_scenes = 3
+    scenes = scenes[:max_scenes]
+    
     comic_strip = []
     
     for idx, scene in enumerate(scenes):
         try:
-            # Generate image for each scene
-            image = pipe(scene, num_inference_steps=50).images[0]
+            # Generate image for each scene with optimizations
+            image = pipe(
+                scene, 
+                num_inference_steps=20,        # Reduced steps for faster generation
+                height=256,                    # Reduced resolution
+                width=256,                     # Reduced resolution
+                guidance_scale=7.5,            # Default guidance scale
+            ).images[0]
             comic_strip.append(image)
         except Exception as e:
             # In case of any error during image generation, append a placeholder image
@@ -73,7 +83,7 @@ def main():
             ["An astronaut lands on a distant planet and encounters alien life forms."]
         ],
         allow_flagging="never",          # Disable flagging of outputs
-        theme="default"                  # You can choose other themes like "huggingface"
+        theme="default",                 # You can choose other themes like "huggingface"
     )
     
     # Launch the Gradio app
@@ -83,10 +93,13 @@ def main():
 pipe = None
 try:
     print("Loading Stable Diffusion model...")
-    # Initialize the Stable Diffusion pipeline
+    # Initialize the Stable Diffusion pipeline with optimizations
     pipe = StableDiffusionPipeline.from_pretrained(
         "stabilityai/stable-diffusion-2-1-base",  # Model name
         torch_dtype=torch.float32,                # Use float32 for CPU
+        low_cpu_mem_usage=True,                    # Optimize for CPU usage
+        safety_checker=None,                       # Disable safety checker to speed up loading
+        force_download=True                        # Force download to avoid resume_download warning
     )
     pipe = pipe.to("cpu")  # Move the model to CPU
     pipe.enable_attention_slicing()  # Reduce memory usage