| | 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) |
| |
|
