allocate 120s for move_to_gpu

adaface/face_id_to_ada_prompt.py

@@ -661,7 +661,8 @@ class Arc2Face_ID2AdaPrompt(FaceID2AdaPrompt):
                                          providers=['CUDAExecutionProvider'],
                                          provider_options=[{"device_id": device_id, 
                                                             "cudnn_conv_algo_search": "HEURISTIC",
-                                                            "gpu_mem_limit": 1 * 1024**3}])
+                                                            "gpu_mem_limit": 2 * 1024**3
+                                                            }])
             self.face_app.prepare(ctx_id=device_id, det_size=(256, 256))
 
         self.device = device
@@ -808,7 +809,8 @@ class ConsistentID_ID2AdaPrompt(FaceID2AdaPrompt):
                                          providers=['CUDAExecutionProvider'],
                                          provider_options=[{"device_id": device_id, 
                                                             "cudnn_conv_algo_search": "HEURISTIC",
-                                                            "gpu_mem_limit": 1 * 1024**3}])
+                                                            "gpu_mem_limit": 2 * 1024**3
+                                                            }])
             self.face_app.prepare(ctx_id=device_id, det_size=(256, 256))
 
         self.device = device

app.py

@@ -23,6 +23,20 @@ def str2bool(v):
 def is_running_on_spaces():
     return os.getenv("SPACE_ID") is not None
 
+@spaces.GPU(duration=120)
+def move_to_gpu(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)
+    
 import argparse
 parser = argparse.ArgumentParser()
 parser.add_argument("--adaface_encoder_types", type=str, nargs="+", default=["consistentID", "arc2face"],
@@ -87,6 +101,7 @@ if not args.test_ui_only:
                              shrink_cross_attn=False,
                              q_lora_updates_query=args.q_lora_updates_query,
                              device='cpu')
+    move_to_gpu(adaface)
 
 def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
     if randomize_seed:
@@ -114,18 +129,6 @@ 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:
         raise gr.Error(f"Cannot find any input face image! Please upload a face image.")