|
|
from .sd_motion import TemporalBlock |
|
|
import torch |
|
|
|
|
|
|
|
|
|
|
|
class SDXLMotionModel(torch.nn.Module): |
|
|
def __init__(self): |
|
|
super().__init__() |
|
|
self.motion_modules = torch.nn.ModuleList([ |
|
|
TemporalBlock(8, 320//8, 320, eps=1e-6), |
|
|
TemporalBlock(8, 320//8, 320, eps=1e-6), |
|
|
|
|
|
TemporalBlock(8, 640//8, 640, eps=1e-6), |
|
|
TemporalBlock(8, 640//8, 640, eps=1e-6), |
|
|
|
|
|
TemporalBlock(8, 1280//8, 1280, eps=1e-6), |
|
|
TemporalBlock(8, 1280//8, 1280, eps=1e-6), |
|
|
|
|
|
TemporalBlock(8, 1280//8, 1280, eps=1e-6), |
|
|
TemporalBlock(8, 1280//8, 1280, eps=1e-6), |
|
|
TemporalBlock(8, 1280//8, 1280, eps=1e-6), |
|
|
|
|
|
TemporalBlock(8, 640//8, 640, eps=1e-6), |
|
|
TemporalBlock(8, 640//8, 640, eps=1e-6), |
|
|
TemporalBlock(8, 640//8, 640, eps=1e-6), |
|
|
|
|
|
TemporalBlock(8, 320//8, 320, eps=1e-6), |
|
|
TemporalBlock(8, 320//8, 320, eps=1e-6), |
|
|
TemporalBlock(8, 320//8, 320, eps=1e-6), |
|
|
]) |
|
|
self.call_block_id = { |
|
|
0: 0, |
|
|
2: 1, |
|
|
7: 2, |
|
|
10: 3, |
|
|
15: 4, |
|
|
18: 5, |
|
|
25: 6, |
|
|
28: 7, |
|
|
31: 8, |
|
|
35: 9, |
|
|
38: 10, |
|
|
41: 11, |
|
|
44: 12, |
|
|
46: 13, |
|
|
48: 14, |
|
|
} |
|
|
|
|
|
def forward(self): |
|
|
pass |
|
|
|
|
|
@staticmethod |
|
|
def state_dict_converter(): |
|
|
return SDMotionModelStateDictConverter() |
|
|
|
|
|
|
|
|
class SDMotionModelStateDictConverter: |
|
|
def __init__(self): |
|
|
pass |
|
|
|
|
|
def from_diffusers(self, state_dict): |
|
|
rename_dict = { |
|
|
"norm": "norm", |
|
|
"proj_in": "proj_in", |
|
|
"transformer_blocks.0.attention_blocks.0.to_q": "transformer_blocks.0.attn1.to_q", |
|
|
"transformer_blocks.0.attention_blocks.0.to_k": "transformer_blocks.0.attn1.to_k", |
|
|
"transformer_blocks.0.attention_blocks.0.to_v": "transformer_blocks.0.attn1.to_v", |
|
|
"transformer_blocks.0.attention_blocks.0.to_out.0": "transformer_blocks.0.attn1.to_out", |
|
|
"transformer_blocks.0.attention_blocks.0.pos_encoder": "transformer_blocks.0.pe1", |
|
|
"transformer_blocks.0.attention_blocks.1.to_q": "transformer_blocks.0.attn2.to_q", |
|
|
"transformer_blocks.0.attention_blocks.1.to_k": "transformer_blocks.0.attn2.to_k", |
|
|
"transformer_blocks.0.attention_blocks.1.to_v": "transformer_blocks.0.attn2.to_v", |
|
|
"transformer_blocks.0.attention_blocks.1.to_out.0": "transformer_blocks.0.attn2.to_out", |
|
|
"transformer_blocks.0.attention_blocks.1.pos_encoder": "transformer_blocks.0.pe2", |
|
|
"transformer_blocks.0.norms.0": "transformer_blocks.0.norm1", |
|
|
"transformer_blocks.0.norms.1": "transformer_blocks.0.norm2", |
|
|
"transformer_blocks.0.ff.net.0.proj": "transformer_blocks.0.act_fn.proj", |
|
|
"transformer_blocks.0.ff.net.2": "transformer_blocks.0.ff", |
|
|
"transformer_blocks.0.ff_norm": "transformer_blocks.0.norm3", |
|
|
"proj_out": "proj_out", |
|
|
} |
|
|
name_list = sorted([i for i in state_dict if i.startswith("down_blocks.")]) |
|
|
name_list += sorted([i for i in state_dict if i.startswith("mid_block.")]) |
|
|
name_list += sorted([i for i in state_dict if i.startswith("up_blocks.")]) |
|
|
state_dict_ = {} |
|
|
last_prefix, module_id = "", -1 |
|
|
for name in name_list: |
|
|
names = name.split(".") |
|
|
prefix_index = names.index("temporal_transformer") + 1 |
|
|
prefix = ".".join(names[:prefix_index]) |
|
|
if prefix != last_prefix: |
|
|
last_prefix = prefix |
|
|
module_id += 1 |
|
|
middle_name = ".".join(names[prefix_index:-1]) |
|
|
suffix = names[-1] |
|
|
if "pos_encoder" in names: |
|
|
rename = ".".join(["motion_modules", str(module_id), rename_dict[middle_name]]) |
|
|
else: |
|
|
rename = ".".join(["motion_modules", str(module_id), rename_dict[middle_name], suffix]) |
|
|
state_dict_[rename] = state_dict[name] |
|
|
return state_dict_ |
|
|
|
|
|
def from_civitai(self, state_dict): |
|
|
return self.from_diffusers(state_dict) |
|
|
|
