Improve device assignment

app.py

@@ -20,6 +20,9 @@ def str2bool(v):
     else:
         raise argparse.ArgumentTypeError("Boolean value expected.")
 
+def is_running_on_spaces():
+    return os.getenv("SPACE_ID") is not None
+
 import argparse
 parser = argparse.ArgumentParser()
 parser.add_argument("--adaface_encoder_types", type=str, nargs="+", default=["consistentID", "arc2face"],
@@ -68,8 +71,6 @@ base_model_path = model_style_type2base_model_path[args.model_style_type]
 
 # global variable
 MAX_SEED = np.iinfo(np.int32).max
-device = "cuda" if args.gpu is None else f"cuda:{args.gpu}"
-print(f"Device: {device}")
 
 global adaface
 adaface = None
@@ -113,6 +114,16 @@ def generate_image(image_paths, image_paths2, guidance_scale, perturb_std,
 
     global adaface
 
+    if is_running_on_spaces():
+        device = 'cuda:0'
+    else:
+        if args.gpu is None:
+            device = "cuda" 
+        else:
+            device = f"cuda:{args.gpu}"
+
+    print(f"Device: {device}")
+
     adaface.to(device)
 
     if image_paths is None or len(image_paths) == 0: