diff --git a/nodes/load_model.py b/nodes/load_model.py index d46af9c..b0d6425 100644 --- a/nodes/load_model.py +++ b/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 ----------------------------------------------------------------- diff --git a/nodes/predict.py b/nodes/predict.py index a89ca07..4372bae 100644 --- a/nodes/predict.py +++ b/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 diff --git a/nodes/predict_depth.py b/nodes/predict_depth.py index 8ab477e..5dbba84 100644 --- a/nodes/predict_depth.py +++ b/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 = [] diff --git a/nodes/sharp/model.py b/nodes/sharp/model.py index 03026e2..bd06949 100644 --- a/nodes/sharp/model.py +++ b/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 diff --git a/requirements.txt b/requirements.txt index b570166..bb2f85b 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,2 +1,2 @@ -comfy-env==0.2.44 +comfy-env==0.2.61 comfy-3d-viewers==0.2.41 \ No newline at end of file