Equimo 0.2.3__tar.gz → 0.3.0__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: Equimo
-Version: 0.2.3
+Version: 0.3.0
 Summary: Implementation of popular vision models in Jax
 Classifier: Development Status :: 4 - Beta
 Classifier: Intended Audience :: Developers
@@ -24,6 +24,13 @@ 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"
+Provides-Extra: text
+Requires-Dist: tensorflow>=2.18.1; extra == "text"
+Requires-Dist: tensorflow-text>=2.18.1; extra == "text"
+Provides-Extra: extras
+Requires-Dist: matplotlib>=3.10.1; extra == "extras"
+Requires-Dist: pillow>=11.1.0; extra == "extras"
+Requires-Dist: scikit-learn>=1.6.1; extra == "extras"
 
 # Equimo: Modern Vision Models in JAX/Equinox
 
@@ -37,6 +44,7 @@ Equimo (Equinox Image Models) provides JAX/Equinox implementations of recent com
 - Focus on recent architectures (2023-2024 papers)
 - Modular design for easy experimentation
 - Extensive documentation and type hints
+- **Experimental** support for text embedding
 
 ## Installation
 
@@ -98,6 +106,76 @@ x = jax.random.normal(key, (3, 224, 224))
 output = model(x, enable_dropout=False, key=key)
 ```
 
+## Working with text embeddings
+
+**Warning: this is experimental, it can break or change at any time**
+
+`equimo.experimental.text` has been added since v0.3.0. It allows working with both text and images. It is especially
+useful for models like SigLIP or TIPS, although only TIPS is currently supported.
+
+Currently, text tokenization relies on tensorflow_text, install equimo with the `text` group such as
+`uv add equimo[text]`.
+
+Here is a very simple example of a 0-shot classification based on the comparison between text and image embeddings:
+
+```python
+import jax
+from einops import rearrange
+
+from equimo.experimental.text import Tokenizer
+from equimo.io import load_image, load_model
+from equimo.utils import PCAVisualizer, normalize, plot_image_and_feature_map
+
+# Random demo inputs
+key = jax.random.PRNGKey(42)
+image = load_image("./demo.jpg", size=448)
+text = [
+    "A baby discovering happiness",
+    "A computer",
+]
+
+# Loading pretrained models
+image_encoder = load_model("vit", "tips_vits14_hr")
+text_encoder = load_model("experimental.textencoder", "tips_vits14_hr_text")
+
+# Encoding text and image
+ids, paddings = Tokenizer(identifier="sentencepiece_tips").tokenize(text, max_len=64)
+
+text_embedding = normalize(
+    jax.vmap(text_encoder, in_axes=(0, 0, None))(ids, paddings, key)
+)
+image_embedding = jax.vmap(image_encoder.norm)(image_encoder.features(image, key))
+cls_token = normalize(image_embedding[0])
+spatial_features = rearrange(
+    image_embedding[2:], "(h w) d -> h w d", h=int(448 / 14), w=int(448 / 14)
+)
+
+# Getting probabilities based on Cosine Similarity
+cos_sim = jax.nn.softmax(
+    ((cls_token[None, :] @ text_embedding.T) / text_encoder.temperature), axis=-1
+)
+
+# Plot the results
+label_idxs = jax.numpy.argmax(cos_sim, axis=-1)
+cos_sim_max = jax.numpy.max(cos_sim, axis=-1)
+label_predicted = text[label_idxs[0]]
+similarity = cos_sim_max[0]
+pca_obj = PCAVisualizer(spatial_features)
+image_pca = pca_obj(spatial_features)
+
+plot_image_and_feature_map(
+    image.transpose(1, 2, 0),
+    image_pca,
+    "./out.png",
+    "Input Image",
+    f"{label_predicted}, prob: {similarity * 100:.2f}%",
+)
+```
+
+Resulting in such a wonderful result:
+
+![Output of TIPS 0-shot classification](./docs/assets/0-shot.png)
+
 ## Saving and Loading Models
 
 Equimo provides utilities for saving models locally and loading pre-trained models from the
@@ -155,9 +233,13 @@ model = load_model(
 
 #### List of pretrained models
 
-Currently, only [DinoV2](https://github.com/facebookresearch/dinov2/) and SigLIP2 have been ported.
+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 allows downloading from equimo's [repository on huggingface](https://huggingface.co/poiretclement/equimo/tree/main/models/default)
+Model identifiers allow downloading from equimo's [repository on huggingface](https://huggingface.co/poiretclement/equimo/tree/main/models/default)
 
 Identifiers are filenames without the extensions, such as:
 
@@ -165,6 +247,7 @@ Identifiers are filenames without the extensions, such as:
 - `dinov2_vits14_reg`
 - `siglip2_vitl16_512`
 - `siglip2_vitso400m16_384`
+- `tips_vitg14_lr`
 
 ## Contributing
 

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

@@ -9,10 +9,6 @@ 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/siglip2.py
-equimo/converters/utils.py
 equimo/layers/__init__.py
 equimo/layers/activation.py
 equimo/layers/attention.py

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

@@ -7,6 +7,15 @@ lz4>=4.4.3
 requests>=2.32.3
 semver>=3.0.4
 
+[extras]
+matplotlib>=3.10.1
+pillow>=11.1.0
+scikit-learn>=1.6.1
+
+[text]
+tensorflow>=2.18.1
+tensorflow-text>=2.18.1
+
 [torch]
 torch>=2.6.0
 timm>=1.0.15

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: Equimo
-Version: 0.2.3
+Version: 0.3.0
 Summary: Implementation of popular vision models in Jax
 Classifier: Development Status :: 4 - Beta
 Classifier: Intended Audience :: Developers
@@ -24,6 +24,13 @@ 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"
+Provides-Extra: text
+Requires-Dist: tensorflow>=2.18.1; extra == "text"
+Requires-Dist: tensorflow-text>=2.18.1; extra == "text"
+Provides-Extra: extras
+Requires-Dist: matplotlib>=3.10.1; extra == "extras"
+Requires-Dist: pillow>=11.1.0; extra == "extras"
+Requires-Dist: scikit-learn>=1.6.1; extra == "extras"
 
 # Equimo: Modern Vision Models in JAX/Equinox
 
@@ -37,6 +44,7 @@ Equimo (Equinox Image Models) provides JAX/Equinox implementations of recent com
 - Focus on recent architectures (2023-2024 papers)
 - Modular design for easy experimentation
 - Extensive documentation and type hints
+- **Experimental** support for text embedding
 
 ## Installation
 
@@ -98,6 +106,76 @@ x = jax.random.normal(key, (3, 224, 224))
 output = model(x, enable_dropout=False, key=key)
 ```
 
+## Working with text embeddings
+
+**Warning: this is experimental, it can break or change at any time**
+
+`equimo.experimental.text` has been added since v0.3.0. It allows working with both text and images. It is especially
+useful for models like SigLIP or TIPS, although only TIPS is currently supported.
+
+Currently, text tokenization relies on tensorflow_text, install equimo with the `text` group such as
+`uv add equimo[text]`.
+
+Here is a very simple example of a 0-shot classification based on the comparison between text and image embeddings:
+
+```python
+import jax
+from einops import rearrange
+
+from equimo.experimental.text import Tokenizer
+from equimo.io import load_image, load_model
+from equimo.utils import PCAVisualizer, normalize, plot_image_and_feature_map
+
+# Random demo inputs
+key = jax.random.PRNGKey(42)
+image = load_image("./demo.jpg", size=448)
+text = [
+    "A baby discovering happiness",
+    "A computer",
+]
+
+# Loading pretrained models
+image_encoder = load_model("vit", "tips_vits14_hr")
+text_encoder = load_model("experimental.textencoder", "tips_vits14_hr_text")
+
+# Encoding text and image
+ids, paddings = Tokenizer(identifier="sentencepiece_tips").tokenize(text, max_len=64)
+
+text_embedding = normalize(
+    jax.vmap(text_encoder, in_axes=(0, 0, None))(ids, paddings, key)
+)
+image_embedding = jax.vmap(image_encoder.norm)(image_encoder.features(image, key))
+cls_token = normalize(image_embedding[0])
+spatial_features = rearrange(
+    image_embedding[2:], "(h w) d -> h w d", h=int(448 / 14), w=int(448 / 14)
+)
+
+# Getting probabilities based on Cosine Similarity
+cos_sim = jax.nn.softmax(
+    ((cls_token[None, :] @ text_embedding.T) / text_encoder.temperature), axis=-1
+)
+
+# Plot the results
+label_idxs = jax.numpy.argmax(cos_sim, axis=-1)
+cos_sim_max = jax.numpy.max(cos_sim, axis=-1)
+label_predicted = text[label_idxs[0]]
+similarity = cos_sim_max[0]
+pca_obj = PCAVisualizer(spatial_features)
+image_pca = pca_obj(spatial_features)
+
+plot_image_and_feature_map(
+    image.transpose(1, 2, 0),
+    image_pca,
+    "./out.png",
+    "Input Image",
+    f"{label_predicted}, prob: {similarity * 100:.2f}%",
+)
+```
+
+Resulting in such a wonderful result:
+
+![Output of TIPS 0-shot classification](./docs/assets/0-shot.png)
+
 ## Saving and Loading Models
 
 Equimo provides utilities for saving models locally and loading pre-trained models from the
@@ -155,9 +233,13 @@ model = load_model(
 
 #### List of pretrained models
 
-Currently, only [DinoV2](https://github.com/facebookresearch/dinov2/) and SigLIP2 have been ported.
+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 allows downloading from equimo's [repository on huggingface](https://huggingface.co/poiretclement/equimo/tree/main/models/default)
+Model identifiers allow downloading from equimo's [repository on huggingface](https://huggingface.co/poiretclement/equimo/tree/main/models/default)
 
 Identifiers are filenames without the extensions, such as:
 
@@ -165,6 +247,7 @@ Identifiers are filenames without the extensions, such as:
 - `dinov2_vits14_reg`
 - `siglip2_vitl16_512`
 - `siglip2_vitso400m16_384`
+- `tips_vitg14_lr`
 
 ## Contributing
 

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

@@ -10,6 +10,7 @@ Equimo (Equinox Image Models) provides JAX/Equinox implementations of recent com
 - Focus on recent architectures (2023-2024 papers)
 - Modular design for easy experimentation
 - Extensive documentation and type hints
+- **Experimental** support for text embedding
 
 ## Installation
 
@@ -71,6 +72,76 @@ x = jax.random.normal(key, (3, 224, 224))
 output = model(x, enable_dropout=False, key=key)
 ```
 
+## Working with text embeddings
+
+**Warning: this is experimental, it can break or change at any time**
+
+`equimo.experimental.text` has been added since v0.3.0. It allows working with both text and images. It is especially
+useful for models like SigLIP or TIPS, although only TIPS is currently supported.
+
+Currently, text tokenization relies on tensorflow_text, install equimo with the `text` group such as
+`uv add equimo[text]`.
+
+Here is a very simple example of a 0-shot classification based on the comparison between text and image embeddings:
+
+```python
+import jax
+from einops import rearrange
+
+from equimo.experimental.text import Tokenizer
+from equimo.io import load_image, load_model
+from equimo.utils import PCAVisualizer, normalize, plot_image_and_feature_map
+
+# Random demo inputs
+key = jax.random.PRNGKey(42)
+image = load_image("./demo.jpg", size=448)
+text = [
+    "A baby discovering happiness",
+    "A computer",
+]
+
+# Loading pretrained models
+image_encoder = load_model("vit", "tips_vits14_hr")
+text_encoder = load_model("experimental.textencoder", "tips_vits14_hr_text")
+
+# Encoding text and image
+ids, paddings = Tokenizer(identifier="sentencepiece_tips").tokenize(text, max_len=64)
+
+text_embedding = normalize(
+    jax.vmap(text_encoder, in_axes=(0, 0, None))(ids, paddings, key)
+)
+image_embedding = jax.vmap(image_encoder.norm)(image_encoder.features(image, key))
+cls_token = normalize(image_embedding[0])
+spatial_features = rearrange(
+    image_embedding[2:], "(h w) d -> h w d", h=int(448 / 14), w=int(448 / 14)
+)
+
+# Getting probabilities based on Cosine Similarity
+cos_sim = jax.nn.softmax(
+    ((cls_token[None, :] @ text_embedding.T) / text_encoder.temperature), axis=-1
+)
+
+# Plot the results
+label_idxs = jax.numpy.argmax(cos_sim, axis=-1)
+cos_sim_max = jax.numpy.max(cos_sim, axis=-1)
+label_predicted = text[label_idxs[0]]
+similarity = cos_sim_max[0]
+pca_obj = PCAVisualizer(spatial_features)
+image_pca = pca_obj(spatial_features)
+
+plot_image_and_feature_map(
+    image.transpose(1, 2, 0),
+    image_pca,
+    "./out.png",
+    "Input Image",
+    f"{label_predicted}, prob: {similarity * 100:.2f}%",
+)
+```
+
+Resulting in such a wonderful result:
+
+![Output of TIPS 0-shot classification](./docs/assets/0-shot.png)
+
 ## Saving and Loading Models
 
 Equimo provides utilities for saving models locally and loading pre-trained models from the
@@ -128,9 +199,13 @@ model = load_model(
 
 #### List of pretrained models
 
-Currently, only [DinoV2](https://github.com/facebookresearch/dinov2/) and SigLIP2 have been ported.
+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 allows downloading from equimo's [repository on huggingface](https://huggingface.co/poiretclement/equimo/tree/main/models/default)
+Model identifiers allow downloading from equimo's [repository on huggingface](https://huggingface.co/poiretclement/equimo/tree/main/models/default)
 
 Identifiers are filenames without the extensions, such as:
 
@@ -138,6 +213,7 @@ Identifiers are filenames without the extensions, such as:
 - `dinov2_vits14_reg`
 - `siglip2_vitl16_512`
 - `siglip2_vitso400m16_384`
+- `tips_vitg14_lr`
 
 ## Contributing
 

equimo-0.3.0/equimo/__init__.py

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

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

@@ -1,10 +1,13 @@
 import json
+import io
 import tarfile
 import tempfile
 from pathlib import Path
+from typing import Optional
 
 import equinox as eqx
 import jax
+import jax.numpy as jnp
 import lz4.frame
 import requests
 from loguru import logger
@@ -72,6 +75,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")
@@ -130,6 +134,7 @@ def load_model(
     identifier: str | None = None,
     path: Path | None = None,
     repository: str = DEFAULT_REPOSITORY_URL,
+    inference_mode: bool = True,
     **model_kwargs,
 ) -> eqx.Module:
     """Load an Equinox model from either a local path or remote repository.
@@ -143,6 +148,8 @@ def load_model(
             Mutually exclusive with identifier. Defaults to None.
         repository (str, optional): Base URL for model download.
             Defaults to DEFAULT_REPOSITORY_URL.
+        inference_mode (bool): Disables dropouts if True.
+            Defaults to True.
         model_kwargs: kwargs passed to model instanciation. Overrides metadatas.
 
     Returns:
@@ -218,6 +225,10 @@ def load_model(
             model_cls = em.FasterViT
         case "partialformer":
             model_cls = em.PartialFormer
+        case "experimental.textencoder":
+            from equimo.experimental.text import TextEncoder
+
+            model_cls = TextEncoder
         case _:
             raise ValueError(f"Unknown model class: {cls}")
 
@@ -225,9 +236,52 @@ def load_model(
     kwargs = metadata["model_config"] | model_kwargs
     model = model_cls(**kwargs, key=jax.random.PRNGKey(42))
 
-    # Load weights
+    # Load weights and set inference mode
     model = eqx.tree_deserialise_leaves(load_path / "weights.eqx", model)
+    model = eqx.nn.inference_mode(model, inference_mode)
 
     logger.info("Model loaded successfully.")
 
     return model
+
+
+def load_image(
+    path: str,
+    mean: Optional[list[float]] = None,
+    std: Optional[list[float]] = None,
+    size: Optional[int] = None,
+):
+    """Load an image and perform minor preprocessing.
+
+    Args:
+        path (str): Path of the image.
+        mean (list, optional): Channel mean for normalization.
+        std (list, optional): Channel std for normalization.
+        size (int, optional): Size to which resize the image.
+
+    Returns:
+        jnp.array: The loaded image.
+
+    Raises:
+        ImportError: If PIL can't be loaded.
+    """
+    try:
+        from PIL import Image
+    except ImportError:
+        raise ImportError("PIL is needed to be able to load images.")
+
+    with open(path, "rb") as fd:
+        image_bytes = io.BytesIO(fd.read())
+        pil_image = Image.open(image_bytes)
+
+        array = jnp.array(pil_image).astype(jnp.float32) / 255.0
+        if size is not None:
+            array = jax.image.resize(array, (size, size, 3), method="bilinear")
+
+        if mean is not None and std is not None:
+            mean = jnp.array(mean)[None, None, :]
+            std = jnp.array(std)[None, None, :]
+
+            array = (array - mean) / std
+
+    return array.transpose(2, 0, 1)

{equimo-0.2.3 → equimo-0.3.0}/equimo/layers/attention.py

@@ -1,5 +1,4 @@
 from typing import Callable, List, Optional, Tuple
-
 import equinox as eqx
 import jax
 import jax.numpy as jnp
@@ -74,6 +73,7 @@ class Attention(eqx.Module):
         x: Float[Array, "seqlen dim"],
         key: PRNGKeyArray,
         inference: Optional[bool] = None,
+        mask: Optional[Float[Array, ""]] = None,
     ) -> Float[Array, "seqlen dim"]:
         key1, key2 = jr.split(key, 2)
 
@@ -90,6 +90,10 @@ class Attention(eqx.Module):
         k = jax.vmap(jax.vmap(self.k_norm))(k)
 
         attn = jnp.einsum("hqd,hkd->hqk", q, k) / jnp.sqrt(self.head_dim)
+
+        if mask is not None:
+            attn = jnp.where(mask == 0, jnp.finfo(attn.dtype).min, attn)
+
         attn = jax.nn.softmax(attn, axis=-1)
         attn = self.attn_drop(attn, inference=inference, key=key1)
 
@@ -312,9 +316,14 @@ class AttentionBlock(eqx.Module):
         x: Float[Array, "seqlen dim"],
         key: PRNGKeyArray,
         inference: Optional[bool] = None,
+        mask: Optional[Float[Array, ""]] = None,
     ) -> Float[Array, "seqlen dim"]:
         key_attn, key_mlp, key_dr1, key_dr2 = jr.split(key, 4)
 
+        # I chose to define extra args here rather than passing mask directly
+        # because not all attention mechanisms support masks as args
+        extra_kwargs = {"mask": mask} if mask is not None else {}
+
         x = self.drop_path1(
             x,
             self.ls1(
@@ -323,6 +332,7 @@ class AttentionBlock(eqx.Module):
                         jax.vmap(self.prenorm)(x),
                         inference=inference,
                         key=key_attn,
+                        **extra_kwargs,
                     )
                 )
             ),
@@ -336,6 +346,7 @@ class AttentionBlock(eqx.Module):
                     jax.vmap(self.norm)(x),
                     inference=inference,
                     key=key_mlp,
+                    **extra_kwargs,
                 )
             ),
             inference=inference,

{equimo-0.2.3 → equimo-0.3.0}/equimo/layers/ffn.py

@@ -210,6 +210,7 @@ class Mlp(eqx.Module):
         self,
         x: Float[Array, "seqlen dim"],
         key: PRNGKeyArray,
+        mask: Optional[Float[Array, ""]] = None,
         inference: Optional[bool] = None,
     ) -> Float[Array, "seqlen dim"]:
         key_dr1, key_dr2 = jr.split(key, 2)
@@ -219,12 +220,19 @@ class Mlp(eqx.Module):
             inference=inference,
             key=key_dr1,
         )
+
+        if mask is not None:
+            x *= mask
+
         x = self.drop2(
             jax.vmap(self.fc2)(x),
             inference=inference,
             key=key_dr2,
         )
 
+        if mask is not None:
+            x *= mask
+
         return x
 
 

{equimo-0.2.3 → equimo-0.3.0}/equimo/utils.py

@@ -4,6 +4,84 @@ import jax
 import jax.numpy as jnp
 from jaxtyping import Array, Float
 
+import typing as t
+import numpy as np
+
+
+_ArrayLike = t.Union[np.ndarray, jnp.ndarray]
+
+
+def normalize(x, order: int = 2):
+    return x / np.linalg.norm(x, ord=order, axis=-1, keepdims=True).clip(min=1e-3)
+
+
+class PCAVisualizer:
+    """PCA visualizer.
+
+    Taken as-is from https://github.com/google-deepmind/tips/blob/72820c1841f973c9543d9c95c5ff2262ec621955/scenic/utils/feature_viz.py#L32
+    """
+
+    def __init__(
+        self, features: _ArrayLike, n_samples: int = 100000, n_components: int = 3
+    ) -> None:
+        """Creates a PCA object for visualizing features of shape [..., F]."""
+        try:
+            from sklearn import decomposition
+        except ImportError:
+            raise ImportError(
+                "You need sklearn to use the PCAVisualizer, install it using `uv add equimo[viz]`"
+            )
+
+        features = np.array(features)
+        pca_object = decomposition.PCA(n_components=n_components)
+        features = features.reshape([-1, features.shape[-1]])
+        features = features[np.random.randint(0, features.shape[0], n_samples), :]
+        pca_object.fit(features)
+        self.pca_object = pca_object
+        self.n_components = n_components
+
+    def __call__(self, features: _ArrayLike) -> np.ndarray:
+        """Apply PCA to features of shape [..., F]."""
+        features = np.array(features)
+        features_pca = self.pca_object.transform(
+            features.reshape([-1, features.shape[-1]])
+        ).reshape(features.shape[:-1] + (self.n_components,))
+        return normalize(features_pca) * 0.5 + 0.5
+
+
+def plot_image_and_feature_map(
+    image,
+    feature_map,
+    save_path,
+    image_title="Image",
+    map_title="Feature Map",
+    fontsize=16,
+):
+    """
+    Plots an image and its feature map side by side and saves the figure.
+    """
+    try:
+        import matplotlib.pyplot as plt
+    except ImportError:
+        raise ImportError("Matplotlib is required to plot an image")
+
+    feature_map_normalized = (feature_map - feature_map.min()) / (
+        feature_map.max() - feature_map.min()
+    )
+
+    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(10, 5))
+    ax1.imshow(image)
+    ax1.set_title(image_title, fontsize=fontsize)
+    ax1.axis("off")
+
+    ax2.imshow(feature_map_normalized, cmap="gray")
+    ax2.set_title(map_title, fontsize=fontsize)
+    ax2.axis("off")
+
+    plt.tight_layout()
+    plt.savefig(save_path)
+    plt.close()
+
 
 def make_2tuple(x):
     """Convert input into a 2-tuple.

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

@@ -1,6 +1,6 @@
 [project]
 name = "Equimo"
-version = "0.2.3"
+version = "0.3.0"
 description = "Implementation of popular vision models in Jax"
 readme = "README.md"
 requires-python = ">=3.11"
@@ -31,10 +31,22 @@ classifiers = [
 ]
 
 [project.optional-dependencies]
+# Useful to convert models from PyTorch
 torch = [
   "torch>=2.6.0",
   "timm>=1.0.15",
 ]
+# Required for equimo.experimental.text (text tokenizer)
+text = [
+  "tensorflow>=2.18.1",
+  "tensorflow-text>=2.18.1",
+]
+# Required for some external utils
+extras = [
+  "matplotlib>=3.10.1",
+  "pillow>=11.1.0",
+  "scikit-learn>=1.6.1",
+]
 
 [dependency-groups]
 dev = [
@@ -43,3 +55,4 @@ dev = [
 
 [tool.setuptools]
 license-files = []
+packages = ["equimo"]

{equimo-0.2.3 → equimo-0.3.0}/tests/test_models.py

@@ -75,7 +75,9 @@ def test_save_load_model_compressed():
         save_model(save_path, model, model_config, torch_hub_cfg, compression=True)
 
         # Load model
-        loaded_model = load_model(cls="vit", path=save_path.with_suffix(".tar.lz4"))
+        loaded_model = load_model(
+            cls="vit", path=save_path.with_suffix(".tar.lz4"), dynamic_img_size=True
+        )
 
         # Test loaded model
         loaded_output = loaded_model.features(x, key=key)
@@ -116,7 +118,7 @@ def test_save_load_model_uncompressed():
 
         save_model(save_path, model, model_config, torch_hub_cfg, compression=False)
 
-        loaded_model = load_model(cls="vit", path=save_path)
+        loaded_model = load_model(cls="vit", path=save_path, dynamic_img_size=True)
         loaded_output = loaded_model.features(x, key=key)
 
         assert jnp.allclose(original_output, loaded_output, atol=1e-5)
@@ -125,7 +127,7 @@ def test_save_load_model_uncompressed():
 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")
+    model = load_model(cls="vit", identifier="dinov2_vits14_reg", dynamic_img_size=True)
 
     # Test inference
     x = jr.normal(key, (3, 224, 224))

equimo-0.2.3/equimo/__init__.py

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

equimo-0.2.3/equimo/converters/__init__.py

@@ -1 +0,0 @@
-

equimo-0.2.3/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.3/equimo/converters/siglip2.py

@@ -1,105 +0,0 @@
-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.3/equimo/converters/utils.py

@@ -1,205 +0,0 @@
-import re
-from typing import Any, Dict, Literal, Optional, 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,
-    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.
-
-    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()`.
-        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
-    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
-    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)
-
-    if return_torch:
-        return loaded_params, torch_model
-    return loaded_params
-
-
-def convert_torch_to_equinox(
-    jax_model: eqx.Module,
-    replace_cfg: dict = {},
-    expand_cfg: dict = {},
-    squeeze_cfg: dict = {},
-    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,
-) -> eqx.Module | Tuple[eqx.Module, Any]:
-    """
-    Convert a PyTorch model from torch.hub to Equinox format.
-
-    Args:
-        jax_model: The Equinox model
-        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)
-    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)