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

Bug fixes. added ability to use l1 loss. varous other tests and improvements

Browse Source
This commit is contained in:
Jaret Burkett
2024-01-31 06:30:54 -07:00
parent 92b9c71d44
commit 1ae1017748
9 changed files with 474 additions and 23 deletions
Download Patch File Download Diff File Expand all files Collapse all files

11
toolkit/models/clip_fusion.py
View File

@@ -172,6 +172,8 @@ class CLIPFusionModule(nn.Module):
dim=self.text_hidden_size,
)
self.alpha = nn.Parameter(torch.zeros([text_tokens]) + 0.01)
def forward(self, text_embeds, vision_embeds):
# text_embeds = (batch_size, 77, 768)
# vision_embeds = (batch_size, 257, 1024)
@@ -186,7 +188,12 @@ class CLIPFusionModule(nn.Module):
x = x + res
# alpha mask
alpha = self.ctx_alpha(text_embeds)
x = alpha * x + (1 - alpha) * text_embeds
ctx_alpha = self.ctx_alpha(text_embeds)
# reshape alpha to (1, 77, 1)
alpha = self.alpha.unsqueeze(0).unsqueeze(-1)
x = ctx_alpha * x * alpha
x = x + text_embeds
return x