RobinsAIWorld/ml-sharp
1
1"""SHARP inference + optional CUDA video rendering utilities.2 3Design goals:4- Reuse SHARP's own predict/render pipeline (no subprocess calls).5- Be robust on Hugging Face Spaces + ZeroGPU.6- Cache model weights and predictor construction across requests.7 8Public API (used by the Gradio app):9- TrajectoryType10- predict_and_maybe_render_gpu(...)11"""12 13from __future__ import annotations14 15import os16import threading17import time18import uuid19from contextlib import contextmanager20from dataclasses import dataclass21from pathlib import Path22from typing import Final, Literal23 24import torch25 26# Optional Spaces GPU support (for HuggingFace Spaces deployment)27try:28 import spaces29 _SPACES_AVAILABLE = True30except ImportError:31 spaces = None # type: ignore[assignment]32 _SPACES_AVAILABLE = False33 34try:35 # Prefer HF cache / Hub downloads (works with Spaces `preload_from_hub`).36 from huggingface_hub import hf_hub_download, try_to_load_from_cache37except Exception: # pragma: no cover38 hf_hub_download = None # type: ignore[assignment]39 try_to_load_from_cache = None # type: ignore[assignment]40 41from sharp.cli.predict import DEFAULT_MODEL_URL, predict_image42from sharp.cli.render import render_gaussians as sharp_render_gaussians43from sharp.models import PredictorParams, create_predictor44from sharp.utils import camera, io45from sharp.utils.gaussians import Gaussians3D, SceneMetaData, save_ply46from sharp.utils.gsplat import GSplatRenderer47 48TrajectoryType = Literal["swipe", "shake", "rotate", "rotate_forward"]49 50# -----------------------------------------------------------------------------51# Helpers52# -----------------------------------------------------------------------------53 54 55def _now_ms() -> int:56 return int(time.time() * 1000)57 58 59def _ensure_dir(path: Path) -> Path:60 path.mkdir(parents=True, exist_ok=True)61 return path62 63 64def _make_even(x: int) -> int:65 return x if x % 2 == 0 else x + 166 67 68def _select_device(preference: str = "auto") -> torch.device:69 """Select the best available device for inference (CPU/CUDA/MPS)."""70 if preference not in {"auto", "cpu", "cuda", "mps"}:71 raise ValueError("device preference must be one of: auto|cpu|cuda|mps")72 73 if preference == "cpu":74 return torch.device("cpu")75 if preference == "cuda":76 return torch.device("cuda" if torch.cuda.is_available() else "cpu")77 if preference == "mps":78 return torch.device("mps" if torch.backends.mps.is_available() else "cpu")79 80 # auto81 if torch.cuda.is_available():82 return torch.device("cuda")83 if torch.backends.mps.is_available():84 return torch.device("mps")85 return torch.device("cpu")86 87 88# -----------------------------------------------------------------------------89# Prediction outputs90# -----------------------------------------------------------------------------91 92 93@dataclass(frozen=True, slots=True)94class PredictionOutputs:95 """Outputs of SHARP inference (plus derived metadata for rendering)."""96 97 ply_path: Path98 gaussians: Gaussians3D99 metadata_for_render: SceneMetaData100 input_resolution_hw: tuple[int, int]101 focal_length_px: float102 103 104# -----------------------------------------------------------------------------105# Patch SHARP VideoWriter to properly close the optional depth writer106# -----------------------------------------------------------------------------107 108 109class _PatchedVideoWriter(io.VideoWriter):110 """Ensure depth writer is closed so files can be safely cleaned up."""111 112 def __init__(113 self, output_path: Path, fps: float = 30.0, render_depth: bool = True114 ) -> None:115 super().__init__(output_path, fps=fps, render_depth=render_depth)116 # Ensure attribute exists for downstream code paths.117 if not hasattr(self, "depth_writer"):118 self.depth_writer = None # type: ignore[attribute-defined-outside-init]119 120 def close(self):121 super().close()122 depth_writer = getattr(self, "depth_writer", None)123 try:124 if depth_writer is not None:125 depth_writer.close()126 except Exception:127 pass128 129 130@contextmanager131def _patched_sharp_videowriter():132 """Temporarily patch `sharp.utils.io.VideoWriter` used by `sharp.cli.render`."""133 original = io.VideoWriter134 io.VideoWriter = _PatchedVideoWriter # type: ignore[assignment]135 try:136 yield137 finally:138 io.VideoWriter = original # type: ignore[assignment]139 140 141# -----------------------------------------------------------------------------142# Model wrapper143# -----------------------------------------------------------------------------144 145 146class ModelWrapper:147 """Cached SHARP model wrapper for Gradio/Spaces."""148 149 def __init__(150 self,151 *,152 outputs_dir: str | Path = "outputs",153 checkpoint_url: str = DEFAULT_MODEL_URL,154 checkpoint_path: str | Path | None = None,155 device_preference: str = "auto",156 keep_model_on_device: bool | None = None,157 hf_repo_id: str | None = None,158 hf_filename: str | None = None,159 hf_revision: str | None = None,160 ) -> None:161 self.outputs_dir = _ensure_dir(Path(outputs_dir))162 self.checkpoint_url = checkpoint_url163 164 env_ckpt = os.getenv("SHARP_CHECKPOINT_PATH") or os.getenv("SHARP_CHECKPOINT")165 if checkpoint_path:166 self.checkpoint_path = Path(checkpoint_path)167 elif env_ckpt:168 self.checkpoint_path = Path(env_ckpt)169 else:170 self.checkpoint_path = None171 172 # Optional Hugging Face Hub fallback (useful when direct CDN download fails).173 self.hf_repo_id = hf_repo_id or os.getenv("SHARP_HF_REPO_ID", "apple/Sharp")174 self.hf_filename = hf_filename or os.getenv(175 "SHARP_HF_FILENAME", "sharp_2572gikvuh.pt"176 )177 self.hf_revision = hf_revision or os.getenv("SHARP_HF_REVISION") or None178 179 self.device_preference = device_preference180 181 # Local CUDA: keep model on device by default for better performance182 if keep_model_on_device is None:183 keep_env = os.getenv("SHARP_KEEP_MODEL_ON_DEVICE", "1")184 self.keep_model_on_device = keep_env != "0"185 else:186 self.keep_model_on_device = keep_model_on_device187 188 # Support CUDA device selection via env var189 cuda_device = os.getenv("CUDA_VISIBLE_DEVICES")190 if cuda_device and device_preference == "auto":191 # Let PyTorch handle device mapping via CUDA_VISIBLE_DEVICES192 pass193 194 self._lock = threading.RLock()195 self._predictor: torch.nn.Module | None = None196 self._predictor_device: torch.device | None = None197 self._state_dict: dict | None = None198 199 def has_cuda(self) -> bool:200 return torch.cuda.is_available()201 202 def _load_state_dict(self) -> dict:203 with self._lock:204 if self._state_dict is not None:205 return self._state_dict206 207 # 1) Explicit local checkpoint path208 if self.checkpoint_path is not None:209 try:210 self._state_dict = torch.load(211 self.checkpoint_path,212 weights_only=True,213 map_location="cpu",214 )215 return self._state_dict216 except Exception as e:217 raise RuntimeError(218 "Failed to load SHARP checkpoint from local path.\n\n"219 f"Path:\n {self.checkpoint_path}\n\n"220 f"Original error:\n {type(e).__name__}: {e}"221 ) from e222 223 # 2) HF cache (no-network): best match for Spaces `preload_from_hub`.224 hf_cache_error: Exception | None = None225 if try_to_load_from_cache is not None:226 try:227 cached = try_to_load_from_cache(228 repo_id=self.hf_repo_id,229 filename=self.hf_filename,230 revision=self.hf_revision,231 repo_type="model",232 )233 except TypeError:234 cached = try_to_load_from_cache(self.hf_repo_id, self.hf_filename) # type: ignore[misc]235 236 try:237 if isinstance(cached, str) and Path(cached).exists():238 self._state_dict = torch.load(239 cached, weights_only=True, map_location="cpu"240 )241 return self._state_dict242 except Exception as e:243 hf_cache_error = e244 245 # 3) HF Hub download (reuse cache when available; may download otherwise).246 hf_error: Exception | None = None247 if hf_hub_download is not None:248 # Attempt "local only" mode if supported (avoids network).249 try:250 import inspect251 252 if "local_files_only" in inspect.signature(hf_hub_download).parameters:253 ckpt_path = hf_hub_download(254 repo_id=self.hf_repo_id,255 filename=self.hf_filename,256 revision=self.hf_revision,257 local_files_only=True,258 )259 if Path(ckpt_path).exists():260 self._state_dict = torch.load(261 ckpt_path, weights_only=True, map_location="cpu"262 )263 return self._state_dict264 except Exception:265 pass266 267 try:268 ckpt_path = hf_hub_download(269 repo_id=self.hf_repo_id,270 filename=self.hf_filename,271 revision=self.hf_revision,272 )273 self._state_dict = torch.load(274 ckpt_path,275 weights_only=True,276 map_location="cpu",277 )278 return self._state_dict279 except Exception as e:280 hf_error = e281 282 # 4) Default upstream CDN (torch hub cache). Last resort.283 url_error: Exception | None = None284 try:285 self._state_dict = torch.hub.load_state_dict_from_url(286 self.checkpoint_url,287 progress=True,288 map_location="cpu",289 )290 return self._state_dict291 except Exception as e:292 url_error = e293 294 # If we got here: all options failed.295 hint_lines = [296 "Failed to load SHARP checkpoint.",297 "",298 "Tried (in order):",299 f" 1) HF cache (preload_from_hub): repo_id={self.hf_repo_id}, filename={self.hf_filename}, revision={self.hf_revision or 'None'}",300 f" 2) HF Hub download: repo_id={self.hf_repo_id}, filename={self.hf_filename}, revision={self.hf_revision or 'None'}",301 f" 3) URL (torch hub): {self.checkpoint_url}",302 "",303 "If network access is restricted, set a local checkpoint path:",304 " - SHARP_CHECKPOINT_PATH=/path/to/sharp_2572gikvuh.pt",305 "",306 "Original errors:",307 ]308 if try_to_load_from_cache is None:309 hint_lines.append(" HF cache: huggingface_hub not installed")310 elif hf_cache_error is not None:311 hint_lines.append(312 f" HF cache: {type(hf_cache_error).__name__}: {hf_cache_error}"313 )314 else:315 hint_lines.append(" HF cache: (not found in cache)")316 317 if hf_hub_download is None:318 hint_lines.append(" HF download: huggingface_hub not installed")319 else:320 hint_lines.append(f" HF download: {type(hf_error).__name__}: {hf_error}")321 322 hint_lines.append(f" URL: {type(url_error).__name__}: {url_error}")323 324 raise RuntimeError("\n".join(hint_lines))325 326 def _get_predictor(self, device: torch.device) -> torch.nn.Module:327 with self._lock:328 if self._predictor is None:329 state_dict = self._load_state_dict()330 predictor = create_predictor(PredictorParams())331 predictor.load_state_dict(state_dict)332 predictor.eval()333 self._predictor = predictor334 self._predictor_device = torch.device("cpu")335 336 assert self._predictor is not None337 assert self._predictor_device is not None338 339 if self._predictor_device != device:340 self._predictor.to(device)341 self._predictor_device = device342 343 return self._predictor344 345 def _maybe_move_model_back_to_cpu(self) -> None:346 if self.keep_model_on_device:347 return348 with self._lock:349 if self._predictor is not None and self._predictor_device is not None:350 if self._predictor_device.type != "cpu":351 self._predictor.to("cpu")352 self._predictor_device = torch.device("cpu")353 if torch.cuda.is_available():354 torch.cuda.empty_cache()355 356 def _make_output_stem(self, input_path: Path) -> str:357 return f"{input_path.stem}-{_now_ms()}-{uuid.uuid4().hex[:8]}"358 359 def predict_to_ply(self, image_path: str | Path) -> PredictionOutputs:360 """Run SHARP inference and export a .ply file."""361 image_path = Path(image_path)362 if not image_path.exists():363 raise FileNotFoundError(f"Image does not exist: {image_path}")364 365 device = _select_device(self.device_preference)366 predictor = self._get_predictor(device)367 368 image_np, _, f_px = io.load_rgb(image_path)369 height, width = image_np.shape[:2]370 371 with torch.no_grad():372 gaussians = predict_image(predictor, image_np, f_px, device)373 374 stem = self._make_output_stem(image_path)375 ply_path = self.outputs_dir / f"{stem}.ply"376 377 # save_ply expects (height, width).378 save_ply(gaussians, f_px, (height, width), ply_path)379 380 # SceneMetaData expects (width, height) for resolution.381 metadata_for_render = SceneMetaData(382 focal_length_px=float(f_px),383 resolution_px=(int(width), int(height)),384 color_space="linearRGB",385 )386 387 self._maybe_move_model_back_to_cpu()388 389 return PredictionOutputs(390 ply_path=ply_path,391 gaussians=gaussians,392 metadata_for_render=metadata_for_render,393 input_resolution_hw=(int(height), int(width)),394 focal_length_px=float(f_px),395 )396 397 def _render_video_impl(398 self,399 *,400 gaussians: Gaussians3D,401 metadata: SceneMetaData,402 output_path: Path,403 trajectory_type: TrajectoryType,404 num_frames: int,405 fps: int,406 output_long_side: int | None,407 ) -> Path:408 if not torch.cuda.is_available():409 raise RuntimeError("Rendering requires CUDA (gsplat).")410 411 if num_frames < 2:412 raise ValueError("num_frames must be >= 2")413 if fps < 1:414 raise ValueError("fps must be >= 1")415 416 # Keep aligned with upstream CLI pipeline where possible.417 if output_long_side is None and int(fps) == 30:418 params = camera.TrajectoryParams(419 type=trajectory_type,420 num_steps=int(num_frames),421 num_repeats=1,422 )423 with _patched_sharp_videowriter():424 sharp_render_gaussians(425 gaussians=gaussians,426 metadata=metadata,427 params=params,428 output_path=output_path,429 )430 depth_path = output_path.with_suffix(".depth.mp4")431 try:432 if depth_path.exists():433 depth_path.unlink()434 except Exception:435 pass436 return output_path437 438 # Adapted pipeline for custom output resolution / FPS.439 src_w, src_h = metadata.resolution_px440 src_f = float(metadata.focal_length_px)441 442 if output_long_side is None:443 out_w, out_h, out_f = src_w, src_h, src_f444 else:445 long_side = max(src_w, src_h)446 scale = float(output_long_side) / float(long_side)447 out_w = _make_even(max(2, int(round(src_w * scale))))448 out_h = _make_even(max(2, int(round(src_h * scale))))449 out_f = src_f * scale450 451 traj_params = camera.TrajectoryParams(452 type=trajectory_type,453 num_steps=int(num_frames),454 num_repeats=1,455 )456 457 device = torch.device("cuda")458 gaussians_cuda = gaussians.to(device)459 460 intrinsics = torch.tensor(461 [462 [out_f, 0.0, (out_w - 1) / 2.0, 0.0],463 [0.0, out_f, (out_h - 1) / 2.0, 0.0],464 [0.0, 0.0, 1.0, 0.0],465 [0.0, 0.0, 0.0, 1.0],466 ],467 device=device,468 dtype=torch.float32,469 )470 471 cam_model = camera.create_camera_model(472 gaussians_cuda,473 intrinsics,474 resolution_px=(out_w, out_h),475 lookat_mode=traj_params.lookat_mode,476 )477 478 trajectory = camera.create_eye_trajectory(479 gaussians_cuda,480 traj_params,481 resolution_px=(out_w, out_h),482 f_px=out_f,483 )484 485 renderer = GSplatRenderer(color_space=metadata.color_space)486 487 # IMPORTANT: Keep render_depth=True (avoids upstream AttributeError).488 video_writer = _PatchedVideoWriter(output_path, fps=float(fps), render_depth=True)489 490 for eye_position in trajectory:491 cam_info = cam_model.compute(eye_position)492 rendering = renderer(493 gaussians_cuda,494 extrinsics=cam_info.extrinsics[None].to(device),495 intrinsics=cam_info.intrinsics[None].to(device),496 image_width=cam_info.width,497 image_height=cam_info.height,498 )499 color = (rendering.color[0].permute(1, 2, 0) * 255.0).to(dtype=torch.uint8)500 depth = rendering.depth[0]501 video_writer.add_frame(color, depth)502 503 video_writer.close()504 505 depth_path = output_path.with_suffix(".depth.mp4")506 try:507 if depth_path.exists():508 depth_path.unlink()509 except Exception:510 pass511 512 return output_path513 514 def render_video(515 self,516 *,517 gaussians: Gaussians3D,518 metadata: SceneMetaData,519 output_stem: str,520 trajectory_type: TrajectoryType = "rotate_forward",521 num_frames: int = 60,522 fps: int = 30,523 output_long_side: int | None = None,524 ) -> Path:525 """Render a camera trajectory as an MP4 (CUDA-only)."""526 output_path = self.outputs_dir / f"{output_stem}.mp4"527 return self._render_video_impl(528 gaussians=gaussians,529 metadata=metadata,530 output_path=output_path,531 trajectory_type=trajectory_type,532 num_frames=num_frames,533 fps=fps,534 output_long_side=output_long_side,535 )536 537 def predict_and_maybe_render(538 self,539 image_path: str | Path,540 *,541 trajectory_type: TrajectoryType,542 num_frames: int,543 fps: int,544 output_long_side: int | None,545 render_video: bool = True,546 ) -> tuple[Path | None, Path]:547 """One-shot helper for the UI: returns (video_path, ply_path)."""548 pred = self.predict_to_ply(image_path)549 550 if not render_video:551 return None, pred.ply_path552 553 if not torch.cuda.is_available():554 return None, pred.ply_path555 556 output_stem = pred.ply_path.with_suffix("").name557 video_path = self.render_video(558 gaussians=pred.gaussians,559 metadata=pred.metadata_for_render,560 output_stem=output_stem,561 trajectory_type=trajectory_type,562 num_frames=num_frames,563 fps=fps,564 output_long_side=output_long_side,565 )566 return video_path, pred.ply_path567 568 569# -----------------------------------------------------------------------------570# Module-level entrypoints571# -----------------------------------------------------------------------------572 573DEFAULT_OUTPUTS_DIR: Final[Path] = _ensure_dir(Path(__file__).resolve().parent / "outputs")574 575_GLOBAL_MODEL: ModelWrapper | None = None576_GLOBAL_MODEL_INIT_LOCK: Final[threading.Lock] = threading.Lock()577 578 579def get_global_model(*, outputs_dir: str | Path = DEFAULT_OUTPUTS_DIR) -> ModelWrapper:580 global _GLOBAL_MODEL581 with _GLOBAL_MODEL_INIT_LOCK:582 if _GLOBAL_MODEL is None:583 _GLOBAL_MODEL = ModelWrapper(outputs_dir=outputs_dir)584 return _GLOBAL_MODEL585 586 587def predict_and_maybe_render(588 image_path: str | Path,589 *,590 trajectory_type: TrajectoryType,591 num_frames: int,592 fps: int,593 output_long_side: int | None,594 render_video: bool = True,595) -> tuple[Path | None, Path]:596 model = get_global_model()597 return model.predict_and_maybe_render(598 image_path,599 trajectory_type=trajectory_type,600 num_frames=num_frames,601 fps=fps,602 output_long_side=output_long_side,603 render_video=render_video,604 )605 606 607# -----------------------------------------------------------------------------608# GPU-wrapped entrypoint (Spaces or local)609# -----------------------------------------------------------------------------610 611 612def _create_spaces_gpu_wrapper(duration: int = 180):613 """Create a Spaces GPU-wrapped version of predict_and_maybe_render.614 615 This is called dynamically based on hardware configuration.616 """617 if spaces is not None and _SPACES_AVAILABLE:618 return spaces.GPU(duration=duration)(predict_and_maybe_render)619 return predict_and_maybe_render620 621 622# Default export: use local CUDA unless explicitly configured for Spaces623# The actual wrapper is created dynamically based on hardware_config624predict_and_maybe_render_gpu = predict_and_maybe_render625 626 627def configure_gpu_mode(use_spaces: bool = False, duration: int = 180) -> None:628 """Configure the GPU mode at runtime.629 630 Args:631 use_spaces: If True and spaces module available, use @spaces.GPU decorator632 duration: Duration for @spaces.GPU decorator (seconds)633 """634 global predict_and_maybe_render_gpu635 636 if use_spaces and _SPACES_AVAILABLE and spaces is not None:637 predict_and_maybe_render_gpu = spaces.GPU(duration=duration)(predict_and_maybe_render)638 else:639 predict_and_maybe_render_gpu = predict_and_maybe_render640 641 642def get_gpu_status() -> dict:643 """Get current GPU status information."""644 import torch645 646 status = {647 "cuda_available": torch.cuda.is_available(),648 "spaces_available": _SPACES_AVAILABLE,649 "device_count": torch.cuda.device_count() if torch.cuda.is_available() else 0,650 "devices": [],651 }652 653 if torch.cuda.is_available():654 for i in range(torch.cuda.device_count()):655 props = torch.cuda.get_device_properties(i)656 status["devices"].append({657 "index": i,658 "name": props.name,659 "total_memory_gb": round(props.total_memory / (1024**3), 2),660 "compute_capability": f"{props.major}.{props.minor}",661 })662 663 return status664