Define model architecture and weights

[Zain-Mahmoud]

Summary

This PR aims to import the model architecture from the original terrain-diffusion repo, and loads the weights for those models. The model architecture is defined under models/ and is connected to the inference interface to be used the model pipeline.

The inference models here are wrappers around the inference models from the original repo and share the same interface as the mock models we created earlier. The purpose of this interface is to abstract away the actual ML implementation from the final pipeline. These models will all be used in the model pipeline where the outputs will be connected to each other.

Validation

• scripts/quality-check.sh passes locally
• Appropriate /test commands were run and are passing (e.g. /test gpu)
• Unit tests were added / e2e tests were added where applicable
• Manual testing, if applicable (describe further):

Checklist

• The change is focused and easy to review.
• Tests were added or updated for behaviour changes.
• No secrets, generated files, or files that only make sense on my machine are included.
• Relevant documentation is updated.

Additional Comments

I ran the GPU tests in a Google colab notebook and they seemed to be working fine.

[coveralls]

Coverage Report for CI Build 38022721650

Coverage decreased (-26.1%) to 72.481%

Details

• Coverage decreased (-26.1%) from the base build.
• Patch coverage: 316 uncovered changes across 5 files (128 of 444 lines covered, 28.83%).
• 1 coverage regression across 1 file.

Uncovered Changes

File	Changed	Covered	%
src/terrain_diffusion/models/mp_layers.py	137	30	21.9%
src/terrain_diffusion/models/edm_unet.py	109	13	11.93%
src/terrain_diffusion/models/unet_block.py	71	8	11.27%
src/terrain_diffusion/inference.py	81	47	58.02%
tests/test_inference.py	44	28	63.64%
Total (6 files)	444	128	28.83%

Coverage Regressions

1 previously-covered line in 1 file lost coverage.

File	Lines Losing Coverage	Coverage
src/terrain_diffusion/inference.py	1	69.63%

---

Coverage Stats

Relevant Lines:	1181
Covered Lines:	856
Line Coverage:	72.48%
Coverage Strength:	0.72 hits per line

---

💛 - Coveralls

[Zain-Mahmoud]

/test all

[github-actions]

Tests passed

• Groups run: python, slow, web
• Commit tested: b505a44
• Full log: Actions run

Group	Result
python	91/91 passed
slow	0/0 passed
web	1/1 passed

[Zain-Mahmoud]

/test gpu