Spaces:
Running
on
Zero
Running
on
Zero
| 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 | |
| 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) | |