https://github.com/pythoncrazy/jimm
JAX Image Modeling of Models contains Computer Vision/Vision Language Model implementations in native flax nnx with proper sharding annotations (allowing you to do easy fully sharded data parallel training with native flax nnx/jax)
https://github.com/pythoncrazy/jimm
computer-vision flax jax machine-learning
Last synced: 6 months ago
JSON representation
JAX Image Modeling of Models contains Computer Vision/Vision Language Model implementations in native flax nnx with proper sharding annotations (allowing you to do easy fully sharded data parallel training with native flax nnx/jax)
- Host: GitHub
- URL: https://github.com/pythoncrazy/jimm
- Owner: pythoncrazy
- License: mit
- Created: 2025-06-05T07:11:31.000Z (about 1 year ago)
- Default Branch: master
- Last Pushed: 2026-01-19T21:19:50.000Z (6 months ago)
- Last Synced: 2026-01-20T01:05:54.950Z (6 months ago)
- Topics: computer-vision, flax, jax, machine-learning
- Language: Python
- Homepage: https://pythoncrazy.github.io/jimm/
- Size: 21.8 MB
- Stars: 5
- Watchers: 0
- Forks: 0
- Open Issues: 5
-
Metadata Files:
- Readme: README.md
- License: LICENSE
Awesome Lists containing this project
README
# Jax Image Modeling of Models (jimm)
Docs are at: [https://pythoncrazy.github.io/jimm](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 pixi.sh:
`pixi add jimm@https://github.com/pythoncrazy/jimm.git --pypi`
### Using uv
`uv add --dev git+https://github.com/pythoncrazy/jimm.git`
or if you prefer to not add as a direct dependency:
`uv pip install git+https://github.com/pythoncrazy/jimm.git`
### Using pip/conda
`pip install git+https://github.com/pythoncrazy/jimm.git`