| """Pure-DVSM render utilities (no viser dependency). |
| |
| Splits the model-loading and NeuralRenderer reconstruction code out of |
| viewer.py so that headless tests and offline rendering scripts can use it |
| without pulling in viser. |
| """ |
| from __future__ import annotations |
|
|
| from pathlib import Path |
|
|
| import torch |
|
|
| from dvsm.config import cfg, update_config |
| from dvsm.models.DVSM import DecoderViewSynthesisModel, NeuralRenderer |
| from dvsm.models.pose_preprocessor import PosePreprocessor |
|
|
|
|
| |
| |
| |
| WORLD_DVSM_TO_VISER = torch.tensor( |
| [[1.0, 0.0, 0.0], [0.0, -1.0, 0.0], [0.0, 0.0, -1.0]], dtype=torch.float32, |
| ) |
|
|
|
|
| def apply_pp_single(pp: PosePreprocessor, c2w_r, c2w_t, apply_scale: bool = False): |
| """DVSM raw pose -> viser world pose: rotate + translate via pp, then Rw. |
| |
| By default we OMIT pp.s. Reason: pp.s shrinks the scene to a unit sphere, |
| which makes viser's orbit/pan/zoom feel too fast (controls are tuned for |
| a "natural" metric scale). Keeping the original metric for viser and |
| applying pp.s only inside the renderer gives natural control feel. |
| |
| c2w_r_n = pp.r @ c2w_r |
| c2w_t_n = pp.r @ c2w_t + pp.t (pp.s applied iff apply_scale) |
| c2w_r_viser = Rw @ c2w_r_n |
| c2w_t_viser = Rw @ c2w_t_n |
| """ |
| r = pp.r[0] |
| t = pp.t[0, 0] |
| r_n = r @ c2w_r |
| t_n = r @ c2w_t + t |
| if apply_scale: |
| s = float(pp.s.reshape(-1)[0]) |
| t_n = t_n * s |
| Rw = WORLD_DVSM_TO_VISER.to(r_n.dtype) |
| r_v = Rw @ r_n |
| t_v = Rw @ t_n |
| return r_v, t_v |
|
|
|
|
| def load_model(ckpt_root: Path | None = None, ckpt_name: str = "latest.pt", |
| cmd_lst: list[str] | None = None, |
| hf_variant: str | None = None) -> DecoderViewSynthesisModel: |
| if hf_variant is not None: |
| model = DecoderViewSynthesisModel.from_hf_pretrained(hf_variant) |
| return model.cuda().eval() |
| assert ckpt_root is not None, "Provide either ckpt_root or hf_variant" |
| update_config([ckpt_root / "config.yaml"], cmd_lst or []) |
| cfg.model.use_compile = True |
| ckpt_path = ckpt_root / ckpt_name |
| assert ckpt_path.exists(), ckpt_path |
| model = DecoderViewSynthesisModel(**cfg.model).cuda().eval() |
| ckpt = torch.load(ckpt_path, weights_only=False, map_location="cpu") |
| model.load_state_dict(ckpt["state_dict"]) |
| del ckpt |
| return model |
|
|
|
|
| def _kv_to_cuda(kv): |
| """Move baked fast-weights to CUDA, upcasting fp16 (how they're stored on |
| disk, see recon_scene._move_kv_to_cpu) back to fp32 for rendering.""" |
| if kv is None: |
| return None |
| if isinstance(kv, torch.Tensor): |
| return (kv.float() if kv.dtype == torch.float16 else kv).cuda() |
| if isinstance(kv, dict): |
| return {k: _kv_to_cuda(v) for k, v in kv.items()} |
| if isinstance(kv, (tuple, list)): |
| return type(kv)(_kv_to_cuda(v) for v in kv) |
| return kv |
|
|
|
|
| def build_renderer(model: DecoderViewSynthesisModel, artifact: dict) -> NeuralRenderer: |
| """Re-materialize a NeuralRenderer from an artifact saved by recon_scene.py.""" |
| pp = PosePreprocessor.__new__(PosePreprocessor) |
| pp.r = artifact["pose_preproc"]["r"].cuda() |
| pp.t = artifact["pose_preproc"]["t"].cuda() |
| pp.s = artifact["pose_preproc"]["s"].cuda() |
| kv_lst = [_kv_to_cuda(kv) for kv in artifact["kv_lst"]] |
| |
| |
| |
| return NeuralRenderer( |
| pose_preprocessor=pp, |
| model=model, |
| kv_lst=kv_lst, |
| ) |
|
|
