Update app.py

app.py

@@ -121,7 +121,15 @@ with gr.Blocks() as demo:
 
 if __name__ == "__main__":
     # Check if CUDA is available
-    device = "cuda" if torch.cuda.is_available() else "cpu"
-    pipe = StableDiffusionPipeline.from_single_file(PATH,torch_dtype=torch.float16).to(device)
+    if torch.cuda.is_available():
+        device = "cuda"
+        dtype = torch.float16  # Use float16 for GPU to save memory
+    else:
+        device = "cpu"
+        dtype = torch.float32  # CPU does not support float16, use float32 instead
+
+
+    
+    pipe = StableDiffusionPipeline.from_single_file(PATH,torch_dtype=dtype).to(device)
     demo.queue(concurrency_count=2,
             ).launch()