diff --git a/InstantID.py b/InstantID.py
index 3810f0f..832d2fc 100644
--- a/InstantID.py
+++ b/InstantID.py
@@ -428,7 +428,7 @@ class ApplyInstantID:
         if face_embed is None:
             raise Exception('Reference Image: No face detected.')
         
-        face_kps = extractFeatures(insightface, image_kps if image_kps is not None else image, extract_kps=True)
+        face_kps = extractFeatures(insightface, image_kps[0].unsqueeze(0) if image_kps is not None else image[0].unsqueeze(0), extract_kps=True)
 
         if face_kps is None:
             face_kps = torch.zeros_like(image) if image_kps is None else image_kps