Fix: Add embedder key remapping for backward compatibility

- Modified load_model_weights to handle old checkpoint format
- Old format: embedder.weight
- New format: encoder.embedder.weight + occupancy_decoder.embedder.weight
- Checkpoint is loaded as dict first, keys are remapped, then loaded to model
- This fixes the AssertionError on HuggingFace Space deployment
- Maintains backward compatibility with both checkpoint formats

code/cube3d/inference/utils.py

@@ -3,7 +3,7 @@ from typing import Any, Optional, Tuple
 
 import torch
 from omegaconf import DictConfig, OmegaConf
-from safetensors.torch import load_model
+from safetensors.torch import load_model, load_file
 
 BOUNDING_BOX_MAX_SIZE = 1.925
 
@@ -49,6 +49,10 @@ def load_model_weights(model: torch.nn.Module, ckpt_path: str) -> None:
     Load a safetensors checkpoint into a PyTorch model.
     The model is updated in place.
 
+    Handles backward compatibility for embedder weight key naming:
+    - Old format: 'embedder.weight'
+    - New format: 'encoder.embedder.weight', 'occupancy_decoder.embedder.weight'
+
     Args:
         model: PyTorch model to load weights into
         ckpt_path: Path to the safetensors checkpoint file
@@ -59,9 +63,21 @@ def load_model_weights(model: torch.nn.Module, ckpt_path: str) -> None:
     assert ckpt_path.endswith(
         ".safetensors"
     ), f"Checkpoint path '{ckpt_path}' is not a safetensors file"
-    
-    #load_model(model, ckpt_path)
-    load_model(model, ckpt_path, strict=False)
+
+    # Load checkpoint as dictionary for key remapping
+    checkpoint = load_file(ckpt_path)
+
+    # Backward compatibility: remap old embedder key format to new format
+    # This handles cases where checkpoint has 'embedder.weight' but model expects
+    # 'encoder.embedder.weight' and 'occupancy_decoder.embedder.weight'
+    if 'embedder.weight' in checkpoint:
+        if 'encoder.embedder.weight' not in checkpoint:
+            checkpoint['encoder.embedder.weight'] = checkpoint['embedder.weight']
+        if 'occupancy_decoder.embedder.weight' not in checkpoint:
+            checkpoint['occupancy_decoder.embedder.weight'] = checkpoint['embedder.weight']
+
+    # Load remapped checkpoint into model with strict=False for flexibility
+    model.load_state_dict(checkpoint, strict=False)
 
 
 def save_model_weights(model: torch.nn.Module, save_path: str) -> None: