Support bare state-dict checkpoint layout

app.py

@@ -96,10 +96,28 @@ print("Fetching MiniWorld-1B (RealEstate10K) ...", flush=True)
 ckpt_path = hf_hub_download(MINIWORLD_REPO, MINIWORLD_CKPT)
 
 _ckpt = torch.load(ckpt_path, map_location="cpu")
-_weights = _ckpt.get("ema_model", _ckpt.get("model"))
-_meta = _ckpt.get("meta", {}) or {}
+_meta: dict = {}
+_weights = None
+if isinstance(_ckpt, dict):
+    # `miniworld/sample.py` expects a training checkpoint wrapper; the *released*
+    # weights are a bare state dict of `net.*` tensors, so support both.
+    for _key in ("ema_model", "model", "ema", "state_dict", "module"):
+        cand = _ckpt.get(_key)
+        if isinstance(cand, dict) and cand:
+            _weights = cand
+            _meta = _ckpt.get("meta") or {}
+            print(f"[Checkpoint] using wrapped weights under {_key!r}", flush=True)
+            break
+    if _weights is None and any(
+        isinstance(k, str) and k.startswith("net.") for k in _ckpt
+    ):
+        _weights = _ckpt
+        print("[Checkpoint] bare state dict (no training wrapper)", flush=True)
 if _weights is None:
-    raise RuntimeError("MiniWorld checkpoint has neither `ema_model` nor `model`")
+    raise RuntimeError(
+        "Unrecognised MiniWorld checkpoint layout; top-level keys: "
+        f"{list(_ckpt)[:8] if isinstance(_ckpt, dict) else type(_ckpt)}"
+    )
 
 
 def _resolve_latent_frames() -> int:
@@ -147,9 +165,15 @@ denoiser = build_denoiser_from_mode(
         df_ardiff_step=DF_ARDIFF_STEP,
     )
 ).eval()
-denoiser.load_state_dict(_weights, strict=True)
+_missing, _unexpected = denoiser.load_state_dict(_weights, strict=False)
+if _missing or _unexpected:
+    raise RuntimeError(
+        f"MiniWorld checkpoint does not match the built model.\n"
+        f"missing ({len(_missing)}): {_missing[:12]}\n"
+        f"unexpected ({len(_unexpected)}): {_unexpected[:12]}"
+    )
 denoiser.trained_num_frames = TRAINED_NUM_FRAMES
-print("[Checkpoint] loaded: all keys matched", flush=True)
+print(f"[Checkpoint] loaded: all {len(_weights)} keys matched", flush=True)
 del _ckpt, _weights
 
 denoiser = denoiser.to("cuda")