|
from dataclasses import dataclass, field |
|
from typing import List |
|
|
|
import torch |
|
import torch.nn as nn |
|
|
|
from sf3d.models.utils import BaseModule |
|
|
|
|
|
class LinearCameraEmbedder(BaseModule): |
|
@dataclass |
|
class Config(BaseModule.Config): |
|
in_channels: int = 25 |
|
out_channels: int = 768 |
|
conditions: List[str] = field(default_factory=list) |
|
|
|
cfg: Config |
|
|
|
def configure(self) -> None: |
|
self.linear = nn.Linear(self.cfg.in_channels, self.cfg.out_channels) |
|
|
|
def forward(self, **kwargs): |
|
cond_tensors = [] |
|
for cond_name in self.cfg.conditions: |
|
assert cond_name in kwargs |
|
cond = kwargs[cond_name] |
|
|
|
cond_tensors.append(cond.view(*cond.shape[:2], -1)) |
|
cond_tensor = torch.cat(cond_tensors, dim=-1) |
|
assert cond_tensor.shape[-1] == self.cfg.in_channels |
|
embedding = self.linear(cond_tensor) |
|
return embedding |
|
|