Merge pull request #4 from allenai/gabi/update_v1.1

Update pyproject.toml

Author: gabrieltseng

.github/workflows/publish.yml

@@ -24,11 +24,11 @@ jobs:
       - uses: actions/checkout@v5
         with:
           ref: ${{ github.event.inputs.tag || github.ref }}
-      
+
       - name: Set up Python
         uses: actions/setup-python@v5
         with:
-          python-version: "3.12"
+          python-version: ".python-version"
 
       - name: Verify version matches tag
         run: |

.github/workflows/tests.yml

@@ -10,7 +10,7 @@ on:
 
 jobs:
   test:
-    name: Test with torch-cpu
+    name: Tests
     runs-on: ubuntu-latest
 
     steps:
@@ -28,18 +28,8 @@ jobs:
 
     - name: Install dependencies
       run: |
-        # Install with CPU-only torch (suitable for CI)
-        # Use uv pip install to bypass uv.sources configuration that causes conflicts
-        uv venv
+        uv sync --python 3.12 --extra dev
         source .venv/bin/activate
-        # Install torch from CPU index
-        uv pip install --index-url https://download.pytorch.org/whl/cpu "torch>=2.8,<2.9" "torchvision>=0.23,<1"
-        # Install base dependencies
-        uv pip install "einops>=0.7.0" "huggingface_hub" "numpy>=1.26.4" "universal-pathlib>=0.2.5"
-        # Install dev dependencies
-        uv pip install "pytest>=7.0.0"
-        # Install the project itself
-        uv pip install -e .
 
     - name: Run tests (without weights)
       run: |

.gitignore

@@ -1 +1,5 @@
 __pycache__/
+**/*.egg-info/
+# we ignore the lockfile since
+# we have gpu and cpu installations.
+uv.lock

.python-version

@@ -0,0 +1 @@
+3.12

README.md

@@ -23,10 +23,7 @@ To install dependencies:
 ```
 git clone git@github.com:allenai/olmoearth_pretrain_minimal.git
 cd olmoearth_pretrain_minimal
-# Install with CPU-only PyTorch (works on all platforms including Linux)
-uv sync --locked --python 3.12 --extra torch-cpu
-# Or install with CUDA 12.8 PyTorch
-uv sync --locked --python 3.12 --extra torch-cu128
+uv sync
 ```
 uv installs everything into a venv, so to keep using python commands you can activate uv's venv: `source .venv/bin/activate`. Otherwise, swap to `uv run python`.
 
@@ -37,7 +34,7 @@ uv installs everything into a venv, so to keep using python commands you can act
 
 <img src="https://raw.githubusercontent.com/allenai/olmoearth_pretrain/main/assets/model.png" alt="Model Architecture Diagram" style="width: 800px; margin-left:'auto' margin-right:'auto' display:'block'"/>
 
-The OlmoEarth models are trained on three satellite modalities (Sentinel 2, Sentinel 1 and Landsat) and six derived maps (OpenStreetMap, WorldCover, USDA Cropland Data Layer, SRTM DEM, WRI Canopy Height Map, and WorldCereal). 
+The OlmoEarth models are trained on three satellite modalities (Sentinel 2, Sentinel 1 and Landsat) and six derived maps (OpenStreetMap, WorldCover, USDA Cropland Data Layer, SRTM DEM, WRI Canopy Height Map, and WorldCereal).
 
 **Note:** The model weights are released under the [OlmoEarth Artifact License](https://github.com/allenai/olmoearth_pretrain/blob/main/LICENSE)
 | Model Size | Weights | Encoder Params | Decoder Params |
@@ -64,7 +61,7 @@ from olmoearth_pretrain_minimal import ModelID, load_model_from_id
 # - ModelID.OLMOEARTH_V1_LARGE - 308M encoder params, 53M decoder params
 model = load_model_from_id(ModelID.OLMOEARTH_V1_BASE, load_weights=True)
 
-# Load with randomly initialized weights 
+# Load with randomly initialized weights
 model_with_weights = load_model_from_id(ModelID.OLMOEARTH_V1_NANO, load_weights=False)
 ```
 

pyproject.toml

@@ -13,18 +13,10 @@ dependencies = [
     "huggingface_hub",
     "numpy>=1.26.4",
     "universal-pathlib>=0.2.5",
+    "torch>=2.8,<2.9",
 ]
 
 [project.optional-dependencies]
-torch-cpu = [
-    "torch>=2.8,<2.9",
-    "torchvision>=0.23,<1",
-]
-torch-cu128 = [
-    "pytorch-triton",
-    "torch>=2.8,<2.9",
-    "torchvision>=0.23,<1",
-]
 dev = [
     "pytest>=7.0.0",
 ]
@@ -39,32 +31,8 @@ include = ["olmoearth_pretrain_minimal*"]
 [tool.setuptools.package-data]
 "olmoearth_pretrain_minimal.olmoearth_pretrain_v1.data.norm_configs" = ["*.json"]
 
-[tool.uv.sources]
-torch = [
-    { index = "pytorch-cpu", extra = "torch-cpu" },
-    { index = "pytorch-cu128", extra = "torch-cu128" },
-]
-torchvision = [
-    { index = "pytorch-cpu", extra = "torch-cpu" },
-    { index = "pytorch-cu128", extra = "torch-cu128" },
-]
-pytorch-triton = [
-    { index = "pytorch-cu128", extra = "torch-cu128" },
-]
-
-[[tool.uv.index]]
-name = "pytorch-cpu"
-url = "https://download.pytorch.org/whl/cpu"
-explicit = true
-
-[[tool.uv.index]]
-name = "pytorch-cu128"
-url = "https://download.pytorch.org/whl/cu128"
-explicit = true
-
 [tool.pytest.ini_options]
 markers = [
     "slow: marks tests as slow (deselect with '-m \"not slow\"')",
 ]
 testpaths = ["tests"]
-