Refactor get devices

app.py

@@ -90,9 +90,14 @@ description = \
 - You can optionally generate a high-resolution sample if the input image is of high resolution. We split the original image into `Vertical Splits` by `Horizontal Splits` patches with some `Overlaps` in between. Due to computation constraints for the online demo, we recommend `Vertical Splits` x `Horizontal Splits` to be no more than 6 and to set 2 for `Overlaps`. The denoising steps should at least be set to 80 for high resolution samples.
 
 """
+@spaces.GPU()
+def get_devices():
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    return device
 
 set_loggers("INFO")
-device = "cuda" if torch.cuda.is_available() else "cpu"
+device = get_devices()
+set_loggers(f"Using devices: {device}")
 
 # Download from model
 logger.info(f"Downloading Models...")
@@ -102,10 +107,10 @@ logger.info(f"Loading Models...")
 model_dict = {
     "Albedo": InferenceModel(ckpt_path="weights/albedo", 
                 use_ddim=True, 
-                gpu_id=0),
+                gpu_id=device),
     "Specular": InferenceModel(ckpt_path="weights/specular", 
                 use_ddim=True, 
-                gpu_id=0),
+                gpu_id=device),
     "remove_bg": rembg.new_session(),
 }
 logger.info(f"All models Loaded!")