la1236567/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 26try:27 import spaces28except Exception: # pragma: no cover29 spaces = None # type: ignore[assignment]30 31try:32 # Prefer HF cache / Hub downloads (works with Spaces `preload_from_hub`).33 from huggingface_hub import hf_hub_download, try_to_load_from_cache34except Exception: # pragma: no cover35 hf_hub_download = None # type: ignore[assignment]36 try_to_load_from_cache = None # type: ignore[assignment]37 38from sharp.cli.predict import DEFAULT_MODEL_URL, predict_image39from sharp.cli.render import render_gaussians as sharp_render_gaussians40from sharp.models import PredictorParams, create_predictor41from sharp.utils import camera, io42from sharp.utils.gaussians import Gaussians3D, SceneMetaData, save_ply43from sharp.utils.gsplat import GSplatRenderer44 45TrajectoryType = Literal["swipe", "shake", "rotate", "rotate_forward"]46 47# -----------------------------------------------------------------------------48# Helpers49# -----------------------------------------------------------------------------50 51 52def _now_ms() -> int:53 return int(time.time() * 1000)54 55 56def _ensure_dir(path: Path) -> Path:57 path.mkdir(parents=True, exist_ok=True)58 return path59 60 61def _make_even(x: int) -> int:62 return x if x % 2 == 0 else x + 163 64 65def _select_device(preference: str = "auto") -> torch.device:66 """Select the best available device for inference (CPU/CUDA/MPS)."""67 if preference not in {"auto", "cpu", "cuda", "mps"}:68 raise ValueError("device preference must be one of: auto|cpu|cuda|mps")69 70 if preference == "cpu":71 return torch.device("cpu")72 if preference == "cuda":73 return torch.device("cuda" if torch.cuda.is_available() else "cpu")74 if preference == "mps":75 return torch.device("mps" if torch.backends.mps.is_available() else "cpu")76 77 # auto78 if torch.cuda.is_available():79 return torch.device("cuda")80 if torch.backends.mps.is_available():81 return torch.device("mps")82 return torch.device("cpu")83 84 85# -----------------------------------------------------------------------------86# Prediction outputs87# -----------------------------------------------------------------------------88 89 90@dataclass(frozen=True, slots=True)91class PredictionOutputs:92 """Outputs of SHARP inference (plus derived metadata for rendering)."""93 94 ply_path: Path95 gaussians: Gaussians3D96 metadata_for_render: SceneMetaData97 input_resolution_hw: tuple[int, int]98 focal_length_px: float99 100 101# -----------------------------------------------------------------------------102# Patch SHARP VideoWriter to properly close the optional depth writer103# -----------------------------------------------------------------------------104 105 106class _PatchedVideoWriter(io.VideoWriter):107 """Ensure depth writer is closed so files can be safely cleaned up."""108 109 def __init__(110 self, output_path: Path, fps: float = 30.0, render_depth: bool = True111 ) -> None:112 super().__init__(output_path, fps=fps, render_depth=render_depth)113 # Ensure attribute exists for downstream code paths.114 if not hasattr(self, "depth_writer"):115 self.depth_writer = None # type: ignore[attribute-defined-outside-init]116 117 def close(self):118 super().close()119 depth_writer = getattr(self, "depth_writer", None)120 try:121 if depth_writer is not None:122 depth_writer.close()123 except Exception:124 pass125 126 127@contextmanager128def _patched_sharp_videowriter():129 """Temporarily patch `sharp.utils.io.VideoWriter` used by `sharp.cli.render`."""130 original = io.VideoWriter131 io.VideoWriter = _PatchedVideoWriter # type: ignore[assignment]132 try:133 yield134 finally:135 io.VideoWriter = original # type: ignore[assignment]136 137 138# -----------------------------------------------------------------------------139# Model wrapper140# -----------------------------------------------------------------------------141 142 143class ModelWrapper:144 """Cached SHARP model wrapper for Gradio/Spaces."""145 146 def __init__(147 self,148 *,149 outputs_dir: str | Path = "outputs",150 checkpoint_url: str = DEFAULT_MODEL_URL,151 checkpoint_path: str | Path | None = None,152 device_preference: str = "auto",153 keep_model_on_device: bool | None = None,154 hf_repo_id: str | None = None,155 hf_filename: str | None = None,156 hf_revision: str | None = None,157 ) -> None:158 self.outputs_dir = _ensure_dir(Path(outputs_dir))159 self.checkpoint_url = checkpoint_url160 161 env_ckpt = os.getenv("SHARP_CHECKPOINT_PATH") or os.getenv("SHARP_CHECKPOINT")162 if checkpoint_path:163 self.checkpoint_path = Path(checkpoint_path)164 elif env_ckpt:165 self.checkpoint_path = Path(env_ckpt)166 else:167 self.checkpoint_path = None168 169 # Optional Hugging Face Hub fallback (useful when direct CDN download fails).170 self.hf_repo_id = hf_repo_id or os.getenv("SHARP_HF_REPO_ID", "apple/Sharp")171 self.hf_filename = hf_filename or os.getenv(172 "SHARP_HF_FILENAME", "sharp_2572gikvuh.pt"173 )174 self.hf_revision = hf_revision or os.getenv("SHARP_HF_REVISION") or None175 176 self.device_preference = device_preference177 178 # For ZeroGPU, it's safer to not keep large tensors on CUDA across calls.179 if keep_model_on_device is None:180 keep_env = (181 os.getenv("SHARP_KEEP_MODEL_ON_DEVICE")182 )183 self.keep_model_on_device = keep_env == "1"184 else:185 self.keep_model_on_device = keep_model_on_device186 187 self._lock = threading.RLock()188 self._predictor: torch.nn.Module | None = None189 self._predictor_device: torch.device | None = None190 self._state_dict: dict | None = None191 192 def has_cuda(self) -> bool:193 return torch.cuda.is_available()194 195 def _load_state_dict(self) -> dict:196 with self._lock:197 if self._state_dict is not None:198 return self._state_dict199 200 # 1) Explicit local checkpoint path201 if self.checkpoint_path is not None:202 try:203 self._state_dict = torch.load(204 self.checkpoint_path,205 weights_only=True,206 map_location="cpu",207 )208 return self._state_dict209 except Exception as e:210 raise RuntimeError(211 "Failed to load SHARP checkpoint from local path.\n\n"212 f"Path:\n {self.checkpoint_path}\n\n"213 f"Original error:\n {type(e).__name__}: {e}"214 ) from e215 216 # 2) HF cache (no-network): best match for Spaces `preload_from_hub`.217 hf_cache_error: Exception | None = None218 if try_to_load_from_cache is not None:219 try:220 cached = try_to_load_from_cache(221 repo_id=self.hf_repo_id,222 filename=self.hf_filename,223 revision=self.hf_revision,224 repo_type="model",225 )226 except TypeError:227 cached = try_to_load_from_cache(self.hf_repo_id, self.hf_filename) # type: ignore[misc]228 229 try:230 if isinstance(cached, str) and Path(cached).exists():231 self._state_dict = torch.load(232 cached, weights_only=True, map_location="cpu"233 )234 return self._state_dict235 except Exception as e:236 hf_cache_error = e237 238 # 3) HF Hub download (reuse cache when available; may download otherwise).239 hf_error: Exception | None = None240 if hf_hub_download is not None:241 # Attempt "local only" mode if supported (avoids network).242 try:243 import inspect244 245 if "local_files_only" in inspect.signature(hf_hub_download).parameters:246 ckpt_path = hf_hub_download(247 repo_id=self.hf_repo_id,248 filename=self.hf_filename,249 revision=self.hf_revision,250 local_files_only=True,251 )252 if Path(ckpt_path).exists():253 self._state_dict = torch.load(254 ckpt_path, weights_only=True, map_location="cpu"255 )256 return self._state_dict257 except Exception:258 pass259 260 try:261 ckpt_path = hf_hub_download(262 repo_id=self.hf_repo_id,263 filename=self.hf_filename,264 revision=self.hf_revision,265 )266 self._state_dict = torch.load(267 ckpt_path,268 weights_only=True,269 map_location="cpu",270 )271 return self._state_dict272 except Exception as e:273 hf_error = e274 275 # 4) Default upstream CDN (torch hub cache). Last resort.276 url_error: Exception | None = None277 try:278 self._state_dict = torch.hub.load_state_dict_from_url(279 self.checkpoint_url,280 progress=True,281 map_location="cpu",282 )283 return self._state_dict284 except Exception as e:285 url_error = e286 287 # If we got here: all options failed.288 hint_lines = [289 "Failed to load SHARP checkpoint.",290 "",291 "Tried (in order):",292 f" 1) HF cache (preload_from_hub): repo_id={self.hf_repo_id}, filename={self.hf_filename}, revision={self.hf_revision or 'None'}",293 f" 2) HF Hub download: repo_id={self.hf_repo_id}, filename={self.hf_filename}, revision={self.hf_revision or 'None'}",294 f" 3) URL (torch hub): {self.checkpoint_url}",295 "",296 "If network access is restricted, set a local checkpoint path:",297 " - SHARP_CHECKPOINT_PATH=/path/to/sharp_2572gikvuh.pt",298 "",299 "Original errors:",300 ]301 if try_to_load_from_cache is None:302 hint_lines.append(" HF cache: huggingface_hub not installed")303 elif hf_cache_error is not None:304 hint_lines.append(305 f" HF cache: {type(hf_cache_error).__name__}: {hf_cache_error}"306 )307 else:308 hint_lines.append(" HF cache: (not found in cache)")309 310 if hf_hub_download is None:311 hint_lines.append(" HF download: huggingface_hub not installed")312 else:313 hint_lines.append(f" HF download: {type(hf_error).__name__}: {hf_error}")314 315 hint_lines.append(f" URL: {type(url_error).__name__}: {url_error}")316 317 raise RuntimeError("\n".join(hint_lines))318 319 def _get_predictor(self, device: torch.device) -> torch.nn.Module:320 with self._lock:321 if self._predictor is None:322 state_dict = self._load_state_dict()323 predictor = create_predictor(PredictorParams())324 predictor.load_state_dict(state_dict)325 predictor.eval()326 self._predictor = predictor327 self._predictor_device = torch.device("cpu")328 329 assert self._predictor is not None330 assert self._predictor_device is not None331 332 if self._predictor_device != device:333 self._predictor.to(device)334 self._predictor_device = device335 336 return self._predictor337 338 def _maybe_move_model_back_to_cpu(self) -> None:339 if self.keep_model_on_device:340 return341 with self._lock:342 if self._predictor is not None and self._predictor_device is not None:343 if self._predictor_device.type != "cpu":344 self._predictor.to("cpu")345 self._predictor_device = torch.device("cpu")346 if torch.cuda.is_available():347 torch.cuda.empty_cache()348 349 def _make_output_stem(self, input_path: Path) -> str:350 return f"{input_path.stem}-{_now_ms()}-{uuid.uuid4().hex[:8]}"351 352 def predict_to_ply(self, image_path: str | Path) -> PredictionOutputs:353 """Run SHARP inference and export a .ply file."""354 image_path = Path(image_path)355 if not image_path.exists():356 raise FileNotFoundError(f"Image does not exist: {image_path}")357 358 device = _select_device(self.device_preference)359 predictor = self._get_predictor(device)360 361 image_np, _, f_px = io.load_rgb(image_path)362 height, width = image_np.shape[:2]363 364 with torch.no_grad():365 gaussians = predict_image(predictor, image_np, f_px, device)366 367 stem = self._make_output_stem(image_path)368 ply_path = self.outputs_dir / f"{stem}.ply"369 370 # save_ply expects (height, width).371 save_ply(gaussians, f_px, (height, width), ply_path)372 373 # SceneMetaData expects (width, height) for resolution.374 metadata_for_render = SceneMetaData(375 focal_length_px=float(f_px),376 resolution_px=(int(width), int(height)),377 color_space="linearRGB",378 )379 380 self._maybe_move_model_back_to_cpu()381 382 return PredictionOutputs(383 ply_path=ply_path,384 gaussians=gaussians,385 metadata_for_render=metadata_for_render,386 input_resolution_hw=(int(height), int(width)),387 focal_length_px=float(f_px),388 )389 390 def _render_video_impl(391 self,392 *,393 gaussians: Gaussians3D,394 metadata: SceneMetaData,395 output_path: Path,396 trajectory_type: TrajectoryType,397 num_frames: int,398 fps: int,399 output_long_side: int | None,400 ) -> Path:401 if not torch.cuda.is_available():402 raise RuntimeError("Rendering requires CUDA (gsplat).")403 404 if num_frames < 2:405 raise ValueError("num_frames must be >= 2")406 if fps < 1:407 raise ValueError("fps must be >= 1")408 409 # Keep aligned with upstream CLI pipeline where possible.410 if output_long_side is None and int(fps) == 30:411 params = camera.TrajectoryParams(412 type=trajectory_type,413 num_steps=int(num_frames),414 num_repeats=1,415 )416 with _patched_sharp_videowriter():417 sharp_render_gaussians(418 gaussians=gaussians,419 metadata=metadata,420 params=params,421 output_path=output_path,422 )423 depth_path = output_path.with_suffix(".depth.mp4")424 try:425 if depth_path.exists():426 depth_path.unlink()427 except Exception:428 pass429 return output_path430 431 # Adapted pipeline for custom output resolution / FPS.432 src_w, src_h = metadata.resolution_px433 src_f = float(metadata.focal_length_px)434 435 if output_long_side is None:436 out_w, out_h, out_f = src_w, src_h, src_f437 else:438 long_side = max(src_w, src_h)439 scale = float(output_long_side) / float(long_side)440 out_w = _make_even(max(2, int(round(src_w * scale))))441 out_h = _make_even(max(2, int(round(src_h * scale))))442 out_f = src_f * scale443 444 traj_params = camera.TrajectoryParams(445 type=trajectory_type,446 num_steps=int(num_frames),447 num_repeats=1,448 )449 450 device = torch.device("cuda")451 gaussians_cuda = gaussians.to(device)452 453 intrinsics = torch.tensor(454 [455 [out_f, 0.0, (out_w - 1) / 2.0, 0.0],456 [0.0, out_f, (out_h - 1) / 2.0, 0.0],457 [0.0, 0.0, 1.0, 0.0],458 [0.0, 0.0, 0.0, 1.0],459 ],460 device=device,461 dtype=torch.float32,462 )463 464 cam_model = camera.create_camera_model(465 gaussians_cuda,466 intrinsics,467 resolution_px=(out_w, out_h),468 lookat_mode=traj_params.lookat_mode,469 )470 471 trajectory = camera.create_eye_trajectory(472 gaussians_cuda,473 traj_params,474 resolution_px=(out_w, out_h),475 f_px=out_f,476 )477 478 renderer = GSplatRenderer(color_space=metadata.color_space)479 480 # IMPORTANT: Keep render_depth=True (avoids upstream AttributeError).481 video_writer = _PatchedVideoWriter(output_path, fps=float(fps), render_depth=True)482 483 for eye_position in trajectory:484 cam_info = cam_model.compute(eye_position)485 rendering = renderer(486 gaussians_cuda,487 extrinsics=cam_info.extrinsics[None].to(device),488 intrinsics=cam_info.intrinsics[None].to(device),489 image_width=cam_info.width,490 image_height=cam_info.height,491 )492 color = (rendering.color[0].permute(1, 2, 0) * 255.0).to(dtype=torch.uint8)493 depth = rendering.depth[0]494 video_writer.add_frame(color, depth)495 496 video_writer.close()497 498 depth_path = output_path.with_suffix(".depth.mp4")499 try:500 if depth_path.exists():501 depth_path.unlink()502 except Exception:503 pass504 505 return output_path506 507 def render_video(508 self,509 *,510 gaussians: Gaussians3D,511 metadata: SceneMetaData,512 output_stem: str,513 trajectory_type: TrajectoryType = "rotate_forward",514 num_frames: int = 60,515 fps: int = 30,516 output_long_side: int | None = None,517 ) -> Path:518 """Render a camera trajectory as an MP4 (CUDA-only)."""519 output_path = self.outputs_dir / f"{output_stem}.mp4"520 return self._render_video_impl(521 gaussians=gaussians,522 metadata=metadata,523 output_path=output_path,524 trajectory_type=trajectory_type,525 num_frames=num_frames,526 fps=fps,527 output_long_side=output_long_side,528 )529 530 def predict_and_maybe_render(531 self,532 image_path: str | Path,533 *,534 trajectory_type: TrajectoryType,535 num_frames: int,536 fps: int,537 output_long_side: int | None,538 render_video: bool = True,539 ) -> tuple[Path | None, Path]:540 """One-shot helper for the UI: returns (video_path, ply_path)."""541 pred = self.predict_to_ply(image_path)542 543 if not render_video:544 return None, pred.ply_path545 546 if not torch.cuda.is_available():547 return None, pred.ply_path548 549 output_stem = pred.ply_path.with_suffix("").name550 video_path = self.render_video(551 gaussians=pred.gaussians,552 metadata=pred.metadata_for_render,553 output_stem=output_stem,554 trajectory_type=trajectory_type,555 num_frames=num_frames,556 fps=fps,557 output_long_side=output_long_side,558 )559 return video_path, pred.ply_path560 561 562# -----------------------------------------------------------------------------563# ZeroGPU entrypoints564# -----------------------------------------------------------------------------565#566# IMPORTANT: Do NOT decorate bound instance methods with `@spaces.GPU` on ZeroGPU.567# The wrapper uses multiprocessing queues and pickles args/kwargs. If `self` is568# included, Python will try to pickle the whole instance. ModelWrapper contains569# a threading.RLock (not pickleable) and the model itself should not be pickled.570#571# Expose module-level functions that accept only pickleable arguments and572# create/cache the ModelWrapper inside the GPU worker process.573 574DEFAULT_OUTPUTS_DIR: Final[Path] = _ensure_dir(Path(__file__).resolve().parent / "outputs")575 576_GLOBAL_MODEL: ModelWrapper | None = None577_GLOBAL_MODEL_INIT_LOCK: Final[threading.Lock] = threading.Lock()578 579 580def get_global_model(*, outputs_dir: str | Path = DEFAULT_OUTPUTS_DIR) -> ModelWrapper:581 global _GLOBAL_MODEL582 with _GLOBAL_MODEL_INIT_LOCK:583 if _GLOBAL_MODEL is None:584 _GLOBAL_MODEL = ModelWrapper(outputs_dir=outputs_dir)585 return _GLOBAL_MODEL586 587 588def predict_and_maybe_render(589 image_path: str | Path,590 *,591 trajectory_type: TrajectoryType,592 num_frames: int,593 fps: int,594 output_long_side: int | None,595 render_video: bool = True,596) -> tuple[Path | None, Path]:597 model = get_global_model()598 return model.predict_and_maybe_render(599 image_path,600 trajectory_type=trajectory_type,601 num_frames=num_frames,602 fps=fps,603 output_long_side=output_long_side,604 render_video=render_video,605 )606 607 608# Export the GPU-wrapped callable (or a no-op wrapper locally).609if spaces is not None:610 predict_and_maybe_render_gpu = spaces.GPU(duration=180)(predict_and_maybe_render)611else: # pragma: no cover612 predict_and_maybe_render_gpu = predict_and_maybe_render613