| 183 | |
| 184 | |
| 185 | class ExportableLingBotDepth(nn.Module): |
| 186 | def __init__(self, model: MDMModel, num_tokens: int): |
| 187 | super().__init__() |
| 188 | self.model = model |
| 189 | self.num_tokens = int(num_tokens) |
| 190 | |
| 191 | def forward(self, image: torch.Tensor, depth: torch.Tensor) -> torch.Tensor: |
| 192 | output = self.model.forward( |
| 193 | image, |
| 194 | num_tokens=self.num_tokens, |
| 195 | depth=depth, |
| 196 | enable_depth_mask=False, |
| 197 | ) |
| 198 | return output["depth_reg"] |
| 199 | |
| 200 | |
| 201 | def load_capture(capture_dir: Path, device: torch.device) -> tuple[torch.Tensor, torch.Tensor, dict[str, Any]]: |