LUCAS M MORGAN
c740b9c7bc
- Use CATEGORY 'Gaussian Splat Shot' for LoadSharpModel, SharpPredict, and GaussianShotRender. - Reorder DOM viewer widget to top of node widget stack in gaussian_shot.js. - Set seed control_after_generate to 'fixed' so UI defaults to seed 0, not randomize. Made-with: Cursor
336 lines
12 KiB
Python
336 lines
12 KiB
Python
"""SharpPredict node for ComfyUI-Sharp."""
|
|
|
|
import hashlib
|
|
import logging
|
|
import os
|
|
import time
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
import torch
|
|
import torch.nn.functional as F
|
|
|
|
log = logging.getLogger("sharp")
|
|
|
|
# Try to import ComfyUI folder_paths for output directory
|
|
try:
|
|
import folder_paths
|
|
OUTPUT_DIR = folder_paths.get_output_directory()
|
|
except ImportError:
|
|
OUTPUT_DIR = os.path.join(os.path.dirname(os.path.dirname(__file__)), "output")
|
|
|
|
from .utils.image import comfy_to_numpy_rgb, convert_focallength
|
|
|
|
|
|
# Global cache for encoded features (single image only)
|
|
_encode_cache = {
|
|
"image_hash": None,
|
|
"monodepth_output": None,
|
|
"image_resized": None,
|
|
"original_shape": None,
|
|
}
|
|
|
|
|
|
def _compute_image_hash(image_np: np.ndarray) -> str:
|
|
"""Compute hash of image for cache key."""
|
|
return hashlib.sha256(image_np.tobytes()).hexdigest()[:16]
|
|
|
|
|
|
class SharpPredict:
|
|
"""Run SHARP inference to generate 3D Gaussians from a single image or batch."""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"model": ("SHARP_MODEL",),
|
|
"image": ("IMAGE",),
|
|
},
|
|
"optional": {
|
|
"focal_length_mm": ("FLOAT", {
|
|
"default": 30.0,
|
|
"min": 0.0,
|
|
"max": 500.0,
|
|
"step": 0.1,
|
|
"tooltip": "Focal length in mm (35mm equivalent). 0 = auto (defaults to 30mm). Ignored if intrinsics provided."
|
|
}),
|
|
"output_prefix": ("STRING", {
|
|
"default": "sharp",
|
|
"tooltip": "Prefix for output PLY filename or folder name for batches."
|
|
}),
|
|
"extrinsics": ("EXTRINSICS", {
|
|
"tooltip": "Camera extrinsics (from SamplePanorama). If batched, must match image batch size."
|
|
}),
|
|
"intrinsics": ("INTRINSICS", {
|
|
"tooltip": "Camera intrinsics (from SamplePanorama). Overrides focal_length_mm if provided."
|
|
}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING", "EXTRINSICS", "INTRINSICS",)
|
|
RETURN_NAMES = ("ply_path", "extrinsics", "intrinsics",)
|
|
FUNCTION = "predict"
|
|
CATEGORY = "Gaussian Splat Shot"
|
|
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."
|
|
|
|
@torch.no_grad()
|
|
def predict(
|
|
self,
|
|
model,
|
|
image: torch.Tensor,
|
|
focal_length_mm: float = 0.0,
|
|
output_prefix: str = "sharp",
|
|
extrinsics: torch.Tensor = None,
|
|
intrinsics: torch.Tensor = None,
|
|
):
|
|
"""Run SHARP inference and save PLY file(s).
|
|
|
|
For single image: saves {prefix}_{timestamp}.ply
|
|
For batch: creates folder {prefix}_{timestamp}/ with 001.ply, 002.ply, etc.
|
|
|
|
Features are cached per image - changing focal_length with same image is instant.
|
|
|
|
If extrinsics/intrinsics are provided (from SamplePanorama), Gaussians are
|
|
unprojected into world coordinates using those camera parameters.
|
|
"""
|
|
import comfy.model_management
|
|
import comfy.utils
|
|
from .sharp.gaussians import save_ply, unproject_gaussians
|
|
|
|
# model is a ModelPatcher from LoadSharpModel
|
|
# Estimate activation memory: SPN processes 35 patches through ViT at 1536x1536
|
|
dtype = model.model.dtype if hasattr(model.model, 'dtype') else torch.float32
|
|
memory_required = (1536 * 1536 * 6) * comfy.model_management.dtype_size(dtype)
|
|
comfy.model_management.load_models_gpu([model], memory_required=memory_required)
|
|
predictor = model.model
|
|
device = model.load_device
|
|
|
|
# Handle batch dimension
|
|
if image.dim() == 3:
|
|
image = image.unsqueeze(0)
|
|
|
|
batch_size = image.shape[0]
|
|
|
|
# Validate extrinsics batch size if provided
|
|
has_camera_params = extrinsics is not None and intrinsics is not None
|
|
if has_camera_params:
|
|
if extrinsics.dim() == 2:
|
|
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})")
|
|
log.info(f"Processing {batch_size} image(s) with provided camera parameters (panorama mode)")
|
|
else:
|
|
log.info(f"Processing {batch_size} image(s)")
|
|
|
|
# Ensure output directory exists
|
|
os.makedirs(OUTPUT_DIR, exist_ok=True)
|
|
timestamp = int(time.time() * 1000)
|
|
|
|
# Determine output path(s)
|
|
if batch_size == 1:
|
|
# Single image: save directly as PLY file
|
|
output_filename = f"{output_prefix}_{timestamp}.ply"
|
|
output_path = os.path.join(OUTPUT_DIR, output_filename)
|
|
is_batch = False
|
|
else:
|
|
# Multiple images: create folder
|
|
folder_name = f"{output_prefix}_{timestamp}"
|
|
output_folder = os.path.join(OUTPUT_DIR, folder_name)
|
|
os.makedirs(output_folder, exist_ok=True)
|
|
output_path = output_folder
|
|
is_batch = True
|
|
|
|
all_ply_paths = []
|
|
all_extrinsics = []
|
|
all_intrinsics = []
|
|
|
|
inference_start = time.time()
|
|
pbar = comfy.utils.ProgressBar(batch_size)
|
|
|
|
for i in range(batch_size):
|
|
comfy.model_management.throw_exception_if_processing_interrupted()
|
|
# Extract single image from batch
|
|
single_image = image[i:i+1]
|
|
image_np = comfy_to_numpy_rgb(single_image)
|
|
height, width = image_np.shape[:2]
|
|
|
|
if i == 0:
|
|
log.info(f"Image size: {width}x{height}")
|
|
|
|
# Get camera parameters for this image
|
|
if has_camera_params:
|
|
# Use provided intrinsics (extract focal length)
|
|
img_intrinsics = intrinsics.to(device)
|
|
img_extrinsics = extrinsics[i].to(device)
|
|
f_px = img_intrinsics[0, 0].item() # fx from intrinsics matrix
|
|
else:
|
|
# Use focal_length_mm parameter
|
|
if focal_length_mm > 0:
|
|
f_px = convert_focallength(width, height, focal_length_mm)
|
|
else:
|
|
f_px = convert_focallength(width, height, 30.0)
|
|
img_extrinsics = None
|
|
img_intrinsics = None
|
|
|
|
# Run inference with caching
|
|
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,
|
|
intrinsics=img_intrinsics,
|
|
)
|
|
|
|
# Determine output filename
|
|
if is_batch:
|
|
ply_filename = f"{i+1:03d}.ply"
|
|
ply_path = os.path.join(output_folder, ply_filename)
|
|
else:
|
|
ply_path = output_path
|
|
|
|
# Save PLY and get metadata
|
|
_, metadata = save_ply(gaussians, f_px, (height, width), Path(ply_path))
|
|
|
|
all_ply_paths.append(ply_path)
|
|
all_extrinsics.append(metadata["extrinsic"])
|
|
all_intrinsics.append(metadata["intrinsic"])
|
|
|
|
log.info(f"Saved: {ply_path} ({metadata['num_gaussians']:,} gaussians)")
|
|
pbar.update(1)
|
|
|
|
inference_time = time.time() - inference_start
|
|
log.info(f"Total inference time: {inference_time:.2f}s ({inference_time/batch_size:.2f}s per image)")
|
|
|
|
# Return values
|
|
if is_batch:
|
|
# For batch: return folder path, and first image's camera params
|
|
# (assuming all images have same camera - user can override)
|
|
return (output_path, all_extrinsics[0], all_intrinsics[0],)
|
|
else:
|
|
# For single image: return PLY path and camera params
|
|
return (output_path, all_extrinsics[0], all_intrinsics[0],)
|
|
|
|
def _predict_image_cached(
|
|
self,
|
|
predictor,
|
|
image: np.ndarray,
|
|
f_px: float,
|
|
device: torch.device,
|
|
extrinsics: torch.Tensor = None,
|
|
intrinsics: torch.Tensor = None,
|
|
):
|
|
"""Predict Gaussians with caching of encoded features.
|
|
|
|
The expensive encode step is cached per image.
|
|
Changing focal_length reuses cached features (instant).
|
|
|
|
Args:
|
|
predictor: SHARP predictor model
|
|
image: Input image as numpy array [H, W, 3]
|
|
f_px: Focal length in pixels
|
|
device: Torch device
|
|
extrinsics: Optional 4x4 camera extrinsics (world-to-camera)
|
|
intrinsics: Optional 4x4 camera intrinsics
|
|
"""
|
|
global _encode_cache
|
|
import comfy.model_management
|
|
from .sharp.gaussians import unproject_gaussians
|
|
|
|
internal_shape = (1536, 1536)
|
|
height, width = image.shape[:2]
|
|
|
|
# Compute image hash for cache
|
|
image_hash = _compute_image_hash(image)
|
|
|
|
# Check cache
|
|
if _encode_cache["image_hash"] == image_hash:
|
|
# Cache hit - reuse encoded features
|
|
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
|
|
log.info("Encoding...")
|
|
|
|
# Clear old cache
|
|
_encode_cache["image_hash"] = None
|
|
_encode_cache["monodepth_output"] = None
|
|
_encode_cache["image_resized"] = None
|
|
_encode_cache["original_shape"] = None
|
|
|
|
# Convert to tensor and normalize
|
|
image_pt = torch.from_numpy(image.copy()).float().to(device).permute(2, 0, 1) / 255.0
|
|
|
|
# Resize to internal resolution
|
|
image_resized_pt = F.interpolate(
|
|
image_pt[None],
|
|
size=(internal_shape[1], internal_shape[0]),
|
|
mode="bilinear",
|
|
align_corners=True,
|
|
)
|
|
|
|
# Encode
|
|
encode_start = time.time()
|
|
monodepth_output, _ = predictor.encode(image_resized_pt)
|
|
log.info(f"Encode time: {time.time() - encode_start:.2f}s")
|
|
|
|
# Release fragmented GPU memory after heavy encode (35 ViT passes)
|
|
comfy.model_management.soft_empty_cache()
|
|
|
|
# Update cache
|
|
_encode_cache["image_hash"] = image_hash
|
|
_encode_cache["monodepth_output"] = monodepth_output
|
|
_encode_cache["image_resized"] = image_resized_pt
|
|
_encode_cache["original_shape"] = (height, width)
|
|
|
|
# Decode - always run with current focal length
|
|
disparity_factor = torch.tensor([f_px / width]).float().to(device)
|
|
|
|
decode_start = time.time()
|
|
gaussians_ndc = predictor.decode(monodepth_output, image_resized_pt, disparity_factor)
|
|
log.info(f"Decode time: {time.time() - decode_start:.2f}s")
|
|
|
|
# Build intrinsics for unprojection (use provided or construct from f_px)
|
|
if intrinsics is not None:
|
|
unproj_intrinsics = intrinsics.clone()
|
|
else:
|
|
unproj_intrinsics = (
|
|
torch.tensor(
|
|
[
|
|
[f_px, 0, width / 2, 0],
|
|
[0, f_px, height / 2, 0],
|
|
[0, 0, 1, 0],
|
|
[0, 0, 0, 1],
|
|
]
|
|
)
|
|
.float()
|
|
.to(device)
|
|
)
|
|
|
|
# Scale intrinsics to internal resolution
|
|
intrinsics_resized = unproj_intrinsics.clone()
|
|
intrinsics_resized[0] *= internal_shape[0] / width
|
|
intrinsics_resized[1] *= internal_shape[1] / height
|
|
|
|
# Use provided extrinsics or identity
|
|
if extrinsics is not None:
|
|
unproj_extrinsics = extrinsics
|
|
else:
|
|
unproj_extrinsics = torch.eye(4, device=device)
|
|
|
|
# Convert Gaussians to world/metric space
|
|
gaussians = unproject_gaussians(
|
|
gaussians_ndc, unproj_extrinsics, intrinsics_resized, internal_shape
|
|
)
|
|
|
|
return gaussians
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"SharpPredict": SharpPredict,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"SharpPredict": "SHARP Predict (Image to PLY)",
|
|
}
|