|
from dataclasses import dataclass |
|
from typing import Optional |
|
|
|
import torch |
|
import torch.nn as nn |
|
from einops import rearrange |
|
from jaxtyping import Float |
|
from torch import Tensor |
|
|
|
from sf3d.models.tokenizers.dinov2 import Dinov2Model |
|
from sf3d.models.transformers.attention import Modulation |
|
from sf3d.models.utils import BaseModule |
|
|
|
|
|
class DINOV2SingleImageTokenizer(BaseModule): |
|
@dataclass |
|
class Config(BaseModule.Config): |
|
pretrained_model_name_or_path: str = "facebook/dinov2-large" |
|
width: int = 512 |
|
height: int = 512 |
|
modulation_cond_dim: int = 768 |
|
|
|
cfg: Config |
|
|
|
def configure(self) -> None: |
|
self.model = Dinov2Model.from_pretrained(self.cfg.pretrained_model_name_or_path) |
|
|
|
for p in self.model.parameters(): |
|
p.requires_grad_(False) |
|
self.model.eval() |
|
|
|
self.model.set_gradient_checkpointing(False) |
|
|
|
|
|
modulations = [] |
|
for layer in self.model.encoder.layer: |
|
norm1_modulation = Modulation( |
|
self.model.config.hidden_size, |
|
self.cfg.modulation_cond_dim, |
|
zero_init=True, |
|
single_layer=True, |
|
) |
|
norm2_modulation = Modulation( |
|
self.model.config.hidden_size, |
|
self.cfg.modulation_cond_dim, |
|
zero_init=True, |
|
single_layer=True, |
|
) |
|
layer.register_ada_norm_modulation(norm1_modulation, norm2_modulation) |
|
modulations += [norm1_modulation, norm2_modulation] |
|
self.modulations = nn.ModuleList(modulations) |
|
|
|
self.register_buffer( |
|
"image_mean", |
|
torch.as_tensor([0.485, 0.456, 0.406]).reshape(1, 1, 3, 1, 1), |
|
persistent=False, |
|
) |
|
self.register_buffer( |
|
"image_std", |
|
torch.as_tensor([0.229, 0.224, 0.225]).reshape(1, 1, 3, 1, 1), |
|
persistent=False, |
|
) |
|
|
|
def forward( |
|
self, |
|
images: Float[Tensor, "B *N C H W"], |
|
modulation_cond: Optional[Float[Tensor, "B *N Cc"]], |
|
**kwargs, |
|
) -> Float[Tensor, "B *N Ct Nt"]: |
|
model = self.model |
|
|
|
packed = False |
|
if images.ndim == 4: |
|
packed = True |
|
images = images.unsqueeze(1) |
|
if modulation_cond is not None: |
|
assert modulation_cond.ndim == 2 |
|
modulation_cond = modulation_cond.unsqueeze(1) |
|
|
|
batch_size, n_input_views = images.shape[:2] |
|
images = (images - self.image_mean) / self.image_std |
|
out = model( |
|
rearrange(images, "B N C H W -> (B N) C H W"), |
|
modulation_cond=rearrange(modulation_cond, "B N Cc -> (B N) Cc") |
|
if modulation_cond is not None |
|
else None, |
|
) |
|
local_features = out.last_hidden_state |
|
local_features = local_features.permute(0, 2, 1) |
|
local_features = rearrange( |
|
local_features, "(B N) Ct Nt -> B N Ct Nt", B=batch_size |
|
) |
|
if packed: |
|
local_features = local_features.squeeze(1) |
|
|
|
return local_features |
|
|
|
def detokenize(self, *args, **kwargs): |
|
raise NotImplementedError |
|
|
