SolitaryThinker
d965f578c9
checkpoint ode checkpoint fix t2v to i2v lint chekpt checkpoint hacked but working ode_init scripts WIP fixing time embedding WIP fixing time embedding checkpoint update fix revert revert revert update update visualize
414 lines
17 KiB
Python
414 lines
17 KiB
Python
# pyright: reportMissingImports=false
|
|
from safetensors.torch import save_file, load_file as safe_load_file
|
|
import argparse
|
|
import os
|
|
import re
|
|
import sys
|
|
from collections import OrderedDict
|
|
from typing import Any, Dict, Mapping, Tuple
|
|
import torch
|
|
|
|
try:
|
|
from huggingface_hub import save_torch_state_dict, load_state_dict_from_file # type: ignore
|
|
except Exception:
|
|
save_torch_state_dict = None # type: ignore[assignment]
|
|
load_state_dict_from_file = None # type: ignore[assignment]
|
|
|
|
_param_names_mapping: dict = {
|
|
r"^text_embedding\.0\.(.*)$":
|
|
r"condition_embedder.text_embedder.linear_1.\1",
|
|
r"^text_embedding\.2\.(.*)$":
|
|
r"condition_embedder.text_embedder.linear_2.\1",
|
|
r"^time_embedding\.0\.(.*)$":
|
|
r"condition_embedder.time_embedder.linear_1.\1",
|
|
r"^time_embedding\.2\.(.*)$":
|
|
r"condition_embedder.time_embedder.linear_2.\1",
|
|
r"^time_projection\.1\.(.*)$":
|
|
r"condition_embedder.time_proj.\1",
|
|
r"^img_emb\.proj\.0\.(.*)$":
|
|
r"condition_embedder.image_embedder.norm1.\1",
|
|
r"^img_emb\.proj\.1\.(.*)$":
|
|
r"condition_embedder.image_embedder.ff.net.0.proj.\1",
|
|
r"^img_emb\.proj\.3\.(.*)$":
|
|
r"condition_embedder.image_embedder.ff.net.2.\1",
|
|
r"^img_emb\.proj\.4\.(.*)$":
|
|
r"condition_embedder.image_embedder.norm2.\1",
|
|
r"^head\.modulation":
|
|
r"scale_shift_table",
|
|
r"^head\.head\.(.*)$":
|
|
r"proj_out.\1",
|
|
r"^blocks\.(\d+)\.self_attn\.q\.(.*)$":
|
|
r"blocks.\1.attn1.to_q.\2",
|
|
r"^blocks\.(\d+)\.self_attn\.k\.(.*)$":
|
|
r"blocks.\1.attn1.to_k.\2",
|
|
r"^blocks\.(\d+)\.self_attn\.v\.(.*)$":
|
|
r"blocks.\1.attn1.to_v.\2",
|
|
r"^blocks\.(\d+)\.self_attn\.o\.(.*)$":
|
|
r"blocks.\1.attn1.to_out.0.\2",
|
|
r"^blocks\.(\d+)\.self_attn\.norm_q\.(.*)$":
|
|
r"blocks.\1.attn1.norm_q.\2",
|
|
r"^blocks\.(\d+)\.self_attn\.norm_k\.(.*)$":
|
|
r"blocks.\1.attn1.norm_k.\2",
|
|
r"^blocks\.(\d+)\.cross_attn\.q\.(.*)$":
|
|
r"blocks.\1.attn2.to_q.\2",
|
|
r"^blocks\.(\d+)\.cross_attn\.k\.(.*)$":
|
|
r"blocks.\1.attn2.to_k.\2",
|
|
r"^blocks\.(\d+)\.cross_attn\.k_img\.(.*)$":
|
|
r"blocks.\1.attn2.add_k_proj.\2",
|
|
r"^blocks\.(\d+)\.cross_attn\.v\.(.*)$":
|
|
r"blocks.\1.attn2.to_v.\2",
|
|
r"^blocks\.(\d+)\.cross_attn\.v_img\.(.*)$":
|
|
r"blocks.\1.attn2.add_v_proj.\2",
|
|
r"^blocks\.(\d+)\.cross_attn\.o\.(.*)$":
|
|
r"blocks.\1.attn2.to_out.0.\2",
|
|
r"^blocks\.(\d+)\.cross_attn\.norm_q\.(.*)$":
|
|
r"blocks.\1.attn2.norm_q.\2",
|
|
r"^blocks\.(\d+)\.cross_attn\.norm_k\.(.*)$":
|
|
r"blocks.\1.attn2.norm_k.\2",
|
|
r"^blocks\.(\d+)\.cross_attn\.norm_k_img\.(.*)$":
|
|
r"blocks.\1.attn2.norm_added_k.\2",
|
|
r"^blocks\.(\d+)\.ffn\.0\.(.*)$":
|
|
r"blocks.\1.ffn.net.0.proj.\2",
|
|
r"^blocks\.(\d+)\.ffn\.2\.(.*)$":
|
|
r"blocks.\1.ffn.net.2.\2",
|
|
r"^blocks\.(\d+)\.modulation":
|
|
r"blocks.\1.scale_shift_table",
|
|
r"^blocks\.(\d+)\.norm3\.(.*)$":
|
|
r"blocks.\1.norm2.\2",
|
|
}
|
|
|
|
# The following mapping has an extra 'patch_embedding' field and also contains
|
|
# the 'model' prefixes
|
|
_self_forcing_to_diffusers_param_names_mapping: dict = {
|
|
r"^model.patch_embedding\.(.*)$":
|
|
r"patch_embedding.\1",
|
|
r"^model.text_embedding\.0\.(.*)$":
|
|
r"condition_embedder.text_embedder.linear_1.\1",
|
|
r"^model.text_embedding\.2\.(.*)$":
|
|
r"condition_embedder.text_embedder.linear_2.\1",
|
|
r"^model.time_embedding\.0\.(.*)$":
|
|
r"condition_embedder.time_embedder.linear_1.\1",
|
|
r"^model.time_embedding\.2\.(.*)$":
|
|
r"condition_embedder.time_embedder.linear_2.\1",
|
|
r"^model.time_projection\.1\.(.*)$":
|
|
r"condition_embedder.time_proj.\1",
|
|
r"^model.img_emb\.proj\.0\.(.*)$":
|
|
r"condition_embedder.image_embedder.norm1.\1",
|
|
r"^model.img_emb\.proj\.1\.(.*)$":
|
|
r"condition_embedder.image_embedder.ff.net.0.proj.\1",
|
|
r"^model.img_emb\.proj\.3\.(.*)$":
|
|
r"condition_embedder.image_embedder.ff.net.2.\1",
|
|
r"^model.img_emb\.proj\.4\.(.*)$":
|
|
r"condition_embedder.image_embedder.norm2.\1",
|
|
r"^model.head\.modulation":
|
|
r"scale_shift_table",
|
|
r"^model.head\.head\.(.*)$":
|
|
r"proj_out.\1",
|
|
r"^model.blocks\.(\d+)\.self_attn\.q\.(.*)$":
|
|
r"blocks.\1.attn1.to_q.\2",
|
|
r"^model.blocks\.(\d+)\.self_attn\.k\.(.*)$":
|
|
r"blocks.\1.attn1.to_k.\2",
|
|
r"^model.blocks\.(\d+)\.self_attn\.v\.(.*)$":
|
|
r"blocks.\1.attn1.to_v.\2",
|
|
r"^model.blocks\.(\d+)\.self_attn\.o\.(.*)$":
|
|
r"blocks.\1.attn1.to_out.0.\2",
|
|
r"^model.blocks\.(\d+)\.self_attn\.norm_q\.(.*)$":
|
|
r"blocks.\1.attn1.norm_q.\2",
|
|
r"^model.blocks\.(\d+)\.self_attn\.norm_k\.(.*)$":
|
|
r"blocks.\1.attn1.norm_k.\2",
|
|
r"^model.blocks\.(\d+)\.cross_attn\.q\.(.*)$":
|
|
r"blocks.\1.attn2.to_q.\2",
|
|
r"^model.blocks\.(\d+)\.cross_attn\.k\.(.*)$":
|
|
r"blocks.\1.attn2.to_k.\2",
|
|
r"^model.blocks\.(\d+)\.cross_attn\.k_img\.(.*)$":
|
|
r"blocks.\1.attn2.add_k_proj.\2",
|
|
r"^model.blocks\.(\d+)\.cross_attn\.v\.(.*)$":
|
|
r"blocks.\1.attn2.to_v.\2",
|
|
r"^model.blocks\.(\d+)\.cross_attn\.v_img\.(.*)$":
|
|
r"blocks.\1.attn2.add_v_proj.\2",
|
|
r"^model.blocks\.(\d+)\.cross_attn\.o\.(.*)$":
|
|
r"blocks.\1.attn2.to_out.0.\2",
|
|
r"^model.blocks\.(\d+)\.cross_attn\.norm_q\.(.*)$":
|
|
r"blocks.\1.attn2.norm_q.\2",
|
|
r"^model.blocks\.(\d+)\.cross_attn\.norm_k\.(.*)$":
|
|
r"blocks.\1.attn2.norm_k.\2",
|
|
r"^model.blocks\.(\d+)\.cross_attn\.norm_k_img\.(.*)$":
|
|
r"blocks.\1.attn2.norm_added_k.\2",
|
|
r"^model.blocks\.(\d+)\.ffn\.0\.(.*)$":
|
|
r"blocks.\1.ffn.net.0.proj.\2",
|
|
r"^model.blocks\.(\d+)\.ffn\.2\.(.*)$":
|
|
r"blocks.\1.ffn.net.2.\2",
|
|
r"^model.blocks\.(\d+)\.modulation":
|
|
r"blocks.\1.scale_shift_table",
|
|
r"^model.blocks\.(\d+)\.norm3\.(.*)$":
|
|
r"blocks.\1.norm2.\2",
|
|
}
|
|
|
|
def _replacement_to_regex_template(replacement: str) -> Tuple[str, int]:
|
|
r"""
|
|
Convert a replacement template like "blocks.\1.attn2.to_q.\2" into a regex pattern
|
|
that can be used to match in the reverse direction: "^blocks\.(.*)\.attn2\.to_q\.(.*)$".
|
|
|
|
Returns the regex template and the number of capture groups.
|
|
"""
|
|
# First, protect placeholders \1..\9
|
|
placeholder_tokens: Dict[str, str] = {}
|
|
group_count = 0
|
|
def _token_for(idx: int) -> str:
|
|
return f"__CAP_{idx}__"
|
|
|
|
out = replacement
|
|
for i in range(1, 10):
|
|
token = _token_for(i)
|
|
if f"\\{i}" in out:
|
|
out = out.replace(f"\\{i}", token)
|
|
placeholder_tokens[token] = f"\\{i}"
|
|
group_count = max(group_count, i)
|
|
|
|
# Escape all regex meta in the literal parts
|
|
out = re.escape(out)
|
|
# Restore placeholders as (.*)
|
|
for token in placeholder_tokens.keys():
|
|
out = out.replace(re.escape(token), "(.*)")
|
|
return out, group_count
|
|
|
|
|
|
def invert_mapping(forward_mapping: Mapping[str, str]) -> OrderedDict:
|
|
"""Create a reverse regex mapping by inverting pattern→replacement pairs.
|
|
|
|
- Maintains order from the forward mapping
|
|
- If the forward pattern is anchored with '$', the reverse is anchored as well
|
|
- If forward pattern is prefix-only (no '$'), reverse is also prefix-only
|
|
"""
|
|
reversed_mapping: "OrderedDict[str, str]" = OrderedDict()
|
|
for pattern, replacement in forward_mapping.items():
|
|
# Build reverse pattern from replacement template
|
|
reverse_pat_core, _ = _replacement_to_regex_template(replacement)
|
|
# Respect anchoring: keep '^' always; add '$' only if original had it
|
|
anchored_end = pattern.endswith('$')
|
|
reverse_pattern = f"^{reverse_pat_core}" + ("$" if anchored_end else "")
|
|
# Reverse replacement must be a literal template with backrefs (\1, \2, ...)
|
|
reverse_replacement = _pattern_to_replacement_template(pattern)
|
|
reversed_mapping[reverse_pattern] = reverse_replacement
|
|
return reversed_mapping
|
|
|
|
|
|
def _pattern_to_replacement_template(pattern: str) -> str:
|
|
r"""
|
|
Convert a regex pattern like "^model.blocks\.(\d+)\.self_attn\.q\.(.*)$" into a replacement
|
|
template suitable for re.sub, e.g., "model.blocks.\1.self_attn.q.\2".
|
|
Only supports simple capturing groups of the form (.*) or (\d+), which
|
|
matches the patterns used in the forward mapping.
|
|
"""
|
|
# strip anchors
|
|
core = pattern
|
|
if core.startswith('^'):
|
|
core = core[1:]
|
|
if core.endswith('$'):
|
|
core = core[:-1]
|
|
|
|
# replace groups (.*) or (\d+) with backref tokens in increasing order
|
|
group_index = 0
|
|
def repl(_m: "re.Match[str]") -> str:
|
|
nonlocal group_index
|
|
group_index += 1
|
|
return f"\\{group_index}"
|
|
|
|
core = re.sub(r"\((?:\.\*|\\d\+)\)", repl, core)
|
|
|
|
# unescape literal dots
|
|
core = core.replace(r"\.", ".")
|
|
return core
|
|
|
|
|
|
def select_inner_state_dict(loaded: Mapping[str, Any], key: str = "") -> Tuple[Mapping[str, Any], str]:
|
|
if key:
|
|
if key not in loaded:
|
|
raise KeyError(f"Key '{key}' not found in loaded object. Available keys: {list(loaded.keys())[:20]}")
|
|
return loaded[key], key
|
|
|
|
# If looks like a state dict (all tensors)
|
|
if len(loaded) > 0 and all(torch.is_tensor(v) for v in loaded.values()):
|
|
return loaded, "<root>"
|
|
|
|
# Common containers
|
|
for candidate in ("state_dict", "generator_ema", "model", "ema", "module"):
|
|
if candidate in loaded and isinstance(loaded[candidate], Mapping):
|
|
inner = loaded[candidate]
|
|
if len(inner) > 0 and all(torch.is_tensor(v) for v in inner.values()):
|
|
return inner, candidate
|
|
|
|
# Fallback: first tensor-dict value
|
|
for v in loaded.values():
|
|
if isinstance(v, Mapping) and len(v) > 0 and all(torch.is_tensor(t) for t in v.values()):
|
|
return v, "<auto>"
|
|
|
|
raise ValueError("Could not locate a state_dict (mapping of tensor parameters) in the loaded file.")
|
|
|
|
|
|
def convert_state_dict(state_dict: Mapping[str, torch.Tensor],
|
|
mapping: Mapping[str, str],
|
|
*,
|
|
strict: bool = True,
|
|
add_norm_added_q_dummy: bool = False) -> Tuple[OrderedDict, Dict[str, int]]:
|
|
new_state_dict: "OrderedDict[str, torch.Tensor]" = OrderedDict()
|
|
matched_count = 0
|
|
unmatched_count = 0
|
|
dummy_added = 0
|
|
examples = [] # type: ignore[var-annotated]
|
|
for k, v in state_dict.items():
|
|
new_key = None
|
|
for pattern, replacement in mapping.items():
|
|
if re.match(pattern, k):
|
|
new_key = re.sub(pattern, replacement, k)
|
|
break
|
|
if new_key is None:
|
|
if strict:
|
|
raise ValueError(f"No mapping rule matched for key: {k}")
|
|
else:
|
|
new_key = k # keep original
|
|
unmatched_count += 1
|
|
else:
|
|
matched_count += 1
|
|
new_state_dict[new_key] = v
|
|
|
|
if len(examples) < 5:
|
|
examples.append((k, new_key))
|
|
|
|
if add_norm_added_q_dummy and "norm_added_k" in new_key:
|
|
dummy_key = new_key.replace("norm_added_k", "norm_added_q")
|
|
dummy_value = torch.zeros_like(v)
|
|
new_state_dict[dummy_key] = dummy_value
|
|
dummy_added += 1
|
|
stats = {"matched": matched_count, "unmatched": unmatched_count, "dummy_added": dummy_added}
|
|
# store examples count-wise in stats by encoding as counts in print time (examples returned separately not typed)
|
|
new_state_dict.__dict__["_examples"] = examples # lightweight attach for printing
|
|
return new_state_dict, stats
|
|
|
|
|
|
def save_output(new_state_dict: Mapping[str, torch.Tensor],
|
|
output: str,
|
|
*,
|
|
shard: bool = True,
|
|
max_shard_size: str = "10GB",
|
|
wrapper_key: str = "",
|
|
force_pt: bool = False) -> None:
|
|
if force_pt:
|
|
out_path = coerce_pt_output_path(output)
|
|
obj: Dict[str, Any]
|
|
if wrapper_key:
|
|
obj = {wrapper_key: OrderedDict(new_state_dict)}
|
|
else:
|
|
# Save raw state_dict mapping
|
|
obj = OrderedDict(new_state_dict) # type: ignore[assignment]
|
|
torch.save(obj, out_path)
|
|
return
|
|
|
|
if shard or output.endswith('/') or os.path.isdir(output):
|
|
if save_torch_state_dict is None:
|
|
raise RuntimeError("Saving shards requires 'huggingface_hub'. Install it or use --single-file.")
|
|
out_dir = output if output.endswith('/') else output + '/'
|
|
os.makedirs(out_dir, exist_ok=True)
|
|
save_torch_state_dict(OrderedDict(new_state_dict), out_dir, max_shard_size=max_shard_size)
|
|
else:
|
|
# Save a single safetensors file
|
|
save_file(OrderedDict(new_state_dict), output)
|
|
|
|
|
|
def coerce_pt_output_path(output: str) -> str:
|
|
"""Ensure output path is a .pt/.pth/.bin file. If a directory or unknown ext, coerce to .pt."""
|
|
if output.endswith('/') or os.path.isdir(output):
|
|
os.makedirs(output, exist_ok=True)
|
|
return os.path.join(output, 'converted_wan.pt')
|
|
lower = output.lower()
|
|
if lower.endswith('.pt') or lower.endswith('.pth') or lower.endswith('.bin'):
|
|
return output
|
|
return output + '.pt'
|
|
|
|
|
|
def load_checkpoint(input_path: str) -> Mapping[str, Any]:
|
|
if load_state_dict_from_file is not None:
|
|
return load_state_dict_from_file(input_path)
|
|
# Fallbacks by extension
|
|
lower = input_path.lower()
|
|
if lower.endswith('.safetensors'):
|
|
return safe_load_file(input_path)
|
|
# torch serialized
|
|
obj = torch.load(input_path, map_location='cpu')
|
|
if isinstance(obj, Mapping):
|
|
return obj
|
|
raise TypeError("Unsupported checkpoint format without huggingface_hub. Provide a mapping-like object.")
|
|
|
|
|
|
def parse_args(argv=None):
|
|
p = argparse.ArgumentParser(description="Convert WAN <-> Diffusers state_dict key names.")
|
|
p.add_argument("--input", "-i", required=True, help="Path to input checkpoint file (.pt/.bin/.safetensors)")
|
|
p.add_argument("--output", "-o", required=True, help="Output directory (for shards) or .safetensors file")
|
|
p.add_argument("--direction", "-d", choices=["wan-to-diffusers", "diffusers-to-wan"], default="wan-to-diffusers",
|
|
help="Conversion direction")
|
|
p.add_argument("--inner-key", "-k", default="", help="WAN->Diffusers: unwrap this key. Diffusers->WAN: wrap output under this key.")
|
|
p.add_argument("--max-shard-size", default="10GB", help="Shard size when saving to a directory")
|
|
p.add_argument("--keep-unmatched", action="store_true", help="Keep keys with no mapping instead of failing")
|
|
p.add_argument("--single-file", action="store_true", help="Save a single .safetensors file instead of shards")
|
|
return p.parse_args(argv)
|
|
|
|
|
|
def main(argv=None):
|
|
args = parse_args(argv)
|
|
|
|
print(f"[conversion] Direction: {args.direction}")
|
|
print(f"[conversion] Input: {args.input}")
|
|
save_mode = "torch .pt (forced)" if args.direction == "diffusers-to-wan" else ("single safetensors" if args.single_file else f"sharded (max_shard_size={args.max_shard_size})")
|
|
print(f"[conversion] Output: {args.output} [{save_mode}]")
|
|
|
|
loaded = load_checkpoint(args.input)
|
|
if not isinstance(loaded, Mapping):
|
|
raise TypeError("Loaded checkpoint is not a mapping.")
|
|
|
|
# Behavior of --inner-key differs by direction
|
|
if args.direction == "wan-to-diffusers":
|
|
inner, inner_source = select_inner_state_dict(loaded, key=args.inner_key)
|
|
print(f"[conversion] Using inner state_dict: {inner_source}")
|
|
else:
|
|
inner, inner_source = select_inner_state_dict(loaded, key="") # do not unwrap; wrap later if -k is provided
|
|
print(f"[conversion] Using inner state_dict: {inner_source} (ignoring --inner-key for unwrap; will wrap on save)")
|
|
print(f"[conversion] Parameters found: {len(inner)}")
|
|
|
|
if args.direction == "wan-to-diffusers":
|
|
mapping = _self_forcing_to_diffusers_param_names_mapping
|
|
add_dummy = True
|
|
else:
|
|
mapping = invert_mapping(_self_forcing_to_diffusers_param_names_mapping)
|
|
add_dummy = False
|
|
print(f"[conversion] Mapping rules: {len(mapping)}")
|
|
|
|
new_state, stats = convert_state_dict(inner, mapping, strict=not args.keep_unmatched, add_norm_added_q_dummy=add_dummy)
|
|
examples = getattr(new_state, "_examples", [])
|
|
if examples:
|
|
print("[conversion] Sample key mappings:")
|
|
for old_k, new_k in examples[:5]:
|
|
print(f" - {old_k} -> {new_k}")
|
|
print(f"[conversion] Converted parameters: {len(new_state)} (matched={stats['matched']}, unmatched_kept={stats['unmatched']})")
|
|
if add_dummy:
|
|
print(f"[conversion] Added dummy norm_added_q tensors: {stats['dummy_added']}")
|
|
|
|
print("[conversion] Saving...")
|
|
wrapper_key = args.inner_key if (args.direction == "diffusers-to-wan" and args.inner_key) else ""
|
|
if args.direction == "diffusers-to-wan":
|
|
if wrapper_key:
|
|
print(f"[conversion] Wrapping output under key: {wrapper_key}")
|
|
out_path = coerce_pt_output_path(args.output)
|
|
print(f"[conversion] Final output path: {out_path}")
|
|
save_output(new_state, out_path, shard=False, max_shard_size=args.max_shard_size, wrapper_key=wrapper_key, force_pt=True)
|
|
else:
|
|
save_output(new_state, args.output, shard=not args.single_file, max_shard_size=args.max_shard_size, wrapper_key=wrapper_key)
|
|
print("[conversion] Done.")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
try:
|
|
main()
|
|
except Exception as e:
|
|
print(f"[conversion] Error: {e}", file=sys.stderr)
|
|
sys.exit(1)
|