File size: 2,842 Bytes
f2c0505 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 | from __future__ import annotations
import re
_BLOCK_RE = re.compile(r"^(blocks\.\d+)\.(?!block\.)")
def _insert_block_wrapper(key: str) -> str:
if key.startswith("blocks.") and ".block." not in key:
return _BLOCK_RE.sub(r"\1.block.", key, count=1)
return key
def _remove_block_wrapper(key: str) -> str:
if key.startswith("blocks."):
return key.replace(".block.", ".", 1)
return key
def diffusers_to_official_key(key: str) -> str:
"""Map WanAnimate2 Diffusers/refactor keyspace to Wan-Video/Wan-Animate-2 keyspace.
OrbitQuant Project-A was repaired into the current Diffusers WanAnimate2 keyspace:
blocks.0.self_attn.to_q.weight
The official Wan-Animate-2 repo uses:
blocks.0.block.self_attn.q.weight
"""
k = key
if k.startswith("model.diffusion_model."):
k = k[len("model.diffusion_model.") :]
# Only attention projection names differ. FFN names stay ffn.0 / ffn.2.
replacements = (
(".self_attn.to_out.0.", ".self_attn.o."),
(".cross_attn.to_out.0.", ".cross_attn.o."),
(".self_attn.to_q.", ".self_attn.q."),
(".self_attn.to_k.", ".self_attn.k."),
(".self_attn.to_v.", ".self_attn.v."),
(".cross_attn.to_q.", ".cross_attn.q."),
(".cross_attn.to_k.", ".cross_attn.k."),
(".cross_attn.to_v.", ".cross_attn.v."),
(".cross_attn.add_k_proj.", ".cross_attn.k_img."),
(".cross_attn.add_v_proj.", ".cross_attn.v_img."),
(".cross_attn.norm_added_k.", ".cross_attn.norm_k_img."),
)
for old, new in replacements:
k = k.replace(old, new)
return _insert_block_wrapper(k)
def official_to_diffusers_key(key: str) -> str:
"""Inverse of diffusers_to_official_key for Wan-Animate-2 transformer keys."""
k = _remove_block_wrapper(key)
replacements = (
(".self_attn.o.", ".self_attn.to_out.0."),
(".cross_attn.o.", ".cross_attn.to_out.0."),
(".self_attn.q.", ".self_attn.to_q."),
(".self_attn.k.", ".self_attn.to_k."),
(".self_attn.v.", ".self_attn.to_v."),
(".cross_attn.q.", ".cross_attn.to_q."),
(".cross_attn.k.", ".cross_attn.to_k."),
(".cross_attn.v.", ".cross_attn.to_v."),
(".cross_attn.k_img.", ".cross_attn.add_k_proj."),
(".cross_attn.v_img.", ".cross_attn.add_v_proj."),
(".cross_attn.norm_k_img.", ".cross_attn.norm_added_k."),
)
for old, new in replacements:
k = k.replace(old, new)
return k
def target_key_for_module(module_name: str, keyspace: str = "official") -> str:
key = f"{module_name}.weight" if module_name else "weight"
if keyspace == "official":
return key
if keyspace == "diffusers":
return official_to_diffusers_key(key)
raise ValueError(f"Unsupported keyspace: {keyspace}")
|