Equimo 0.2.1__tar.gz → 0.2.2__tar.gz

{equimo-0.2.1 → equimo-0.2.2}/Equimo.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: Equimo
-Version: 0.2.1
+Version: 0.2.2
 Summary: Implementation of popular vision models in Jax
 Classifier: Development Status :: 4 - Beta
 Classifier: Intended Audience :: Developers

{equimo-0.2.1 → equimo-0.2.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: Equimo
-Version: 0.2.1
+Version: 0.2.2
 Summary: Implementation of popular vision models in Jax
 Classifier: Development Status :: 4 - Beta
 Classifier: Intended Audience :: Developers

{equimo-0.2.1 → equimo-0.2.2}/equimo/io.py

@@ -102,6 +102,7 @@ def download(identifier: str, repository: str) -> Path:
 
     url = f"{repository}/{identifier}.tar.lz4"
     path = Path(f"~/.cache/equimo/{identifier}.tar.lz4").expanduser()
+    path.parent.mkdir(parents=True, exist_ok=True)
 
     if path.exists():
         logger.info("Archive already downloaded, using cached file.")

{equimo-0.2.1 → equimo-0.2.2}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "Equimo"
-version = "0.2.1"
+version = "0.2.2"
 description = "Implementation of popular vision models in Jax"
 readme = "README.md"
 requires-python = ">=3.11"

{equimo-0.2.1 → equimo-0.2.2}/tests/test_models.py

@@ -120,3 +120,16 @@ def test_save_load_model_uncompressed():
         loaded_output = loaded_model.features(x, key=key)
 
         assert jnp.allclose(original_output, loaded_output, atol=1e-5)
+
+
+def test_load_pretrained_model():
+    """Test loading a pretrained model from the repository."""
+    key = jr.PRNGKey(42)
+    model = load_model(cls="vit", identifier="dinov2_vits14_reg")
+
+    # Test inference
+    x = jr.normal(key, (3, 224, 224))
+    features = model.features(x, key=key)
+
+    assert features.shape[-1] == 384  # DINOv2-S has embedding dimension of 384
+    assert jnp.all(jnp.isfinite(features))  # Check for NaN/Inf values