Spaces:
Paused
Paused
File size: 2,141 Bytes
1c72248 |
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 |
import torch
import torch.nn as nn
class SrefImageEncoder(torch.nn.Module):
def __init__(
self,
input_features: int = 1152,
input_tokens: int = 512,
output_tokens: int = 512,
output_features: int = 4096,
intermediate_size: int = 4096,
num_digits: int = 10,
device=None,
dtype=None,
) -> None:
super().__init__()
self.input_features = input_features
self.device = device
self.dtype = dtype
self.input_tokens = input_tokens
self.output_tokens = output_tokens
self.output_features = output_features
self.intermediate_size = intermediate_size
self.num_digits = num_digits
self.proj_in = nn.Linear(
input_features, intermediate_size, dtype=dtype)
# (bs, num_digits, intermediate_size)
self.conv_pool = nn.Conv1d(input_tokens, num_digits, 1, dtype=dtype)
self.linear_pool = nn.Linear(
intermediate_size, 1, dtype=dtype) # (bs, num_digits, 1)
# do sigmoid for digits 0.0-1.0 = (0 to 10) Always floor when rounding digits so you get 0-9
self.flatten = nn.Flatten() # (bs, num_digits * intermediate_size)
# a numeric sref would come in here with num_digits
self.sref_in = nn.Linear(num_digits, intermediate_size, dtype=dtype)
self.fc1 = nn.Linear(intermediate_size, intermediate_size, dtype=dtype)
self.fc2 = nn.Linear(intermediate_size, intermediate_size, dtype=dtype)
self.proj_out = nn.Linear(
intermediate_size, output_features * output_tokens, dtype=dtype)
def forward(self, siglip_embeds) -> torch.Tensor:
x = self.proj_in(siglip_embeds)
x = torch.nn.functional.silu(x)
x = self.conv_pool(x)
x = self.linear_pool(x)
x = torch.sigmoid(x)
sref = self.flatten(x)
x = self.sref_in(sref)
x = torch.nn.functional.silu(x)
x = self.fc1(x)
x = torch.nn.functional.silu(x)
x = self.fc2(x)
x = torch.nn.functional.silu(x)
x = self.proj_out(x)
return x
|