mirror of
https://github.com/ostris/ai-toolkit.git
synced 2026-01-26 16:39:47 +00:00
56 lines
1.6 KiB
Python
56 lines
1.6 KiB
Python
import torch
|
|
import os
|
|
from torch import nn
|
|
from safetensors.torch import load_file
|
|
|
|
|
|
class DFEBlock(nn.Module):
|
|
def __init__(self, channels):
|
|
super().__init__()
|
|
self.conv1 = nn.Conv2d(channels, channels, 3, padding=1)
|
|
self.conv2 = nn.Conv2d(channels, channels, 3, padding=1)
|
|
self.act = nn.GELU()
|
|
|
|
def forward(self, x):
|
|
x_in = x
|
|
x = self.conv1(x)
|
|
x = self.conv2(x)
|
|
x = self.act(x)
|
|
x = x + x_in
|
|
return x
|
|
|
|
|
|
class DiffusionFeatureExtractor(nn.Module):
|
|
def __init__(self, in_channels=32):
|
|
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)
|
|
return x
|
|
|
|
|
|
def load_dfe(model_path) -> DiffusionFeatureExtractor:
|
|
dfe = DiffusionFeatureExtractor()
|
|
if not os.path.exists(model_path):
|
|
raise FileNotFoundError(f"Model file not found: {model_path}")
|
|
# if it ende with safetensors
|
|
if model_path.endswith('.safetensors'):
|
|
state_dict = load_file(model_path)
|
|
else:
|
|
state_dict = torch.load(model_path, weights_only=True)
|
|
if 'model_state_dict' in state_dict:
|
|
state_dict = state_dict['model_state_dict']
|
|
|
|
dfe.load_state_dict(state_dict)
|
|
dfe.eval()
|
|
return dfe
|