75 lines
3.2 KiB
Python
75 lines
3.2 KiB
Python
# --- START OF FILE latent_slice.py ---
|
|
|
|
import torch
|
|
|
|
class CleanLatentSlice:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"latent": ("LATENT",),
|
|
"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)."}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("LATENT",)
|
|
RETURN_NAMES = ("latent",)
|
|
FUNCTION = "slice_latent"
|
|
CATEGORY = "WhatDreamsCost"
|
|
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, start, length):
|
|
new_latent = latent.copy()
|
|
|
|
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:
|
|
# torch.narrow is the safest low-level C++ slice method to bypass NestedTensor bugs
|
|
if dims == 5:
|
|
# [Batch, Channels, Frames, Height, Width] -> Slice dimension 2
|
|
return torch.narrow(tensor, 2, actual_start, actual_len)
|
|
elif dims == 4:
|
|
# [Frames, Channels, Height, Width] -> Slice dimension 0
|
|
return torch.narrow(tensor, 0, actual_start, actual_len)
|
|
elif dims == 3:
|
|
# [Frames, Height, Width] -> Slice dimension 0
|
|
return torch.narrow(tensor, 0, actual_start, actual_len)
|
|
except Exception as e:
|
|
# Fallback if narrow fails
|
|
if dims == 5:
|
|
return tensor[:, :, actual_start : actual_start + actual_len]
|
|
elif dims == 4:
|
|
return tensor[actual_start : actual_start + actual_len]
|
|
elif dims == 3:
|
|
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"], 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"], start, length)
|
|
|
|
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'] |