Update app.py

app.py

@@ -2,24 +2,24 @@ import gradio as gr
 from diffusers import StableDiffusionPipeline
 import torch
 
-# Load the Stable Diffusion model
-model_id = "runwayml/stable-diffusion-v1-5"
-pipeline = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
-pipeline.to("cuda" if torch.cuda.is_available() else "cpu")
+# Load a smaller model if running on CPU
+model_id = "CompVis/stable-diffusion-v1-4"
+pipeline = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float32)
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+pipeline.to(device)
 
 def generate_image(prompt):
     image = pipeline(prompt).images[0]
     return image
 
-# Define the Gradio interface
 interface = gr.Interface(
     fn=generate_image,
     inputs=gr.Textbox(label="Enter your text prompt"),
     outputs=gr.Image(type="pil"),
     title="Text-to-Image Generator",
-    description="Enter a prompt, and the AI will generate an image based on it.",
+    description="Enter a prompt, and the AI will generate an image based on it."
 )
 
-# Launch the app
 if __name__ == "__main__":
     interface.launch()