Revert move_to_gpu, let face_app stays on cpu

ConsistentID/lib/pipeline_ConsistentID.py

@@ -72,15 +72,11 @@ class ConsistentIDPipeline(StableDiffusionPipeline):
         self.id_image_processor = CLIPImageProcessor()
         self.crop_size = 512
 
-        '''
         # face_app: FaceAnalysis object
         self.face_app = FaceAnalysis(name="buffalo_l", root='models/insightface', 
                                      providers=['CPUExecutionProvider'])
         # The original det_size=(640, 640) is too large and face_app often fails to detect faces.
-        self.face_app.prepare(ctx_id=0, det_size=(256, 256))
-        '''
-        # Will be initialized within adaface.
-        self.face_app = None
+        self.face_app.prepare(ctx_id=0, det_size=(512, 512))
 
         if not os.path.exists(consistentID_weight_path):
             ### Download pretrained models

adaface/face_id_to_ada_prompt.py

@@ -606,7 +606,7 @@ class Arc2Face_ID2AdaPrompt(FaceID2AdaPrompt):
         # Note there are two "models" in the path.        
         self.face_app = FaceAnalysis(name='antelopev2', root='models/insightface', 
                                             providers=['CPUExecutionProvider'])
-        self.face_app.prepare(ctx_id=0, det_size=(256, 256))
+        self.face_app.prepare(ctx_id=0, det_size=(512, 512))
         print(f'Arc2Face Face encoder loaded on CPU.')
 
         self.text_to_image_prompt_encoder = CLIPTextModelWrapper.from_pretrained(
@@ -642,6 +642,7 @@ class Arc2Face_ID2AdaPrompt(FaceID2AdaPrompt):
 
     def _apply(self, fn):
         super()._apply(fn)  # Call the parent _apply to handle parameters and buffers
+        return
         # A dirty hack to get the device of the model, passed from 
         # parent.model.to(self.root_device) => parent._apply(convert) => module._apply(fn)
         test_tensor = torch.zeros(1)  # Create a test tensor
@@ -654,7 +655,7 @@ class Arc2Face_ID2AdaPrompt(FaceID2AdaPrompt):
         if str(device) == 'cpu':
             self.face_app = FaceAnalysis(name='antelopev2', root='models/insightface', 
                                         providers=['CPUExecutionProvider'])
-            self.face_app.prepare(ctx_id=0, det_size=(256, 256))
+            self.face_app.prepare(ctx_id=0, det_size=(512, 512))
         else:
             device_id = device.index
             self.face_app = FaceAnalysis(name='antelopev2', root='models/insightface', 
@@ -663,7 +664,7 @@ class Arc2Face_ID2AdaPrompt(FaceID2AdaPrompt):
                                                             "cudnn_conv_algo_search": "HEURISTIC",
                                                             "gpu_mem_limit": 2 * 1024**3
                                                             }])
-            self.face_app.prepare(ctx_id=device_id, det_size=(256, 256))
+            self.face_app.prepare(ctx_id=device_id, det_size=(512, 512))
 
         self.device = device
         print(f'Arc2Face Face encoder reloaded on {device}.')
@@ -790,6 +791,7 @@ class ConsistentID_ID2AdaPrompt(FaceID2AdaPrompt):
 
     def _apply(self, fn):
         super()._apply(fn)  # Call the parent _apply to handle parameters and buffers
+        return
         # A dirty hack to get the device of the model, passed from 
         # parent.model.to(self.root_device) => parent._apply(convert) => module._apply(fn)
         test_tensor = torch.zeros(1)  # Create a test tensor
@@ -802,7 +804,7 @@ class ConsistentID_ID2AdaPrompt(FaceID2AdaPrompt):
         if str(device) == 'cpu':
             self.face_app = FaceAnalysis(name='buffalo_l', root='models/insightface', 
                                          providers=['CPUExecutionProvider'])
-            self.face_app.prepare(ctx_id=0, det_size=(256, 256))
+            self.face_app.prepare(ctx_id=0, det_size=(512, 512))
         else:
             device_id = device.index
             self.face_app = FaceAnalysis(name='buffalo_l', root='models/insightface', 
@@ -811,7 +813,7 @@ class ConsistentID_ID2AdaPrompt(FaceID2AdaPrompt):
                                                             "cudnn_conv_algo_search": "HEURISTIC",
                                                             "gpu_mem_limit": 2 * 1024**3
                                                             }])
-            self.face_app.prepare(ctx_id=device_id, det_size=(256, 256))
+            self.face_app.prepare(ctx_id=device_id, det_size=(512, 512))
 
         self.device = device
         self.pipe.face_app = self.face_app

app.py

@@ -23,20 +23,6 @@ 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"],
@@ -101,7 +87,6 @@ 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:
@@ -127,8 +112,21 @@ def generate_image(image_paths, image_paths2, guidance_scale, perturb_std,
                    ablate_prompt_embed_type, nonmix_prompt_emb_weight,
                    composition_level, seed, disable_adaface, subj_name_sig, progress=gr.Progress(track_tqdm=True)):
 
-    global adaface
+    global adaface, args
+    
+    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)
+    args.device = 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.")
     
@@ -149,10 +147,10 @@ def generate_image(image_paths, image_paths2, guidance_scale, perturb_std,
 
     # Sometimes the pipeline is on CPU, although we've put it on CUDA (due to some offloading mechanism).
     # Therefore we set the generator to the correct device.
-    generator = torch.Generator(device=device).manual_seed(seed)
+    generator = torch.Generator(device=args.device).manual_seed(seed)
     print(f"Manual seed: {seed}.")
     # Generate two images each time for the user to select from.
-    noise = torch.randn(num_images, 3, 512, 512, device=device, generator=generator)
+    noise = torch.randn(num_images, 3, 512, 512, device=args.device, generator=generator)
     #print(noise.abs().sum())
     # samples: A list of PIL Image instances.
     if highlight_face: