Update app.py

app.py

@@ -2,6 +2,8 @@
 import torch
 from diffusers import StableDiffusionPipeline, EulerAncestralDiscreteScheduler
 import gradio as gr
+from PIL import Image
+import io
 import time
 
 # Force CPU usage
@@ -9,21 +11,20 @@ device = "cpu"
 print(f"Using device: {device}")
 
 # Load a smaller, CPU-friendly PUBLIC model
-# Using 'dreamlike-art/dreamlike-diffusion-1.0' - a public fine-tuned model
 model_id = "dreamlike-art/dreamlike-diffusion-1.0"
 
 print("Loading pipeline... This may take a few minutes.")
 try:
-    # Use torch.float32 for CPU compatibility. DO NOT use float16.
+    # Use torch.float32 for CPU compatibility
     pipe = StableDiffusionPipeline.from_pretrained(
         model_id,
-        torch_dtype=torch.float32, # Critical for CPU
+        torch_dtype=torch.float32,
         use_safetensors=True,
-        safety_checker=None,  # Disable safety checker to avoid NSFW filters that can cause issues
-        requires_safety_checker=False  # Newer versions of diffusers need this
+        safety_checker=None,
+        requires_safety_checker=False
     )
     
-    # Optional: Use a faster scheduler for quicker generation
+    # Use a faster scheduler for quicker generation
     pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
     
     # Move the pipeline to the CPU
@@ -32,7 +33,6 @@ try:
 
 except Exception as e:
     print(f"Error loading model: {e}")
-    # Provide a more helpful error message
     raise e
 
 # Define the image generation function
@@ -40,21 +40,24 @@ def generate_image(prompt):
     """
     This function takes a text prompt and returns a generated image.
     """
-    # Add a consistent style to all prompts to get a children's book look
+    # Add a consistent style to all prompts
     enhanced_prompt = f"children's book illustration, watercolor style, cute, whimsical, {prompt}"
-
     print(f"Generating image for prompt: {enhanced_prompt}")
     
-    # Generate the image. Use a lower resolution for speed on CPU.
+    # Generate the image
     image = pipe(
         prompt=enhanced_prompt,
-        width=512,   # Smaller image = faster generation
+        width=512,
         height=512,
         guidance_scale=7.5,
-        num_inference_steps=20,  # Fewer steps = much faster
-        generator=torch.Generator(device=device) # Ensure generator is on CPU
+        num_inference_steps=20,
+        generator=torch.Generator(device=device)
     ).images[0]
-
+    
+    # Convert to RGB to ensure proper color format
+    if image.mode != 'RGB':
+        image = image.convert('RGB')
+    
     print("Image generated successfully!")
     return image
 
@@ -71,5 +74,10 @@ demo = gr.Interface(
     description="This free version runs on CPU. It's slower but gets the job done! Enter a scene description."
 )
 
-# Launch the app
-demo.launch(debug=True, server_name="0.0.0.0")
+# Launch the app with more robust settings
+demo.launch(
+    debug=True,
+    server_name="0.0.0.0",
+    share=False,
+    enable_queue=True  # This helps handle multiple requests properly
+)