fixed memory estimation

nodes/load_model.py

@@ -109,12 +109,7 @@ def _load_sharp_model(config):
         _model_patcher = patcher
         _model_config = config
 
-    # Load to GPU via ComfyUI VRAM management
-    # SPN processes ~35 patches through ViT + merge + upsample features + decode
-    # ~3 GB activation memory at 1536x1536 with chunked processing
-    memory_required = 3 * 1024 * 1024 * 1024
-    comfy.model_management.load_models_gpu([_model_patcher], memory_required=memory_required)
-    return _model_patcher.model, _model_patcher.load_device
+    return _model_patcher
 
 
 # -- Node -----------------------------------------------------------------

nodes/predict.py

@@ -108,7 +108,9 @@ class SharpPredict(io.ComfyNode):
         from .sharp.gaussians import save_ply, unproject_gaussians
 
         # model is a config dict from LoadSharpModel — load on-demand
-        predictor, device = _load_sharp_model(model)
+        patcher = _load_sharp_model(model)
+        predictor = patcher.model
+        device = patcher.load_device
 
         # Handle batch dimension
         if image.dim() == 3:
@@ -180,7 +182,7 @@ class SharpPredict(io.ComfyNode):
             # Run inference with caching
             log.info(f"Running inference on image {i+1}/{batch_size}...")
             gaussians = _predict_image_cached(
-                predictor, image_np, f_px, device,
+                patcher, predictor, image_np, f_px, device,
                 extrinsics=img_extrinsics,
                 intrinsics=img_intrinsics,
             )
@@ -216,6 +218,7 @@ class SharpPredict(io.ComfyNode):
 
 
 def _predict_image_cached(
+    patcher,
     predictor,
     image: np.ndarray,
     f_px: float,
@@ -229,6 +232,7 @@ def _predict_image_cached(
     Changing focal_length reuses cached features (instant).
 
     Args:
+        patcher: ComfyUI ModelPatcher wrapping the predictor
         predictor: SHARP predictor model
         image: Input image as numpy array [H, W, 3]
         f_px: Focal length in pixels
@@ -241,6 +245,11 @@ def _predict_image_cached(
     from .sharp.gaussians import unproject_gaussians
 
     internal_shape = (1536, 1536)
+
+    # Load to GPU with dynamic memory budget based on input shape
+    input_shape = [1, 3, internal_shape[0], internal_shape[1]]
+    memory_required = patcher.memory_required(input_shape)
+    comfy.model_management.load_models_gpu([patcher], memory_required=memory_required)
     height, width = image.shape[:2]
 
     # Compute image hash for cache

nodes/predict_depth.py

@@ -68,7 +68,9 @@ class SharpPredictDepth(io.ComfyNode):
         from .load_model import _load_sharp_model
 
         # model is a config dict from LoadSharpModel — load on-demand
-        predictor, device = _load_sharp_model(model)
+        patcher = _load_sharp_model(model)
+        predictor = patcher.model
+        device = patcher.load_device
 
         # Handle batch dimension
         if image.dim() == 3:
@@ -99,6 +101,11 @@ class SharpPredictDepth(io.ComfyNode):
         # We output depth at native disparity resolution (1536x1536)
         internal_shape = (1536, 1536)
 
+        # Load to GPU with dynamic memory budget based on input shape
+        input_shape = [1, 3, internal_shape[0], internal_shape[1]]
+        memory_required = patcher.memory_required(input_shape)
+        comfy.model_management.load_models_gpu([patcher], memory_required=memory_required)
+
         all_depth_maps = []
         all_alignment_maps = []
 

nodes/sharp/model.py

@@ -585,6 +585,34 @@ class SlidingPyramidNetwork(nn.Module):
     def internal_resolution(self) -> int:
         return self.patch_size * 4
 
+    def memory_required(self, input_shape):
+        """Estimate peak activation VRAM for SPN encoder + Monodepth decoder.
+
+        The peak occurs during the decoder fusion phase when SPN feature maps
+        and decoder intermediates (ResNet Conv2d) are all live.  Memory scales
+        with the number of pyramid patches (depends on input spatial dims),
+        the ViT embedding width, and dtype — following the same empirical-factor
+        pattern as ComfyUI's ``memory_usage_factor`` and VAE estimators.
+        """
+        import comfy.model_management
+        H, W = input_shape[-2], input_shape[-1]
+        dtype = next(self.parameters()).dtype
+        dtype_size = comfy.model_management.dtype_size(dtype)
+
+        # Patch count from the 3-level pyramid (mirrors forward() logic)
+        ps = self.patch_size
+        n_x0 = (math.ceil((H - ps) / int(ps * 0.75)) + 1) ** 2 if H > ps else 1
+        n_x1 = (math.ceil((H // 2 - ps) / int(ps * 0.5)) + 1) ** 2 if H // 2 > ps else 1
+        n_total = n_x0 + n_x1 + 1
+
+        # Per-patch memory covers: chunked ViT forward, encoding storage,
+        # merged feature maps, ConvTranspose2d upsample intermediates,
+        # Monodepth decoder fusion+ResNet Conv2d, and PyTorch allocator overhead.
+        # Calibrated at 1536x1536 bf16: 35 patches, ~6.8 GB peak reserved.
+        mem_per_patch = 200 * 1024 * 1024 * (dtype_size / 2) * (self.patch_encoder.embed_dim / 1024)
+
+        return int(n_total * mem_per_patch)
+
     def _create_pyramid(
         self, x: torch.Tensor
     ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
@@ -1970,6 +1998,10 @@ class RGBGaussianPredictor(nn.Module):
         self.gaussian_composer = gaussian_composer
         self.depth_alignment = DepthAlignment(scale_map_estimator)
 
+    def memory_required(self, input_shape):
+        """Estimate peak activation VRAM (delegates to SPN encoder)."""
+        return self.monodepth_model.monodepth_predictor.encoder.memory_required(input_shape)
+
     def encode(self, image: torch.Tensor):
         monodepth_output = self.monodepth_model(image)
         return monodepth_output, image

requirements.txt

@@ -1,2 +1,2 @@
-comfy-env==0.2.44
+comfy-env==0.2.61
 comfy-3d-viewers==0.2.41