xsl
24c526af68
- 模型 17/17 文件经镜像全量缓存（gated+token，直连本机不通；sfx.md/setup.md/asset-pipeline.md 镜像结论同步更新） - 修复迁移遗留：stable_audio_3/models/ 源码子包被 .gitignore models/ 规则误伤丢失，从上游恢复并加白名单 - radio：arg_muffled/brake_screech 等叙事件落地（接线原有 L111/L154 零改动）+ phone_ring SYN→SA3 升级；dungeon 13 源 A/B 对齐（RMS 偏差 ≤0.3dB）全量同名替换，SYN 原版备份 tools/asset_fix_backup/sa3_ab_20260902/ - 关键玩法语义达标：echo 距离档 0.83<1.43<2.18<2.98s 严格递增；one-shot 真峰值全 ≤-3dBFS（解码回读法） - check_all_assets 8 游戏全绿（radio 254/254、dungeon 303/303）；playtest radio 221/dungeon 1610 全绿
112 lines
5.0 KiB
Python
112 lines
5.0 KiB
Python
"""Helpers for loading pre-trained LoRA checkpoints onto a live model."""
|
|
import os
|
|
from functools import partial
|
|
|
|
import torch
|
|
|
|
from .model import LoRAParametrization, add_lora
|
|
from .utils import (
|
|
get_lora_layers,
|
|
infer_global_rank,
|
|
load_lora_checkpoint,
|
|
prepare_dora_state_dict,
|
|
remap_lora_state_dict,
|
|
resolve_adapter_type,
|
|
)
|
|
from ...verbose import vprint
|
|
|
|
|
|
def load_and_apply_loras(model, lora_ckpt_paths, model_type, svd_bases_path=None):
|
|
"""Load LoRA checkpoints from disk and attach them to `model`.
|
|
|
|
- Uses a two-pass approach: first resolves each LoRA's adapter type, then
|
|
loads SVD bases (only if any LoRA is -XS), then applies each LoRA.
|
|
- Handles legacy "dora" adapter_type via `resolve_adapter_type`.
|
|
- Only passes svd_bases to -XS LoRAs to avoid spurious "no -XS layers"
|
|
messages for non-XS LoRAs.
|
|
- Sets `model.use_lora = True` and `model.lora_names = [...]` as a
|
|
convenience for downstream UI code.
|
|
|
|
Args:
|
|
model: The loaded model. For `diffusion_cond` / `diffusion_cond_inpaint`
|
|
types, LoRAs are applied to `model.model` and `model.conditioner`;
|
|
otherwise they are applied to `model` directly.
|
|
lora_ckpt_paths: List of paths to `.ckpt` or `.safetensors` files.
|
|
model_type: String from the model config, e.g. "diffusion_cond".
|
|
svd_bases_path: Optional path to a precomputed SVD bases .pt file
|
|
(used only by -XS adapter types).
|
|
|
|
Returns:
|
|
List of display names (file stems) in the same order as the input paths.
|
|
"""
|
|
if not lora_ckpt_paths:
|
|
model.use_lora = False
|
|
model.lora_names = []
|
|
return []
|
|
|
|
is_cond = model_type in ("diffusion_cond", "diffusion_cond_inpaint")
|
|
|
|
# Pass 1: load each checkpoint and resolve adapter type.
|
|
lora_entries = [] # (path, state_dict, config_dict, adapter_type)
|
|
for i, lora_path in enumerate(lora_ckpt_paths):
|
|
print(f"Loading LoRA {i} from {lora_path}")
|
|
state_dict, config_dict = load_lora_checkpoint(lora_path)
|
|
adapter_type_raw = config_dict.get("adapter_type", "lora")
|
|
adapter_type = resolve_adapter_type(adapter_type_raw, state_dict)
|
|
if adapter_type != adapter_type_raw:
|
|
vprint(f"Resolved legacy '{adapter_type_raw}' -> {adapter_type}")
|
|
lora_entries.append((lora_path, state_dict, config_dict, adapter_type))
|
|
|
|
# Load SVD bases only if at least one LoRA is -XS.
|
|
model_svd_bases = None
|
|
cond_svd_bases = None
|
|
any_xs = any(t.endswith("-xs") for _, _, _, t in lora_entries)
|
|
if any_xs:
|
|
if svd_bases_path is not None:
|
|
vprint(f"Loading SVD bases from {svd_bases_path}")
|
|
svd_bases_all = torch.load(svd_bases_path, map_location="cpu", weights_only=True)
|
|
model_svd_bases = {k[len("model."):]: v for k, v in svd_bases_all.items() if k.startswith("model.")}
|
|
cond_svd_bases = {k[len("conditioner."):]: v for k, v in svd_bases_all.items() if k.startswith("conditioner.")}
|
|
vprint(f" model bases: {len(model_svd_bases)}, conditioner bases: {len(cond_svd_bases)}")
|
|
else:
|
|
print("WARNING: -XS adapter type present without svd_bases_path -- SVD will be computed on current device")
|
|
|
|
# Pass 2: apply each LoRA.
|
|
lora_names = []
|
|
for i, (lora_path, state_dict, config_dict, adapter_type) in enumerate(lora_entries):
|
|
rank = config_dict.get("rank", infer_global_rank(state_dict))
|
|
alpha = config_dict.get("alpha", rank)
|
|
include = config_dict.get("include", None)
|
|
exclude = config_dict.get("exclude", None)
|
|
is_xs = adapter_type.endswith("-xs")
|
|
|
|
lora_config = {
|
|
torch.nn.Linear: {
|
|
"weight": partial(LoRAParametrization.from_linear, rank=rank, lora_alpha=alpha, adapter_type=adapter_type, lora_index=i),
|
|
},
|
|
torch.nn.Conv1d: {
|
|
"weight": partial(LoRAParametrization.from_conv1d, rank=rank, lora_alpha=alpha, adapter_type=adapter_type, lora_index=i),
|
|
},
|
|
}
|
|
if is_cond:
|
|
add_lora(model.model, lora_config, include=include, exclude=exclude,
|
|
svd_bases=model_svd_bases if is_xs else None)
|
|
add_lora(model.conditioner, lora_config, include=include, exclude=exclude,
|
|
svd_bases=cond_svd_bases if is_xs else None)
|
|
else:
|
|
add_lora(model, lora_config, include=include, exclude=exclude,
|
|
svd_bases=model_svd_bases if is_xs else None)
|
|
|
|
prepare_dora_state_dict(state_dict)
|
|
remapped_sd = remap_lora_state_dict(state_dict, i)
|
|
if is_cond:
|
|
model.model.load_state_dict(remapped_sd, strict=False)
|
|
model.conditioner.load_state_dict(remapped_sd, strict=False)
|
|
else:
|
|
model.load_state_dict(remapped_sd, strict=False)
|
|
lora_names.append(os.path.splitext(os.path.basename(lora_path))[0])
|
|
|
|
vprint("lora layers:", len(get_lora_layers(model)))
|
|
model.use_lora = True
|
|
model.lora_names = lora_names
|
|
return lora_names |