Add files via upload

latent_slice.py

@@ -8,7 +8,8 @@ class CleanLatentSlice:
         return {
             "required": {
                 "latent": ("LATENT",),
-                "length": ("INT", {"default": 1, "min": 1, "max": 100000, "step": 1, "tooltip": "The target length in latent frames."}),
+                "start": ("INT", {"default": 0, "min": 0, "max": 100000, "step": 1, "tooltip": "The starting frame index to slice from (plug 'latent_start_index' here)."}),
+                "length": ("INT", {"default": 1, "min": 1, "max": 100000, "step": 1, "tooltip": "The number of frames to keep (plug 'clean_latent_frames' here)."}),
             }
         }
 
@@ -16,37 +17,59 @@ class CleanLatentSlice:
     RETURN_NAMES = ("latent",)
     FUNCTION = "slice_latent"
     CATEGORY = "WhatDreamsCost"
-    DESCRIPTION = "Safely slices a video latent to a specific length. Uses torch.narrow to bypass PyTorch NestedTensor slicing bugs."
+    DESCRIPTION = "Safely slices a video latent starting from an offset index for a specific length. Uses torch.narrow to bypass PyTorch NestedTensor slicing bugs."
 
-    def slice_latent(self, latent, length):
+    def slice_latent(self, latent, start, length):
         new_latent = latent.copy()
         
-        def safe_slice(tensor, target_len):
+        def safe_slice(tensor, target_start, target_len):
             dims = tensor.ndim if hasattr(tensor, "ndim") else len(tensor.shape)
             
+            # Constrain starting index and length to the actual size of the tensor
+            max_size = tensor.size(2) if dims == 5 else tensor.size(0)
+            actual_start = min(target_start, max_size - 1) if max_size > 0 else 0
+            actual_len = min(target_len, max_size - actual_start)
+            
             try:
-                # Bypasses NestedTensor slicing bugs
+                # torch.narrow is the safest low-level C++ slice method to bypass NestedTensor bugs
                 if dims == 5:
-                    return torch.narrow(tensor, 2, 0, target_len)
+                    # [Batch, Channels, Frames, Height, Width] -> Slice dimension 2
+                    return torch.narrow(tensor, 2, actual_start, actual_len)
                 elif dims == 4:
-                    return torch.narrow(tensor, 0, 0, target_len)
+                    # [Frames, Channels, Height, Width] -> Slice dimension 0
+                    return torch.narrow(tensor, 0, actual_start, actual_len)
                 elif dims == 3:
-                    return torch.narrow(tensor, 0, 0, target_len)
+                    # [Frames, Height, Width] -> Slice dimension 0
+                    return torch.narrow(tensor, 0, actual_start, actual_len)
             except Exception as e:
-                # Extreme fallback if narrow fails
+                # Fallback if narrow fails
                 if dims == 5:
-                    return tensor[:, :, :target_len]
+                    return tensor[:, :, actual_start : actual_start + actual_len]
                 elif dims == 4:
-                    return tensor[:target_len]
+                    return tensor[actual_start : actual_start + actual_len]
                 elif dims == 3:
-                    return tensor[:target_len]
+                    return tensor[actual_start : actual_start + actual_len]
                     
             return tensor
 
+        # Safely slice video samples
         if "samples" in new_latent:
-            new_latent["samples"] = safe_slice(new_latent["samples"], length)
+            new_latent["samples"] = safe_slice(new_latent["samples"], start, length)
 
+        # Safely slice video noise mask (if it exists)
         if "noise_mask" in new_latent:
-            new_latent["noise_mask"] = safe_slice(new_latent["noise_mask"], length)
+            new_latent["noise_mask"] = safe_slice(new_latent["noise_mask"], start, length)
 
-        return (new_latent,)
+        return (new_latent,)
+
+
+# Register the node with ComfyUI
+NODE_CLASS_MAPPINGS = {
+    "CleanLatentSlice": CleanLatentSlice
+}
+
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "CleanLatentSlice": "Clean Latent Slice"
+}
+
+__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS']