fix: 🐛 Fix COSEM model loading.

src/cellmap_models/pytorch/__init__.py

@@ -1,2 +1 @@
-from . import cosem
-from . import cellpose
+from . import cosem, cellpose, untrained_models

src/cellmap_models/pytorch/cosem/README.md

@@ -25,7 +25,7 @@ Each model has a separate backbone and single layer prediction head. The `backbo
 ```python
 import cellmap_models.cosem as cosem_models
 model = cosem_models.load_model('setup04/1820500')
-backnone = model.backbone
+backbone = model.backbone
 head = model.prediction_head
 ```
 

src/cellmap_models/pytorch/cosem/load_model.py

@@ -81,6 +81,7 @@ def load_model(checkpoint_name: str) -> torch.nn.Module:
             new_checkpoint["model"].pop(key)
             continue
         new_key = key.replace("architecture.", "")
+        new_key = new_key.replace("unet.", "backbone.")
         new_checkpoint["model"][new_key] = new_checkpoint["model"].pop(key)
     model.load_state_dict(new_checkpoint["model"])
     model.eval()