initial commit

Author: gaarutyunov

.dockerignore

@@ -0,0 +1,58 @@
+# Git files
+.git
+.gitignore
+.github
+
+# Python cache
+__pycache__
+*.py[cod]
+*$py.class
+*.so
+.pytest_cache
+
+# Virtual environments
+venv/
+ENV/
+env/
+.venv
+
+# IDEs
+.vscode/
+.idea/
+*.swp
+*.swo
+*~
+
+# Build artifacts
+build/
+dist/
+*.egg-info/
+
+# Notebooks
+.ipynb_checkpoints
+
+# Test outputs
+htmlcov/
+.coverage
+.tox/
+
+# Training outputs
+output_*/
+wandb/
+logs/
+checkpoints/
+*.log
+
+# Large files (weights should be downloaded in container)
+weights/
+videos/
+data/
+datasets/
+
+# Temporary files
+tmp/
+temp/
+*.tmp
+
+# macOS
+.DS_Store

.github/workflows/test.yml

@@ -0,0 +1,84 @@
+name: Model Comparison Tests
+
+on:
+  push:
+    branches: [ main, develop ]
+  pull_request:
+    branches: [ main, develop ]
+  workflow_dispatch:
+
+jobs:
+  test:
+    runs-on: ubuntu-latest
+    
+    strategy:
+      matrix:
+        python-version: ["3.11", "3.12"]
+    
+    steps:
+    - uses: actions/checkout@v4
+    
+    - name: Set up Python ${{ matrix.python-version }}
+      uses: actions/setup-python@v4
+      with:
+        python-version: ${{ matrix.python-version }}
+    
+    - name: Cache pip packages
+      uses: actions/cache@v3
+      with:
+        path: ~/.cache/pip
+        key: ${{ runner.os }}-pip-${{ hashFiles('**/requirements*.txt') }}
+        restore-keys: |
+          ${{ runner.os }}-pip-
+    
+    - name: Cache model weights
+      uses: actions/cache@v3
+      with:
+        path: weights
+        key: ${{ runner.os }}-model-weights-vitl-v1
+        restore-keys: |
+          ${{ runner.os }}-model-weights-
+    
+    - name: Cache MLX build
+      uses: actions/cache@v3
+      with:
+        path: ../mlx
+        key: ${{ runner.os }}-mlx-${{ hashFiles('**/requirements.txt') }}
+        restore-keys: |
+          ${{ runner.os }}-mlx-
+    
+    - name: Install system dependencies
+      run: |
+        sudo apt-get update
+        sudo apt-get install -y cmake ninja-build libblas-dev liblapack-dev liblapacke-dev libopenblas-dev gfortran
+    
+    - name: Build and install MLX from source
+      run: |
+        cd ..
+        if [ ! -d "mlx" ]; then
+          git clone https://github.com/ml-explore/mlx.git
+        fi
+        cd mlx
+        pip install .
+        cd ../vjepa2-mlx
+    
+    - name: Install dependencies
+      run: |
+        python -m pip install --upgrade pip
+        pip install -r requirements.txt
+        pip install -r requirements-test.txt
+        pip install -e .
+    
+    - name: Clone and install vjepa2 repository
+      run: |
+        cd ..
+        if [ ! -d "vjepa2" ]; then
+          git clone https://github.com/facebookresearch/vjepa2.git
+        fi
+        cd vjepa2
+        pip install torch torchvision --index-url https://download.pytorch.org/whl/cpu
+        pip install -e .
+    
+    - name: Run tests
+      run: |
+        pytest tests/test_model_comparison.py -v -s --tb=short

.gitignore

@@ -0,0 +1,93 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+pip-wheel-metadata/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+*.manifest
+*.spec
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+*.ipynb_checkpoints/
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+.python-version
+
+# Virtual environments
+venv/
+ENV/
+env/
+.venv
+
+# IDEs
+.vscode/
+.idea/
+*.swp
+*.swo
+*~
+.DS_Store
+
+# MLX specific
+*.safetensors
+*.npz
+
+# Training outputs
+output_*/
+wandb/
+logs/
+checkpoints/
+*.log
+
+# Data directories
+data/
+datasets/
+videos/
+weights/
+
+# Temporary files
+tmp/
+temp/
+*.tmp

CHANGELOG.md

@@ -0,0 +1,41 @@
+# Changelog
+
+All notable changes to the V-JEPA 2 MLX project will be documented in this file.
+
+The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
+and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
+
+## [Unreleased]
+
+## [0.0.1] - 2025-11-12
+
+### Added
+- Initial MLX implementation of V-JEPA 2 models
+- Vision Transformer (ViT) with 3D patch embedding for video input
+- Attentive classifier for video action recognition
+- SSv2 classifier training script with frozen encoder
+- Training configuration system with YAML support
+- Interactive Jupyter notebooks for demonstrations
+- Pretrained weight conversion utilities
+- Comprehensive documentation and installation guide
+- Support for gradient accumulation for memory efficiency
+- Weights & Biases integration for experiment tracking
+- Training scripts with sensible defaults
+
+### Features
+- MLX-optimized models for Apple Silicon (M1/M2/M3)
+- Frozen encoder fine-tuning for efficient transfer learning
+- Something-Something-v2 action recognition training pipeline
+- Configurable training hyperparameters
+- Checkpoint saving and resuming
+- Learning rate scheduling with cosine decay and warmup
+- Top-1 and Top-5 accuracy metrics
+
+### Documentation
+- README with quick start guide
+- Detailed installation instructions (INSTALL.md)
+- Training configuration examples
+- Notebook tutorials for model usage
+
+[Unreleased]: https://github.com/facebookresearch/vjepa2-mlx/compare/v0.0.1...HEAD
+[0.0.1]: https://github.com/facebookresearch/vjepa2-mlx/releases/tag/v0.0.1