From e9a3170c72864cfc1e1c7562c082669229b5e3a4 Mon Sep 17 00:00:00 2001
From: kijai <40791699+kijai@users.noreply.github.com>
Date: Sun, 2 Feb 2025 00:16:57 +0200
Subject: [PATCH] Update nodes.py

---
 nodes/nodes.py | 13 +++++++++----
 1 file changed, 9 insertions(+), 4 deletions(-)

diff --git a/nodes/nodes.py b/nodes/nodes.py
index d6cf23d..730106b 100644
--- a/nodes/nodes.py
+++ b/nodes/nodes.py
@@ -2233,6 +2233,8 @@ class LeapfusionHunyuanI2V:
                 "model": ("MODEL",),
                 "latent": ("LATENT",),
                 "index": ("INT", {"default": 0, "min": -1, "max": 1000, "step": 1,"tooltip": "The index of the latent to be replaced. 0 for first frame and -1 for last"}),
+                "start_percent": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.01, "tooltip": "The start percentage of steps to apply"}),
+                "end_percent": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01, "tooltip": "The end percentage of steps to apply"}),
                 "strength": ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.001}),
             }
         }
@@ -2242,19 +2244,22 @@ class LeapfusionHunyuanI2V:
 
     CATEGORY = "KJNodes/experimental"
 
-    def patch(self, model, latent, index, strength):
+    def patch(self, model, latent, index, strength, start_percent, end_percent):
 
-        def outer_wrapper(samples, index):
+        def outer_wrapper(samples, index, start_percent, end_percent):
             def unet_wrapper(apply_model, args):
+                steps = args["c"]["transformer_options"]["sample_sigmas"]
                 inp, timestep, c = args["input"], args["timestep"], args["c"]
-                if samples is not None:
+                current_step_index = (steps == timestep).nonzero().item()
+                current_percent = current_step_index / (len(steps) - 1)
+                if samples is not None and start_percent <= current_percent <= end_percent:
                     inp[:, :, [index], :, :] = samples[:, :, [0], :, :].to(inp)
                 return apply_model(inp, timestep, **c)
             return unet_wrapper
         
         samples = latent["samples"] * 0.476986 * strength
         m = model.clone()
-        m.set_model_unet_function_wrapper(outer_wrapper(samples, index))
+        m.set_model_unet_function_wrapper(outer_wrapper(samples, index, start_percent, end_percent))
 
         return (m,)