From afcc321771d0d7eb04fef22aeb345530fc2220cc Mon Sep 17 00:00:00 2001 From: PozzettiAndrea Date: Fri, 20 Feb 2026 20:38:46 +0100 Subject: [PATCH] made it mps compatible --- nodes/align_depth_maps.py | 73 ++++++++-------- nodes/load_image_exif.py | 9 +- nodes/load_model.py | 133 ++++++++++++++--------------- nodes/merge_gaussians.py | 15 ++-- nodes/predict.py | 48 ++++++++--- nodes/predict_depth.py | 55 ++++++++---- nodes/project_depth_to_panorama.py | 23 ++--- nodes/sample_panorama.py | 11 ++- requirements.txt | 4 +- 9 files changed, 210 insertions(+), 161 deletions(-) diff --git a/nodes/align_depth_maps.py b/nodes/align_depth_maps.py index 544ec70..f9b6aec 100644 --- a/nodes/align_depth_maps.py +++ b/nodes/align_depth_maps.py @@ -5,12 +5,15 @@ across multiple perspective views before panorama projection. """ import json +import logging import math from collections import deque import torch import torch.nn.functional as F +log = logging.getLogger("sharp") + def extrinsics_to_direction(extrinsics: torch.Tensor) -> torch.Tensor: """Get the viewing direction (forward vector) from extrinsics. @@ -327,15 +330,15 @@ def compute_global_alignments( ) # Detailed debug output - print(f"[AlignDepthMaps] View {neighbor} -> {current}:") - print(f" Overlap: {debug_info['overlap_pixels']} pixels, {debug_info['valid_pixels']} valid") - print(f" Depth1 (ref): {debug_info['d1_range'][0]:.2f} - {debug_info['d1_range'][1]:.2f}") - print(f" Depth2 (src): {debug_info['d2_range'][0]:.2f} - {debug_info['d2_range'][1]:.2f}") - print(f" Ratio range: {debug_info['ratio_range'][0]:.3f} - {debug_info['ratio_range'][1]:.3f}") - print(f" Scale: {scale:.4f} (median={debug_info['ratio_median']:.4f}, std={debug_info['ratio_std']:.4f})") + log.debug(f"View {neighbor} -> {current}:") + log.debug(f" Overlap: {debug_info['overlap_pixels']} pixels, {debug_info['valid_pixels']} valid") + log.debug(f" Depth1 (ref): {debug_info['d1_range'][0]:.2f} - {debug_info['d1_range'][1]:.2f}") + log.debug(f" Depth2 (src): {debug_info['d2_range'][0]:.2f} - {debug_info['d2_range'][1]:.2f}") + log.debug(f" Ratio range: {debug_info['ratio_range'][0]:.3f} - {debug_info['ratio_range'][1]:.3f}") + log.debug(f" Scale: {scale:.4f} (median={debug_info['ratio_median']:.4f}, std={debug_info['ratio_std']:.4f})") if debug_info.get('scale_clamped'): - print(f" WARNING: Scale clamped from {debug_info['original_scale']:.4f} to {scale:.4f}") - print(f" Inlier ratio: {debug_info['inlier_ratio']:.1%}") + log.warning(f"Scale clamped from {debug_info['original_scale']:.4f} to {scale:.4f}") + log.debug(f" Inlier ratio: {debug_info['inlier_ratio']:.1%}") alignments[neighbor] = (scale, shift) visited.add(neighbor) @@ -344,7 +347,7 @@ def compute_global_alignments( # Handle disconnected views (no overlap with any aligned view) for i in range(num_views): if i not in alignments: - print(f"[AlignDepthMaps] Warning: View {i} has no overlap with aligned views, using identity") + log.warning(f"View {i} has no overlap with aligned views, using identity") alignments[i] = (1.0, 0.0) return alignments @@ -416,13 +419,13 @@ def compute_global_alignments_optimized( depth_h, depth_w = depth_maps.shape[1:3] device = depth_maps.device - print(f"[GlobalOpt] Computing globally optimal alignment for {num_views} views") + log.info(f"Computing globally optimal alignment for {num_views} views") # Step 1: Compute ALL pairwise scale ratios adjacency = build_adjacency_graph(extrinsics, intrinsics) pairwise_ratios = {} # (i, j) -> ratio where i < j - print(f"[GlobalOpt] Computing pairwise ratios...") + log.info("Computing pairwise ratios...") for i in range(num_views): for j in adjacency[i]: if i < j: # Only compute once per pair @@ -433,7 +436,7 @@ def compute_global_alignments_optimized( ) if coords_i.shape[0] < 10: - print(f"[GlobalOpt] Pair ({i},{j}): insufficient overlap ({coords_i.shape[0]} pixels)") + log.debug(f"Pair ({i},{j}): insufficient overlap ({coords_i.shape[0]} pixels)") continue depth_i = depth_maps[i, :, :, 0] @@ -446,25 +449,25 @@ def compute_global_alignments_optimized( if n_valid > 100: pairwise_ratios[(i, j)] = scale - print(f"[GlobalOpt] Pair ({i},{j}): ratio={scale:.4f}, " - f"valid={n_valid}, inlier_ratio={debug_info['inlier_ratio']:.1%}") + log.debug(f"Pair ({i},{j}): ratio={scale:.4f}, " + f"valid={n_valid}, inlier_ratio={debug_info['inlier_ratio']:.1%}") else: - print(f"[GlobalOpt] Pair ({i},{j}): insufficient valid pixels ({n_valid})") + log.debug(f"Pair ({i},{j}): insufficient valid pixels ({n_valid})") # Compute initial loss (with identity scales) identity_scales = [1.0] * num_views initial_loss, initial_rmse, _ = compute_global_misalignment(identity_scales, pairwise_ratios) - print(f"[GlobalOpt] Initial misalignment (identity scales): loss={initial_loss:.4f}, RMSE={initial_rmse:.4f}") + log.info(f"Initial misalignment (identity scales): loss={initial_loss:.4f}, RMSE={initial_rmse:.4f}") # Step 2: Build linear system for log-scale optimization # x_i - x_j = log(r_ij) for each pair # Fix x_0 = 0 (reference) num_pairs = len(pairwise_ratios) - print(f"[GlobalOpt] Found {num_pairs} valid pairs") + log.info(f"Found {num_pairs} valid pairs") if num_pairs == 0: - print("[GlobalOpt] No valid pairs, using identity scales") + log.warning("No valid pairs, using identity scales") return {i: (1.0, 0.0) for i in range(num_views)} # Build A matrix and b vector @@ -501,7 +504,7 @@ def compute_global_alignments_optimized( b[k] = log_ratio # Step 3: Solve least squares: x = (A^T A)^{-1} A^T b - print(f"[GlobalOpt] Solving least squares system ({num_pairs} equations, {num_views-1} unknowns)") + log.info(f"Solving least squares system ({num_pairs} equations, {num_views-1} unknowns)") try: # Use torch.linalg.lstsq for numerical stability @@ -512,14 +515,14 @@ def compute_global_alignments_optimized( for i in range(num_views - 1): scales.append(math.exp(solution[i].item())) - print(f"[GlobalOpt] Raw scales: {[f'{s:.4f}' for s in scales]}") + log.debug(f"Raw scales: {[f'{s:.4f}' for s in scales]}") # Compute loss after optimization (before normalization) opt_loss, opt_rmse, per_pair = compute_global_misalignment(scales, pairwise_ratios) - print(f"[GlobalOpt] Optimized misalignment: loss={opt_loss:.4f}, RMSE={opt_rmse:.4f}") + log.info(f"Optimized misalignment: loss={opt_loss:.4f}, RMSE={opt_rmse:.4f}") except Exception as e: - print(f"[GlobalOpt] Least squares failed: {e}, using identity scales") + log.error(f"Least squares failed: {e}, using identity scales") scales = [1.0] * num_views # Step 4: Normalize so median scale = 1.0 (prevents overall drift) @@ -527,29 +530,29 @@ def compute_global_alignments_optimized( if median_scale > 0: scales = [s / median_scale for s in scales] - print(f"[GlobalOpt] Normalized scales: {[f'{s:.4f}' for s in scales]}") + log.debug(f"Normalized scales: {[f'{s:.4f}' for s in scales]}") # Compute final loss after normalization final_loss, final_rmse, per_pair = compute_global_misalignment(scales, pairwise_ratios) - print(f"[GlobalOpt] Final misalignment: loss={final_loss:.4f}, RMSE={final_rmse:.4f}") + log.info(f"Final misalignment: loss={final_loss:.4f}, RMSE={final_rmse:.4f}") # Print per-pair residuals - print(f"[GlobalOpt] Per-pair log-errors:") + log.debug("Per-pair log-errors:") for (i, j), err in sorted(per_pair.items()): status = "OK" if err < 0.1 else "WARN" if err < 0.2 else "BAD" - print(f" Pair ({i},{j}): {err:.4f} [{status}]") + log.debug(f" Pair ({i},{j}): {err:.4f} [{status}]") # Summary min_scale, max_scale = min(scales), max(scales) variation = max_scale / min_scale if min_scale > 0 else float('inf') - print(f"[GlobalOpt] Scale range: {min_scale:.4f} - {max_scale:.4f} ({variation:.2f}x variation)") + log.info(f"Scale range: {min_scale:.4f} - {max_scale:.4f} ({variation:.2f}x variation)") improvement = (initial_rmse - final_rmse) / initial_rmse * 100 if initial_rmse > 0 else 0 - print(f"[GlobalOpt] Improvement: {improvement:.1f}% reduction in RMSE") + log.info(f"Improvement: {improvement:.1f}% reduction in RMSE") # Warn if variation is still high if variation > 2.0: - print(f"[GlobalOpt] WARNING: High scale variation ({variation:.2f}x) - alignment may be unreliable") + log.warning(f"High scale variation ({variation:.2f}x) - alignment may be unreliable") return {i: (scales[i], 0.0) for i in range(num_views)} @@ -602,8 +605,8 @@ class AlignDepthMaps: """Align depth maps using global optimization or BFS chain.""" num_views = depth_maps.shape[0] - print(f"[AlignDepthMaps] Aligning {num_views} depth maps") - print(f"[AlignDepthMaps] Method: {method}") + log.info(f"Aligning {num_views} depth maps") + log.info(f"Method: {method}") # Clamp reference view to valid range reference_view = min(reference_view, num_views - 1) @@ -615,7 +618,7 @@ class AlignDepthMaps: ) else: # Legacy BFS method (kept for comparison) - print(f"[AlignDepthMaps] Using BFS chain (reference view: {reference_view})") + log.info(f"Using BFS chain (reference view: {reference_view})") alignments = compute_global_alignments( depth_maps, extrinsics, intrinsics, reference_view, num_iterations=1000, inlier_threshold=0.05 @@ -645,12 +648,12 @@ class AlignDepthMaps: min_scale, max_scale = min(scales), max(scales) variation = max_scale / min_scale if min_scale > 0 else float('inf') - print(f"[AlignDepthMaps] Alignment complete") - print(f"[AlignDepthMaps] Final scale range: {min_scale:.4f} - {max_scale:.4f} ({variation:.2f}x variation)") + log.info("Alignment complete") + log.info(f"Final scale range: {min_scale:.4f} - {max_scale:.4f} ({variation:.2f}x variation)") # Check for problematic alignments if variation > 2.0: - print(f"[AlignDepthMaps] WARNING: Large scale variation ({variation:.2f}x) - alignment may be unreliable.") + log.warning(f"Large scale variation ({variation:.2f}x) - alignment may be unreliable.") return (aligned_batch, extrinsics, intrinsics, info_str,) diff --git a/nodes/load_image_exif.py b/nodes/load_image_exif.py index 1782b26..8307f3c 100644 --- a/nodes/load_image_exif.py +++ b/nodes/load_image_exif.py @@ -4,12 +4,15 @@ Loads an image and extracts focal length from EXIF metadata. """ import hashlib +import logging import os import numpy as np import torch from PIL import Image, ImageOps, ImageSequence, ExifTags, TiffTags +log = logging.getLogger("sharp") + try: import folder_paths import node_helpers @@ -73,16 +76,16 @@ def extract_focal_length_mm(img_pil: Image.Image, default_mm: float = 30.0) -> f f_35mm = img_exif.get("FocalLength", None) if f_35mm is None: - print(f"[LoadImageWithExif] No focal length in EXIF - using default {default_mm}mm") + log.info(f"No focal length in EXIF - using default {default_mm}mm") return default_mm # If focal length is very small, it's probably not for 35mm equivalent if f_35mm < 10.0: - print(f"[LoadImageWithExif] Found focal length {f_35mm}mm < 10mm, assuming not 35mm equivalent") + log.info(f"Found focal length {f_35mm}mm < 10mm, assuming not 35mm equivalent") # This is a crude approximation (assumes typical smartphone sensor crop factor) f_35mm *= 8.4 - print(f"[LoadImageWithExif] Extracted focal length: {f_35mm}mm (35mm equivalent)") + log.info(f"Extracted focal length: {f_35mm}mm (35mm equivalent)") return float(f_35mm) diff --git a/nodes/load_model.py b/nodes/load_model.py index dc07772..fc6fdc0 100644 --- a/nodes/load_model.py +++ b/nodes/load_model.py @@ -1,10 +1,14 @@ """LoadSharpModel node for ComfyUI-Sharp.""" import os -from urllib.request import urlopen +import logging import torch -from tqdm import tqdm +from huggingface_hub import hf_hub_download +import comfy.model_management +import comfy.model_patcher + +log = logging.getLogger("sharp") # Try to get ComfyUI models directory try: @@ -13,27 +17,47 @@ try: except ImportError: MODELS_DIR = os.path.join(os.path.dirname(os.path.dirname(__file__)), "models", "sharp") -# Model cache -_MODEL_CACHE = {} +SHARP_REPO_ID = "apple/Sharp" +SHARP_FILENAME = "sharp_2572gikvuh.pt" -# Default model URL and filename -DEFAULT_MODEL_URL = "https://huggingface.co/apple/Sharp/resolve/main/sharp_2572gikvuh.pt" -DEFAULT_MODEL_FILENAME = "sharp_2572gikvuh.pt" + +def _build_sharp_model(model_path, dtype): + """Build and load the SHARP predictor model. + + Called lazily by inference nodes on first use. + Returns a loaded nn.Module on CPU with the given dtype. + """ + from .sharp.models import PredictorParams, create_predictor + import comfy.utils + + log.info(f"Loading checkpoint from {model_path}") + state_dict = comfy.utils.load_torch_file(model_path) + + log.info("Initializing model...") + with torch.device("meta"): + predictor = create_predictor(PredictorParams()) + predictor.load_state_dict(state_dict, assign=True) + predictor.eval() + predictor.to(dtype=dtype) # set dtype, stay on CPU + log.info(f"Model ready ({dtype})") + return predictor class LoadSharpModel: - """Load and cache the SHARP model for 3D Gaussian prediction.""" + """Configure the SHARP model. Downloads checkpoint if needed. + + Model weights are loaded lazily by inference nodes on first use. + """ @classmethod def INPUT_TYPES(cls): return { - "required": { - "device": (["auto", "cuda", "mps", "cpu"], { - "default": "auto", - "tooltip": "Device to run inference on. 'auto' selects best available." - }), - }, + "required": {}, "optional": { + "precision": (["auto", "bf16", "fp16", "fp32"], { + "default": "auto", + "tooltip": "Model precision. auto: best for your GPU (bf16 on Ampere+, fp16 on Volta/Turing, fp32 on older)." + }), "checkpoint_path": ("STRING", { "default": "", "tooltip": "Path to .pt checkpoint. Leave empty to auto-download from Hugging Face." @@ -45,72 +69,39 @@ class LoadSharpModel: RETURN_NAMES = ("model",) FUNCTION = "load_model" CATEGORY = "SHARP" - DESCRIPTION = "Load the SHARP model for monocular 3D Gaussian Splatting prediction." + DESCRIPTION = "Configure the SHARP model for monocular 3D Gaussian Splatting prediction." - def load_model(self, device: str, checkpoint_path: str = ""): - """Load and cache the SHARP model.""" - from .sharp.models import PredictorParams, create_predictor + def load_model(self, precision: str = "auto", checkpoint_path: str = ""): + """Resolve config and download model if needed. No weight loading.""" + device = comfy.model_management.get_torch_device() - # Resolve device - if device == "auto": - if torch.cuda.is_available(): - device = "cuda" - elif hasattr(torch, "mps") and torch.mps.is_available(): - device = "mps" + # Resolve dtype + if precision == "auto": + if comfy.model_management.should_use_bf16(device): + dtype = torch.bfloat16 + elif comfy.model_management.should_use_fp16(device): + dtype = torch.float16 else: - device = "cpu" + dtype = torch.float32 + elif precision == "bf16": + dtype = torch.bfloat16 + elif precision == "fp16": + dtype = torch.float16 + else: + dtype = torch.float32 - # Create cache key - cache_key = f"{checkpoint_path or 'default'}_{device}" - - if cache_key in _MODEL_CACHE: - return (_MODEL_CACHE[cache_key],) - - # Load or download checkpoint + # Resolve / download checkpoint if checkpoint_path and os.path.exists(checkpoint_path): model_path = checkpoint_path else: - # Use default model in ComfyUI models/sharp directory os.makedirs(MODELS_DIR, exist_ok=True) - model_path = os.path.join(MODELS_DIR, DEFAULT_MODEL_FILENAME) + model_path = hf_hub_download( + repo_id=SHARP_REPO_ID, + filename=SHARP_FILENAME, + local_dir=MODELS_DIR, + ) - if not os.path.exists(model_path): - print(f"[SHARP] Downloading model to {model_path}") - response = urlopen(DEFAULT_MODEL_URL) - total_size = int(response.headers.get('content-length', 0)) - - with open(model_path, 'wb') as f: - with tqdm(total=total_size, unit='B', unit_scale=True, desc="Downloading SHARP model") as pbar: - while True: - chunk = response.read(8192) - if not chunk: - break - f.write(chunk) - pbar.update(len(chunk)) - print(f"[SHARP] Download complete.") - else: - print(f"[SHARP] Loading model from {model_path}") - - print(f"[SHARP] Loading checkpoint...") - state_dict = torch.load(model_path, weights_only=True) - - # Create predictor - print(f"[SHARP] Initializing model...") - predictor = create_predictor(PredictorParams()) - predictor.load_state_dict(state_dict) - predictor.eval() - print(f"[SHARP] Moving model to {device}...") - predictor.to(device) - print(f"[SHARP] Model ready!") - - model_dict = { - "predictor": predictor, - "device": device, - } - - _MODEL_CACHE[cache_key] = model_dict - - return (model_dict,) + return ({"model_path": model_path, "dtype": dtype},) NODE_CLASS_MAPPINGS = { diff --git a/nodes/merge_gaussians.py b/nodes/merge_gaussians.py index a49fe0c..9d2c4b6 100644 --- a/nodes/merge_gaussians.py +++ b/nodes/merge_gaussians.py @@ -3,6 +3,7 @@ Merges multiple PLY files (from panorama samples) into a single unified Gaussian scene. """ +import logging import os import time from pathlib import Path @@ -11,6 +12,8 @@ import numpy as np import torch from plyfile import PlyData, PlyElement +log = logging.getLogger("sharp") + def load_ply_simple(path: str) -> dict: """Load Gaussian data from PLY file. @@ -162,7 +165,7 @@ class MergeGaussians: if not ply_files: raise ValueError(f"No PLY files found in: {ply_folder}") - print(f"[MergeGaussians] Found {len(ply_files)} PLY files to merge") + log.info(f"Found {len(ply_files)} PLY files to merge") # Load all PLY files all_positions = [] @@ -172,7 +175,7 @@ class MergeGaussians: all_opacities = [] for i, ply_path in enumerate(ply_files): - print(f"[MergeGaussians] Loading {ply_path.name} ({i+1}/{len(ply_files)})") + log.info(f"Loading {ply_path.name} ({i+1}/{len(ply_files)})") data = load_ply_simple(str(ply_path)) positions = data['positions'] @@ -194,7 +197,7 @@ class MergeGaussians: filtered_count = (~mask).sum() if filtered_count > 0: - print(f"[MergeGaussians] Filtered out {filtered_count:,} Gaussians") + log.info(f" Filtered out {filtered_count:,} Gaussians") all_positions.append(positions[mask]) all_colors.append(colors[mask]) @@ -210,14 +213,14 @@ class MergeGaussians: merged_opacities = np.concatenate(all_opacities, axis=0) num_gaussians = len(merged_positions) - print(f"[MergeGaussians] Total Gaussians after merge: {num_gaussians:,}") + log.info(f"Total Gaussians after merge: {num_gaussians:,}") # Save merged PLY timestamp = int(time.time() * 1000) output_filename = f"{output_prefix}_{timestamp}.ply" output_path = ply_folder.parent / output_filename - print(f"[MergeGaussians] Saving to {output_path}") + log.info(f"Saving to {output_path}") save_merged_ply( merged_positions, merged_colors, @@ -227,7 +230,7 @@ class MergeGaussians: str(output_path), ) - print(f"[MergeGaussians] Done! Merged {len(ply_files)} files into {num_gaussians:,} Gaussians") + log.info(f"Done! Merged {len(ply_files)} files into {num_gaussians:,} Gaussians") return (str(output_path), num_gaussians,) diff --git a/nodes/predict.py b/nodes/predict.py index cd746dc..e5d893e 100644 --- a/nodes/predict.py +++ b/nodes/predict.py @@ -1,6 +1,7 @@ """SharpPredict node for ComfyUI-Sharp.""" import hashlib +import logging import os import time from pathlib import Path @@ -8,6 +9,12 @@ from pathlib import Path import numpy as np import torch import torch.nn.functional as F +import comfy.model_management +import comfy.model_patcher + +from .load_model import _build_sharp_model + +log = logging.getLogger("sharp") # Try to import ComfyUI folder_paths for output directory try: @@ -71,10 +78,23 @@ class SharpPredict: OUTPUT_NODE = True DESCRIPTION = "Generate 3D Gaussian Splatting PLY file(s) from image(s) using SHARP. Batch input creates a folder with numbered PLY files." + def _get_patcher(self, config): + """Lazily build and cache ModelPatcher from config dict.""" + key = (config["model_path"], config["dtype"]) + if not hasattr(self, '_patcher') or self._config_key != key: + predictor = _build_sharp_model(config["model_path"], config["dtype"]) + self._patcher = comfy.model_patcher.ModelPatcher( + predictor, + load_device=comfy.model_management.get_torch_device(), + offload_device=comfy.model_management.unet_offload_device(), + ) + self._config_key = key + return self._patcher + @torch.no_grad() def predict( self, - model: dict, + model, image: torch.Tensor, focal_length_mm: float = 0.0, output_prefix: str = "sharp", @@ -93,8 +113,10 @@ class SharpPredict: """ from .sharp.utils.gaussians import save_ply, unproject_gaussians - predictor = model["predictor"] - device = torch.device(model["device"]) + patcher = self._get_patcher(model) + comfy.model_management.load_models_gpu([patcher]) + predictor = patcher.model + device = patcher.load_device # Handle batch dimension if image.dim() == 3: @@ -109,9 +131,9 @@ class SharpPredict: extrinsics = extrinsics.unsqueeze(0) if extrinsics.shape[0] != batch_size: raise ValueError(f"Extrinsics batch size ({extrinsics.shape[0]}) must match image batch size ({batch_size})") - print(f"[SHARP] Processing {batch_size} image(s) with provided camera parameters (panorama mode)") + log.info(f"Processing {batch_size} image(s) with provided camera parameters (panorama mode)") else: - print(f"[SHARP] Processing {batch_size} image(s)") + log.info(f"Processing {batch_size} image(s)") # Ensure output directory exists os.makedirs(OUTPUT_DIR, exist_ok=True) @@ -144,7 +166,7 @@ class SharpPredict: height, width = image_np.shape[:2] if i == 0: - print(f"[SHARP] Image size: {width}x{height}") + log.info(f"Image size: {width}x{height}") # Get camera parameters for this image if has_camera_params: @@ -162,7 +184,7 @@ class SharpPredict: img_intrinsics = None # Run inference with caching - print(f"[SHARP] Running inference on image {i+1}/{batch_size}...") + log.info(f"Running inference on image {i+1}/{batch_size}...") gaussians = self._predict_image_cached( predictor, image_np, f_px, device, extrinsics=img_extrinsics, @@ -183,10 +205,10 @@ class SharpPredict: all_extrinsics.append(metadata["extrinsic"]) all_intrinsics.append(metadata["intrinsic"]) - print(f"[SHARP] Saved: {ply_path} ({metadata['num_gaussians']:,} gaussians)") + log.info(f"Saved: {ply_path} ({metadata['num_gaussians']:,} gaussians)") inference_time = time.time() - inference_start - print(f"[SHARP] Total inference time: {inference_time:.2f}s ({inference_time/batch_size:.2f}s per image)") + log.info(f"Total inference time: {inference_time:.2f}s ({inference_time/batch_size:.2f}s per image)") # Return values if is_batch: @@ -231,12 +253,12 @@ class SharpPredict: # Check cache if _encode_cache["image_hash"] == image_hash: # Cache hit - reuse encoded features - print(f"[SHARP] Cache hit - reusing encoded features (focal_length change is instant)") + log.info("Cache hit - reusing encoded features (focal_length change is instant)") monodepth_output = _encode_cache["monodepth_output"] image_resized_pt = _encode_cache["image_resized"] else: # Cache miss - need to encode - print(f"[SHARP] Cache miss - running encoder (this is the slow part)...") + log.info("Cache miss - running encoder (this is the slow part)...") # Clear old cache _encode_cache["image_hash"] = None @@ -259,7 +281,7 @@ class SharpPredict: encode_start = time.time() monodepth_output, _ = predictor.encode(image_resized_pt) encode_time = time.time() - encode_start - print(f"[SHARP] Encode time: {encode_time:.2f}s") + log.info(f"Encode time: {encode_time:.2f}s") # Update cache _encode_cache["image_hash"] = image_hash @@ -273,7 +295,7 @@ class SharpPredict: decode_start = time.time() gaussians_ndc = predictor.decode(monodepth_output, image_resized_pt, disparity_factor) decode_time = time.time() - decode_start - print(f"[SHARP] Decode time: {decode_time:.2f}s") + log.info(f"Decode time: {decode_time:.2f}s") # Build intrinsics for unprojection (use provided or construct from f_px) if intrinsics is not None: diff --git a/nodes/predict_depth.py b/nodes/predict_depth.py index 6271310..c729876 100644 --- a/nodes/predict_depth.py +++ b/nodes/predict_depth.py @@ -4,14 +4,20 @@ Runs SHARP inference but outputs depth maps instead of PLY files. Used for panorama depth visualization and verification before Gaussian creation. """ +import logging import time import numpy as np import torch import torch.nn.functional as F +import comfy.model_management +import comfy.model_patcher +from .load_model import _build_sharp_model from .utils.image import comfy_to_numpy_rgb +log = logging.getLogger("sharp") + class SharpPredictDepth: """Run SHARP inference to generate depth maps from images. @@ -46,10 +52,23 @@ class SharpPredictDepth: CATEGORY = "SHARP" DESCRIPTION = "Generate depth maps from images using SHARP. Optionally align to reference depth (e.g., DepthAnythingV3) using learned dense alignment." + def _get_patcher(self, config): + """Lazily build and cache ModelPatcher from config dict.""" + key = (config["model_path"], config["dtype"]) + if not hasattr(self, '_patcher') or self._config_key != key: + predictor = _build_sharp_model(config["model_path"], config["dtype"]) + self._patcher = comfy.model_patcher.ModelPatcher( + predictor, + load_device=comfy.model_management.get_torch_device(), + offload_device=comfy.model_management.unet_offload_device(), + ) + self._config_key = key + return self._patcher + @torch.no_grad() def predict_depth( self, - model: dict, + model, image: torch.Tensor, extrinsics: torch.Tensor = None, intrinsics: torch.Tensor = None, @@ -60,15 +79,17 @@ class SharpPredictDepth: Returns depth maps as [N, H, W, 1] tensor (grayscale images). If reference_depth is provided, uses SHARP's learned dense alignment. """ - predictor = model["predictor"] - device = torch.device(model["device"]) + patcher = self._get_patcher(model) + comfy.model_management.load_models_gpu([patcher]) + predictor = patcher.model + device = patcher.load_device # Handle batch dimension if image.dim() == 3: image = image.unsqueeze(0) batch_size = image.shape[0] - print(f"[SharpPredictDepth] Processing {batch_size} image(s)") + log.info(f"Processing {batch_size} image(s)") # Check if we have reference depth for alignment use_alignment = reference_depth is not None @@ -76,17 +97,17 @@ class SharpPredictDepth: if reference_depth.dim() == 3: reference_depth = reference_depth.unsqueeze(0) if reference_depth.shape[0] != batch_size: - print(f"[SharpPredictDepth] WARNING: reference_depth batch size ({reference_depth.shape[0]}) " - f"doesn't match image batch size ({batch_size}). Disabling alignment.") + log.warning(f"reference_depth batch size ({reference_depth.shape[0]}) " + f"doesn't match image batch size ({batch_size}). Disabling alignment.") use_alignment = False else: # Check if scale_map_estimator is available scale_map_estimator = predictor.depth_alignment.scale_map_estimator if scale_map_estimator is None: - print(f"[SharpPredictDepth] WARNING: scale_map_estimator not available in model. Disabling alignment.") + log.warning("scale_map_estimator not available in model. Disabling alignment.") use_alignment = False else: - print(f"[SharpPredictDepth] Using learned dense alignment with reference depth") + log.info("Using learned dense alignment with reference depth") # SHARP processes at 1536x1536 internally # We output depth at native disparity resolution (1536x1536) @@ -104,7 +125,7 @@ class SharpPredictDepth: height, width = image_np.shape[:2] if i == 0: - print(f"[SharpPredictDepth] Image size: {width}x{height}") + log.info(f"Image size: {width}x{height}") # Get focal length from intrinsics or use default if intrinsics is not None: @@ -127,7 +148,7 @@ class SharpPredictDepth: ) # Run SHARP encode/decode - print(f"[SharpPredictDepth] Running inference on image {i+1}/{batch_size}...") + log.info(f"Running inference on image {i+1}/{batch_size}...") encode_start = time.time() monodepth_output, depth_decoder_features = predictor.encode(image_resized_pt) @@ -136,7 +157,7 @@ class SharpPredictDepth: # Compute disparity factor (consistent across all views with same FOV) disparity_factor = f_px / width - print(f"[SharpPredictDepth] Encode: {encode_time:.2f}s, disparity_factor: {disparity_factor:.4f}") + log.info(f" Encode: {encode_time:.2f}s, disparity_factor: {disparity_factor:.4f}") # Extract depth from RAW DISPARITY (not Gaussians!) # monodepth_output.disparity is [1, 1, 1536, 1536] at internal resolution @@ -180,7 +201,7 @@ class SharpPredictDepth: # Convert back to disparity for consistent output aligned_disparity = disparity_factor / aligned_depth.clamp(min=1e-4) - print(f"[SharpPredictDepth] Alignment map range: {alignment_map.min():.3f} - {alignment_map.max():.3f}") + log.info(f" Alignment map range: {alignment_map.min():.3f} - {alignment_map.max():.3f}") # Store alignment map for visualization alignment_map_vis = alignment_map[0, 0].cpu() # [H, W] @@ -205,7 +226,7 @@ class SharpPredictDepth: depth_map_raw = disparity_factor / disparity_resized.clamp(min=1e-4) # [H, W] if i == 0: - print(f"[SharpPredictDepth] Depth range: {depth_map_raw.min():.3f} - {depth_map_raw.max():.3f}") + log.debug(f" Depth range: {depth_map_raw.min():.3f} - {depth_map_raw.max():.3f}") # Convert to [H, W, 1] format for ComfyUI depth_map = depth_map_raw.unsqueeze(-1).cpu() # [H, W, 1] @@ -216,13 +237,13 @@ class SharpPredictDepth: all_depth_maps.append(depth_map_rgb) inference_time = time.time() - inference_start - print(f"[SharpPredictDepth] Total inference time: {inference_time:.2f}s") + log.info(f"Total inference time: {inference_time:.2f}s") # Stack all depth maps depth_maps_batch = torch.stack(all_depth_maps, dim=0) # [N, H, W, 3] alignment_maps_batch = torch.stack(all_alignment_maps, dim=0) # [N, H, W, 3] - print(f"[SharpPredictDepth] Output shape: {depth_maps_batch.shape}") + log.info(f"Output shape: {depth_maps_batch.shape}") # Scale intrinsics to match output depth map size (1536x1536 internal resolution) # Input intrinsics are for the original image size @@ -233,8 +254,8 @@ class SharpPredictDepth: scale_factor = output_size / input_size if abs(scale_factor - 1.0) > 0.01: # Only scale if significantly different - print(f"[SharpPredictDepth] Scaling intrinsics by {scale_factor:.3f} " - f"(input {input_size} -> output {output_size})") + log.info(f"Scaling intrinsics by {scale_factor:.3f} " + f"(input {input_size} -> output {output_size})") # Scale f_px, f_py, cx, cy intrinsics_scaled = intrinsics.clone() intrinsics_scaled[0, 0] *= scale_factor # f_px diff --git a/nodes/project_depth_to_panorama.py b/nodes/project_depth_to_panorama.py index 247bc33..8bc8678 100644 --- a/nodes/project_depth_to_panorama.py +++ b/nodes/project_depth_to_panorama.py @@ -4,12 +4,15 @@ Projects depth maps from perspective views back to equirectangular panorama with blending and debug visualization. """ +import logging import math import numpy as np import torch import torch.nn.functional as F +log = logging.getLogger("sharp") + def extrinsics_to_yaw_pitch(extrinsics: torch.Tensor) -> tuple[float, float]: """Extract yaw and pitch from extrinsics matrix. @@ -209,14 +212,14 @@ def project_depth_maps_to_panorama_with_disagreement( if overlap_mask.any(): overlap_std = std_dev[overlap_mask] overlap_rel = relative_disagreement[overlap_mask] - print(f"[ProjectDepthToPanorama] Overlap regions ({overlap_mask.sum().item()} pixels):") - print(f" Absolute disagreement: mean={overlap_std.mean():.4f}, max={overlap_std.max():.4f}") - print(f" Relative disagreement: mean={overlap_rel.mean():.2%}, max={overlap_rel.max():.2%}") + log.info(f"Overlap regions ({overlap_mask.sum().item()} pixels):") + log.info(f" Absolute disagreement: mean={overlap_std.mean():.4f}, max={overlap_std.max():.4f}") + log.info(f" Relative disagreement: mean={overlap_rel.mean():.2%}, max={overlap_rel.max():.2%}") # Count pixels with high disagreement (>10% relative) high_disagreement = (relative_disagreement > 0.1) & overlap_mask - print(f" High disagreement (>10%): {high_disagreement.sum().item()} pixels " - f"({high_disagreement.sum().item() / overlap_mask.sum().item() * 100:.1f}% of overlap)") + log.info(f" High disagreement (>10%): {high_disagreement.sum().item()} pixels " + f"({high_disagreement.sum().item() / overlap_mask.sum().item() * 100:.1f}% of overlap)") # Global normalization for visualization valid_mask = panorama_weight > 0.01 @@ -225,7 +228,7 @@ def project_depth_maps_to_panorama_with_disagreement( valid_depths = panorama_depth[valid_mask] depth_min = valid_depths.min() depth_max = valid_depths.max() - print(f"[ProjectDepthToPanorama] Depth range: {depth_min:.4f} to {depth_max:.4f}") + log.info(f"Depth range: {depth_min:.4f} to {depth_max:.4f}") if depth_max > depth_min: panorama_depth = torch.where( @@ -643,9 +646,9 @@ class ProjectDepthToPanorama: ): """Project depth maps to panorama with disagreement visualization.""" - print(f"[ProjectDepthToPanorama] Input: {depth_maps.shape[0]} depth maps") - print(f"[ProjectDepthToPanorama] Output panorama: {panorama_width} x {panorama_width // 2}") - print(f"[ProjectDepthToPanorama] Blend mode: {blend_mode}") + log.info(f"Input: {depth_maps.shape[0]} depth maps") + log.info(f"Output panorama: {panorama_width} x {panorama_width // 2}") + log.info(f"Blend mode: {blend_mode}") # Project depth maps with disagreement tracking panoramic_depth, disagreement, overlap_count = project_depth_maps_to_panorama_with_disagreement( @@ -699,7 +702,7 @@ class ProjectDepthToPanorama: debug_overlay_out = debug_overlay.unsqueeze(0) # [1, H, W, 3] disagreement_out = disagreement_heatmap.unsqueeze(0) # [1, H, W, 3] - print(f"[ProjectDepthToPanorama] Output shape: {panoramic_depth_out.shape}") + log.info(f"Output shape: {panoramic_depth_out.shape}") return (panoramic_depth_out, debug_overlay_out, disagreement_out,) diff --git a/nodes/sample_panorama.py b/nodes/sample_panorama.py index 4f8594b..0c8c953 100644 --- a/nodes/sample_panorama.py +++ b/nodes/sample_panorama.py @@ -3,12 +3,15 @@ Samples perspective cutouts from an equirectangular panorama for use with SHARP. """ +import logging import math import numpy as np import torch import torch.nn.functional as F +log = logging.getLogger("sharp") + def create_rotation_matrix(yaw: float, pitch: float) -> torch.Tensor: """Create a rotation matrix from yaw and pitch angles. @@ -221,9 +224,9 @@ class SamplePanorama: num_vertical = max(1, math.ceil(vertical_range / step_degrees)) - print(f"[SamplePanorama] FOV: {fov_degrees}°, Overlap: {overlap_percent}%") - print(f"[SamplePanorama] Step angle: {step_degrees:.1f}°") - print(f"[SamplePanorama] Samples: {num_horizontal} horizontal × {num_vertical} vertical = {num_horizontal * num_vertical} total") + log.info(f"FOV: {fov_degrees}, Overlap: {overlap_percent}%") + log.info(f"Step angle: {step_degrees:.1f}") + log.info(f"Samples: {num_horizontal} horizontal x {num_vertical} vertical = {num_horizontal * num_vertical} total") # Handle batch dimension if panorama.dim() == 3: @@ -265,7 +268,7 @@ class SamplePanorama: images_batch = torch.stack(all_images, dim=0) # [N, H, W, C] extrinsics_batch = torch.stack(all_extrinsics, dim=0) # [N, 4, 4] - print(f"[SamplePanorama] Output shape: {images_batch.shape}") + log.info(f"Output shape: {images_batch.shape}") return (images_batch, extrinsics_batch, intrinsics, num_horizontal, num_vertical) diff --git a/requirements.txt b/requirements.txt index ac252c6..26d0d8d 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,5 +1,5 @@ -comfy-env==0.1.92 -comfy-3d-viewers==0.2.25 +comfy-env==0.1.98 +comfy-3d-viewers==0.2.27 # Core dependencies for SHARP torch>=2.0.0