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密钥已脱敏为环境变量，原文件备份在本地
111 lines
3.8 KiB
Python
111 lines
3.8 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
from collections import OrderedDict
|
|
import numpy as np
|
|
import os
|
|
# from hairstyle_model import modelRoot
|
|
modelRoot = "./weights"
|
|
# modelRoot = "/home/yangchaojie/Desktop/hairstyle/hairstyle_infer/weights"
|
|
class Flatten(nn.Module):
|
|
def __init__(self):
|
|
super(Flatten, self).__init__()
|
|
def forward(self, x):
|
|
x = x.transpose(3, 2).contiguous()
|
|
return x.view(x.size(0), -1)
|
|
|
|
class PNet(nn.Module):
|
|
def __init__(self):
|
|
super(PNet, self).__init__()
|
|
self.model_path = modelRoot
|
|
|
|
self.features = nn.Sequential(OrderedDict([
|
|
('conv1', nn.Conv2d(3, 10, 3, 1)),
|
|
('prelu1', nn.PReLU(10)),
|
|
('pool1', nn.MaxPool2d(2, 2, ceil_mode=True)),
|
|
('conv2', nn.Conv2d(10, 16, 3, 1)),
|
|
('prelu2', nn.PReLU(16)),
|
|
('conv3', nn.Conv2d(16, 32, 3, 1)),
|
|
('prelu3', nn.PReLU(32))
|
|
]))
|
|
self.conv4_1 = nn.Conv2d(32, 2, 1, 1)
|
|
self.conv4_2 = nn.Conv2d(32, 4, 1, 1)
|
|
weights = np.load(os.path.join(self.model_path, 'pnet.npy'), allow_pickle=True)[()]
|
|
for n, p in self.named_parameters():
|
|
p.data = torch.FloatTensor(weights[n])
|
|
|
|
def forward(self, x):
|
|
x = self.features(x)
|
|
a = self.conv4_1(x)
|
|
b = self.conv4_2(x)
|
|
a = F.softmax(a, dim=1)
|
|
return b, a
|
|
|
|
class RNet(nn.Module):
|
|
def __init__(self):
|
|
super(RNet, self).__init__()
|
|
self.model_path = modelRoot
|
|
|
|
self.features = nn.Sequential(OrderedDict([
|
|
('conv1', nn.Conv2d(3, 28, 3, 1)),
|
|
('prelu1', nn.PReLU(28)),
|
|
('pool1', nn.MaxPool2d(3, 2, ceil_mode=True)),
|
|
('conv2', nn.Conv2d(28, 48, 3, 1)),
|
|
('prelu2', nn.PReLU(48)),
|
|
('pool2', nn.MaxPool2d(3, 2, ceil_mode=True)),
|
|
('conv3', nn.Conv2d(48, 64, 2, 1)),
|
|
('prelu3', nn.PReLU(64)),
|
|
('flatten', Flatten()),
|
|
('conv4', nn.Linear(576, 128)),
|
|
('prelu4', nn.PReLU(128))
|
|
]))
|
|
self.conv5_1 = nn.Linear(128, 2)
|
|
self.conv5_2 = nn.Linear(128, 4)
|
|
weights = np.load(os.path.join(self.model_path, 'rnet.npy'), allow_pickle=True)[()]
|
|
for n, p in self.named_parameters():
|
|
p.data = torch.FloatTensor(weights[n])
|
|
|
|
def forward(self, x):
|
|
x = self.features(x)
|
|
a = self.conv5_1(x)
|
|
b = self.conv5_2(x)
|
|
a = F.softmax(a, dim=1)
|
|
return b, a
|
|
|
|
class ONet(nn.Module):
|
|
def __init__(self):
|
|
super(ONet, self).__init__()
|
|
self.model_path = modelRoot
|
|
|
|
self.features = nn.Sequential(OrderedDict([
|
|
('conv1', nn.Conv2d(3, 32, 3, 1)),
|
|
('prelu1', nn.PReLU(32)),
|
|
('pool1', nn.MaxPool2d(3, 2, ceil_mode=True)),
|
|
('conv2', nn.Conv2d(32, 64, 3, 1)),
|
|
('prelu2', nn.PReLU(64)),
|
|
('pool2', nn.MaxPool2d(3, 2, ceil_mode=True)),
|
|
('conv3', nn.Conv2d(64, 64, 3, 1)),
|
|
('prelu3', nn.PReLU(64)),
|
|
('pool3', nn.MaxPool2d(2, 2, ceil_mode=True)),
|
|
('conv4', nn.Conv2d(64, 128, 2, 1)),
|
|
('prelu4', nn.PReLU(128)),
|
|
('flatten', Flatten()),
|
|
('conv5', nn.Linear(1152, 256)),
|
|
('drop5', nn.Dropout(0.25)),
|
|
('prelu5', nn.PReLU(256)),
|
|
]))
|
|
self.conv6_1 = nn.Linear(256, 2)
|
|
self.conv6_2 = nn.Linear(256, 4)
|
|
self.conv6_3 = nn.Linear(256, 10)
|
|
weights = np.load(os.path.join(self.model_path, 'onet.npy'), allow_pickle=True)[()]
|
|
for n, p in self.named_parameters():
|
|
p.data = torch.FloatTensor(weights[n])
|
|
|
|
def forward(self, x):
|
|
x = self.features(x)
|
|
a = self.conv6_1(x)
|
|
b = self.conv6_2(x)
|
|
c = self.conv6_3(x)
|
|
a = F.softmax(a, dim=1)
|
|
return c, b, a
|