Update app.py

app.py

@@ -32,7 +32,14 @@ JS = """function () {
   }
 }"""
 
-pipe = FluxPipeline.from_pretrained(model, torch_dtype=torch.bfloat16)
+if torch.cuda.is_available():
+    device = "cuda"
+    print("Using GPU")
+else:
+    device = "cpu"
+    print("Using CPU")
+
+pipe = FluxPipeline.from_pretrained(model, torch_dtype=torch.bfloat16).to(device)
 pipe.load_lora_weights(default_lora, weight_name = default_weight_name) # default load lora
 pipe.fuse_lora(lora_scale=0.9)