xsl
443cfa298f
包含： - hair_service_sd: 主服务（换发型/换发色/生发，端口8801） - photo_service: LoRA调度+训练（端口32678） - hair_grow_service: 调试测试页（端口8888，含4个测试页） - 批量训练脚本（batch_train_hairstyles.py） - 发际线mask自动识别（hairline_mask.py，4种方案） - 手绘mask换发型（hair_swap_manual.py） - 文档：README.md + LARGE_FILES.md + docs/ 大文件（模型权重200G、训练数据123G）已排除，见 LARGE_FILES.md OSS/COS密钥已脱敏为环境变量，原文件备份在本地
103 lines
3.8 KiB
Python
103 lines
3.8 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
import math
|
|
import torch.optim as optim
|
|
import sys
|
|
import os
|
|
|
|
def op_name(op_name, m):
|
|
m.op_name = op_name
|
|
return m
|
|
|
|
class Flatten(nn.Module):
|
|
def __init__(self, axis):
|
|
super(Flatten, self).__init__()
|
|
|
|
def forward(self, x):
|
|
x = x.view(-1, 2560)
|
|
return x
|
|
|
|
def flatten(name, axis):
|
|
return op_name(name, Flatten(axis))
|
|
|
|
def conv_bn_relu(name, in_channels, out_channels, kernel_size, stride=1, padding=0, dilation=1, groups=1, momentum = 0.1, track_running_stats=True):
|
|
return nn.Sequential(
|
|
op_name(name, nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding, dilation, groups, False)),
|
|
op_name(name + '/bn', nn.BatchNorm2d(out_channels, momentum = momentum, track_running_stats = track_running_stats)),
|
|
op_name(name + '/relu', nn.ReLU(inplace=True)),
|
|
)
|
|
|
|
def linear_bn_relu(name, in_features, out_features, momentum = 0.1, track_running_stats=True):
|
|
return nn.Sequential(
|
|
op_name(name + '/FC', nn.Linear(in_features, out_features)),
|
|
op_name(name + '/bn', nn.BatchNorm1d(out_features, momentum = momentum, track_running_stats = track_running_stats)),
|
|
op_name(name + '/relu', nn.ReLU(inplace=True)),
|
|
)
|
|
|
|
# Define a Basic resnet block
|
|
class BasicResnetBlock(nn.Module):
|
|
def __init__(self, name, in_channels, out_channels, kernel_size=3, stride=2, padding=1, direct_plus=False):
|
|
super(BasicResnetBlock, self).__init__()
|
|
|
|
self.op_name = name
|
|
|
|
self.conv_block = nn.Sequential(
|
|
conv_bn_relu(name=name + '/conv1', in_channels=in_channels, out_channels=out_channels,
|
|
kernel_size=kernel_size, stride=stride, padding=padding),
|
|
|
|
conv_bn_relu(name=name + '/conv2', in_channels=out_channels, out_channels=out_channels,
|
|
kernel_size=3, stride=1, padding=1)
|
|
)
|
|
|
|
if direct_plus and (in_channels == out_channels) and (stride == 1):
|
|
self.sc_block = None
|
|
else:
|
|
self.sc_block = conv_bn_relu(name=name + '/sc_conv', in_channels=in_channels, out_channels=out_channels,
|
|
kernel_size=1, stride=stride, padding=0)
|
|
|
|
def forward(self, x):
|
|
if self.sc_block is None:
|
|
out = x + self.conv_block(x)
|
|
else:
|
|
out = self.conv_block(x) + self.sc_block(x)
|
|
return out
|
|
|
|
class SuperBigResNetV2(nn.Module):
|
|
def __init__(self, name, in_channels, out_channels):
|
|
super(SuperBigResNetV2, self).__init__()
|
|
self.op_name = name
|
|
|
|
input_size = 256
|
|
|
|
op_list = []
|
|
|
|
stride_cnt = 4
|
|
|
|
op_list += [conv_bn_relu(name + '/first_conv', in_channels, 48, kernel_size=5, stride=4, padding=2)]
|
|
|
|
ch_num = [48, 64, 96, 128, 160]
|
|
for i in range(len(ch_num) - 1):
|
|
if ch_num[i] == ch_num[i+1]:
|
|
op_list += [BasicResnetBlock(name + '/stage%d'%(i + 1), ch_num[i], ch_num[i+1], kernel_size=3, stride=1, direct_plus=True)]
|
|
else:
|
|
stride_cnt = stride_cnt * 2
|
|
op_list += [BasicResnetBlock(name + '/stage%d'%(i + 1), ch_num[i], ch_num[i+1], kernel_size=3, stride=2)]
|
|
|
|
item_cnt = int((input_size / stride_cnt) * (input_size / stride_cnt) * ch_num[-1])
|
|
op_list += [flatten(name + '/flatten', 1),
|
|
linear_bn_relu(name + '/FC1', item_cnt, 512),
|
|
op_name(name + '/FC2', nn.Linear(512, out_channels))]
|
|
|
|
self.conv_block = nn.Sequential(*op_list)
|
|
|
|
self.model_path, _ = os.path.split(os.path.realpath(__file__))
|
|
|
|
self.weights = torch.load(os.path.join(self.model_path, 'SuperBigResNetV2_latest.pth'),
|
|
map_location=lambda storage, loc: storage)
|
|
|
|
def forward(self, x):
|
|
return self.conv_block(x)
|
|
|
|
if __name__=='__main__':
|
|
pass
|