ci: add Ruff config, pre-commit hooks, and GitHub Actions lint workflow

.github/workflows/lint.yml

@@ -0,0 +1,16 @@
+name: Lint
+
+on:
+  push:
+    branches: [main]
+  pull_request:
+
+jobs:
+  lint:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - uses: astral-sh/setup-uv@v5
+      - run: uv python install 3.10
+      - run: uv pip install ruff
+      - run: uv run ruff check .

.pre-commit-config.yaml

@@ -0,0 +1,6 @@
+repos:
+  - repo: https://github.com/astral-sh/ruff-pre-commit
+    rev: v0.12.2
+    hooks:
+      - id: ruff
+        args: [--fix]

app/vjepa_2_1/models/predictor.py

@@ -10,11 +10,10 @@
 import torch
 import torch.nn as nn
 
+from app.vjepa_2_1.models.utils.modules import Block
 from src.masks.utils import apply_masks
 from src.utils.tensors import repeat_interleave_batch, trunc_normal_
 
-from app.vjepa_2_1.models.utils.modules import Block
-
 
 class VisionTransformerPredictor(nn.Module):
     """Vision Transformer Predictor"""

app/vjepa_2_1/models/utils/masks_dist.py

@@ -1,6 +1,4 @@
 import torch
-import torch.nn as nn
-import torchvision
 
 
 def _get_frame_pos(ids, H_patches=None, W_patches=None, grid_size=None):

app/vjepa_2_1/models/utils/modules.py

@@ -7,7 +7,6 @@
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-
 from timm.models.layers import drop_path
 
 

app/vjepa_2_1/models/utils/patch_embed.py

@@ -5,9 +5,8 @@
 # LICENSE file in the root directory of this source tree.
 #
 
-from einops import rearrange
-
 import torch.nn as nn
+from einops import rearrange
 
 
 class AudioPatchEmbed(nn.Module):

app/vjepa_2_1/models/vision_transformer.py

@@ -10,11 +10,10 @@
 import torch
 import torch.nn as nn
 
-from src.masks.utils import apply_masks
-from src.utils.tensors import trunc_normal_
-
 from app.vjepa_2_1.models.utils.modules import Block
 from app.vjepa_2_1.models.utils.patch_embed import PatchEmbed, PatchEmbed3D
+from src.masks.utils import apply_masks
+from src.utils.tensors import trunc_normal_
 
 
 class VisionTransformer(nn.Module):

app/vjepa_2_1/train.py

@@ -21,6 +21,8 @@
 import torch
 import torch.multiprocessing as mp
 import torch.nn.functional as F
+from torch.nn.parallel import DistributedDataParallel
+
 from app.vjepa_2_1.models.utils.masks_dist import compute_mask_distance
 from app.vjepa_2_1.models.utils.modules import Lambda_LinearWarmupHold
 from app.vjepa_2_1.transforms import make_transforms
@@ -35,8 +37,6 @@
 from src.masks.utils import apply_masks
 from src.utils.distributed import init_distributed
 from src.utils.logging import AverageMeter, CSVLogger, get_logger, gpu_timer
-from torch.nn.parallel import DistributedDataParallel
-
 
 log_timings = True
 log_freq = 10

app/vjepa_2_1/transforms.py

@@ -5,11 +5,11 @@
 # LICENSE file in the root directory of this source tree.
 
 import numpy as np
-
-import src.datasets.utils.video.transforms as video_transforms
 import torch
 import torchvision.transforms as transforms
 from PIL import Image
+
+import src.datasets.utils.video.transforms as video_transforms
 from src.datasets.utils.video.randerase import RandomErasing
 
 

app/vjepa_2_1/utils.py

@@ -7,11 +7,12 @@
 import logging
 import sys
 
-import app.vjepa_2_1.models.predictor as vit_pred
-import app.vjepa_2_1.models.vision_transformer as video_vit
 import torch
 import torch.nn.functional as F
 import yaml
+
+import app.vjepa_2_1.models.predictor as vit_pred
+import app.vjepa_2_1.models.vision_transformer as video_vit
 from app.vjepa_2_1.wrappers import MultiSeqWrapper, PredictorMultiSeqWrapper
 from src.utils.checkpoint_loader import robust_checkpoint_loader
 from src.utils.schedulers import (

evals/image_classification_frozen/eval.py

@@ -29,7 +29,6 @@
 from evals.image_classification_frozen.models import init_module
 from src.datasets.data_manager import init_data
 from src.models.attentive_pooler import AttentiveClassifier
-from src.models.utils.modules import Block, CrossAttentionBlock
 from src.utils.checkpoint_loader import robust_checkpoint_loader
 from src.utils.distributed import AllReduce, init_distributed
 from src.utils.logging import AverageMeter, CSVLogger

hubconf.py

@@ -3,17 +3,5 @@
 # This source code is licensed under the MIT license found in the
 # LICENSE file in the root directory of this source tree.
 
-from evals.hub.preprocessor import vjepa2_preprocessor
-from src.hub.backbones import (
-    vjepa2_ac_vit_giant,
-    vjepa2_vit_giant,
-    vjepa2_vit_giant_384,
-    vjepa2_vit_huge,
-    vjepa2_vit_large,
-    vjepa2_1_vit_base_384,
-    vjepa2_1_vit_large_384,
-    vjepa2_1_vit_giant_384,
-    vjepa2_1_vit_gigantic_384,
-)
 
 dependencies = ["torch", "timm", "einops"]

pyproject.toml

@@ -4,3 +4,14 @@ line_length=119
 
 [tool.black]
 line-length = 119
+
+[tool.ruff]
+line-length = 119
+extend-exclude = ["*.ipynb"]
+
+[tool.ruff.lint]
+select = ["E", "F", "I"]
+ignore = ["E501", "E402", "E731", "E722", "E741", "E721", "E701", "F841", "F403", "F405"]
+
+[tool.ruff.lint.per-file-ignores]
+"__init__.py" = ["F401"]

src/datasets/imagenet1k.py

@@ -4,8 +4,6 @@
 # LICENSE file in the root directory of this source tree.
 
 import os
-import subprocess
-import time
 from logging import getLogger
 
 import numpy as np

src/datasets/utils/utils.py

@@ -4,6 +4,7 @@
 # LICENSE file in the root directory of this source tree.
 
 from src.utils.cluster import dataset_paths
+
 from src.utils.logging import get_logger
 
 logger = get_logger("Datasets utils")

src/datasets/utils/weighted_sampler.py

@@ -8,9 +8,9 @@
 
 import numpy as np
 import torch
+from torch.utils.data import DistributedSampler, RandomSampler
 
 from src.utils.logging import get_logger
-from torch.utils.data import DistributedSampler, RandomSampler
 
 logger = get_logger("WeightedSampler")
 

src/datasets/video_dataset.py

@@ -13,7 +13,7 @@
 import pandas as pd
 import torch
 import torchvision
-from decord import cpu, VideoReader
+from decord import VideoReader, cpu
 
 from src.datasets.utils.dataloader import (
     ConcatIndices,

src/hub/backbones.py

@@ -46,6 +46,8 @@ def _make_vjepa2_ac_model(
 ):
     from ..models import (
         ac_predictor as vit_ac_predictor,
+    )
+    from ..models import (
         vision_transformer as vit_encoder,
     )
 
@@ -101,7 +103,8 @@ def _make_vjepa2_model(
     pretrained: bool = True,
     **kwargs,
 ):
-    from ..models import predictor as vit_predictor, vision_transformer as vit_encoder
+    from ..models import predictor as vit_predictor
+    from ..models import vision_transformer as vit_encoder
 
     vit_encoder_kwargs = dict(
         patch_size=patch_size,
@@ -224,7 +227,8 @@ def _make_vjepa2_1_model(
     pretrained: bool = True,
     **kwargs,
 ):
-    from app.vjepa_2_1.models import predictor as vit_predictor, vision_transformer as vit_encoder
+    from app.vjepa_2_1.models import predictor as vit_predictor
+    from app.vjepa_2_1.models import vision_transformer as vit_encoder
 
     vit_encoder_kwargs = dict(
         patch_size=patch_size,

src/models/utils/patch_embed.py

@@ -4,7 +4,6 @@
 # LICENSE file in the root directory of this source tree.
 
 import torch.nn as nn
-from einops import rearrange
 
 
 class PatchEmbed(nn.Module):