Add files via upload

latent_slice.py

@@ -0,0 +1,52 @@
+# --- START OF FILE latent_slice.py ---
+
+import torch
+
+class CleanLatentSlice:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "latent": ("LATENT",),
+                "length": ("INT", {"default": 1, "min": 1, "max": 100000, "step": 1, "tooltip": "The target length in latent frames."}),
+            }
+        }
+
+    RETURN_TYPES = ("LATENT",)
+    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."
+
+    def slice_latent(self, latent, length):
+        new_latent = latent.copy()
+        
+        def safe_slice(tensor, target_len):
+            dims = tensor.ndim if hasattr(tensor, "ndim") else len(tensor.shape)
+            
+            try:
+                # Bypasses NestedTensor slicing bugs
+                if dims == 5:
+                    return torch.narrow(tensor, 2, 0, target_len)
+                elif dims == 4:
+                    return torch.narrow(tensor, 0, 0, target_len)
+                elif dims == 3:
+                    return torch.narrow(tensor, 0, 0, target_len)
+            except Exception as e:
+                # Extreme fallback if narrow fails
+                if dims == 5:
+                    return tensor[:, :, :target_len]
+                elif dims == 4:
+                    return tensor[:target_len]
+                elif dims == 3:
+                    return tensor[:target_len]
+                    
+            return tensor
+
+        if "samples" in new_latent:
+            new_latent["samples"] = safe_slice(new_latent["samples"], length)
+
+        if "noise_mask" in new_latent:
+            new_latent["noise_mask"] = safe_slice(new_latent["noise_mask"], length)
+
+        return (new_latent,)