support animatediff on sdxl

2026-03-21 16:18:13 +00:00 · 2024-05-11 22:48:02 +08:00
parent 8fa03aa997
commit 3b5bbb5773
7 changed files with 403 additions and 10 deletions
--- a/diffsynth/models/sdxl_motion.py
+++ b/diffsynth/models/sdxl_motion.py
@@ -0,0 +1,103 @@
+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(self):
+        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)