From c65f81d089512d8cbe8dc436febab885b250fcbd Mon Sep 17 00:00:00 2001
From: kijai <40791699+kijai@users.noreply.github.com>
Date: Sun, 24 Aug 2025 20:50:11 +0300
Subject: [PATCH] Handle missing faces better in FantasyPortrait detection

---
 fantasyportrait/nodes.py      | 33 ++++++++++++++++++++++++++++-
 fantasyportrait/pd_fgc/pdf.py | 40 ++++++++++++++++++++---------------
 2 files changed, 55 insertions(+), 18 deletions(-)

diff --git a/fantasyportrait/nodes.py b/fantasyportrait/nodes.py
index 4fd74b6..23d0312 100644
--- a/fantasyportrait/nodes.py
+++ b/fantasyportrait/nodes.py
@@ -43,7 +43,38 @@ def get_emo_feature(frame_list, face_aligner, pd_fpg_motion, device):
     
 
     comfy_pbar = ProgressBar(3)
-    _, landmark_list, rect_list = det_landmarks(face_aligner, frame_list, comfy_pbar) 
+    _, landmark_list, rect_list = det_landmarks(face_aligner, frame_list, comfy_pbar)
+
+    
+    # Fill missing landmarks and rects with previous valid one
+    last_valid_landmark = None
+    last_valid_rect = None
+    for i in range(len(landmark_list)):
+        if landmark_list[i] is None:
+            landmark_list[i] = last_valid_landmark
+        else:
+            last_valid_landmark = landmark_list[i]
+        if rect_list[i] is None:
+            rect_list[i] = last_valid_rect
+        else:
+            last_valid_rect = rect_list[i]
+
+    # Forward fill for leading None values
+    if landmark_list[0] is None:
+        first_valid = next((l for l in landmark_list if l is not None), None)
+        for i in range(len(landmark_list)):
+            if landmark_list[i] is None:
+                landmark_list[i] = first_valid
+            else:
+                break
+    if rect_list[0] is None:
+        first_valid = next((r for r in rect_list if r is not None), None)
+        for i in range(len(rect_list)):
+            if rect_list[i] is None:
+                rect_list[i] = first_valid
+            else:
+                break
+
     emo_list = get_drive_expression_pd_fgc(pd_fpg_motion, frame_list, landmark_list, device)
     comfy_pbar.update(1)
 
diff --git a/fantasyportrait/pd_fgc/pdf.py b/fantasyportrait/pd_fgc/pdf.py
index 1468210..d92d074 100644
--- a/fantasyportrait/pd_fgc/pdf.py
+++ b/fantasyportrait/pd_fgc/pdf.py
@@ -156,7 +156,7 @@ def det_landmarks(face_aligner, frame_list, comfy_pbar):
     face_aligner.reset_track()
 
     with tqdm(total=len(frame_list)) as pbar:
-        for frame in frame_list:
+        for i, frame in enumerate(frame_list):
             faces = face_aligner.forward(frame)
             if len(faces) > 0:
                 face = sorted(
@@ -167,36 +167,42 @@ def det_landmarks(face_aligner, frame_list, comfy_pbar):
                 rect_list.append(face["face_rect"])
                 new_frame_list.append(frame)
             else:
-                log.warning(f"No face detected in the frame {frame}, skipping.")
+                log.warning(f"No face detected in the frame {i}, inserting empty frame.")
+                rect_list.append(None)  # Add placeholder
+                new_frame_list.append(None)  # Add placeholder
             pbar.set_description("DET stage1")
             pbar.update()
     comfy_pbar.update(1)
 
-    assert len(new_frame_list) > 0
     face_aligner.reset_track()
     save_frame_list = []
     save_landmark_list = []
     with tqdm(total=len(new_frame_list)) as pbar:
-        for frame, rect in zip(new_frame_list, rect_list):
-            faces = face_aligner.forward(frame, pre_rect=rect)
-            if len(faces) > 0:
-                face = sorted(
-                    faces,
-                    key=lambda x: (x["face_rect"][2] - x["face_rect"][0])
-                    * (x["face_rect"][3] - x["face_rect"][1]),
-                )[-1]
-                landmarks = face["pre_kpt_222"]
-                save_frame_list.append(frame)
-                save_landmark_list.append(landmarks)
+        for i, (frame, rect) in enumerate(zip(new_frame_list, rect_list)):
+            if frame is None or rect is None:
+                save_frame_list.append(None)
+                save_landmark_list.append(None)
+                log.warning(f"No face detected in the frame {i}, inserting empty landmark.")
             else:
-                log.warning(f"No face detected in the frame {frame}, skipping.")
+                faces = face_aligner.forward(frame, pre_rect=rect)
+                if len(faces) > 0:
+                    face = sorted(
+                        faces,
+                        key=lambda x: (x["face_rect"][2] - x["face_rect"][0])
+                        * (x["face_rect"][3] - x["face_rect"][1]),
+                    )[-1]
+                    landmarks = face["pre_kpt_222"]
+                    save_frame_list.append(frame)
+                    save_landmark_list.append(landmarks)
+                else:
+                    save_frame_list.append(None)
+                    save_landmark_list.append(None)
+                    log.warning(f"No face detected in the frame {i}, inserting empty landmark.")
             pbar.set_description("DET stage2")
             pbar.update()
 
     comfy_pbar.update(1)
 
-    assert len(save_frame_list) > 0
-    save_landmark_list = np.stack(save_landmark_list, axis=0)
     face_aligner.reset_track()
     return save_frame_list, save_landmark_list, rect_list