Update app.py

app.py

@@ -1,23 +1,37 @@
 # Import necessary libraries
 import torch
-from diffusers import StableDiffusionXLPipeline
+from diffusers import StableDiffusionPipeline, EulerAncestralDiscreteScheduler
 import gradio as gr
 import time
 
-# Check if we have a GPU (CUDA) or need to use the CPU
-device = "cuda" if torch.cuda.is_available() else "cpu"
+# Force CPU usage
+device = "cpu"
+print(f"Using device: {device}")
 
-# Load the Stable Diffusion XL pipeline
-# We use torch_dtype=torch.float16 for faster generation and less memory usage
-print("Loading Stable Diffusion XL pipeline... This may take a few minutes.")
-pipe = StableDiffusionXLPipeline.from_pretrained(
-    "stabilityai/stable-diffusion-xl-base-1.0",
-    torch_dtype=torch.float16,
-    use_safetensors=True
-)
-# Move the pipeline to the chosen device (GPU or CPU)
-pipe = pipe.to(device)
-print("Model loaded successfully!")
+# Load a smaller, CPU-friendly model
+# Using 'Oppenheimer/DALL-E_1.5' - a fine-tuned SD 1.5 model known for good results
+model_id = "Oppenheimer/DALL-E_1.5"
+
+print("Loading pipeline... This may take a few minutes.")
+try:
+    # Use torch.float32 for CPU compatibility. DO NOT use float16.
+    pipe = StableDiffusionPipeline.from_pretrained(
+        model_id,
+        torch_dtype=torch.float32, # Critical for CPU
+        use_safetensors=True
+    )
+    
+    # Optional: Use a faster scheduler for quicker generation
+    pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
+    
+    # Move the pipeline to the CPU
+    pipe = pipe.to(device)
+    print("Model loaded successfully on CPU!")
+
+except Exception as e:
+    print(f"Error loading model: {e}")
+    # Provide a more helpful error message
+    raise e
 
 # Define the image generation function
 def generate_image(prompt):
@@ -27,31 +41,33 @@ def generate_image(prompt):
     # Add a consistent style to all prompts to get a children's book look
     enhanced_prompt = f"children's book illustration, watercolor style, cute, whimsical, {prompt}"
 
-    # Generate the image with some default parameters
-    # guidance_scale controls how closely the image follows the prompt
+    print(f"Generating image for prompt: {enhanced_prompt}")
+    
+    # Generate the image. Use a lower resolution for speed on CPU.
     image = pipe(
         prompt=enhanced_prompt,
-        guidance_scale=9.5,
-        num_inference_steps=25  # More steps can mean higher quality, but is slower
-    ).images[0]  # We get the first (and only) image from the result
+        width=512,   # Smaller image = faster generation
+        height=512,
+        guidance_scale=7.5,
+        num_inference_steps=20,  # Fewer steps = much faster
+        generator=torch.Generator(device=device) # Ensure generator is on CPU
+    ).images[0]
 
+    print("Image generated successfully!")
     return image
 
 # Create the Gradio Interface
-# We are creating a simple one-input (prompt) one-output (image) interface
 demo = gr.Interface(
-    fn=generate_image,  # The function to call
-    inputs=gr.Textbox(  # The input is a text box for the prompt
+    fn=generate_image,
+    inputs=gr.Textbox(
         label="Enter your scene description",
         lines=2,
         placeholder="A brave little mouse exploring a giant forest..."
     ),
-    outputs=gr.Image(label="Generated Illustration", type="pil"), # The output is an image
-    title="Children's Book Illustrator 🤖🎨",
-    description="Generate beautiful illustrations for your children's story. Enter a description of a scene."
+    outputs=gr.Image(label="Generated Illustration", type="pil"),
+    title="Children's Book Illustrator (CPU Edition) 🤖🎨",
+    description="This free version runs on CPU. It's slower but gets the job done! Enter a scene description."
 )
 
-# This is the key part for making the API work correctly with n8n.
-# We launch the Gradio app with a custom `api_name`.
-# Setting `api_name="generate"` creates an API endpoint at `/api/generate/`
-demo.launch(debug=True, server_name="0.0.0.0", api_name="generate")
+# Launch the app
+demo.launch(debug=True, server_name="0.0.0.0")