i2v_chain_helper.py

import torch
import numpy as np
from PIL import Image
import torchvision.transforms.v2 as T
from comfy.utils import ProgressBar

def tensor_to_pil(img):
    return T.ToPILImage()(img.permute(2, 0, 1)).convert('RGB')

def calculate_ear(eye_landmarks):
    """
    Calculate Eye Aspect Ratio (EAR)
    For DLib (6 points): 0-3 horizontal, 1-5 and 2-4 vertical
    For InsightFace (10 points): we'll take average vertical / horizontal
    """
    if len(eye_landmarks) == 6:
        # DLib style
        p = eye_landmarks
        v1 = np.linalg.norm(p[1] - p[5])
        v2 = np.linalg.norm(p[2] - p[4])
        h = np.linalg.norm(p[0] - p[3])
        return (v1 + v2) / (2.0 * h)
    elif len(eye_landmarks) == 10:
        # InsightFace style (usually perimeter)
        # Assuming order is somewhat standard, we'll use bounding box ratio as fallback or pick points
        # For simplicity, let's use the bounding box of the eye points
        min_p = np.min(eye_landmarks, axis=0)
        max_p = np.max(eye_landmarks, axis=0)
        width = max_p[0] - min_p[0]
        height = max_p[1] - min_p[1]
        return height / width if width > 0 else 0
    return 0

class I2VChainHelper:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "images": ("IMAGE",),
                "analysis_models": ("ANALYSIS_MODELS",),
                "min_face_similarity": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01}),
                "min_eyes_openness": ("FLOAT", {"default": 0.4, "min": 0.0, "max": 1.0, "step": 0.01}),
            },
            "optional": {
                "reference_image": ("IMAGE",),
            }
        }

    RETURN_TYPES = ("IMAGE", "INT", "IMAGE", "IMAGE", "FLOAT", "FLOAT")
    RETURN_NAMES = ("trimmed_images", "frame_count", "first_frame", "last_frame", "face_similarity", "eyes_openness")
    FUNCTION = "execute"
    CATEGORY = "I2VChain"

    def execute(self, images, analysis_models, min_face_similarity, min_eyes_openness, reference_image=None):
        if images.shape[0] == 0:
            return (images, 0, images, images, 0.0, 0.0)

        # 1. Get reference embedding
        if reference_image is not None and reference_image.shape[0] > 0:
            ref_img = tensor_to_pil(reference_image[0])
            print("I2VChainHelper: Using provided reference_image for identity.")
        else:
            ref_img = tensor_to_pil(images[0])
            print("I2VChainHelper: No reference_image provided. Using first frame as identity reference.")
            
        ref_embed = analysis_models.get_embeds(np.array(ref_img))
        
        if ref_embed is None:
            print("I2VChainHelper: No face detected in reference. Returning empty batch.")
            return (images[:0], 0, images[:0], images[:0], 0.0, 0.0)

        ref_embed = ref_embed / np.linalg.norm(ref_embed)

        last_good_index = 0
        final_similarity = 0.0
        final_openness = 0.0
        pbar = ProgressBar(images.shape[0])
        
        # Scan from last to first to save time as requested
        for i in range(images.shape[0] - 1, -1, -1):
            img_pil = tensor_to_pil(images[i])
            img_np = np.array(img_pil)
            
            # Check similarity
            curr_embed = analysis_models.get_embeds(img_np)
            if curr_embed is None:
                pbar.update(1)
                continue
            
            curr_embed = curr_embed / np.linalg.norm(curr_embed)
            similarity = np.dot(ref_embed, curr_embed)
            
            if similarity < min_face_similarity:
                pbar.update(1)
                continue
            
            # Check eyes
            landmarks = analysis_models.get_landmarks(img_np)
            if landmarks is None:
                pbar.update(1)
                continue
                
            left_eye = landmarks[3]
            right_eye = landmarks[4]
            
            ear_l = calculate_ear(left_eye)
            ear_r = calculate_ear(right_eye)
            avg_ear = (ear_l + ear_r) / 2.0
            
            if avg_ear < min_eyes_openness:
                pbar.update(1)
                continue
                
            # Found the last good frame!
            last_good_index = i
            final_similarity = float(similarity)
            final_openness = float(avg_ear)
            pbar.update(images.shape[0]) # Mark as complete for the UI
            break
            
        trimmed_images = images[:last_good_index + 1]
        first_frame = trimmed_images[0:1]
        last_frame = trimmed_images[-1:]
        
        return (trimmed_images, trimmed_images.shape[0], first_frame, last_frame, final_similarity, final_openness)

NODE_CLASS_MAPPINGS = {
    "I2VChainHelper": I2VChainHelper
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "I2VChainHelper": "I2V Chain Helper"
}