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

gaussian support

Browse Source
This commit is contained in:
Andrea Pozzetti
2025-12-19 21:31:46 +01:00
parent 29892e4b3a
commit d0e8c17fd6
4 changed files with 90 additions and 79 deletions
Download Patch File Download Diff File Expand all files Collapse all files
nodes/load_model.py
+38 -4
View File
@@ -1,14 +1,24 @@
"""LoadSharpModel node for ComfyUI-Sharp."""
import os
from urllib.request import urlopen
import torch
from tqdm import tqdm
# Try to get ComfyUI models directory
try:
import folder_paths
MODELS_DIR = os.path.join(folder_paths.models_dir, "sharp")
except ImportError:
MODELS_DIR = os.path.join(os.path.dirname(os.path.dirname(__file__)), "models", "sharp")
# Model cache
_MODEL_CACHE = {}
# Default model URL (Hugging Face mirror)
# Default model URL and filename
DEFAULT_MODEL_URL = "https://huggingface.co/apple/Sharp/resolve/main/sharp_2572gikvuh.pt"
DEFAULT_MODEL_FILENAME = "sharp_2572gikvuh.pt"
class LoadSharpModel:
@@ -58,16 +68,40 @@ class LoadSharpModel:
# Load or download checkpoint
if checkpoint_path and os.path.exists(checkpoint_path):
state_dict = torch.load(checkpoint_path, weights_only=True)
model_path = checkpoint_path
else:
print(f"[SHARP] Downloading model from Hugging Face...")
state_dict = torch.hub.load_state_dict_from_url(DEFAULT_MODEL_URL, progress=True)
# 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)
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,
nodes/predict.py
+17 -6
View File
@@ -43,8 +43,8 @@ class SharpPredict:
}
}
RETURN_TYPES = ("STRING",)
RETURN_NAMES = ("ply_path",)
RETURN_TYPES = ("STRING", "EXTRINSICS", "INTRINSICS",)
RETURN_NAMES = ("ply_path", "extrinsics", "intrinsics",)
FUNCTION = "predict"
CATEGORY = "SHARP"
OUTPUT_NODE = True
@@ -67,6 +67,7 @@ class SharpPredict:
# Convert ComfyUI image to numpy RGB
image_np = comfy_to_numpy_rgb(image)
height, width = image_np.shape[:2]
print(f"[SHARP] Input image: {width}x{height}")
# Determine focal length in pixels
if focal_length_mm > 0:
@@ -76,7 +77,11 @@ class SharpPredict:
f_px = convert_focallength(width, height, 30.0)
# Run inference
print(f"[SHARP] Running inference...")
inference_start = time.time()
gaussians = self._predict_image(predictor, image_np, f_px, device)
inference_time = time.time() - inference_start
print(f"[SHARP] Inference complete in {inference_time:.2f}s")
# Ensure output directory exists
os.makedirs(OUTPUT_DIR, exist_ok=True)
@@ -86,12 +91,18 @@ class SharpPredict:
output_filename = f"{output_prefix}_{timestamp}.ply"
output_path = os.path.join(OUTPUT_DIR, output_filename)
# Save PLY
save_ply(gaussians, f_px, (height, width), Path(output_path))
# Save PLY and get metadata
print(f"[SHARP] Saving PLY file...")
_, metadata = save_ply(gaussians, f_px, (height, width), Path(output_path))
print(f"[SHARP] Saved PLY to: {output_path}")
print(f"[SHARP] Saved: {output_path}")
print(f"[SHARP] Result: {metadata['num_gaussians']:,} gaussians")
return (output_path,)
# Extract extrinsics (4x4) and intrinsics (3x3) for camera setup
extrinsics = metadata["extrinsic"] # 4x4 identity matrix
intrinsics = metadata["intrinsic"] # 3x3 camera matrix
return (output_path, extrinsics, intrinsics,)
def _predict_image(
self,
requirements.txt
+6
View File
@@ -11,3 +11,9 @@ scipy>=1.10.0
pillow-heif>=0.10.0
imageio>=2.0.0
imageio-ffmpeg>=0.4.0
# Hugging Face transfer extension (faster downloads on Windows)
hf_xet
# Progress bars
tqdm>=4.60.0
sharp/utils/gaussians.py
+29 -69
View File
@@ -346,8 +346,12 @@ def load_ply(path: Path) -> tuple[Gaussians3D, SceneMetaData]:
@torch.no_grad()
def save_ply(
gaussians: Gaussians3D, f_px: float, image_shape: tuple[int, int], path: Path
) -> PlyData:
"""Save a predicted Gaussian3D to a ply file."""
) -> tuple[PlyData, dict]:
"""Save a predicted Gaussian3D to a ply file.
Returns:
tuple: (PlyData, metadata_dict) where metadata_dict contains camera and scene info.
"""
def _inverse_sigmoid(tensor: torch.Tensor) -> torch.Tensor:
return torch.log(tensor / (1.0 - tensor))
@@ -401,50 +405,10 @@ def save_ply(
elements[:] = list(map(tuple, attributes.detach().cpu().numpy()))
vertex_elements = PlyElement.describe(elements, "vertex")
# Load image-wise metadata.
# Build metadata dictionary instead of embedding in PLY
image_height, image_width = image_shape
# Export image size.
dtype_image_size = [("image_size", "u4")]
image_size_array = np.empty(2, dtype=dtype_image_size)
image_size_array[:] = np.array([image_width, image_height])
image_size_element = PlyElement.describe(image_size_array, "image_size")
# Export intrinsics.
dtype_intrinsic = [("intrinsic", "f4")]
intrinsic_array = np.empty(9, dtype=dtype_intrinsic)
intrinsic = np.array(
[
f_px,
0,
image_width * 0.5,
0,
f_px,
image_height * 0.5,
0,
0,
1,
]
)
intrinsic_array[:] = intrinsic.flatten()
intrinsic_element = PlyElement.describe(intrinsic_array, "intrinsic")
# Export dummy extrinsics.
dtype_extrinsic = [("extrinsic", "f4")]
extrinsic_array = np.empty(16, dtype=dtype_extrinsic)
extrinsic_array[:] = np.eye(4).flatten()
extrinsic_element = PlyElement.describe(extrinsic_array, "extrinsic")
# Export number of frames and particles per frame.
dtype_frames = [("frame", "i4")]
frame_array = np.empty(2, dtype=dtype_frames)
frame_array[:] = np.array([1, num_gaussians], dtype=np.int32)
frame_element = PlyElement.describe(frame_array, "frame")
# Export disparity ranges for transform.
dtype_disparity = [("disparity", "f4")]
disparity_array = np.empty(2, dtype=dtype_disparity)
# Calculate disparity quantiles
disparity = 1.0 / gaussians.mean_vectors[0, ..., -1]
quantiles = (
torch.quantile(disparity, q=torch.tensor([0.1, 0.9], device=disparity.device))
@@ -452,32 +416,28 @@ def save_ply(
.cpu()
.numpy()
)
disparity_array[:] = quantiles
disparity_element = PlyElement.describe(disparity_array, "disparity")
# Export colorspace.
dtype_color_space = [("color_space", "u1")]
color_space_array = np.empty(1, dtype=dtype_color_space)
color_space_array[:] = np.array([color_space_index]).flatten()
color_space_element = PlyElement.describe(color_space_array, "color_space")
# Build intrinsic matrix
intrinsic = [
[f_px, 0, image_width * 0.5],
[0, f_px, image_height * 0.5],
[0, 0, 1],
]
dtype_version = [("version", "u1")]
version_array = np.empty(3, dtype=dtype_version)
version_array[:] = np.array([1, 5, 0], dtype=np.uint8).flatten()
version_element = PlyElement.describe(version_array, "version")
plydata = PlyData(
[
vertex_elements,
extrinsic_element,
intrinsic_element,
image_size_element,
frame_element,
disparity_element,
color_space_element,
version_element,
]
)
metadata = {
"num_gaussians": num_gaussians,
"image_width": image_width,
"image_height": image_height,
"focal_length_px": f_px,
"intrinsic": intrinsic,
"extrinsic": [[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 1, 0], [0, 0, 0, 1]],
"disparity_range": [float(quantiles[0]), float(quantiles[1])],
"color_space": "sRGB",
"version": [1, 5, 0],
}
# Save only vertex data (standard 3DGS PLY format for compatibility)
plydata = PlyData([vertex_elements])
plydata.write(path)
return plydata
return plydata, metadata