From 8a3b8f77f38ae7d0da2c6493381b5e75cfd9dcc4 Mon Sep 17 00:00:00 2001 From: Jedrzej Kosinski Date: Wed, 6 Aug 2025 14:15:29 -0700 Subject: [PATCH] Made dim slicing cleaner --- comfy/context_windows.py | 24 ++++++------------------ 1 file changed, 6 insertions(+), 18 deletions(-) diff --git a/comfy/context_windows.py b/comfy/context_windows.py index a546fae12..fffd75fc9 100644 --- a/comfy/context_windows.py +++ b/comfy/context_windows.py @@ -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