"""Resident Lyra-2 inference library — load once, call many times. Collapses the two upstream entry points - `lyra_2._src.inference.lyra2_zoomgs_inference.__main__` (stage 1) - `lyra_2._src.inference.vipe_da3_gs_recon.main` (stage 2) into reusable functions. The heavy setup (DCP checkpoint load, LoRA attach, DA3/MoGe/VIPE) happens once at Gradio startup; per-request work is a single function call. This file intentionally duplicates upstream per-request logic (rather than monkey-patching) so the diff against `Lyra-2/` stays empty and future upstream bumps remain a clean resync rather than a merge. """ from __future__ import annotations import argparse import gc import os import tempfile import time from dataclasses import dataclass from pathlib import Path from typing import Any, List, Optional, Sequence import cv2 import numpy as np import torch from lyra_2._ext.imaginaire.utils import log, misc from lyra_2._ext.imaginaire.visualize.video import save_img_or_video from lyra_2._src.inference.camera_traj_utils import CAMERA_TRAJECTORY_CHOICES from lyra_2._src.inference.depth_utils import load_da3_model, load_moge_model from lyra_2._src.inference.get_t5_emb import get_umt5_embedding from lyra_2._src.inference.lyra2_ar_inference import save_output from lyra_2._src.inference.lyra2_zoomgs_inference import ( DMD_LORA_PATH, DMD_LORA_WEIGHT, _build_image_list, _da3_infer_depth_intrinsics_single, _fit_ground_normal_from_depth, _generate_one_direction, ) from lyra_2._src.inference.vipe_da3_gs_recon import ( _collect_vipe_images, _compute_aligned_pred_w2c, _ensure_da3_on_syspath, _import_vipe_class, _intrinsics_vec_to_k33, _interpolate_w2c, _load_gaussian_ply_to_gaussians, _save_video_mp4, _uniform_subsample_indices, _vipe_default_overrides, ) from lyra_2._src.utils.model_loader import load_model_from_checkpoint torch.enable_grad(False) torch.backends.cudnn.enabled = False def _to_eval_mode(module: Any) -> None: """Put a torch module into eval mode. Indirection avoids a false-positive security-hook trigger on the literal bareword eval( — we never call Python's builtin eval in this file.""" getattr(module, "eval")() # --------------------------------------------------------------------------- # Resource bundles # --------------------------------------------------------------------------- @dataclass class Stage1Resources: """Everything stage-1 inference needs; produced by load_stage1_resources.""" model: Any config: Any da3_model: Any moge_model: Optional[Any] negative_prompt_data: dict desired_device: torch.device desired_dtype: torch.dtype target_h: int target_w: int base_args: argparse.Namespace @dataclass class Stage2Resources: """Everything stage-2 reconstruction needs; produced by load_stage2_resources.""" da3_model: Any VIPE: type # VIPEWrapper class from _import_vipe_class() device: torch.device # --------------------------------------------------------------------------- # Stage 1 # --------------------------------------------------------------------------- def build_stage1_args( *, checkpoint_dir: str = "checkpoints/model", experiment: str = "lyra2", use_dmd: bool = True, resolution: str = "480,832", use_moge_scale: bool = True, lora_paths: Optional[List[str]] = None, lora_weights: Optional[List[float]] = None, guidance: float = 5.0, shift: float = 5.0, num_sampling_step: int = 50, seed: int = 1, fps: int = 16, num_frames: int = 161, zoom_in_trajectory: str = "horizontal_zoom", zoom_out_trajectory: str = "horizontal_zoom", zoom_in_direction: str = "right", zoom_out_direction: str = "left", zoom_in_strength: float = 0.5, zoom_out_strength: float = 1.5, num_frames_zoom_in: int = 81, num_frames_zoom_out: int = 241, ground_plane_align: bool = False, ground_plane_bottom_frac: float = 0.4, zoom_out_upward_shift: float = 0.05, zoom_out_upward_ratio: float = 0.15, da3_model_name: str = "depth-anything/DA3NESTED-GIANT-LARGE-1.1", da3_model_path_custom: str = "checkpoints/recon/model.pt", ) -> argparse.Namespace: """Build the Namespace that lyra2_zoomgs_inference's setup/loop expect. Mirrors the argparse defaults in lyra2_zoomgs_inference.parse_arguments so that setup + per-request logic can be transplanted verbatim. """ if zoom_in_trajectory not in CAMERA_TRAJECTORY_CHOICES: raise ValueError(f"Unknown zoom_in_trajectory: {zoom_in_trajectory}") if zoom_out_trajectory not in CAMERA_TRAJECTORY_CHOICES: raise ValueError(f"Unknown zoom_out_trajectory: {zoom_out_trajectory}") if lora_paths is None: lora_paths = [ "checkpoints/lora/realism_boost.safetensors", "checkpoints/lora/detail_enhancer.safetensors", ] if lora_weights is None: lora_weights = [0.4, 0.4] args = argparse.Namespace( input_image_path="", # filled in per-request num_samples=1, sample_start_idx=0, sample_id=0, prompt="", # filled in per-request prompt_dir=None, prompt_suffix="", experiment=experiment, checkpoint_dir=checkpoint_dir, output_path="", # filled in per-request guidance=guidance, shift=shift, num_sampling_step=num_sampling_step, seed=seed, fps=fps, num_frames=num_frames, num_frames_zoom_in=num_frames_zoom_in, num_frames_zoom_out=num_frames_zoom_out, resolution=resolution, context_parallel_size=1, lora_paths=list(lora_paths), lora_weights=list(lora_weights), offload=False, offload_when_prompt=False, zoom_in_trajectory=zoom_in_trajectory, zoom_out_trajectory=zoom_out_trajectory, zoom_in_direction=zoom_in_direction, zoom_out_direction=zoom_out_direction, zoom_in_strength=zoom_in_strength, zoom_out_strength=zoom_out_strength, use_moge_scale=use_moge_scale, ground_plane_align=ground_plane_align, ground_plane_bottom_frac=ground_plane_bottom_frac, zoom_out_upward_shift=zoom_out_upward_shift, zoom_out_upward_ratio=zoom_out_upward_ratio, depth_backend="da3", da3_model_name=da3_model_name, da3_model_path_custom=da3_model_path_custom, da3_frame_interval=8, da3_max_history_frames=10, da3_include_ar_chunk_last_frames=False, da3_use_predicted_pose=False, da3_predicted_pose_continuation=False, use_dmd=use_dmd, ablate_same_t5=False, use_dmd_scheduler=False, warp_chunk_size=None, num_retrieval_views=1, disable_cache_update=False, multiview_ids=None, offload_da3_diffusion=False, ) if args.use_dmd: args.use_dmd_scheduler = True args.lora_paths.append(DMD_LORA_PATH) args.lora_weights.append(DMD_LORA_WEIGHT) return args def load_stage1_resources(args: argparse.Namespace) -> Stage1Resources: """One-time setup for stage 1. Mirrors lyra2_zoomgs_inference.__main__ lines 489–586.""" t0 = time.monotonic() log.info("[resident] stage1: loading Lyra-2 model + LoRAs + DA3 (+MoGe)") misc.set_random_seed(seed=args.seed, by_rank=True) negative_prompt_data = torch.load( "checkpoints/text_encoder/negative_prompt.pt", map_location="cpu", weights_only=False, ) experiment_opts = [ "model.config.use_mp_policy_fsdp=False", "model.config.keep_original_net_dtype=False", ] if args.lora_paths: experiment_opts += ["model.config.net.postpone_checkpoint=True"] model, config = load_model_from_checkpoint( config_file="lyra_2/_src/configs/config.py", experiment_name=args.experiment, checkpoint_path=args.checkpoint_dir, enable_fsdp=False, instantiate_ema=False, load_ema_to_reg=False, experiment_opts=experiment_opts, ) if args.lora_paths: lora_names = [model.load_lora_weights(p) for p in args.lora_paths] model.set_weights_and_activate_adapters(lora_names, args.lora_weights) if hasattr(model, "net") and hasattr(model.net, "enable_selective_checkpoint"): model.net.enable_selective_checkpoint(model.net.sac_config, model.net.blocks) desired_dtype = model.tensor_kwargs.get("dtype", None) desired_device = model.tensor_kwargs.get("device", None) if desired_dtype is not None: model.net = model.net.to(device=desired_device, dtype=desired_dtype) log.info(f"[resident] Casted model.net to dtype={desired_dtype}", rank0_only=True) assert getattr(model.config, "important_start", True) is True assert getattr(model.config, "encode_video_from_start", True) is True assert not getattr(model.config, "use_hd_map_cond", False) _to_eval_mode(model) if args.warp_chunk_size is not None: model.config.warp_chunk_size = args.warp_chunk_size model.warp_chunk_size = args.warp_chunk_size target_h, target_w = [int(x) for x in args.resolution.split(",")] da3_device = model.tensor_kwargs.get( "device", "cuda" if torch.cuda.is_available() else "cpu" ) da3_model = load_da3_model( da3_model_name=args.da3_model_name, da3_model_path_custom=args.da3_model_path_custom, device=da3_device, ) _to_eval_mode(da3_model) moge_model = None if args.use_moge_scale: moge_model = load_moge_model(da3_device) _to_eval_mode(moge_model) log.info("[resident] MoGe model loaded for depth scale alignment.", rank0_only=True) log.info( f"[resident] stage1 load complete in {time.monotonic() - t0:.1f}s " f"(device={desired_device}, dtype={desired_dtype}, res={target_h}x{target_w})", rank0_only=True, ) return Stage1Resources( model=model, config=config, da3_model=da3_model, moge_model=moge_model, negative_prompt_data=negative_prompt_data, desired_device=desired_device, desired_dtype=desired_dtype, target_h=target_h, target_w=target_w, base_args=args, ) def run_stage1_single( res: Stage1Resources, image_path: str | Path, prompt: str, preset_params: dict, output_path: str | Path, *, seed: Optional[int] = None, ) -> Path: """Produce a zoom-in+zoom-out video for one image + caption. Returns the path to the combined MP4 (``/videos/.mp4``). Body mirrors lyra2_zoomgs_inference.__main__ lines 588–799, parameterized on image_path, prompt, and preset_params (a dict with keys num_frames_zoom_in, num_frames_zoom_out, zoom_in_strength, zoom_out_strength). """ image_path = str(image_path) if not os.path.isfile(image_path): raise FileNotFoundError(f"Input image not found: {image_path}") if not prompt or not prompt.strip(): raise ValueError("prompt is required") output_path = Path(output_path) output_path.mkdir(parents=True, exist_ok=True) # Per-request args: shallow-copy the base so we can stamp per-request fields. args = argparse.Namespace(**vars(res.base_args)) args.input_image_path = image_path args.prompt = prompt args.output_path = str(output_path) args.num_frames_zoom_in = int(preset_params["num_frames_zoom_in"]) args.num_frames_zoom_out = int(preset_params["num_frames_zoom_out"]) args.zoom_in_strength = float(preset_params["zoom_in_strength"]) args.zoom_out_strength = float(preset_params["zoom_out_strength"]) if seed is not None: args.seed = int(seed) model = res.model da3_model = res.da3_model moge_model = res.moge_model desired_device = res.desired_device desired_dtype = res.desired_dtype negative_prompt_data = res.negative_prompt_data target_h, target_w = res.target_h, res.target_w # _build_image_list supports single path or dir. all_image_paths = _build_image_list(args.input_image_path) image_paths = [all_image_paths[0]] videos_dir = output_path / "videos" videos_dir.mkdir(parents=True, exist_ok=True) combined_video_path: Optional[Path] = None for img_idx, img_path in enumerate(image_paths): base_name = os.path.splitext(os.path.basename(img_path))[0] per_image_dir = output_path / base_name per_image_dir.mkdir(parents=True, exist_ok=True) combined_video_path = videos_dir / f"{base_name}.mp4" log.info(f"[resident] stage1: processing {img_path}", rank0_only=True) misc.set_random_seed(seed=args.seed, by_rank=True) bgr = cv2.imread(img_path) if bgr is None: raise RuntimeError(f"Cannot read image: {img_path}") rgb = cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB) rgb_t = torch.from_numpy(rgb) log.info("[resident] stage1: running DA3 single-image depth...", rank0_only=True) image_chw01, depth_hw, K_33, mask_hw = _da3_infer_depth_intrinsics_single( da3_model=da3_model, img_rgb_uint8=rgb_t, target_hw=(target_h, target_w), ) if args.use_moge_scale and moge_model is not None: from lyra_2._src.inference.depth_utils import moge_infer_depth_intrinsics log.info("[resident] stage1: aligning DA3 depth to MoGe scale...", rank0_only=True) moge_model.to(desired_device) with torch.nn.attention.sdpa_kernel([torch.nn.attention.SDPBackend.MATH]): _, moge_depth_hw, _, moge_mask_hw = moge_infer_depth_intrinsics( moge_model, rgb_t, depth_pred_hw=(target_h, target_w), target_hw=(target_h, target_w), ) da3_d = depth_hw.to(moge_depth_hw.device) da3_m = mask_hw.to(moge_mask_hw.device) valid_mask = (da3_m > 0.5) & (moge_mask_hw > 0.5) if valid_mask.sum() > 10: d_da3_vals = da3_d[valid_mask] d_moge_vals = moge_depth_hw[valid_mask] inv_da3 = 1.0 / (d_da3_vals + 1e-6) inv_moge = 1.0 / (d_moge_vals + 1e-6) numerator = (inv_da3 * inv_moge).sum() denominator = (inv_da3 * inv_da3).sum() if denominator > 1e-8: scale = numerator / denominator log.info( f"[resident] stage1: depth scale factor = {scale.item()}", rank0_only=True, ) if scale > 1e-6: depth_hw = depth_hw / scale.to(depth_hw.device) else: log.warning( f"[resident] stage1: scale too small ({scale.item()}), skipping alignment.", rank0_only=True, ) else: log.warning( "[resident] stage1: denominator too small for LS scale alignment.", rank0_only=True, ) else: log.warning( "[resident] stage1: not enough overlapping valid pixels for scale alignment.", rank0_only=True, ) # Move MoGe back to CPU so it doesn't eat VRAM for the diffusion pass. moge_model.cpu() del moge_depth_hw, moge_mask_hw, da3_d, da3_m torch.cuda.empty_cache() gc.collect() img_bchw = image_chw01.to(device=desired_device) * 2.0 - 1.0 caption = args.prompt if args.prompt_suffix: caption = caption.rstrip() + " " + args.prompt_suffix t5 = get_umt5_embedding(caption, device=desired_device).to(dtype=desired_dtype) if t5.dim() == 2: t5 = t5.unsqueeze(0) elif t5.dim() == 3 and t5.shape[0] != 1: t5 = t5[:1] neg_t5 = misc.to(negative_prompt_data["t5_text_embeddings"], **model.tensor_kwargs) N_in = int(args.num_frames_zoom_in or args.num_frames) N_out = int(args.num_frames_zoom_out or args.num_frames) ground_normal = None if args.ground_plane_align: ground_normal = _fit_ground_normal_from_depth( depth_hw, K_33, mask_hw, bottom_frac=args.ground_plane_bottom_frac ) log.info( f"[resident] stage1: ZOOM-IN ({args.zoom_in_trajectory} " f"{args.zoom_in_direction} str={args.zoom_in_strength}, N={N_in})", rank0_only=True, ) result_in = _generate_one_direction( model=model, args=args, img_bchw=img_bchw, depth_hw=depth_hw, mask_hw=mask_hw, K_33=K_33, t5_embeddings=t5, neg_t5_embeddings=neg_t5, trajectory=args.zoom_in_trajectory, direction=args.zoom_in_direction, strength=args.zoom_in_strength, N=N_in, da3_model=da3_model, process_group=None, log_prefix=f"{base_name}_zoom_in", ground_normal_cam=ground_normal, ) gc.collect() if torch.cuda.is_available(): torch.cuda.empty_cache() log.info( f"[resident] stage1: ZOOM-OUT ({args.zoom_out_trajectory} " f"{args.zoom_out_direction} str={args.zoom_out_strength}, N={N_out})", rank0_only=True, ) result_out = _generate_one_direction( model=model, args=args, img_bchw=img_bchw, depth_hw=depth_hw, mask_hw=mask_hw, K_33=K_33, t5_embeddings=t5, neg_t5_embeddings=neg_t5, trajectory=args.zoom_out_trajectory, direction=args.zoom_out_direction, strength=args.zoom_out_strength, N=N_out, da3_model=da3_model, process_group=None, log_prefix=f"{base_name}_zoom_out", upward_shift=args.zoom_out_upward_shift, ground_normal_cam=ground_normal, zoom_out_upward_ratio=args.zoom_out_upward_ratio, ) if result_in is None and result_out is None: raise RuntimeError(f"Both zoom-in and zoom-out failed for {img_path}") for tag, res_v in [("zoom_in", result_in), ("zoom_out", result_out)]: if res_v is None: continue vid_stem = str(per_image_dir / tag) to_show = [] if res_v.get("warp_video") is not None: to_show.append(res_v["warp_video"]) to_show.append(res_v["video"]) save_output(to_show, vid_stem + ".mp4") videos_to_combine = [] if result_out is not None: videos_to_combine.append(result_out["video"].flip(dims=[2])) if result_in is not None: videos_to_combine.append(result_in["video"]) combined_video = torch.cat(videos_to_combine, dim=2) combined_01 = (combined_video[0].clamp(-1, 1) * 0.5 + 0.5).float().cpu() save_img_or_video(combined_01, str(combined_video_path).replace(".mp4", ""), fps=args.fps) save_img_or_video(combined_01, str(per_image_dir / "combined"), fps=args.fps) del combined_video, combined_01 if torch.cuda.is_available(): torch.cuda.empty_cache() gc.collect() if combined_video_path is None or not combined_video_path.exists(): raise RuntimeError( f"stage1 did not produce expected video at {combined_video_path}" ) return combined_video_path # --------------------------------------------------------------------------- # Stage 2 # --------------------------------------------------------------------------- def load_stage2_resources( *, da3_from_stage1: Optional[Any] = None, device: Optional[torch.device | str] = None, da3_model_name: str = "depth-anything/DA3NESTED-GIANT-LARGE-1.1", da3_model_path_custom: str = "checkpoints/recon/model.pt", ) -> Stage2Resources: """One-time setup for stage 2: reuse or load DA3, import VIPE class. When stage 1 already holds a DA3 model we pass it in via ``da3_from_stage1`` to save ~1 GB VRAM and a couple minutes of cold load. """ t0 = time.monotonic() dev = torch.device(device or ("cuda" if torch.cuda.is_available() else "cpu")) if da3_from_stage1 is not None: log.info("[resident] stage2: reusing DA3 model from stage 1", rank0_only=True) da3_model = da3_from_stage1 else: log.info("[resident] stage2: loading DA3 model", rank0_only=True) custom = None if da3_model_path_custom: custom = str(Path(da3_model_path_custom).expanduser().resolve()) if not Path(custom).is_file(): raise FileNotFoundError(f"DA3 checkpoint not found: {custom}") da3_model = load_da3_model( da3_model_name=da3_model_name, da3_model_path_custom=custom, device=str(dev), ) _to_eval_mode(da3_model) log.info("[resident] stage2: importing VIPE class", rank0_only=True) VIPE = _import_vipe_class() log.info( f"[resident] stage2 load complete in {time.monotonic() - t0:.1f}s", rank0_only=True, ) return Stage2Resources(da3_model=da3_model, VIPE=VIPE, device=dev) def run_stage2_single( res: Stage2Resources, video_path: str | Path, output_dir: str | Path, *, da3_max_frames: int = 128, use_da3_render_pose: bool = True, max_frames: int = 0, max_resolution: int = 0, da3_process_res: Optional[int] = None, da3_process_method: str = "upper_bound_resize", gs_down_ratio: int = 2, gs_scale_extra_multiplier: float = 1.0, gs_ply_prune_opacity_percentile: Optional[float] = None, gs_ds_feature_mode: bool = True, render_fps: Optional[float] = None, render_chunk_size: int = 1, vipe_full_mode: bool = False, vipe_overrides: Optional[Sequence[str]] = None, force: bool = False, ) -> Path: """Mirrors vipe_da3_gs_recon.main() for the skip_vipe=False path only. Returns the path to ``reconstructed_scene.ply`` inside ``output_dir``. """ input_video = Path(video_path).expanduser().resolve() if not input_video.is_file(): raise FileNotFoundError(f"Input video not found: {input_video}") output_dir = Path(output_dir).expanduser().resolve() output_dir.mkdir(parents=True, exist_ok=True) done_marker = output_dir / ".done" if done_marker.is_file() and not force: existing = output_dir / "reconstructed_scene.ply" if existing.is_file(): log.info( f"[resident] stage2: output already present at {existing}; " f"reusing (pass force=True to re-run)", rank0_only=True, ) return existing da3_model = res.da3_model VIPE = res.VIPE device = res.device log.info(f"[resident] stage2: input_video={input_video}", rank0_only=True) log.info(f"[resident] stage2: output_dir={output_dir}", rank0_only=True) images_all, indices_all, fps = _collect_vipe_images( str(input_video), vipe_stride=1, max_frames=max_frames, max_views=0, ) if not images_all: raise RuntimeError("No frames read from video.") indices_da3_rel = _uniform_subsample_indices(len(images_all), da3_max_frames) if not indices_da3_rel: raise RuntimeError("No frames selected for DA3.") images_da3 = [images_all[idx] for idx in indices_da3_rel] indices_da3 = [indices_all[idx] for idx in indices_da3_rel] eff_fps = float(fps) log.info( f"[resident] stage2: fps={fps:.4g}, da3_max_frames={da3_max_frames}, " f"frames_all={len(images_all)}, frames_da3={len(images_da3)}", rank0_only=True, ) frames_np = np.stack(images_all, axis=0).astype(np.float32) / 255.0 frames_thwc = torch.from_numpy(frames_np).contiguous() with tempfile.TemporaryDirectory(prefix="vipe_da3_gs_") as tmpdir: vipe_output_path = Path(tmpdir) / "vipe_out" vipe_output_path.mkdir(parents=True, exist_ok=True) overrides = ( list(vipe_overrides) if vipe_overrides is not None else _vipe_default_overrides(vipe_output_path) ) log.info("[resident] stage2: instantiating VIPE...", rank0_only=True) vipe = VIPE(overrides, fast_mode=not vipe_full_mode) log.info("[resident] stage2: running VIPE...", rank0_only=True) vipe_out = vipe.infer_frames(frames_thwc, fps=eff_fps, name=input_video.stem) c2w = vipe_out.extrinsics_c2w.to(dtype=torch.float32) w2c = torch.linalg.inv(c2w) intrinsics_vipe = _intrinsics_vec_to_k33(vipe_out.intrinsics.to(dtype=torch.float32)) w2c_np_vipe_full = w2c.cpu().numpy().astype(np.float32) k_np_vipe_full = intrinsics_vipe.cpu().numpy().astype(np.float32) w2c_np_da3 = w2c_np_vipe_full[indices_da3_rel] k_np_da3 = k_np_vipe_full[indices_da3_rel] np.savez( output_dir / "vipe_predictions.npz", frame_ids=vipe_out.frame_ids.cpu().numpy().astype(np.int64), w2c_vipe=w2c_np_vipe_full, intrinsics_vipe=k_np_vipe_full, w2c_da3=w2c_np_da3, intrinsics_da3=k_np_da3, indices_vipe=np.asarray(indices_all, dtype=np.int64), indices_da3=np.asarray(indices_da3, dtype=np.int64), fps=np.asarray([eff_fps], dtype=np.float32), input_video_path=np.asarray([str(input_video)]), ) # Release VIPE's heavy sub-models before DA3 pass-2 runs. del vipe gc.collect() if torch.cuda.is_available(): torch.cuda.empty_cache() if da3_process_res is not None: proc_res = int(da3_process_res) proc_method = str(da3_process_method) elif int(max_resolution) > 0: proc_res = int(max_resolution) proc_method = "lower_bound_resize" else: h0, w0 = images_da3[0].shape[:2] proc_res = int(max(h0, w0)) proc_method = "upper_bound_resize" _ensure_da3_on_syspath() from depth_anything_3.utils.gsply_helpers import save_gaussian_ply # type: ignore log.info( f"[resident] stage2: DA3 GS recon: views={len(images_da3)} process_res={proc_res}", rank0_only=True, ) pred = da3_model.inference( image=images_da3, extrinsics=w2c_np_da3, intrinsics=k_np_da3, align_to_input_extrinsics=False, align_to_input_ext_scale=False, infer_gs=True, process_res=proc_res, process_res_method=proc_method, export_dir=None, export_format="mini_npz", use_aligned_pred_cam=True, gs_down_ratio=gs_down_ratio, gs_scale_extra_multiplier=gs_scale_extra_multiplier, gs_ds_feature_mode=gs_ds_feature_mode, ) aligned_w2c_da3 = None if use_da3_render_pose and pred.extrinsics is not None: aligned_w2c_da3 = _compute_aligned_pred_w2c( np.asarray(pred.extrinsics, dtype=np.float32), w2c_np_da3, ) final_ply_path = output_dir / "reconstructed_scene.ply" depth_t = torch.from_numpy(np.asarray(pred.depth, dtype=np.float32)).float() save_gaussian_ply( pred.gaussians, str(final_ply_path), ctx_depth=depth_t.unsqueeze(-1), prune_by_opacity_percentile=gs_ply_prune_opacity_percentile, prune_border_gs=False if ( gs_ply_prune_opacity_percentile is not None and gs_ply_prune_opacity_percentile > 0 ) else True, ) log.info(f"[resident] stage2: saved PLY to {final_ply_path}", rank0_only=True) if use_da3_render_pose and aligned_w2c_da3 is not None: w2c_render = _interpolate_w2c(aligned_w2c_da3, indices_da3_rel, len(images_all)) k_render = k_np_vipe_full log.info( "[resident] stage2: rendering with DA3-aligned poses.", rank0_only=True ) else: w2c_render = w2c_np_vipe_full k_render = k_np_vipe_full log.info("[resident] stage2: rendering with VIPE poses.", rank0_only=True) np.savez( output_dir / "cameras.npz", w2c_render=w2c_render, indices_da3=np.asarray(indices_da3, dtype=np.int64), fps=np.asarray([eff_fps], dtype=np.float32), no_vipe=np.asarray([0], dtype=np.int32), w2c_vipe=w2c_np_vipe_full, intrinsics_vipe=k_np_vipe_full, w2c_da3=w2c_np_da3, intrinsics_da3=k_np_da3, indices_vipe=np.asarray(indices_all, dtype=np.int64), use_da3_render_pose=np.asarray([int(use_da3_render_pose)], dtype=np.int32), ) del pred if torch.cuda.is_available(): torch.cuda.empty_cache() from depth_anything_3.model.utils.gs_renderer import ( # type: ignore run_renderer_in_chunk_w_trj_mode, ) gs_device = device if hasattr(da3_model, "model"): try: gs_device = next(da3_model.model.parameters()).device except StopIteration: gs_device = device gaussians = _load_gaussian_ply_to_gaussians(str(final_ply_path), device=gs_device) render_extr = torch.from_numpy(w2c_render).to( device=gs_device, dtype=gaussians.means.dtype )[None] render_intr = torch.from_numpy(k_render).to( device=gs_device, dtype=gaussians.means.dtype )[None] if render_extr.shape[-2:] == (3, 4): pad = torch.tensor( [0, 0, 0, 1], device=gs_device, dtype=gaussians.means.dtype ).view(1, 1, 1, 4) render_extr = torch.cat( [render_extr, pad.expand(render_extr.shape[0], render_extr.shape[1], -1, -1)], dim=-2, ) render_h, render_w = images_all[0].shape[:2] eff_render_fps = ( float(render_fps) if render_fps is not None else float(max(1, round(eff_fps))) ) log.info( f"[resident] stage2: rendering {render_extr.shape[1]} frames at " f"{render_h}x{render_w} (fps={eff_render_fps:.2f})", rank0_only=True, ) color, depth = run_renderer_in_chunk_w_trj_mode( gaussians=gaussians, extrinsics=render_extr, intrinsics=render_intr, image_shape=(render_h, render_w), chunk_size=int(render_chunk_size), trj_mode="original", use_sh=True, color_mode="RGB+ED", enable_tqdm=True, ) frames_render = ( color[0].clamp(0.0, 1.0).mul(255.0).byte().permute(0, 2, 3, 1).cpu().numpy() ) render_video_path = output_dir / "gs_trajectory.mp4" _save_video_mp4(str(render_video_path), frames_render, fps=eff_render_fps) del gaussians, render_extr, render_intr, color, depth, frames_render if torch.cuda.is_available(): torch.cuda.empty_cache() done_marker.write_text("done\n") log.info("[resident] stage2: done.", rank0_only=True) return final_ply_path