初始化换发型项目：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,287 @@
|
||||
import os
|
||||
import cv2
|
||||
import torch
|
||||
import logging
|
||||
import numpy as np
|
||||
# from utils.config import CONFIG
|
||||
# import torch.distributed as dist
|
||||
|
||||
def mkdirs(paths):
|
||||
"""create empty directories if they don't exist
|
||||
Parameters:
|
||||
paths (str list) -- a list of directory paths
|
||||
"""
|
||||
if isinstance(paths, list) and not isinstance(paths, str):
|
||||
for path in paths:
|
||||
os.makedirs(path)
|
||||
else:
|
||||
os.makedirs(paths)
|
||||
|
||||
def make_dir(target_dir):
|
||||
"""
|
||||
Create dir if not exists
|
||||
"""
|
||||
if not os.path.exists(target_dir):
|
||||
os.makedirs(target_dir)
|
||||
|
||||
|
||||
def print_network(model, name):
|
||||
"""
|
||||
Print out the network information
|
||||
"""
|
||||
logger = logging.getLogger("Logger")
|
||||
num_params = 0
|
||||
for p in model.parameters():
|
||||
num_params += p.numel()
|
||||
|
||||
logger.info(model)
|
||||
logger.info(name)
|
||||
logger.info("Number of parameters: {}".format(num_params))
|
||||
|
||||
|
||||
def update_lr(lr, optimizer):
|
||||
"""
|
||||
update learning rates
|
||||
"""
|
||||
for param_group in optimizer.param_groups:
|
||||
param_group['lr'] = lr
|
||||
|
||||
|
||||
def warmup_lr(init_lr, step, iter_num):
|
||||
"""
|
||||
Warm up learning rate
|
||||
"""
|
||||
return step/iter_num*init_lr
|
||||
|
||||
|
||||
def add_prefix_state_dict(state_dict, prefix="module"):
|
||||
"""
|
||||
add prefix from the key of pretrained state dict for Data-Parallel
|
||||
"""
|
||||
new_state_dict = {}
|
||||
first_state_name = list(state_dict.keys())[0]
|
||||
if not first_state_name.startswith(prefix):
|
||||
for key, value in state_dict.items():
|
||||
new_state_dict[prefix+"."+key] = state_dict[key].float()
|
||||
else:
|
||||
for key, value in state_dict.items():
|
||||
new_state_dict[key] = state_dict[key].float()
|
||||
return new_state_dict
|
||||
|
||||
|
||||
def remove_prefix_state_dict(state_dict, prefix="module"):
|
||||
"""
|
||||
remove prefix from the key of pretrained state dict for Data-Parallel
|
||||
"""
|
||||
new_state_dict = {}
|
||||
first_state_name = list(state_dict.keys())[0]
|
||||
if not first_state_name.startswith(prefix):
|
||||
for key, value in state_dict.items():
|
||||
new_state_dict[key] = state_dict[key].float()
|
||||
else:
|
||||
for key, value in state_dict.items():
|
||||
new_state_dict[key[len(prefix)+1:]] = state_dict[key].float()
|
||||
return new_state_dict
|
||||
|
||||
#
|
||||
# def load_imagenet_pretrain(model, checkpoint_file):
|
||||
# """
|
||||
# Load imagenet pretrained resnet
|
||||
# Add zeros channel to the first convolution layer
|
||||
# Since we have the spectral normalization, we need to do a little more
|
||||
# """
|
||||
# checkpoint = torch.load(checkpoint_file, map_location = lambda storage, loc: storage.cuda(CONFIG.gpu))
|
||||
# state_dict = remove_prefix_state_dict(checkpoint['state_dict'])
|
||||
# for key, value in state_dict.items():
|
||||
# state_dict[key] = state_dict[key].float()
|
||||
#
|
||||
# logger = logging.getLogger("Logger")
|
||||
# logger.debug("Imagenet pretrained keys:")
|
||||
# logger.debug(state_dict.keys())
|
||||
# logger.debug("Generator keys:")
|
||||
# logger.debug(model.module.encoder.state_dict().keys())
|
||||
# logger.debug("Intersection keys:")
|
||||
# logger.debug(set(model.module.encoder.state_dict().keys())&set(state_dict.keys()))
|
||||
#
|
||||
# weight_u = state_dict["conv1.module.weight_u"]
|
||||
# weight_v = state_dict["conv1.module.weight_v"]
|
||||
# weight_bar = state_dict["conv1.module.weight_bar"]
|
||||
#
|
||||
# logger.debug("weight_v: {}".format(weight_v))
|
||||
# logger.debug("weight_bar: {}".format(weight_bar.view(32, -1)))
|
||||
# logger.debug("sigma: {}".format(weight_u.dot(weight_bar.view(32, -1).mv(weight_v))))
|
||||
#
|
||||
# new_weight_v = torch.zeros(6, 3, 3).cuda()
|
||||
# new_weight_bar = torch.zeros(32, 6, 3, 3).cuda()
|
||||
#
|
||||
# new_weight_v[:3, :, :].copy_(weight_v.view(3, 3, 3))
|
||||
# new_weight_bar[:, :3, :, :].copy_(weight_bar)
|
||||
#
|
||||
# logger.debug("new weight_v: {}".format(new_weight_v.view(-1)))
|
||||
# logger.debug("new weight_bar: {}".format(new_weight_bar.view(32, -1)))
|
||||
# logger.debug("new sigma: {}".format(weight_u.dot(new_weight_bar.view(32, -1).mv(new_weight_v.view(-1)))))
|
||||
#
|
||||
# state_dict["conv1.module.weight_v"] = new_weight_v.view(-1)
|
||||
# state_dict["conv1.module.weight_bar"] = new_weight_bar
|
||||
#
|
||||
# model.module.encoder.load_state_dict(state_dict, strict=False)
|
||||
|
||||
|
||||
def load_VGG_pretrain(model, checkpoint_file):
|
||||
"""
|
||||
Load imagenet pretrained resnet
|
||||
Add zeros channel to the first convolution layer
|
||||
Since we have the spectral normalization, we need to do a little more
|
||||
"""
|
||||
checkpoint = torch.load(checkpoint_file, map_location = lambda storage, loc: storage.cuda())
|
||||
backbone_state_dict = remove_prefix_state_dict(checkpoint['state_dict'])
|
||||
|
||||
model.module.encoder.load_state_dict(backbone_state_dict, strict=False)
|
||||
|
||||
|
||||
def get_unknown_tensor(trimap):
|
||||
"""
|
||||
get 1-channel unknown area tensor from the 3-channel/1-channel trimap tensor
|
||||
"""
|
||||
# if CONFIG.model.trimap_channel == 3:
|
||||
weight = trimap[:, 1:2, :, :].float()
|
||||
# else:
|
||||
# weight = trimap.eq(1).float()
|
||||
return weight
|
||||
|
||||
|
||||
def get_gaborfilter(angles):
|
||||
"""
|
||||
generate gabor filter as the conv kernel
|
||||
:param angles: number of different angles
|
||||
"""
|
||||
gabor_filter = []
|
||||
for angle in range(angles):
|
||||
gabor_filter.append(cv2.getGaborKernel(ksize=(5,5), sigma=0.5, theta=angle*np.pi/8, lambd=5, gamma=0.5))
|
||||
gabor_filter = np.array(gabor_filter)
|
||||
gabor_filter = np.expand_dims(gabor_filter, axis=1)
|
||||
return gabor_filter.astype(np.float32)
|
||||
|
||||
|
||||
def get_gradfilter():
|
||||
"""
|
||||
generate gradient filter as the conv kernel
|
||||
"""
|
||||
grad_filter = []
|
||||
grad_filter.append([[-1, -2, -1], [0, 0, 0], [1, 2, 1]])
|
||||
grad_filter.append([[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]])
|
||||
grad_filter = np.array(grad_filter)
|
||||
grad_filter = np.expand_dims(grad_filter, axis=1)
|
||||
return grad_filter.astype(np.float32)
|
||||
|
||||
|
||||
# def reduce_tensor_dict(tensor_dict, mode='mean'):
|
||||
# """
|
||||
# average tensor dict over different GPUs
|
||||
# """
|
||||
# for key, tensor in tensor_dict.items():
|
||||
# if tensor is not None:
|
||||
# tensor_dict[key] = reduce_tensor(tensor, mode)
|
||||
# return tensor_dict
|
||||
#
|
||||
#
|
||||
# def reduce_tensor(tensor, mode='mean'):
|
||||
# """
|
||||
# average tensor over different GPUs
|
||||
# """
|
||||
# rt = tensor.clone()
|
||||
# dist.all_reduce(rt, op=dist.ReduceOp.SUM)
|
||||
# if mode == 'mean':
|
||||
# rt /= CONFIG.world_size
|
||||
# elif mode == 'sum':
|
||||
# pass
|
||||
# else:
|
||||
# raise NotImplementedError("reduce mode can only be 'mean' or 'sum'")
|
||||
# return rt
|
||||
|
||||
def make_color_wheel():
|
||||
# from https://github.com/JiahuiYu/generative_inpainting/blob/master/inpaint_ops.py
|
||||
RY, YG, GC, CB, BM, MR = (15, 6, 4, 11, 13, 6)
|
||||
ncols = RY + YG + GC + CB + BM + MR
|
||||
colorwheel = np.zeros([ncols, 3])
|
||||
col = 0
|
||||
# RY
|
||||
colorwheel[0:RY, 0] = 255
|
||||
colorwheel[0:RY, 1] = np.transpose(np.floor(255*np.arange(0, RY) / RY))
|
||||
col += RY
|
||||
# YG
|
||||
colorwheel[col:col+YG, 0] = 255 - np.transpose(np.floor(255*np.arange(0, YG) / YG))
|
||||
colorwheel[col:col+YG, 1] = 255
|
||||
col += YG
|
||||
# GC
|
||||
colorwheel[col:col+GC, 1] = 255
|
||||
colorwheel[col:col+GC, 2] = np.transpose(np.floor(255*np.arange(0, GC) / GC))
|
||||
col += GC
|
||||
# CB
|
||||
colorwheel[col:col+CB, 1] = 255 - np.transpose(np.floor(255*np.arange(0, CB) / CB))
|
||||
colorwheel[col:col+CB, 2] = 255
|
||||
col += CB
|
||||
# BM
|
||||
colorwheel[col:col+BM, 2] = 255
|
||||
colorwheel[col:col+BM, 0] = np.transpose(np.floor(255*np.arange(0, BM) / BM))
|
||||
col += + BM
|
||||
# MR
|
||||
colorwheel[col:col+MR, 2] = 255 - np.transpose(np.floor(255 * np.arange(0, MR) / MR))
|
||||
colorwheel[col:col+MR, 0] = 255
|
||||
return colorwheel
|
||||
|
||||
|
||||
COLORWHEEL = make_color_wheel()
|
||||
|
||||
def compute_color(u,v):
|
||||
# from https://github.com/JiahuiYu/generative_inpainting/blob/master/inpaint_ops.py
|
||||
h, w = u.shape
|
||||
img = np.zeros([h, w, 3])
|
||||
nanIdx = np.isnan(u) | np.isnan(v)
|
||||
u[nanIdx] = 0
|
||||
v[nanIdx] = 0
|
||||
colorwheel = COLORWHEEL
|
||||
# colorwheel = make_color_wheel()
|
||||
ncols = np.size(colorwheel, 0)
|
||||
rad = np.sqrt(u**2+v**2)
|
||||
a = np.arctan2(-v, -u) / np.pi
|
||||
fk = (a+1) / 2 * (ncols - 1) + 1
|
||||
k0 = np.floor(fk).astype(int)
|
||||
k1 = k0 + 1
|
||||
k1[k1 == ncols+1] = 1
|
||||
f = fk - k0
|
||||
for i in range(np.size(colorwheel,1)):
|
||||
tmp = colorwheel[:, i]
|
||||
col0 = tmp[k0-1] / 255
|
||||
col1 = tmp[k1-1] / 255
|
||||
col = (1-f) * col0 + f * col1
|
||||
idx = rad <= 1
|
||||
col[idx] = 1-rad[idx]*(1-col[idx])
|
||||
notidx = np.logical_not(idx)
|
||||
col[notidx] *= 0.75
|
||||
img[:, :, i] = np.uint8(np.floor(255 * col*(1-nanIdx)))
|
||||
return img
|
||||
|
||||
def flow_to_image(flow):
|
||||
# part from https://github.com/JiahuiYu/generative_inpainting/blob/master/inpaint_ops.py
|
||||
maxrad = -1
|
||||
u = flow[0, :, :]
|
||||
v = flow[1, :, :]
|
||||
rad = np.sqrt(u ** 2 + v ** 2)
|
||||
maxrad = max(maxrad, np.max(rad))
|
||||
u = u/(maxrad + np.finfo(float).eps)
|
||||
v = v/(maxrad + np.finfo(float).eps)
|
||||
img = compute_color(u, v)
|
||||
|
||||
return img
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import networks
|
||||
logging.basicConfig(level=logging.DEBUG, format='[%(asctime)s] %(levelname)s: %(message)s',
|
||||
datefmt='%m-%d %H:%M:%S')
|
||||
G = networks.get_generator().cuda()
|
||||
# load_imagenet_pretrain(G, CONFIG.model.imagenet_pretrain_path)
|
||||
x = torch.randn(4,3,512,512).cuda()
|
||||
y = torch.randn(4,3,512,512).cuda()
|
||||
z = G(x, y)
|
||||
Reference in New Issue
Block a user