From 0c3346dbf2d17e1fd773c459de88a835edf99a4c Mon Sep 17 00:00:00 2001
From: Denarius40 <112815664+Denarius40@users.noreply.github.com>
Date: Wed, 17 Jun 2026 23:30:00 +0200
Subject: [PATCH] fix: move mask to correct device in LTXVDilateVideoMask

Max-pooling ops fail when the input tensor is on CPU but ComfyUI is
running on CUDA. Move the mask to the active torch device before the
pooling passes and return it to CPU afterwards so downstream nodes
receive a CPU tensor as expected by ComfyUI conventions.
---
 vanish_nodes.py | 11 ++++++++---
 1 file changed, 8 insertions(+), 3 deletions(-)

diff --git a/vanish_nodes.py b/vanish_nodes.py
index 2baa6ee..7de6a9e 100644
--- a/vanish_nodes.py
+++ b/vanish_nodes.py
@@ -59,19 +59,23 @@ def execute(
     ) -> io.NodeOutput:
         if mask is None and image_as_mask is None:
             raise ValueError("Either 'mask' or 'image_as_mask' must be provided.")
-
+        
         if mask is None:
             mask = image_as_mask.mean(dim=-1)
-
+            
         if mask.ndim == 4:
             mask = mask[:, :, :, 0]
 
+        from comfy import model_management
+        device = model_management.get_torch_device()
+        mask = mask.to(device)
+
         s_kernel = spatial_radius * 2 + 1
         t_kernel = temporal_radius * 2 + 1
 
         # Separable dilation: 2D spatial + 1D temporal (much faster than 3D pooling)
         if s_kernel > 1:
-            mask = mask.unsqueeze(1)  # (B, 1, H, W)
+            mask = mask.unsqueeze(1) # (B, 1, H, W)
             mask = F.max_pool2d(
                 mask, kernel_size=s_kernel, stride=1, padding=spatial_radius
             )
@@ -86,6 +90,7 @@ def execute(
             mask = mask.reshape(H, W, B).permute(2, 0, 1)
 
         mask = (mask > 0.5).float()
+        mask = mask.cpu()
         return io.NodeOutput(mask)