ostris/ai-toolkit
1
0
Fork 0
mirror of https://github.com/ostris/ai-toolkit.git synced 2026-03-13 14:39:50 +00:00
Code Issues Packages Projects Releases Wiki Activity

Update dfe model arch

Browse Source
This commit is contained in:
Jaret Burkett
2025-01-22 10:37:23 -07:00
parent 04abe57c76
commit e1549ad54d
1 changed files with 0 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
toolkit/models/diffusion_feature_extraction.py
View File

@@ -25,13 +25,11 @@ class DiffusionFeatureExtractor(nn.Module):
super().__init__()
num_blocks = 6
self.conv_in = nn.Conv2d(in_channels, 512, 1)
self.conv_pool = nn.Conv2d(512, 512, 3, stride=2, padding=1)
self.blocks = nn.ModuleList([DFEBlock(512) for _ in range(num_blocks)])
self.conv_out = nn.Conv2d(512, 512, 1)
def forward(self, x):
x = self.conv_in(x)
x = self.conv_pool(x)
for block in self.blocks:
x = block(x)
x = self.conv_out(x)