Equimo 0.2.3__tar.gz → 0.2.4__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: Equimo
-Version: 0.2.3
+Version: 0.2.4
 Summary: Implementation of popular vision models in Jax
 Classifier: Development Status :: 4 - Beta
 Classifier: Intended Audience :: Developers
@@ -155,9 +155,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:
 
-Model identifiers allows downloading from equimo's [repository on huggingface](https://huggingface.co/poiretclement/equimo/tree/main/models/default)
+- [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)
 
 Identifiers are filenames without the extensions, such as:
 
@@ -165,6 +169,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.2.4}/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.2.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: Equimo
-Version: 0.2.3
+Version: 0.2.4
 Summary: Implementation of popular vision models in Jax
 Classifier: Development Status :: 4 - Beta
 Classifier: Intended Audience :: Developers
@@ -155,9 +155,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:
 
-Model identifiers allows downloading from equimo's [repository on huggingface](https://huggingface.co/poiretclement/equimo/tree/main/models/default)
+- [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)
 
 Identifiers are filenames without the extensions, such as:
 
@@ -165,6 +169,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.2.4}/README.md

@@ -128,9 +128,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:
 
-Model identifiers allows downloading from equimo's [repository on huggingface](https://huggingface.co/poiretclement/equimo/tree/main/models/default)
+- [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)
 
 Identifiers are filenames without the extensions, such as:
 
@@ -138,6 +142,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.4/equimo/__init__.py

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

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

@@ -72,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")

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

@@ -1,6 +1,6 @@
 [project]
 name = "Equimo"
-version = "0.2.3"
+version = "0.2.4"
 description = "Implementation of popular vision models in Jax"
 readme = "README.md"
 requires-python = ">=3.11"
@@ -43,3 +43,4 @@ dev = [
 
 [tool.setuptools]
 license-files = []
+packages = ["equimo"]

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)