diff --git a/nodes/load_model.py b/nodes/load_model.py index f36262f..bba063e 100644 --- a/nodes/load_model.py +++ b/nodes/load_model.py @@ -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, diff --git a/nodes/predict.py b/nodes/predict.py index 8bfa153..bfd07ef 100644 --- a/nodes/predict.py +++ b/nodes/predict.py @@ -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, diff --git a/requirements.txt b/requirements.txt index d70454a..b5d6a69 100644 --- a/requirements.txt +++ b/requirements.txt @@ -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 diff --git a/sharp/utils/gaussians.py b/sharp/utils/gaussians.py index ed73de8..62b4181 100644 --- a/sharp/utils/gaussians.py +++ b/sharp/utils/gaussians.py @@ -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