初始化:换发型/换发色/训练发型服务
包含: - 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密钥已脱敏为环境变量,原文件备份在本地
This commit is contained in:
@@ -0,0 +1,25 @@
|
||||
import os
|
||||
import torch
|
||||
|
||||
def load_model_by_path(model_save_path, model, gpu_id = None):
|
||||
if not os.path.exists(model_save_path): return
|
||||
loc = 'cpu' if gpu_id is None else 'cuda:{}'.format(gpu_id)
|
||||
pretrained_dict = torch.load(model_save_path, map_location=loc)
|
||||
if 'state_dict' in pretrained_dict: pretrained_dict = pretrained_dict['state_dict']
|
||||
model_dict = model.state_dict()
|
||||
# pretrained_dict.pop('netG.model.1.weight', '404')
|
||||
pretrained_dict_new = {k: v for k, v in pretrained_dict.items()
|
||||
if (k in model_dict and model_dict[k].data.shape == v.data.shape)}
|
||||
# if 'netG.model.1.weight' not in pretrained_dict_new and 'netG.model.1.weight' in model_dict:
|
||||
# pretrained_dict_new['netG.model.1.weight'] = model_dict['netG.model.1.weight']
|
||||
# pretrained_dict_new['netG.model.1.weight'][:,:12] = pretrained_dict['netG.model.1.weight']
|
||||
|
||||
model_dict.update(pretrained_dict_new)
|
||||
model.load_state_dict(model_dict)
|
||||
print("load model: ", model_save_path)
|
||||
|
||||
def save_model_by_path(model_save_path, model):
|
||||
save_dir, _ = os.path.split(model_save_path)
|
||||
if not os.path.isdir(save_dir): os.makedirs(save_dir)
|
||||
model_dic={k.replace('.module', ''):v for k,v in model.state_dict().items()}
|
||||
torch.save(model_dic, model_save_path)
|
||||
Reference in New Issue
Block a user