52 lines
1.9 KiB
Python
52 lines
1.9 KiB
Python
# --- 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,) |