Spaces:
Running on Zero
Running on Zero
File size: 7,026 Bytes
9368ee7 | 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 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 | """Shared student-model loader for the inference scripts.
``inference_segmentwise.py`` and ``inference_streaming.py`` both import
``load_diffusion_model`` from here so the 3-source assembly (base Wan 2.1 T2V +
OmniAvatar-LS adapter + Self-Forcing student) and the post-load LoRA merge live
in exactly one place.
This module intentionally keeps only light top-level deps (``os``, ``torch``);
the heavy model import is lazy (inside the function), so tools that merely need
the loader (e.g. ``scripts/export_merged_checkpoint.py``) don't pull in the
video/audio stack.
"""
import os
import torch
def load_diffusion_model(args, device, dtype):
"""Load the CausalOmniAvatarWan student model.
Constructs the model, loads base Wan + OmniAvatar weights (if paths given),
then overlays the Self-Forcing trained checkpoint on top.
Returns:
CausalOmniAvatarWan model in eval mode on *device*.
"""
from lipforcing.networks.OmniAvatar.network_causal import CausalOmniAvatarWan
# 14B SF LoRA: construct with merge_lora=False so PEFT is injected and
# the model exposes lora_A/lora_B keys that the trainable-filtered SF
# state_dict expects. After load_state_dict the LoRA values are
# populated; we then merge them into base for inference speed (unless
# --no_merge_lora_post_load is passed).
#
# For 1.3B, merge_lora=True folds the OmniAvatar V2V LoRA into base at
# construction, so the SF checkpoint is expected to have plain
# (already-merged) keys.
constructor_merge_lora = (args.model_size == "1.3B")
model = CausalOmniAvatarWan(
model_size=args.model_size,
in_dim=65,
mode="v2v",
use_audio=True,
audio_hidden_size=32,
chunk_size=args.chunk_size,
total_num_frames=21,
base_model_paths=args.base_model_paths,
omniavatar_ckpt_path=args.omniavatar_ckpt_path,
merge_lora=constructor_merge_lora,
lora_rank=128,
lora_alpha=64,
net_pred_type="flow",
schedule_type="rf",
mask_all_frames=True,
dtype=args.dtype,
local_attn_size=args.local_attn_size,
sink_size=args.sink_size,
use_dynamic_rope=args.use_dynamic_rope,
)
# Load Self-Forcing checkpoint on top
# Supports: regular .pt/.pth, FSDP distcp directory, or .pth + adjacent distcp dir
print(f"Loading SF checkpoint from {args.ckpt_path} ...")
state_dict = None
# Check for FSDP distributed checkpoint: look for .net_model/ directory
ckpt_stem = args.ckpt_path.replace(".pth", "")
fsdp_net_dir = ckpt_stem + ".net_model"
if os.path.isdir(fsdp_net_dir):
# FSDP2 distributed checkpoint — load via torch.distributed.checkpoint
print(f" Loading FSDP distributed checkpoint from {fsdp_net_dir} ...")
from torch.distributed.checkpoint import FileSystemReader
from torch.distributed.checkpoint.state_dict_loader import load as dcp_load
reader = FileSystemReader(fsdp_net_dir)
md = reader.read_metadata()
state_dict = {}
for key, meta in md.state_dict_metadata.items():
if hasattr(meta, "size"):
state_dict[key] = torch.empty(meta.size)
dcp_load(state_dict, storage_reader=reader, no_dist=True)
print(f" Loaded {len(state_dict)} tensors from FSDP distcp")
elif os.path.isdir(args.ckpt_path):
# Direct distcp directory path
from torch.distributed.checkpoint import FileSystemReader
from torch.distributed.checkpoint.state_dict_loader import load as dcp_load
reader = FileSystemReader(args.ckpt_path)
md = reader.read_metadata()
state_dict = {}
for key, meta in md.state_dict_metadata.items():
if hasattr(meta, "size"):
state_dict[key] = torch.empty(meta.size)
dcp_load(state_dict, storage_reader=reader, no_dist=True)
print(f" Loaded {len(state_dict)} tensors from distcp directory")
else:
# Regular .pt/.pth checkpoint
ckpt = torch.load(args.ckpt_path, map_location="cpu", weights_only=False)
if isinstance(ckpt, dict):
if "model" in ckpt and isinstance(ckpt["model"], dict) and "net" in ckpt["model"]:
state_dict = ckpt["model"]["net"]
elif "net" in ckpt:
state_dict = ckpt["net"]
else:
state_dict = ckpt
else:
state_dict = ckpt
# Keys in checkpoint use plain names (e.g. "patch_embedding.weight")
# Model wraps everything under _core, so keys are "_core.xxx"
# Try loading directly first, if too many missing, try adding _core prefix
missing, unexpected = model.load_state_dict(state_dict, strict=False)
if len(missing) > len(state_dict) * 0.5 and not any(k.startswith("_core.") for k in state_dict):
# Try with _core prefix
prefixed_sd = {"_core." + k: v for k, v in state_dict.items()}
missing2, unexpected2 = model.load_state_dict(prefixed_sd, strict=False)
if len(missing2) < len(missing):
missing, unexpected = missing2, unexpected2
print(" Applied _core. prefix for key matching")
print(f" SF checkpoint: {len(state_dict)} params, {len(missing)} missing, {len(unexpected)} unexpected")
# Merge LoRA into base post-load (14B path only; the 1.3B path constructs
# with merge_lora=True so there's nothing to merge here).
#
# Uses PEFT's per-layer merge() which fuses LoRA A/B into the wrapped
# base_layer.weight in-place. After merge, LoraLinear.forward dispatches
# to base_layer.forward and the LoRA path is a no-op — same numeric
# result as not having LoRA at all. We keep the LoraLinear wrappers in
# place (no unload) because that's a structural change and not necessary
# for correctness; only the merge_count is reported.
if args.model_size == "14B" and args.merge_lora_post_load:
print(" Merging LoRA into base for inference speed...")
from peft.tuners.lora import LoraLayer
merge_count = 0
for name, module in model.named_modules():
if isinstance(module, LoraLayer):
module.merge()
merge_count += 1
if merge_count > 0:
print(f" Merged {merge_count} LoRA layers.")
elif args.base_model_paths is None and args.omniavatar_ckpt_path is None:
# Consolidated single-file checkpoint: no base/adapter supplied and
# the .pth already carries fully-merged plain weights, so there are
# no PEFT layers to merge. This is the expected published-weights path.
print(" Checkpoint is already consolidated (merged); nothing to merge.")
else:
print(" WARN: no LoraLayer instances found; model has no PEFT "
"adapters to merge (was --model_size set correctly?).")
model = model.to(device=device, dtype=dtype)
model.eval()
return model
|