Spaces:
Running on Zero
Running on Zero
File size: 2,625 Bytes
dc09dcc | 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 | # Modified for SplatAnything: inference-only release.
"""Load complete inference checkpoints; reject missing or mismatched model tensors."""
from pathlib import Path
import torch
from safetensors.torch import load_file
def load_checkpoint(path):
path = Path(path)
if not path.is_file():
raise FileNotFoundError(f"Checkpoint not found: {path}")
if path.suffix == '.safetensors':
state = load_file(str(path))
else:
checkpoint = torch.load(str(path), map_location='cpu', mmap=True, weights_only=True)
state = checkpoint.get('state_dict', checkpoint)
state = {k.removeprefix('model.'): v for k, v in state.items()}
return state
def load_model_weights(model, path):
state = load_checkpoint(path)
adapt_depth_mask_checkpoint(model, state)
expected = model.state_dict()
missing = [k for k in expected if k not in state]
mismatched = [k for k in expected if k in state and expected[k].shape != state[k].shape]
if missing or mismatched:
raise ValueError(f"Incomplete/incompatible checkpoint {path}: "
f"missing={missing[:10]}, shape mismatches={mismatched[:10]}. "
"Use complete stage-1 / WorldMirror-2 geometry weights.")
model.load_state_dict({k: state[k] for k in expected}, strict=True)
print(f"[Weights] {type(model).__name__}: {len(expected)} tensors loaded from {path}")
def adapt_depth_mask_checkpoint(model: torch.nn.Module, state_dict: dict) -> None:
"""Expand a 2-channel pretrained depth output conv to the model's 3-channel mask head.
Mirrors wrapper._adapt_depth_mask_checkpoint: when the model was built with
enable_depth_mask=True but the base pretrained weights only have the 2-channel
(depth, conf) output conv, copy those two channels into the 3-channel target and leave
the extra mask channel at init. Without this the whole conv shape-mismatches and is
skipped -> depth channels stay random and predicted depth is garbage. In-place on
state_dict; no-op if the model has no depth mask.
"""
if not getattr(model, "enable_depth_mask", False):
return
current = model.state_dict()
for suffix in ("weight", "bias"):
key = f"depth_head.scratch.output_conv2.2.{suffix}"
if key not in state_dict or key not in current:
continue
ckpt_t, cur_t = state_dict[key], current[key]
if ckpt_t.shape[0] == cur_t.shape[0] - 1 and ckpt_t.shape[1:] == cur_t.shape[1:]:
new_t = cur_t.clone()
new_t[: ckpt_t.shape[0]] = ckpt_t
state_dict[key] = new_t
|