Equimo 0.2.4__tar.gz → 0.3.0__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: Equimo
-Version: 0.2.4
+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

{equimo-0.2.4 → 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.4 → equimo-0.3.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: Equimo
-Version: 0.2.4
+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

{equimo-0.2.4 → 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

equimo-0.3.0/equimo/__init__.py

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

{equimo-0.2.4 → 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
@@ -131,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.
@@ -144,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:
@@ -219,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}")
 
@@ -226,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.4 → 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.4 → 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.4 → 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.4 → equimo-0.3.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "Equimo"
-version = "0.2.4"
+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 = [

{equimo-0.2.4 → 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.4/equimo/__init__.py

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