Don't put onnx face encoders on gpu when using hf space

adaface/adaface_wrapper.py

@@ -31,7 +31,7 @@ class AdaFaceWrapper(nn.Module):
                  main_unet_filepath=None, unet_types=None, extra_unet_dirpaths=None, unet_weights_in_ensemble=None,
                  enable_static_img_suffix_embs=None, unet_uses_attn_lora=False,
                  attn_lora_layer_names=['q', 'k', 'v', 'out'], normalize_cross_attn=False, q_lora_updates_query=False,
-                 device='cuda', is_training=False):
+                 device='cuda', is_training=False, is_on_hf_space=False):
         '''
         pipeline_name: "text2img", "text2imgxl", "img2img", "text2img3", "flux", or None. 
         If None, it's used only as a face encoder, and the unet and vae are
@@ -64,6 +64,7 @@ class AdaFaceWrapper(nn.Module):
         self.unet_weights_in_ensemble = unet_weights_in_ensemble
         self.device = device
         self.is_training = is_training
+        self.is_on_hf_space = is_on_hf_space
 
         if negative_prompt is None:
             self.negative_prompt = \
@@ -99,6 +100,7 @@ class AdaFaceWrapper(nn.Module):
                                                                   self.adaface_ckpt_paths,
                                                                   self.adaface_encoder_cfg_scales,
                                                                   self.enabled_encoders,
+                                                                  is_on_hf_space=self.is_on_hf_space,
                                                                   num_static_img_suffix_embs=4)
 
         self.id2ada_prompt_encoder.to(self.device)

adaface/face_id_to_ada_prompt.py

@@ -26,7 +26,7 @@ def create_id2ada_prompt_encoder(adaface_encoder_types, adaface_ckpt_paths=None,
         if adaface_encoder_type == 'arc2face':
             id2ada_prompt_encoder = \
                 Arc2Face_ID2AdaPrompt(adaface_ckpt_path=adaface_ckpt_path, 
-                                    *args, **kwargs)
+                                      *args, **kwargs)
         elif adaface_encoder_type == 'consistentID':
             id2ada_prompt_encoder = \
                 ConsistentID_ID2AdaPrompt(pipe=None,
@@ -64,6 +64,7 @@ class FaceID2AdaPrompt(nn.Module):
         # i.e., 6 for arc2face and 1 for consistentID.
         self.out_id_embs_cfg_scale          = kwargs.get('out_id_embs_cfg_scale', -1)
         self.is_training                    = kwargs.get('is_training', False)
+        self.is_on_hf_space                 = kwargs.get('is_on_hf_space', False)
         # extend_prompt2token_proj_attention_multiplier is an integer >= 1.
         # TODO: extend_prompt2token_proj_attention_multiplier should be a list of integers.        
         self.extend_prompt2token_proj_attention_multiplier = kwargs.get('extend_prompt2token_proj_attention_multiplier', 1)
@@ -655,7 +656,7 @@ class Arc2Face_ID2AdaPrompt(FaceID2AdaPrompt):
         if device == self.device:
             return
         
-        if str(device) == 'cpu':
+        if str(device) == 'cpu' or self.is_on_hf_space:
             self.face_app = FaceAnalysis(name='antelopev2', root='models/insightface', 
                                          providers=['CPUExecutionProvider'])
             self.face_app.prepare(ctx_id=0, det_size=(512, 512))
@@ -801,7 +802,7 @@ class ConsistentID_ID2AdaPrompt(FaceID2AdaPrompt):
         if device == self.device:
             return
 
-        if str(device) == 'cpu':
+        if str(device) == 'cpu' or self.is_on_hf_space:
             self.face_app = FaceAnalysis(name='buffalo_l', root='models/insightface', 
                                          providers=['CPUExecutionProvider'])
             self.face_app.prepare(ctx_id=0, det_size=(512, 512))

app.py

@@ -44,7 +44,7 @@ def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
         seed = random.randint(0, MAX_SEED)
     return seed
 
-def is_running_on_spaces():
+def is_running_on_hf_space():
     return os.getenv("SPACE_ID") is not None
 
 from huggingface_hub import snapshot_download
@@ -59,13 +59,15 @@ os.makedirs("/tmp/gradio", exist_ok=True)
 app = FaceAnalysis(name="buffalo_l", root='models/insightface', providers=['CPUExecutionProvider'])
 app.prepare(ctx_id=0, det_size=(320, 320))
 
-if is_running_on_spaces():
+if is_running_on_hf_space():
     device = 'cuda:0'
+    is_on_hf_space = True
 else:
     if args.gpu is None:
         device = "cuda" 
     else:
         device = f"cuda:{args.gpu}"
+    is_on_hf_space = False
 
 print(f"Device: {device}")
 
@@ -76,7 +78,8 @@ id_animator = load_model(model_style_type=args.model_style_type, device='cpu')
 adaface = AdaFaceWrapper(pipeline_name="text2img", base_model_path=adaface_base_model_path,
                          adaface_encoder_types=args.adaface_encoder_types,
                          adaface_ckpt_paths=args.adaface_ckpt_path, device='cpu',
-                         num_inference_steps=args.num_inference_steps)
+                         num_inference_steps=args.num_inference_steps,
+                         is_on_hf_space=is_on_hf_space)
 
 basedir = os.getcwd()
 savedir = os.path.join(basedir,'samples')

requirements.txt

@@ -11,7 +11,7 @@ insightface
 omegaconf
 opencv-python
 onnx>=1.16.0
-onnxruntime-gpu
+onnxruntime
 safetensors
 spaces
 ftfy