deploy changes

Author: github-actions[bot]

asparagus/functional/huggingface.py

@@ -0,0 +1,62 @@
+"""Download and load models from HuggingFace Hub."""
+
+from huggingface_hub import hf_hub_download, list_repo_files
+
+
+def download_hf_checkpoint(repo_id: str) -> str:
+    """Download a checkpoint file from HuggingFace Hub. Returns local path."""
+    all_files = list_repo_files(repo_id)
+    checkpoint_files = [f for f in all_files if f.endswith((".ckpt", ".pt", ".pth"))]
+
+    if not checkpoint_files:
+        raise ValueError(f"No checkpoint files found in {repo_id}. Available: {all_files}")
+
+    filename = checkpoint_files[0]
+    print(f"Downloading {filename} from {repo_id}...")
+
+    return hf_hub_download(repo_id, filename)
+
+
+class HuggingFaceWeightMapper:
+    """Remaps weights to asparagus format."""
+
+    def __init__(self, state_dict: dict):
+        self.state_dict = state_dict
+
+    def remap_keys(self) -> dict:
+        """Add 'model.' prefix for Lightning module compatibility. Subclasses should override and call super()."""
+        first_key = next(iter(self.state_dict.keys()))
+        if not first_key.startswith("model."):
+            return {f"model.{k}": v for k, v in self.state_dict.items()}
+        return self.state_dict
+
+
+class OpenMindResEncWeightMapper(HuggingFaceWeightMapper):
+    """Remaps OpenMind ResEncUNet weights to asparagus format."""
+
+    def remap_keys(self) -> dict:
+        original_keys = self.state_dict.keys()
+        self.state_dict = {
+            k.replace(".convs.0.", ".conv1.").replace(".norm.", ".norm_op."): v for k, v in self.state_dict.items()
+        }
+        if self.state_dict.keys() != original_keys:
+            print("Remapped OpenMind ResEncUNet keys to asparagus naming conventions.")
+
+        return super().remap_keys()
+
+
+class OpenMindPrimusWeightMapper(HuggingFaceWeightMapper):
+    """Remaps OpenMind Primus weights to asparagus format."""
+
+    def remap_keys(self) -> dict:
+        original_keys = self.state_dict.keys()
+        self.state_dict = {
+            k.replace("encoder.eva.", "eva.")
+            .replace("encoder.down_projection.proj.", "encoder.proj.")
+            .replace("encoder.mask_token", "mask_token"): v
+            for k, v in self.state_dict.items()
+        }
+        if self.state_dict.keys() != original_keys:
+            print("Remapped OpenMind Primus keys to asparagus naming conventions.")
+
+        return super().remap_keys()

asparagus/modules/dataclasses/pathing.py

@@ -1,10 +1,11 @@
 from dataclasses import dataclass
+from typing import Optional
 
 
 @dataclass
 class PathingConfig:
     run_dir: str
     ckpt_save_dir: str
-    ckpt_path: str
-    ckpt_parent_folder: str
+    ckpt_path: Optional[str]
+    ckpt_parent_folder: Optional[str]
     dataset_json_path: str

asparagus/modules/lightning_modules/base_module.py

@@ -29,16 +29,16 @@ def __init__(
         decoder_warmup_epochs: int = 0,
         cosine_period_ratio: float = 1,
         compile_mode: str = None,
-        weights: str = None,
+        weights: dict = None,
         load_decoder: bool = True,
-        repeat_stem_weights: bool = True,
         optimizer: str = "SGD",
         train_transforms: Optional[transforms.Compose] = None,
         test_transforms: Optional[transforms.Compose] = None,
         val_transforms: Optional[transforms.Compose] = None,
         weight_decay: float = 3e-5,
         nesterov: bool = True,
         momentum: float = 0.99,
+        repeat_stem_weights: bool = True,
     ):
         super().__init__()
         self.learning_rate = learning_rate
@@ -60,11 +60,11 @@ def __init__(
         self.repeat_stem_weights = repeat_stem_weights
         assert 0 < cosine_period_ratio <= 1
 
-        self.save_hyperparameters(ignore=["model", "train_transforms", "val_transforms", "test_transforms"])
+        self.save_hyperparameters(ignore=["model", "weights", "train_transforms", "val_transforms", "test_transforms"])
         self.model = model
 
         if weights is not None:
-            self.load_weights(weights, load_decoder=load_decoder)
+            self.load_state_dict(weights, load_decoder=load_decoder, strict=False)
 
         self.model = torch.compile(model, mode=compile_mode) if compile_mode is not None else model
 
@@ -143,11 +143,6 @@ def configure_optimizers(self):
 
         return [optimizer], [scheduler_config]
 
-    def load_weights(self, weights, load_decoder=True):
-        ckpt = torch.load(weights, map_location="cpu", weights_only=False)
-        print(f"Loading weights trained for {ckpt['global_step']} steps / {ckpt['epoch']} epochs.")
-        self.load_state_dict(ckpt["state_dict"], load_decoder=load_decoder, strict=False)
-
     def load_state_dict(self, state_dict, load_decoder=True, *args, **kwargs):
         old_params = copy.deepcopy(self.state_dict())
 
@@ -161,10 +156,10 @@ def load_state_dict(self, state_dict, load_decoder=True, *args, **kwargs):
             state_dict = {k.replace("_orig_mod.", ""): v for k, v in state_dict.items()}
 
         # Repeat stem weights when state_dict num_channels is smaller than new_state_dict num_channels
-        if self.model.stem_weight_name is not None and self.repeat_stem_weights:
+        if hasattr(self.model, "stem_weight_name") and self.model.stem_weight_name is not None and self.repeat_stem_weights:
             prefix = "model._orig_mod." if "_orig_mod" in list(state_dict.keys())[0] else "model."
             stem_name = f"{prefix}{self.model.stem_weight_name}"
-            pt_input_channels = state_dict[stem_name].shape[1]  # (N, C, H, W, Z) where N is num tokens.
+            pt_input_channels = state_dict[stem_name].shape[1]
             ft_input_channels = old_params[stem_name].shape[1]
             if pt_input_channels < ft_input_channels:
                 assert pt_input_channels == 1, (

asparagus/modules/lightning_modules/clsreg_module.py

@@ -23,7 +23,7 @@ def __init__(
         decoder_warmup_epochs: int = 0,
         cosine_period_ratio: float = 1,
         compile_mode: str = None,
-        weights: str = None,
+        weights: dict = None,
         optimizer: str = "SGD",
         train_transforms: Optional[transforms.Compose] = None,
         test_transforms: Optional[transforms.Compose] = None,
@@ -75,6 +75,7 @@ def training_step(self, batch, batch_idx):
 
         pred = self.model(x)
         loss = self.loss(pred, y)
+
         self.log(
             "train/loss", loss, on_step=False, on_epoch=True, sync_dist=True, batch_size=self.trainer.datamodule.batch_size
         )

asparagus/modules/lightning_modules/segmentation_module.py

@@ -44,7 +44,7 @@ def __init__(
         decoder_warmup_epochs: int = 0,
         cosine_period_ratio: float = 1,
         compile_mode: str = None,
-        weights: str = None,
+        weights: dict = None,
         deep_supervision: bool = False,
         train_transforms: Optional[transforms.Compose] = None,
         test_transforms: Optional[transforms.Compose] = None,

asparagus/pipeline/auto_configuration/checkpoint.py

@@ -0,0 +1,50 @@
+import os
+import torch
+from asparagus.functional.huggingface import download_hf_checkpoint
+from asparagus.functional.versioning import detect_id
+from hydra.utils import get_class
+
+
+def load_checkpoint_state_dict(path):
+    """Load a checkpoint file and return the state_dict."""
+    ckpt = torch.load(path, map_location="cpu", weights_only=False)
+
+    if "state_dict" in ckpt:
+        print(f"Loading weights trained for {ckpt.get('global_step', '?')} steps / {ckpt.get('epoch', '?')} epochs.")
+        return ckpt["state_dict"]
+    elif "network_weights" in ckpt:
+        print("Loading weights from external checkpoint (network_weights key).")
+        return ckpt["network_weights"]
+    else:
+        raise ValueError("Unsupported checkpoint format. Expected 'state_dict' or 'network_weights' key.")
+
+
+def resolve_checkpoint_path(cfg):
+    """Resolve checkpoint file path from config. Returns path or None."""
+    if cfg.checkpoint_run_id:
+        folder = detect_id(cfg.checkpoint_run_id)
+        return os.path.join(folder, "checkpoints", cfg.load_checkpoint_name)
+    if cfg.checkpoint_path:
+        return cfg.checkpoint_path
+    return None
+
+
+def resolve_checkpoint(cfg):
+    """Resolve and load checkpoint from config. Returns a state_dict or None."""
+    hf_id = getattr(cfg, "hf_model_id", None) or None
+    ckpt_path = resolve_checkpoint_path(cfg)
+
+    sources = [s for s in [ckpt_path, hf_id] if s]
+    if len(sources) > 1:
+        raise ValueError("Provide only one of: checkpoint_run_id, checkpoint_path, hf_model_id")
+    if len(sources) == 0:
+        return None
+
+    if ckpt_path:
+        return load_checkpoint_state_dict(ckpt_path)
+
+    path = download_hf_checkpoint(hf_id)
+    state_dict = load_checkpoint_state_dict(path)
+
+    weight_mapper = get_class(cfg.hf_weight_format)
+    return weight_mapper(state_dict).remap_keys()

asparagus/pipeline/auto_configuration/versioning.py

@@ -1,6 +1,7 @@
 import os
 from asparagus.functional.versioning import detect_id, detect_mlflow_id, detect_wandb_id
 from asparagus.modules.dataclasses import PathingConfig, VersioningConfig
+from asparagus.pipeline.auto_configuration.checkpoint import resolve_checkpoint_path
 from hydra.core.hydra_config import HydraConfig
 
 
@@ -21,25 +22,18 @@ def pathing(cfg, train=True):
     run_dir = HydraConfig.get().runtime.output_dir
     os.makedirs(run_dir, exist_ok=True)
 
-    if cfg.checkpoint_run_id is not None and cfg.checkpoint_run_id != "":
-        model_folder = detect_id(cfg.checkpoint_run_id)
-        pretrained_ckpt = os.path.join(model_folder, "checkpoints", cfg.load_checkpoint_name)
-        assert cfg.checkpoint_path is None, "You cannot provide both a checkpoint path and a checkpoint run id"
-    elif cfg.checkpoint_path is not None and cfg.checkpoint_path != "":
-        model_folder = None
-        pretrained_ckpt = cfg.checkpoint_path
-    else:
-        model_folder, pretrained_ckpt = None, None
+    ckpt_path = resolve_checkpoint_path(cfg)
+    ckpt_parent_folder = detect_id(cfg.checkpoint_run_id) if cfg.checkpoint_run_id else None
 
     if train:
         dataset_json_path = cfg.data.data_path + "/dataset.json"
     else:
         dataset_json_path = cfg.data.test_data_path + "/dataset.json"
-    pathingcfg = PathingConfig(
+
+    return PathingConfig(
         run_dir=run_dir,
         ckpt_save_dir=os.path.join(run_dir, "checkpoints"),
-        ckpt_parent_folder=model_folder,
-        ckpt_path=pretrained_ckpt,
+        ckpt_parent_folder=ckpt_parent_folder,
+        ckpt_path=ckpt_path,
         dataset_json_path=dataset_json_path,
     )
-    return pathingcfg

asparagus/pipeline/run/finetune_cls.py

@@ -6,6 +6,7 @@
 from asparagus.modules.hydra.plugins.searchpath_plugins import FinetuneSearchpathPlugin
 from asparagus.modules.transforms.presets import CPU_clsreg_val_test_transforms_crop
 from asparagus.paths import get_config_path
+from asparagus.pipeline.auto_configuration.checkpoint import resolve_checkpoint
 from asparagus.pipeline.auto_configuration.experiment_setup import (
     prepare_standard_experiment,
 )
@@ -39,6 +40,7 @@ def main(cfg: DictConfig) -> None:
     print(f"{OmegaConf.to_yaml(cfg)}\n Version: {cfg.run_id}\n Run dir: {HydraConfig.get().run.dir}\n")
     logging_safe_cfg = OmegaConf.to_container(cfg, resolve=True, throw_on_missing=True)
     file_store, path_store, version_store = prepare_standard_experiment(cfg)
+    weights = resolve_checkpoint(cfg)
     pl.seed_everything(seed=cfg.training.seed, workers=True)
 
     loggers = logging(
@@ -107,7 +109,7 @@ def main(cfg: DictConfig) -> None:
         decoder_warmup_epochs=cfg.training.decoder_warmup_epochs,
         train_transforms=gpu_tr_transforms,
         val_transforms=None,
-        weights=path_store.ckpt_path,
+        weights=weights,
         log_image_every_n_epochs=cfg.logger.log_images_every_n_epoch,
         optimizer=cfg.model.finetune_optim,
         learning_rate=cfg.model.finetune_lr,

asparagus/pipeline/run/finetune_reg.py

@@ -6,6 +6,7 @@
 from asparagus.modules.hydra.plugins.searchpath_plugins import FinetuneSearchpathPlugin
 from asparagus.modules.transforms.presets import CPU_clsreg_val_test_transforms_crop
 from asparagus.paths import get_config_path
+from asparagus.pipeline.auto_configuration.checkpoint import resolve_checkpoint
 from asparagus.pipeline.auto_configuration.experiment_setup import (
     prepare_standard_experiment,
 )
@@ -38,6 +39,7 @@
 def main(cfg: DictConfig) -> None:
     print(f"{OmegaConf.to_yaml(cfg)}\n Version: {cfg.run_id}\n Run dir: {HydraConfig.get().run.dir}\n")
     file_store, path_store, version_store = prepare_standard_experiment(cfg)
+    weights = resolve_checkpoint(cfg)
 
     pl.seed_everything(seed=cfg.training.seed, workers=True)
 
@@ -105,7 +107,7 @@ def main(cfg: DictConfig) -> None:
         model=model,
         train_transforms=gpu_tr_transforms,
         val_transforms=None,
-        weights=path_store.ckpt_path,
+        weights=weights,
         log_image_every_n_epochs=cfg.logger.log_images_every_n_epoch,
         optimizer=cfg.model.finetune_optim,
         learning_rate=cfg.model.finetune_lr,

asparagus/pipeline/run/finetune_seg.py

@@ -6,6 +6,7 @@
 from asparagus.modules.hydra.plugins.searchpath_plugins import FinetuneSearchpathPlugin
 from asparagus.modules.transforms.presets import CPU_seg_test_transforms
 from asparagus.paths import get_config_path
+from asparagus.pipeline.auto_configuration.checkpoint import resolve_checkpoint
 from asparagus.pipeline.auto_configuration.experiment_setup import (
     prepare_standard_experiment,
 )
@@ -39,10 +40,10 @@ def main(cfg: DictConfig) -> None:
     print(f"{OmegaConf.to_yaml(cfg)}\n Version: {cfg.run_id}\n Run dir: {HydraConfig.get().run.dir}\n")
     logging_safe_cfg = OmegaConf.to_container(cfg, resolve=True, throw_on_missing=True)
     file_store, path_store, version_store = prepare_standard_experiment(cfg)
+    weights = resolve_checkpoint(cfg)
     pl.seed_everything(seed=cfg.training.seed, workers=True)
 
     assert "load_checkpoint_name" in cfg.keys(), "load_checkpoint_name not in config. Did you supply a scratch config?"
-    assert path_store.ckpt_path is not None, "Checkpoint must be provided for finetuning."
 
     loggers = logging(
         ckpt_wandb_id=version_store.wandb_id,
@@ -112,7 +113,7 @@ def main(cfg: DictConfig) -> None:
         model=model,
         warmup_epochs=cfg.training.warmup_epochs,
         decoder_warmup_epochs=cfg.training.decoder_warmup_epochs,
-        weights=path_store.ckpt_path,
+        weights=weights,
         train_transforms=gpu_tr_transforms,
         val_transforms=None,
         optimizer=cfg.model.finetune_optim,