Equimo 0.2.2__tar.gz → 0.2.3__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: Equimo
-Version: 0.2.2
+Version: 0.2.3
 Summary: Implementation of popular vision models in Jax
 Classifier: Development Status :: 4 - Beta
 Classifier: Intended Audience :: Developers
@@ -20,8 +20,10 @@ Requires-Dist: jaxlib>=0.4.25
 Requires-Dist: loguru>=0.7.3
 Requires-Dist: lz4>=4.4.3
 Requires-Dist: requests>=2.32.3
+Requires-Dist: semver>=3.0.4
 Provides-Extra: torch
 Requires-Dist: torch>=2.6.0; extra == "torch"
+Requires-Dist: timm>=1.0.15; extra == "torch"
 
 # Equimo: Modern Vision Models in JAX/Equinox
 
@@ -54,6 +56,8 @@ pip install -e .
 
 ## Implemented Models
 
+Beyond normal ViT (e.g., dinov2 or siglip), equimo proposes other SotA architectures:
+
 | Model         | Paper                                                                                                                                                           | Year | Status    |
 | ------------- | --------------------------------------------------------------------------------------------------------------------------------------------------------------- | ---- | --------- |
 | FasterViT     | [FasterViT: Fast Vision Transformers with Hierarchical Attention](https://arxiv.org/abs/2306.06189)                                                             | 2023 | ✅        |
@@ -139,19 +143,28 @@ model = load_model(cls="vit", path=Path("path/to/model.tar.lz4"))
 model = load_model(cls="vit", path=Path("path/to/model/"))
 ```
 
+Parameters passed to models can be overridden such as:
+
+```python
+model = load_model(
+    cls="vit",
+    identifier="siglip2_vitb16_256",
+    dynamic_img_size=True,  # passed to the VisionTransformer class
+)
+```
+
 #### List of pretrained models
 
-Currently, only [DinoV2](https://github.com/facebookresearch/dinov2/) have been ported, therefore the following
-identifiers are available:
+Currently, only [DinoV2](https://github.com/facebookresearch/dinov2/) and SigLIP2 have been ported.
+
+Model identifiers allows downloading from equimo's [repository on huggingface](https://huggingface.co/poiretclement/equimo/tree/main/models/default)
 
-- `dinov2_vitb14.tar.lz4`
-- `dinov2_vitb14_reg.tar.lz4`
-- `dinov2_vitl14.tar.lz4`
-- `dinov2_vitl14_reg.tar.lz4`
-- `dinov2_vits14.tar.lz4`
-- `dinov2_vits14_reg.tar.lz4`
+Identifiers are filenames without the extensions, such as:
 
-(`giant` version coming soon)
+- `dinov2_vitb14`
+- `dinov2_vits14_reg`
+- `siglip2_vitl16_512`
+- `siglip2_vitso400m16_384`
 
 ## Contributing
 

{equimo-0.2.2 → equimo-0.2.3}/Equimo.egg-info/SOURCES.txt

@@ -11,8 +11,10 @@ equimo/io.py
 equimo/utils.py
 equimo/converters/__init__.py
 equimo/converters/dinov2.py
+equimo/converters/siglip2.py
 equimo/converters/utils.py
 equimo/layers/__init__.py
+equimo/layers/activation.py
 equimo/layers/attention.py
 equimo/layers/convolution.py
 equimo/layers/downsample.py

{equimo-0.2.2 → equimo-0.2.3}/Equimo.egg-info/requires.txt

@@ -5,6 +5,8 @@ jaxlib>=0.4.25
 loguru>=0.7.3
 lz4>=4.4.3
 requests>=2.32.3
+semver>=3.0.4
 
 [torch]
 torch>=2.6.0
+timm>=1.0.15

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: Equimo
-Version: 0.2.2
+Version: 0.2.3
 Summary: Implementation of popular vision models in Jax
 Classifier: Development Status :: 4 - Beta
 Classifier: Intended Audience :: Developers
@@ -20,8 +20,10 @@ Requires-Dist: jaxlib>=0.4.25
 Requires-Dist: loguru>=0.7.3
 Requires-Dist: lz4>=4.4.3
 Requires-Dist: requests>=2.32.3
+Requires-Dist: semver>=3.0.4
 Provides-Extra: torch
 Requires-Dist: torch>=2.6.0; extra == "torch"
+Requires-Dist: timm>=1.0.15; extra == "torch"
 
 # Equimo: Modern Vision Models in JAX/Equinox
 
@@ -54,6 +56,8 @@ pip install -e .
 
 ## Implemented Models
 
+Beyond normal ViT (e.g., dinov2 or siglip), equimo proposes other SotA architectures:
+
 | Model         | Paper                                                                                                                                                           | Year | Status    |
 | ------------- | --------------------------------------------------------------------------------------------------------------------------------------------------------------- | ---- | --------- |
 | FasterViT     | [FasterViT: Fast Vision Transformers with Hierarchical Attention](https://arxiv.org/abs/2306.06189)                                                             | 2023 | ✅        |
@@ -139,19 +143,28 @@ model = load_model(cls="vit", path=Path("path/to/model.tar.lz4"))
 model = load_model(cls="vit", path=Path("path/to/model/"))
 ```
 
+Parameters passed to models can be overridden such as:
+
+```python
+model = load_model(
+    cls="vit",
+    identifier="siglip2_vitb16_256",
+    dynamic_img_size=True,  # passed to the VisionTransformer class
+)
+```
+
 #### List of pretrained models
 
-Currently, only [DinoV2](https://github.com/facebookresearch/dinov2/) have been ported, therefore the following
-identifiers are available:
+Currently, only [DinoV2](https://github.com/facebookresearch/dinov2/) and SigLIP2 have been ported.
+
+Model identifiers allows downloading from equimo's [repository on huggingface](https://huggingface.co/poiretclement/equimo/tree/main/models/default)
 
-- `dinov2_vitb14.tar.lz4`
-- `dinov2_vitb14_reg.tar.lz4`
-- `dinov2_vitl14.tar.lz4`
-- `dinov2_vitl14_reg.tar.lz4`
-- `dinov2_vits14.tar.lz4`
-- `dinov2_vits14_reg.tar.lz4`
+Identifiers are filenames without the extensions, such as:
 
-(`giant` version coming soon)
+- `dinov2_vitb14`
+- `dinov2_vits14_reg`
+- `siglip2_vitl16_512`
+- `siglip2_vitso400m16_384`
 
 ## Contributing
 

{equimo-0.2.2 → equimo-0.2.3}/README.md

@@ -29,6 +29,8 @@ pip install -e .
 
 ## Implemented Models
 
+Beyond normal ViT (e.g., dinov2 or siglip), equimo proposes other SotA architectures:
+
 | Model         | Paper                                                                                                                                                           | Year | Status    |
 | ------------- | --------------------------------------------------------------------------------------------------------------------------------------------------------------- | ---- | --------- |
 | FasterViT     | [FasterViT: Fast Vision Transformers with Hierarchical Attention](https://arxiv.org/abs/2306.06189)                                                             | 2023 | ✅        |
@@ -114,19 +116,28 @@ model = load_model(cls="vit", path=Path("path/to/model.tar.lz4"))
 model = load_model(cls="vit", path=Path("path/to/model/"))
 ```
 
+Parameters passed to models can be overridden such as:
+
+```python
+model = load_model(
+    cls="vit",
+    identifier="siglip2_vitb16_256",
+    dynamic_img_size=True,  # passed to the VisionTransformer class
+)
+```
+
 #### List of pretrained models
 
-Currently, only [DinoV2](https://github.com/facebookresearch/dinov2/) have been ported, therefore the following
-identifiers are available:
+Currently, only [DinoV2](https://github.com/facebookresearch/dinov2/) and SigLIP2 have been ported.
+
+Model identifiers allows downloading from equimo's [repository on huggingface](https://huggingface.co/poiretclement/equimo/tree/main/models/default)
 
-- `dinov2_vitb14.tar.lz4`
-- `dinov2_vitb14_reg.tar.lz4`
-- `dinov2_vitl14.tar.lz4`
-- `dinov2_vitl14_reg.tar.lz4`
-- `dinov2_vits14.tar.lz4`
-- `dinov2_vits14_reg.tar.lz4`
+Identifiers are filenames without the extensions, such as:
 
-(`giant` version coming soon)
+- `dinov2_vitb14`
+- `dinov2_vits14_reg`
+- `siglip2_vitl16_512`
+- `siglip2_vitso400m16_384`
 
 ## Contributing
 

equimo-0.2.3/equimo/__init__.py

@@ -0,0 +1 @@
+__version__ = "0.2.3"

equimo-0.2.3/equimo/converters/siglip2.py

@@ -0,0 +1,105 @@
+from pathlib import Path
+
+import jax
+import jax.numpy as jnp
+import numpy as np
+
+import equimo.models as em
+from equimo.converters.utils import convert_torch_to_equinox
+from equimo.io import load_model, save_model
+
+
+def compare(j, t) -> float:
+    j = np.array(j)
+    t = t.squeeze().detach().numpy()
+    return float(np.mean(np.abs(j - t)))
+
+
+def main():
+    try:
+        import torch
+    except:
+        raise ImportError("`torch` not available")
+
+    key = jax.random.PRNGKey(42)
+    siglip2_config = {
+        "img_size": 384,
+        "in_channels": 3,
+        "dim": 1536,
+        "patch_size": 16,
+        "num_heads": [16],
+        "depths": [40],
+        "mlp_ratio": 4,
+        "num_classes": 0,
+        "use_mask_token": False,
+        "reg_tokens": 0,
+        "class_token": False,
+        "no_embed_class": True,
+        "init_values": None,
+        "eps": 1e-6,
+        "dynamic_img_size": False,
+        # "act_layer": "exactgelu",
+        "act_layer": "gelu",
+    }
+
+    siglip2 = em.VisionTransformer(
+        **siglip2_config,
+        key=key,
+    )
+
+    timm_cfg = ["vit_giantopt_patch16_siglip_gap_384.v2_webli", True]
+    _n = timm_cfg[0].split("_")
+    mname = f"siglip2_vitgiantopt{_n[2][-2:]}_{_n[-2][:3]}"
+    # mname = f"siglip2_vit{_n[1][0]}{_n[2][-2:]}_{_n[-2][:3]}"
+
+    print(f"Converting {mname}...")
+
+    replace_cfg = {
+        "blocks.0.blocks": "blocks",
+        ".prenorm.": ".norm1.",
+        ".norm.": ".norm2.",
+    }
+    expand_cfg = {"patch_embed.proj.bias": ["after", 2]}
+    squeeze_cfg = {
+        "pos_embed": 0,
+        "cls_token": 0,
+        "register_tokens": 0,
+    }
+    whitelist = []
+
+    siglip2, torch_model = convert_torch_to_equinox(
+        siglip2,
+        replace_cfg,
+        expand_cfg,
+        squeeze_cfg,
+        whitelist,
+        strict=True,
+        source="timm",
+        timm_cfg=timm_cfg,
+        return_torch=True,
+    )
+
+    arr = np.random.randn(3, siglip2_config["img_size"], siglip2_config["img_size"])
+    jax_arr = jnp.array(arr)
+    torch_arr = torch.tensor(arr).unsqueeze(0).float()
+
+    assert (
+        compare(
+            jax.vmap(siglip2.norm)(siglip2.features(jax_arr, key)),
+            torch_model.forward_features(torch_arr),
+        )
+        < 1e-5
+    )
+
+    save_model(
+        Path(f"~/.cache/equimo/{mname}").expanduser(),
+        siglip2,
+        siglip2_config,
+        timm_cfg=timm_cfg,
+        compression=True,
+    )
+
+    _ = load_model(
+        cls="vit",
+        path=Path(f"~/.cache/equimo/{mname}.tar.lz4").expanduser(),
+    )

{equimo-0.2.2 → equimo-0.2.3}/equimo/converters/utils.py

@@ -1,5 +1,5 @@
 import re
-from typing import Dict, Tuple
+from typing import Any, Dict, Literal, Optional, Tuple
 
 import equinox as eqx
 import jax
@@ -33,12 +33,15 @@ def expand_torch_tensor(tensor, pos: str, n: int):
 
 def convert_params_from_torch_hub(
     jax_model: eqx.Module,
-    torch_hub_cfg: list[str],
     replace_cfg: Dict[str, str],
     expand_cfg: Dict[str, list],
     squeeze_cfg: Dict[str, int | None],
     whitelist: list[str],
     strict: bool = True,
+    source: Literal["torchhub", "timm"] = "torchhub",
+    torch_hub_cfg: Optional[list[str]] = None,
+    timm_cfg: Optional[list] = None,
+    return_torch: bool = False,
 ):
     """
     Load weights from a torch hub model into an Equinox module.
@@ -46,16 +49,37 @@ def convert_params_from_torch_hub(
     Args:
         jax_model (eqx.Module): A preexisting Jax model corresponding to the checkpoint to download.
         torch_hub_cfg (Tuple[str]): Arguments passed to `torch.hub.load()`.
-        transpose_whitelist (Set[str]): Parameters to exclude from format conversion.
+        replace_cfg (Dict[str, str]): Rename parameters from key to value.
+        expand_cfg (Dict[str, list]): Config to reshape params, see `expand_torch_tensor`
+        sqeeze_cfg (Dict[str, int|None]): Config to squeeze tensors, opposite of expand.
+        whitelist (Set[str]): Parameters to exclude from format conversion.
         strict (bool): Whether to crash on missing parameters one of the models.
+        source (str): Torch Hub or timm.
+        torch_hub_cfg (Optional[list]): args to pass to `torch.hub.load`.
+        timm_cfg (Optional[list]): args to pass to `timm.create_model`.
+        return_torch (bool): Return both jax and torch models.
     """
     try:
+        import timm
         import torch
     except:
         raise ImportError("`torch` not available")
 
-    # Load the pytorch model from torch hub
-    torch_model = torch.hub.load(*torch_hub_cfg)
+    # Load the pytorch model
+    match source:
+        case "torchhub":
+            if torch_hub_cfg is None:
+                raise ValueError(
+                    "The `torchhub` source is selected but `torch_hub_cfg` is None."
+                )
+            torch_model = torch.hub.load(*torch_hub_cfg)
+        case "timm":
+            if timm_cfg is None:
+                raise ValueError(
+                    "The `timm` source is selected but `timm_cfg` is None."
+                )
+            torch_model = timm.create_model(*timm_cfg)
+
     torch_params = dict(torch_model.named_parameters())
 
     # Extract the parameters from the defined Jax model
@@ -111,42 +135,71 @@ def convert_params_from_torch_hub(
             logger.error(_msg)
             raise AttributeError(_msg)
 
+    if return_torch:
+        return loaded_params, torch_model
     return loaded_params
 
 
 def convert_torch_to_equinox(
     jax_model: eqx.Module,
-    torch_hub_cfg: list[str],
     replace_cfg: dict = {},
     expand_cfg: dict = {},
     squeeze_cfg: dict = {},
     whitelist: list[str] = [],
     strict: bool = True,
-) -> eqx.Module:
+    source: Literal["torchhub", "timm"] = "torchhub",
+    torch_hub_cfg: Optional[list[str]] = None,
+    timm_cfg: Optional[list] = None,
+    return_torch: bool = False,
+) -> eqx.Module | Tuple[eqx.Module, Any]:
     """
     Convert a PyTorch model from torch.hub to Equinox format.
 
     Args:
         jax_model: The Equinox model
-        torch_hub_cfg: [repo, model_name] for torch.hub.load
         replace_cfg: Dict of parameter name replacements
         expand_cfg: Dict of dimensions to expand
         squeeze_cfg: Dict of dimensions to squeeze
         whitelist: List of parameters to keep from JAX model
         strict: Wether to raise an issue if not all weights are converted
+        source (str): Torch Hub or timm.
+        torch_hub_cfg: [repo, model_name] for torch.hub.load
+        timm_cfg (Optional[list]): args to pass to `timm.create_model`.
+        return_torch (bool): Return both jax and torch models.
 
     Returns:
         eqx.Module: Converted Equinox model in inference mode
     """
     dynamic, static = eqx.partition(jax_model, eqx.is_array)
-    converted_params = convert_params_from_torch_hub(
-        dynamic,
-        torch_hub_cfg,
-        replace_cfg,
-        expand_cfg,
-        squeeze_cfg,
-        whitelist,
-        strict,
-    )
-
-    return eqx.nn.inference_mode(eqx.combine(converted_params, static), value=True)
+    if return_torch:
+        converted_params, torch_model = convert_params_from_torch_hub(
+            dynamic,
+            replace_cfg,
+            expand_cfg,
+            squeeze_cfg,
+            whitelist,
+            strict,
+            source,
+            torch_hub_cfg,
+            timm_cfg,
+            return_torch,
+        )
+
+        return eqx.nn.inference_mode(
+            eqx.combine(converted_params, static), value=True
+        ), torch_model.eval()
+    else:
+        converted_params = convert_params_from_torch_hub(
+            dynamic,
+            replace_cfg,
+            expand_cfg,
+            squeeze_cfg,
+            whitelist,
+            strict,
+            source,
+            torch_hub_cfg,
+            timm_cfg,
+            return_torch,
+        )
+
+        return eqx.nn.inference_mode(eqx.combine(converted_params, static), value=True)

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

@@ -8,8 +8,10 @@ import jax
 import lz4.frame
 import requests
 from loguru import logger
+from semver import Version
 
 import equimo.models as em
+from equimo import __version__
 
 DEFAULT_REPOSITORY_URL = (
     "https://huggingface.co/poiretclement/equimo/resolve/main/models/default"
@@ -20,7 +22,8 @@ def save_model(
     path: Path,
     model: eqx.Module,
     model_config: dict,
-    torch_hub_cfg: list[str],
+    torch_hub_cfg: list[str] = [],
+    timm_cfg: list = [],
     compression: bool = True,
 ):
     """Save an Equinox model with its configuration and metadata to disk.
@@ -31,6 +34,7 @@ def save_model(
         model (eqx.Module): The Equinox model to be saved.
         model_config (dict): Configuration dictionary containing model hyperparameters.
         torch_hub_cfg (list[str]): List of torch hub configuration parameters.
+        timm_cfg (list[str]): List of timm configuration parameters.
         compression (bool, optional): Whether to compress the saved model using LZ4.
             Defaults to True.
 
@@ -46,8 +50,10 @@ def save_model(
     metadata = {
         "model_config": model_config,
         "torch_hub_cfg": torch_hub_cfg,
+        "timm": timm_cfg,
         "jax_version": jax.__version__,
         "equinox_version": eqx.__version__,
+        "equimo_version": __version__,
     }
 
     logger.debug(f"Metadata: {metadata}")
@@ -100,8 +106,9 @@ def download(identifier: str, repository: str) -> Path:
 
     logger.info(f"Downloading {identifier}...")
 
-    url = f"{repository}/{identifier}.tar.lz4"
-    path = Path(f"~/.cache/equimo/{identifier}.tar.lz4").expanduser()
+    model = identifier.split("_")[0]
+    url = f"{repository}/{model}/{identifier}.tar.lz4"
+    path = Path(f"~/.cache/equimo/{model}/{identifier}.tar.lz4").expanduser()
     path.parent.mkdir(parents=True, exist_ok=True)
 
     if path.exists():
@@ -123,6 +130,7 @@ def load_model(
     identifier: str | None = None,
     path: Path | None = None,
     repository: str = DEFAULT_REPOSITORY_URL,
+    **model_kwargs,
 ) -> eqx.Module:
     """Load an Equinox model from either a local path or remote repository.
 
@@ -135,6 +143,7 @@ def load_model(
             Mutually exclusive with identifier. Defaults to None.
         repository (str, optional): Base URL for model download.
             Defaults to DEFAULT_REPOSITORY_URL.
+        model_kwargs: kwargs passed to model instanciation. Overrides metadatas.
 
     Returns:
         eqx.Module: Loaded and initialized model with weights.
@@ -189,6 +198,12 @@ def load_model(
 
     logger.debug(f"Metadata: {metadata}")
 
+    model_eqm_version = metadata.get("equimo_version", "0.2.0")
+    if Version.parse(model_eqm_version) > Version.parse(__version__):
+        logger.warning(
+            f"The model you are importing was packaged with equimo v{model_eqm_version}, but you have equimo v{__version__}. You may face unexpected errors. Please consider updating equimo."
+        )
+
     # Class resolution
     match cls:
         case "vit":
@@ -207,7 +222,8 @@ def load_model(
             raise ValueError(f"Unknown model class: {cls}")
 
     # Reconstruct model skeleton
-    model = model_cls(**metadata["model_config"], key=jax.random.PRNGKey(42))
+    kwargs = metadata["model_config"] | model_kwargs
+    model = model_cls(**kwargs, key=jax.random.PRNGKey(42))
 
     # Load weights
     model = eqx.tree_deserialise_leaves(load_path / "weights.eqx", model)

equimo-0.2.3/equimo/layers/activation.py

@@ -0,0 +1,33 @@
+from typing import Callable
+import jax
+from functools import partial
+
+
+def get_act(activation: str | Callable) -> Callable:
+    """Get an activation function from its common name.
+
+    This is necessary because configs have to be stringified and stored as
+    json files to allow (de)serialization.
+    """
+    if not isinstance(activation, str):
+        return activation
+
+    match activation:
+        case "relu":
+            return jax.nn.relu
+        case "gelu":
+            return jax.nn.gelu
+        case "exactgelu":
+            return partial(jax.nn.gelu, approximate=False)
+        case "silu":
+            return jax.nn.silu
+        case "elu":
+            return jax.nn.elu
+        case "sigmoid":
+            return jax.nn.sigmoid
+        case "hard_sigmoid":
+            return jax.nn.hard_sigmoid
+        case "softmax":
+            return jax.nn.softmax
+        case _:
+            raise ValueError(f"Got an unknown activation string: {activation}")

{equimo-0.2.2 → equimo-0.2.3}/equimo/models/vit.py

@@ -8,6 +8,7 @@ import jax.random as jr
 from einops import rearrange
 from jaxtyping import Array, Float, Int, PRNGKeyArray
 
+from equimo.layers.activation import get_act
 from equimo.layers.attention import (
     Attention,
     AttentionBlock,
@@ -190,7 +191,7 @@ class VisionTransformer(eqx.Module):
         qk_norm: bool = False,
         attn_drop: float = 0.0,
         proj_drop: float = 0.0,
-        act_layer: Callable = jax.nn.gelu,
+        act_layer: Callable | str = jax.nn.gelu,
         attn_layer: str | eqx.Module = Attention,
         ffn_layer: str | eqx.Module = Mlp,
         ffn_bias: bool = True,
@@ -222,6 +223,7 @@ class VisionTransformer(eqx.Module):
         attn_layer = get_attention(attn_layer)
         ffn_layer = get_ffn(ffn_layer)
         norm_layer = get_norm(norm_layer)
+        act_layer = get_act(act_layer)
 
         self.patch_embed = PatchEmbedding(
             in_channels,
@@ -234,7 +236,7 @@ class VisionTransformer(eqx.Module):
             key=key_patchemb,
         )
         self.num_patches = self.patch_embed.num_patches
-        self.cls_token = jr.normal(key_cls, (1, dim))
+        self.cls_token = jr.normal(key_cls, (1, dim)) if class_token else None
         self.reg_tokens = (
             jr.normal(key_reg, (reg_tokens, dim)) if reg_tokens > 0 else None
         )

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

@@ -1,6 +1,6 @@
 [project]
 name = "Equimo"
-version = "0.2.2"
+version = "0.2.3"
 description = "Implementation of popular vision models in Jax"
 readme = "README.md"
 requires-python = ">=3.11"
@@ -12,6 +12,7 @@ dependencies = [
   "loguru>=0.7.3",
   "lz4>=4.4.3",
   "requests>=2.32.3",
+  "semver>=3.0.4",
 ]
 classifiers = [
   "Development Status :: 4 - Beta",
@@ -32,6 +33,7 @@ classifiers = [
 [project.optional-dependencies]
 torch = [
   "torch>=2.6.0",
+  "timm>=1.0.15",
 ]
 
 [dependency-groups]