mirror of
https://github.com/lllyasviel/ControlNet.git
synced 2026-01-09 14:08:03 -05:00
51 lines
1.3 KiB
Python
51 lines
1.3 KiB
Python
import sys
|
|
import os
|
|
|
|
assert len(sys.argv) == 3, 'Args are wrong.'
|
|
|
|
input_path = sys.argv[1]
|
|
output_path = sys.argv[2]
|
|
|
|
assert os.path.exists(input_path), 'Input model does not exist.'
|
|
assert not os.path.exists(output_path), 'Output filename already exists.'
|
|
assert os.path.exists(os.path.dirname(output_path)), 'Output path is not valid.'
|
|
|
|
import torch
|
|
from share import *
|
|
from cldm.model import create_model
|
|
|
|
|
|
def get_node_name(name, parent_name):
|
|
if len(name) <= len(parent_name):
|
|
return False, ''
|
|
p = name[:len(parent_name)]
|
|
if p != parent_name:
|
|
return False, ''
|
|
return True, name[len(parent_name):]
|
|
|
|
|
|
model = create_model(config_path='./models/cldm_v15.yaml')
|
|
|
|
pretrained_weights = torch.load(input_path)
|
|
if 'state_dict' in pretrained_weights:
|
|
pretrained_weights = pretrained_weights['state_dict']
|
|
|
|
scratch_dict = model.state_dict()
|
|
|
|
target_dict = {}
|
|
for k in scratch_dict.keys():
|
|
is_control, name = get_node_name(k, 'control_')
|
|
if is_control:
|
|
copy_k = 'model.diffusion_' + name
|
|
else:
|
|
copy_k = k
|
|
if copy_k in pretrained_weights:
|
|
target_dict[k] = pretrained_weights[copy_k].clone()
|
|
else:
|
|
target_dict[k] = scratch_dict[k].clone()
|
|
print(f'These weights are newly added: {k}')
|
|
|
|
model.load_state_dict(target_dict, strict=True)
|
|
torch.save(model.state_dict(), output_path)
|
|
print('Done.')
|