Add support to pass config to model and processor

src/colipri/model/multimodal.py

@@ -8,6 +8,8 @@ from accelerate import init_empty_weights
 from accelerate import load_checkpoint_and_dispatch
 from einops import rearrange
 from hydra.utils import instantiate
+from omegaconf import DictConfig
+from omegaconf import OmegaConf
 from safetensors.torch import load_model
 from safetensors.torch import save_model
 from torch import nn
@@ -38,17 +40,42 @@ from .text import TextEncoder
 def get_model(
     checkpoint_path: TypePath | None = None,
     *,
+    config: TypePath | DictConfig | None = None,
     pretrained: bool = True,
     image_only: bool = False,
     **kwargs,
 ) -> Model:
+    """Create a :class:`Model` instance.
+
+    Args:
+        checkpoint_path: Path to a ``.safetensors`` checkpoint. When ``None``
+            and ``pretrained`` is ``True``, the pretrained weights are
+            downloaded from the Hugging Face Hub.
+        config: Optional custom model config. Can be a path to a YAML file or
+            a ``DictConfig``. When ``None``, the built-in default config is
+            used.
+        pretrained: If ``True`` and ``checkpoint_path`` is ``None``, download
+            the pretrained weights.
+        image_only: If ``True``, the text encoder is disabled.
+        **kwargs: Hydra-style dot-list overrides applied on top of the config.
+
+    Returns:
+        A configured :class:`Model` instance.
+    """
     if pretrained and checkpoint_path is None:
         checkpoint_path = download_weights()
 
-    overrides = []
-    for key, value in kwargs.items():
-        overrides.append(f"{key}={value}")
-    config = load_model_config(overrides=overrides)
+    overrides = [f"{key}={value}" for key, value in kwargs.items()]
+
+    if config is None:
+        config = load_model_config(overrides=overrides or None)
+    else:
+        if isinstance(config, (str, Path)):
+            config = OmegaConf.load(config)
+        assert isinstance(config, DictConfig)
+        if overrides:
+            config = OmegaConf.merge(config, OmegaConf.from_dotlist(overrides))
+            assert isinstance(config, DictConfig)
 
     if image_only:
         config.text_encoder = None

src/colipri/processor.py

@@ -1,8 +1,12 @@
 from __future__ import annotations
 
+from pathlib import Path
+
 import torch
 import torchio as tio
 from hydra.utils import instantiate
+from omegaconf import DictConfig
+from omegaconf import OmegaConf
 from transformers import BertTokenizer
 from transformers.tokenization_utils_base import BatchEncoding
 
@@ -15,12 +19,50 @@ from .types import TypeStringOrStrings
 from .types import TypeTextAttentionMask
 from .types import TypeTextTokenIds
 
+PROCESSOR_TARGET = "colipri.processor.Processor"
+
+
+def get_processor(
+    *,
+    config: TypePath | DictConfig | None = None,
+    image_only: bool = False,
+    **kwargs,
+) -> Processor:
+    """Create a :class:`Processor` instance.
+
+    Args:
+        config: Optional custom config. Can be a path to a YAML file or a
+            ``DictConfig``. If the config is a transform-only config (e.g., as
+            exported by estereo), it is automatically wrapped into a full
+            processor config using the default tokenizer. When ``None``, the
+            built-in default config is used.
+        image_only: If ``True``, the tokenizer is disabled.
+        **kwargs: Hydra-style dot-list overrides applied on top of the config.
+
+    Returns:
+        A configured :class:`Processor` instance.
+    """
+    overrides = [f"{key}={value}" for key, value in kwargs.items()]
+
+    if config is None:
+        config = load_processor_config(overrides=overrides or None)
+    else:
+        if isinstance(config, (str, Path)):
+            config = OmegaConf.load(config)
+        assert isinstance(config, DictConfig)
+
+        # If the config is a transform (not a full Processor config), wrap it
+        # into a complete processor config using the default tokenizer.
+        is_full_processor_config = config.get("_target_") == PROCESSOR_TARGET
+        if not is_full_processor_config:
+            default_config = load_processor_config()
+            default_config.image_transform = config
+            config = default_config
+
+        if overrides:
+            config = OmegaConf.merge(config, OmegaConf.from_dotlist(overrides))
+            assert isinstance(config, DictConfig)
 
-def get_processor(*, image_only: bool = False, **kwargs) -> Processor:
-    overrides = []
-    for key, value in kwargs.items():
-        overrides.append(f"{key}={value}")
-    config = load_processor_config(overrides=overrides)
     if image_only:
         config.tokenizer = None
     return instantiate(config)

tests/test_custom_config.py

@@ -0,0 +1,125 @@
+"""Tests for custom config support in get_model and get_processor."""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+import pytest
+import torchio as tio
+from omegaconf import DictConfig
+from omegaconf import OmegaConf
+
+from colipri.checkpoint import load_model_config
+from colipri.checkpoint import load_processor_config
+from colipri.processor import Processor
+from colipri.processor import get_processor
+
+TRANSFORM_YAML_CONTENT = """\
+_target_: torchio.transforms.augmentation.composition.Compose
+transforms:
+  - _target_: torchio.transforms.preprocessing.intensity.clamp.Clamp
+    out_min: -500
+    out_max: 500
+"""
+
+
+@pytest.fixture
+def transform_yaml(tmp_path: Path) -> Path:
+    """A self-contained transform YAML with no interpolation variables."""
+    path = tmp_path / "transform.yaml"
+    path.write_text(TRANSFORM_YAML_CONTENT)
+    return path
+
+
+@pytest.fixture
+def resolved_processor_config() -> DictConfig:
+    """Default processor config with all interpolations resolved."""
+    config = load_processor_config()
+    resolved = OmegaConf.to_container(config, resolve=True)
+    assert isinstance(resolved, dict)
+    return OmegaConf.create(resolved)
+
+
+@pytest.fixture
+def resolved_model_config() -> DictConfig:
+    """Default model config with all interpolations resolved."""
+    config = load_model_config()
+    resolved = OmegaConf.to_container(config, resolve=True)
+    assert isinstance(resolved, dict)
+    return OmegaConf.create(resolved)
+
+
+class TestGetProcessorCustomConfig:
+    def test_with_transform_yaml_path(self, transform_yaml: Path) -> None:
+        """Transform YAML path → Processor with custom transform."""
+        processor = get_processor(config=transform_yaml, image_only=True)
+        assert isinstance(processor, Processor)
+        transform = processor._image_transform
+        assert isinstance(transform, tio.Compose)
+        assert len(transform.transforms) == 1
+        clamp = transform.transforms[0]
+        assert clamp.out_min == -500
+        assert clamp.out_max == 500
+
+    def test_with_transform_dictconfig(self) -> None:
+        """Transform DictConfig object → Processor with custom transform."""
+        config = OmegaConf.create(TRANSFORM_YAML_CONTENT)
+        processor = get_processor(config=config, image_only=True)
+        assert isinstance(processor, Processor)
+        transform = processor._image_transform
+        assert isinstance(transform, tio.Compose)
+        assert transform.transforms[0].out_min == -500
+
+    def test_with_full_processor_config(
+        self,
+        resolved_processor_config: DictConfig,
+    ) -> None:
+        """Full processor DictConfig → Processor matching that config."""
+        # Remove all but the first transform to distinguish from default (5 transforms)
+        resolved_processor_config.image_transform.transforms = (
+            resolved_processor_config.image_transform.transforms[:1]
+        )
+        processor = get_processor(
+            config=resolved_processor_config,
+            image_only=True,
+        )
+        assert isinstance(processor, Processor)
+        assert isinstance(processor._image_transform, tio.Compose)
+        assert len(processor._image_transform.transforms) == 1
+
+    def test_transform_yaml_wraps_with_default_tokenizer(
+        self,
+        transform_yaml: Path,
+    ) -> None:
+        """Transform-only config is wrapped with default tokenizer config."""
+        # Without image_only, the transform YAML should be wrapped into a full
+        # processor config that includes the default tokenizer.
+        processor = get_processor(config=transform_yaml)
+        assert isinstance(processor, Processor)
+        # Should have both custom transform and default tokenizer
+        assert isinstance(processor._image_transform, tio.Compose)
+        assert processor._text_tokenizer is not None
+
+    def test_default_unchanged(self) -> None:
+        """get_processor() without config still works (backward compat)."""
+        processor = get_processor(image_only=True)
+        assert isinstance(processor, Processor)
+        assert isinstance(processor._image_transform, tio.Compose)
+
+
+class TestGetModelCustomConfig:
+    def test_with_config(self, resolved_model_config: DictConfig) -> None:
+        """Pass a model DictConfig → Model with that config."""
+        from colipri.model.multimodal import Model
+        from colipri.model.multimodal import get_model
+
+        model = get_model(pretrained=False, config=resolved_model_config)
+        assert isinstance(model, Model)
+
+    def test_default_unchanged(self) -> None:
+        """get_model(pretrained=False) without config still works."""
+        from colipri.model.multimodal import Model
+        from colipri.model.multimodal import get_model
+
+        model = get_model(pretrained=False)
+        assert isinstance(model, Model)