333 lines
11 KiB
Python
333 lines
11 KiB
Python
import math
|
|
from typing import Optional
|
|
|
|
import torch
|
|
from huggingface_hub import hf_hub_download
|
|
from safetensors import safe_open
|
|
from torch import nn
|
|
from transformers.models.dinov3_vit.configuration_dinov3_vit import DINOv3ViTConfig
|
|
|
|
from src.model.dino import (
|
|
DINOv3ViTEmbeddings,
|
|
DINOv3ViTLayer,
|
|
DINOv3ViTLayerScale,
|
|
DINOv3ViTRopePositionEmbedding,
|
|
)
|
|
|
|
|
|
class TimestepEmbedder(nn.Module):
|
|
def __init__(self, hidden_size: int, frequency_embedding_size: int = 256):
|
|
super().__init__()
|
|
self.mlp = nn.Sequential(
|
|
nn.Linear(frequency_embedding_size, hidden_size),
|
|
nn.SiLU(),
|
|
nn.Linear(hidden_size, hidden_size),
|
|
)
|
|
self.frequency_embedding_size = frequency_embedding_size
|
|
|
|
@staticmethod
|
|
def timestep_embedding(t, dim, max_period=10000):
|
|
half = dim // 2
|
|
freqs = torch.exp(
|
|
-math.log(max_period) * torch.arange(start=0, end=half) / half
|
|
).to(t.device)
|
|
args = t[:, None] * freqs[None]
|
|
embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
|
|
if dim % 2:
|
|
embedding = torch.cat(
|
|
[embedding, torch.zeros_like(embedding[:, :1])], dim=-1
|
|
)
|
|
return embedding
|
|
|
|
def forward(self, t):
|
|
t_freq = self.timestep_embedding(t, self.frequency_embedding_size).to(
|
|
dtype=next(self.parameters()).dtype
|
|
)
|
|
t_emb = self.mlp(t_freq)
|
|
return t_emb
|
|
|
|
|
|
class LabelEmbedder(nn.Module):
|
|
def __init__(self, num_classes, hidden_size, dropout_prob):
|
|
super().__init__()
|
|
use_cfg_embedding = int(dropout_prob > 0)
|
|
self.embedding_table = nn.Embedding(
|
|
num_classes + use_cfg_embedding, hidden_size
|
|
)
|
|
self.num_classes = num_classes
|
|
self.dropout_prob = dropout_prob
|
|
|
|
def token_drop(self, labels, force_drop_ids=None):
|
|
if force_drop_ids is None:
|
|
drop_ids = torch.rand(labels.shape[0]) < self.dropout_prob
|
|
drop_ids = drop_ids.cuda()
|
|
drop_ids = drop_ids.to(labels.device)
|
|
else:
|
|
drop_ids = force_drop_ids == 1
|
|
labels = torch.where(drop_ids, self.num_classes, labels)
|
|
return labels
|
|
|
|
def forward(self, labels, train, force_drop_ids=None):
|
|
use_dropout = self.dropout_prob > 0
|
|
if (train and use_dropout) or (force_drop_ids is not None):
|
|
labels = self.token_drop(labels, force_drop_ids)
|
|
embeddings = self.embedding_table(labels)
|
|
return embeddings
|
|
|
|
|
|
class DinoConditionedLayer(DINOv3ViTLayer):
|
|
def __init__(self, config: DINOv3ViTConfig, is_encoder: bool = False):
|
|
super().__init__(config)
|
|
|
|
self.norm_cond = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
|
|
self.cond = nn.MultiheadAttention(
|
|
config.hidden_size,
|
|
config.num_attention_heads,
|
|
config.drop_path_rate,
|
|
batch_first=True,
|
|
)
|
|
self.layer_scale_cond = DINOv3ViTLayerScale(config)
|
|
|
|
# no init zeros!
|
|
if is_encoder:
|
|
nn.init.constant_(self.layer_scale_cond.lambda1, 0)
|
|
self.norm1.requires_grad_(False)
|
|
self.norm2.requires_grad_(False)
|
|
self.attention.requires_grad_(False)
|
|
self.mlp.requires_grad_(False)
|
|
self.layer_scale1.requires_grad_(False)
|
|
self.layer_scale2.requires_grad_(False)
|
|
|
|
def forward(
|
|
self,
|
|
hidden_states: torch.Tensor,
|
|
*,
|
|
conditioning_input: Optional[torch.Tensor] = None,
|
|
attention_mask: Optional[torch.Tensor] = None,
|
|
position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None,
|
|
**kwargs,
|
|
) -> torch.Tensor:
|
|
assert position_embeddings is not None
|
|
assert conditioning_input is not None
|
|
|
|
residual = hidden_states
|
|
hidden_states = self.norm1(hidden_states)
|
|
hidden_states, _ = self.attention(
|
|
hidden_states,
|
|
attention_mask=attention_mask,
|
|
position_embeddings=position_embeddings,
|
|
)
|
|
hidden_states = self.layer_scale1(hidden_states)
|
|
hidden_states = self.drop_path(hidden_states) + residual
|
|
|
|
residual = hidden_states
|
|
hidden_states = self.norm_cond(hidden_states)
|
|
hidden_states, _ = self.cond(
|
|
hidden_states, conditioning_input, conditioning_input
|
|
)
|
|
hidden_states = self.layer_scale_cond(hidden_states)
|
|
hidden_states = self.drop_path(hidden_states) + residual
|
|
|
|
residual = hidden_states
|
|
hidden_states = self.norm2(hidden_states)
|
|
hidden_states = self.mlp(hidden_states)
|
|
hidden_states = self.layer_scale2(hidden_states)
|
|
hidden_states = self.drop_path(hidden_states) + residual
|
|
|
|
return hidden_states
|
|
|
|
|
|
# class DinoV3ViTDecoder(nn.Module):
|
|
# def __init__(self, config: DINOv3ViTConfig):
|
|
# super().__init__()
|
|
# self.config = config
|
|
# self.num_channels_out = config.num_channels
|
|
|
|
# self.projection = nn.Linear(
|
|
# config.hidden_size,
|
|
# self.num_channels_out * config.patch_size * config.patch_size,
|
|
# bias=True,
|
|
# )
|
|
|
|
# def forward(self, x: torch.Tensor, image_size: tuple[int, int]) -> torch.Tensor:
|
|
# batch_size = x.shape[0]
|
|
|
|
# num_special_tokens = 1 + self.config.num_register_tokens
|
|
# patch_tokens = x[:, num_special_tokens:, :]
|
|
|
|
# projected_tokens = self.projection(patch_tokens)
|
|
|
|
# p = self.config.patch_size
|
|
# c = self.num_channels_out
|
|
# h_grid = image_size[0] // p
|
|
# w_grid = image_size[1] // p
|
|
|
|
# assert patch_tokens.shape[1] == h_grid * w_grid, (
|
|
# "Number of patches does not match image size."
|
|
# )
|
|
|
|
# x_reshaped = projected_tokens.reshape(batch_size, h_grid, w_grid, p, p, c)
|
|
|
|
# x_permuted = torch.einsum("nhwpqc->nchpwq", x_reshaped)
|
|
|
|
# reconstructed_image = x_permuted.reshape(batch_size, c, h_grid * p, w_grid * p)
|
|
|
|
# return reconstructed_image
|
|
|
|
# lets try conv decoder
|
|
|
|
|
|
class DinoV3ViTDecoder(nn.Module):
|
|
def __init__(self, config: DINOv3ViTConfig):
|
|
super().__init__()
|
|
self.config = config
|
|
self.num_channels_out = config.num_channels
|
|
hidden_dim = config.hidden_size
|
|
patch_size = config.patch_size
|
|
|
|
self.projection = nn.Linear(hidden_dim, hidden_dim)
|
|
|
|
if patch_size == 14:
|
|
final_upsample = 7
|
|
elif patch_size == 16:
|
|
final_upsample = 8
|
|
elif patch_size == 8:
|
|
final_upsample = 4
|
|
else:
|
|
raise ValueError("invalid")
|
|
|
|
self.decoder = nn.Sequential(
|
|
nn.Conv2d(hidden_dim, 256, kernel_size=3, padding=1),
|
|
nn.BatchNorm2d(256),
|
|
nn.ReLU(inplace=True),
|
|
nn.Upsample(scale_factor=2, mode="bilinear", align_corners=False),
|
|
nn.Conv2d(256, 128, kernel_size=3, padding=1),
|
|
nn.BatchNorm2d(128),
|
|
nn.ReLU(inplace=True),
|
|
nn.Upsample(
|
|
scale_factor=final_upsample, mode="bilinear", align_corners=False
|
|
),
|
|
nn.Conv2d(128, 64, kernel_size=3, padding=1),
|
|
nn.BatchNorm2d(64),
|
|
nn.ReLU(inplace=True),
|
|
nn.Conv2d(64, 32, kernel_size=3, padding=1),
|
|
nn.ReLU(inplace=True),
|
|
nn.Conv2d(32, self.num_channels_out, kernel_size=1),
|
|
)
|
|
|
|
def forward(self, x: torch.Tensor, image_size: tuple[int, int]) -> torch.Tensor:
|
|
batch_size = x.shape[0]
|
|
|
|
patch_tokens = x[:, 1 + self.config.num_register_tokens :, :]
|
|
|
|
projected_tokens = self.projection(patch_tokens)
|
|
|
|
p = self.config.patch_size
|
|
h_grid = image_size[0] // p
|
|
w_grid = image_size[1] // p
|
|
|
|
assert patch_tokens.shape[1] == h_grid * w_grid
|
|
|
|
x_spatial = projected_tokens.reshape(
|
|
batch_size, h_grid, w_grid, self.config.hidden_size
|
|
)
|
|
|
|
x_spatial = x_spatial.permute(0, 3, 1, 2)
|
|
reconstructed_image = self.decoder(x_spatial)
|
|
|
|
return reconstructed_image
|
|
|
|
|
|
class UTransformer(nn.Module):
|
|
def __init__(self, config: DINOv3ViTConfig, num_classes: int):
|
|
super().__init__()
|
|
self.config = config
|
|
|
|
self.embeddings = DINOv3ViTEmbeddings(config)
|
|
self.rope_embeddings = DINOv3ViTRopePositionEmbedding(config)
|
|
self.t_embedder = TimestepEmbedder(config.hidden_size)
|
|
# self.y_embedder = LabelEmbedder(
|
|
# num_classes, config.hidden_size, config.drop_path_rate
|
|
# ) # disable cond for now
|
|
|
|
self.encoder_layers = nn.ModuleList(
|
|
[
|
|
DinoConditionedLayer(config, True)
|
|
for _ in range(config.num_hidden_layers)
|
|
]
|
|
)
|
|
self.encoder_norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
|
|
|
|
self.decoder_layers = nn.ModuleList(
|
|
[
|
|
DinoConditionedLayer(config, False)
|
|
for _ in range(config.num_hidden_layers)
|
|
]
|
|
)
|
|
self.decoder = DinoV3ViTDecoder(config)
|
|
|
|
# freeze pretrained
|
|
self.embeddings.requires_grad_(False)
|
|
self.rope_embeddings.requires_grad_(False)
|
|
self.encoder_norm.requires_grad_(False)
|
|
|
|
def forward(
|
|
self,
|
|
pixel_values: torch.Tensor,
|
|
time: torch.Tensor,
|
|
# cond: torch.Tensor,
|
|
bool_masked_pos: Optional[torch.Tensor] = None,
|
|
head_mask: Optional[torch.Tensor] = None,
|
|
):
|
|
pixel_values = pixel_values.to(self.embeddings.patch_embeddings.weight.dtype)
|
|
position_embeddings = self.rope_embeddings(pixel_values)
|
|
|
|
x = self.embeddings(pixel_values, bool_masked_pos=bool_masked_pos)
|
|
t = self.t_embedder(time).unsqueeze(1)
|
|
# y = self.y_embedder(cond, self.training).unsqueeze(1)
|
|
# conditioning_input = t.to(x.dtype) + y.to(x.dtype)
|
|
conditioning_input = t.to(x.dtype)
|
|
|
|
residual = []
|
|
for i, layer_module in enumerate(self.encoder_layers):
|
|
residual.append(x)
|
|
layer_head_mask = head_mask[i] if head_mask is not None else None
|
|
x = layer_module(
|
|
x,
|
|
conditioning_input=conditioning_input,
|
|
attention_mask=layer_head_mask,
|
|
position_embeddings=position_embeddings,
|
|
)
|
|
x = self.encoder_norm(x)
|
|
|
|
for i, layer_module in enumerate(self.decoder_layers):
|
|
layer_head_mask = head_mask[i] if head_mask is not None else None
|
|
x = layer_module(
|
|
x,
|
|
conditioning_input=conditioning_input,
|
|
attention_mask=layer_head_mask,
|
|
position_embeddings=position_embeddings,
|
|
)
|
|
x = x + residual.pop()
|
|
|
|
return self.decoder(x, image_size=pixel_values.shape[-2:])
|
|
|
|
@staticmethod
|
|
def from_pretrained_backbone(name: str):
|
|
config = DINOv3ViTConfig.from_pretrained(name)
|
|
instance = UTransformer(config, 0).to("cuda:3")
|
|
|
|
weight_dict = {}
|
|
with safe_open(
|
|
hf_hub_download(name, "model.safetensors"), framework="pt", device="cuda:3"
|
|
) as f:
|
|
for key in f.keys():
|
|
new_key = key.replace("layer.", "encoder_layers.").replace(
|
|
"norm.", "encoder_norm."
|
|
)
|
|
weight_dict[new_key] = f.get_tensor(key)
|
|
|
|
instance.load_state_dict(weight_dict, strict=False)
|
|
|
|
return instance
|