nodes/PozzettiAndrea-ComfyUI-Sharp
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

made it mps compatible

Browse Source
This commit is contained in:
PozzettiAndrea
2026-02-20 20:38:46 +01:00
parent a927304aa3
commit afcc321771
9 changed files with 210 additions and 161 deletions
Download Patch File Download Diff File Expand all files Collapse all files
nodes/align_depth_maps.py
+38 -35
View File
@@ -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,)
nodes/load_image_exif.py
+6 -3
View File
@@ -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)
nodes/load_model.py
+62 -71
View File
@@ -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 = {
nodes/merge_gaussians.py
+9 -6
View File
@@ -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,)
nodes/predict.py
+35 -13
View File
@@ -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:
nodes/predict_depth.py
+38 -17
View File
@@ -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
nodes/project_depth_to_panorama.py
+13 -10
View File
@@ -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,)
nodes/sample_panorama.py
+7 -4
View File
@@ -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)
requirements.txt
+2 -2
View File
@@ -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