Basic model inference

pyproject.toml

@@ -8,7 +8,9 @@ requires-python = ">=3.14"
 # Packages the project itself needs in order to run
 # TODO: add entries here as the project starts needing them,
 # then run `uv sync` to install and update uv.lock
-dependencies = []
+dependencies = [
+    "numpy>=2.5.2",
+]
 
 # The commands the project installs. `uv run terrain-diffusion` runs main() in
 # src/terrain_diffusion/cli.py.

src/terrain_diffusion/inference.py

@@ -16,3 +16,138 @@
 - It loads weights from the external model weights download.
 - It runs on the GPU compute node.
 """
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import ClassVar
+
+import numpy as np
+
+PATCH_SIZE = (512, 512)
+LATENT_MAP_SIZE = (3, 50, 100)  # placeholder
+
+
+@dataclass
+class ModelOutput(ABC):
+    @abstractmethod
+    def __init__(self):
+        raise NotImplementedError
+
+
+@dataclass
+class ModelInput(ABC):
+    @abstractmethod
+    def __init__(self):
+        raise NotImplementedError
+
+
+class TerrainModel[InputT: ModelInput, OutputT: ModelOutput](ABC):
+    @abstractmethod
+    def predict(self, patch: InputT) -> OutputT:
+        raise NotImplementedError
+
+    @abstractmethod
+    def load_weights(self, model_path: str):
+        """
+        Load a model stored in model_path
+        """
+        raise NotImplementedError
+
+
+@dataclass
+class MockCoreModelInput(ModelInput):
+    patch: np.ndarray
+    patch_shape: ClassVar[tuple] = PATCH_SIZE
+
+    def __post_init__(self):
+        assert self.patch.shape == self.patch_shape, "invalid input patch shape"
+
+    def __eq__(self, other: MockCoreModelInput):
+        return np.array_equal(self.patch, other.patch)
+
+
+@dataclass
+class MockCoreModelOutput(ModelOutput):
+    low_res_grid: np.ndarray
+    latent_map: np.ndarray
+    low_res_grid_shape: ClassVar[tuple] = (PATCH_SIZE[0] // 8, PATCH_SIZE[1] // 8)
+    latent_map_shape: ClassVar[tuple] = LATENT_MAP_SIZE
+
+    def __post_init__(self):
+        assert self.latent_map.shape == self.latent_map_shape, "invalid latent map size"
+        assert self.low_res_grid.shape == self.low_res_grid_shape, (
+            "invalid low resolution grid shape"
+        )
+
+    def __eq__(self, other: MockCoreModelInput):
+        return np.array_equal(self.low_res_grid, other.low_res_grid) and np.array_equal(
+            self.latent_map, other.latent_map
+        )
+
+
+@dataclass
+class MockDecoderModelInput(ModelInput):
+    latent_map: np.ndarray
+    latent_map_shape: ClassVar[tuple] = LATENT_MAP_SIZE
+
+    def __post_init__(self):
+        assert self.latent_map.shape == self.latent_map_shape, "invalid latent map size"
+
+    def __eq__(self, other: MockDecoderModelInput):
+        return np.array_equal(self.latent_map, other.latent_map)
+
+
+@dataclass
+class MockDecoderModelOutput(ModelOutput):
+    full_res_grid: np.ndarray
+    full_res_grid_shape: ClassVar[tuple] = PATCH_SIZE
+
+    def __post_init__(self):
+        assert self.full_res_grid.shape == self.full_res_grid_shape, (
+            "invalid full resolution grid size"
+        )
+
+    def __eq__(self, other: MockCoreModelOutput):
+        return np.array_equal(self.full_res_grid, other.full_res_grid)
+
+
+class MockCoreModel(TerrainModel[MockCoreModelInput, MockCoreModelOutput]):
+    weights: np.ndarray
+
+    def predict(self, input: MockCoreModelInput) -> MockCoreModelOutput:
+
+        double = input.patch * 2
+        low_res_grid = np.resize(double, (PATCH_SIZE[0] // 8, PATCH_SIZE[1] // 8))
+        latent_map = np.resize(double, LATENT_MAP_SIZE)
+        output = MockCoreModelOutput(low_res_grid, latent_map)
+
+        return output
+
+    def load_weights(self, model_path: str):
+        self.weights = np.ones((3, 4, 5))
+
+
+class MockDecoderModel(TerrainModel[MockDecoderModelInput, MockDecoderModelOutput]):
+    weights: np.ndarray
+
+    def predict(self, input: MockDecoderModelInput) -> MockDecoderModelOutput:
+
+        double = input.latent_map * 2
+        full_res_grid = np.resize(double, PATCH_SIZE)
+        output = MockDecoderModelOutput(full_res_grid)
+
+        return output
+
+    def load_weights(self, model_path: str):
+        self.weights = np.ones((3, 4, 5))
+
+
+MODELS = {"decoder": MockDecoderModel, "core": MockCoreModel}
+
+
+def load_model(model_name: str) -> TerrainModel:
+    if model_name not in MODELS:
+        raise ValueError("invalid model name")
+    return MODELS[model_name]()

tests/test_inference.py

@@ -0,0 +1,91 @@
+"""
+Tests inference interface by instantiating mock models
+and implementing the `generate` function
+"""
+
+import numpy as np
+import pytest
+
+from terrain_diffusion.inference import (
+    LATENT_MAP_SIZE,
+    PATCH_SIZE,
+    MockCoreModel,
+    MockCoreModelInput,
+    MockCoreModelOutput,
+    MockDecoderModel,
+    MockDecoderModelInput,
+    MockDecoderModelOutput,
+    load_model,
+)
+
+
+class TestTerrainModel:
+    @pytest.fixture
+    def initial_decoder(self) -> MockDecoderModel:
+        return MockDecoderModel()
+
+    @pytest.fixture
+    def initial_core(self) -> MockCoreModel:
+        return MockCoreModel()
+
+    def test_core_input_generation(self):
+        with pytest.raises(AssertionError):
+            MockCoreModelInput(np.ones((1, 1)))
+
+    def test_core_output_generation(self):
+        with pytest.raises(AssertionError):
+            MockCoreModelOutput(np.ones((1, 1)), np.ones(LATENT_MAP_SIZE))
+        with pytest.raises(AssertionError):
+            MockCoreModelOutput(np.ones((PATCH_SIZE[0] // 8, PATCH_SIZE[1] // 8)), np.ones((1, 1)))
+
+    def test_decoder_input_generation(self):
+        with pytest.raises(AssertionError):
+            MockDecoderModelInput(np.ones((1, 1)))
+
+    def test_decoder_output_generation(self):
+        with pytest.raises(AssertionError):
+            MockDecoderModelOutput(np.ones((1, 1)))
+
+    def test_load_model(self):
+        assert isinstance(load_model("decoder"), MockDecoderModel), (
+            "model does not load correct decoder"
+        )
+        assert isinstance(load_model("core"), MockCoreModel), "model does not load correct core"
+
+    def test_predict_core(self, initial_core):
+        input = MockCoreModelInput(np.ones(PATCH_SIZE))
+
+        actual = initial_core.predict(input)
+        actual_2 = initial_core.predict(input)
+
+        expected_low_res = np.ndarray((PATCH_SIZE[0] // 8, PATCH_SIZE[1] // 8))
+        expected_latent = np.ndarray(LATENT_MAP_SIZE)
+
+        expected_low_res.fill(2)
+        expected_latent.fill(2)
+
+        expected = MockCoreModelOutput(expected_low_res, expected_latent)
+
+        assert actual == expected, "predictions are not equal"
+        assert actual == actual_2, "model return different predictions on same input"
+        assert actual.low_res_grid.shape == (PATCH_SIZE[0] // 8, PATCH_SIZE[1] // 8), (
+            "low resolution map shape is not patch size // 8"
+        )
+        assert actual.latent_map.shape == LATENT_MAP_SIZE, "latent map size is not correct"
+
+    def test_predict_decoder(self, initial_decoder):
+        input = MockDecoderModelInput(np.ones(LATENT_MAP_SIZE))
+
+        actual = initial_decoder.predict(input)
+        actual_2 = initial_decoder.predict(input)
+
+        expected_full_res = np.ndarray(PATCH_SIZE)
+        expected_full_res.fill(2)
+
+        expected = MockDecoderModelOutput(expected_full_res)
+
+        assert actual == expected, "predictions are not equal"
+        assert actual == actual_2, "model return different predictions on same input"
+        assert actual.full_res_grid.shape == PATCH_SIZE, (
+            "full resolution map does not match patch size"
+        )