Equimo 0.2.2__tar.gz → 0.2.4__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: Equimo
-Version: 0.2.2
+Version: 0.2.4
 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,33 @@ 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:
+The following models have pretrained weights available in Equimo:
+
+- [DinoV2](https://arxiv.org/abs/2304.07193),
+- [SigLIP2](https://arxiv.org/abs/2502.14786),
+- [TIPS](https://arxiv.org/abs/2410.16512).
+
+Model identifiers allow 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`
+- `tips_vitg14_lr`
 
 ## Contributing
 

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

@@ -9,10 +9,8 @@ Equimo.egg-info/top_level.txt
 equimo/__init__.py
 equimo/io.py
 equimo/utils.py
-equimo/converters/__init__.py
-equimo/converters/dinov2.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.4}/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.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: Equimo
-Version: 0.2.2
+Version: 0.2.4
 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,33 @@ 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:
+The following models have pretrained weights available in Equimo:
+
+- [DinoV2](https://arxiv.org/abs/2304.07193),
+- [SigLIP2](https://arxiv.org/abs/2502.14786),
+- [TIPS](https://arxiv.org/abs/2410.16512).
+
+Model identifiers allow 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`
+- `tips_vitg14_lr`
 
 ## Contributing
 

{equimo-0.2.2 → equimo-0.2.4}/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,33 @@ 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:
+The following models have pretrained weights available in Equimo:
+
+- [DinoV2](https://arxiv.org/abs/2304.07193),
+- [SigLIP2](https://arxiv.org/abs/2502.14786),
+- [TIPS](https://arxiv.org/abs/2410.16512).
+
+Model identifiers allow 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`
+- `tips_vitg14_lr`
 
 ## Contributing
 

equimo-0.2.4/equimo/__init__.py

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

{equimo-0.2.2 → equimo-0.2.4}/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}")
@@ -66,6 +72,7 @@ def save_model(
             eqx.tree_serialise_leaves(tmp_path / "weights.eqx", model)
 
             # Create compressed archive
+            path.parent.mkdir(parents=True, exist_ok=True)
             with lz4.frame.open(path, "wb") as f_out:
                 with tarfile.open(fileobj=f_out, mode="w") as tar:
                     tar.add(tmp_path / "metadata.json", arcname="metadata.json")
@@ -100,8 +107,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 +131,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 +144,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 +199,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 +223,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.4/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.4}/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.4}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "Equimo"
-version = "0.2.2"
+version = "0.2.4"
 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]
@@ -41,3 +43,4 @@ dev = [
 
 [tool.setuptools]
 license-files = []
+packages = ["equimo"]

equimo-0.2.2/equimo/converters/__init__.py

@@ -1 +0,0 @@
-

equimo-0.2.2/equimo/converters/dinov2.py

@@ -1,69 +0,0 @@
-from pathlib import Path
-
-import jax
-
-import equimo.models as em
-from equimo.converters.utils import convert_torch_to_equinox
-from equimo.io import load_model, save_model
-
-
-def main():
-    key = jax.random.PRNGKey(42)
-    dinov2_vits14_reg_config = {
-        "img_size": 518,
-        "in_channels": 3,
-        "dim": 384,
-        "patch_size": 14,
-        "num_heads": [6],
-        "depths": [12],
-        "num_classes": 0,
-        "use_mask_token": True,
-        "reg_tokens": 4,
-        "init_values": 1e-5,
-        "eps": 1e-6,
-        "dynamic_img_size": True,
-    }
-
-    dinov2_vits14_reg = em.VisionTransformer(
-        **dinov2_vits14_reg_config,
-        key=key,
-    )
-
-    torch_hub_cfg = ["facebookresearch/dinov2", "dinov2_vits14_reg"]
-
-    replace_cfg = {
-        "reg_tokens": "register_tokens",
-        "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 = []
-
-    dinov2_vits14_reg = convert_torch_to_equinox(
-        dinov2_vits14_reg,
-        torch_hub_cfg,
-        replace_cfg,
-        expand_cfg,
-        squeeze_cfg,
-        whitelist,
-        strict=True,
-    )
-
-    save_model(
-        Path("~/.cache/equimo/dinov2_vits14_reg").expanduser(),
-        dinov2_vits14_reg,
-        dinov2_vits14_reg_config,
-        torch_hub_cfg,
-        compression=True,
-    )
-
-    _ = load_model(
-        cls="vit",
-        path=Path("~/.cache/equimo/dinov2_vits14_reg.tar.lz4").expanduser(),
-    )

equimo-0.2.2/equimo/converters/utils.py

@@ -1,152 +0,0 @@
-import re
-from typing import Dict, Tuple
-
-import equinox as eqx
-import jax
-import jax.numpy as jnp
-from jax.tree_util import GetAttrKey, SequenceKey
-from loguru import logger
-
-
-def stringify_name(path: Tuple) -> str:
-    stringified = []
-    for p in path:
-        if isinstance(p, GetAttrKey):
-            stringified.append(p.name)
-        if isinstance(p, SequenceKey):
-            stringified.append(str(p.idx))
-    return ".".join(stringified)
-
-
-def expand_torch_tensor(tensor, pos: str, n: int):
-    padding = [None] * n
-    match pos:
-        case "before":
-            return tensor[*padding, ...]
-        case "after":
-            return tensor[..., *padding]
-        case _:
-            raise ValueError(
-                f"Invalid `pos`, expected one of [`before`, `after`], got: {pos}"
-            )
-
-
-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,
-):
-    """
-    Load weights from a torch hub model into an Equinox module.
-
-    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.
-        strict (bool): Whether to crash on missing parameters one of the models.
-    """
-    try:
-        import torch
-    except:
-        raise ImportError("`torch` not available")
-
-    # Load the pytorch model from torch hub
-    torch_model = torch.hub.load(*torch_hub_cfg)
-    torch_params = dict(torch_model.named_parameters())
-
-    # Extract the parameters from the defined Jax model
-    jax_params = eqx.filter(jax_model, eqx.is_array)
-    # _, jax_params, _ = nnx.split(jax_model, nnx.Param, ...)
-    jax_params_flat, jax_param_pytree = jax.tree_util.tree_flatten_with_path(jax_params)
-
-    torch_params_flat = []
-    for path, param in jax_params_flat:
-        # Match the parameters' path of pytorch
-        param_path = stringify_name(path)
-        param_path = re.sub(r"\.scale|.kernel", ".weight", param_path)
-
-        for old, new in replace_cfg.items():
-            param_path = param_path.replace(old, new)
-
-        shape = param.shape
-
-        if param_path not in torch_params:
-            _msg = f"{param_path} ({shape}) not found in PyTorch model."
-            if strict:
-                logger.error(_msg)
-                raise AttributeError(_msg)
-
-            logger.warning(f"{_msg} Appending `None` to flat param list.")
-            torch_params_flat.append(None)
-            continue
-
-        logger.info(f"Converting {param_path}...")
-        torch_param = torch_params[param_path]
-
-        if param_path in expand_cfg:
-            torch_param = expand_torch_tensor(torch_param, *expand_cfg[param_path])
-        if param_path in squeeze_cfg:
-            torch_param = torch.squeeze(torch_param, dim=squeeze_cfg[param_path])
-
-        if shape != torch_param.shape:
-            _msg = f"`{param_path}`: expected shape ({shape}) does not match its pytorch implementation ({torch_param.shape})."
-            logger.error(_msg)
-            raise ValueError(_msg)
-
-        torch_params_flat.append(jnp.asarray(torch_param.detach().numpy()))
-        _ = torch_params.pop(param_path)
-
-    loaded_params = jax.tree_util.tree_unflatten(jax_param_pytree, torch_params_flat)
-
-    for path, param in torch_params.items():
-        logger.warning(
-            f"PyTorch parameters `{path}` ({param.shape}) were not converted."
-        )
-        if strict and path not in whitelist:
-            _msg = f"The PyTorch model contains parameters ({path}) that do not have a Jax counterpart."
-            logger.error(_msg)
-            raise AttributeError(_msg)
-
-    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:
-    """
-    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
-
-    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)