-
Notifications
You must be signed in to change notification settings - Fork 4
/
03_export_combine.py
118 lines (89 loc) · 4.33 KB
/
03_export_combine.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
'''
导出controlnet中各部分的torch模型
xj 2023-07-17
'''
import os
import torch
import torch.nn as nn
from omegaconf import OmegaConf
from ldm.util import instantiate_from_config
import torchvision
import onnx # 用于验证onnx模型
# import onnxruntime as ort # 也可以使用onnxruntime 来做推断
import time
device=torch.device("cuda")
# class CombineUnetControlModel(nn.Module):
# def __init__(self, unet, control_model):
# super().__init__()
# self.unet = unet
# self.control_model = control_model
# # (self, x, hint, timesteps, context, **kwargs):
# # self, x, timesteps=None, context=None, control=None, only_mid_control=False, **kwargs
# # ["x_noisy","hint","timesteps","context"] + [f"control_{i}" for i in range(13)]
# def forward(self, x, hint, timesteps, context, control, **kwargs):
# unet_out = self.unet(x, timesteps=timesteps, context=context, control=control,
# only_mid_control=False, **kwargs)
# control_out = self.control_model(x, hint, timesteps, context, **kwargs)
# return unet_out, control_out
class CombineUnetControlModel2(nn.Module):
def __init__(self, unet, control_model):
super().__init__()
self.unet = unet
self.control_model = control_model
# (self, x, hint, timesteps, context, **kwargs):
# self, x, timesteps=None, context=None, control=None, only_mid_control=False, **kwargs
# ["x_noisy","hint","timesteps","context"] + [f"control_{i}" for i in range(13)]
def forward(self, x, hint, timesteps, context, **kwargs):
control = self.control_model(x, hint, timesteps, context, **kwargs)
unet_out = self.unet(x, timesteps=timesteps, context=context, control=control,
only_mid_control=False, **kwargs)
return unet_out
def get_state_dict(d):
return d.get('state_dict', d)
def get_state_dicts(ckpt_path, location='cuda'):
_, extension = os.path.splitext(ckpt_path)
if extension.lower() == ".safetensors":
import safetensors.torch
state_dict = safetensors.torch.load_file(ckpt_path, device=location)
else:
state_dict = get_state_dict(torch.load(ckpt_path, map_location=torch.device(location)))
state_dict = get_state_dict(state_dict)
print(f'Loaded state_dict from [{ckpt_path}]')
return state_dict
def create_pt_model(config_path):
config = OmegaConf.load(config_path)
model = instantiate_from_config(config.model).cuda()
# model.half()
state_dicts = get_state_dicts('/home/player/ControlNet/models/control_sd15_canny.pth', location='cuda')
# CombineModel = CombineUnetControlModel(unet_model, controlnet_model)
# torch.save(CombineModel,"./models/combine.pth")
# print("save CombineNet success!!!")
#controlNet
controlnet_config = config["model"]['params']['control_stage_config']
controlnet_model = instantiate_from_config(controlnet_config)
controlnet_dicts = {k:state_dicts["control_model."+k] for k in controlnet_model.state_dict()}
controlnet_model.load_state_dict(controlnet_dicts)
controlnet_model = controlnet_model.to(device)
# unet
unet_config = config["model"]['params']['unet_config']
unet_model = instantiate_from_config(unet_config)
unet_dicts = {k:state_dicts["model.diffusion_model."+k] for k in unet_model.state_dict()}
unet_model.load_state_dict(unet_dicts)
unet_model = unet_model.to(device)
with torch.no_grad():
CombineModel = CombineUnetControlModel2(unet_model, controlnet_model)
torch.save(CombineModel,"./models/combine.pth")
print("save CombineNet success!!!")
combine_model = CombineModel
combine_model.to(device)
combine_model.eval()
input_names = ["x_noisy", "hint", "timesteps", "context"]
output_names = ["unet_out"]
dummpy_inputs = (torch.randn(2,4,32,48).cuda(),torch.randn(2,3, 256, 384).cuda(), torch.tensor([2]).cuda(),torch.randn(2, 77, 768).cuda())
torch.onnx.export(combine_model, dummpy_inputs, "./models/combine.onnx", verbose=True,input_names=input_names, output_names=output_names,
opset_version=13)
# print("combine_model torch2onnx success!!!")
def save_model():
pass
if __name__ == "__main__":
model = create_pt_model('./models/cldm_v15.yaml')