Script to convert OlmoEarth v1 model variants from .pth to .pt2 format

[weiji14]

What I am changing

• Portable .pt2 format for storing OlmoEarth model weights, including the model definition

How I did it

TODO:

• Create a PEP 723 compatible standalone script to do the .pth to .pt2 conversion

  uv init --script scripts/export_olmoearth.py
  uv add --script scripts/export_olmoearth.py 'olmoearth-pretrain-minimal' 'torch'
  uv lock --script scripts/export_olmoearth.py

• Upload *.pt2 weights to https://huggingface.co/developmentseed?

References:

• https://docs.astral.sh/uv/guides/scripts/#running-scripts
• https://docs.pytorch.org/docs/2.11/user_guide/torch_compiler/export/pt2_archive.html

How you can test it

• Run uv run --no-project scripts/export_olmoearth.py --variant tiny script locally

Related Issues

Address #32

[martham93]

This looks good thank you!

I think as a follow up, we should consider expanding and generating .pt2 files to support other modality combinations we think we would want to use? (I can ticket + help work on if we think this would be useful)

These weights will only allow us to run inference with S2 data which is a reasonable starting point but if we wanted to gather other data as well it would fail.

[martham93]

cool! I didn't know how to separately manage script environments outside of the main project dependencies!

Thanks for including the uv running scripts documentation in the PR description, super helpful

[weiji14]

Yeah, I was aware of PEP 723 scripts for a while, but didn't realize uv was able to generate lockfiles for scripts too! It's also pretty cool that you could add more dependencies using uv add --script *.py 'some-dep' and the lockfile updates accordingly, essentially a mini pyproject.toml+uv.lock kind of situation.

[weiji14]

Btw, I've only tested converting the nano and tiny weights, not the base or large ones yet. Edit: Have tested converting base and large now, seems to work (though haven't ran inference on them).

I think as a follow up, we should consider expanding and generating .pt2 files to support other modality combinations we think we would want to use? (I can ticket + help work on if we think this would be useful)

Sure, that would be good! Should I maybe rename the weights from olmoearth_v1_tiny.pt2 to olmoearth_v1_tiny_s2.pt2 or some other naming scheme? Maybe you could open the ticket and we can discuss there before we upload these .pt2 weights to Huggingface.

[martham93]

+1 to that naming scheme, but I'll get a ticket up soon so we can discuss there