Exception in minimal inference-only environment

[jleben]

I am trying to use OlmoEarth for inference only in a minimal environment - without ai2-olmo-core installed.

Following the example here, the following code raises an exception:

import torch
from olmoearth_pretrain.model_loader import ModelID, load_model_from_id
from olmoearth_pretrain.datatypes import MaskedOlmoEarthSample, MaskValue

# Load model from HuggingFace
model = load_model_from_id(ModelID.OLMOEARTH_V1_BASE)

File c:\Users\jakob\Projects\ml\olmoearth\olmoearth_pretrain\model_loader.py:116, in _load_model_from_config(path)
    114     config_dict = json.load(f)
    115     model_config = Config.from_dict(config_dict["model"])
--> [116](file:///C:/Users/jakob/Projects/ml/olmoearth/olmoearth_pretrain/model_loader.py:116) return model_config.build()

File c:\Users\jakob\Projects\ml\olmoearth\olmoearth_pretrain\config.py:244, in _StandaloneConfig.build(self)
...
   (...)    242         NotImplementedError: Always, unless overridden by subclass.
    243     """
--> [244](file:///C:/Users/jakob/Projects/ml/olmoearth/olmoearth_pretrain/config.py:244)     raise NotImplementedError("Subclasses must implement build()")

NotImplementedError: Subclasses must implement build()

[Hgherzog]

Hi! Thanks for trying out our model! I tried to reproduce your bug but wan unable to. How did you install the package? If you give a little more information on your setup, I can try to help further.

I followed the instructions in a freshly cloned repo (off of both main and the latest pypi release):

downloading uv 0.9.28 aarch64-apple-darwin
no checksums to verify
installing to /Users/henryh/.local/bin
  uv
  uvx
everything's installed!
➜  olmoearth_pretrain git:(main) uv sync
Using CPython 3.12.12
Creating virtual environment at: .venv
warning: The `extra-build-dependencies` option is experimental and may change without warning. Pass `--preview-features extra-build-dependencies` to disable this warning.
Resolved 269 packages in 6ms
      Built olmoearth-pretrain @ file:///Users/henryh/Desktop/eai-repos/olmoearth_pretrain
Prepared 1 package in 572ms
Installed 49 packages in 823ms
 + attrs==25.4.0
 + certifi==2025.10.5
 + cfgv==3.4.0
 + charset-normalizer==3.4.4
 + click==8.3.0
 + colorama==0.4.6
 + distlib==0.4.0
 + einops==0.7.0
 + execnet==2.1.1
 + filelock==3.20.0
 + fsspec==2025.9.0
 + gitdb==4.0.12
 + gitpython==3.1.45
 + hf-xet==1.1.10
 + huggingface-hub==0.35.3
 + identify==2.6.15
 + idna==3.11
 + iniconfig==2.3.0
 + interrogate==1.7.0
 + jinja2==3.1.6
 + markupsafe==3.0.3
 + mpmath==1.3.0
 + networkx==3.5
 + nodeenv==1.9.1
 + numpy==1.26.4
 + olmoearth-pretrain==0.0.2 (from file:///Users/henryh/Desktop/eai-repos/olmoearth_pretrain)
 + packaging==24.2
 + pathlib-abc==0.5.2
 + pip==25.2
 + platformdirs==4.5.0
 + pluggy==1.6.0
 + pre-commit==4.3.0
 + py==1.11.0
 + pygments==2.19.2
 + pytest==8.4.2
 + pytest-xdist==3.8.0
 + pyyaml==6.0.3
 + requests==2.32.5
 + ruff==0.14.1
 + setuptools==80.9.0
 + smmap==5.0.2
 + sympy==1.14.0
 + tabulate==0.9.0
 + torch==2.7.1
 + tqdm==4.67.1
 + typing-extensions==4.15.0
 + universal-pathlib==0.3.4
 + urllib3==2.5.0
 + virtualenv==20.35.3
➜  olmoearth_pretrain git:(main) uv run --with olmoearth-pretrain --with torch --python 3.12 - <<'EOF'
import torch
from olmoearth_pretrain.model_loader import ModelID, load_model_from_id
from olmoearth_pretrain.datatypes import MaskedOlmoEarthSample, MaskValue

# Load model from HuggingFace
model = load_model_from_id(ModelID.OLMOEARTH_V1_BASE)
model.eval()

# Create synthetic input (B=1, H=64, W=64, T=1, C=12 for Sentinel-2)
dummy_image = torch.randn(1, 64, 64, 1, 12)
dummy_mask = torch.ones(1, 64, 64, 1, 3) * MaskValue.ONLINE_ENCODER.value
dummy_timestamps = torch.tensor([[[15, 6, 2024]]])  # day, month (0-indexed), year

sample = MaskedOlmoEarthSample(
    sentinel2_l2a=dummy_image,
    sentinel2_l2a_mask=dummy_mask,
    timestamps=dummy_timestamps,
)

with torch.no_grad():
    output = model.encoder(sample, fast_pass=True, patch_size=4)
    features = output["tokens_and_masks"].sentinel2_l2a
    print(f"Output shape: {features.shape}")  # (B, H', W', T, S, D)
EOF
warning: The `extra-build-dependencies` option is experimental and may change without warning. Pass `--preview-features extra-build-dependencies` to disable this warning.
/Users/henryh/Desktop/eai-repos/olmoearth_pretrain/olmoearth_pretrain/model_loader.py:37: UserWarning: olmo-core not installed. Running in inference-only mode. For training: pip install olmoearth-pretrain[training]
  from olmoearth_pretrain.config import Config
/Users/henryh/.local/share/uv/python/cpython-3.12.12-macos-aarch64-none/lib/python3.12/importlib/__init__.py:90: FutureWarning: The 'helios' package has been renamed to 'olmoearth_pretrain'. Please update your imports; this compatibility shim will be removed in a future release.
  return _bootstrap._gcd_import(name[level:], package, level)
/Users/henryh/Desktop/eai-repos/olmoearth_pretrain/.venv/lib/python3.12/site-packages/einops/einops.py:847: SyntaxWarning: invalid escape sequence '\s'
  \sum_{c, d, g} x[a, b, c] * y[c, b, d] * z[a, g, k]
Output shape: torch.Size([1, 16, 16, 1, 3, 768])```

[Hgherzog]

I also tried to reproduce it with a slightly different setup and was unable to reproduce

Setup:
uv venv .venv-minimal && source .venv-minimal/bin/activate && uv pip install -e.
File to run: python test_hf_load.py

import torch

from olmoearth_pretrain.datatypes import MaskedOlmoEarthSample, MaskValue
from olmoearth_pretrain.model_loader import ModelID, load_model_from_id


def create_sample_input(batch_size: int = 1) -> MaskedOlmoEarthSample:
    """Create a minimal sample input for testing."""
    # Dimensions: B=batch, H=height, W=width, T=timesteps, C=channels
    H, W, T = 16, 16, 4
    S2_CHANNELS = 12  # Sentinel-2 L2A has 12 channels (4 + 6 + 2 bands)

    # Sentinel-2 data: [B, H, W, T, C]
    sentinel2_l2a = torch.randn(batch_size, H, W, T, S2_CHANNELS)
    sentinel2_l2a_mask = torch.full(
        (batch_size, H, W, T, S2_CHANNELS),
        fill_value=MaskValue.ONLINE_ENCODER.value,
        dtype=torch.long,
    )

    # Lat/lon: [B, 2]
    latlon = torch.tensor([[37.7749, -122.4194]] * batch_size)  # San Francisco
    latlon_mask = torch.full(
        (batch_size, 2),
        fill_value=MaskValue.ONLINE_ENCODER.value,
        dtype=torch.long,
    )

    # Timestamps: [B, T, 3] where each row is [day, month, year]
    timestamps = torch.tensor(
        [[[15, 0, 2023], [15, 3, 2023], [15, 6, 2023], [15, 9, 2023]]] * batch_size,
        dtype=torch.long,
    )

    return MaskedOlmoEarthSample(
        timestamps=timestamps,
        sentinel2_l2a=sentinel2_l2a,
        sentinel2_l2a_mask=sentinel2_l2a_mask,
        latlon=latlon,
        latlon_mask=latlon_mask,
    )


def main() -> None:
    """Load model and run minimal inference test."""
    print("Loading OlmoEarth-v1-Base from HuggingFace...")
    model = load_model_from_id(ModelID.OLMOEARTH_V1_BASE)
    model.eval()

    print(f"Model loaded: {type(model).__name__}")
    print(f"Model parameters: {sum(p.numel() for p in model.parameters()):,}")

    print("\nCreating sample input...")
    sample = create_sample_input(batch_size=1)
    print(f"Sentinel-2 shape: {sample.sentinel2_l2a.shape}")
    print(f"Timestamps shape: {sample.timestamps.shape}")

    print("\nRunning forward pass...")
    with torch.no_grad():
        output = model(sample, patch_size=4)

    print("\nForward pass successful!")
    print(f"Output type: {type(output)}")
    if isinstance(output, dict):
        for key, value in output.items():
            if hasattr(value, "shape"):
                print(f"  {key}: {value.shape}")
            else:
                print(f"  {key}: {type(value)}")


if __name__ == "__main__":
    main()

[jleben]

@Hgherzog Thanks for trying to reproduce this on your end.

I am trying to integrate OlmoEarth into my own project, so I am not using its pyproject.toml file. Instead, I am manually installing a minimal set of dependencies that I think it needs (based on studying its pyproject.toml).

I solved the issue though.

I pinpointed it to cls._resolve_class(class_name) returning None here.

Further investigation revealed that import helios.nn.latent_mim raised an exception because I was missing the einops package. That in turn caused the config class to not be resolved successfullly. After installing einops, I am able to load the model with load_model_from_id(ModelID.OLMOEARTH_V1_NANO).

I would have got to this point quicker and without so much effort in debugging, had that exception not been suppressed in _resolve_class - here. Any particular reason to suppress those exceptions there?

[Hgherzog]

@jleben Glad you were able to resolve this! Is there anything making it difficult to install from pypi or use the pyproject.toml directly?

Also, I made a PR to update the error handling to not hide the error.

[jleben]

@Hgherzog Thanks for that PR.

Is there anything making it difficult to install from pypi or use the pyproject.toml directly?

Mainly package requirement incompatibilities.

I tried integrating the whole olmo_pretrain repo into my project as a git submodule and using olmo's pyproject.toml as a workspace member in my pyproject.toml. In that case, there is no way around uv considering all the packages mentioned in all of olmo's dependency groups for version resolution - even those groups that are not installed by uv sync. That constrains the versions a lot, including, critically, requiring numpy <2 whereas I require numpy 2.

I also tried installing the olmoearth-pretrain package. In that case, the package version constraints are more relaxed. However, olmoearth still constrains torch to a single minor version of 2.7 which causes some difficulty for me.

By using olmoearth source code as a submodule and ignoring its pyproject.toml, I was able to run inference using dummy input data, and using torch version 2.10. I also have another project where I'd like to use olmoearth that is constrained to pytorch 2.5 at the moment. I gues I am exposing myself to some risk this way, since I'm not sure olmoearth was tested with this torch version.

Could you comment on the olmoearth's compatibility with different torch versions?