xinyun/ComfyUI
1
0
Fork 0
mirror of https://git.datalinker.icu/comfyanonymous/ComfyUI synced 2026-08-15 21:36:42 +08:00
Code Issues Packages Projects Releases Wiki Activity

Made dim slicing cleaner

Browse Source
This commit is contained in:
Jedrzej Kosinski 2025-08-06 14:15:29 -07:00
parent af74959150
commit 8a3b8f77f3
1 changed files with 6 additions and 18 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

24
comfy/context_windows.py
View File

@ -56,28 +56,16 @@ class IndexListContextWindow(ContextWindowABC):
def get_tensor(self, full: torch.Tensor, device=None, dim=None) -> torch.Tensor:
if dim is None:
dim = self.dim
if dim == 0:
if full.shape[dim] == 1:
return full
return full[self.index_list].to(device)
elif dim == 1:
return full[:, self.index_list].to(device)
elif dim == 2:
return full[:, :, self.index_list].to(device)
else:
raise ValueError(f"Invalid dimension: {dim}")
if dim == 0 and full.shape[dim] == 1:
return full
idx = [slice(None)] * dim + [self.index_list]
return full[idx].to(device)
def add_window(self, full: torch.Tensor, to_window: torch.Tensor, dim=None) -> torch.Tensor:
if dim is None:
dim = self.dim
if dim == 0:
full[self.index_list] += to_window
elif dim == 1:
full[:, self.index_list] += to_window
elif dim == 2:
full[:, :, self.index_list] += to_window
else:
raise ValueError(f"Invalid dimension: {dim}")
idx = [slice(None)] * dim + [self.index_list]
full[idx] += to_window
return full