Update app.py

app.py

@@ -2,10 +2,11 @@ import gradio as gr
 from diffusers import StableDiffusionPipeline
 import torch
 
-# Load the fine-tuned Stable Diffusion model from your Hugging Face repository
+# Fine-tuned Stable Diffusion model
 model_id = "sk2003/room-styler"
 pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
-# pipe.to("cuda")  # Uncomment if using GPU
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+pipe.to(device)
 
 def generate_image(prompt):
     image = pipe(prompt).images[0]