github/ai-toolkit
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

Work on additional image embedding methods. Finalized zipper resampler. It works amazing

Browse Source
This commit is contained in:
Jaret Burkett
2024-02-10 09:00:05 -07:00
parent a8481c1670
commit e074058faa
7 changed files with 261 additions and 47 deletions
Download Patch File Download Diff File Expand all files Collapse all files

41
toolkit/models/clip_fusion.py
View File

@@ -1,6 +1,8 @@
import torch
import torch.nn as nn
from toolkit.models.zipper_resampler import ContextualAlphaMask
# Conv1d MLP
# MLP that can alternately be used as a conv1d on dim 1
@@ -86,46 +88,7 @@ class ZipperBlock(nn.Module):
return x
class ContextualAlphaMask(nn.Module):
def __init__(
self,
dim: int = 768,
):
super(ContextualAlphaMask, self).__init__()
self.dim = dim
half_dim = dim // 2
quarter_dim = dim // 4
self.fc1 = nn.Linear(self.dim, self.dim)
self.fc2 = nn.Linear(self.dim, half_dim)
self.norm1 = nn.LayerNorm(half_dim)
self.fc3 = nn.Linear(half_dim, half_dim)
self.fc4 = nn.Linear(half_dim, quarter_dim)
self.norm2 = nn.LayerNorm(quarter_dim)
self.fc5 = nn.Linear(quarter_dim, quarter_dim)
self.fc6 = nn.Linear(quarter_dim, 1)
# set fc6 weights to near zero
self.fc6.weight.data.normal_(mean=0.0, std=0.0001)
self.act_fn = nn.GELU()
def forward(self, x):
# x = (batch_size, 77, 768)
x = self.fc1(x)
x = self.act_fn(x)
x = self.fc2(x)
x = self.norm1(x)
x = self.act_fn(x)
x = self.fc3(x)
x = self.act_fn(x)
x = self.fc4(x)
x = self.norm2(x)
x = self.act_fn(x)
x = self.fc5(x)
x = self.act_fn(x)
x = self.fc6(x)
x = torch.sigmoid(x)
return x