From 15ff49a4d77f3548ba8443ddfa21c0a9c37eeba8 Mon Sep 17 00:00:00 2001 From: Your Name Date: Fri, 23 Jan 2026 00:42:30 +0700 Subject: [PATCH] Refine: update defaults and add optional reference_image input --- README.md | 5 +++-- i2v_chain_helper.py | 21 +++++++++++++++------ 2 files changed, 18 insertions(+), 8 deletions(-) diff --git a/README.md b/README.md index aa9273c..7b61120 100644 --- a/README.md +++ b/README.md @@ -25,8 +25,9 @@ pip install -r requirements.txt - **Inputs**: - `images`: Image batch. - `analysis_models`: Face analysis models (InsightFace or DLib). - - `min_face_similarity`: Minimum cosine similarity to the reference face (default: 0.6). - - `min_eyes_openness`: Minimum EAR score for open eyes (default: 0.2). + - `min_face_similarity`: Minimum cosine similarity to the reference face (default: 0.5). + - `min_eyes_openness`: Minimum EAR score for open eyes (default: 0.4). + - `reference_image` (optional): If provided, use this image as the identity reference instead of the first frame of the batch. - **Outputs**: - `trimmed_images`: The filtered and potentially shortened batch. - `frame_count`: Number of images in the trimmed batch. diff --git a/i2v_chain_helper.py b/i2v_chain_helper.py index ab12ec5..da1ed89 100644 --- a/i2v_chain_helper.py +++ b/i2v_chain_helper.py @@ -38,8 +38,11 @@ class I2VChainHelper: "required": { "images": ("IMAGE",), "analysis_models": ("ANALYSIS_MODELS",), - "min_face_similarity": ("FLOAT", {"default": 0.6, "min": 0.0, "max": 1.0, "step": 0.01}), - "min_eyes_openness": ("FLOAT", {"default": 0.2, "min": 0.0, "max": 1.0, "step": 0.01}), + "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",), } } @@ -48,16 +51,22 @@ class I2VChainHelper: FUNCTION = "execute" CATEGORY = "I2VChain" - def execute(self, images, analysis_models, min_face_similarity, min_eyes_openness): + 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 from the first frame - ref_img = tensor_to_pil(images[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 the first frame. Returning empty batch.") + 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)