Jax Image Modeling of Models
Project description
Jax Image Modeling of Models (jimm)
Docs are at: https://pythoncrazy.github.io/jimm
- This aims to be the jax counterpart to timm, with the exception that for image-text models (CLIP, SigLIP, etc), we support the text model entirely.
- Made with flax nnx, supports weight loading from pytorch_model.bin and safetensors (as well as both methods from huggingface).
Models Supported:
- Vision Transformers
- Both with a classification linear layer, or not
- Using a CLS Token for pooling, or using Multihead Attention Pooling
- Can load any standard variant of Vision Transformers of any size/resolution(e.g. "google/vit-base-patch16-224" or "google/vit-large-patch16-384")
- CLIP
- Can load from any checkpoints of the clip model on github (such as "openai/clip-vit-base-patch32" or "geolocal/StreetCLIP")
- SigLIP
- Can load any non-naflex version of the SigLIP model, from both siglipv1 and siglipv2 (eg "google/siglip-base-patch16-256" or "google/siglip2-large-patch16-512" from huggingface or locally)
Installation
Using uv
uv add jax-image-models
Using pip/conda
pip install jax-image-models
Release history Release notifications | RSS feed
Download files
Download the file for your platform. If you're not sure which to choose, learn more about installing packages.
Source Distribution
Built Distribution
Filter files by name, interpreter, ABI, and platform.
If you're not sure about the file name format, learn more about wheel file names.
Copy a direct link to the current filters
File details
Details for the file jax_image_models-0.4.0.tar.gz.
File metadata
- Download URL: jax_image_models-0.4.0.tar.gz
- Upload date:
- Size: 7.7 MB
- Tags: Source
- Uploaded using Trusted Publishing? Yes
- Uploaded via: uv/0.11.21 {"installer":{"name":"uv","version":"0.11.21","subcommand":["publish"]},"python":null,"implementation":{"name":null,"version":null},"distro":{"name":"Ubuntu","version":"24.04","id":"noble","libc":null},"system":{"name":null,"release":null},"cpu":null,"openssl_version":null,"setuptools_version":null,"rustc_version":null,"ci":true}
File hashes
| Algorithm | Hash digest | |
|---|---|---|
| SHA256 |
6444431fd49d0803f5e62edeb5ed311d06ef232e8a7b8ab9dea8b221edf73470
|
|
| MD5 |
3506430f048e505ca30dce7d7ed0537e
|
|
| BLAKE2b-256 |
bf4fed238079b59c15fb4a1a5a301079dcd51d199c6e8bb41428dd43e56e634a
|
File details
Details for the file jax_image_models-0.4.0-py3-none-any.whl.
File metadata
- Download URL: jax_image_models-0.4.0-py3-none-any.whl
- Upload date:
- Size: 60.6 kB
- Tags: Python 3
- Uploaded using Trusted Publishing? Yes
- Uploaded via: uv/0.11.21 {"installer":{"name":"uv","version":"0.11.21","subcommand":["publish"]},"python":null,"implementation":{"name":null,"version":null},"distro":{"name":"Ubuntu","version":"24.04","id":"noble","libc":null},"system":{"name":null,"release":null},"cpu":null,"openssl_version":null,"setuptools_version":null,"rustc_version":null,"ci":true}
File hashes
| Algorithm | Hash digest | |
|---|---|---|
| SHA256 |
597dfdb4276350b143c080b2a8cd03d18dd9b97a1ebfbe15be3a785adbd489c1
|
|
| MD5 |
3c2d62d7a49ce779c4957e3ac3bc0d07
|
|
| BLAKE2b-256 |
3d6b04c06b6b1f68905b8267deb24979430a076de80301e4906acf8ecfb595b0
|
