fix vjepa code

Author: aparnabg

src/sailsprep/action_model_testing/vjepa/clips_fixed_length/vjepa_clip_level_ablation.py

@@ -235,6 +235,9 @@ def _load_clip(self, s):
         total_vid = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
 
         bbox_map  = load_bbox_map(s["h5_path"])
+        if not bbox_map:
+            cap.release()
+            raise ValueError("empty bbox map")
         bbox_keys = np.array(sorted(bbox_map.keys()))
 
         ann_frames = np.arange(s["start_frame"], s["end_frame"] + 1)
@@ -423,8 +426,7 @@ def extract_features(encoder, samples, label_map, device, batch_size=4):
             all_labels.append(labels)
         except Exception as e:
             print(f"    batch {i} error: {e}")
-            all_feats.append(torch.zeros(clips.shape[0], 1, EMBED_DIM))
-            all_labels.append(labels)
+            raise
 
     return torch.cat(all_feats, 0), torch.cat(all_labels, 0)
 

src/sailsprep/action_model_testing/vjepa/clips_without_coi_crop/extract_features.py

@@ -121,7 +121,7 @@ def extract_all_features(model, processor, video_paths, labels, device):
             all_labels.append(batch_labels)
         except Exception as e:
             print(f"  Error at batch {batch_idx}: {e}")
-            errors.append(batch_idx)
+            raise
 
     all_features = torch.cat(all_features, dim=0)
     all_labels   = torch.cat(
@@ -152,9 +152,10 @@ def main():
 
     # Load encoder
     processor = AutoVideoProcessor.from_pretrained(HF_MODEL_NAME)
+    dtype = torch.float16 if torch.cuda.is_available() else torch.float32
     encoder   = AutoModel.from_pretrained(
         HF_MODEL_NAME,
-        torch_dtype=torch.float16,
+        torch_dtype=dtype,
         attn_implementation="sdpa",
     )
     encoder.eval()

src/sailsprep/action_model_testing/vjepa/clips_without_coi_crop/rmm/extract_features.py

@@ -124,7 +124,7 @@ def extract_all_features(model, processor, video_paths, labels, device):
             all_labels.append(batch_labels)
         except Exception as e:
             print(f"  Error at batch {batch_idx}: {e}", flush=True)
-            errors.append(batch_idx)
+            raise
 
     all_features = torch.cat(all_features, dim=0)
     all_labels   = torch.cat(
@@ -155,9 +155,10 @@ def main():
 
     # Load encoder
     processor = AutoVideoProcessor.from_pretrained(HF_MODEL_NAME)
+    dtype = torch.float16 if torch.cuda.is_available() else torch.float32
     encoder   = AutoModel.from_pretrained(
         HF_MODEL_NAME,
-        torch_dtype=torch.float16,
+        torch_dtype=dtype,
         attn_implementation="sdpa",
     )
     encoder.eval()

src/sailsprep/action_model_testing/vjepa/coi_crop/finetune_vjepa2_h5bbox.py

@@ -202,8 +202,9 @@ def __init__(self, num_classes=NUM_CLASSES, full_finetune=False, class_weights=N
         self.save_hyperparameters(ignore=["class_weights"])
         self.full_finetune = full_finetune
 
+        dtype = torch.float16 if torch.cuda.is_available() else torch.float32
         self.encoder = AutoModel.from_pretrained(
-            MODEL_NAME, torch_dtype=torch.float16, attn_implementation="sdpa",
+            MODEL_NAME, torch_dtype=dtype, attn_implementation="sdpa",
         )
         if not full_finetune:
             for p in self.encoder.parameters(): p.requires_grad = False
@@ -230,8 +231,8 @@ def forward(self, inputs):
 
     def _step(self, batch, stage):
         inputs, labels = batch
-        # Cast float inputs to half to match encoder
-        inputs = {k: (v.half() if v.is_floating_point() else v) for k, v in inputs.items()}
+        # Cast float inputs to match encoder dtype
+        inputs = {k: (v.to(self.encoder.dtype) if v.is_floating_point() else v) for k, v in inputs.items()}
         logits = self(inputs)
         loss = F.cross_entropy(logits, labels,
                                weight=self.class_weights if stage == "train" else None)
@@ -324,7 +325,7 @@ def run_inference(model, test_samples, label_map, processor, device):
         try:
             frames, label = ds[i]
             inputs = processor([frames], return_tensors="pt")
-            inputs = {k: (v.half().to(device) if v.is_floating_point() else v.to(device))
+            inputs = {k: (v.to(device=device, dtype=model.encoder.dtype) if v.is_floating_point() else v.to(device))
                       for k, v in inputs.items()}
             with torch.no_grad():
                 logits = model(inputs)

src/sailsprep/action_model_testing/vjepa/full_video/rerun_window_inference.py

@@ -29,7 +29,7 @@
 SPLITS_CSV    = "/home/aparnabg/orcd/scratch/latest_split_csv_new.csv"
 BASE_DIR      = "/orcd/data/satra/002/projects/SAILS/vjepa_features/models_output_seeds/vjepa/"
 
-EMBED_DIM     = 1408
+EMBED_DIM     = 1024
 WINDOW_FRAMES = 30
 STRIDE_FRAMES = 15
 BATCH_SIZE    = 256

src/sailsprep/action_model_testing/vjepa/full_video/train_probe_framelevel.py

@@ -28,7 +28,7 @@
 SPLITS_CSV    = "/home/aparnabg/orcd/scratch/latest_split_csv_new.csv"
 OUTPUT_BASE   = "/orcd/data/satra/002/projects/SAILS/vjepa_features/models_output_seeds/vjepa/framelevel_onest/"
 
-EMBED_DIM     = 1408
+EMBED_DIM     = 1024
 CONTEXT       = 5          # ±5 frames → window of 11
 WINDOW_SIZE   = 2 * CONTEXT + 1   # 11
 BATCH_SIZE    = 256        # larger batch fine — frames are cheap

src/sailsprep/action_model_testing/vjepa/full_video/train_probe_framelevel_hierarchical.py

@@ -30,7 +30,7 @@
 SPLITS_CSV    = "/home/aparnabg/orcd/scratch/latest_split_csv_new.csv"
 OUTPUT_BASE   = "/orcd/data/satra/002/projects/SAILS/vjepa_features/models_output_seeds/vjepa/framelevel_hierarchical/"
 
-EMBED_DIM     = 1408
+EMBED_DIM     = 1024
 CONTEXT       = 5
 WINDOW_SIZE   = 2 * CONTEXT + 1    # 11
 BATCH_SIZE    = 256

src/sailsprep/action_model_testing/vjepa/full_video/two_stage.py

@@ -29,7 +29,7 @@
 SPLITS_CSV    = "/home/aparnabg/orcd/scratch/latest_split_csv_new.csv"
 OUTPUT_BASE   = "/orcd/data/satra/002/projects/SAILS/vjepa_features/models_output_seeds/vjepa/window_two_st/"
 
-EMBED_DIM     = 1408
+EMBED_DIM     = 1024
 WINDOW_FRAMES = 30         # 2s at 15fps
 STRIDE_FRAMES = 15         # 1s stride
 BATCH_SIZE    = 64