| 11 | |
| 12 | |
| 13 | class SplatModel(StaticEncoder): |
| 14 | def __init__(self, opt: Options, **model_kwargs): |
| 15 | super().__init__(opt) |
| 16 | self.opt = opt |
| 17 | self.model = SplatPredictor(opt, **model_kwargs) |
| 18 | if hasattr(opt, 'compile') and opt.compile: |
| 19 | self.model = torch.compile(self.model) |
| 20 | self.gaussian_renderer = gaussian_renderer_dynamic.render |
| 21 | self.background = torch.tensor(opt.background_color, dtype=torch.float32, device="cuda") |
| 22 | self.lpips_loss = None |
| 23 | |
| 24 | def load_state_dict(self, state_dict, strict=True): |
| 25 | # if opt.use_dino, remove missing keys related to condition_encoder |
| 26 | missing_keys, unexpected_keys = super().load_state_dict(state_dict, strict=strict) |
| 27 | if self.opt.use_dino: |
| 28 | missing_keys = [k for k in missing_keys if "condition_encoder" not in k] |
| 29 | return missing_keys, unexpected_keys |
| 30 | |
| 31 | def forward_gaussians(self, frames, depths, cond_times=None): |
| 32 | # frames: [B, V, C, H, W] |
| 33 | # return: gaussians: [B, N, D] |
| 34 | decoder_out = self.model(frames, depths, cond_times=cond_times) |
| 35 | return decoder_out |
| 36 | |
| 37 | def forward(self, data, step_ratio=0.0): |
| 38 | # data: [B, V, C, H, W] |
| 39 | input_frames = data['frames'] # [B, V, C, H, W], input features |
| 40 | input_depths = data['depths'] # [B, V, C, H, W], input features |
| 41 | timestamps = data['timestamps'] # [B, V], input timestamps |
| 42 | timestamps = torch.as_tensor(timestamps, dtype=torch.float32, device=input_frames.device) |
| 43 | timestamps = timestamps / timestamps[..., -1].unsqueeze(-1) |
| 44 | anchor_time = torch.tensor([0.0, 1.0], device=input_frames.device) |
| 45 | results = {} |
| 46 | decoder_out = self.forward_gaussians(input_frames, input_depths, timestamps) # dict |
| 47 | with autocast('cuda', enabled=False): |
| 48 | render_pkg = self.gaussian_renderer(decoder_out["pred_gs"], self.background, |
| 49 | opt=self.opt, timestamps=timestamps, |
| 50 | anchor_time=anchor_time, |
| 51 | override_opacity=False, training=self.training, |
| 52 | ) |
| 53 | |
| 54 | results['pred_frames'] = render_pkg["render"] |
| 55 | results['input_frames'] = input_frames |
| 56 | results['timestamps'] = timestamps |
| 57 | |
| 58 | return results |
| 59 | |