From 81b7d7145402e46aa42efd9a3450a23089d04b7f Mon Sep 17 00:00:00 2001
From: kijai <40791699+kijai@users.noreply.github.com>
Date: Sat, 1 Feb 2025 22:15:13 +0200
Subject: [PATCH] Update nodes.py

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

diff --git a/nodes/nodes.py b/nodes/nodes.py
index 8d4fdcc..c66a6d9 100644
--- a/nodes/nodes.py
+++ b/nodes/nodes.py
@@ -2232,6 +2232,7 @@ class LeapfusionHunyuanI2V:
             "required": {
                 "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"}),
             }
         }
 
@@ -2240,18 +2241,18 @@ class LeapfusionHunyuanI2V:
 
     CATEGORY = "KJNodes/experimental"
 
-    def patch(self, model, latent):
+    def patch(self, model, latent, index):
 
-        def outer_wrapper(samples):
+        def outer_wrapper(samples, index):
             def unet_wrapper(apply_model, args):
                 inp, timestep, c = args["input"], args["timestep"], args["c"]
                 if samples is not None:
-                    inp[:, :, [0], :, :] = samples[:, :, [0], :, :].to(inp)
+                    inp[:, :, [index], :, :] = samples[:, :, [0], :, :].to(inp)
                 return apply_model(inp, timestep, **c)
             return unet_wrapper
         
         samples = latent["samples"] * 0.476986
         m = model.clone()
-        m.set_model_unet_function_wrapper(outer_wrapper(samples))
+        m.set_model_unet_function_wrapper(outer_wrapper(samples, index))
 
         return (m,)
\ No newline at end of file