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

added icosahedron merge

Browse Source
This commit is contained in:
PozzettiAndrea
2026-05-26 22:50:11 +02:00
parent 695987087b
commit 344051f3e2
12 changed files with 4020 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files
nodes/__init__.py
+20
View File
@@ -16,6 +16,16 @@ from .project_depth_to_panorama import NODE_CLASS_MAPPINGS as PROJECT_DEPTH_MAPP
from .project_depth_to_panorama import NODE_DISPLAY_NAME_MAPPINGS as PROJECT_DEPTH_DISPLAY_MAPPINGS
from .align_depth_maps import NODE_CLASS_MAPPINGS as ALIGN_DEPTH_MAPPINGS
from .align_depth_maps import NODE_DISPLAY_NAME_MAPPINGS as ALIGN_DEPTH_DISPLAY_MAPPINGS
from .panorama_cube_split import NODE_CLASS_MAPPINGS as PANORAMA_CUBE_SPLIT_MAPPINGS
from .panorama_cube_split import NODE_DISPLAY_NAME_MAPPINGS as PANORAMA_CUBE_SPLIT_DISPLAY_MAPPINGS
from .depth_merge import NODE_CLASS_MAPPINGS as DEPTH_MERGE_MAPPINGS
from .depth_merge import NODE_DISPLAY_NAME_MAPPINGS as DEPTH_MERGE_DISPLAY_MAPPINGS
from .refined_depth_from_gaussians import NODE_CLASS_MAPPINGS as REFINED_DEPTH_MAPPINGS
from .refined_depth_from_gaussians import NODE_DISPLAY_NAME_MAPPINGS as REFINED_DEPTH_DISPLAY_MAPPINGS
from .predict_gaussian_attrs import NODE_CLASS_MAPPINGS as PREDICT_ATTRS_MAPPINGS
from .predict_gaussian_attrs import NODE_DISPLAY_NAME_MAPPINGS as PREDICT_ATTRS_DISPLAY_MAPPINGS
from .panorama_icosahedron_split import NODE_CLASS_MAPPINGS as ICOSAHEDRON_SPLIT_MAPPINGS
from .panorama_icosahedron_split import NODE_DISPLAY_NAME_MAPPINGS as ICOSAHEDRON_SPLIT_DISPLAY_MAPPINGS
# Aggregate all node mappings
NODE_CLASS_MAPPINGS = {}
@@ -29,6 +39,11 @@ NODE_CLASS_MAPPINGS.update(MERGE_GAUSSIANS_MAPPINGS)
NODE_CLASS_MAPPINGS.update(PREDICT_DEPTH_MAPPINGS)
NODE_CLASS_MAPPINGS.update(PROJECT_DEPTH_MAPPINGS)
NODE_CLASS_MAPPINGS.update(ALIGN_DEPTH_MAPPINGS)
NODE_CLASS_MAPPINGS.update(PANORAMA_CUBE_SPLIT_MAPPINGS)
NODE_CLASS_MAPPINGS.update(DEPTH_MERGE_MAPPINGS)
NODE_CLASS_MAPPINGS.update(REFINED_DEPTH_MAPPINGS)
NODE_CLASS_MAPPINGS.update(PREDICT_ATTRS_MAPPINGS)
NODE_CLASS_MAPPINGS.update(ICOSAHEDRON_SPLIT_MAPPINGS)
NODE_DISPLAY_NAME_MAPPINGS.update(LOAD_MODEL_DISPLAY_MAPPINGS)
NODE_DISPLAY_NAME_MAPPINGS.update(PREDICT_DISPLAY_MAPPINGS)
@@ -38,5 +53,10 @@ NODE_DISPLAY_NAME_MAPPINGS.update(MERGE_GAUSSIANS_DISPLAY_MAPPINGS)
NODE_DISPLAY_NAME_MAPPINGS.update(PREDICT_DEPTH_DISPLAY_MAPPINGS)
NODE_DISPLAY_NAME_MAPPINGS.update(PROJECT_DEPTH_DISPLAY_MAPPINGS)
NODE_DISPLAY_NAME_MAPPINGS.update(ALIGN_DEPTH_DISPLAY_MAPPINGS)
NODE_DISPLAY_NAME_MAPPINGS.update(PANORAMA_CUBE_SPLIT_DISPLAY_MAPPINGS)
NODE_DISPLAY_NAME_MAPPINGS.update(DEPTH_MERGE_DISPLAY_MAPPINGS)
NODE_DISPLAY_NAME_MAPPINGS.update(REFINED_DEPTH_DISPLAY_MAPPINGS)
NODE_DISPLAY_NAME_MAPPINGS.update(PREDICT_ATTRS_DISPLAY_MAPPINGS)
NODE_DISPLAY_NAME_MAPPINGS.update(ICOSAHEDRON_SPLIT_DISPLAY_MAPPINGS)
__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]
nodes/_vendor/__init__.py
View File
nodes/_vendor/moge_panorama.py
+354
View File
@@ -0,0 +1,354 @@
import os
os.environ['OPENCV_IO_ENABLE_OPENEXR'] = '1'
from pathlib import Path
from typing import *
import itertools
import json
import warnings
import cv2
import numpy as np
import torch
from numpy import ndarray
from tqdm import tqdm, trange
from scipy.sparse import csr_array, hstack, vstack
from scipy.ndimage import convolve
from scipy.sparse.linalg import lsmr
import utils3d
def get_panorama_cameras():
vertices, _ = utils3d.np.create_icosahedron_mesh()
intrinsics = utils3d.np.intrinsics_from_fov(fov_x=np.deg2rad(90), fov_y=np.deg2rad(90))
extrinsics = utils3d.np.extrinsics_look_at([0, 0, 0], vertices, [0, 0, 1]).astype(np.float32)
return extrinsics, [intrinsics] * len(vertices)
def spherical_uv_to_directions(uv: np.ndarray):
theta, phi = (1 - uv[..., 0]) * (2 * np.pi), uv[..., 1] * np.pi
directions = np.stack([np.sin(phi) * np.cos(theta), np.sin(phi) * np.sin(theta), np.cos(phi)], axis=-1)
return directions
def directions_to_spherical_uv(directions: np.ndarray):
directions = directions / np.linalg.norm(directions, axis=-1, keepdims=True)
u = 1 - np.arctan2(directions[..., 1], directions[..., 0]) / (2 * np.pi) % 1.0
v = np.arccos(directions[..., 2]) / np.pi
return np.stack([u, v], axis=-1)
def split_panorama_image(image: np.ndarray, extrinsics: np.ndarray, intrinsics: np.ndarray, resolution: int):
height, width = image.shape[:2]
uv = utils3d.np.uv_map((resolution, resolution))
splitted_images = []
for i in range(len(extrinsics)):
spherical_uv = directions_to_spherical_uv(utils3d.np.unproject_cv(uv, np.ones_like(uv[..., 0]), extrinsics=extrinsics[i], intrinsics=intrinsics[i]))
pixels = utils3d.np.uv_to_pixel(spherical_uv, (height, width)).astype(np.float32)
splitted_image = cv2.remap(image, pixels[..., 0], pixels[..., 1], interpolation=cv2.INTER_LINEAR)
splitted_images.append(splitted_image)
return splitted_images
def split_panorama_image_gpu(image: np.ndarray, extrinsics: np.ndarray, intrinsics: np.ndarray, resolution: int):
"""Batched-GPU equivalent of `split_panorama_image`.
Same math: for each of N face cameras, build a (resolution, resolution)
lookup of spherical UVs from the panorama, and bilinearly resample.
Difference: the lookup is computed once for all N faces as a batched
(N, R, R, 2) grid (np vectorized over the face dim) and the per-face
resampling collapses to one F.grid_sample call. cv2.remap is single-
threaded per call and was the dominant cost on N=42 face splits.
Bilinear matches cv2.INTER_LINEAR to within fp32 round-off when we use
`align_corners=False` and convert the cv2 pixel coordinate `p` (where
`p=0` is the centre of pixel 0) to the grid_sample coordinate
`(2p + 1) / W - 1`. Panoramic wrap in U is handled by the
`(% 1.0)` inside `directions_to_spherical_uv` so the lookups never
exit the image in U; in V we clip to [0, 1] which is what cv2 does
by default (BORDER_CONSTANT 0 below the south pole, but those samples
are unreachable because the icosahedron forwards are constrained to
the sphere).
Falls back to the CPU implementation if CUDA isn't available.
"""
if not torch.cuda.is_available():
return split_panorama_image(image, extrinsics, intrinsics, resolution)
height, width = image.shape[:2]
N = len(extrinsics)
device = torch.device("cuda")
# Build the (R, R, 2) UV map once on CPU (cheap, sub-ms).
uv = utils3d.np.uv_map((resolution, resolution))
# Per-face lookup pixels in panorama pixel coordinates (cv2 convention:
# pixel centre = integer). Stack into (N, R, R, 2) before going to GPU.
pixel_lookups = np.empty((N, resolution, resolution, 2), dtype=np.float32)
for i in range(N):
spherical_uv = directions_to_spherical_uv(
utils3d.np.unproject_cv(
uv, np.ones_like(uv[..., 0]),
extrinsics=extrinsics[i], intrinsics=intrinsics[i],
)
)
pixel_lookups[i] = utils3d.np.uv_to_pixel(
spherical_uv, (height, width)
).astype(np.float32)
# Convert pixel coords -> grid_sample normalized coords. Use the
# half-pixel-centre convention so this lines up with cv2.remap's
# INTER_LINEAR.
grid = np.empty_like(pixel_lookups)
grid[..., 0] = (2.0 * pixel_lookups[..., 0] + 1.0) / float(width) - 1.0
grid[..., 1] = (2.0 * pixel_lookups[..., 1] + 1.0) / float(height) - 1.0
grid_t = torch.from_numpy(grid).to(device) # (N, R, R, 2) fp32
# Image to (1, C, H, W) → broadcast to (N, C, H, W) for batched sample.
img_chw = torch.from_numpy(np.ascontiguousarray(image)).to(device)
if img_chw.ndim == 2:
img_chw = img_chw.unsqueeze(-1) # (H, W) → (H, W, 1)
img_chw = img_chw.permute(2, 0, 1).contiguous().to(torch.float32) # (C, H, W)
img_batched = img_chw.unsqueeze(0).expand(N, -1, -1, -1) # (N, C, H, W)
sampled = torch.nn.functional.grid_sample(
img_batched, grid_t,
mode="bilinear", padding_mode="zeros", align_corners=False,
) # (N, C, R, R)
sampled = sampled.permute(0, 2, 3, 1).contiguous() # (N, R, R, C)
arr = sampled.cpu().numpy()
if image.dtype == np.uint8:
arr = np.clip(np.rint(arr), 0, 255).astype(np.uint8)
elif image.dtype == np.float32:
arr = arr.astype(np.float32)
else:
arr = arr.astype(image.dtype)
# Drop the C dim if input was 2-D (single-channel mask etc.)
if image.ndim == 2:
arr = arr[..., 0]
return [arr[i] for i in range(N)]
def poisson_equation(width: int, height: int, wrap_x: bool = False, wrap_y: bool = False) -> Tuple[csr_array, ndarray]:
grid_index = np.arange(height * width).reshape(height, width)
grid_index = np.pad(grid_index, ((0, 0), (1, 1)), mode='wrap' if wrap_x else 'edge')
grid_index = np.pad(grid_index, ((1, 1), (0, 0)), mode='wrap' if wrap_y else 'edge')
data = np.array([[-4, 1, 1, 1, 1]], dtype=np.float32).repeat(height * width, axis=0).reshape(-1)
indices = np.stack([
grid_index[1:-1, 1:-1],
grid_index[:-2, 1:-1], # up
grid_index[2:, 1:-1], # down
grid_index[1:-1, :-2], # left
grid_index[1:-1, 2:] # right
], axis=-1).reshape(-1)
indptr = np.arange(0, height * width * 5 + 1, 5)
A = csr_array((data, indices, indptr), shape=(height * width, height * width))
return A
def grad_equation(width: int, height: int, wrap_x: bool = False, wrap_y: bool = False) -> Tuple[csr_array, np.ndarray]:
grid_index = np.arange(width * height).reshape(height, width)
if wrap_x:
grid_index = np.pad(grid_index, ((0, 0), (0, 1)), mode='wrap')
if wrap_y:
grid_index = np.pad(grid_index, ((0, 1), (0, 0)), mode='wrap')
data = np.concatenate([
np.concatenate([
np.ones((grid_index.shape[0], grid_index.shape[1] - 1), dtype=np.float32).reshape(-1, 1), # x[i,j]
-np.ones((grid_index.shape[0], grid_index.shape[1] - 1), dtype=np.float32).reshape(-1, 1), # x[i,j-1]
], axis=1).reshape(-1),
np.concatenate([
np.ones((grid_index.shape[0] - 1, grid_index.shape[1]), dtype=np.float32).reshape(-1, 1), # x[i,j]
-np.ones((grid_index.shape[0] - 1, grid_index.shape[1]), dtype=np.float32).reshape(-1, 1), # x[i-1,j]
], axis=1).reshape(-1),
])
indices = np.concatenate([
np.concatenate([
grid_index[:, :-1].reshape(-1, 1),
grid_index[:, 1:].reshape(-1, 1),
], axis=1).reshape(-1),
np.concatenate([
grid_index[:-1, :].reshape(-1, 1),
grid_index[1:, :].reshape(-1, 1),
], axis=1).reshape(-1),
])
indptr = np.arange(0, grid_index.shape[0] * (grid_index.shape[1] - 1) * 2 + (grid_index.shape[0] - 1) * grid_index.shape[1] * 2 + 1, 2)
A = csr_array((data, indices, indptr), shape=(grid_index.shape[0] * (grid_index.shape[1] - 1) + (grid_index.shape[0] - 1) * grid_index.shape[1], height * width))
return A
def merge_panorama_depth(width: int, height: int, distance_maps: List[np.ndarray], pred_masks: List[np.ndarray], extrinsics: List[np.ndarray], intrinsics: List[np.ndarray], *, pbar=None, _top_size=None):
"""Recursive multi-resolution merge of N rectilinear face depth maps → equirect depth.
Optional `pbar` is a ComfyUI ProgressBar-like object (any obj with
`update_absolute(value, total)`). Progress is weighted by `pixels^1.5`
per level to match LSMR's asymptotic cost — the top level dominates,
so the bar advances slowly until the unwind reaches it.
"""
import sys
import time
def _p(msg):
print(f"[merge_panorama_depth {width:>4}x{height:<4}] {msg}", file=sys.stderr, flush=True)
# --- Progress-bar setup (one-time on top-level call) ---------------------
if _top_size is None:
_top_size = (width, height)
if pbar is not None and not hasattr(pbar, "_wn_levels"):
levels = []
w_, h_ = _top_size
while True:
levels.append((w_, h_))
if max(w_, h_) <= 256:
break
w_, h_ = w_ // 2, h_ // 2
levels.reverse() # smallest → largest, matches recursion unwind order
weights = np.array([(lw * lh) ** 1.5 for lw, lh in levels], dtype=np.float64)
weights = weights / weights.sum()
pbar._wn_levels = levels
pbar._wn_weights = weights
_p(f"progress bar: {len(levels)} levels — weights {[f'{x:.1%}' for x in weights]}")
def _bump(level_frac: float):
"""Advance the global progress bar to (cum % of prior levels) + (level_frac × this level's share)."""
if pbar is None or not hasattr(pbar, "_wn_levels"):
return
try:
i = pbar._wn_levels.index((width, height))
except ValueError:
return
pct_floor = float(pbar._wn_weights[:i].sum())
pct_ceil = float(pbar._wn_weights[:i + 1].sum())
pct = pct_floor + (pct_ceil - pct_floor) * max(0.0, min(1.0, level_frac))
pbar.update_absolute(int(min(99, pct * 100)), 100)
t_lvl_start = time.perf_counter()
_bump(0.0)
if max(width, height) > 256:
_p(f"recurse into {width // 2}x{height // 2}…")
panorama_depth_init, _ = merge_panorama_depth(
width // 2, height // 2, distance_maps, pred_masks, extrinsics, intrinsics,
pbar=pbar, _top_size=_top_size,
)
_p(f"recursion returned; upsampling {width // 2}x{height // 2} → {width}x{height}")
panorama_depth_init = cv2.resize(panorama_depth_init, (width, height), cv2.INTER_LINEAR)
else:
_p(f"base case (≤256 on each side) → no recursion, no init x0 for LSMR")
panorama_depth_init = None
_bump(0.05)
_p(f"building sphere uv + ray directions ({width}×{height})…")
uv = utils3d.np.uv_map(height, width)
spherical_directions = spherical_uv_to_directions(uv)
# Warp each view to the panorama
_p(f"warping {len(distance_maps)} face depth maps onto the equirect…")
t_warp_start = time.perf_counter()
panorama_log_distance_grad_maps, panorama_grad_masks = [], []
panorama_log_distance_laplacian_maps, panorama_laplacian_masks = [], []
panorama_pred_masks = []
for i in range(len(distance_maps)):
projected_uv, projected_depth = utils3d.np.project_cv(spherical_directions, extrinsics=extrinsics[i], intrinsics=intrinsics[i])
projection_valid_mask = (projected_depth > 0) & (projected_uv > 0).all(axis=-1) & (projected_uv < 1).all(axis=-1)
projected_pixels = utils3d.np.uv_to_pixel(np.clip(projected_uv, 0, 1), distance_maps[i].shape).astype(np.float32)
log_splitted_distance = np.log(distance_maps[i])
panorama_log_distance_map = np.where(projection_valid_mask, cv2.remap(log_splitted_distance, projected_pixels[..., 0], projected_pixels[..., 1], cv2.INTER_LINEAR, borderMode=cv2.BORDER_REPLICATE), 0)
panorama_pred_mask = projection_valid_mask & (cv2.remap(pred_masks[i].astype(np.uint8), projected_pixels[..., 0], projected_pixels[..., 1], cv2.INTER_NEAREST, borderMode=cv2.BORDER_REPLICATE) > 0)
# calculate gradient map
padded = np.pad(panorama_log_distance_map, ((0, 0), (0, 1)), mode='wrap')
grad_x, grad_y = padded[:, :-1] - padded[:, 1:], padded[:-1, :] - padded[1:, :]
padded = np.pad(panorama_pred_mask, ((0, 0), (0, 1)), mode='wrap')
mask_x, mask_y = padded[:, :-1] & padded[:, 1:], padded[:-1, :] & padded[1:, :]
panorama_log_distance_grad_maps.append((grad_x, grad_y))
panorama_grad_masks.append((mask_x, mask_y))
# calculate laplacian map
padded = np.pad(panorama_log_distance_map, ((1, 1), (0, 0)), mode='edge')
padded = np.pad(padded, ((0, 0), (1, 1)), mode='wrap')
laplacian = convolve(padded, np.array([[0, 1, 0], [1, -4, 1], [0, 1, 0]], dtype=np.float32))[1:-1, 1:-1]
padded = np.pad(panorama_pred_mask, ((1, 1), (0, 0)), mode='edge')
padded = np.pad(padded, ((0, 0), (1, 1)), mode='wrap')
mask = convolve(padded.astype(np.uint8), np.array([[0, 1, 0], [1, 1, 1], [0, 1, 0]], dtype=np.uint8))[1:-1, 1:-1] == 5
panorama_log_distance_laplacian_maps.append(laplacian)
panorama_laplacian_masks.append(mask)
panorama_pred_masks.append(panorama_pred_mask)
if (i + 1) % 10 == 0 or i == len(distance_maps) - 1:
_p(f" warped {i + 1}/{len(distance_maps)} faces")
_bump(0.45)
_p(f"warp stage done in {time.perf_counter() - t_warp_start:.2f}s; "
f"stacking + averaging gradient/laplacian maps…")
panorama_log_distance_grad_x = np.stack([grad_map[0] for grad_map in panorama_log_distance_grad_maps], axis=0)
panorama_log_distance_grad_y = np.stack([grad_map[1] for grad_map in panorama_log_distance_grad_maps], axis=0)
panorama_grad_mask_x = np.stack([mask_map[0] for mask_map in panorama_grad_masks], axis=0)
panorama_grad_mask_y = np.stack([mask_map[1] for mask_map in panorama_grad_masks], axis=0)
panorama_log_distance_grad_x = np.sum(panorama_log_distance_grad_x * panorama_grad_mask_x, axis=0) / np.sum(panorama_grad_mask_x, axis=0).clip(1e-3)
panorama_log_distance_grad_y = np.sum(panorama_log_distance_grad_y * panorama_grad_mask_y, axis=0) / np.sum(panorama_grad_mask_y, axis=0).clip(1e-3)
panorama_laplacian_maps = np.stack(panorama_log_distance_laplacian_maps, axis=0)
panorama_laplacian_masks = np.stack(panorama_laplacian_masks, axis=0)
panorama_laplacian_map = np.sum(panorama_laplacian_maps * panorama_laplacian_masks, axis=0) / np.sum(panorama_laplacian_masks, axis=0).clip(1e-3)
grad_x_mask = np.any(panorama_grad_mask_x, axis=0).reshape(-1)
grad_y_mask = np.any(panorama_grad_mask_y, axis=0).reshape(-1)
grad_mask = np.concatenate([grad_x_mask, grad_y_mask])
laplacian_mask = np.any(panorama_laplacian_masks, axis=0).reshape(-1)
_bump(0.55)
# Solve overdetermined system
_p(f"building sparse linear system (grad_equation + poisson_equation)…")
t_build = time.perf_counter()
A = vstack([
grad_equation(width, height, wrap_x=True, wrap_y=False)[grad_mask],
poisson_equation(width, height, wrap_x=True, wrap_y=False)[laplacian_mask],
])
b = np.concatenate([
panorama_log_distance_grad_x.reshape(-1)[grad_x_mask],
panorama_log_distance_grad_y.reshape(-1)[grad_y_mask],
panorama_laplacian_map.reshape(-1)[laplacian_mask]
])
_p(f"linear system: A {A.shape} ({A.nnz} nnz), b {b.shape} — built in "
f"{time.perf_counter() - t_build:.2f}s")
_bump(0.65)
_p(f"LSMR solve (atol=1e-5, btol=1e-5, "
f"x0={'provided' if panorama_depth_init is not None else 'None'})…")
t_solve = time.perf_counter()
x, istop, itn, normr, normar, norma, conda, normx = lsmr(
A, b,
atol=1e-5, btol=1e-5,
x0=np.log(panorama_depth_init).reshape(-1) if panorama_depth_init is not None else None,
show=False,
)
_p(f"LSMR done in {time.perf_counter() - t_solve:.2f}s: "
f"{itn} iterations, residual={normr:.4f}, istop={istop}")
_bump(0.95)
panorama_depth = np.exp(x).reshape(height, width).astype(np.float32)
panorama_mask = np.any(panorama_pred_masks, axis=0)
_p(f"level done in {time.perf_counter() - t_lvl_start:.2f}s; "
f"depth: min={float(panorama_depth.min()):.3f}, "
f"median={float(np.median(panorama_depth)):.3f}, "
f"max={float(panorama_depth.max()):.3f}; "
f"valid pixels: {int(panorama_mask.sum())}/{panorama_mask.size}")
_bump(1.0)
return panorama_depth, panorama_mask
nodes/_vendor/worldgen/__init__.py
View File
nodes/_vendor/worldgen/src/__init__.py
View File
nodes/_vendor/worldgen/src/panorama_utils.py
+2074
View File
File diff suppressed because it is too large Load Diff
nodes/comfy-env.toml
+10 -1
View File
@@ -11,4 +11,13 @@ safetensors = "*"
einops = "*"
comfy-aimdo = "*"
scipy = "*"
transformers = "*"
transformers = "*"
tqdm = "*"
opencv-python = "*"
# utils3d: used by the vendored MoGe LSMR depth-merge solver
# (_vendor/worldgen/src/panorama_utils.py + _vendor/moge_panorama.py).
# Same pin as WorldNav/WorldStereo so behavior is identical.
[pypi-dependencies.utils3d]
git = "https://github.com/EasternJournalist/utils3d.git"
rev = "3fab839f0be9931dac7c8488eb0e1600c236e183"
nodes/depth_merge.py
+275
View File
@@ -0,0 +1,275 @@
"""SharpDepthMerge node for ComfyUI-Sharp.
Per-face rectilinear depth maps -> a single equirectangular distance map,
via MoGe's sparse-LSMR formulation in log-distance space (Laplacian +
gradient terms). 1:1 schema mirror of WorldNavDepthMerge — wiring is
identical so workflows port over by editing only the node-class name.
Backed by the same solver as WorldNav: vendored at
`nodes/_vendor/moge_panorama.py` (CPU LSMR via scipy) + vendored at
`nodes/_vendor/worldgen/src/panorama_utils.py` (GPU LSMR via
solve_lsmr_gpu + post-merge sky inpaint + south-pole smooth).
"""
from __future__ import annotations
import logging
import sys
import numpy as np
import torch
from comfy_api.latest import io
log = logging.getLogger("sharp")
def _p(msg: str) -> None:
print(f"[SharpDepthMerge] {msg}", file=sys.stderr, flush=True)
class SharpDepthMerge(io.ComfyNode):
"""Per-face depth maps + per-face extrinsics/intrinsics -> equirect depth IMAGE.
Schema is a 1:1 mirror of WorldNavDepthMerge so workflows port over
without rewiring. Internally uses the same vendored MoGe LSMR solver.
"""
@classmethod
def define_schema(cls):
return io.Schema(
node_id="SharpDepthMerge",
display_name="Sharp Depth Merge",
category="SHARP",
description=(
"Stitch per-face rectilinear depth maps from a panorama "
"split (SharpPanoramaIcosahedronSplit etc.) back into a "
"single equirect distance map. Sparse-LSMR formulation in "
"log-distance space (Laplacian + gradient terms in MoGe "
"style). 1:1 schema mirror of WorldNavDepthMerge."
),
inputs=[
io.Image.Input(
"face_depths",
tooltip="Per-face depth maps from a rectilinear depth "
"model (MoGe2 / SHARP / etc.). Shape "
"(N, h, w, C); only the first channel is used "
"as the distance value."),
io.Mask.Input(
"face_valid_masks",
optional=True,
tooltip="Per-face valid masks (typically the depth "
"model's valid_mask output). Shape (N, h, w) "
"bool/float. Load-bearing: upstream "
"merge_panorama_depth uses this to gate which "
"pixels enter the gradient + laplacian LSMR "
"system. Without it, sky / invalid pixels "
"corrupt the solve. If unwired, all pixels "
"are treated as valid (degraded fallback)."),
io.Custom("EXTRINSICS").Input(
"extrinsics",
tooltip="From the panorama-split node. (N, 4, 4) "
"world-to-camera per face."),
io.Custom("INTRINSICS").Input(
"intrinsics",
tooltip="From the panorama-split node. (N, 3, 3) "
"pixel-space per face."),
io.Int.Input(
"out_width", default=1920, min=512, max=4096, step=64,
tooltip="Output equirect width. 2:1 aspect (height = "
"width / 2 is the convention; provide "
"out_height to override)."),
io.Int.Input(
"out_height", default=960, min=256, max=2048, step=64,
tooltip="Output equirect height. Should be width/2 "
"for 2:1 aspect."),
io.Boolean.Input(
"use_gpu", default=True,
tooltip="DEPRECATED. Kept for back-compat. Equivalent "
"to strategy='cpu_fallback' when False. "
"Otherwise respects strategy."),
io.Combo.Input(
"strategy",
options=["auto", "full", "vcycle", "cpu_fallback"],
default="auto",
tooltip=(
"Solver strategy for the equirect depth output:\n"
" auto: full if out<=1920x960 else vcycle (recommended).\n"
" full: solve LSMR at out_resolution (original). OOMs above ~2 Mpix on 24 GB.\n"
" vcycle: solve LSMR at 1920x960 then bilinear-upsample depth+mask "
"to out_resolution. ~4x faster + fits memory at 3840x1920. Per-face "
"info saturates at 1920x960 so quality is indistinguishable.\n"
" cpu_fallback: scipy LSMR on CPU end-to-end. Slow (~30-60s at "
"1920x960) but works at any resolution."
)),
],
outputs=[
io.Image.Output(display_name="depth"),
io.Mask.Output(display_name="valid_mask"),
],
)
@classmethod
def execute(cls, face_depths, extrinsics, intrinsics,
face_valid_masks=None, out_width=1920, out_height=960,
use_gpu=True, strategy="auto"):
import cv2
# Vendored MoGe LSMR solvers — imported here so depth_merge.py
# loads cleanly even if scipy / utils3d / etc. aren't installed.
# (They ARE in Sharp's pixi env per comfy-env.toml; this lazy
# import keeps the node-class registration unconditional.)
from ._vendor.moge_panorama import merge_panorama_depth
from ._vendor.worldgen.src.panorama_utils import (
merge_panorama_depth_gpu,
smooth_south_pole_depth,
)
# --- face_depths: (N, h, w, C) → list of (h, w) float ---
d = (face_depths.detach().cpu().numpy()
if isinstance(face_depths, torch.Tensor) else np.asarray(face_depths))
if d.ndim != 4:
raise ValueError(
f"SharpDepthMerge: face_depths must be (N, h, w, C), got {d.shape}"
)
N, fh, fw, _ = d.shape
distance_maps = [d[i, ..., 0].astype(np.float32) for i in range(N)]
# --- extrinsics / intrinsics: (N, 4, 4) and (N, 3, 3) ---
ex = (extrinsics.detach().cpu().numpy()
if isinstance(extrinsics, torch.Tensor) else np.asarray(extrinsics))
intr = (intrinsics.detach().cpu().numpy()
if isinstance(intrinsics, torch.Tensor) else np.asarray(intrinsics))
if ex.shape != (N, 4, 4) or intr.shape != (N, 3, 3):
raise ValueError(
f"SharpDepthMerge: shape mismatch — extrinsics {ex.shape}, "
f"intrinsics {intr.shape}, expected ({N},4,4) and ({N},3,3)"
)
extr_list = [ex[i].astype(np.float32) for i in range(N)]
intr_list = [intr[i].astype(np.float32) for i in range(N)]
# --- valid_masks: (N, h, w, C) or None → list of (h, w) bool ---
if face_valid_masks is None:
pred_masks = [np.ones((fh, fw), dtype=bool) for _ in range(N)]
else:
v = (face_valid_masks.detach().cpu().numpy()
if isinstance(face_valid_masks, torch.Tensor)
else np.asarray(face_valid_masks))
if v.ndim == 4:
v = v[..., 0]
if v.shape != (N, fh, fw):
raise ValueError(
f"SharpDepthMerge: face_valid_masks shape {v.shape} "
f"doesn't match face_depths ({N}, {fh}, {fw})"
)
pred_masks = [(v[i] > 0.5).astype(bool) for i in range(N)]
_p(f"merging {N} faces @ {fh}x{fw} -> equirect {out_height}x{out_width} "
f"(use_gpu={use_gpu})")
# Progress bar — visible in ComfyUI if available.
try:
import comfy.utils
pbar = comfy.utils.ProgressBar(100)
except Exception:
class _Noop:
def update_absolute(self, *_a, **_k):
pass
pbar = _Noop()
pbar.update_absolute(1, 100)
# Resolve strategy. use_gpu=False overrides → cpu_fallback (back-compat
# for the equivalent toggle in WorldNavDepthMerge).
out_w, out_h = int(out_width), int(out_height)
resolved = strategy
if not use_gpu:
resolved = "cpu_fallback"
if resolved == "auto":
resolved = "vcycle" if (out_w * out_h > 1920 * 960) else "full"
_p(f"resolved strategy={resolved} (requested={strategy}, "
f"use_gpu={use_gpu}, out={out_w}x{out_h})")
if resolved == "full":
# Recursive LSMR all the way up to out_resolution. Reference;
# OOMs above ~2 Mpix on 24 GB cards.
depth_np, mask_np = merge_panorama_depth_gpu(
out_w, out_h,
distance_maps, pred_masks,
extr_list, intr_list,
)
elif resolved == "vcycle":
# V-cycle: solve the recursive pyramid up to 1920x960 (the
# information-content ceiling for typical 42x ~512² per-face
# data), then bilinear-upsample depth + nearest-upsample mask
# to out_resolution. Skips the OOM-prone outermost LSMR.
depth_low, mask_low = merge_panorama_depth_gpu(
1920, 960,
distance_maps, pred_masks,
extr_list, intr_list,
)
if (out_w, out_h) == (1920, 960):
depth_np, mask_np = depth_low, mask_low
else:
depth_np = cv2.resize(
depth_low.astype(np.float32),
(out_w, out_h),
interpolation=cv2.INTER_LINEAR,
)
mask_np = cv2.resize(
mask_low.astype(np.uint8),
(out_w, out_h),
interpolation=cv2.INTER_NEAREST,
).astype(bool)
_p(f"vcycle: upsampled 1920x960 -> {out_w}x{out_h} "
f"(bilinear depth, nearest mask)")
elif resolved == "cpu_fallback":
# scipy LSMR end-to-end on CPU. Slow but no VRAM ceiling.
depth_np, mask_np = merge_panorama_depth(
out_w, out_h,
distance_maps, pred_masks,
extr_list, intr_list,
pbar=pbar,
)
else:
raise ValueError(f"SharpDepthMerge: unknown strategy {resolved!r}")
# Match upstream HY-World 2.0 post-merge processing
# (panorama_utils.py:572-580). Two steps:
# 1. Inpaint invalid pixels with the 99.9th-percentile (sky)
# depth instead of zero. Zero-depth puts those pixels right
# on the camera origin and breaks downstream point-cloud /
# mesh geometry.
# 2. Smooth the south-pole strip (bottom 5%) to fix left-right
# seams the LSMR solver can't enforce due to the equirect
# parameterization stretching infinitely at the pole.
depth_np = depth_np.astype(np.float32)
valid = mask_np.astype(bool)
if valid.any() and (~valid).any():
sky_depth = float(np.nanquantile(depth_np[valid], 0.999))
depth_np[~valid] = sky_depth
_p(f" inpainted {int((~valid).sum())} invalid pixels with "
f"sky_depth={sky_depth:.3f}")
depth_np = smooth_south_pole_depth(depth_np, smooth_height_ratio=0.05)
# Belt-and-braces: trap any residual non-finite values that escaped.
depth_np = np.nan_to_num(depth_np, nan=0.0, posinf=0.0, neginf=0.0)
_p(f"merged depth: shape {depth_np.shape}, min={float(depth_np.min()):.3f}, "
f"median={float(np.median(depth_np)):.3f}, max={float(depth_np.max()):.3f}; "
f"valid pixels: {int(mask_np.sum())} / {mask_np.size}")
# IMAGE convention: (B, H, W, C). Broadcast depth across 3 channels so
# it composes with regular depth-viz nodes.
depth_img = (
torch.from_numpy(depth_np)
.unsqueeze(-1).expand(-1, -1, 3).unsqueeze(0).contiguous()
)
valid_mask = torch.from_numpy(mask_np.astype(np.float32)).unsqueeze(0)
pbar.update_absolute(100, 100)
return io.NodeOutput(depth_img, valid_mask)
NODE_CLASS_MAPPINGS = {
"SharpDepthMerge": SharpDepthMerge,
}
NODE_DISPLAY_NAME_MAPPINGS = {
"SharpDepthMerge": "Sharp Depth Merge",
}
nodes/panorama_cube_split.py
+306
View File
@@ -0,0 +1,306 @@
"""SharpPanoramaCubeSplit node for ComfyUI-Sharp.
Equirect panorama -> N perspective crops at a fixed pitch x yaw grid.
Parallel to WorldStereo's WorldStereoPanoramaCubeSplit:
3 pitch levels x N yaw steps at configurable rectangular FOV / aspect.
Default: 3 x 9 = 27 views at 832x480, 120 deg horizontal x 90 deg
vertical FOV.
Different from Sharp's existing SamplePanorama (sample_panorama.py):
- SamplePanorama: square output, FOV-derived dense grid, blends to
cover the full sphere — flexible for arbitrary FOV/overlap.
- SharpPanoramaCubeSplit: fixed 3-pitch grid, rectangular aspect,
matches WorldStereo's anchor-view geometry — for feeding a
rectilinear depth model (MoGe2 etc.) before merging back to equirect
via SharpDepthMerge.
Output convention matches the rest of ComfyUI-Sharp:
face_images (IMAGE [N, H, W, 3]) + extrinsics (EXTRINSICS [N, 4, 4]) +
intrinsics (INTRINSICS [N, 3, 3]).
"""
import logging
import math
import numpy as np
import torch
import torch.nn.functional as F
from comfy_api.latest import io
log = logging.getLogger("sharp")
def _rot_around_z(deg: float) -> torch.Tensor:
"""3x3 rotation matrix around the Z axis. World up = +Z convention.
Matches WorldStereo's `rotate_around_z_axis` in panorama_utils.py so
the pitch-up / pitch-down hemispheres land at the same world points.
"""
rad = math.radians(deg)
c, s = math.cos(rad), math.sin(rad)
return torch.tensor([
[c, -s, 0],
[s, c, 0],
[0, 0, 1],
], dtype=torch.float32)
def _look_at_w2c(eye: torch.Tensor, target: torch.Tensor, up: torch.Tensor) -> torch.Tensor:
"""world-to-camera 4x4 matrix from a look-at triple (numpy-free).
Convention: camera looks down +Z in camera space (matches utils3d
extrinsics_look_at). Returns a w2c matrix.
"""
# forward (in world): from eye to target, normalized
fwd = (target - eye)
fwd = fwd / (fwd.norm() + 1e-12)
# right = forward x up
right = torch.linalg.cross(fwd, up)
right = right / (right.norm() + 1e-12)
# recomputed up = right x forward
new_up = torch.linalg.cross(right, fwd)
# Camera-to-world rotation: columns are the camera basis in world coords.
# We want world-to-camera, which is the transpose for pure rotation.
# In camera space: +X right, +Y down (image convention), +Z forward.
# So R_c2w has columns = [right, -new_up, fwd]; w2c = transpose.
R_c2w = torch.stack([right, -new_up, fwd], dim=1) # [3, 3]
R_w2c = R_c2w.T
ext = torch.eye(4, dtype=torch.float32)
ext[:3, :3] = R_w2c
ext[:3, 3] = -R_w2c @ eye
return ext
def _intrinsics_from_fov(fov_x_rad: float, fov_y_rad: float, w: int, h: int) -> torch.Tensor:
"""Pixel-space 3x3 intrinsics from horizontal + vertical FOV."""
fx = (w / 2.0) / math.tan(fov_x_rad / 2.0)
fy = (h / 2.0) / math.tan(fov_y_rad / 2.0)
cx = (w - 1) / 2.0
cy = (h - 1) / 2.0
K = torch.tensor([
[fx, 0, cx],
[0, fy, cy],
[0, 0, 1.0],
], dtype=torch.float32)
return K
def _sample_perspective_from_equirect(
panorama: torch.Tensor,
ext_w2c: torch.Tensor,
K_pixel: torch.Tensor,
out_h: int,
out_w: int,
) -> torch.Tensor:
"""Backward warp: equirect panorama -> a single perspective view.
Mirrors sample_panorama.py's geometry but with rectangular output
and an explicit extrinsics/intrinsics input (so the same look-at
targets used to build ext_w2c here can be passed downstream
unchanged).
"""
H_pano, W_pano, _ = panorama.shape
device = panorama.device
# Pixel grid in the output perspective image.
u = torch.arange(out_w, dtype=torch.float32, device=device)
v = torch.arange(out_h, dtype=torch.float32, device=device)
uu, vv = torch.meshgrid(u, v, indexing='xy')
fx, fy = K_pixel[0, 0].item(), K_pixel[1, 1].item()
cx, cy = K_pixel[0, 2].item(), K_pixel[1, 2].item()
# Camera-space rays for each output pixel.
dx = (uu - cx) / fx
dy = (vv - cy) / fy
dz = torch.ones_like(dx)
rays_cam = torch.stack([dx, dy, dz], dim=-1) # [H, W, 3]
rays_cam = F.normalize(rays_cam, dim=-1)
# World rays via camera-to-world rotation.
R_w2c = ext_w2c[:3, :3].to(device)
R_c2w = R_w2c.T
rays_world = torch.einsum('ij,hwj->hwi', R_c2w, rays_cam) # [H, W, 3]
# World rays -> equirect sample coords.
# WorldStereo / HY-World convention: world up = +Z, panorama columns
# = yaw around Z. Match that here so this node's extrinsics line up
# with WorldStereo's pano_bank geometry.
rx, ry, rz = rays_world[..., 0], rays_world[..., 1], rays_world[..., 2]
yaw = torch.atan2(ry, rx) # [-pi, pi]
pitch = torch.asin(torch.clamp(rz, -1.0, 1.0)) # [-pi/2, pi/2]
eq_x = (yaw / math.pi + 1.0) * (W_pano - 1) / 2.0 # 0..W-1
eq_y = (0.5 - pitch / math.pi) * (H_pano - 1) # 0..H-1
grid_x = eq_x / (W_pano - 1) * 2.0 - 1.0
grid_y = eq_y / (H_pano - 1) * 2.0 - 1.0
grid = torch.stack([grid_x, grid_y], dim=-1).unsqueeze(0) # [1, H, W, 2]
pano_nchw = panorama.permute(2, 0, 1).unsqueeze(0) # [1, C, H, W]
sampled = F.grid_sample(
pano_nchw, grid,
mode='bilinear', padding_mode='border', align_corners=True,
)
return sampled[0].permute(1, 2, 0) # [H, W, C]
class SharpPanoramaCubeSplit(io.ComfyNode):
"""Equirect panorama -> N perspective crops at a fixed pitch x yaw grid."""
@classmethod
def define_schema(cls):
return io.Schema(
node_id="SharpPanoramaCubeSplit",
display_name="Sharp Panorama Cube Split",
category="SHARP",
description=(
"Sample N perspective crops from a 2:1 equirectangular panorama "
"at a fixed pitch x yaw grid (3 pitch x 9 yaw = 27 views by "
"default). Matches WorldStereoPanoramaCubeSplit's geometry — "
"useful as the input stage for a per-face depth model whose "
"outputs will be stitched back via SharpDepthMerge.\n\n"
"Different from Sharp's SamplePanorama (which derives the grid "
"from FOV+overlap, square output): this node has a fixed pitch "
"grid with rectangular aspect (default 832x480), and matches "
"the camera/intrinsics convention WorldStereo's pano_bank uses."
),
inputs=[
io.Image.Input(
"panorama",
tooltip="Equirectangular RGB panorama, 2:1 aspect. "
"[B, H, W, 3] or [H, W, 3] in [0, 1]. First image "
"of batch is used."),
io.Int.Input(
"image_w", default=832, min=224, max=2048, step=16,
tooltip="Crop width. 832 matches WorldStereo's anchor width."),
io.Int.Input(
"image_h", default=480, min=224, max=2048, step=16,
tooltip="Crop height. 480 matches WorldStereo's anchor height."),
io.Float.Input(
"fov_x_deg", default=120.0, min=30.0, max=170.0, step=1.0,
tooltip="Horizontal FOV. 120 matches WorldStereo upstream."),
io.Float.Input(
"fov_y_deg", default=90.0, min=30.0, max=170.0, step=1.0,
tooltip="Vertical FOV. 90 matches WorldStereo upstream."),
io.Float.Input(
"rot_deg", default=40.0, min=10.0, max=120.0, step=5.0,
tooltip="Yaw step between adjacent views. 40 -> 9 views "
"per pitch -> 27 total at 3 pitches."),
io.Float.Input(
"pitch_up", default=0.5, min=0.0, max=1.0, step=0.05,
tooltip="Z component of the upper-hemisphere look-at "
"targets. 0.5 matches WorldStereo upstream."),
io.Float.Input(
"pitch_down", default=-0.5, min=-1.0, max=0.0, step=0.05,
tooltip="Z component of the lower-hemisphere look-at "
"targets. -0.5 matches WorldStereo upstream."),
],
outputs=[
io.Image.Output(
display_name="face_images",
tooltip="N perspective crops [N, H, W, 3] in [0, 1]. "
"Wire to a depth model (SharpPredictDepth / MoGe2 "
"/ etc.)."),
io.Custom("EXTRINSICS").Output(
display_name="extrinsics",
tooltip="[N, 4, 4] world-to-camera matrices per face. "
"Camera at world origin, look-at along the pitch/"
"yaw direction. Pass to SharpDepthMerge."),
io.Custom("INTRINSICS").Output(
display_name="intrinsics",
tooltip="[N, 3, 3] pixel-space intrinsics per face. "
"Identical across faces (same FOV + size). Pass "
"to SharpDepthMerge."),
io.Float.Output(
display_name="fov_x_deg",
tooltip="Pass-through. Wire into MoGe2 / depth model's "
"fov_x widget so it doesn't have to estimate."),
io.Int.Output(
display_name="num_faces",
tooltip="N = 3 * (360 / rot_deg). 27 at defaults."),
],
)
@classmethod
def execute(
cls, panorama: torch.Tensor,
image_w: int = 832, image_h: int = 480,
fov_x_deg: float = 120.0, fov_y_deg: float = 90.0,
rot_deg: float = 40.0,
pitch_up: float = 0.5, pitch_down: float = -0.5,
):
# Normalize input shape -> [H, W, 3], float in [0, 1].
if panorama.dim() == 4:
panorama = panorama[0]
if panorama.dtype != torch.float32:
panorama = panorama.float()
if panorama.shape[-1] == 4:
panorama = panorama[..., :3]
# Sharp convention is float [0,1]; guard against accidental uint8 input.
if panorama.max() > 2.0:
panorama = panorama / 255.0
device = panorama.device
H_pano, W_pano, _ = panorama.shape
# Build the (3 pitch x N yaw) direction grid.
# WorldStereo convention: world up = +Z, base directions look at +X
# rotated around Z. (Eye is world origin; targets are unit-ish
# points on a sphere.)
# NOTE: WorldStereo uses [-1, 0, 0] as the base "forward" with +Z up;
# we mirror that exactly so panorama columns line up the same way.
base_targets = [
torch.tensor([-1.0, 0.0, 0.0], dtype=torch.float32),
torch.tensor([-1.0, 0.0, float(pitch_up)], dtype=torch.float32),
torch.tensor([-1.0, 0.0, float(pitch_down)], dtype=torch.float32),
]
N_yaw = max(1, int(round(360.0 / max(1.0, float(rot_deg)))))
targets = []
for base in base_targets:
targets.append(base)
for i in range(1, N_yaw):
R = _rot_around_z(rot_deg * i)
targets.append((R @ base).contiguous())
targets = torch.stack(targets, dim=0) # [N, 3]
N = targets.shape[0]
# Same intrinsics for every face (same FOV + same output size).
fov_x_rad = math.radians(fov_x_deg)
fov_y_rad = math.radians(fov_y_deg)
K = _intrinsics_from_fov(fov_x_rad, fov_y_rad, int(image_w), int(image_h))
# Sample each face.
eye = torch.zeros(3, dtype=torch.float32)
up = torch.tensor([0.0, 0.0, 1.0], dtype=torch.float32)
face_images = []
extrinsics = []
for i in range(N):
ext_w2c = _look_at_w2c(eye, targets[i], up)
img = _sample_perspective_from_equirect(
panorama, ext_w2c, K, int(image_h), int(image_w),
)
face_images.append(img)
extrinsics.append(ext_w2c)
face_images_t = torch.stack(face_images, dim=0).contiguous() # [N, H, W, 3]
ext_t = torch.stack(extrinsics, dim=0).contiguous() # [N, 4, 4]
intr_t = K.unsqueeze(0).expand(N, -1, -1).contiguous() # [N, 3, 3]
log.info(
f"[SharpPanoramaCubeSplit] {N} faces ({N_yaw}/pitch x 3 pitches) "
f"@ {image_w}x{image_h}, fov=({fov_x_deg:.0f},{fov_y_deg:.0f})"
)
return io.NodeOutput(face_images_t, ext_t, intr_t, float(fov_x_deg), int(N))
NODE_CLASS_MAPPINGS = {
"SharpPanoramaCubeSplit": SharpPanoramaCubeSplit,
}
NODE_DISPLAY_NAME_MAPPINGS = {
"SharpPanoramaCubeSplit": "Sharp Panorama Cube Split",
}
nodes/panorama_icosahedron_split.py
+296
View File
@@ -0,0 +1,296 @@
"""SharpPanoramaIcosahedronSplit node for ComfyUI-Sharp.
Equirect panorama -> N square perspective crops at icosahedron vertices,
90° FOV each. Schema mirrors WorldNavPanoramaSplit 1:1 (modulo
`panorama` being plain IMAGE here instead of WorldNav's
WORLDSTEREO_PANORAMA custom type) so workflows port over by editing
only the node-class name.
Geometry:
- icosahedron_12: 12 base vertices via golden-ratio coords
- icosahedron_42: subdivide once (add edge midpoints, normalize to
sphere) -> 12 + 30 = 42 vertices
Why square / icosahedron for SHARP: SHARP's internal resolution is
1536x1536 square and was trained on cube-face-like inputs. The
icosahedron tiling gives uniform sphere coverage (no pole over/under-
sampling like a pitch/yaw grid) at 90° FOV which matches cube-face
geometry exactly.
Implementation is pure numpy + torch. Reuses three helpers from
`panorama_cube_split.py` (`_look_at_w2c`, `_intrinsics_from_fov`,
`_sample_perspective_from_equirect`) so no math is duplicated.
Pole-singularity fix (mirrors WorldNav's): two vertices of the
subdivided icosahedron land at (0, 0, ±1). With the default world-up
of +Z, `forward × up = 0` for those, producing NaN extrinsics. We
detect `|forward · up| > 0.999` per vertex and swap to up=+Y just for
those.
"""
from __future__ import annotations
import logging
import math
import numpy as np
import torch
import torch.nn.functional as F
from comfy_api.latest import io
# Re-use the math helpers from the sibling cube-split node — no copy.
from .panorama_cube_split import (
_look_at_w2c,
_intrinsics_from_fov,
_sample_perspective_from_equirect,
)
log = logging.getLogger("sharp")
def _icosahedron_vertices() -> np.ndarray:
"""12 unit-sphere vertices of a regular icosahedron.
Standard golden-ratio coords. After normalizing to the unit sphere,
edge length squared between adjacent vertices = (2 / phi) ≈ 1.236.
"""
phi = (1.0 + math.sqrt(5.0)) / 2.0
verts = np.array(
[
( 0, +1, +phi), ( 0, +1, -phi), ( 0, -1, +phi), ( 0, -1, -phi),
(+1, +phi, 0), (+1, -phi, 0), (-1, +phi, 0), (-1, -phi, 0),
(+phi, 0, +1), (+phi, 0, -1), (-phi, 0, +1), (-phi, 0, -1),
],
dtype=np.float32,
)
verts /= np.linalg.norm(verts, axis=1, keepdims=True)
return verts
def _icosahedron_edges(verts: np.ndarray) -> np.ndarray:
"""30 edge (i, j) pairs of the regular icosahedron.
Derived by picking the 30 shortest pairwise distances out of C(12,2)=66
— equivalently, the pairs with squared distance ≈ 2/phi after the
unit-sphere normalization. No hardcoded indices, so this stays correct
regardless of the vertex ordering in `_icosahedron_vertices`.
"""
N = len(verts)
pairs = []
for i in range(N):
for j in range(i + 1, N):
d2 = float(((verts[i] - verts[j]) ** 2).sum())
pairs.append((d2, i, j))
pairs.sort()
# The icosahedron has 30 edges; threshold by taking exactly the 30
# closest pairs. Sanity-check: the gap between edge-distance and
# non-edge-distance should be large (>1.5x).
edge_d2 = pairs[29][0]
non_edge_d2 = pairs[30][0]
assert non_edge_d2 > edge_d2 * 1.5, (
f"icosahedron edge classification ambiguous: "
f"edge_d2={edge_d2:.4f}, non_edge_d2={non_edge_d2:.4f}"
)
return np.array([(i, j) for _, i, j in pairs[:30]], dtype=np.int32)
def _subdivide_icosahedron(verts: np.ndarray, edges: np.ndarray) -> np.ndarray:
"""Add edge midpoints (normalized to the unit sphere) → 42 vertices."""
midpoints = (verts[edges[:, 0]] + verts[edges[:, 1]]) / 2.0
midpoints /= np.linalg.norm(midpoints, axis=1, keepdims=True)
return np.concatenate([verts, midpoints.astype(np.float32)], axis=0)
def _icosahedron_split_vertices(subdivision: str) -> np.ndarray:
"""Return [N, 3] unit-sphere target points for the chosen subdivision."""
base = _icosahedron_vertices()
if subdivision == "icosahedron_12":
return base
if subdivision == "icosahedron_42":
edges = _icosahedron_edges(base)
return _subdivide_icosahedron(base, edges)
raise ValueError(f"unknown subdivision={subdivision!r}")
def _build_extrinsics_with_pole_fix(targets: torch.Tensor) -> torch.Tensor:
"""Build per-vertex world-to-camera matrices, swapping up at poles.
`targets`: [N, 3] unit-sphere vectors (look-at targets from world
origin). World up default = +Z. For vertices within ~2.5° of ±Z,
swap up to +Y to avoid the cross-product singularity.
"""
N = targets.shape[0]
eye = torch.zeros(3, dtype=torch.float32)
up_default = torch.tensor([0.0, 0.0, 1.0], dtype=torch.float32)
up_fallback = torch.tensor([0.0, 1.0, 0.0], dtype=torch.float32)
cos_thresh = 0.999
out = torch.empty(N, 4, 4, dtype=torch.float32)
cos_with_up = torch.abs(targets @ up_default)
for i in range(N):
up = up_fallback if float(cos_with_up[i]) > cos_thresh else up_default
out[i] = _look_at_w2c(eye, targets[i], up)
return out
def _make_debug_grid(face_images: torch.Tensor, tile_size: int = 96) -> torch.Tensor:
"""Tile N face crops into a single [H, W, 3] preview image.
Output is a roughly-square grid (cols = ceil(sqrt(N)) by default).
Each tile is `tile_size`x`tile_size`, downsampled via bilinear. Used
to satisfy WorldNavPanoramaSplit's `debug_image` output convention.
"""
N, H, W, _ = face_images.shape
cols = int(math.ceil(math.sqrt(N)))
rows = int(math.ceil(N / cols))
# Bilinear-resize each face to (tile_size, tile_size).
nchw = face_images.permute(0, 3, 1, 2) # [N, 3, H, W]
tiled = F.interpolate(
nchw, size=(tile_size, tile_size), mode="bilinear", align_corners=False,
) # [N, 3, T, T]
tiled = tiled.permute(0, 2, 3, 1) # [N, T, T, 3]
grid = torch.zeros(rows * tile_size, cols * tile_size, 3, dtype=face_images.dtype)
for i in range(N):
r, c = divmod(i, cols)
grid[
r * tile_size:(r + 1) * tile_size,
c * tile_size:(c + 1) * tile_size,
:,
] = tiled[i]
return grid
class SharpPanoramaIcosahedronSplit(io.ComfyNode):
"""Equirect panorama -> N square perspective crops at icosahedron vertices."""
@classmethod
def define_schema(cls):
return io.Schema(
node_id="SharpPanoramaIcosahedronSplit",
display_name="Sharp Panorama Icosahedron Split",
category="SHARP",
description=(
"Split an equirectangular panorama into N rectilinear views "
"(90° fov each) at icosahedron vertices. Pipe the views "
"through any rectilinear depth model (e.g. SHARP / MoGe2), "
"then merge the per-view depths back into an equirect depth "
"map via SharpDepthMerge.\n\n"
"1:1 schema mirror of WorldNavPanoramaSplit (modulo "
"panorama input being plain IMAGE here, not WorldStereo's "
"WORLDSTEREO_PANORAMA custom type)."
),
inputs=[
io.Image.Input(
"panorama",
tooltip="Equirectangular RGB panorama (2:1). [B, H, W, 3] "
"or [H, W, 3] in [0, 1]."),
io.Int.Input(
"resolution", default=512, min=128, max=2048, step=64,
tooltip="Per-face image resolution (square). Default 512 "
"matches WorldNavPanoramaSplit."),
io.Combo.Input(
"subdivision",
options=["icosahedron_12", "icosahedron_42"],
default="icosahedron_42",
tooltip="Tiling density. 12 = base icosahedron (faster), "
"42 = subdivided (better polar coverage)."),
io.Boolean.Input(
"use_gpu", default=True, optional=True,
tooltip="Batched panorama->face resampling via "
"torch.nn.functional.grid_sample (bilinear) on "
"CUDA. Falls back to CPU if CUDA isn't "
"available. Math identical within fp32 round-off."),
],
outputs=[
io.Image.Output(display_name="face_images"),
io.Custom("EXTRINSICS").Output(display_name="extrinsics"),
io.Custom("INTRINSICS").Output(display_name="intrinsics"),
io.Float.Output(display_name="fov_x_deg"),
io.Image.Output(display_name="debug_image"),
],
)
@classmethod
def execute(
cls, panorama: torch.Tensor,
resolution: int = 512,
subdivision: str = "icosahedron_42",
use_gpu: bool = True,
):
# Normalize input shape -> [H, W, 3], float in [0, 1].
pano = panorama
if pano.dim() == 4:
pano = pano[0]
if pano.dtype != torch.float32:
pano = pano.float()
if pano.shape[-1] == 4:
pano = pano[..., :3]
if pano.max() > 2.0:
pano = pano / 255.0
H_pano, W_pano, _ = pano.shape
# Hard 2:1 aspect check — same as WorldNavPanoramaSplit.
aspect = W_pano / H_pano
if abs(aspect - 2.0) > 0.01:
raise ValueError(
f"SharpPanoramaIcosahedronSplit: panorama must be 2:1 "
f"equirect, got {W_pano}x{H_pano} (ratio {aspect:.3f}). "
f"Crop / pad to 2:1 first."
)
device = pano.device
if use_gpu and torch.cuda.is_available() and device.type != "cuda":
pano = pano.cuda()
device = pano.device
targets_np = _icosahedron_split_vertices(subdivision)
targets = torch.from_numpy(targets_np).to(device=device, dtype=torch.float32)
N = targets.shape[0]
# Per-vertex extrinsics with pole-singularity fix.
extrinsics = _build_extrinsics_with_pole_fix(targets).to(device)
# Same intrinsics for every face (90° FOV, square).
fov_x_deg = 90.0
fov_y_deg = 90.0
K = _intrinsics_from_fov(
math.radians(fov_x_deg), math.radians(fov_y_deg),
int(resolution), int(resolution),
).to(device)
face_images = []
for i in range(N):
img = _sample_perspective_from_equirect(
pano, extrinsics[i], K, int(resolution), int(resolution),
)
face_images.append(img)
face_images_t = torch.stack(face_images, dim=0).contiguous() # [N, R, R, 3]
intrinsics_t = K.unsqueeze(0).expand(N, -1, -1).contiguous() # [N, 3, 3]
extrinsics_t = extrinsics.contiguous() # [N, 4, 4]
# Debug grid — small thumbnails so the preview node renders fast.
debug_grid = _make_debug_grid(face_images_t.cpu(), tile_size=96)
debug_image = debug_grid.unsqueeze(0) # [1, H, W, 3]
log.info(
f"[SharpPanoramaIcosahedronSplit] {N} faces "
f"({subdivision}) @ {resolution}x{resolution}, fov=90°"
)
return io.NodeOutput(
face_images_t, extrinsics_t, intrinsics_t,
float(fov_x_deg), debug_image,
)
NODE_CLASS_MAPPINGS = {
"SharpPanoramaIcosahedronSplit": SharpPanoramaIcosahedronSplit,
}
NODE_DISPLAY_NAME_MAPPINGS = {
"SharpPanoramaIcosahedronSplit": "Sharp Panorama Icosahedron Split",
}
nodes/predict_gaussian_attrs.py
+394
View File
@@ -0,0 +1,394 @@
"""SharpPredictGaussianAttrs node for ComfyUI-Sharp.
Runs SHARP inference and exposes the per-pixel gaussian decoder output as
structured per-layer tensors instead of flattening to a PLY.
Why: SHARP outputs `num_layers = 2` gaussians per stride-pixel (layer 0 =
visible/near surface, layer 1 = back/occluded surface). The gaussian
composer's pre-flatten shape is:
mean_vectors [B, 3, num_layers, H_grid, W_grid]
singular_values [B, 3, num_layers, H_grid, W_grid]
quaternions [B, 4, num_layers, H_grid, W_grid]
colors [B, 3, num_layers, H_grid, W_grid]
opacities [B, num_layers, H_grid, W_grid]
Where H_grid = W_grid = `predictor.output_resolution` (typically 768 when
SHARP internal_resolution = 1536). The `gaussian_composer` then flattens
to [B, num_layers*H*W, C] for save_ply; the post-flatten order is
contiguous in (layer, h, w) — layer 0 first, layer 1 second.
This node un-flattens that, separates layers, and emits:
first_layer_depth_refined IMAGE [B, H, W, 3] -- ndc-z of layer 0
first_layer_gaussian_attrs MULTIBAND_IMAGE [B, 14, H, W]
second_layer_depth_refined IMAGE [B, H, W, 3] -- ndc-z of layer 1
second_layer_gaussian_attrs MULTIBAND_IMAGE [B, 14, H, W]
The "refined depth" is the gaussian's ndc-z. Since
`ndc_matrix[2, 2] = 1.0` in `unproject_gaussians`, ndc-z equals
camera-space z. That's the post-decoder refined depth — tighter than
the disparity head's pre-refinement output exposed by SharpPredictDepth.
Multiband channel order (per layer):
0: position_x (ndc, in [-1, 1])
1: position_y (ndc, in [-1, 1])
2: position_z (ndc / camera z, same as depth_refined output;
duplicated here so multiband is self-contained)
3: scale_x (Σ-singular value)
4: scale_y
5: scale_z
6: quaternion_w
7: quaternion_x
8: quaternion_y
9: quaternion_z
10: color_r (linear or SH-DC; see gaussians.py)
11: color_g
12: color_b
13: opacity (post-activation, in [0, 1])
Downstream: feed `first_layer_depth_refined` (across N faces) into
AlignDepthMaps to recover per-face scales, then multiply the scales into
each multiband's position channels before reconstructing a PLY.
"""
from __future__ import annotations
import hashlib
import logging
import time
import numpy as np
import torch
import torch.nn.functional as F
from comfy_api.latest import io
log = logging.getLogger("sharp")
# Channel order for the multiband output. Mirrors the gaussian composer's
# concat order in the un-flattened tensor we assemble.
GAUSS_ATTR_CHANNEL_NAMES = [
"position_x", "position_y", "position_z",
"scale_x", "scale_y", "scale_z",
"quaternion_w", "quaternion_x", "quaternion_y", "quaternion_z",
"color_r", "color_g", "color_b",
"opacity",
]
NUM_GAUSS_ATTR_CHANNELS = len(GAUSS_ATTR_CHANNEL_NAMES) # 14
# Same encode cache as SharpPredict — feature-extraction is the heaviest
# step, and a user iterating on focal length should hit the cache.
_encode_cache: dict[str, object] = {
"image_hash": None,
"monodepth_output": None,
"image_resized": None,
"original_shape": None,
}
def _compute_image_hash(image_np: np.ndarray) -> str:
return hashlib.sha256(image_np.tobytes()).hexdigest()[:16]
def _monodepth_to(output, device):
"""Move all tensors in a MonodepthOutput to the given device."""
from .sharp.model import MonodepthOutput
return MonodepthOutput(
disparity=output.disparity.to(device),
encoder_features=[t.to(device) for t in output.encoder_features],
decoder_features=output.decoder_features.to(device),
output_features=[t.to(device) for t in output.output_features],
intermediate_features=[t.to(device) for t in output.intermediate_features],
)
def _split_layers_from_flat_gaussians(
gaussians,
num_layers: int,
H_grid: int,
W_grid: int,
) -> tuple[torch.Tensor, torch.Tensor]:
"""Flatten Gaussians3D -> per-layer per-pixel attribute stack.
The composer's flatten was:
x.permute(0, 2, 3, 4, 1).flatten(1, 3)
starting from [B, C, num_layers, H, W]. The flatten dim ordering is
(num_layers -> H -> W), so reshaping back to [B, num_layers, H, W, C]
recovers the original grid layout.
Returns:
layer0_attrs: [B, H, W, 14]
layer1_attrs: [B, H, W, 14] (for num_layers=2; if num_layers
differs, raises)
"""
mean = gaussians.mean_vectors # [B, N, 3]
scale = gaussians.singular_values # [B, N, 3]
quat = gaussians.quaternions # [B, N, 4]
color = gaussians.colors # [B, N, 3]
opac = gaussians.opacities # [B, N]
B, N, _ = mean.shape
expected_N = num_layers * H_grid * W_grid
if N != expected_N:
raise ValueError(
f"gaussian count mismatch: got N={N}, expected "
f"num_layers={num_layers} * H={H_grid} * W={W_grid} = {expected_N}. "
f"Check predictor.output_resolution / num_layers."
)
# Un-flatten: [B, N, C] -> [B, num_layers, H, W, C]
mean5 = mean.view(B, num_layers, H_grid, W_grid, 3)
scale5 = scale.view(B, num_layers, H_grid, W_grid, 3)
quat5 = quat.view(B, num_layers, H_grid, W_grid, 4)
color5 = color.view(B, num_layers, H_grid, W_grid, 3)
opac5 = opac.view(B, num_layers, H_grid, W_grid, 1)
# Concat into a single [B, num_layers, H, W, 14] tensor in the order
# the multiband channel names declare.
stacked = torch.cat([mean5, scale5, quat5, color5, opac5], dim=-1)
if stacked.shape[-1] != NUM_GAUSS_ATTR_CHANNELS:
raise RuntimeError(
f"channel count mismatch: stacked={stacked.shape[-1]}, "
f"expected {NUM_GAUSS_ATTR_CHANNELS}"
)
if num_layers < 2:
raise ValueError(
f"expected num_layers>=2 for SHARP's two-layer output, got {num_layers}"
)
return stacked[:, 0], stacked[:, 1]
def _build_multiband(attrs: torch.Tensor) -> dict:
"""Convert [B, H, W, 14] -> MULTIBAND_IMAGE dict ([B, 14, H, W])."""
if attrs.dim() != 4 or attrs.shape[-1] != NUM_GAUSS_ATTR_CHANNELS:
raise ValueError(
f"expected [B, H, W, {NUM_GAUSS_ATTR_CHANNELS}], got {tuple(attrs.shape)}"
)
samples = attrs.permute(0, 3, 1, 2).contiguous().float() # [B, C, H, W]
return {
"samples": samples,
"channel_names": list(GAUSS_ATTR_CHANNEL_NAMES),
"metadata": {
"source": "SharpPredictGaussianAttrs",
"ndc_space": True,
"note": (
"position channels are in NDC space (ndc_matrix[2,2]=1.0 "
"so position_z == camera-space z). Apply per-face scale "
"to positions and run unproject_gaussians to recover "
"world coords."
),
},
}
class SharpPredictGaussianAttrs(io.ComfyNode):
"""SHARP inference -> per-layer depth + multiband gaussian attributes."""
@classmethod
def define_schema(cls):
return io.Schema(
node_id="SharpPredictGaussianAttrs",
display_name="SHARP Predict (Depth + Gaussian Attrs)",
category="SHARP",
description=(
"Run SHARP inference and expose the per-pixel gaussian "
"decoder output as structured per-layer tensors (depth + "
"MULTIBAND_IMAGE of all gaussian attributes) instead of "
"flattening to a PLY file.\n\n"
"Use the depth outputs to drive AlignDepthMaps for per-"
"face scale recovery, then multiply scales into the "
"multiband position channels before reconstructing PLYs."
),
inputs=[
io.Custom("SHARP_MODEL_CONFIG").Input("model"),
io.Image.Input("image"),
io.Float.Input(
"focal_length_mm", default=30.0, min=0.0, max=500.0,
step=0.1, optional=True,
tooltip="Focal length in mm (35mm equiv). 0 = 30mm. "
"Ignored if intrinsics provided."),
io.Custom("EXTRINSICS").Input(
"extrinsics", optional=True,
tooltip="Pass-through (e.g. from SamplePanorama / "
"SharpPanoramaCubeSplit). Re-emitted on the "
"extrinsics output so downstream nodes "
"(AlignDepthMaps) can wire them."),
io.Custom("INTRINSICS").Input(
"intrinsics", optional=True,
tooltip="Pass-through. If provided, overrides "
"focal_length_mm. Re-emitted on the intrinsics "
"output."),
],
outputs=[
io.Image.Output(
display_name="first_layer_depth_refined",
tooltip="[B, H, W, 3] depth broadcast across 3 channels. "
"Layer-0 (visible surface) post-decoder z."),
io.Custom("MULTIBAND_IMAGE").Output(
display_name="first_layer_gaussian_attrs",
tooltip="[B, 14, H, W] all layer-0 gaussian attributes. "
"Channels: position_xyz, scale_xyz, quat_wxyz, "
"color_rgb, opacity."),
io.Image.Output(
display_name="second_layer_depth_refined",
tooltip="[B, H, W, 3] layer-1 (back/occluded surface) "
"post-decoder z."),
io.Custom("MULTIBAND_IMAGE").Output(
display_name="second_layer_gaussian_attrs",
tooltip="[B, 14, H, W] all layer-1 gaussian attributes."),
io.Custom("EXTRINSICS").Output(display_name="extrinsics"),
io.Custom("INTRINSICS").Output(display_name="intrinsics"),
],
)
@classmethod
@torch.no_grad()
def execute(
cls,
model,
image: torch.Tensor,
focal_length_mm: float = 30.0,
extrinsics: torch.Tensor | None = None,
intrinsics: torch.Tensor | None = None,
):
global _encode_cache
import comfy.model_management
from .load_model import _load_sharp_model
patcher = _load_sharp_model(model)
predictor = patcher.model
device = patcher.load_device
if image.dim() == 3:
image = image.unsqueeze(0)
B = image.shape[0]
log.info(
f"[SharpPredictGaussianAttrs] processing {B} image(s); "
f"output_resolution={predictor.output_resolution}, "
f"num_layers=2 (SHARP convention)"
)
# Process per-image (matches existing SharpPredict structure).
first_depths = []
first_attrs = []
second_depths = []
second_attrs = []
# Load model to GPU with the budget the existing predict path uses.
internal_shape = (1536, 1536)
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,
)
H_grid = W_grid = int(predictor.output_resolution)
num_layers = 2 # SHARP fixed; see model.py:__init__ default
for b in range(B):
img_np = image[b].cpu().numpy() if isinstance(image, torch.Tensor) else np.asarray(image[b])
if img_np.dtype != np.uint8:
img_np = (np.clip(img_np, 0, 1) * 255 + 0.5).astype(np.uint8)
height, width = img_np.shape[:2]
image_hash = _compute_image_hash(img_np)
# Cache logic — same as SharpPredict.
if _encode_cache["image_hash"] == image_hash:
log.info(f" [{b}] cache hit")
monodepth_output = _monodepth_to(_encode_cache["monodepth_output"], device)
image_resized_pt = _encode_cache["image_resized"].to(device)
else:
log.info(f" [{b}] encoding...")
_encode_cache["image_hash"] = None
image_pt = (
torch.from_numpy(img_np.copy()).float().to(device).permute(2, 0, 1) / 255.0
)
image_resized_pt = F.interpolate(
image_pt[None],
size=(internal_shape[1], internal_shape[0]),
mode="bilinear", align_corners=True,
)
t0 = time.time()
monodepth_output, _ = predictor.encode(image_resized_pt)
log.info(f" [{b}] encode time: {time.time() - t0:.2f}s")
_encode_cache["image_hash"] = image_hash
_encode_cache["monodepth_output"] = _monodepth_to(monodepth_output, "cpu")
_encode_cache["image_resized"] = image_resized_pt.cpu()
_encode_cache["original_shape"] = (height, width)
comfy.model_management.soft_empty_cache()
# f_px on the original image, projected to internal_shape's coords.
# Matches SharpPredict's disparity_factor convention so the decoder
# produces the same gaussians.
if intrinsics is not None:
intr_b = intrinsics[b] if intrinsics.dim() == 3 else intrinsics
f_px = float(intr_b[0, 0]) * (internal_shape[0] / width)
else:
# 30mm 35mm-equiv focal length: f_px = (width / 36) * fl_mm
f_px = (width / 36.0) * max(0.1, float(focal_length_mm or 30.0))
disparity_factor = torch.tensor([f_px / width]).float().to(device)
t1 = time.time()
gaussians_ndc = predictor.decode(monodepth_output, image_resized_pt, disparity_factor)
log.info(f" [{b}] decode time: {time.time() - t1:.2f}s")
# Un-flatten to per-pixel grid + split layers.
l0, l1 = _split_layers_from_flat_gaussians(
gaussians_ndc, num_layers, H_grid, W_grid,
)
# l0, l1 shape: [1, H, W, 14]
# Depth = position_z = channel index 2.
d0 = l0[..., 2] # [1, H, W]
d1 = l1[..., 2]
first_depths.append(d0.cpu())
second_depths.append(d1.cpu())
first_attrs.append(l0.cpu())
second_attrs.append(l1.cpu())
# Stack across batch.
d0_batch = torch.cat(first_depths, dim=0) # [B, H, W]
d1_batch = torch.cat(second_depths, dim=0)
attrs0_batch = torch.cat(first_attrs, dim=0) # [B, H, W, 14]
attrs1_batch = torch.cat(second_attrs, dim=0)
# IMAGE format: [B, H, W, 3] depth broadcast over channels.
d0_img = d0_batch.unsqueeze(-1).expand(-1, -1, -1, 3).contiguous()
d1_img = d1_batch.unsqueeze(-1).expand(-1, -1, -1, 3).contiguous()
# MULTIBAND_IMAGE: dict with samples [B, 14, H, W] + channel_names.
l0_mb = _build_multiband(attrs0_batch)
l1_mb = _build_multiband(attrs1_batch)
n_valid_l0 = int(d0_batch.gt(0).sum())
n_valid_l1 = int(d1_batch.gt(0).sum())
log.info(
f"[SharpPredictGaussianAttrs] done. "
f"shape: [{B}, {H_grid}, {W_grid}], "
f"layer0 depth: min={float(d0_batch.min()):.3f} "
f"median={float(d0_batch.median()):.3f} "
f"max={float(d0_batch.max()):.3f} (valid={n_valid_l0}), "
f"layer1 depth: min={float(d1_batch.min()):.3f} "
f"median={float(d1_batch.median()):.3f} "
f"max={float(d1_batch.max()):.3f} (valid={n_valid_l1})"
)
return io.NodeOutput(
d0_img, l0_mb, d1_img, l1_mb,
extrinsics, intrinsics,
)
NODE_CLASS_MAPPINGS = {
"SharpPredictGaussianAttrs": SharpPredictGaussianAttrs,
}
NODE_DISPLAY_NAME_MAPPINGS = {
"SharpPredictGaussianAttrs": "SHARP Predict (Depth + Gaussian Attrs)",
}
nodes/refined_depth_from_gaussians.py
+291
View File
@@ -0,0 +1,291 @@
"""SharpRefinedDepthFromGaussians node for ComfyUI-Sharp.
Recovers a per-face "refined depth map" from the gaussian PLYs produced by
SharpPredict, by reprojecting each face's layer-0 (visible-surface)
gaussians back into that face's camera frustum.
Why: SHARP's `SharpPredictDepth` exposes the pre-refinement disparity head
output. The full gaussian decoder produces a refined per-gaussian z that's
typically tighter than the disparity head's prediction. Those refined z's
ARE in the saved PLY (in world coords after `unproject_gaussians`); this
node inverts that unproject to recover the camera-space depth they
represent.
Pipeline shape:
SamplePanorama / SharpPanoramaCubeSplit
-> images, ext, intr
SharpPredict (batch)
-> ply_path (folder with N PLYs), ext, intr
SharpRefinedDepthFromGaussians (this node)
-> depth_maps [N, H, W, 3] # refined per-face depth
AlignDepthMaps
-> globally scaled depth_maps
(multiply scales back into the gaussians + MergeGaussians)
Algorithm per face i:
g_world = layer-0 gaussian positions in world coords (loaded from PLY)
cam = R_w2c[i] @ g_world + t_w2c[i] # world -> camera
z = cam[:, 2] # depth along view axis
valid = z > 0 # in front of camera
u = fx[i] * cam[:, 0] / z + cx[i] # face pixel coords
v = fy[i] * cam[:, 1] / z + cy[i]
depth_map[v, u] = min(depth_map[v, u], z) # nearest gaussian per pixel
# = layer-0 by construction
# (layer-1 is behind layer-0)
Output is sparse — only pixels where a layer-0 gaussian landed get a depth.
Other pixels are 0 + masked invalid.
"""
import logging
import os
from pathlib import Path
import numpy as np
import torch
from comfy_api.latest import io
log = logging.getLogger("sharp")
def _load_gaussian_positions(ply_path: Path) -> np.ndarray:
"""Read just the x/y/z fields from a Sharp-format PLY. Returns [N, 3]."""
from plyfile import PlyData
plydata = PlyData.read(str(ply_path))
vertex = plydata['vertex']
return np.stack([
np.asarray(vertex['x']),
np.asarray(vertex['y']),
np.asarray(vertex['z']),
], axis=-1).astype(np.float32)
def _list_ply_paths(ply_path_str: str, expected_n: int) -> list[Path]:
"""Resolve the ply_path input (SharpPredict output) to a list of PLYs.
SharpPredict writes:
- single image: <prefix>_<ts>.ply
- batch: <prefix>_<ts>/001.ply, 002.ply, ...
"""
p = Path(ply_path_str)
if p.is_file() and p.suffix.lower() == ".ply":
if expected_n != 1:
raise ValueError(
f"ply_path is a single file ({p.name}) but extrinsics has "
f"N={expected_n}. Pass the SharpPredict batch folder instead."
)
return [p]
if p.is_dir():
plies = sorted(
child for child in p.iterdir()
if child.is_file() and child.suffix.lower() == ".ply"
)
if len(plies) != expected_n:
raise ValueError(
f"ply_path directory {p} has {len(plies)} PLYs but extrinsics "
f"has N={expected_n}. Did SharpPredict run on the same batch?"
)
return plies
raise FileNotFoundError(f"ply_path {ply_path_str!r} is not a file or directory")
def _project_gaussians_to_depth(
gauss_world: torch.Tensor,
ext_w2c: torch.Tensor,
K_pixel: torch.Tensor,
out_h: int,
out_w: int,
) -> tuple[torch.Tensor, torch.Tensor]:
"""Project N gaussian positions into one face camera, return depth map.
Args:
gauss_world: [G, 3] world-space gaussian positions
ext_w2c: [4, 4] world-to-camera matrix for this face
K_pixel: [3, 3] pixel-space intrinsics
out_h, out_w: output depth map size
Returns:
depth_map: [H, W] float; depth at each pixel where a gaussian landed,
0 elsewhere. Per-pixel = min over all gaussians that hit
that pixel (so layer-0 is taken automatically since
layer-1 sits behind layer-0).
valid: [H, W] bool; True where at least one gaussian landed.
"""
device = gauss_world.device
# world -> camera (homogeneous)
R = ext_w2c[:3, :3]
t = ext_w2c[:3, 3]
cam = gauss_world @ R.T + t # [G, 3]
z = cam[:, 2]
# Project to pixel coords. Only consider gaussians in front of camera.
fx = K_pixel[0, 0]
fy = K_pixel[1, 1]
cx = K_pixel[0, 2]
cy = K_pixel[1, 2]
in_front = z > 1e-6
z_safe = torch.where(in_front, z, torch.ones_like(z))
u = fx * cam[:, 0] / z_safe + cx
v = fy * cam[:, 1] / z_safe + cy
# Round to nearest pixel; require in-bounds.
u_i = u.round().long()
v_i = v.round().long()
in_bounds = (u_i >= 0) & (u_i < out_w) & (v_i >= 0) & (v_i < out_h) & in_front
# Scatter: keep MIN z per pixel (-> layer-0, visible surface).
# 1-D flat pixel index for scatter_reduce.
flat_idx = v_i * out_w + u_i
# For invalid gaussians, point at index 0 with z = +inf so they lose every reduce.
safe_idx = torch.where(in_bounds, flat_idx, torch.zeros_like(flat_idx))
safe_z = torch.where(in_bounds, z, torch.full_like(z, float("inf")))
depth_flat = torch.full((out_h * out_w,), float("inf"), dtype=torch.float32, device=device)
depth_flat.scatter_reduce_(0, safe_idx, safe_z, reduce="amin", include_self=True)
depth_map = depth_flat.view(out_h, out_w)
valid = torch.isfinite(depth_map)
depth_map = torch.where(valid, depth_map, torch.zeros_like(depth_map))
return depth_map, valid
class SharpRefinedDepthFromGaussians(io.ComfyNode):
"""Per-face gaussian PLYs -> per-face refined depth maps (layer-0)."""
@classmethod
def define_schema(cls):
return io.Schema(
node_id="SharpRefinedDepthFromGaussians",
display_name="SHARP Refined Depth from Gaussians",
category="SHARP",
description=(
"Reproject SharpPredict's PLY gaussians back into per-face "
"camera space to recover the refined per-pixel depth (the "
"actual z's the decoder produced, NOT the pre-refinement "
"disparity head). Min-depth per pixel = layer 0 (visible "
"surface).\n\n"
"Feed depth_maps into AlignDepthMaps for log-scale "
"alignment, then apply the resulting per-face scales back "
"to the gaussians before MergeGaussians."
),
inputs=[
io.String.Input(
"ply_path",
tooltip="From SharpPredict.ply_path. Single .ply (for "
"batch_size=1) OR a folder with one .ply per "
"face (for batch). Folder file count must match "
"extrinsics N."),
io.Custom("EXTRINSICS").Input(
"extrinsics",
tooltip="[N, 4, 4] world-to-camera per face. Same matrix "
"you passed to SharpPredict — typically from "
"SamplePanorama or SharpPanoramaCubeSplit."),
io.Custom("INTRINSICS").Input(
"intrinsics",
tooltip="[N, 3, 3] (or [N, 4, 4]) pixel-space intrinsics "
"per face. The 3x3 K is used; any extra row/col "
"is ignored."),
io.Int.Input(
"out_size", default=1536, min=256, max=2048, step=64,
tooltip="Output depth-map resolution per face (square). "
"1536 matches SHARP's internal resolution."),
],
outputs=[
io.Image.Output(
display_name="depth_maps",
tooltip="[N, H, W, 3] float depth maps. Depth value "
"broadcast across 3 channels so it composes "
"with regular IMAGE-consuming nodes. 0 where "
"no gaussian landed."),
io.Mask.Output(
display_name="valid_mask",
tooltip="[N, H, W] float 0/1; 1 where a gaussian landed."),
io.Custom("EXTRINSICS").Output(
display_name="extrinsics",
tooltip="Pass-through, so this node can sit upstream of "
"AlignDepthMaps without rewiring."),
io.Custom("INTRINSICS").Output(
display_name="intrinsics",
tooltip="Pass-through."),
],
)
@classmethod
def execute(
cls, ply_path: str,
extrinsics: torch.Tensor, intrinsics: torch.Tensor,
out_size: int = 1536,
):
# Normalize ext/intr shapes.
ext = extrinsics.float()
if ext.dim() == 2:
ext = ext.unsqueeze(0)
N = ext.shape[0]
if ext.shape[1:] != (4, 4):
raise ValueError(f"extrinsics must be [N, 4, 4], got {tuple(ext.shape)}")
intr = intrinsics.float()
if intr.dim() == 2:
intr = intr.unsqueeze(0)
if intr.shape[0] != N:
raise ValueError(
f"intrinsics N ({intr.shape[0]}) doesn't match extrinsics N ({N})"
)
# Accept 4x4 or 3x3.
K_pixel = intr[:, :3, :3].contiguous()
# Resolve to per-face PLY paths.
ply_files = _list_ply_paths(ply_path, expected_n=N)
log.info(
f"[SharpRefinedDepthFromGaussians] loading {len(ply_files)} PLY(s) "
f"-> projecting to {out_size}x{out_size} per face"
)
device = ext.device
H_out = W_out = int(out_size)
depth_maps = []
valid_masks = []
for i, ply in enumerate(ply_files):
pos_np = _load_gaussian_positions(ply)
pos = torch.from_numpy(pos_np).to(device=device, dtype=torch.float32)
depth_i, valid_i = _project_gaussians_to_depth(
pos, ext[i].to(device), K_pixel[i].to(device),
out_h=H_out, out_w=W_out,
)
depth_maps.append(depth_i)
valid_masks.append(valid_i)
n_g = pos.shape[0]
n_v = int(valid_i.sum())
d_min = float(depth_i[valid_i].min()) if n_v > 0 else 0.0
d_med = float(torch.median(depth_i[valid_i])) if n_v > 0 else 0.0
d_max = float(depth_i[valid_i].max()) if n_v > 0 else 0.0
log.info(
f" face {i}: {n_g} gaussians -> {n_v}/{H_out*W_out} pixels covered "
f"({100*n_v/(H_out*W_out):.1f}%), depth min={d_min:.3f} "
f"median={d_med:.3f} max={d_max:.3f}"
)
depth_stack = torch.stack(depth_maps, dim=0) # [N, H, W]
valid_stack = torch.stack(valid_masks, dim=0) # [N, H, W]
# IMAGE convention: [N, H, W, 3], broadcast depth across channels.
depth_img = depth_stack.unsqueeze(-1).expand(-1, -1, -1, 3).contiguous()
valid_img = valid_stack.float().contiguous()
return io.NodeOutput(depth_img, valid_img, extrinsics, intrinsics)
NODE_CLASS_MAPPINGS = {
"SharpRefinedDepthFromGaussians": SharpRefinedDepthFromGaussians,
}
NODE_DISPLAY_NAME_MAPPINGS = {
"SharpRefinedDepthFromGaussians": "SHARP Refined Depth from Gaussians",
}