初始化换发型项目：3个微服务代码 + 部署脚本
包含： - hair_service_sd: 换发型/换发色算法服务 (端口 8801) - photo_service: LoRA 训练调度服务 (端口 32678) - stable-diffusion-webui: SD WebUI 推理服务 (端口 57860) - kohya_ss_home: 训练环境代码 - meidaojia: 监控测试脚本 - setup.sh: 一键部署脚本 (conda环境恢复 + 配置生成 + 完整性检查) - start_all_services.sh: 启动3个服务 - configure.ini.template: 路径模板化 (BASE_DIR自动推导) - conda_envs/py310.yml: py310 环境定义 大文件 (weights/, models/, data/, conda_envs/*.tar.gz 等) 通过 .gitignore 排除， 由网盘单独上传。
This commit is contained in:
colomi
@@ -0,0 +1,39 @@
|
||||
import logging
|
||||
from .resnet_enc import ResNet_D, BasicBlock
|
||||
from .res_shortcut_enc import ResShortCut_D
|
||||
from .res_gca_enc import ResGuidedCxtAtten
|
||||
|
||||
|
||||
__all__ = ['res_shortcut_encoder_29', 'resnet_gca_encoder_29']
|
||||
|
||||
|
||||
def _res_shortcut_D(block, layers, **kwargs):
|
||||
model = ResShortCut_D(block, layers, **kwargs)
|
||||
return model
|
||||
|
||||
|
||||
def _res_gca_D(block, layers, **kwargs):
|
||||
model = ResGuidedCxtAtten(block, layers, **kwargs)
|
||||
return model
|
||||
|
||||
|
||||
def resnet_gca_encoder_29(**kwargs):
|
||||
"""Constructs a resnet_encoder_29 model.
|
||||
"""
|
||||
return _res_gca_D(BasicBlock, [3, 4, 4, 2], **kwargs)
|
||||
|
||||
|
||||
def res_shortcut_encoder_29(**kwargs):
|
||||
"""Constructs a resnet_encoder_25 model.
|
||||
"""
|
||||
return _res_shortcut_D(BasicBlock, [3, 4, 4, 2], **kwargs)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import torch
|
||||
logging.basicConfig(level=logging.DEBUG, format='[%(asctime)s] %(levelname)s: %(message)s',
|
||||
datefmt='%m-%d %H:%M:%S')
|
||||
resnet_encoder = res_shortcut_encoder_29()
|
||||
x = torch.randn(4,6,512,512)
|
||||
z = resnet_encoder(x)
|
||||
print(z[0].shape)
|
||||
@@ -0,0 +1,97 @@
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
# from utils import CONFIG
|
||||
from matting.networks.encoders.resnet_enc import ResNet_D
|
||||
from matting.networks.ops import GuidedCxtAtten, SpectralNorm
|
||||
|
||||
class ResGuidedCxtAtten(ResNet_D):
|
||||
|
||||
def __init__(self, block, layers, norm_layer=None, late_downsample=False):
|
||||
super(ResGuidedCxtAtten, self).__init__(block, layers, norm_layer, late_downsample=late_downsample)
|
||||
first_inplane = 3 + 3
|
||||
self.shortcut_inplane = [first_inplane, self.midplanes, 64, 128, 256]
|
||||
self.shortcut_plane = [32, self.midplanes, 64, 128, 256]
|
||||
|
||||
self.shortcut = nn.ModuleList()
|
||||
for stage, inplane in enumerate(self.shortcut_inplane):
|
||||
self.shortcut.append(self._make_shortcut(inplane, self.shortcut_plane[stage]))
|
||||
|
||||
self.guidance_head = nn.Sequential(
|
||||
nn.ReflectionPad2d(1),
|
||||
SpectralNorm(nn.Conv2d(3, 16, kernel_size=3, padding=0, stride=2, bias=False)),
|
||||
nn.ReLU(inplace=True),
|
||||
self._norm_layer(16),
|
||||
nn.ReflectionPad2d(1),
|
||||
SpectralNorm(nn.Conv2d(16, 32, kernel_size=3, padding=0, stride=2, bias=False)),
|
||||
nn.ReLU(inplace=True),
|
||||
self._norm_layer(32),
|
||||
nn.ReflectionPad2d(1),
|
||||
SpectralNorm(nn.Conv2d(32, 128, kernel_size=3, padding=0, stride=2, bias=False)),
|
||||
nn.ReLU(inplace=True),
|
||||
self._norm_layer(128)
|
||||
)
|
||||
|
||||
self.gca = GuidedCxtAtten(128, 128)
|
||||
|
||||
# initialize guidance head
|
||||
for layers in range(len(self.guidance_head)):
|
||||
m = self.guidance_head[layers]
|
||||
if isinstance(m, nn.Conv2d):
|
||||
if hasattr(m, "weight_bar"):
|
||||
nn.init.xavier_uniform_(m.weight_bar)
|
||||
elif isinstance(m, nn.BatchNorm2d):
|
||||
nn.init.constant_(m.weight, 1)
|
||||
nn.init.constant_(m.bias, 0)
|
||||
|
||||
def _make_shortcut(self, inplane, planes):
|
||||
return nn.Sequential(
|
||||
SpectralNorm(nn.Conv2d(inplane, planes, kernel_size=3, padding=1, bias=False)),
|
||||
nn.ReLU(inplace=True),
|
||||
self._norm_layer(planes),
|
||||
SpectralNorm(nn.Conv2d(planes, planes, kernel_size=3, padding=1, bias=False)),
|
||||
nn.ReLU(inplace=True),
|
||||
self._norm_layer(planes)
|
||||
)
|
||||
|
||||
def forward(self, x):
|
||||
|
||||
out = self.conv1(x)
|
||||
out = self.bn1(out)
|
||||
out = self.activation(out)
|
||||
out = self.conv2(out)
|
||||
out = self.bn2(out)
|
||||
x1 = self.activation(out) # N x 32 x 256 x 256
|
||||
out = self.conv3(x1)
|
||||
out = self.bn3(out)
|
||||
out = self.activation(out)
|
||||
|
||||
im_fea = self.guidance_head(x[:, :3, ...]) # downsample origin image and extract features
|
||||
# if CONFIG.model.trimap_channel == 3:
|
||||
unknown = F.interpolate(x[:, 4:5, ...], scale_factor=1/8, mode='nearest')
|
||||
# else:
|
||||
# unknown = F.interpolate(x[:,3:,...].eq(1.).float(), scale_factor=1/8, mode='nearest')
|
||||
|
||||
x2 = self.layer1(out) # N x 64 x 128 x 128
|
||||
x3= self.layer2(x2) # N x 128 x 64 x 64
|
||||
x3, offset = self.gca(im_fea, x3, unknown) # contextual attention
|
||||
x4 = self.layer3(x3) # N x 256 x 32 x 32
|
||||
out = self.layer_bottleneck(x4) # N x 512 x 16 x 16
|
||||
|
||||
fea1 = self.shortcut[0](x) # input image and trimap
|
||||
fea2 = self.shortcut[1](x1)
|
||||
fea3 = self.shortcut[2](x2)
|
||||
fea4 = self.shortcut[3](x3)
|
||||
fea5 = self.shortcut[4](x4)
|
||||
|
||||
return out, {'shortcut': (fea1, fea2, fea3, fea4, fea5),
|
||||
'image_fea': im_fea,
|
||||
'unknown': unknown,
|
||||
'offset_1': offset}
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
from matting.networks.encoders.resnet_enc import BasicBlock
|
||||
m = ResGuidedCxtAtten(BasicBlock, [3, 4, 4, 2])
|
||||
for m in m.modules():
|
||||
print(m)
|
||||
@@ -0,0 +1,51 @@
|
||||
import torch.nn as nn
|
||||
# from utils import CONFIG
|
||||
from matting.networks.encoders.resnet_enc import ResNet_D
|
||||
from matting.networks.ops import SpectralNorm
|
||||
|
||||
class ResShortCut_D(ResNet_D):
|
||||
|
||||
def __init__(self, block, layers, norm_layer=None, late_downsample=False):
|
||||
super(ResShortCut_D, self).__init__(block, layers, norm_layer, late_downsample=late_downsample)
|
||||
first_inplane = 3 + 3
|
||||
self.shortcut_inplane = [first_inplane, self.midplanes, 64, 128, 256]
|
||||
self.shortcut_plane = [32, self.midplanes, 64, 128, 256]
|
||||
|
||||
self.shortcut = nn.ModuleList()
|
||||
for stage, inplane in enumerate(self.shortcut_inplane):
|
||||
self.shortcut.append(self._make_shortcut(inplane, self.shortcut_plane[stage]))
|
||||
|
||||
def _make_shortcut(self, inplane, planes):
|
||||
return nn.Sequential(
|
||||
SpectralNorm(nn.Conv2d(inplane, planes, kernel_size=3, padding=1, bias=False)),
|
||||
nn.ReLU(inplace=True),
|
||||
self._norm_layer(planes),
|
||||
SpectralNorm(nn.Conv2d(planes, planes, kernel_size=3, padding=1, bias=False)),
|
||||
nn.ReLU(inplace=True),
|
||||
self._norm_layer(planes)
|
||||
)
|
||||
|
||||
def forward(self, x):
|
||||
|
||||
out = self.conv1(x)
|
||||
out = self.bn1(out)
|
||||
out = self.activation(out)
|
||||
out = self.conv2(out)
|
||||
out = self.bn2(out)
|
||||
x1 = self.activation(out) # N x 32 x 256 x 256
|
||||
out = self.conv3(x1)
|
||||
out = self.bn3(out)
|
||||
out = self.activation(out)
|
||||
|
||||
x2 = self.layer1(out) # N x 64 x 128 x 128
|
||||
x3= self.layer2(x2) # N x 128 x 64 x 64
|
||||
x4 = self.layer3(x3) # N x 256 x 32 x 32
|
||||
out = self.layer_bottleneck(x4) # N x 512 x 16 x 16
|
||||
|
||||
fea1 = self.shortcut[0](x) # input image and trimap
|
||||
fea2 = self.shortcut[1](x1)
|
||||
fea3 = self.shortcut[2](x2)
|
||||
fea4 = self.shortcut[3](x3)
|
||||
fea5 = self.shortcut[4](x4)
|
||||
|
||||
return out, {'shortcut':(fea1, fea2, fea3, fea4, fea5), 'image':x[:,:3,...]}
|
||||
@@ -0,0 +1,150 @@
|
||||
import logging
|
||||
import torch.nn as nn
|
||||
from matting.networks.ops import SpectralNorm
|
||||
|
||||
def conv3x3(in_planes, out_planes, stride=1, groups=1, dilation=1):
|
||||
"""3x3 convolution with padding"""
|
||||
return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
|
||||
padding=dilation, groups=groups, bias=False, dilation=dilation)
|
||||
|
||||
|
||||
def conv1x1(in_planes, out_planes, stride=1):
|
||||
"""1x1 convolution"""
|
||||
return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, bias=False)
|
||||
|
||||
|
||||
class BasicBlock(nn.Module):
|
||||
expansion = 1
|
||||
|
||||
def __init__(self, inplanes, planes, stride=1, downsample=None, norm_layer=None):
|
||||
super(BasicBlock, self).__init__()
|
||||
if norm_layer is None:
|
||||
norm_layer = nn.BatchNorm2d
|
||||
# Both self.conv1 and self.downsample layers downsample the input when stride != 1
|
||||
self.conv1 = SpectralNorm(conv3x3(inplanes, planes, stride))
|
||||
self.bn1 = norm_layer(planes)
|
||||
self.activation = nn.ReLU(inplace=True)
|
||||
self.conv2 = SpectralNorm(conv3x3(planes, planes))
|
||||
self.bn2 = norm_layer(planes)
|
||||
self.downsample = downsample
|
||||
self.stride = stride
|
||||
|
||||
def forward(self, x):
|
||||
identity = x
|
||||
|
||||
out = self.conv1(x)
|
||||
out = self.bn1(out)
|
||||
out = self.activation(out)
|
||||
|
||||
out = self.conv2(out)
|
||||
out = self.bn2(out)
|
||||
|
||||
if self.downsample is not None:
|
||||
identity = self.downsample(x)
|
||||
|
||||
out += identity
|
||||
out = self.activation(out)
|
||||
|
||||
return out
|
||||
|
||||
|
||||
class ResNet_D(nn.Module):
|
||||
"""
|
||||
Implement and pre-train on ImageNet with the tricks from
|
||||
https://arxiv.org/abs/1812.01187
|
||||
without the mix-up part.
|
||||
"""
|
||||
|
||||
def __init__(self, block, layers, norm_layer=None, late_downsample=False):
|
||||
super(ResNet_D, self).__init__()
|
||||
self.logger = logging.getLogger("Logger")
|
||||
if norm_layer is None:
|
||||
norm_layer = nn.BatchNorm2d
|
||||
self._norm_layer = norm_layer
|
||||
|
||||
self.inplanes = 64
|
||||
self.late_downsample = late_downsample
|
||||
self.midplanes = 64 if late_downsample else 32
|
||||
self.start_stride = [1, 2, 1, 2] if late_downsample else [2, 1, 2, 1]
|
||||
self.conv1 = SpectralNorm(nn.Conv2d(3 + 3, 32, kernel_size=3,
|
||||
stride=self.start_stride[0], padding=1, bias=False))
|
||||
self.conv2 = SpectralNorm(nn.Conv2d(32, self.midplanes, kernel_size=3, stride=self.start_stride[1], padding=1,
|
||||
bias=False))
|
||||
self.conv3 = SpectralNorm(nn.Conv2d(self.midplanes, self.inplanes, kernel_size=3, stride=self.start_stride[2],
|
||||
padding=1, bias=False))
|
||||
self.bn1 = norm_layer(32)
|
||||
self.bn2 = norm_layer(self.midplanes)
|
||||
self.bn3 = norm_layer(self.inplanes)
|
||||
self.activation = nn.ReLU(inplace=True)
|
||||
self.layer1 = self._make_layer(block, 64, layers[0], stride=self.start_stride[3])
|
||||
self.layer2 = self._make_layer(block, 128, layers[1], stride=2)
|
||||
self.layer3 = self._make_layer(block, 256, layers[2], stride=2)
|
||||
self.layer_bottleneck = self._make_layer(block, 512, layers[3], stride=2)
|
||||
|
||||
for m in self.modules():
|
||||
if isinstance(m, nn.Conv2d):
|
||||
nn.init.xavier_uniform_(m.weight_bar)
|
||||
elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)):
|
||||
nn.init.constant_(m.weight, 1)
|
||||
nn.init.constant_(m.bias, 0)
|
||||
|
||||
# Zero-initialize the last BN in each residual branch,
|
||||
# so that the residual branch starts with zeros, and each residual block behaves like an identity.
|
||||
# This improves the model by 0.2~0.3% according to https://arxiv.org/abs/1706.02677
|
||||
|
||||
for m in self.modules():
|
||||
if isinstance(m, BasicBlock):
|
||||
nn.init.constant_(m.bn2.weight, 0)
|
||||
|
||||
self.logger.debug("encoder conv1 weight shape: {}".format(str(self.conv1.module.weight_bar.data.shape)))
|
||||
self.conv1.module.weight_bar.data[:,3:,:,:] = 0
|
||||
|
||||
self.logger.debug(self)
|
||||
|
||||
def _make_layer(self, block, planes, blocks, stride=1):
|
||||
if blocks == 0:
|
||||
return nn.Sequential(nn.Identity())
|
||||
norm_layer = self._norm_layer
|
||||
downsample = None
|
||||
if stride != 1:
|
||||
downsample = nn.Sequential(
|
||||
nn.AvgPool2d(2, stride),
|
||||
SpectralNorm(conv1x1(self.inplanes, planes * block.expansion)),
|
||||
norm_layer(planes * block.expansion),
|
||||
)
|
||||
elif self.inplanes != planes * block.expansion:
|
||||
downsample = nn.Sequential(
|
||||
SpectralNorm(conv1x1(self.inplanes, planes * block.expansion, stride)),
|
||||
norm_layer(planes * block.expansion),
|
||||
)
|
||||
|
||||
layers = [block(self.inplanes, planes, stride, downsample, norm_layer)]
|
||||
self.inplanes = planes * block.expansion
|
||||
for _ in range(1, blocks):
|
||||
layers.append(block(self.inplanes, planes, norm_layer=norm_layer))
|
||||
|
||||
return nn.Sequential(*layers)
|
||||
|
||||
def forward(self, x):
|
||||
x = self.conv1(x)
|
||||
x = self.bn1(x)
|
||||
x = self.activation(x)
|
||||
x = self.conv2(x)
|
||||
x = self.bn2(x)
|
||||
x1 = self.activation(x) # N x 32 x 256 x 256
|
||||
x = self.conv3(x1)
|
||||
x = self.bn3(x)
|
||||
x2 = self.activation(x) # N x 64 x 128 x 128
|
||||
|
||||
x3 = self.layer1(x2) # N x 64 x 128 x 128
|
||||
x4 = self.layer2(x3) # N x 128 x 64 x 64
|
||||
x5 = self.layer3(x4) # N x 256 x 32 x 32
|
||||
x = self.layer_bottleneck(x5) # N x 512 x 16 x 16
|
||||
|
||||
return x, (x1, x2, x3, x4, x5)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
m = ResNet_D(BasicBlock, [3, 4, 4, 2])
|
||||
for m in m.modules():
|
||||
print(m._get_name())
|
||||
Reference in New Issue
Block a user