Spaces:
Running on Zero
Running on Zero
File size: 2,631 Bytes
1d998f6 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 | from dataclasses import dataclass
from typing import Literal
import torch
from ..types import LVSPM
from .encoder import Encoder
from .vit.lvspm import LVSPMNet
from ..ttt.config import TTTLayerCfg
from .common.pose_enc import pose_encoding_to_extri_intri
@dataclass
class LVSPMCfg:
patch_size: int
in_channels: int
transformer_d: int
transformer_d_head: int
transformer_n_layer: int
grad_checkpoint_every: int
use_flashatt: bool
qk_norm: bool
fix_attention: bool
rope_freq: float
max_rope_freq: int
compile_model: bool
@dataclass
class EncoderLVSPMCfg:
name: Literal["lvspm"]
lvspm: LVSPMCfg
ttt_layer: TTTLayerCfg
class EncoderLVSPM(Encoder[EncoderLVSPMCfg]):
def __init__(self, cfg: EncoderLVSPMCfg) -> None:
super().__init__(cfg)
self.view_predictor = LVSPMNet(
cfg=cfg.lvspm,
output_channel=3,
ttt_layer_cfg=cfg.ttt_layer,
)
def prefill_context(self, context: dict):
return self.view_predictor.prefill_context({"images": context["image"]})
@staticmethod
def decode_context_pose(prefill_state):
return EncoderLVSPM._decode_pose(
prefill_state.pose_encoding, prefill_state.height, prefill_state.width
)
@staticmethod
def _decode_pose(pose_encoding, height, width):
extrinsic, intrinsic = pose_encoding_to_extri_intri(
pose_encoding, (height, width),
)
final_row = torch.tensor(
[0, 0, 0, 1], dtype=extrinsic.dtype, device=extrinsic.device
).expand(extrinsic.size(0), extrinsic.size(1), 1, 4)
return torch.cat((extrinsic, final_row), dim=2), intrinsic
def predict_pose(self, context: dict):
prefill_state = self.prefill_context(context)
extrinsic, intrinsic = self.decode_context_pose(prefill_state)
return extrinsic, intrinsic, prefill_state
def render_targets(self, prefill_state, target: dict) -> torch.Tensor:
return self.view_predictor.render_targets(
prefill_state,
target["intrinsics"],
target["extrinsics"],
)
def forward(
self,
context: dict,
target: dict,
) -> LVSPM:
pred, pose_encoding = self.view_predictor(
context["image"], target["intrinsics"], target["extrinsics"]
)
extrinsic, intrinsic = self._decode_pose(
pose_encoding, *context["image"].shape[-2:]
)
return LVSPM(
pred,
extrinsic,
intrinsic,
pose_encoding,
)
|