Update app.py

app.py

@@ -2,20 +2,24 @@ import gradio as gr
 from diffusers import StableDiffusionPipeline
 import torch
 
-# Load model from Hugging Face Hub
+# Automatically detect device
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+# Load Stable Diffusion model
 pipe = StableDiffusionPipeline.from_pretrained(
     "stabilityai/stable-diffusion-2",
-    torch_dtype=torch.float16
-).to("cuda")
+    torch_dtype=torch.float16 if device == "cuda" else torch.float32
+).to(device)
 
+# Define image generation function
 def generate_image(prompt):
     image = pipe(prompt).images[0]
     return image
 
-# Gradio UI
+# Gradio Interface
 gr.Interface(
     fn=generate_image,
-    inputs=gr.Textbox(label="Enter your prompt"),
+    inputs=gr.Textbox(label="Enter your prompt here"),
     outputs="image",
-    title="AI Image Generator"
+    title="Text to Image Generator with Stable Diffusion 2"
 ).launch()