Skip to content

VJEPA2 model, classifier works well! need some help with fm though #121

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
jenriver merged 30 commits into from
Jan 16, 2026
Merged

VJEPA2 model, classifier works well! need some help with fm though #121

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
30 commits
Select commit Hold shift + click to select a range
de7432f
adding readme, modeling and params files for dinov3
coder0143 Dec 11, 2025
0f5f9d4
added run and testing
coder0143 Dec 11, 2025
2c24d24
added readme and removed some basic imports
coder0143 Dec 11, 2025
4507250
removed ignore statement for ide
coder0143 Dec 11, 2025
078e933
Added colab link in markdown
coder0143 Dec 11, 2025
6ec2690
Merge branch 'main' into main
coder0143 Dec 12, 2025
5b1a81e
changed config and made changes wrt it, added all outputs, removed co…
coder0143 Dec 16, 2025
d818a67
adding more tests and some final changes
coder0143 Dec 16, 2025
0273b54
final changes to params
coder0143 Dec 16, 2025
a2e0241
fixed ruff based text formatting
coder0143 Dec 17, 2025
5c3d6c6
using random init torch model in test_outputs and removed run_model d…
coder0143 Dec 17, 2025
b0bb52f
using hf_hub constants for path and pre-commit done for ruff
coder0143 Dec 18, 2025
23d3c0a
updating tests
coder0143 Dec 20, 2025
ba596c5
rename test_outputs to be model specific
jenriver Dec 22, 2025
e410e74
Merge branch 'main' into main
jenriver Dec 22, 2025
e0d33b1
inital commit for qwen3vl
coder0143 Jan 1, 2026
039881c
removed initial commit
coder0143 Jan 1, 2026
0b8e32e
vjepa2 base fm, classifier and params done
coder0143 Jan 6, 2026
0e1b4ca
fixed params and testing for model porting
coder0143 Jan 6, 2026
259c8a2
updating for additional testing and stability
coder0143 Jan 6, 2026
542350b
more testing, ruff, classifier works well
coder0143 Jan 6, 2026
48592d1
Merge branch 'main' into main
coder0143 Jan 6, 2026
b66d3dd
Merge branch 'jax-ml:main' into main
coder0143 Jan 8, 2026
63a53ea
fixed foundation model, reformatted modeling.py, fixed testing
coder0143 Jan 9, 2026
ff2a152
adding readme
coder0143 Jan 9, 2026
8da697d
name in readme 😅
coder0143 Jan 9, 2026
41f345d
Merge branch 'jax-ml:main' into main
coder0143 Jan 13, 2026
41120be
removed unnecessary self.config = config
coder0143 Jan 14, 2026
cfcd29c
changes and fixes
coder0143 Jan 15, 2026
bb30d5d
Merge branch 'main' into main
jenriver Jan 16, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
35 changes: 35 additions & 0 deletions bonsai/models/vjepa2/Readme.md
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,35 @@
# VJEPA-2 in JAX

This directory contains a pure JAX implementation of the [VJEPA-2 foundation world model](https://ai.meta.com/vjepa/), using the [Flax NNX](https://flax.readthedocs.io/en/v0.8.3/experimental/nnx/index.html) API.


## Model Configuration Support Status

| Model Name | Size | Config Support Status |
| :--- | :--- | :--- |
| **Foundation models** | | |
| [vitl-fpc64-256](https://huggingface.co/facebook/vjepa2-vitl-fpc64-256) | 0.3B | **✅ Supported** |
| [vith-fpc64-256](https://huggingface.co/facebook/vjepa2-vith-fpc64-256) | 0.7B | **✅ Supported** |
| [vitg-fpc64-256](https://huggingface.co/facebook/vjepa2-vitg-fpc64-256) | 1B | **✅ Supported** |
| [vitg-fpc64-384](https://huggingface.co/facebook/vjepa2-vitg-fpc64-384) | 1B | **✅ Supported** |
| **Video Classifiers** | | |
| [vitl-fpc16-256-ssv2](https://huggingface.co/facebook/vjepa2-vitl-fpc16-256-ssv2) | 0.4B | **✅ Supported** |
| [vitg-fpc64-384-ssv2](https://huggingface.co/facebook/vjepa2-vitg-fpc64-384-ssv2) | 1B | **✅ Supported** |
| [vitl-fpc32-256-diving48](https://huggingface.co/facebook/vjepa2-vitl-fpc32-256-diving48) | 0.4B | **✅ Supported** |
| [vitg-fpc32-384-diving48](https://huggingface.co/facebook/vjepa2-vitg-fpc32-384-diving48) | 1B | **✅ Supported** |


### Running this model

Run VJEPA-2 in action, implemented in [500 lines of code](bonsai/models/vjepa2/modeling.py) in JAX.

```sh
python3 -m bonsai.models.vjepa2.tests.run_model
```


## How to contribute to this model

We welcome contributions! You can contribute to this model via the following:
* Add a model config variant from the above `🟡 Not started` to `class ModelConfig` in [modeling.py](modeling.py). Make sure your code is runnable on at least one hardware before creating a PR.
* Got some hardware? Run [run_model.py](tests/run_model.py) the existing configs above on hardwares marked `❔ Needs check`. Mark as `✅ Runs` or `⛔️ Not supported`.
Loading