Update app.py

app.py

@@ -181,19 +181,38 @@ def generate_image_prompts(story_text):
 
 @torch.inference_mode()
 @spaces.GPU(duration=60)
-def generate_story_image(prompt):
+def generate_story_image(prompt, seed=-1):
+    """Generate an image using Stable Diffusion with LoRA temporarily loaded."""
     torch.cuda.empty_cache()
     
-    pipe.load_lora_weights("Prof-Hunt/lora-bulldog")
+    generator = torch.Generator("cuda")
+    if seed != -1:
+        generator.manual_seed(seed)
+    else:
+        generator.manual_seed(torch.randint(0, 2**32 - 1, (1,)).item())
+
     enhanced_prompt = f"{prompt}, watercolor style, children's book illustration, soft colors"
-    
-    image = pipe(
-        prompt=enhanced_prompt,
-        negative_prompt="deformed, ugly, blurry, bad art, poor quality, distorted",
-        num_inference_steps=50,
-        guidance_scale=15,
-    ).images[0]
-    
+
+    try:
+        # Load LoRA only for this function
+        pipe.load_lora_weights("Prof-Hunt/lora-bulldog")
+
+        image = pipe(
+            prompt=enhanced_prompt,
+            negative_prompt="deformed, ugly, blurry, bad art, poor quality, distorted",
+            num_inference_steps=50,
+            guidance_scale=15,
+            generator=generator
+        ).images[0]
+
+        # Unload LoRA properly
+        pipe.unload_lora_weights()
+        torch.cuda.empty_cache()
+
+    except Exception as e:
+        print(f"Error generating image: {e}")
+        return None
+
     return image
 
 @torch.inference_mode()
@@ -221,20 +240,28 @@ def generate_all_scenes(prompts_text):
         
         if scene_prompt:
             try:
-                torch.cuda.empty_cache()  # Clear GPU memory before generation
+                torch.cuda.empty_cache()
                 print(f"Generating image for scene: {scene_prompt}")  # Debugging
+                
                 image = generate_story_image(scene_prompt)
                 
                 if image is not None:
-                    generated_images.append(np.array(image))
+                    img_array = np.array(image)
+                    
+                    # Ensure the image is valid
+                    if img_array.shape[0] > 0:
+                        generated_images.append(img_array)
                 
-                torch.cuda.empty_cache()  # Clear memory after each image
+                torch.cuda.empty_cache()
             except Exception as e:
                 print(f"Error generating image: {str(e)}")
                 continue
     
+    print(f"Generated {len(generated_images)} images.")
     return generated_images, "\n\n".join(formatted_prompts)
 
+
+
 # Helper functions without GPU usage
 def clean_story_output(story):
     story = story.replace("<|im_end|>", "")