Update TARA to latest Tarsier2 checkpoint and runnable demo.

Replace weights and tokenizer/config artifacts, refresh README/demo/runtime code for the new stack, and ensure large files (including tokenizer.json) are tracked via LFS for Hub compatibility.

Made-with: Cursor

.gitattributes

@@ -6,6 +6,7 @@
 *.ftz filter=lfs diff=lfs merge=lfs -text
 *.gz filter=lfs diff=lfs merge=lfs -text
 *.h5 filter=lfs diff=lfs merge=lfs -text
+*.jar filter=lfs diff=lfs merge=lfs -text
 *.joblib filter=lfs diff=lfs merge=lfs -text
 *.lfs.* filter=lfs diff=lfs merge=lfs -text
 *.mlmodel filter=lfs diff=lfs merge=lfs -text
@@ -33,6 +34,7 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+tokenizer.json filter=lfs diff=lfs merge=lfs -text
 assets/cat.png filter=lfs diff=lfs merge=lfs -text
 assets/dog+cat.png filter=lfs diff=lfs merge=lfs -text
 assets/source-27375787.mp4 filter=lfs diff=lfs merge=lfs -text

README.md

@@ -1,27 +1,56 @@
----
-license: apache-2.0
-datasets:
-- sentence-transformers/all-nli
-language:
-- en
-metrics:
-- accuracy
-base_model:
-- omni-research/Tarsier-7b
-tags:
-- video-retrieval
-- text-to-video-retrieval
-- time-awareness
-- video-models
----
-
-# ![](assets/tara-logo.png) TARA: Time-Aware Retrieval Adaptation for Video Understanding
+# ![](assets/tara-logo.png) TARA: *T*ext *A*dapted *R*etrieval *A*lignment for Nuanced Video Retrieval
 <!-- # <img src="./assets/logo.png" width="24"> TARA: Time-Aware Retrieval Adaptation for Video Understanding -->
 
-TARA (Time-Aware Retrieval Adaptation) is a multimodal model for video and text understanding.
+This repository contains inference and evaluation code for the TARA model based on the paper:
+[Adapting MLLMs for Nuanced Video Retrieval](https://arxiv.org/abs/2512.13511)
+
+<p align="center">
+  <a href="https://bpiyush.github.io/tara-website/" target="_blank">
+    <img src="https://img.shields.io/badge/Project-Page-blue" alt="Project Page">
+  </a>
+  &nbsp;&nbsp;&nbsp;
+  <a href="https://github.com/bpiyush/TARA" target="_blank">
+    <img src="https://img.shields.io/badge/GitHub-Code-black?logo=github" alt="GitHub Code">
+  </a>
+  &nbsp;&nbsp;&nbsp;
+  <a href="https://arxiv.org/abs/2512.13511" target="_blank">
+    <img src="https://img.shields.io/badge/arXiv-Paper-b31b1b?logo=arxiv&logoColor=white" alt="arXiv">
+  </a>
+  &nbsp;&nbsp;&nbsp;
+  <a href="https://huggingface.co/datasets/bpiyush/chirality-in-action" target="_blank">
+    <img src="https://huggingface.co/datasets/huggingface/badges/resolve/main/dataset-on-hf-md-dark.svg" alt="Dataset on Hugging Face">
+  </a>
+</p>
+
+<!-- Show arch fig in 75% of the screen and center it -->
+<p align="center">
+  <img src="./assets/arch.png" width="75%" alt="TARA architecture">
+</p>
+<!-- Add a caption with small font size and center it such that it align with image width and center it-->
+<p style="text-align: left; font-size: 13px; width: 75%; display: block; margin: 0 auto;"><b>TARA Architecture:</b> We use EOL prompt to embed videos using an MLLM (Tarsier2-7B). We train the LLM weights with contrastive loss on carefully crafted hard-negatives to instill (i) temporal, (ii) negation and (iii) multimodal nuances in the embedding space.</p>
+
+
+<!-- Add a Table of Contents here -->
+## Table of Contents
+- [Installation & Setup](#installation--setup)
+- [Quick Start](#quick-start)
+- [Evaluation](#evaluation)
+  - [Data Preparation](#data-preparation)
+  - [Embedding Computation](#embedding-computation)
+  - [General evaluation: MMEB-V2 (Meng et al.)](#general-evaluation-mmeb-v2-meng-et-al)
+- [Citation](#citation)
+- [License](#license)
+
 
 ## Installation & Setup
 
+First, clone the repository:
+```bash
+git clone https://github.com/bpiyush/tara.git
+cd tara
+```
+
+
 ### 1. Install Git LFS (if not already installed)
 
 Git LFS is required to download the model weights.
@@ -43,9 +72,9 @@ Git LFS initialized.
 ```
 
 
-### 2. Clone the Repository
+### 2. Download the Model Weights
 ```bash
-git clone https://huggingface.co/bpiyush/TARA
+git clone https://huggingface.co/bpiyush/TARA /path/to/download/tara
 cd TARA
 ```
 
@@ -76,10 +105,37 @@ This will download all model weights (may take a few minutes depending on your c
 
 ## Quick Start
 
-See the script at [demo_usage.py](demo_usage.py) for a quick start. You can run it:
+TARA is primarily designed to encode videos and texts in a joint embedding space under an MLLM.
+
+```python
+import torch
+from modeling_tara import TARA
+
+model = TARA.from_pretrained(
+    "/path/to/download/tara",  # Load from current directory
+    device_map='auto',
+    torch_dtype=torch.bfloat16,
+)
+n_params = sum(p.numel() for p in model.model.parameters())
+print(f"Number of parameters: {round(n_params/1e9, 3)}B")
+
+# Embed a video
+video_path = "./assets/folding_paper.mp4"
+with torch.no_grad():
+    video_emb = model.encode_vision(video_path).cpu().squeeze(0).float()
+print(f"Video embedding shape: {video_emb.shape}")  # torch.Size([3584])
+
+# Embed a text
+text = ['someone is folding a paper', 'cutting a paper', 'someone is folding a paper']
+with torch.no_grad():
+    text_emb = model.encode_text(text).cpu().float()
+print(f"Text embedding shape: {text_emb.shape}")  # torch.Size([3, 3584])
+```
+
+For a more detailed demo, see the script at [demo_usage.py](demo_usage.py). You can run it:
 
 ```sh
-python demo_usage.py
+python demo_usage.py --model_path /path/to/download/tara
 ```
 The output should look something like this:
 
@@ -88,45 +144,50 @@ The output should look something like this:
 TARA Model Demo
 ============================================================
 
-[1/6] Loading model...
-[ MODEL ] Loading TARA from /work/piyush/pretrained_checkpoints/TARA/ [..............]
-### do_image_padding is set as False, images will be resized directly!
+[1/5] Loading model...
+The argument `trust_remote_code` is to be used with Auto classes. It has no effect here and is ignored.
+Unrecognized keys in `rope_scaling` for 'rope_type'='default': {'mrope_section'}
+The argument `trust_remote_code` is to be used with Auto classes. It has no effect here and is ignored.
 The model weights are not tied. Please use the `tie_weights` method before using the `infer_auto_device` function.
-Loading checkpoint shards: 100%|██████████████████████████████████████████████████████████████████████████████████████| 3/3 [00:03<00:00,  1.05s/it]
+Loading checkpoint shards: 100%|████████████████████████████████████████████████████████████████████████████████████████████████| 4/4 [00:03<00:00,  1.07it/s]
 ✓ Model loaded successfully!
-Number of parameters: 7.063B
+Number of parameters: 8.291B
 ----------------------------------------------------------------------------------------------------
 
-[2/6] Testing video encoding and captioning ...
+[2/5] Testing video encoding ...
+From v4.47 onwards, when a model cache is to be returned, `generate` will return a `Cache` instance instead by default (as opposed to the legacy tuple of tuples format). If you want to keep returning the legacy format, please set `return_legacy_cache=True`.
 ✓ Video encoded successfully!
-Video shape: torch.Size([1, 16, 3, 240, 426])
-Video embedding shape: torch.Size([4096])
-Video caption: A hand is seen folding a white paper on a gray carpeted floor. The paper is opened flat on the surface, and then the hand folds it in half vertically, creating a crease in the middle. The hand continues to fold the paper further, resulting in a smaller, more compact size. The background remains a consistent gray carpet throughout the video.
+Video embedding shape: torch.Size([3584])
 ----------------------------------------------------------------------------------------------------
 
-[3/6] Testing text encoding...
+[3/5] Testing text encoding...
+Setting `pad_token_id` to `eos_token_id`:None for open-end generation.
+Setting `pad_token_id` to `eos_token_id`:None for open-end generation.
+Setting `pad_token_id` to `eos_token_id`:None for open-end generation.
 ✓ Text encoded successfully!
 Text: ['someone is folding a paper', 'cutting a paper', 'someone is unfolding a paper']
-Text embedding shape: torch.Size([3, 4096])
+Text embedding shape: torch.Size([3, 3584])
 
-[4/6] Computing video-text similarities...
+[4/5] Computing video-text similarities...
 ✓ Similarities computed!
-  'someone is folding a paper': 0.5039
-  'cutting a paper': 0.3022
-  'someone is unfolding a paper': 0.3877
+  'someone is folding a paper': 0.6488
+  'cutting a paper': 0.3952
+  'someone is unfolding a paper': 0.3009
 ----------------------------------------------------------------------------------------------------
 
-[5/6] Testing negation example...
-Image embedding shape: torch.Size([2, 4096])
+[5/5] Testing negation example...
+Image embedding shape: torch.Size([2, 3584])
+Setting `pad_token_id` to `eos_token_id`:None for open-end generation.
 Text query:  ['an image of a cat but there is no dog in it']
-Text-Image similarity: tensor([[0.2585, 0.1449]])
+Text-Image similarity: tensor([[0.5169, 0.3659]])
 - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - 
+Setting `pad_token_id` to `eos_token_id`:None for open-end generation.
 Text query:  ['an image of a cat and a dog together']
-Text-Image similarity: tensor([[0.2815, 0.4399]])
+Text-Image similarity: tensor([[0.4364, 0.6004]])
 ----------------------------------------------------------------------------------------------------
 
-[6/6] Testing composed video retrieval...
-Source-Target similarity with edit: 0.6476313471794128
+[Bonus] Testing composed video retrieval...
+Source-Target similarity with edit: 0.757888674736023
 
 ============================================================
 Demo completed successfully! 🎉
@@ -134,45 +195,91 @@ Demo completed successfully! 🎉
 ```
 
 
-OR use the snippet below:
+## Evaluation
 
-```python
-import torch
-from modeling_tara import TARA, read_frames_decord
 
-model = TARA.from_pretrained(
-    ".",  # Load from current directory
-    device_map='auto',
-    torch_dtype=torch.bfloat16,
-)
-n_params = sum(p.numel() for p in model.model.parameters())
-print(f"Number of parameters: {round(n_params/1e9, 3)}B")
+### Data Preparation
 
-# Embed a video
-video_path = "./assets/folding_paper.mp4"
-video_tensor = read_frames_decord(video_path, num_frames=16)
-video_tensor = video_tensor.unsqueeze(0)
-video_tensor = video_tensor.to(model.model.device)
-with torch.no_grad():
-    video_emb = model.encode_vision(video_tensor).cpu().squeeze(0).float()
-print(f"Video shape: {video_tensor.shape}")  # torch.Size([1, 16, 3, 240, 426])
-print(f"Video embedding shape: {video_emb.shape}")  # torch.Size([4096])
+We release the nuanced video retrieval splits used in the dataset in [data/](data/) folder.
+For ease of use, we have combined all the data for (i) temporal, (ii) negation and (iii) multimodal
+nuance into a single file where each entry is a video/text/video-text/image, etc.
 
-# Embed a text
-text = ['someone is folding a paper', 'cutting a paper', 'someone is folding a paper']
-with torch.no_grad():
-    text_emb = model.encode_text(text).cpu().float()
-print(f"Text embedding shape: {text_emb.shape}")  # torch.Size([3, 4096])
+```sh
+data
+├── nuanced_retrieval_inputs-test.csv # List of examples to embed (video, text, composed video-text, etc.) for test set
+├── nuanced_retrieval_inputs-val.csv # List of examples to embed (video, text, composed video-text, etc.) for validation set
+├── nuanced_retrieval_labels-test.json # Labels for test set
+└── nuanced_retrieval_labels-val.json # Labels for validation set
 ```
 
+An example input row looks like this:
+```json
+{
+  'id': '138629', 
+  'value': '138629',
+  'nuance': 'time',
+  'source': 'cia-ssv2',
+  'modality': 'video',
+}
+```
+where `id`is the unique identified, `value` is actual value (e.g., for a text caption, the ID can be different and value stores the actual caption), `nuance` is the type of nuance, 
+`source` is the source of the example (e.g., `cia-ssv2` for SSv2), and `modality` is the modality of the example (e.g., `video` or `text`).
+
+
+The coresponding label looks like this:
+```json
+['12055391_1.0']
+```
+which denotes the `id` of the text associated with the video.
+
+Finally, set the right paths to the data directories in [evals/compute_embeddings.py](evals/compute_embeddings.py) 
+based on your local setup.
+
+### Embedding Computation
+
+First, you need to compute the embeddings for the entire dataset. You can do this by running the following script:
+
+```bash
+python evals/compute_embeddings.py \
+--model_path /path/to/download/tara \
+--csv_path ./data/nuanced_retrieval_inputs-val.csv \
+--model_name tara_7b
+```
+
+Then, run the script to compute retrieval metrics.
+
+```bash
+python evals/compute_metrics.py \
+--model_path /path/to/download/tara \
+--csv_path ./data/nuanced_retrieval_inputs-val.csv \
+--lab_path ./data/nuanced_retrieval_labels-val.json \
+--model_name tara_7b
+```
+
+### General evaluation: MMEB-V2 ([Meng et al.](https://arxiv.org/abs/2507.04590))
+
+We evaluate on the video classification and video retrieval tasks in MMEB-V2 to demonstrate the generalizability of TARA.
+
+TODO
+
 ## Citation
 
 If you use this model, please cite:
 ```bibtex
-@misc{tara2025,
-  title={TARA: Simple and Efficient Time Aware Retrieval Adaptation of MLLMs for Video Understanding},
+@article{tara2025,
+  title={Adapting MLLMs for Nuanced Video Retrieval},
   author={Piyush Bagad and Andrew Zisserman},
   year={2025}
+  journal={arXiv preprint arXiv:2512.13511}
+}
+```
+
+```bibtex
+@article{bagad2025chirality,
+  title={Chirality in Action: Time-Aware Video Representation Learning by Latent Straightening},
+  author={Bagad, Piyush and Zisserman, Andrew},
+  journal={arXiv preprint arXiv:2509.08502},
+  year={2025}
 }
 ```
 

added_tokens.json

@@ -1,3 +1,16 @@
 {
-  "<image>": 32000
+  "<|box_end|>": 151649,
+  "<|box_start|>": 151648,
+  "<|endoftext|>": 151643,
+  "<|im_end|>": 151645,
+  "<|im_start|>": 151644,
+  "<|image_pad|>": 151655,
+  "<|object_ref_end|>": 151647,
+  "<|object_ref_start|>": 151646,
+  "<|quad_end|>": 151651,
+  "<|quad_start|>": 151650,
+  "<|video_pad|>": 151656,
+  "<|vision_end|>": 151653,
+  "<|vision_pad|>": 151654,
+  "<|vision_start|>": 151652
 }

chat_template.json

@@ -0,0 +1,3 @@
+{
+  "chat_template": "{% set image_count = namespace(value=0) %}{% set video_count = namespace(value=0) %}{% for message in messages %}{% if loop.first and message['role'] != 'system' %}<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n{% endif %}{% if (message['role'] != 'assistant') %}<|im_start|>{{ message['role'] }}\n{% if message['content'] is string %}{{ message['content'] }}<|im_end|>\n{% else %}{% for content in message['content'] %}{% if content['type'] == 'image' %}{% set image_count.value = image_count.value + 1 %}{% if add_vision_id %}Picture {{ image_count.value }}: {% endif %}<|vision_start|>{{'<|image_pad|>' * content['num_vision_tokens']}}<|vision_end|>{% elif content['type'] == 'video' %}{% set video_count.value = video_count.value + 1 %}{% if add_vision_id %}Video {{ video_count.value }}: {% endif %}<|vision_start|>{{'<|image_pad|>' * content['num_vision_tokens']}}<|vision_end|>{% elif content['type'] == 'text' %}{{ content['text'] }}{% endif %}{% endfor %}<|im_end|>\n{% endif %}{% elif (message['role'] == 'assistant') %}<|im_start|>{{ message['role'] }}\n{% generation %}{% if message['content'] is string %}{{ message['content'] }}<|im_end|>\n{% else %}{% for content in message['content'] %}{% if content['type'] == 'image' %}{% set image_count.value = image_count.value + 1 %}{% if add_vision_id %}Picture {{ image_count.value }}: {% endif %}<|vision_start|>{{'<|image_pad|>' * content['num_vision_tokens']}}<|vision_end|>{% elif content['type'] == 'video' %}{% set video_count.value = video_count.value + 1 %}{% if add_vision_id %}Video {{ video_count.value }}: {% endif %}<|vision_start|>{{'<|image_pad|>' * content['num_vision_tokens']}}<|vision_end|>{% elif content['type'] == 'text' %}{{ content['text'] }}{% endif %}{% endfor %}{% if not loop.last or not strip_final_eos %}<|im_end|>\n{% endif %}{% endif %}{% endgeneration %}{% endif %}{% endfor %}{% if add_generation_prompt %}<|im_start|>assistant\n{% endif %}"
+}

config.json

@@ -1,27 +1,25 @@
 {
-  "_name_or_path": "/work/piyush/pretrained_checkpoints/Tarsier-7b",
-  "add_transformer_projector": false,
+  "_name_or_path": "/work/piyush/pretrained_checkpoints/Tarsier2-7b-0115",
   "architectures": [
-    "TarsierForConditionalGeneration"
+    "Tarsier2ForConditionalGeneration"
   ],
   "ignore_index": -100,
   "image_new_idx": 32003,
   "image_newline_idx": 32002,
   "image_seq_length": 576,
-  "image_token_index": 32000,
+  "image_token_index": 151655,
   "model_type": "llava",
+  "projection_head": null,
   "projector_hidden_act": "gelu",
   "text_config": {
-    "_name_or_path": "lmsys/vicuna-7b-v1.5",
+    "_attn_implementation_autoset": true,
+    "_name_or_path": "",
     "add_cross_attention": false,
-    "architectures": [
-      "LlamaForCausalLM"
-    ],
-    "attention_bias": false,
+    "architectures": null,
     "attention_dropout": 0.0,
     "bad_words_ids": null,
     "begin_suppress_tokens": null,
-    "bos_token_id": 1,
+    "bos_token_id": 151643,
     "chunk_size_feed_forward": 0,
     "cross_attention_hidden_size": null,
     "decoder_start_token_id": null,
@@ -29,20 +27,20 @@
     "do_sample": false,
     "early_stopping": false,
     "encoder_no_repeat_ngram_size": 0,
-    "eos_token_id": 2,
+    "eos_token_id": 151645,
     "exponential_decay_length_penalty": null,
     "finetuning_task": null,
     "forced_bos_token_id": null,
     "forced_eos_token_id": null,
-    "head_dim": 128,
     "hidden_act": "silu",
-    "hidden_size": 4096,
+    "hidden_size": 3584,
     "id2label": {
       "0": "LABEL_0",
       "1": "LABEL_1"
     },
+    "image_token_id": 151655,
     "initializer_range": 0.02,
-    "intermediate_size": 11008,
+    "intermediate_size": 18944,
     "is_decoder": false,
     "is_encoder_decoder": false,
     "label2id": {
@@ -51,33 +49,42 @@
     },
     "length_penalty": 1.0,
     "max_length": 20,
-    "max_position_embeddings": 4096,
+    "max_position_embeddings": 32768,
+    "max_window_layers": 28,
     "min_length": 0,
-    "mlp_bias": false,
-    "model_type": "llama",
+    "model_type": "qwen2_vl",
     "no_repeat_ngram_size": 0,
-    "num_attention_heads": 32,
+    "num_attention_heads": 28,
     "num_beam_groups": 1,
     "num_beams": 1,
-    "num_hidden_layers": 32,
-    "num_key_value_heads": 32,
+    "num_hidden_layers": 28,
+    "num_key_value_heads": 4,
     "num_return_sequences": 1,
     "output_attentions": false,
     "output_hidden_states": false,
     "output_scores": false,
-    "pad_token_id": 0,
+    "pad_token_id": null,
     "prefix": null,
-    "pretraining_tp": 1,
     "problem_type": null,
     "pruned_heads": {},
     "remove_invalid_values": false,
     "repetition_penalty": 1.0,
     "return_dict": true,
     "return_dict_in_generate": false,
-    "rms_norm_eps": 1e-05,
-    "rope_scaling": null,
-    "rope_theta": 10000.0,
+    "rms_norm_eps": 1e-06,
+    "rope_scaling": {
+      "mrope_section": [
+        16,
+        24,
+        24
+      ],
+      "rope_type": "default",
+      "type": "default"
+    },
+    "rope_theta": 1000000.0,
     "sep_token_id": null,
+    "sliding_window": 32768,
+    "spatial_merge_size": 2,
     "suppress_tokens": null,
     "task_specific_params": null,
     "temperature": 1.0,
@@ -87,32 +94,112 @@
     "tokenizer_class": null,
     "top_k": 50,
     "top_p": 1.0,
-    "torch_dtype": "float16",
+    "torch_dtype": "bfloat16",
     "torchscript": false,
     "typical_p": 1.0,
     "use_bfloat16": false,
-    "use_bpex_rmsnorm": true,
-    "use_bpex_rotary": true,
+    "use_bpex_rotary": false,
     "use_cache": false,
-    "use_ft_flash_attn": true,
-    "vocab_size": 32064
+    "use_sliding_window": false,
+    "video_token_id": 151656,
+    "vision_config": {
+      "_name_or_path": "",
+      "add_cross_attention": false,
+      "architectures": null,
+      "bad_words_ids": null,
+      "begin_suppress_tokens": null,
+      "bos_token_id": null,
+      "chunk_size_feed_forward": 0,
+      "cross_attention_hidden_size": null,
+      "decoder_start_token_id": null,
+      "depth": 32,
+      "diversity_penalty": 0.0,
+      "do_sample": false,
+      "early_stopping": false,
+      "embed_dim": 1280,
+      "encoder_no_repeat_ngram_size": 0,
+      "eos_token_id": null,
+      "exponential_decay_length_penalty": null,
+      "finetuning_task": null,
+      "forced_bos_token_id": null,
+      "forced_eos_token_id": null,
+      "hidden_act": "quick_gelu",
+      "hidden_size": 3584,
+      "id2label": {
+        "0": "LABEL_0",
+        "1": "LABEL_1"
+      },
+      "in_channels": 3,
+      "is_decoder": false,
+      "is_encoder_decoder": false,
+      "label2id": {
+        "LABEL_0": 0,
+        "LABEL_1": 1
+      },
+      "length_penalty": 1.0,
+      "max_length": 20,
+      "min_length": 0,
+      "mlp_ratio": 4,
+      "model_type": "qwen2_vl",
+      "no_repeat_ngram_size": 0,
+      "num_beam_groups": 1,
+      "num_beams": 1,
+      "num_heads": 16,
+      "num_return_sequences": 1,
+      "output_attentions": false,
+      "output_hidden_states": false,
+      "output_scores": false,
+      "pad_token_id": null,
+      "patch_size": 14,
+      "prefix": null,
+      "problem_type": null,
+      "pruned_heads": {},
+      "remove_invalid_values": false,
+      "repetition_penalty": 1.0,
+      "return_dict": true,
+      "return_dict_in_generate": false,
+      "sep_token_id": null,
+      "spatial_merge_size": 2,
+      "suppress_tokens": null,
+      "task_specific_params": null,
+      "temperature": 1.0,
+      "temporal_patch_size": 2,
+      "tf_legacy_loss": false,
+      "tie_encoder_decoder": false,
+      "tie_word_embeddings": true,
+      "tokenizer_class": null,
+      "top_k": 50,
+      "top_p": 1.0,
+      "torch_dtype": null,
+      "torchscript": false,
+      "typical_p": 1.0,
+      "use_bfloat16": false
+    },
+    "vision_end_token_id": 151653,
+    "vision_start_token_id": 151652,
+    "vision_token_id": 151654,
+    "vocab_size": 152064
   },
   "torch_dtype": "bfloat16",
-  "transformers_version": "4.45.2",
+  "transformers_version": "4.45.0",
   "vision_config": {
+    "_attn_implementation_autoset": false,
     "_name_or_path": "",
     "add_cross_attention": false,
     "architectures": null,
     "attention_dropout": 0.0,
+    "attn_implementation": "flash_attention_2",
     "bad_words_ids": null,
     "begin_suppress_tokens": null,
     "bos_token_id": null,
     "chunk_size_feed_forward": 0,
     "cross_attention_hidden_size": null,
     "decoder_start_token_id": null,
+    "depth": 32,
     "diversity_penalty": 0.0,
     "do_sample": false,
     "early_stopping": false,
+    "embed_dim": 1280,
     "encoder_no_repeat_ngram_size": 0,
     "eos_token_id": null,
     "exponential_decay_length_penalty": null,
@@ -120,32 +207,35 @@
     "forced_bos_token_id": null,
     "forced_eos_token_id": null,
     "hidden_act": "quick_gelu",
-    "hidden_size": 1024,
+    "hidden_size": 3584,
     "id2label": {
       "0": "LABEL_0",
       "1": "LABEL_1"
     },
-    "image_size": 336,
-    "initializer_factor": 1.0,
+    "in_channels": 3,
+    "in_chans": 3,
     "initializer_range": 0.02,
-    "intermediate_size": 4096,
+    "intermediate_size": 29568,
     "is_decoder": false,
     "is_encoder_decoder": false,
     "label2id": {
       "LABEL_0": 0,
       "LABEL_1": 1
     },
-    "layer_norm_eps": 1e-05,
     "length_penalty": 1.0,
     "max_length": 20,
+    "max_position_embeddings": 32768,
+    "max_window_layers": 80,
     "min_length": 0,
-    "model_type": "clip_vision_model",
+    "mlp_ratio": 4,
+    "model_type": "qwen2_vl",
     "no_repeat_ngram_size": 0,
-    "num_attention_heads": 16,
+    "num_attention_heads": 64,
     "num_beam_groups": 1,
     "num_beams": 1,
-    "num_channels": 3,
-    "num_hidden_layers": 24,
+    "num_heads": 16,
+    "num_hidden_layers": 80,
+    "num_key_value_heads": 8,
     "num_return_sequences": 1,
     "output_attentions": false,
     "output_hidden_states": false,
@@ -154,16 +244,22 @@
     "patch_size": 14,
     "prefix": null,
     "problem_type": null,
-    "projection_dim": 768,
     "pruned_heads": {},
     "remove_invalid_values": false,
     "repetition_penalty": 1.0,
     "return_dict": true,
     "return_dict_in_generate": false,
+    "rms_norm_eps": 1e-05,
+    "rope_scaling": null,
+    "rope_theta": 1000000.0,
     "sep_token_id": null,
+    "sliding_window": 4096,
+    "spatial_merge_size": 2,
+    "spatial_patch_size": 14,
     "suppress_tokens": null,
     "task_specific_params": null,
     "temperature": 1.0,
+    "temporal_patch_size": 2,
     "tf_legacy_loss": false,
     "tie_encoder_decoder": false,
     "tie_word_embeddings": true,
@@ -174,9 +270,83 @@
     "torchscript": false,
     "typical_p": 1.0,
     "use_bfloat16": false,
-    "vocab_size": 32000
+    "use_cache": true,
+    "use_sliding_window": false,
+    "vision_config": {
+      "_name_or_path": "",
+      "add_cross_attention": false,
+      "architectures": null,
+      "bad_words_ids": null,
+      "begin_suppress_tokens": null,
+      "bos_token_id": null,
+      "chunk_size_feed_forward": 0,
+      "cross_attention_hidden_size": null,
+      "decoder_start_token_id": null,
+      "depth": 32,
+      "diversity_penalty": 0.0,
+      "do_sample": false,
+      "early_stopping": false,
+      "embed_dim": 1280,
+      "encoder_no_repeat_ngram_size": 0,
+      "eos_token_id": null,
+      "exponential_decay_length_penalty": null,
+      "finetuning_task": null,
+      "forced_bos_token_id": null,
+      "forced_eos_token_id": null,
+      "hidden_act": "quick_gelu",
+      "hidden_size": 3584,
+      "id2label": {
+        "0": "LABEL_0",
+        "1": "LABEL_1"
+      },
+      "in_channels": 3,
+      "is_decoder": false,
+      "is_encoder_decoder": false,
+      "label2id": {
+        "LABEL_0": 0,
+        "LABEL_1": 1
+      },
+      "length_penalty": 1.0,
+      "max_length": 20,
+      "min_length": 0,
+      "mlp_ratio": 4,
+      "model_type": "qwen2_vl",
+      "no_repeat_ngram_size": 0,
+      "num_beam_groups": 1,
+      "num_beams": 1,
+      "num_heads": 16,
+      "num_return_sequences": 1,
+      "output_attentions": false,
+      "output_hidden_states": false,
+      "output_scores": false,
+      "pad_token_id": null,
+      "patch_size": 14,
+      "prefix": null,
+      "problem_type": null,
+      "pruned_heads": {},
+      "remove_invalid_values": false,
+      "repetition_penalty": 1.0,
+      "return_dict": true,
+      "return_dict_in_generate": false,
+      "sep_token_id": null,
+      "spatial_merge_size": 2,
+      "suppress_tokens": null,
+      "task_specific_params": null,
+      "temperature": 1.0,
+      "temporal_patch_size": 2,
+      "tf_legacy_loss": false,
+      "tie_encoder_decoder": false,
+      "tie_word_embeddings": true,
+      "tokenizer_class": null,
+      "top_k": 50,
+      "top_p": 1.0,
+      "torch_dtype": null,
+      "torchscript": false,
+      "typical_p": 1.0,
+      "use_bfloat16": false
+    },
+    "vocab_size": 152064
   },
   "vision_feature_layer": -2,
-  "vision_feature_select_strategy": "default",
-  "vocab_size": 32064
+  "vision_feature_select_strategy": "default"
 }

demo_usage.py

@@ -1,6 +1,6 @@
 import torch
 from termcolor import colored
-from modeling_tara import TARA, read_frames_decord, read_images_decord
+
 
 import warnings
 warnings.filterwarnings("ignore")
@@ -12,36 +12,28 @@ def main(model_path: str = "."):
     print(colored("="*60, 'yellow'))
     
     # Load model from current directory
-    print(colored("\n[1/6] Loading model...", 'cyan'))
+    print(colored("\n[1/5] Loading model...", 'cyan'))
     model = TARA.from_pretrained(
         model_path,  # Load from current directory
         device_map='auto',
         torch_dtype=torch.bfloat16,
+        attn_implementation='flash_attention_2',
     )
     
     n_params = sum(p.numel() for p in model.model.parameters())
-    print(colored(f"✓ Model loaded successfully!", 'green'))
+    print(colored("✓ Model loaded successfully!", 'green'))
     print(f"Number of parameters: {round(n_params/1e9, 3)}B")
     print("-" * 100)
     
     # Encode a sample video
-    print(colored("\n[2/6] Testing video encoding and captioning ...", 'cyan'))
+    print(colored("\n[2/5] Testing video encoding ...", 'cyan'))
     video_path = "./assets/folding_paper.mp4"
     try:
-        video_tensor = read_frames_decord(video_path, num_frames=16)
-        video_tensor = video_tensor.unsqueeze(0)
-        video_tensor = video_tensor.to(model.model.device)
-        
         with torch.no_grad():
-            video_emb = model.encode_vision(video_tensor).cpu().squeeze(0).float()
-            
-            # Get caption for the video
-            video_caption = model.describe(video_tensor)[0]
+            video_emb = model.encode_vision(video_path).cpu().squeeze(0).float()
         
         print(colored("✓ Video encoded successfully!", 'green'))
-        print(f"Video shape: {video_tensor.shape}")  # torch.Size([1, 16, 3, 240, 426])
         print(f"Video embedding shape: {video_emb.shape}")  # torch.Size([4096])
-        print(colored(f"Video caption: {video_caption}", 'magenta'))
     except FileNotFoundError:
         print(colored(f"⚠ Video file not found: {video_path}", 'red'))
         print(colored("  Please add a video file or update the path in demo_usage.py", 'yellow'))
@@ -49,7 +41,7 @@ def main(model_path: str = "."):
     print("-" * 100)
     
     # Encode sample texts
-    print(colored("\n[3/6] Testing text encoding...", 'cyan'))
+    print(colored("\n[3/5] Testing text encoding...", 'cyan'))
     text = ['someone is folding a paper', 'cutting a paper', 'someone is unfolding a paper']
     # NOTE: It can also take a single string
     
@@ -62,7 +54,7 @@ def main(model_path: str = "."):
     
     # Compute similarities if video was encoded
     if video_emb is not None:
-        print(colored("\n[4/6] Computing video-text similarities...", 'cyan'))
+        print(colored("\n[4/5] Computing video-text similarities...", 'cyan'))
         similarities = torch.cosine_similarity(
             video_emb.unsqueeze(0).unsqueeze(0),  # [1, 1, 4096]
             text_emb.unsqueeze(0),                # [1, 3, 4096]
@@ -76,15 +68,18 @@ def main(model_path: str = "."):
     
     # Negation example: a negation in text query should result
     # in retrieval of images without the neg. object in the query
-    print(colored("\n[5/6] Testing negation example...", 'cyan'))
+    print(colored("\n[5/5] Testing negation example...", 'cyan'))
     image_paths = [
         './assets/cat.png',
         './assets/dog+cat.png',
     ]
-    image_tensors = read_images_decord(image_paths)
-    with torch.no_grad():
-        image_embs = model.encode_vision(image_tensors.to(model.model.device)).cpu().float()
-        image_embs = torch.nn.functional.normalize(image_embs, dim=-1)
+    image_embs = []
+    for image_path in image_paths:
+        with torch.no_grad():
+            image_emb = model.encode_vision(image_path).cpu().float()
+            image_embs.append(image_emb)
+    image_embs = torch.cat(image_embs, dim=0)
+    image_embs = torch.nn.functional.normalize(image_embs, dim=-1)
     print(f"Image embedding shape: {image_embs.shape}")
     
     texts = ['an image of a cat but there is no dog in it']
@@ -107,19 +102,17 @@ def main(model_path: str = "."):
     
     
     # Composed video retrieval example
-    print(colored("\n[6/6] Testing composed video retrieval...", 'cyan'))
+    print(colored("\n[Bonus] Testing composed video retrieval...", 'cyan'))
     # source_video_path = './assets/source-27375787.mp4'
     # target_video_path = './assets/target-27387901.mp4'
     # edit_text = "Make the billboard blank"
     source_video_path = "./assets/5369546.mp4"
     target_video_path = "./assets/1006630957.mp4"
     edit_text ="make the tree lit up"
-    source_video_tensor = read_frames_decord(source_video_path, num_frames=4)
-    target_video_tensor = read_frames_decord(target_video_path, num_frames=16)
     with torch.no_grad():
-        source_video_emb = model.encode_vision(source_video_tensor.unsqueeze(0), edit_text).cpu().squeeze(0).float()
+        source_video_emb = model.encode_vision_with_text(source_video_path, edit_text).cpu().squeeze(0).float()
         source_video_emb = torch.nn.functional.normalize(source_video_emb, dim=-1)
-        target_video_emb = model.encode_vision(target_video_tensor.unsqueeze(0)).cpu().squeeze(0).float()
+        target_video_emb = model.encode_vision(target_video_path).cpu().squeeze(0).float()
         target_video_emb = torch.nn.functional.normalize(target_video_emb, dim=-1)
     sim_with_edit = source_video_emb @ target_video_emb.t()
     print(f"Source-Target similarity with edit: {sim_with_edit}")
@@ -133,7 +126,11 @@ def main(model_path: str = "."):
 if __name__ == "__main__":
     import argparse
     parser = argparse.ArgumentParser()
-    parser.add_argument("--model_path", type=str, default=".")
+    parser.add_argument("--model_path", type=str, default="/work/piyush/pretrained_checkpoints/Tarsier2-7b-0115/")
     args = parser.parse_args()
 
-    main(args.model_path)
+    # import sys
+    # sys.path.append(args.model_path)
+    from modeling_tara import TARA
+
+    main(args.model_path)

generation_config.json

@@ -1,8 +1,6 @@
 {
   "_from_model_config": true,
-  "bos_token_id": 1,
-  "eos_token_id": 2,
-  "pad_token_id": 0,
-  "transformers_version": "4.45.2",
-  "use_cache": false
+  "bos_token_id": 151643,
+  "eos_token_id": 151645,
+  "transformers_version": "4.45.0"
 }

model-00001-of-00003.safetensors → model-00001-of-00004.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:21e95f82edbfe05e3e2d55e1c76a3b56a4b29f5da9f7d7f7f7e8ed6ef29c9602
-size 4992930688
+oid sha256:f2c855a92271b7d4f6510eb0274ab3c5049ab98871c501da10845d0a6b83a8b2
+size 4966663320

model-00002-of-00003.safetensors → model-00002-of-00004.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:12847ca907e67d2bc044f2c1741c7b43e25aecd2ac32a663f4de8d16dbaaed0f
-size 4957878552
+oid sha256:96ae9d84f6e61e177d4f6bdaf7a7ee16e21af31e213f357b2862cc66b78fe481
+size 4991497784

model-00003-of-00003.safetensors → model-00003-of-00004.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:58d5999ca44b53916bc916d3d135376b420f5db31763e6b623b1705207865837
-size 4176137496
+oid sha256:8c51f1448e8650e932981b25398957691c0b3a16926ba944f6d4aff2a8599342
+size 4932752872

tokenizer.model → model-00004-of-00004.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:9e556afd44213b6bd1be2b850ebbbd98f5481437a8021afaf58ee7fb1818d347
-size 499723
+oid sha256:c98c3be29d5f182a6b18bfe72139c87a53aa0ff7287678d63faa10af4d09f0d4
+size 1691924640

modeling_tara.py

@@ -2,57 +2,34 @@ import os
 from abc import ABCMeta, abstractmethod
 from typing import Optional, Union, Dict, List
 from termcolor import colored
-import random
 
-
-import numpy as np
 import torch
 from transformers import (
-    AutoProcessor,
-    AutoTokenizer,
     LlavaConfig, 
-    LlamaForCausalLM,
-)
-from torchvision.transforms.v2 import (
-    ToPILImage,
 )
 import decord
-from decord import VideoReader
+import PIL.Image
+from tarsier2.dataset.utils import format_one_sample
+from tarsier2.modeling_tarsier2 import Tarsier2ForConditionalGeneration
+from tarsier2.modeling_qwen2_vl_fast import Qwen2VLForCausalLM
+from tarsier2.dataset.tarsier_datamodule import init_processor
 
 decord.bridge.set_bridge("torch")
 
-# TODO: need to use these directly
-from tarsier.modeling_tarsier import TarsierForConditionalGeneration
-from tarsier.processor import Processor
-# from utils.model import transform_pixel_values
-
 
 EOL_PROMPTS = {
     'text': '<sent>\nSummary above sentence in one word:',
     'image': '<image>\nSummary above image in one word:',
     'video': '<video>\nSummary above video in one word:',
+    "video_edit": "USER: Source video: <video>\nEdit instruction: <sent>\n"\
+    "Look at the attached video carefully. The provided text is instruction to edit the video. "\
+    "Imagine this edit instruction being applied to the provided video frame.\n"\
+    "Summarize the resulting edited video in one word: ASSISTANT:"
 }
+base_registry = {}
+encoder_registry = {}
 
 
-def transform_pixel_values(pixel_values: torch.Tensor | List[torch.Tensor]) -> torch.Tensor:
-    # NOTE: this function doesn't accept unbatched inputs
-    # pixel_values should be uint8 of (B, T, C, H, W)
-    if isinstance(pixel_values, list):
-        pixel_values = torch.stack(pixel_values)
-
-    if pixel_values.ndim == 4:
-        # pixel_values is (B, C, H, W)
-        # (B, C, H, W) -> (B, 1, C, H, W)
-        pixel_values = pixel_values.unsqueeze(1)
-    elif pixel_values.ndim == 5:
-        # pixel_values is (B, T, C, H, W)
-        pass
-    else:
-        raise ValueError(f"pixel_values should be 4D or 5D, got {pixel_values.ndim}D")
-    return pixel_values
-
-
-base_registry = {}
 class BaseModel(metaclass=ABCMeta):
     def __init_subclass__(cls, **kwargs):
         super().__init_subclass__(**kwargs)
@@ -67,16 +44,65 @@ class BaseModel(metaclass=ABCMeta):
         load_llm: bool = False,
         device_map: Optional[Union[str, Dict[str, int]]] = None,
         **kwargs):
-        print(colored(f'[ MODEL ] Loading {cls.__name__} from {model_name_or_path} [..............]', 'yellow'))
+        colored(f'Loading {cls.__name__} from {model_name_or_path}')
 
         return cls(model_name_or_path, load_llm=load_llm, device_map=device_map, **kwargs)
 
 
+class EncodeMixin(metaclass=ABCMeta):
+    def __init_subclass__(cls, **kwargs):
+        super().__init_subclass__(**kwargs)
+        # register model architecture
+        if hasattr(cls, 'ARCHITECTURE'):
+            encoder_registry[cls.ARCHITECTURE] = cls
+
+    @abstractmethod
+    def encode_vision(self, pixel_values: torch.Tensor | List[torch.Tensor]) -> torch.Tensor:
+        """
+        Encodes vision data (images or videos) into a tensor representation.
+
+        Args:
+            pixel_values (torch.Tensor | List[torch.Tensor]): The input pixel values. 
+                - If a tensor, it should be of shape (B, C, H, W) for images or (B, T, C, H, W) for videos.
+                - If a list, it will be stacked into a tensor.
+
+        Returns:
+            torch.Tensor: The encoded tensor representation of the input vision data.
+
+        Raises:
+            ValueError: If `pixel_values` is not 4D or 5D.
+
+        ## Notes:
+            - This function does not accept unbatched inputs.
+            - `pixel_values` should be of type uint8.
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    def encode_text(self, text: str | List[str]) -> torch.Tensor:
+        """
+        Encodes the given text(s) into a tensor representation using the model.
+
+        Args:
+            text (str | List[str]): A single string or a list of strings to be encoded.
+
+        Returns:
+            torch.Tensor: The tensor representation of the encoded text(s).
+
+        ## Notes:
+            - The method uses a prompt to encode the text.
+            - If a single string is provided, it is converted into a list containing that string.
+            - The method processes the prompts and generates the tensor representation using the model.
+            - The output tensor contains the hidden states of the last token for each input text.
+        """
+        raise NotImplementedError
+
+
 class BaseModelForTARA(BaseModel):
     
-    ARCHITECTURE = "TarsierForConditionalGeneration"
-    LLM_CLASS = LlamaForCausalLM
-    MLLM_CLASS = TarsierForConditionalGeneration
+    ARCHITECTURE = "Tarsier2ForConditionalGeneration"
+    LLM_CLASS = Qwen2VLForCausalLM
+    MLLM_CLASS = Tarsier2ForConditionalGeneration
 
     @property
     def describe_prompt(self):
@@ -84,27 +110,38 @@ class BaseModelForTARA(BaseModel):
 
     @property
     def text_eol_prompt(self):
-        prompt = f'USER: {EOL_PROMPTS["text"]} ASSISTANT: '
+        # prompt = f'USER: {EOL_PROMPTS["text"]} ASSISTANT: '
+        prompt = EOL_PROMPTS["text"]
         return prompt
     
     @property
     def image_eol_prompt(self):
-        prompt = f'USER: {EOL_PROMPTS["image"]} ASSISTANT: '
+        # prompt = f'USER: {EOL_PROMPTS["image"]} ASSISTANT: '
+        prompt = EOL_PROMPTS["image"]
         return prompt
     
     @property
     def video_eol_prompt(self):
-        prompt = f'USER: {EOL_PROMPTS["video"]} ASSISTANT: '
+        # prompt = f'USER: {EOL_PROMPTS["video"]} ASSISTANT: '
+        prompt = EOL_PROMPTS["video"]
         return prompt
 
-    @property
-    def video_edit_eol_prompt(self):
-        prompt = "Source video: <video>\nEdit instruction: <sent>\n"\
-        "Look at the attached video carefully. The provided text is instruction to edit the video. "\
-        "Imagine this edit instruction being applied to the provided video frame.\n"\
-        "Summarize the resulting edited video in one word:"
-        prompt = f"USER: {prompt} ASSISTANT: "
-        return prompt
+    @staticmethod
+    def _resolve_attn_implementation(requested_attn_impl: Optional[str] = None) -> str:
+        attn_impl = requested_attn_impl or "flash_attention_2"
+        if attn_impl != "flash_attention_2":
+            return attn_impl
+        if not torch.cuda.is_available():
+            print("CUDA is unavailable; falling back attn_implementation to 'eager'.")
+            return "eager"
+        major, _ = torch.cuda.get_device_capability(torch.cuda.current_device())
+        if major < 8:
+            print(
+                f"GPU compute capability {major}.x does not support FlashAttention-2; "
+                "falling back attn_implementation to 'eager'."
+            )
+            return "eager"
+        return "flash_attention_2"
 
     def __init__(
             self, 
@@ -120,27 +157,61 @@ class BaseModelForTARA(BaseModel):
             self.split_weights(model_name_or_path, model_name_or_path + '-llm')
             model_name_or_path += '-llm'
             model_config = None
-            self.processor = AutoProcessor.from_pretrained(model_name_or_path, use_fast=False)
+            
+            # from tarsier2.tarsier2_processor import TarsierProcessor
+            # self.processor = TarsierProcessor.from_pretrained(model_name_or_path, use_fast=False)
+            # self.tokenizer = self.processor.tokenizer
+            
+            import shared.utils as su
+            self.base_config = su.io.load_yml(
+                os.path.join(
+                    su.log.repo_path, 'tarsier2/default_config.yaml'
+                )
+            )
+            self.super_processor = init_processor(model_name_or_path, self.base_config)
+            self.processor = self.super_processor.processor
+            self.tokenizer = self.processor.tokenizer
+
         else:
             model_config = LlavaConfig.from_pretrained(
                 model_name_or_path,
-                # trust_remote_code=True,
+                trust_remote_code=True,
             )
-            self.processor = Processor(
-                model_name_or_path,
-                max_n_frames=32,
+            # from tarsier2.tarsier2_processor import TarsierProcessor
+            # self.processor = TarsierProcessor.from_pretrained(
+            #     model_name_or_path,
+            #     padding_side='left',
+            #     trust_remote_code=True,
+            # )
+            # Load base config
+            import shared.utils as su
+            self.base_config = su.io.load_yml(
+                os.path.join(
+                    su.log.repo_path, 'tarsier2/default_config.yaml'
+                )
             )
-        
-        self.tokenizer = self.processor.tokenizer
+            self.super_processor = init_processor(model_name_or_path, self.base_config)
+            self.processor = self.super_processor.processor
+            self.tokenizer = self.processor.tokenizer
+
+        attn_implementation = self._resolve_attn_implementation(
+            kwargs.get("attn_implementation", "flash_attention_2")
+        )
 
         self.model = MODEL_CLASS.from_pretrained(
             model_name_or_path,
             config=model_config,
-            torch_dtype=kwargs.get("torch_dtype", torch.bfloat16),
+            attn_implementation=attn_implementation,
+            # torch_dtype=kwargs.get("torch_dtype", torch.bfloat16),
+            torch_dtype=torch.bfloat16,
             device_map=device_map,
-            # trust_remote_code=True
+            trust_remote_code=True,
+            low_cpu_mem_usage=kwargs.get("low_cpu_mem_usage", True),  # Default to True for large models
         )
         
+        # self.processor.patch_size = self.model.config.vision_config.patch_size
+        # self.processor.vision_feature_select_strategy = self.model.config.vision_feature_select_strategy
+        
         self.model.eval()
 
     def split_weights(self, mllm_path, llm_path):
@@ -148,176 +219,167 @@ class BaseModelForTARA(BaseModel):
             print(f'{llm_path} already exists. Skip splitting weights.')
             return
         print('Splitting LLM weights from MLLM.')
-        model = self.MLLM_CLASS.from_pretrained(mllm_path)
+        attn_implementation = self._resolve_attn_implementation("flash_attention_2")
+        model = self.MLLM_CLASS.from_pretrained(
+            mllm_path,
+            attn_implementation=attn_implementation,
+            torch_dtype=torch.bfloat16,
+        )
         llm = model.language_model
-        processor = AutoProcessor.from_pretrained(mllm_path)
-        tokenizer = AutoTokenizer.from_pretrained(mllm_path)
         llm.save_pretrained(llm_path)
-        processor.save_pretrained(llm_path)
-        tokenizer.save_pretrained(llm_path)
-
-
-encoder_registry = {}
-class EncodeMixin(metaclass=ABCMeta):
-    def __init_subclass__(cls, **kwargs):
-        super().__init_subclass__(**kwargs)
-        # register model architecture
-        if hasattr(cls, 'ARCHITECTURE'):
-            encoder_registry[cls.ARCHITECTURE] = cls
-
-    @abstractmethod
-    def encode_vision(self, pixel_values: torch.Tensor | List[torch.Tensor]) -> torch.Tensor:
-        """
-        Encodes vision data (images or videos) into a tensor representation.
-
-        Args:
-            pixel_values (torch.Tensor | List[torch.Tensor]): The input pixel values. 
-                - If a tensor, it should be of shape (B, C, H, W) for images or (B, T, C, H, W) for videos.
-                - If a list, it will be stacked into a tensor.
-
-        Returns:
-            torch.Tensor: The encoded tensor representation of the input vision data.
-
-        Raises:
-            ValueError: If `pixel_values` is not 4D or 5D.
-
-        ## Notes:
-            - This function does not accept unbatched inputs.
-            - `pixel_values` should be of type uint8.
-        """
-        raise NotImplementedError
-
-    @abstractmethod
-    def encode_text(self, text: str | List[str]) -> torch.Tensor:
-        """
-        Encodes the given text(s) into a tensor representation using the model.
-
-        Args:
-            text (str | List[str]): A single string or a list of strings to be encoded.
-
-        Returns:
-            torch.Tensor: The tensor representation of the encoded text(s).
-
-        ## Notes:
-            - The method uses a prompt to encode the text.
-            - If a single string is provided, it is converted into a list containing that string.
-            - The method processes the prompts and generates the tensor representation using the model.
-            - The output tensor contains the hidden states of the last token for each input text.
-        """
-        raise NotImplementedError
+        
+        import shared.utils as su
+        from tarsier2.dataset.tarsier_datamodule import init_processor
+        base_config = su.io.load_yml(
+            os.path.join(su.log.repo_path, 'models/tarsier2/default_config.yaml'),
+        )
+        super_processor = init_processor(
+            mllm_path,
+            base_config,
+        )
+        super_processor.processor.save_pretrained(llm_path)
+        super_processor.processor.tokenizer.save_pretrained(llm_path)
 
 
 class TARA(BaseModelForTARA, EncodeMixin):
-
-    def encode_vision(self, pixel_values: torch.Tensor | List[torch.Tensor], edit_text: Optional[str] = None) -> torch.Tensor:
-
-        pixel_values = transform_pixel_values(pixel_values) # [B, T, C, H, W]
-        nframes = pixel_values.shape[1]
+    
+    def encode_vision(self, video_path: str, prompt=None) -> torch.Tensor:
         
-        if edit_text is not None:
-            # For composed video retrieval, we need to embed a video with the given text
-            prompt = self.video_edit_eol_prompt.replace('<sent>', edit_text)
+        ext = video_path.split('.')[-1]
+        if ext in ['mp4', 'avi', 'mov', 'mkv', 'webm']:
+            is_video = True
         else:
-            prompt = self.image_eol_prompt if nframes == 1 else self.video_eol_prompt
+            is_video = False
         
-        to_image = ToPILImage()
-        batched_frames = []
-        for batch in pixel_values:
-            frames = [to_image(v) for v in batch]
-            batched_frames.append(frames)
-
-        generate_kwargs = {
-            "max_new_tokens": 1,
-            "output_hidden_states": True,
-            "return_dict_in_generate": True,
-        }
-
-        vision_embs = []
-
-        for frames in batched_frames:
-            input_prompt = prompt.replace("<video>", "<image>"*len(frames))
-            input_ids = self.processor.get_text_inputs(input_prompt)
-            frames = self.processor.get_pixel_values(frames)
-            inputs = {
-                "input_ids": input_ids,
-                "pixel_values": frames
-            }
-            inputs = {k:v.to(self.model.device) for k,v in inputs.items() if v is not None}
-            outputs = self.model.generate(
-                **inputs,
-                **generate_kwargs,
+        if prompt is None:
+            if is_video:
+                prompt = self.video_eol_prompt
+            else:
+                prompt = self.image_eol_prompt
+        else:
+            assert "<video>" in prompt or "<image>" in prompt
+        sample = format_one_sample(media_file=video_path, prompt=prompt)
+        sample = self.super_processor(sample)
+        model_inputs = {}
+        for k, v in sample.items():
+            if not isinstance(v, torch.Tensor):
+                continue
+            model_inputs[k] = v.to(self.model.device)
+        with torch.inference_mode():
+            output = self.model.generate(
+                **model_inputs,
+                max_new_tokens=1,
+                output_hidden_states=True,
+                return_dict_in_generate=True,
+                pad_token_id=self.processor.tokenizer.eos_token_id
             )
-            vision_embs.append(outputs.hidden_states[0][-1][:, -1, :])
-        
-        vision_embs = torch.cat(vision_embs)
-        return vision_embs
+            emb = output.hidden_states[0][-1][:, -1, :]
+        return emb
     
-    def encode_text(self, text: str | List[str]) -> torch.Tensor:
-
-        prompt = self.text_eol_prompt
-
-        if isinstance(text, str):
-            text = [text]
-        
-        prompts = [prompt.replace('<sent>', t) for t in text]
-
-        generate_kwargs = {
-            "max_new_tokens": 1,
-            "output_hidden_states": True,
-            "return_dict_in_generate": True,
-        }
-
-        text_embs = []
-
-        for p in prompts:
-            text_inputs = self.processor.get_text_inputs(p)
-            inputs = {
-                "input_ids": text_inputs,
-            }
-            inputs = {k:v.to(self.model.device) for k,v in inputs.items() if v is not None}
-            outputs = self.model.generate(
-                **inputs,
-                **generate_kwargs,
+    def encode_vision_with_text(self, video_path: str, text: str) -> torch.Tensor:
+        ext = video_path.split('.')[-1]
+        # if ext in ['mp4', 'avi', 'mov', 'mkv', 'webm']:
+        #     is_video = True
+        # else:
+        #     is_video = False
+        # assert not is_video
+        prompt = EOL_PROMPTS["video_edit"].replace('<sent>', text)
+        sample = format_one_sample(media_file=video_path, prompt=prompt)
+        sample = self.super_processor(sample)
+        model_inputs = {}
+        for k, v in sample.items():
+            if not isinstance(v, torch.Tensor):
+                continue
+            model_inputs[k] = v.to(self.model.device)
+        with torch.inference_mode():
+            output = self.model.generate(
+                **model_inputs,
+                max_new_tokens=1,
+                output_hidden_states=True,
+                return_dict_in_generate=True,
+                pad_token_id=self.processor.tokenizer.eos_token_id
             )
-            text_embs.append(outputs.hidden_states[0][-1][:, -1, :])
-        
-        text_embs = torch.cat(text_embs)
-        return text_embs
-
-    def describe(self, pixel_values: torch.Tensor | List[torch.Tensor]) -> List[str]:
-
-        pixel_values = transform_pixel_values(pixel_values) # [B, T, C, H, W]
-        to_image = ToPILImage()
-        batched_frames = []
-        for batch in pixel_values:
-            frames = [to_image(v) for v in batch]
-            batched_frames.append(frames)
-        descriptions = []
-        generate_kwargs = {
-            "do_sample": False,
-            "max_new_tokens": 2048,
-            "top_p": 1,
-            "temperature": 0,
-            "use_cache": True
-        }
-
-        for frames in batched_frames:
-            text_inputs = f"<video>\n{self.describe_prompt}"
-            text_inputs = self.processor.process_prompt(text_inputs, frames)
-            text_inputs = self.processor.get_text_inputs(text_inputs)
-            frames = self.processor.get_pixel_values(frames)
-            inputs = {
-                "input_ids": text_inputs,
-                "pixel_values": frames
-            }
-            inputs = {k:v.to(self.model.device) for k,v in inputs.items() if v is not None}
-            outputs = self.model.generate(
-                **inputs,
-                **generate_kwargs,
+            emb = output.hidden_states[0][-1][:, -1, :]
+        return emb
+    
+    def encode_image(self, image_path: str, prompt=None):
+        ext = image_path.split('.')[-1]
+        if ext in ['mp4', 'avi', 'mov', 'mkv', 'webm']:
+            is_video = True
+        else:
+            is_video = False
+        assert not is_video
+        if prompt is None:
+            prompt = self.image_eol_prompt
+        else:
+            assert "<image>" in prompt
+        sample = format_one_sample(media_file=image_path, prompt=prompt)
+        sample = self.super_processor(sample)
+        model_inputs = {}
+        for k, v in sample.items():
+            if not isinstance(v, torch.Tensor):
+                continue
+            model_inputs[k] = v.to(self.model.device)
+        with torch.inference_mode():
+            output = self.model.generate(
+                **model_inputs,
+                max_new_tokens=1,
+                output_hidden_states=True,
+                return_dict_in_generate=True,
+                pad_token_id=self.processor.tokenizer.eos_token_id
             )
-            output_text = self.processor.tokenizer.decode(outputs[0][inputs['input_ids'][0].shape[0]:], skip_special_tokens=True)
-            descriptions.append(output_text)
-        return descriptions
+            emb = output.hidden_states[0][-1][:, -1, :]
+        return emb
+    
+    def encode_text(self, text: str, prompt=None) -> torch.Tensor:
+        
+        if prompt is None:
+            prompt = self.text_eol_prompt
+        else:
+            assert "<sent>" in prompt
+
+        if isinstance(text, str):
+            prompt = prompt.replace('<sent>', text)
+            sample = format_one_sample(media_file=None, prompt=prompt)
+            sample = self.super_processor(sample)
+            model_inputs = {}
+            for k, v in sample.items():
+                if not isinstance(v, torch.Tensor):
+                    continue
+                model_inputs[k] = v.to(self.model.device)
+            with torch.inference_mode():
+                output = self.model.generate(
+                    **model_inputs,
+                    max_new_tokens=1,
+                    output_hidden_states=True,
+                    return_dict_in_generate=True,
+                    pad_token_id=self.processor.tokenizer.eos_token_id
+                )
+                emb = output.hidden_states[0][-1][:, -1, :]
+            return emb
+        elif isinstance(text, list):
+            text_embs = []
+            for t in text:
+                prompt = self.text_eol_prompt.replace('<sent>', t)
+                sample = format_one_sample(media_file=None, prompt=prompt)
+                sample = self.super_processor(sample)
+                model_inputs = {}
+                for k, v in sample.items():
+                    if not isinstance(v, torch.Tensor):
+                        continue
+                    model_inputs[k] = v.to(self.model.device)
+                with torch.inference_mode():
+                    output = self.model.generate(
+                        **model_inputs,
+                        max_new_tokens=1,
+                        output_hidden_states=True,
+                        return_dict_in_generate=True,
+                    )
+                    emb = output.hidden_states[0][-1][:, -1, :]
+                    text_embs.append(emb)
+            return torch.cat(text_embs)
+        else:
+            raise ValueError(f"Invalid type for text: {type(text)}")
 
 
 def get_frame_indices(num_frames, vlen, sample='rand', fix_start=None, input_fps=1, max_num_frames=-1):
@@ -396,7 +458,6 @@ def read_frames_decord(
         del video_reader
 
 
-import PIL.Image
 def read_image_decord(image_path):
     image = PIL.Image.open(image_path)
     image = image.convert('RGB')
@@ -417,10 +478,14 @@ def read_images_decord(image_paths):
 
 
 if __name__ == "__main__":
+    from termcolor import colored
+    import random
+    import numpy as np
+    from decord import VideoReader
 
     # Load model
     model = TARA.from_pretrained(
-        "/work/piyush/experiments/CaRe/Tarsier-7b/final-10112025/nli_9000+ego_1000+subj_replaced-seed_42/merged_checkpoint",
+        "/work/piyush/experiments/CaRe/Tarsier2-7b-0115/covr/chiral10k-covr10k/merged_checkpoint/",
         device_map='auto',
         dtype=torch.bfloat16,
     )
@@ -430,12 +495,12 @@ if __name__ == "__main__":
     # Let's encode a sample video
     print(colored("Testing video encoding...", 'cyan'))
     video_path = "./assets/folding_paper.mp4"
-    video_tensor = read_frames_decord(video_path, num_frames=16)
-    video_tensor = video_tensor.unsqueeze(0)
-    video_tensor = video_tensor.to(model.model.device)
+    # video_tensor = read_frames_decord(video_path, num_frames=16)
+    # video_tensor = video_tensor.unsqueeze(0)
+    # video_tensor = video_tensor.to(model.model.device)
     with torch.no_grad():
-        video_emb = model.encode_vision(video_tensor).cpu().squeeze(0).float()
-    print("Video shape:", video_tensor.shape) # torch.Size([1, 16, 3, 240, 426])
+        video_emb = model.encode_vision(video_path).cpu().squeeze(0).float()
+    # print("Video shape:", video_tensor.shape) # torch.Size([1, 16, 3, 240, 426])
     print("Video embedding shape:", video_emb.shape) # torch.Size([4096])
     
     # Let's encode a sample text

preprocessor_config.json

@@ -1,9 +1,4 @@
 {
-  "crop_size": {
-    "height": 336,
-    "width": 336
-  },
-  "do_center_crop": true,
   "do_convert_rgb": true,
   "do_normalize": true,
   "do_rescale": true,
@@ -13,16 +8,22 @@
     0.4578275,
     0.40821073
   ],
-  "image_processor_type": "CLIPImageProcessor",
+  "image_processor_type": "Qwen2VLImageProcessor",
   "image_std": [
     0.26862954,
     0.26130258,
     0.27577711
   ],
-  "processor_class": "LlavaProcessor",
+  "max_pixels": 2073600,
+  "merge_size": 2,
+  "min_pixels": 3136,
+  "patch_size": 14,
+  "processor_class": "TarsierProcessor",
   "resample": 3,
   "rescale_factor": 0.00392156862745098,
   "size": {
-    "shortest_edge": 336
-  }
+    "max_pixels": 2073600,
+    "min_pixels": 3136
+  },
+  "temporal_patch_size": 2
 }

processor_config.json

@@ -1,6 +1,8 @@
 {
-  "image_token": "<image>",
-  "patch_size": null,
-  "processor_class": "LlavaProcessor",
-  "vision_feature_select_strategy": null
+  "image_token": "<|image_pad|>",
+  "max_seq_len": 16384,
+  "merge_size": 2,
+  "patch_size": 14,
+  "processor_class": "TarsierProcessor",
+  "temporal_patch_size": 2
 }

shared/run/cut_clips_ego4d.sh

@@ -0,0 +1,35 @@
+video_dir=/scratch/shared/beegfs/shared-datasets/EGO4D/ego4d_data_v1/full_scale/
+cut_dir=/scratch/shared/beegfs/piyush/datasets/Ego4D-HCap/cut_full_scale/
+# csv=/scratch/shared/beegfs/piyush/datasets/Ego4D-HCap/ego4d_chiral_subset-v1-with_reverse_captions-650K.csv
+# csv="/scratch/shared/beegfs/piyush/datasets/Ego4D-HCap/ego4d_chiral_subset-v1-with_reverse_captions-560K_>=0.5s_buffer=0.2.csv"
+# csv="/scratch/shared/beegfs/piyush/datasets/Ego4D-HCap/ego4d_chiral_subset-v1-with_reverse_captions-490K_>=0.5s_buffer=0.2.csv"
+csv="/scratch/shared/beegfs/piyush/datasets/Ego4D-HCap/metadata/cleaned_chiral+general-850K-2025-07-17_18:21:38.csv"
+si=$1
+ei=$2
+
+echo "CSV: $csv"
+echo "Start index: $si"
+echo "End index: $ei"
+echo "--------------------------------"
+
+# First, cut the original videos into clips based on the CSV
+echo "Cutting clips..."
+python shared/scripts/cut_clips_fast.py \
+    --csv $csv \
+    --video_dir $video_dir \
+    --cut_dir $cut_dir \
+    --video_id_key video_id \
+    --start_time_key start_sec \
+    --end_time_key stop_sec \
+    --si $si \
+    --ei $ei
+
+echo "Done cutting clips."
+echo "--------------------------------"
+
+# # Then, downsize the videos to 360 width
+# echo "Downsizing videos..."
+# python shared/scripts/downsize_videos_simple.py --video_dir $cut_dir --remove_old --width 360
+# echo "Done downsizing videos."
+# echo "--------------------------------"
+

shared/run/extract_feat_dinov2_ego4d.sh

@@ -0,0 +1,12 @@
+# DATA_DIR="/scratch/shared/beegfs/piyush/datasets/Ego4D-HCap"
+# csv="${DATA_DIR}/ego4d_chiral_subset-v1-with_reverse_captions-560K_>=0.5s_buffer=0.2.csv"
+# video_dir="${DATA_DIR}/cut_full_scale"
+dataset='ego4d_subset'
+si=$1
+ei=$2
+
+python adapt4change/scripts/compute_dino_features.py \
+    --dataset $dataset \
+    --no_filter_chiral \
+    --si $si \
+    --ei $ei

shared/run/extract_feat_pe_ego4d.sh

@@ -0,0 +1,17 @@
+# csv="/scratch/shared/beegfs/piyush/datasets/Ego4D-HCap/metadata/cleaned_chiral_subset_with_reverse_captions-425K-2025-07-17_18:21:38.csv"
+# csv="/scratch/shared/beegfs/piyush/datasets/Ego4D-HCap/metadata/cleaned_chiral+general-850K-2025-07-17_18:21:38.csv"
+csv="/scratch/shared/beegfs/piyush/datasets/Ego4D-HCap/metadata/cleaned_chiral+general-816K-2025-07-17_18:21:38.csv"
+output_dir="/scratch/shared/beegfs/piyush/datasets/Ego4D-HCap/features/"
+text_col="caption_forward"
+id_col="id"
+si=$1
+ei=$2
+
+python chiral_retrieval/scripts/compute_video_text_features.py \
+--csv $csv  \
+--output_dir $output_dir \
+--text_col $text_col \
+--id_col $id_col \
+--si $si \
+--ei $ei \
+--devices 1

shared/run/extract_feat_pe_ego4d_reverse.sh

@@ -0,0 +1,16 @@
+csv="/scratch/shared/beegfs/piyush/datasets/Ego4D-HCap/metadata/cleaned_chiral_subset_with_reverse_captions-425K-2025-07-17_18:21:38.csv"
+output_dir="/scratch/shared/beegfs/piyush/datasets/Ego4D-HCap/features/"
+text_col="caption_reverse"
+id_col="id"
+si=$1
+ei=$2
+
+python chiral_retrieval/scripts/compute_video_text_features.py \
+--csv $csv  \
+--output_dir $output_dir \
+--text_col $text_col \
+--id_col $id_col \
+--si $si \
+--ei $ei \
+--devices 1 \
+--reverse

shared/run/generate_water.py

@@ -0,0 +1,24 @@
+from diffusers import AudioLDM2Pipeline
+import torch
+import scipy.io.wavfile as wavfile
+
+def main():
+    model_id = "cvssp/audioldm2"
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+
+    # Load the pipeline with half-precision (optional if on GPU)
+    pipe = AudioLDM2Pipeline.from_pretrained(model_id, torch_dtype=torch.float16 if device=="cuda" else torch.float32)
+    pipe = pipe.to(device)
+
+    prompt = "High-quality sound of water being poured into a glass in a quiet room"
+    audio_length = 5.0  # seconds
+    steps = 200
+
+    outputs = pipe(prompt, num_inference_steps=steps, audio_length_in_s=audio_length)
+    audio = outputs.audios[0]
+
+    wavfile.write("pouring_water.wav", rate=16000, data=audio)
+    print("Saved to pouring_water.wav")
+
+if __name__ == "__main__":
+    main()

shared/run/generate_water_v2.py

@@ -0,0 +1,41 @@
+import torch
+import torchaudio
+from einops import rearrange
+from stable_audio_tools import get_pretrained_model
+from stable_audio_tools.inference.generation import generate_diffusion_cond
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+# Download model
+model, model_config = get_pretrained_model("stabilityai/stable-audio-open-1.0")
+sample_rate = model_config["sample_rate"]
+sample_size = model_config["sample_size"]
+
+model = model.to(device)
+
+# Set up text and timing conditioning
+conditioning = [{
+    "prompt": "The sound of a glass being filled with hot boiling water.",
+    "seconds_start": 0, 
+    "seconds_total": 10,
+}]
+
+# Generate stereo audio
+output = generate_diffusion_cond(
+    model,
+    steps=100,
+    cfg_scale=7,
+    conditioning=conditioning,
+    sample_size=sample_size,
+    sigma_min=0.3,
+    sigma_max=500,
+    sampler_type="dpmpp-3m-sde",
+    device=device
+)
+
+# Rearrange audio batch to a single sequence
+output = rearrange(output, "b d n -> d (b n)")
+
+# Peak normalize, clip, convert to int16, and save to file
+output = output.to(torch.float32).div(torch.max(torch.abs(output))).clamp(-1, 1).mul(32767).to(torch.int16).cpu()
+torchaudio.save("pouring_water_hot.wav", output, sample_rate)

shared/run/run.sh

@@ -0,0 +1,4 @@
+CUDA_VISIBLE_DEVICES=0 bash shared/run/extract_feat_dinov2_ego4d.sh 0 110000 &
+CUDA_VISIBLE_DEVICES=1 bash shared/run/extract_feat_dinov2_ego4d.sh 110000 220000 &
+CUDA_VISIBLE_DEVICES=2 bash shared/run/extract_feat_dinov2_ego4d.sh 220000 330000 &
+CUDA_VISIBLE_DEVICES=3 bash shared/run/extract_feat_dinov2_ego4d.sh 330000 440000 & wait

shared/scripts/avi_to_mp4.py

@@ -0,0 +1,50 @@
+import os
+import glob
+import argparse
+from moviepy.editor import VideoFileClip
+import sys
+
+def convert_avi_to_mp4(src_dir, dst_dir, ext, start_idx, end_idx):
+    # Ensure destination directory exists
+    os.makedirs(dst_dir, exist_ok=True)
+
+    # Get list of all .avi files in subdirectories
+    avi_files = glob.glob(os.path.join(src_dir, f'**/*.{ext}'), recursive=True)
+    
+    if not avi_files:
+        print(f"No .{ext} files found in {src_dir}")
+        return
+
+    # Apply start and end index filtering
+    avi_files = avi_files[start_idx:end_idx]
+    
+    for avi_file in avi_files:
+        # Get file ID without extension and parent path
+        file_id = os.path.splitext(os.path.relpath(avi_file, src_dir))[0]
+        mp4_file = os.path.join(dst_dir, f"{file_id}.mp4")
+        
+        # Create any necessary subdirectories in dst_dir
+        os.makedirs(os.path.dirname(mp4_file), exist_ok=True)
+        
+        try:
+            # Suppress moviepy verbose output
+            sys.stdout = open(os.devnull, 'w')
+            clip = VideoFileClip(avi_file)
+            clip.write_videofile(mp4_file, codec="libx264", audio_codec="aac", verbose=False, logger=None)
+            sys.stdout = sys.__stdout__
+            print(f"Converted: {avi_file} -> {mp4_file}")
+        except Exception as e:
+            sys.stdout = sys.__stdout__
+            print(f"Error converting {avi_file}: {e}")
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Convert .avi files to .mp4")
+    parser.add_argument("--src_dir", required=True, help="Source directory containing .avi files")
+    parser.add_argument("--dst_dir", required=True, help="Destination directory for .mp4 files")
+    parser.add_argument("--ext", default="avi", help="Extension of the source files (default: avi)")
+    parser.add_argument("--si", type=int, default=0, help="Start index of files to process (default: 0)")
+    parser.add_argument("--ei", type=int, default=None, help="End index of files to process (default: None)")
+    
+    args = parser.parse_args()
+
+    convert_avi_to_mp4(args.src_dir, args.dst_dir, args.ext, args.si, args.ei)

shared/scripts/check_cut_files.py

@@ -0,0 +1,50 @@
+"""Checks cut files."""
+import os
+import sys
+from glob import glob
+from tqdm import tqdm
+from joblib import Parallel, delayed
+
+import decord
+import numpy as np
+import pandas as pd
+
+
+if __name__ == "__main__":
+    video_dir = "/work/piyush/from_nfs2/datasets/EPIC-Kitchens-100/cut_clips"
+    files = glob(os.path.join(video_dir, "*/*/*.MP4"))
+    print("Total files:", len(files))
+
+    parallel = True
+
+    if not parallel:
+        failed = []
+        iterator = tqdm(files, desc="Checking files")
+        for f in iterator:
+            try:
+                vr = decord.VideoReader(f, ctx=decord.cpu(), num_threads=1)
+                random_frame = np.random.randint(0, len(vr))
+                random_frame = vr.get_batch([random_frame]).asnumpy()
+            except Exception as e:
+                failed.append(f)
+        import ipdb; ipdb.set_trace()
+    else:
+        def check_file(f):
+            try:
+                vr = decord.VideoReader(f, ctx=decord.cpu(), num_threads=1)
+                random_frame = np.random.randint(0, len(vr))
+                random_frame = len(vr) - 1
+                random_frame = vr.get_batch([random_frame]).asnumpy()
+                return None
+            except Exception as e:
+                return f
+
+        status = Parallel(n_jobs=24)(
+            delayed(check_file)(f) for f in tqdm(files, desc="Checking files")
+        )
+        failed = [f for f in status if f is not None]
+        print("Number of files on which loading failed:", len(failed))
+        import ipdb; ipdb.set_trace()
+
+        for f in failed: os.remove(f)
+    

shared/scripts/check_video_health.py

@@ -0,0 +1,109 @@
+"""Checks videos in a CSV for health."""
+import os
+import sys
+from joblib import Parallel, delayed
+
+import pandas as pd
+import numpy as np
+from torchcodec.decoders import SimpleVideoDecoder
+import decord
+
+import shared.utils as su
+
+
+def get_video_width(path):
+    try:
+        return SimpleVideoDecoder(path).metadata.width 
+    except:
+        return -1
+
+
+def check_decord_videoreader(path):
+    try:
+        vr = decord.VideoReader(path)
+        return True
+    except:
+        return False
+
+
+def check_decord_random_frame(path):
+    try:
+        vr = decord.VideoReader(path)
+        i = np.random.randint(len(vr))
+        frame = vr[i]
+        return True
+    except:
+        return False
+
+
+if __name__ == "__main__":
+    # Configure video_dir and csv_path
+    data_dir = "/scratch/shared/beegfs/piyush/datasets/Ego4D-HCap"
+    video_dir = f"{data_dir}/cut_full_scale"
+    csv_path = f"{data_dir}/metadata/"\
+        "cleaned_chiral_subset_with_reverse_captions-425K-2025-07-17_18:21:38.csv"
+    id_col = "id"
+    save_dir = "./outputs/ego4d_video_health"
+    os.makedirs(save_dir, exist_ok=True)
+    
+    health_checks = [
+        # "video_exists",
+        "video_widths",
+        # "decord_videoreader",
+        # "decord_random_frame",
+    ]
+    
+    
+    assert os.path.isdir(video_dir), f"Video directory {video_dir} does not exist"
+    assert os.path.exists(csv_path), f"CSV file {csv_path} does not exist"
+    
+    # Load the CSV
+    su.log.print_update(f"Loading CSV: {csv_path}")
+    df = pd.read_csv(csv_path)
+    print("Number of rows in CSV:", len(df))
+    
+    # Add video_path column
+    df['video_path'] = df[id_col].apply(lambda x: f"{video_dir}/{x}.mp4")
+    
+    
+    # 1. Check if the video files exist
+    if "video_exists" in health_checks:
+        su.log.print_update("Checking if the video files exist")
+        iterator = su.log.tqdm_iterator(df.video_path.tolist())
+        video_exists = Parallel(n_jobs=-1)(delayed(os.path.exists)(f) for f in iterator)
+        print("Fraction of videos that exist: ", np.mean(video_exists))
+        ids_with_missing_videos = df.loc[~np.array(video_exists), id_col].tolist()
+        np.save(f"{save_dir}/ids_with_missing_videos.npy", ids_with_missing_videos)
+    
+    # 2. Check if the video widths are valid
+    if "video_widths" in health_checks:
+        su.log.print_update("Checking if the video widths are valid")
+        iterator = su.log.tqdm_iterator(df.video_path.tolist())
+        video_widths = Parallel(n_jobs=-1)(delayed(get_video_width)(f) for f in iterator)
+        video_widths = np.array(video_widths)
+        import ipdb; ipdb.set_trace()
+        print("Fraction of videos that are valid: ", np.mean(video_widths != -1))
+        ids_with_invalid_widths = df.loc[np.where(video_widths == -1), id_col].tolist()
+        np.save(f"{save_dir}/ids_with_invalid_widths.npy", np.array(ids_with_invalid_widths))
+    
+    # 3. Check if the video files are decodable by decord
+    if "decord_videoreader" in health_checks:
+        su.log.print_update("Checking if the video files are decodable by decord")
+        iterator = su.log.tqdm_iterator(df.video_path.tolist())
+        decord_videoreader = Parallel(n_jobs=-1)(delayed(check_decord_videoreader)(f) for f in iterator)
+        decord_videoreader = np.array(decord_videoreader)
+        import ipdb; ipdb.set_trace()
+        print("Fraction of videos that are decodable by decord: ", np.mean(decord_videoreader))
+        ids_with_invalid_decord = df.loc[~decord_videoreader, id_col].tolist()
+        np.save(f"{save_dir}/ids_with_invalid_decord.npy", np.array(ids_with_invalid_decord))
+    
+    
+    # 4. Check if a random frame can be decoded by decord
+    if "decord_random_frame" in health_checks:
+        su.log.print_update("Checking if a random frame can be decoded by decord")
+        iterator = su.log.tqdm_iterator(df.video_path.tolist())
+        decord_random_frame = Parallel(n_jobs=-1)(delayed(check_decord_random_frame)(f) for f in iterator)
+        decord_random_frame = np.array(decord_random_frame)
+        print("Fraction of videos that have a random frame that can be decoded by decord: ", np.mean(decord_random_frame))
+        ids_with_invalid_decord_random_frame = df.loc[~decord_random_frame, id_col].tolist()
+        np.save(f"{save_dir}/ids_with_invalid_decord_random_frame.npy", np.array(ids_with_invalid_decord_random_frame))

shared/scripts/check_webdataset.py

@@ -0,0 +1,69 @@
+"""Loads tar files using webdataset."""
+import os
+import webdataset as wds
+import decord
+from torch.utils.data import DataLoader
+import numpy as np
+import einops
+
+
+# Define a function to decode videos using Decord
+def decode_video(video_bytes):
+    # Save the video bytes to a temporary file and decode with decord
+    vr = decord.VideoReader(video_bytes)
+    frames = [vr[i].asnumpy() for i in range(0, len(vr), 5)]  # Frame skip example
+    return frames
+
+
+def convert_bytes_to_frames(video_bytes):
+    vr = decord.VideoReader(video_bytes)
+    frames = [vr[i].asnumpy() for i in range(0, len(vr), 5)]  # Frame skip example
+    return frames
+
+
+
+def decode_video(video_bytes):
+    """
+    Given video bytes, decode them into frames.
+    """
+    pass
+
+
+if __name__ == "__main__":
+    shard_folder = "/work/piyush/from_nfs2/datasets/SSv2/ssv2_shards/"
+    shard_path = os.path.join(shard_folder, "shard-0000.tar")
+
+    # Define your WebDataset path pattern (all shards)
+    dataset_path = os.path.join(shard_folder, "shard-{0000..0002}.tar")
+
+    ds = wds.WebDataset(dataset_path)
+    sample = next(iter(ds))
+    dl = DataLoader(ds, batch_size=16, num_workers=8)
+    batch = next(iter(dl))
+
+    # Create a WebDataset loader
+    dataset = (
+        wds.WebDataset(dataset_path)
+        # .decode("rgb")  # Ensure that we decode the video bytes into RGB images
+        .to_tuple("webm")  # Ensure that we get the video bytes and metadata
+        # .map_tuple(decode_video)  # Apply your video decoding function
+    )
+    # dataloader = DataLoader(dataset, batch_size=16, num_workers=8)
+    # batch = next(iter(dataloader))
+    # print(batch[0].shape)  # (16, 32, 256, 256, 3) for example
+    # H, W, C = 256, 256, 3
+    H, W, C = (240, 427, 3)
+    for (video,) in dataset:
+        # Get file name
+
+
+        # print(len(video[0]))
+        print(type(video))
+        video_array = np.frombuffer(video, dtype=np.uint8)
+        reshaped_video = einops.rearrange(video_array, "(t h w c) -> t c h w", h=H, w=W, c=C)
+        # print(len(video))
+        # print(video[1])
+
+        # np_video_bytes = np.frombuffer(video[0], np.uint8)
+        break
+    import ipdb; ipdb.set_trace()

shared/scripts/convert_frames_to_videos.py

@@ -0,0 +1,70 @@
+"""
+Converts frames to videos.
+
+input_dir=/scratch/shared/beegfs/piyush/datasets/Jester/images/
+output_dir=/scratch/shared/beegfs/piyush/datasets/Jester/videos/
+python shared/scripts/convert_frames_to_videos.py --input_dir $input_dir --output_dir $output_dir
+"""
+import os
+import sys
+import decord
+import moviepy.editor as mpy
+from glob import glob
+from moviepy.editor import ImageSequenceClip
+
+import shared.utils as su
+
+
+def create_video(image_paths, output_file, fps=12):
+    """
+    Convert a list of image files into an MP4 video
+    
+    Parameters:
+    - image_paths: List of image file paths (sorted in desired order)
+    - output_file: Output filename (e.g., 'output.mp4')
+    - fps: Frames per second (default 12)
+    """
+    clip = ImageSequenceClip(image_paths, fps=fps, load_images=True)
+    clip.write_videofile(output_file, codec='libx264', logger=None)
+
+
+
+if __name__ == "__main__":
+    # Read arguments
+    import argparse
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--input_dir", type=str, required=True)
+    parser.add_argument("--output_dir", type=str, required=True)
+    parser.add_argument("--fps", type=int, default=12)
+    parser.add_argument("--extension", type=str, default="jpg")
+    parser.add_argument("--debug", action="store_true")
+    parser.add_argument("--overwrite", action="store_true")
+    args = parser.parse_args()
+
+    # Create output directory
+    os.makedirs(args.output_dir, exist_ok=True)
+
+    # Get list of all frame folders
+    print("Getting list of all frame folders...")
+    folders = os.listdir(args.input_dir)
+
+    # Run the process for each folder
+    iterator = su.log.tqdm_iterator(folders, desc="Converting frames to videos")
+    for folder in iterator:
+
+        # Get output path
+        save_path = os.path.join(args.output_dir, f"{folder}.mp4")
+        if os.path.exists(save_path) and not args.overwrite:
+            continue
+
+        # Get list of frames
+        frame_paths = glob(
+            os.path.join(args.input_dir, folder, f"*.{args.extension}"),
+        )
+
+        # Save it as a video
+        create_video(frame_paths, save_path, fps=args.fps)
+
+        if args.debug:
+            print("Debugging mode. Exiting after processing one folder.")
+            break

shared/scripts/convert_webm_to_mp4.py

@@ -0,0 +1,44 @@
+"""
+Script to conver webm to mp4
+
+Example:
+video_dir=/scratch/shared/beegfs/shared-datasets/SomethingSomething-V2/videos/
+out_dir=/scratch/shared/nfs2/piyush/datasets/SSv2/videos/
+python shared/scripts/convert_webm_to_mp4.py --video_dir $video_dir --out_dir $out_dir
+"""
+
+import os
+from subprocess import call
+from glob import glob
+from tqdm import tqdm
+
+
+if __name__ == "__main__":
+    import argparse
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--video_dir", type=str, required=True)
+    parser.add_argument("--out_dir", type=str, required=True)
+    args = parser.parse_args()
+
+    ext = ".webm"
+    
+    in_files = glob(os.path.join(args.video_dir, f"*/*{ext}"))
+    out_files = [f.replace(f"{ext}", ".mp4") for f in in_files]
+    out_files = [f.replace(args.video_dir, args.out_dir) for f in out_files]
+
+    print("> Set to convert", len(in_files), f"files from {ext} to mp4.")
+    iterator = tqdm(range(len(in_files)), desc="Converting videos")
+    for i in iterator:
+        in_file = in_files[i]
+        out_file = out_files[i]
+        os.makedirs(os.path.dirname(out_file), exist_ok=True)
+        
+        if not os.path.exists(out_file):
+            command = f"ffmpeg -i {in_file} -c:v copy -c:a copy -strict -2 {out_file} -loglevel quiet"
+            call(command, shell=True)
+        else:
+            print(f"Skipping {in_file} as {out_file} already exists.")
+
+        # os.remove(in_file)
+        break
+    print("> Done converting.")

shared/scripts/create_webdataset.py

@@ -0,0 +1,265 @@
+import os
+import csv
+import argparse
+import multiprocessing as mp
+from pathlib import Path
+from typing import List, Dict
+from functools import partial
+import webdataset as wds
+import torch
+import numpy as np
+from tqdm import tqdm
+from decord import VideoReader
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Convert CSV file to WebDataset format with video data")
+    parser.add_argument("--csv_path", type=str, required=True, help="Path to the CSV file")
+    parser.add_argument("--output_dir", type=str, required=True, help="Output directory for WebDataset shards")
+    parser.add_argument("--num_shards", type=int, default=128, help="Number of shards to create")
+    parser.add_argument("--samples_per_shard", type=int, default=None, 
+                        help="Max samples per shard (overrides num_shards if specified)")
+    parser.add_argument("--worker_count", type=int, default=mp.cpu_count(), 
+                        help="Number of worker processes")
+    parser.add_argument("--shard_prefix", type=str, default="shard", 
+                        help="Prefix for shard filenames")
+    parser.add_argument("--video_extension", type=str, default=".webm", 
+                        help="Extension of video files (default: .webm)")
+    parser.add_argument("--debug", action="store_true", 
+                        help="Debug mode: create a shard_debug.tar file with max 1000 videos")
+    parser.add_argument("--si", type=int, default=0, help="Start index")
+    parser.add_argument("--ei", type=int, default=None, help="End index")
+    return parser.parse_args()
+
+
+def read_csv_data(csv_path: str, debug: bool = False) -> List[Dict]:
+    """Read the CSV file and return a list of samples."""
+    samples = []
+    with open(csv_path, 'r') as f:
+        reader = csv.DictReader(f)
+        for i, row in enumerate(reader):
+            samples.append(row)
+            # In debug mode, limit to 1000 samples
+            if debug and i >= 999:
+                break
+    
+    # Select start and end index if specified
+    si = args.si
+    ei = args.ei if args.ei is not None else len(samples)
+    print("Selected samples from index", si, "to", ei)
+    samples = samples[si:ei]
+
+    return samples
+
+
+def distribute_samples(samples: List[Dict], num_shards: int) -> List[List[Dict]]:
+    """Distribute samples across shards."""
+    samples_per_shard = len(samples) // num_shards
+    remainder = len(samples) % num_shards
+    
+    distributed_samples = []
+    start_idx = 0
+    
+    for i in range(num_shards):
+        # Add one extra sample for the first 'remainder' shards
+        shard_size = samples_per_shard + (1 if i < remainder else 0)
+        end_idx = start_idx + shard_size
+        
+        distributed_samples.append(samples[start_idx:end_idx])
+        start_idx = end_idx
+    
+    return distributed_samples
+
+
+def process_shard(shard_samples: List[Dict], shard_path: str, video_extension: str = ".webm"):
+    """Process and write a single shard with actual video data."""
+    with wds.TarWriter(shard_path) as sink:
+        for sample in tqdm(shard_samples, desc=f"Processing {shard_path}"):
+            video_path = sample['video_path']
+
+            vr = VideoReader(video_path, num_threads=1)
+            n_frames = len(vr)
+            fps = vr.get_avg_fps()
+            H, W, _ = vr[0].shape
+            
+            try:
+                # Read video file as binary data
+                with open(video_path, 'rb') as f:
+                    video_data = f.read()
+                
+                # Get filename without path for the key
+                filename = Path(video_path).stem
+
+                # Create sample with the actual video data
+                sample_dict = {
+                    "__key__": filename,
+                    "video": video_data,  # Actual video binary data
+                    "video.extension": video_extension.lstrip('.'),  # Store extension without dot
+                    "target": str(sample['target']), # Target/label
+                    # "split": str(sample['split']),  # Train/val/test split
+                    "json": dict(n_frames=n_frames, fps=fps, H=H, W=W),  # Additional metadata
+                }
+                
+                sink.write(sample_dict)
+            except Exception as e:
+                print(f"Error processing {video_path}: {str(e)}")
+
+
+import io
+import torchvision
+
+def encode_tensor(tensor):
+    """
+    Convert tensor to bytes in memory.
+    """
+    # Convert the tensor to bytes in memory
+    with io.BytesIO() as buf:
+        if isinstance(tensor, torch.Tensor):
+            torch.save(tensor, buf)  # Save tensor to the buffer
+        return buf.getvalue()  # Return the byte data
+
+
+def process_shard_tensor(shard_samples: List[Dict], shard_path: str, video_extension: str = ".webm"):
+    """Process and write a single shard with actual video data (actual tensor)."""
+    with wds.TarWriter(shard_path) as sink:
+        for sample in tqdm(shard_samples, desc=f"Processing {shard_path}"):
+            video_path = sample['video_path']
+
+            # Load the entire video as a tensor
+            video, audio, info = torchvision.io.read_video(video_path, pts_unit='sec')
+            video_data = encode_tensor(video)
+            n_frames = len(video)
+            fps = info['video_fps']
+            H, W = video.shape[1:-1]
+            
+            try:
+                # # Read video file as binary data
+                # with open(video_path, 'rb') as f:
+                #     video_data = f.read()
+                
+                # Get filename without path for the key
+                filename = Path(video_path).stem
+
+                # Create sample with the actual video data
+                sample_dict = {
+                    "__key__": filename,
+                    "video": video_data,  # Actual video binary data
+                    "video.extension": video_extension.lstrip('.'),  # Store extension without dot
+                    "target": str(sample['target']), # Target/label
+                    # "split": str(sample['split']),  # Train/val/test split
+                    "json": dict(n_frames=n_frames, fps=fps, H=H, W=W),  # Additional metadata
+                }
+
+                sink.write(sample_dict)
+            except Exception as e:
+                print(f"Error processing {video_path}: {str(e)}")
+
+
+def create_webdataset(csv_path: str, output_dir: str, num_shards: int, 
+                      samples_per_shard: int = None, worker_count: int = None,
+                      shard_prefix: str = "shard", video_extension: str = ".webm", 
+                      debug: bool = False):
+    """Convert CSV to WebDataset format with video data and parallel processing."""
+    os.makedirs(output_dir, exist_ok=True)
+    
+    # Read all samples from the CSV
+    print(f"Reading samples from {csv_path}...")
+    samples = read_csv_data(csv_path, debug=debug)
+    total_samples = len(samples)
+    print(f"Found {total_samples} samples in the CSV file")
+    
+    # Handle debug mode with a specific shard_debug.tar file
+    if debug:
+        print("Debug mode enabled: Creating shard_debug.tar with max 1000 videos")
+        debug_shard_path = os.path.join(output_dir, "shard_debug.tar")
+        # process_shard(samples, debug_shard_path, video_extension)
+        process_shard_tensor(samples, debug_shard_path, video_extension)
+        
+        # Calculate and display file size
+        file_size = os.path.getsize(debug_shard_path)
+        print(f"Created debug shard: {debug_shard_path}")
+        print(f"Debug shard size: {file_size / (1024**2):.2f} MB")
+        
+        # Test the debug shard
+        test_dataset(output_dir, debug_pattern="shard_debug.tar")
+        return
+    
+    # Determine number of shards based on samples_per_shard if provided
+    if samples_per_shard is not None:
+        num_shards = (total_samples + samples_per_shard - 1) // samples_per_shard
+        print(f"Creating {num_shards} shards with max {samples_per_shard} samples per shard")
+    else:
+        print(f"Creating {num_shards} shards")
+    
+    # Distribute samples across shards
+    shard_samples = distribute_samples(samples, num_shards)
+    
+    # Prepare shard paths
+    shard_paths = [
+        os.path.join(output_dir, f"{shard_prefix}_{i:05d}.tar") 
+        for i in range(num_shards)
+    ]
+    
+    # Use all available cores if worker_count is not specified
+    if worker_count is None:
+        worker_count = mp.cpu_count()
+        
+    worker_count = min(worker_count, num_shards)  # Don't use more workers than shards
+    
+    print(f"Using {worker_count} worker processes")
+    
+    # Process shards in parallel with video extension
+    # process_func = partial(process_shard, video_extension=video_extension)
+    process_func = partial(process_shard_tensor, video_extension=video_extension)
+    
+    with mp.Pool(worker_count) as pool:
+        list(tqdm(
+            pool.starmap(process_func, zip(shard_samples, shard_paths)),
+            total=num_shards,
+            desc="Creating WebDataset shards with video data"
+        ))
+    
+    print(f"Successfully created {num_shards} WebDataset shards in {output_dir}")
+    
+    # Calculate and display total dataset size
+    total_size = sum(os.path.getsize(path) for path in shard_paths)
+    print(f"Total dataset size: {total_size / (1024**2):.2f} MB")
+
+
+def test_dataset(output_dir: str, shard_prefix: str = "shard", debug_pattern: str = None):
+    """Test reading from the created WebDataset."""
+    # Find all shard files or use debug pattern
+    if debug_pattern:
+        shard_pattern = os.path.join(output_dir, debug_pattern)
+    else:
+        shard_pattern = os.path.join(output_dir, f"{shard_prefix}_*.tar")
+    
+    # Create a dataset
+    dataset = wds.WebDataset(shard_pattern)
+    
+    # Display sample info
+    print("\nTesting dataset:")
+    for i, sample in enumerate(dataset):
+        print(f"Sample {i}:")
+        for key, value in sample.items():
+            if key == "video":
+                print(f"  {key}: <binary data of length {len(value)}>")
+            else:
+                print(f"  {key}: {value}")
+        
+        if i >= 2:  # Just show a few samples
+            break
+
+
+if __name__ == "__main__":
+    args = parse_args()
+    create_webdataset(
+        csv_path=args.csv_path,
+        output_dir=args.output_dir,
+        num_shards=args.num_shards,
+        samples_per_shard=args.samples_per_shard,
+        worker_count=args.worker_count,
+        shard_prefix=args.shard_prefix,
+        video_extension=args.video_extension,
+        debug=args.debug
+    )

shared/scripts/cut_clips.py

@@ -0,0 +1,203 @@
+"""
+Cut video clips from downloaded videos.
+
+Example.
+D=/work/piyush/from_nfs2/datasets/Charades
+video_dir=$D/Charades_v1_480/
+
+EPIC
+
+S=/datasets/EpicKitchens-100/
+D=/work/piyush/from_nfs2/datasets/EPIC-Kitchens-100/cut_clips
+csv=$D/../epic-kitchens-100-annotations/EPIC_100_train_with_id.csv
+python shared/scripts/cut_clips.py --csv $csv --video_id_key path_id --start_time_key start_sec --end_time_key stop_sec --video_dir $S/ --cut_dir $D/  --ext MP4
+"""
+import os
+from os.path import join, exists
+from subprocess import call
+import time
+
+import numpy as np
+import pandas as pd
+from tqdm import tqdm
+
+import shared.utils.io as io
+import shared.utils.log as log
+from video_language.datasets.charades import get_paths, load_main_csv
+
+
+def time_float_to_str(time_in_seconds):
+    import datetime
+
+    # Calculate hours, minutes, seconds, and milliseconds
+    hours, remainder = divmod(time_in_seconds, 3600)
+    minutes, seconds_with_ms = divmod(remainder, 60)
+    seconds, milliseconds = divmod(int(seconds_with_ms * 1000), 1000)
+
+    # Create a timedelta object
+    time_delta = datetime.timedelta(hours=hours, minutes=minutes, seconds=seconds, milliseconds=milliseconds)
+
+    # Format the time as HH:MM:SS.mmm
+    formatted_time = str(time_delta)
+    
+    return formatted_time
+
+
+if __name__ == "__main__":
+
+    import argparse
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--csv", type=str, required=True,
+        help="Path to CSV file containing video IDs and timestamps",
+    )
+    parser.add_argument(
+        "--video_id_key", type=str, default="video_id",
+    )
+    parser.add_argument(
+        "--start_time_key", type=str, default="start_time",
+    )
+    parser.add_argument(
+        "--end_time_key", type=str, default="end_time",
+    )
+    parser.add_argument(
+        "--video_dir", type=str, required=True,
+        help="Path to directory containing downloaded videos",
+    )
+    parser.add_argument(
+        "--cut_dir", type=str, required=True,
+        help="Path to directory where cut videos will be saved",
+    )
+    parser.add_argument(
+        "--overwrite", action="store_true",
+        help="Whether to overwrite existing cut videos",
+    )
+    parser.add_argument(
+        "--verbose", action="store_true",
+    )
+    parser.add_argument(
+        "--no_round_times", action="store_true",
+        help="Whether to round start and end times to nearest second in filenames",
+    )
+    parser.add_argument(
+        "--debug", action="store_true",
+    )
+    parser.add_argument(
+        "--ext", type=str, default="mp4",
+    )
+    parser.add_argument(
+        "--si", type=int, default=0,
+    )
+    parser.add_argument(
+        "--ei", type=int, default=None,
+    )
+    parser.add_argument(
+        "--filter_csv", type=str, default=None, required=False,
+    )
+    parser.add_argument(
+        "--filter_key", type=str, default=None, required=False,
+    )
+    args = parser.parse_args()
+    
+    # Make cut_dir
+    os.makedirs(args.cut_dir, exist_ok=True)
+    
+    # Load csv
+    assert os.path.exists(args.csv), f"CSV file {args.csv} does not exist."
+    df = pd.read_csv(args.csv)
+    print(">>> Loaded CSV file with shape", df.shape)
+    assert {args.video_id_key, args.start_time_key, args.end_time_key}.issubset(df.columns), \
+        f"CSV file must contain columns {args.video_id_key}, {args.start_time_key}, and {args.end_time_key}."
+
+    # Filter CSV
+    if args.filter_csv is not None:
+        path = args.filter_csv
+        assert os.path.exists(path), f"CSV file {path} does not exist."
+
+        key = args.filter_key
+        df_filter = pd.read_csv(path)
+        assert key in df_filter.columns, f"CSV file must contain column {key}."
+
+        # Only keep the rows in df that match on key with df_filter
+        keep_values = df_filter[key].unique()
+
+        df = df[df[key].isin(keep_values)]
+        print(">>> Filtered CSV file with shape", df.shape)
+    
+    # Filter out videos that don't exist
+    df["video_path"] = df[args.video_id_key].apply(
+        lambda video_id: join(args.video_dir, f"{video_id}.{args.ext}"),
+    )
+    df["check_video"] = df["video_path"].apply(exists)
+    df = df[df["check_video"]]
+    del df["check_video"]
+    print(">>> Found videos for", df.shape[0], "rows.")
+
+    if len(df) == 0:
+        print(">>> No videos to cut.")
+        exit()
+
+
+    si = args.si
+    ei = args.ei if args.ei is not None else len(df)
+    df = df.iloc[si:ei]
+    print("Start index:", si, "End index:", ei)
+
+    # Custom filter
+    # df = df[df.split == "validation"]
+    # print(">>> Filtered videos for", df.shape[0], "rows.")
+
+    if args.debug:
+        args.verbose = True
+
+
+    # Cut videos
+    ext = args.ext
+    iterator = tqdm(range(len(df)), desc="Cutting clips")
+    for i in iterator:
+
+        row = df.iloc[i].to_dict()
+        f = row["video_path"]
+        v, s, e = row[args.video_id_key], row[args.start_time_key], row[args.end_time_key]
+        s = float(s)
+        e = float(e)
+
+        if args.no_round_times:
+            clip_filename = f"{v}_{s}_{e}.{ext}"
+        else:
+            clip_filename = f"{v}_{np.round(s, 1)}_{np.round(e, 1)}.{ext}"
+        clip_filepath = join(args.cut_dir, clip_filename)
+        os.makedirs(os.path.dirname(clip_filepath), exist_ok=True)
+
+        if os.path.exists(clip_filepath) and not args.overwrite:
+            continue
+        
+        # bring s in HH:MM:SS.mmm format with milliseconds
+        s = time_float_to_str(s)
+        e = time_float_to_str(e)
+        # # bring s in HH:MM:SS. format
+        # s = time.strftime("%H:%M:%S", time.gmtime(s))
+        # e = time.strftime("%H:%M:%S", time.gmtime(e))
+
+        # ffmpeg code
+        # ffmpeg_source = "/users/piyush/install/ffmpeg-06092024/ffmpeg-7.0.2-i686-static/ffmpeg"
+        ffmpeg_source = " /users/piyush/install/ffmpeg/ffmpeg-7.0.2-i686-static/ffmpeg"
+        # print("FFMpeg version: ", call(f"{ffmpeg_source} -version", shell=True))
+        # use ffmpeg to cut the clip + change spatial resolution to have max height
+        # NOTE: also changes spatial resolution to have max width as 480
+        command = f"{ffmpeg_source} -i {f} -ss {s} -to {e} -strict -2 -c:v libx264 "\
+            f"-pix_fmt yuv420p -c:a copy"\
+            " -vf 'scale=480:-1' "\
+            f"{clip_filepath} "\
+            f"-y -format {ext}"
+        if not args.verbose:
+            command += " -loglevel quiet"
+        else:
+            print(">>> Cutting clip", clip_filepath)
+        call(command, shell=True)
+
+        if args.debug:
+            print(command)
+            break
+    
+    print(">>> Number of cut files:", len(os.listdir(args.cut_dir)))

shared/scripts/cut_clips_fast.py

@@ -0,0 +1,200 @@
+"""
+Faster clip cutting script generated by Claude.
+
+S=/datasets/EpicKitchens-100/
+D=/work/piyush/from_nfs2/datasets/EPIC-Kitchens-100/cut_clips
+csv=$D/../epic-kitchens-100-annotations/EPIC_100_train_with_id.csv
+python shared/scripts/cut_clips_fast.py --csv $csv --video_id_key path_id --start_time_key start_sec --end_time_key stop_sec --video_dir $S/ --cut_dir $D/ --ext MP4 --max_workers 4
+
+"""
+import os
+from os.path import join, exists
+import time
+from concurrent.futures import ThreadPoolExecutor, as_completed
+
+import numpy as np
+import pandas as pd
+from tqdm import tqdm
+from moviepy.editor import VideoFileClip
+from moviepy.video.fx.resize import resize
+
+def time_float_to_str(time_in_seconds):
+    import datetime
+    hours, remainder = divmod(time_in_seconds, 3600)
+    minutes, seconds_with_ms = divmod(remainder, 60)
+    seconds, milliseconds = divmod(int(seconds_with_ms * 1000), 1000)
+    time_delta = datetime.timedelta(hours=hours, minutes=minutes, seconds=seconds, milliseconds=milliseconds)
+    return str(time_delta)
+
+def process_video(row, args):
+    """Process a single video clip"""
+    try:
+        f = row["video_path"]
+        v, s, e = row[args.video_id_key], float(row[args.start_time_key]), float(row[args.end_time_key])
+        
+        if args.no_round_times:
+            clip_filename = f"{v}_{s}_{e}.{args.ext}"
+        else:
+            clip_filename = f"{v}_{np.round(s, 1)}_{np.round(e, 1)}.{args.ext}"
+            
+        clip_filepath = join(args.cut_dir, clip_filename)
+        os.makedirs(os.path.dirname(clip_filepath), exist_ok=True)
+
+        if os.path.exists(clip_filepath) and not args.overwrite:
+            return None
+
+        # Load video and extract clip
+        with VideoFileClip(f) as video:
+            # Calculate target width maintaining aspect ratio with max height 480
+            aspect_ratio = video.w / video.h
+            target_height = 480
+            target_width = int(target_height * aspect_ratio)
+            
+            # Extract and resize clip
+            clip = video.subclip(s, e)
+            clip = clip.resize(width=target_width, height=target_height)
+            
+            # Write clip with optimized settings
+            clip.write_videofile(
+                clip_filepath,
+                codec='libx264',
+                audio_codec='aac',
+                preset='faster',  # Faster encoding
+                threads=2,  # Use multiple threads for encoding
+                logger=None if not args.verbose else None
+            )
+            
+        return clip_filepath
+    except Exception as e:
+        if args.verbose:
+            print(f"Error processing {row[args.video_id_key]}: {str(e)}")
+        return None
+
+if __name__ == "__main__":
+    import argparse
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--csv", type=str, required=True,
+        help="Path to CSV file containing video IDs and timestamps",
+    )
+    parser.add_argument(
+        "--video_id_key", type=str, default="video_id",
+    )
+    parser.add_argument(
+        "--start_time_key", type=str, default="start_time",
+    )
+    parser.add_argument(
+        "--end_time_key", type=str, default="end_time",
+    )
+    parser.add_argument(
+        "--video_dir", type=str, required=True,
+        help="Path to directory containing downloaded videos",
+    )
+    parser.add_argument(
+        "--cut_dir", type=str, required=True,
+        help="Path to directory where cut videos will be saved",
+    )
+    parser.add_argument(
+        "--overwrite", action="store_true",
+        help="Whether to overwrite existing cut videos",
+    )
+    parser.add_argument(
+        "--verbose", action="store_true",
+    )
+    parser.add_argument(
+        "--no_round_times", action="store_true",
+        help="Whether to round start and end times to nearest second in filenames",
+    )
+    parser.add_argument(
+        "--debug", action="store_true",
+    )
+    parser.add_argument(
+        "--ext", type=str, default="mp4",
+    )
+    parser.add_argument(
+        "--si", type=int, default=0,
+    )
+    parser.add_argument(
+        "--ei", type=int, default=None,
+    )
+    parser.add_argument(
+        "--filter_csv", type=str, default=None, required=False,
+    )
+    parser.add_argument(
+        "--filter_key", type=str, default=None, required=False,
+    )
+    parser.add_argument(
+        "--max_workers", type=int, default=4,
+        help="Number of parallel workers for processing videos",
+    )
+    args = parser.parse_args()
+    
+    # Make cut_dir
+    os.makedirs(args.cut_dir, exist_ok=True)
+    
+    # Load and filter CSV
+    assert os.path.exists(args.csv), f"CSV file {args.csv} does not exist."
+    df = pd.read_csv(args.csv)
+    print(">>> Loaded CSV file with shape", df.shape)
+    assert {args.video_id_key, args.start_time_key, args.end_time_key}.issubset(df.columns)
+
+    # Filter CSV if needed
+    if args.filter_csv is not None:
+        path = args.filter_csv
+        assert os.path.exists(path), f"CSV file {path} does not exist."
+        key = args.filter_key
+        df_filter = pd.read_csv(path)
+        assert key in df_filter.columns, f"CSV file must contain column {key}."
+        keep_values = df_filter[key].unique()
+        df = df[df[key].isin(keep_values)]
+        print(">>> Filtered CSV file with shape", df.shape)
+    
+    # Apply index slicing
+    si = args.si
+    ei = args.ei if args.ei is not None else len(df)
+    df = df.iloc[si:ei]
+    print("Start index:", si, "End index:", ei)
+    
+    # More efficient way to add video path
+    print(">>> Adding video paths to dataframe")
+    video_ids = df[args.video_id_key].unique()
+    video_paths = [join(args.video_dir, f"{video_id}.{args.ext}") for video_id in video_ids]
+    video_id_to_path = {video_id: path for video_id, path in zip(video_ids, video_paths)}
+    df["video_path"] = df[args.video_id_key].map(video_id_to_path)
+    # df = df[df["video_path"].apply(exists)]
+    df['check_video'] = df['video_path'].apply(exists)
+    df = df[df['check_video']]
+    del df['check_video']
+    print(">>> Found videos for", df.shape[0], "rows.")
+
+    # # Filter out videos that don't exist
+    # df["video_path"] = df[args.video_id_key].apply(
+    #     lambda video_id: join(args.video_dir, f"{video_id}.{args.ext}"),
+    # )
+    # df["check_video"] = df["video_path"].apply(exists)
+    # df = df[df["check_video"]]
+    # del df["check_video"]
+    # print(">>> Found videos for", df.shape[0], "rows.")
+
+    if len(df) == 0:
+        print(">>> No videos to cut.")
+        exit()
+
+
+    if args.debug:
+        args.verbose = True
+        # Process only one video in debug mode
+        process_video(df.iloc[0], args)
+    else:
+        # Process videos in parallel
+        with ThreadPoolExecutor(max_workers=args.max_workers) as executor:
+            futures = [executor.submit(process_video, row, args) 
+                      for _, row in df.iterrows()]
+            
+            # Show progress bar
+            with tqdm(total=len(futures), desc="Cutting clips") as pbar:
+                for future in as_completed(futures):
+                    result = future.result()
+                    pbar.update(1)
+    
+    print(">>> Number of cut files:", len(os.listdir(args.cut_dir)))

shared/scripts/cut_multiple_clips.py

@@ -0,0 +1,357 @@
+"""Cuts multiple clips from a single video using ffmpeg."""
+import os
+from os.path import join, exists
+import numpy as np
+import pandas as pd
+from subprocess import call
+from tqdm import tqdm
+
+
+def cut_multiple_clips_video_only(
+        video_path: str,
+        start_times: list,
+        end_times: list,
+        save_dir: str,
+        video_id=None,
+        ext=None,
+        verbose=False,
+    ):
+    """Cuts multiple clips from a single video using ffmpeg.
+
+    Args:
+        video_path: Path to the video file.
+        start_times: List of start times for each clip.
+        end_times: List of end times for each clip.
+    """
+    os.makedirs(save_dir, exist_ok=True)
+    assert len(start_times) == len(end_times), \
+        'start_times and end_times must have the same length.'
+
+    if video_id is None:
+        video_id = os.path.basename(video_path).split(".")[0]
+    if ext is None:
+        ext = os.path.basename(video_path).split(".")[1]
+
+    item_ids = [
+        f"{video_id}_{np.round(s, 1)}_{np.round(e, 1)}" \
+            for (s, e) in zip(start_times, end_times)
+    ]
+
+    ins = [
+        f"[0:v]trim=start={s}:end={e},setpts=PTS-STARTPTS,scale=480:-1[v{i}]" \
+        for i, (s, e) in enumerate(zip(start_times, end_times))
+    ]
+    ins = ";".join(ins)
+    outs = [
+        f"-map [v{i}] {save_dir}/{item_ids[i]}.{ext}" \
+        for i in range(len(start_times))
+    ]
+    outs = " ".join(outs)
+    if not verbose:
+        suffix = "-loglevel panic"
+    else:
+        suffix = ""
+    command = f"""
+        ffmpeg -i {video_path} -filter_complex "{ins}" {outs} -y {suffix}
+    """
+    call(command, shell=True)
+
+
+def cut_multiple_clips_audio_and_video(
+        video_path: str,
+        start_times: list,
+        end_times: list,
+        save_dir: str,
+        video_id=None,
+        ext=None,
+        verbose=False,
+    ):
+    """Cuts multiple clips from a single video using ffmpeg.
+
+    Args:
+        video_path: Path to the video file.
+        start_times: List of start times for each clip.
+        end_times: List of end times for each clip.
+    """
+
+    if args.verbose:
+        print("[:::] Cutting clips from video: ", video_path)
+        print("[:::] Number of clips to cut: ", len(start_times))
+
+    os.makedirs(save_dir, exist_ok=True)
+    assert len(start_times) == len(end_times), \
+        'start_times and end_times must have the same length.'
+
+    if video_id is None:
+        video_id = os.path.basename(video_path).split(".")[0]
+    if ext is None:
+        ext = os.path.basename(video_path).split(".")[1]
+
+    item_ids = [
+        f"{video_id}_{np.round(s, 1)}_{np.round(e, 1)}" \
+            for (s, e) in zip(start_times, end_times)
+    ]
+
+    ins = [
+        f"[0:v]trim=start={s}:end={e},setpts=PTS-STARTPTS,scale=480:-1[v{i}];"\
+        f"[0:a:0]atrim=start={s}:end={e},asetpts=PTS-STARTPTS[a{i}]" \
+        for i, (s, e) in enumerate(zip(start_times, end_times))
+    ]
+    ins = ";".join(ins)
+    outs = [
+        f"-map [v{i}] -map [a{i}] {save_dir}/{item_ids[i]}.{ext}" \
+        for i in range(len(start_times))
+    ]
+    outs = " ".join(outs)
+    if not verbose:
+        suffix = "-loglevel panic"
+    else:
+        suffix = ""
+    command = f"""
+        ffmpeg -i {video_path} -filter_complex "{ins}" {outs} -y {suffix}
+    """
+    call(command, shell=True)
+
+
+
+
+def cut_multiple_clips_audio_and_video_v2(
+        video_path: str,
+        start_times: list,
+        end_times: list,
+        save_dir: str,
+        video_id=None,
+        ext=None,
+        verbose=False,
+    ):
+    """Cuts multiple clips from a single video using ffmpeg.
+
+    Args:
+        video_path: Path to the video file.
+        start_times: List of start times for each clip.
+        end_times: List of end times for each clip.
+    """
+
+    if args.verbose:
+        print("[:::] Cutting clips from video: ", video_path)
+        print("[:::] Number of clips to cut: ", len(start_times))
+
+    os.makedirs(save_dir, exist_ok=True)
+    assert len(start_times) == len(end_times), \
+        'start_times and end_times must have the same length.'
+
+    if video_id is None:
+        video_id = os.path.basename(video_path).split(".")[0]
+    if ext is None:
+        ext = os.path.basename(video_path).split(".")[1]
+
+    item_ids = [
+        f"{video_id}_{np.round(s, 1)}_{np.round(e, 1)}" \
+            for (s, e) in zip(start_times, end_times)
+    ]
+
+    ins = [
+        f"[0:v]trim=start={s}:end={e},setpts=PTS-STARTPTS,scale=480:-1[v{i}];"\
+        f"[0:a:0]atrim=start={s}:end={e},asetpts=PTS-STARTPTS[a{i}]" \
+        for i, (s, e) in enumerate(zip(start_times, end_times))
+    ]
+    # ins = ";".join(ins)
+    outs = [
+        f"-map [v{i}] -map [a{i}] {save_dir}/{item_ids[i]}.{ext}" \
+        for i in range(len(start_times))
+    ]
+    # outs = " ".join(outs)
+    if not verbose:
+        suffix = "-loglevel panic"
+    else:
+        suffix = ""
+
+    iterator = tqdm(range(len(start_times)), desc="Cutting clips for {}".format(video_id))
+    for i in iterator:
+        ins_ = ins[i]
+        outs_ = outs[i]
+        save_path = f"{save_dir}/{item_ids[i]}.{ext}"
+        os.makedirs(os.path.dirname(save_path), exist_ok=True)
+        if os.path.exists(save_path):
+            continue
+        command = f"""
+            ffmpeg -i {video_path} -filter_complex "{ins_}" {outs_} -y {suffix}
+        """
+        call(command, shell=True)
+
+
+if __name__ == "__main__":
+    import argparse
+    parser = argparse.ArgumentParser()
+    # General arguments
+    parser.add_argument("--sanity", action="store_true")
+    parser.add_argument("--debug", action="store_true")
+    parser.add_argument("--verbose", action="store_true")
+    parser.add_argument(
+        "--ext", type=str, default="mp4",
+    )
+    # Arguments for input CSV
+    parser.add_argument(
+        "--csv", type=str, required=True,
+        help="Path to CSV file containing video IDs and timestamps",
+    )
+    parser.add_argument(
+        "--video_id_key", type=str, default="video_id",
+    )
+    parser.add_argument(
+        "--start_time_key", type=str, default="start_time",
+    )
+    parser.add_argument(
+        "--end_time_key", type=str, default="end_time",
+    )
+    parser.add_argument(
+        "--video_dir", type=str, required=True,
+        help="Path to directory containing downloaded videos",
+    )
+    parser.add_argument(
+        "--cut_dir", type=str, required=True,
+        help="Path to directory where cut videos will be saved",
+    )
+    parser.add_argument(
+        "--overwrite", action="store_true",
+        help="Whether to overwrite existing cut videos",
+    )
+    parser.add_argument(
+        "--video_only", action="store_true",
+    )
+    parser.add_argument(
+        "--si", type=int, default=0,
+    )
+    parser.add_argument(
+        "--ei", type=int, default=None,
+    )
+    args = parser.parse_args()
+
+    if args.sanity:
+
+        # Test without audio
+        video_path = "sample_data/folding_paper.mp4"
+        start_times = [0, 5, 10]
+        end_times = [5, 10, 15]
+        cut_multiple_clips_video_only(
+            video_path,
+            start_times,
+            end_times,
+            "./sample_data/clips",
+            verbose=args.verbose,
+            ext=args.ext,
+        )
+
+        # Test with audio
+        video_path = "sample_data/pouring_water_youtube.mp4"
+        start_times = [0, 5, 10]
+        end_times = [5, 10, 15]
+        cut_multiple_clips_audio_and_video(
+            video_path,
+            start_times,
+            end_times,
+            "./sample_data/clips",
+            verbose=args.verbose,
+            ext=args.ext,
+        )
+
+    else:
+
+        # Make cut_dir
+        os.makedirs(args.cut_dir, exist_ok=True)
+        
+        # Load csv
+        assert os.path.exists(args.csv), f"CSV file {args.csv} does not exist."
+        df = pd.read_csv(args.csv)
+        print(">>> Loaded CSV file with shape", df.shape)
+        keys = [args.video_id_key, args.start_time_key, args.end_time_key]
+        assert set(keys).issubset(df.columns), \
+            f"CSV file must contain columns {keys}."
+        
+        # Filter out videos that don't exist
+        df["video_path"] = df[args.video_id_key].apply(
+            lambda video_id: join(args.video_dir, f"{video_id}.{args.ext}"),
+        )
+        df["check_video"] = df["video_path"].apply(exists)
+        df = df[df["check_video"]]
+        del df["check_video"]
+        print(">>> Found videos for", df.shape[0], "rows.")
+
+        si = args.si
+        ei = args.ei if args.ei is not None else df.shape[0]
+        print("Running from indices", si, "to", ei)
+        df = df.iloc[si:ei]
+
+        if args.debug:
+            args.verbose = True
+        ext = args.ext
+
+        # Iterate over each video
+        video_paths = df["video_path"].unique()
+        # iterator = tqdm(range(len(video_paths)), desc="Cutting clips")
+        print("Number of unique videos:", len(video_paths))
+        for i in range(len(video_paths)):
+            video_path = video_paths[i]
+
+            # Find rows corresponding to this video
+            df_video = df[df["video_path"] == video_path]
+            print("Number of clips to cut from video", video_path, ":", df_video.shape[0])
+            start_times = df_video[args.start_time_key].values
+            end_times = df_video[args.end_time_key].values
+            video_id = df_video[args.video_id_key].values[0]
+
+            """
+            # Cut to MAXLEN clips per video
+            MAX_LEN = 10
+            start_times_batches = np.array_split(start_times, MAX_LEN)
+            end_times_batches = np.array_split(end_times, MAX_LEN)
+            for start_times_, end_times_ in zip(start_times_batches, end_times_batches):
+                if args.video_only:
+                    cut_multiple_clips_video_only(
+                        video_path,
+                        start_times_,
+                        end_times_,
+                        args.cut_dir,
+                        video_id=video_id,
+                        ext=ext,
+                        verbose=args.verbose,
+                        # verbose=True,
+                    )
+                else:
+                    cut_multiple_clips_audio_and_video_v2(
+                        video_path,
+                        start_times_,
+                        end_times_,
+                        args.cut_dir,
+                        video_id=video_id,
+                        ext=ext,
+                        verbose=args.verbose,
+                        # verbose=True,
+                    )
+            """
+            # """
+            # Cut videos
+            if args.video_only:
+                cut_multiple_clips_video_only(
+                    video_path,
+                    start_times,
+                    end_times,
+                    args.cut_dir,
+                    video_id=video_id,
+                    ext=ext,
+                    verbose=args.verbose,
+                )
+            else:
+                cut_multiple_clips_audio_and_video_v2(
+                    video_path,
+                    start_times,
+                    end_times,
+                    args.cut_dir,
+                    video_id=video_id,
+                    ext=ext,
+                    verbose=args.verbose,
+                )
+            # """
+            
+            if args.debug:
+                break

shared/scripts/downscale_videos.py

@@ -0,0 +1,274 @@
+#!/usr/bin/env python3
+
+import argparse
+import sys
+from pathlib import Path
+from typing import List
+
+import cv2
+from tqdm import tqdm
+from joblib import Parallel, delayed
+from contextlib import contextmanager
+from joblib import parallel as joblib_parallel
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(
+        description="Downscale videos and save them to a new folder with a given extension."
+    )
+    parser.add_argument(
+        "--video_dir",
+        type=Path,
+        default=Path("/scratch/shared/beegfs/piyush/datasets/NTU/nturgb+d_rgb/"),
+        help="Root directory containing videos (searched recursively)",
+    )
+    parser.add_argument(
+        "--ext",
+        type=str,
+        default="avi",
+        help="Video file extension to search for (without dot)",
+    )
+    parser.add_argument(
+        "--downscale_factor",
+        type=float,
+        default=0.4,
+        help="Factor by which to downscale width and height (e.g., 0.4)",
+    )
+    parser.add_argument(
+        "--save_dir",
+        type=Path,
+        default=None,
+        help="Directory to save downscaled videos. Defaults to video_dir + '-downscaled={factor}'",
+    )
+    parser.add_argument(
+        "--save_ext",
+        type=str,
+        default="mp4",
+        help="Extension to save resulting videos with (without dot)",
+    )
+    parser.add_argument(
+        "--debug",
+        action="store_true",
+        help="Process only the first 10 videos and print saved paths",
+    )
+    parser.add_argument(
+        "--si",
+        type=int,
+        default=0,
+        help="Start index (inclusive) in the sorted video list",
+    )
+    parser.add_argument(
+        "--ei",
+        type=int,
+        default=None,
+        help="End index (exclusive) in the sorted video list; None means till end",
+    )
+    parser.add_argument(
+        "--n_jobs",
+        type=int,
+        default=-1,
+        help="Number of parallel jobs (-1 uses all cores)",
+    )
+    args = parser.parse_args()
+
+    if args.save_dir is None:
+        args.save_dir = Path(f"{args.video_dir}-downscaled={args.downscale_factor}")
+
+    # Normalize extensions (strip leading dots)
+    args.ext = args.ext.lstrip('.')
+    args.save_ext = args.save_ext.lstrip('.')
+    return args
+
+
+def ensure_even_dimension(value: int) -> int:
+    if value < 1:
+        return 1
+    return value if value % 2 == 0 else value - 1 if value > 1 else 1
+
+
+def list_videos(video_dir: Path, ext: str) -> List[Path]:
+    pattern = f"**/*.{ext}"
+    return sorted(video_dir.rglob(pattern))
+
+
+def change_extension(path: Path, new_ext: str) -> Path:
+    return path.with_suffix('.' + new_ext)
+
+
+def downscale_video(
+    src_path: Path,
+    dst_path: Path,
+    downscale_factor: float,
+    save_ext: str,
+) -> bool:
+    cap = cv2.VideoCapture(str(src_path))
+    if not cap.isOpened():
+        return False
+
+    # Read properties
+    fps = cap.get(cv2.CAP_PROP_FPS)
+    if fps <= 0 or fps != fps:  # NaN check
+        fps = 30.0
+    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+
+    out_w = ensure_even_dimension(max(1, int(width * downscale_factor)))
+    out_h = ensure_even_dimension(max(1, int(height * downscale_factor)))
+
+    # Choose FOURCC based on extension
+    save_ext_lower = save_ext.lower()
+    if save_ext_lower in {"mp4", "m4v"}:
+        fourcc = cv2.VideoWriter_fourcc(*"mp4v")
+    elif save_ext_lower in {"avi"}:
+        fourcc = cv2.VideoWriter_fourcc(*"XVID")
+    elif save_ext_lower in {"mov"}:
+        fourcc = cv2.VideoWriter_fourcc(*"mp4v")
+    else:
+        # Fallback
+        fourcc = cv2.VideoWriter_fourcc(*"mp4v")
+
+    dst_path.parent.mkdir(parents=True, exist_ok=True)
+    writer = cv2.VideoWriter(str(dst_path), fourcc, fps, (out_w, out_h))
+    if not writer.isOpened():
+        cap.release()
+        return False
+
+    ok = True
+    try:
+        while True:
+            ret, frame = cap.read()
+            if not ret:
+                break
+            resized = cv2.resize(frame, (out_w, out_h), interpolation=cv2.INTER_AREA)
+            writer.write(resized)
+    except Exception:
+        ok = False
+    finally:
+        writer.release()
+        cap.release()
+
+    # If nothing written, treat as failure
+    if dst_path.exists() and dst_path.stat().st_size > 0 and ok:
+        return True
+    try:
+        if dst_path.exists():
+            dst_path.unlink()
+    except Exception:
+        pass
+    return False
+
+
+def process_one(
+    src: Path,
+    dst: Path,
+    downscale_factor: float,
+    save_ext: str,
+) -> tuple:
+    success = downscale_video(src, dst, downscale_factor, save_ext)
+    return success, dst
+
+
+@contextmanager
+def tqdm_joblib(tqdm_object):
+    """Context manager to patch joblib to report into tqdm progress bar."""
+    class TqdmBatchCompletionCallback(joblib_parallel.BatchCompletionCallBack):
+        def __call__(self, *args, **kwargs):
+            tqdm_object.update(n=self.batch_size)
+            return super().__call__(*args, **kwargs)
+
+    old_cb = joblib_parallel.BatchCompletionCallBack
+    joblib_parallel.BatchCompletionCallBack = TqdmBatchCompletionCallback
+    try:
+        yield tqdm_object
+    finally:
+        joblib_parallel.BatchCompletionCallBack = old_cb
+        try:
+            tqdm_object.close()
+        except Exception:
+            pass
+
+
+def main() -> int:
+    args = parse_args()
+
+    video_dir: Path = args.video_dir
+    save_dir: Path = args.save_dir
+    ext: str = args.ext
+    save_ext: str = args.save_ext
+    downscale_factor: float = args.downscale_factor
+    si: int = max(0, int(args.si))
+    ei = args.ei if args.ei is None else max(0, int(args.ei))
+    n_jobs: int = int(args.n_jobs)
+
+    if not video_dir.exists() or not video_dir.is_dir():
+        print(f"ERROR: video_dir does not exist or is not a directory: {video_dir}", file=sys.stderr)
+        return 1
+
+    videos = list_videos(video_dir, ext)
+
+    # Slice by [si:ei]
+    try:
+        videos = videos[si:ei]
+    except Exception:
+        # Fallback if indexing fails
+        videos = []
+
+    if len(videos) == 0:
+        print("No videos found.")
+        return 0
+
+    # In debug mode, limit to first 10 from the sliced list
+    if args.debug:
+        videos = videos[:10]
+
+    # Build (src, dst) pairs and count already existing
+    tasks = []
+    skipped_count = 0
+    for src in videos:
+        try:
+            rel = src.relative_to(video_dir)
+        except ValueError:
+            rel = src.name
+        rel_path = Path(rel)
+        rel_with_new_ext = change_extension(rel_path, save_ext)
+        dst = (save_dir / rel_with_new_ext).resolve()
+        if dst.exists() and dst.stat().st_size > 0:
+            skipped_count += 1
+            continue
+        tasks.append((src, dst))
+
+    total_count = len(videos)
+    saved_paths: List[Path] = []
+    errors: int = 0
+
+    with tqdm(total=total_count, desc="Downscaling videos", unit="vid") as pbar:
+        # account for already existing outputs
+        if skipped_count:
+            pbar.update(skipped_count)
+
+        if len(tasks) > 0:
+            with tqdm_joblib(pbar):
+                results = Parallel(n_jobs=n_jobs, backend="loky")( \
+                    delayed(process_one)(src, dst, downscale_factor, save_ext) for (src, dst) in tasks
+                )
+            for success, dst in results:
+                if success:
+                    saved_paths.append(dst)
+                else:
+                    errors += 1
+    
+    if args.debug:
+        print("Saved (debug mode):")
+        for p in saved_paths:
+            print(str(p))
+
+    if errors > 0:
+        print(f"Completed with {errors} failures out of {len(videos)}", file=sys.stderr)
+        return 2
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())
+
+

shared/scripts/downsize_videos.py

@@ -0,0 +1,194 @@
+"""Downsize videos preserving aspect resolution."""
+import torch
+import torchvision
+from decord import VideoReader
+from glob import glob
+import os
+from os.path import join, basename, exists
+import subprocess
+import numpy as np
+import pandas as pd
+import subprocess
+import ffmpeg
+import time
+import librosa
+from moviepy.editor import VideoFileClip
+
+import shared.utils.log as log
+import shared.utils.io as io
+
+
+# def downsize(input_path, output_path, width=480, height=None, maintain_aspect_ratio=True):
+#     """Downsizes a given video."""
+
+#     # Check if the video exists
+#     assert exists(input_path), f"Video {input_path} does not exist."
+    
+#     # Define ffmpeg command to downsize video with width=480 maintaining aspect ratio
+#     # And save it at output_path
+#     if maintain_aspect_ratio:
+#         assert height is None, "Cannot specify height when maintaining aspect ratio."
+#         height = -1
+
+#     (
+#         ffmpeg
+#         .input(input_path)
+#         .output(output_path, preset="ultrafast", vf=f"scale={width}:{height}", loglevel="quiet")
+#         .run()
+#     )
+
+def resize_video_maintain_aspect_ratio(input_path, output_path):
+    (
+        ffmpeg
+        .input(input_path)
+        .filter("scale", w=480, h=-2)
+        .output(output_path, crf=18, preset="ultrafast", loglevel="quiet")
+        .run()
+    )
+
+
+def resize_video_maintain_aspect_ratio_faster(input_path, output_path, width=480):
+    (
+        ffmpeg
+        .input(input_path)
+        .filter('scale', width, -1)
+        .output(output_path, vcodec='h264_nvenc', preset='fast', pix_fmt='yuv420p')
+        .overwrite_output()
+        .run(capture_stdout=True)
+    )
+
+
+def resize_video_maintain_aspect_ratio_vanilla(input_path, output_path, width=480):
+    (
+        ffmpeg
+        .input(input_path)
+        .filter('scale', width, -1)
+        .output(output_path, c="copy")
+        .overwrite_output()
+        .run(capture_stdout=True)
+    )
+
+
+def resize_video_moviepy(input_path, output_path, width=480):
+    # Load the input video
+    video = VideoFileClip(input_path)
+    
+    # Resize the video
+    video_resized = video.resize(width=width)
+    
+    # Save the resized video
+    video_resized.write_videofile(output_path)
+
+
+
+def resize_video_simple(input_path, output_path, width=480):
+    command = f"""ffmpeg -loglevel quiet -i {input_path} -vf "scale={width}:-1" -c:a copy {output_path} -y"""
+    subprocess.call(command, shell=True)
+
+
+if __name__ == "__main__":
+    import argparse
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--csv", type=str, required=True,
+        help="Path to csv file containing in/out video paths."
+    )
+    parser.add_argument(
+        "--in_colname", type=str, default="input",
+        help="column name of input videos.",
+    )
+    parser.add_argument(
+        "--out_colname", type=str, default="output",
+        help="column name of output videos.",
+    )
+    parser.add_argument(
+        "--remove_old", action="store_true",
+    )
+    parser.add_argument(
+        "--width", type=int, default=480,
+    )
+    parser.add_argument(
+        "--debug", action="store_true",
+    )
+    parser.add_argument(
+        "--si", type=int, default=None,
+        help="Start index.",
+    )
+    parser.add_argument(
+        "--ei", type=int, default=None,
+        help="End index.",
+    )
+    parser.add_argument('--overwrite', action='store_true')
+    args = parser.parse_args()
+    
+    print("Width:", args.width)
+    assert exists(args.csv), f"File {args.csv} does not exist."
+    df = pd.read_csv(args.csv)
+    print("> Number of videos:", len(df))
+    
+    si = args.si if args.si is not None else 0
+    ei = args.ei if args.ei is not None else len(df)
+    print("> Start index:", si)
+    print("> End index:", ei)
+    df = df.iloc[si:ei]
+    print("> Number of videos to downsize:", len(df))
+    
+    ifiles = df[args.in_colname].tolist()
+    ofiles = df[args.out_colname].tolist()
+    assert len(ifiles) == len(ofiles), \
+        "Number of input and output videos must be the same."
+
+    iterator = log.tqdm_iterator(
+        range(len(ifiles)), total=len(ifiles), desc="Downsizing videos",
+    )
+    for i in iterator:
+        ifile, ofile = ifiles[i], ofiles[i]
+        
+        # If ofile == ifile (i.e., edit the same file),
+        # then we need to operate on a temporary file
+        # which will then be moved to the original file
+        replace = ofile == ifile
+        if replace:
+            ofile_actual = ofile
+            ofile = ofile.replace(".mp4", "_temp.mp4")
+
+        # Check if the video exists
+        assert exists(ifile), f"Video {ifile} does not exist."
+
+        # If output file already exists, skip
+        if exists(ofile) and not args.overwrite:
+            continue
+
+        # Make sure output directory exists
+        os.makedirs(os.path.dirname(ofile), exist_ok=True)
+
+        # resize
+        start_time = time.time()
+        resize_video_simple(ifile, ofile, width=args.width)
+        end_time = time.time()
+        time_taken = end_time - start_time
+        desc = "Time taken {:.2f}s for {}".format(time_taken, basename(ifile))
+        iterator.set_description(desc)
+        
+        # If replace, move the temporary file to the original file
+        if replace:
+            if exists(ofile):
+                os.rename(ofile, ofile_actual)
+            ofile = ofile_actual
+        
+        if args.debug:
+            yold, srold = librosa.load(ifile, offset=1.0, duration=1.0)
+            ynew, srnew = librosa.load(ofile, offset=1.0, duration=1.0)
+            assert srold == srnew, "Sampling rate mismatch."
+            assert len(yold) == len(ynew), "Length mismatch."
+            assert np.allclose(yold, ynew), "Audio mismatch."
+            
+            # Try loading the new video file
+            vr = VideoReader(ofile)
+            frames = vr.get_batch(range(0, 10)).asnumpy()
+            assert frames.shape[0] == 10, "Length mismatch."
+            assert frames.shape[2] == 480, "Width mismatch."
+        
+        # If remove_old, remove old video
+        if args.remove_old:
+            os.remove(ifile)

shared/scripts/downsize_videos_simple.py

@@ -0,0 +1,150 @@
+import os
+from glob import glob
+import time
+import argparse
+import numpy as np
+import librosa
+from decord import VideoReader
+import shared.utils.log as log
+from os.path import exists, basename
+from natsort import natsorted
+
+def resize_video_simple(input_path, output_path, width=480):
+    import subprocess
+    command = f"""ffmpeg -loglevel quiet -i {input_path} -vf \"scale={width}:-1\" -c:a copy {output_path}"""
+    subprocess.call(command, shell=True)
+
+def load_pending_videos(tracker_file):
+    """Load list of pending videos from tracker file."""
+    if not exists(tracker_file):
+        return []
+    
+    with open(tracker_file, 'r') as f:
+        return [line.strip() for line in f.readlines() if line.strip()]
+
+def save_pending_videos(tracker_file, video_paths):
+    """Save list of pending videos to tracker file."""
+    with open(tracker_file, 'w') as f:
+        for path in video_paths:
+            f.write(f"{path}\n")
+
+def remove_completed_video(tracker_file, completed_video):
+    """Remove a completed video from the tracker file."""
+    pending_videos = load_pending_videos(tracker_file)
+    if completed_video in pending_videos:
+        pending_videos.remove(completed_video)
+        save_pending_videos(tracker_file, pending_videos)
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--video_dir", type=str, required=True,
+        help="Directory containing videos to downsize."
+    )
+    parser.add_argument(
+        "--ext", type=str, default="mp4",
+        help="File extension to search for (default: mp4)."
+    )
+    parser.add_argument(
+        "--remove_old", action="store_true",
+        help="Remove original video after downsizing."
+    )
+    parser.add_argument(
+        "--width", type=int, default=480,
+        help="Width to resize videos to (default: 480)."
+    )
+    parser.add_argument(
+        "--debug", action="store_true",
+        help="Run debug checks after downsizing."
+    )
+    parser.add_argument("--si", type=int, default=None)
+    parser.add_argument("--ei", type=int, default=None)
+    parser.add_argument(
+        "--tracker_file", type=str, default="video_resize_tracker.txt",
+        help="Tracker file to keep track of pending videos (default: video_resize_tracker.txt)."
+    )
+    parser.add_argument(
+        "--reset_tracker", action="store_true",
+        help="Reset the tracker file and start fresh."
+    )
+    args = parser.parse_args()
+
+    assert os.path.isdir(args.video_dir), f"Directory {args.video_dir} does not exist."
+    
+    # Handle tracker file
+    if args.reset_tracker and exists(args.tracker_file):
+        os.remove(args.tracker_file)
+        print(f"> Reset tracker file: {args.tracker_file}")
+    
+    # Check if we have pending videos from previous run
+    pending_videos = load_pending_videos(args.tracker_file)
+    
+    if pending_videos:
+        print(f"> Found {len(pending_videos)} pending videos from previous run")
+        ifiles = pending_videos
+        ofiles = pending_videos  # In-place replacement
+    else:
+        # Start fresh - find all videos
+        pattern = os.path.join(args.video_dir, f"**/*.{args.ext}")
+        ifiles = glob(pattern, recursive=True)
+        ifiles = natsorted(ifiles)
+        ofiles = ifiles  # In-place replacement
+        print("> Number of videos in the directory:", len(ifiles))
+        
+        # Apply start/end index filtering
+        si = args.si if args.si is not None else 0
+        ei = args.ei if args.ei is not None else len(ifiles)
+        print("> Start index:", si)
+        print("> End index:", ei)
+        ifiles = ifiles[si:ei]
+        ofiles = ofiles[si:ei]
+        
+        # Save to tracker file for future runs
+        save_pending_videos(args.tracker_file, ifiles)
+        print(f"> Saved {len(ifiles)} videos to tracker file: {args.tracker_file}")
+    
+    print("> Number of videos to downsize:", len(ifiles))
+
+    iterator = log.tqdm_iterator(
+        range(len(ifiles)), total=len(ifiles), desc="Downsizing videos",
+    )
+    for i in iterator:
+        ifile = ifiles[i]
+        ofile = ifile
+        assert exists(ifile), f"Video {ifile} does not exist."
+        tmp_ofile = ifile + ".tmp.mp4"
+        start_time = time.time()
+        resize_video_simple(ifile, tmp_ofile, width=args.width)
+        end_time = time.time()
+        time_taken = end_time - start_time
+        desc = f"Time taken {time_taken:.2f}s for {basename(ifile)}"
+        iterator.set_description(desc)
+
+        if args.debug:
+            yold, srold = librosa.load(ifile, offset=1.0, duration=1.0)
+            ynew, srnew = librosa.load(tmp_ofile, offset=1.0, duration=1.0)
+            assert srold == srnew, "Sampling rate mismatch."
+            assert len(yold) == len(ynew), "Length mismatch."
+            assert np.allclose(yold, ynew), "Audio mismatch."
+            vr = VideoReader(tmp_ofile)
+            frames = vr.get_batch(range(0, 10)).asnumpy()
+            assert frames.shape[0] == 10, "Length mismatch."
+            assert frames.shape[2] == args.width, "Width mismatch."
+
+        # Replace original file
+        os.replace(tmp_ofile, ofile)
+        
+        # Remove completed video from tracker
+        remove_completed_video(args.tracker_file, ifile)
+        
+        if args.remove_old:
+            # Already replaced, so nothing to remove
+            pass
+    
+    # Clean up tracker file if all videos are done
+    if not load_pending_videos(args.tracker_file):
+        os.remove(args.tracker_file)
+        print(f"> All videos completed. Removed tracker file: {args.tracker_file}")
+    else:
+        remaining = len(load_pending_videos(args.tracker_file))
+        print(f"> {remaining} videos remaining. Tracker file preserved: {args.tracker_file}") 

shared/scripts/extract_speed_clips.py

@@ -0,0 +1,126 @@
+import os
+import sys
+from glob import glob 
+from collections import defaultdict
+
+import json
+from hydra import compose, initialize
+from omegaconf import OmegaConf
+import PIL, PIL.Image
+import decord
+
+from adapt4change.utils.speednet import *
+import shared.utils as su
+
+classes_selected = [
+    "skipping rope",
+    "gymnastics tumbling",
+    "somersaulting",
+    "cartwheeling",
+    "trampolines bouncing",
+    "swinging on something",
+    "vault",
+    "deadlifting",
+    "clean and jerk",
+    "diving cliff",
+]
+
+SAVE_DIR = "/scratch/shared/beegfs/piyush/datasets/SpeedyKinetics/clips"
+os.makedirs(SAVE_DIR, exist_ok=True)
+
+def check_all_files_exist(save_paths):
+    for path in save_paths:
+        if not os.path.exists(path):
+            return False
+    return True
+
+def save_clips_for_single_video(video_path, show=False):
+    """
+    Note that I randomly sample 3s clips out of the original video.
+    """
+    video_id = os.path.basename(video_path).split(".mp4")[0]
+    save_paths = [
+        f"{SAVE_DIR}/{video_id}-normal.mp4",
+        f"{SAVE_DIR}/{video_id}-spedup.mp4",
+        f"{SAVE_DIR}/{video_id}-slowdn.mp4",
+    ]
+    if check_all_files_exist(save_paths):
+        return
+
+    try:
+        vr = decord.VideoReader(video_path)
+    except Exception as e:
+        print(f"Error opening video {video_path}: {e}")
+        return
+
+    total_frames = len(vr)
+    fps = vr.get_avg_fps()
+
+    # Initialize sampler for a video 
+    sampler = FrameIndexSampler(total_frames=total_frames)
+    
+    # Sample clips 
+    clip_duration = 3.
+    T = int(clip_duration * fps)
+    start_frame = sampler.get_valid_start_frame(T)
+
+    # Get all clip indices
+    normal_indices, sped_up_indices, slowed_down_indices = sampler.sample_all_clip_indices(start_frame, T)
+
+    try:
+        frames_normal = [PIL.Image.fromarray(f) for f in vr.get_batch(normal_indices).asnumpy()]
+        frames_spedup = [PIL.Image.fromarray(f) for f in vr.get_batch(sped_up_indices).asnumpy()]
+        frames_slowdn = [PIL.Image.fromarray(f) for f in vr.get_batch(slowed_down_indices).asnumpy()]
+    except Exception as e:
+        print(f"Error processing video {video_path}: {e}")
+        return
+
+    su.io.save_video(frames_normal, save_paths[0], fps=vr.get_avg_fps())
+    su.io.save_video(frames_spedup, save_paths[1], fps=vr.get_avg_fps())
+    su.io.save_video(frames_slowdn, save_paths[2], fps=vr.get_avg_fps())
+    if show:
+        su.visualize.show_grid_of_videos(
+            files=save_paths,
+            labels=["Normal", "Sped up", "Slowed down"],
+        )
+
+
+if __name__ == "__main__":
+    import argparse
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--start_index", type=int, default=0)
+    parser.add_argument("--end_index", type=int, default=1000000)
+    args = parser.parse_args()
+
+    data_dir = "/datasets/KineticsClean/"
+    verbose = True
+    total_train = []
+    total_valid = []
+    for c in classes_selected:
+        files_train = glob(f"{data_dir}/train_split/{c}/*.mp4")
+        files_valid = glob(f"{data_dir}/val_split/{c}/*.mp4")
+        if verbose:
+            print(c)
+            print("Train videos: ", len(files_train))
+            print("Valid videos: ", len(files_valid))
+            print("-" * 80)
+        total_train.extend(files_train)
+        total_valid.extend(files_valid)
+    print("Total train files: ", len(total_train))
+    print("Total valid files: ", len(total_valid))
+    
+    files = total_train + total_valid
+    print(f"Total files: {len(files)}")
+    print(f"Start index: {args.start_index}")
+    print(f"End index: {args.end_index}")
+    files = files[args.start_index:args.end_index]
+    print(f"Total files to process: {len(files)}")
+    
+    parallelize = True
+    if not parallelize:
+        for file in su.log.tqdm_iterator(files, desc="Processing files"):
+            save_clips_for_single_video(file)
+    else:
+        from joblib import Parallel, delayed
+        iterator = su.log.tqdm_iterator(files, desc="Processing files")
+        Parallel(n_jobs=16)(delayed(save_clips_for_single_video)(file) for file in iterator)

shared/scripts/save_grid_of_videos.py

@@ -0,0 +1,203 @@
+#!/usr/bin/env python3
+"""
+Video Grid Visualizer (imageio-based with Grid)
+
+Creates a grid layout of multiple videos with visual separators and saves it
+as an MP4 or GIF. This version uses the imageio library for robust video
+writing, adds customizable grid gaps, and intelligently resizes cells for a
+balanced look.
+"""
+
+import cv2
+import numpy as np
+import argparse
+import os
+import imageio.v2 as imageio # Use imageio.v2 to avoid deprecation warnings
+
+def get_video_info(video_path):
+    """Get video information using imageio with a fallback to OpenCV."""
+    try:
+        with imageio.get_reader(video_path) as reader:
+            meta = reader.get_meta_data()
+            fps = meta.get('fps', 30)
+            duration = meta.get('duration', 0)
+            size = meta.get('size', (0, 0))
+            if duration == 0 and fps > 0:
+                duration = reader.count_frames() / fps
+            frame_count = int(duration * fps) if duration and fps else reader.count_frames()
+            
+            return {
+                'duration': duration, 'fps': fps, 'frame_count': frame_count,
+                'width': size[0], 'height': size[1]
+            }
+    except Exception:
+        cap = cv2.VideoCapture(video_path)
+        if not cap.isOpened():
+            raise ValueError(f"Cannot open video: {video_path}")
+        
+        fps = cap.get(cv2.CAP_PROP_FPS)
+        frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+        width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+        height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+        duration = frame_count / fps if fps > 0 else 0
+        cap.release()
+        return {
+            'duration': duration, 'fps': fps, 'frame_count': frame_count,
+            'width': width, 'height': height
+        }
+
+def resize_frame(frame, target_width, target_height):
+    """Resize frame to fit inside target dimensions, maintaining aspect ratio."""
+    h, w = frame.shape[:2]
+    aspect = w / h
+    target_aspect = target_width / target_height
+
+    if aspect > target_aspect:
+        new_w = target_width
+        new_h = int(new_w / aspect)
+    else:
+        new_h = target_height
+        new_w = int(new_h * aspect)
+        
+    resized = cv2.resize(frame, (new_w, new_h), interpolation=cv2.INTER_AREA)
+    
+    canvas = np.zeros((target_height, target_width, 3), dtype=np.uint8)
+    y_offset = (target_height - new_h) // 2
+    x_offset = (target_width - new_w) // 2
+    canvas[y_offset:y_offset + new_h, x_offset:x_offset + new_w] = resized
+    
+    return canvas
+
+def create_video_grid(video_paths, n_rows=1, n_cols=None, save_path="output.mp4", 
+                      gap=10, grid_color_name="white", max_cell_width=640):
+    """
+    Create a grid of videos with separators and save as MP4 or GIF.
+    """
+    if not video_paths:
+        raise ValueError("No video paths provided")
+    
+    if n_cols is None:
+        n_cols = len(video_paths)
+    
+    if n_rows * n_cols < len(video_paths):
+        raise ValueError(f"Grid size ({n_rows}x{n_cols}) is too small for {len(video_paths)} videos")
+
+    grid_color = (255, 255, 255) if grid_color_name.lower() == "white" else (0, 0, 0)
+
+    print("Gathering video information...")
+    video_infos = [get_video_info(path) for path in video_paths]
+    
+    # --- Intelligent Resizing and Dimension Calculation ---
+    valid_infos = [info for info in video_infos if info['width'] > 0 and info['height'] > 0]
+    if not valid_infos:
+        raise ValueError("Could not get valid dimensions from any input video.")
+        
+    avg_aspect_ratio = sum(info['width'] / info['height'] for info in valid_infos) / len(valid_infos)
+    
+    cell_width = min(max_cell_width, 1920 // n_cols) # Don't let cells get too big
+    cell_height = int(cell_width / avg_aspect_ratio)
+
+    output_width = (cell_width * n_cols) + (gap * (n_cols + 1))
+    output_height = (cell_height * n_rows) + (gap * (n_rows + 1))
+    
+    # Ensure dimensions are even, as required by many video codecs
+    output_width += output_width % 2
+    output_height += output_height % 2
+    # ---
+
+    max_duration = max(info['duration'] for info in video_infos if info)
+    target_fps = max(info['fps'] for info in video_infos if info and info['fps'])
+    if not target_fps or target_fps <= 0:
+        target_fps = 30
+    
+    total_frames = int(max_duration * target_fps)
+    
+    print(f"Grid: {n_rows}x{n_cols} with {gap}px {grid_color_name} gaps")
+    print(f"Calculated Cell Size: {cell_width}x{cell_height}")
+    print(f"Final Output Size: {output_width}x{output_height}")
+    print(f"Max duration: {max_duration:.2f}s | Target FPS: {target_fps} | Total frames: {total_frames}")
+
+    caps = [cv2.VideoCapture(path) for path in video_paths]
+    last_frames = [None] * len(video_paths)
+
+    writer = imageio.get_writer(save_path, fps=target_fps, codec='libx264', macro_block_size=None)
+
+    try:
+        for frame_idx in range(total_frames):
+            # Initialize the master frame with the grid color
+            output_frame = np.full((output_height, output_width, 3), grid_color, dtype=np.uint8)
+            
+            for i, cap in enumerate(caps):
+                ret, frame = cap.read()
+                if ret:
+                    last_frames[i] = frame
+                else:
+                    if last_frames[i] is None:
+                        info = video_infos[i]
+                        h = info['height'] if info['height'] > 0 else cell_height
+                        w = info['width'] if info['width'] > 0 else cell_width
+                        last_frames[i] = np.zeros((h, w, 3), dtype=np.uint8)
+                    frame = last_frames[i]
+                
+                resized_frame = resize_frame(frame, cell_width, cell_height)
+                
+                row = i // n_cols
+                col = i % n_cols
+                
+                # Calculate position with gaps
+                y_start = gap + row * (cell_height + gap)
+                x_start = gap + col * (cell_width + gap)
+                
+                output_frame[y_start : y_start + cell_height, x_start : x_start + cell_width] = resized_frame
+            
+            rgb_frame = cv2.cvtColor(output_frame, cv2.COLOR_BGR2RGB)
+            writer.append_data(rgb_frame)
+            
+            if frame_idx % int(target_fps) == 0 or frame_idx == total_frames - 1:
+                progress = (frame_idx + 1) / total_frames * 100
+                print(f"Processing... {progress:.1f}% complete", end='\r')
+
+    finally:
+        print("\nCleaning up resources...")
+        for cap in caps:
+            cap.release()
+        writer.close()
+    
+    print(f"Video grid successfully saved to: {save_path}")
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Create a grid visualization of multiple videos with separators.",
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter
+    )
+    parser.add_argument("video_paths", nargs="+", help="One or more paths to input videos.")
+    parser.add_argument("--n_rows", type=int, default=1, help="Number of rows in the grid.")
+    parser.add_argument("--n_cols", type=int, help="Number of columns. Defaults to the number of videos if n_rows is 1.")
+    parser.add_argument("--save_path", default="output.mp4", help="Output path for the video (e.g., 'output.mp4' or 'output.gif').")
+    parser.add_argument("--gap", type=int, default=10, help="Size of the gap between videos in pixels.")
+    parser.add_argument("--grid_color", default="white", choices=["white", "black"], help="Color of the grid gaps.")
+    parser.add_argument("--max_cell_width", type=int, default=640, help="Maximum width for each video cell in the grid.")
+
+    args = parser.parse_args()
+    
+    # Set default n_cols if not provided
+    if args.n_cols is None:
+        args.n_cols = len(args.video_paths) // args.n_rows
+        if len(args.video_paths) % args.n_rows != 0:
+            args.n_cols += 1
+
+    try:
+        create_video_grid(
+            video_paths=args.video_paths,
+            n_rows=args.n_rows,
+            n_cols=args.n_cols,
+            save_path=args.save_path,
+            gap=args.gap,
+            grid_color_name=args.grid_color,
+            max_cell_width=args.max_cell_width
+        )
+    except Exception as e:
+        print(f"\nAn error occurred: {e}")
+
+if __name__ == "__main__":
+    main()

shared/scripts/shard_video_dataset.py

@@ -0,0 +1,71 @@
+"""
+Creates multiple shards of videos in a dataset.
+
+Help:
+
+input_dir=/work/piyush/from_nfs2/datasets/SSv2/20bn-something-something-v2/
+output_dir=/work/piyush/from_nfs2/datasets/SSv2/ssv2_shards/
+python shared/scripts/shard_video_dataset.py -i $input_dir -o $output_dir
+"""
+import os
+import tarfile
+from glob import glob
+
+import shared.utils as su
+
+
+def read_args():
+    import argparse
+    parser = argparse.ArgumentParser()
+    # Input video directory
+    parser.add_argument("-i", "--input_dir", type=str, required=True)
+    parser.add_argument("--ext", type=str, default="webm")
+    # Output shard directory
+    parser.add_argument("-o", "--output_dir", type=str, required=True)
+    # Max size of each shard (n.o. videos)
+    parser.add_argument("--shard_size", type=int, default=10000)
+    args = parser.parse_args()
+    return args
+
+
+# Iterate over videos in your dataset
+def write_to_shard(tar, video_path, video_name):
+    with open(video_path, 'rb') as f:
+        tarinfo = tarfile.TarInfo(name=video_name)
+        tarinfo.size = os.path.getsize(video_path)
+        tar.addfile(tarinfo, f)
+
+
+if __name__ == "__main__":
+
+    # Read arguments
+    args = read_args()
+
+    # Create output directory
+    os.makedirs(args.output_dir, exist_ok=True)
+
+    # Get all video files
+    su.log.print_update("Processing files at " + args.input_dir)
+    video_files = glob(os.path.join(args.input_dir, f"*.{args.ext}"))
+    print(f"Found {len(video_files)} video files")
+
+    # Iterate
+    iterator = su.log.tqdm_iterator(video_files, desc="Sharding videos")
+
+    shard_size = args.shard_size
+    output_dir = args.output_dir
+    shard_id = 0
+    i = 0
+    for video_file in iterator:
+        if i % shard_size == 0:
+            if i > 0:
+                tar.close()
+            shard_path = os.path.join(output_dir, f"shard-{shard_id:04d}.tar")
+            print(f"Creating shard {shard_path}")
+            tar = tarfile.open(shard_path, 'w')
+            shard_id += 1
+        write_to_shard(tar, video_file, os.path.basename(video_file))
+        i += 1
+
+    if tar:
+        tar.close()

shared/utils/__init__.py

@@ -0,0 +1,16 @@
+import shared.utils.paths as paths
+import shared.utils.log as log
+import shared.utils.io as io
+# import shared.utils.audio as audio
+import shared.utils.image as image
+# import shared.utils.av as av
+import shared.utils.pandas_utils as pd_utils
+import shared.utils.visualize as visualize
+import shared.utils.metrics as metrics
+import shared.utils.misc as misc
+# import shared.utils.keypoint_matching as keypoint_matching
+import shared.utils.physics as physics
+import shared.utils.video as video
+import shared.utils.visual_prompts as visual_prompts
+import shared.utils.gif as gif
+

shared/utils/audio.py

@@ -0,0 +1,227 @@
+"""Audio utils"""
+import numpy as np
+import matplotlib.pyplot as plt
+
+
+def load_audio(audio_path: str, sr: int = None, max_duration: int = 10., start: int = 0, stop: int = None):
+    """Loads audio and pads/trims it to max_duration"""
+    import librosa
+    data, sr = librosa.load(audio_path, sr=sr)
+    
+    if stop is not None:
+        start = int(start * sr)
+        stop = int(stop * sr)
+        data = data[start:stop]
+    
+    # Convert to mono
+    if len(data.shape) > 1:
+        data = np.mean(data, axis=1)
+    
+    n_frames = int(max_duration * sr)
+    if len(data) > n_frames:
+        data = data[:n_frames]
+    elif len(data) < n_frames:
+        data = np.pad(data, (0, n_frames - len(data)), "constant")
+    return data, sr
+    
+
+# def compute_spectrogram(data: np.ndarray, sr: int):
+#     D = librosa.stft(data)  # STFT of y
+#     S_db = librosa.amplitude_to_db(np.abs(D), ref=np.max)
+#     return S_db
+
+
+def compute_spec_freq_mean(S_db: np.ndarray, eps=1e-5):
+    # Compute mean of spectrogram over frequency axis
+    S_db_normalized = (S_db - S_db.mean(axis=1)[:, None]) / (S_db.std(axis=1)[:, None] + eps)
+    S_db_over_time = S_db_normalized.sum(axis=0)
+    return S_db_over_time
+
+
+def process_audiofile(audio_path, functions=["load_audio", "compute_spectrogram", "compute_spec_freq_mean"]):
+    """Processes audio file with a list of functions"""
+    data, sr = load_audio(audio_path)
+    for function in functions:
+        if function == "load_audio":
+            pass
+        elif function == "compute_spectrogram":
+            data = compute_spectrogram(data, sr)
+        elif function == "compute_spec_freq_mean":
+            data = compute_spec_freq_mean(data)
+        else:
+            raise ValueError(f"Unknown function {function}")
+    return data
+
+
+
+"""PyDub's silence detection is based on the energy of the audio signal."""
+import numpy as np
+
+
+def sigmoid(x):
+    return 1 / (1 + np.exp(-x))
+
+
+class SilenceDetector:
+    
+
+    def __init__(self, silence_thresh=-36) -> None:
+        self.silence_thresh = silence_thresh
+    
+    def __call__(self, audio_path: str, start=None, end=None):
+        
+        import pydub
+        from pydub.utils import db_to_float
+        
+        try:
+            waveform = pydub.AudioSegment.from_file(audio_path)
+        except:
+            print("Error loading audio file: ", audio_path)
+            return 100.0
+        
+        start_ms = int(start * 1000) if start else 0
+        end_ms = int(end * 1000) if end else len(waveform)
+        waveform = waveform[start_ms:end_ms]
+        
+        # convert silence threshold to a float value (so we can compare it to rms)
+        silence_thresh = db_to_float(self.silence_thresh) * waveform.max_possible_amplitude
+        
+        if waveform.rms == 0:
+            return 100.0
+
+        silence_prob = sigmoid((silence_thresh - waveform.rms) / waveform.rms)
+
+        # return waveform.rms <= silence_thresh
+        return np.round(100 * silence_prob, 2)
+
+
+def frequency_bin_to_value(bin_index, sr, n_fft):
+    return int(bin_index * sr / n_fft)
+
+
+def time_bin_to_value(bin_index, hop_length, sr):
+    return (bin_index) * (hop_length / sr)
+
+
+def add_time_annotations(ax, nt_bins, hop_length, sr, skip=50):
+    # Show time (s) values on the x-axis
+    t_bins = np.arange(nt_bins)
+    t_vals = np.round(np.array([time_bin_to_value(tb, hop_length, sr) for tb in t_bins]), 1)
+    try:
+        ax.set_xticks(t_bins[::skip], t_vals[::skip])
+    except:
+        pass
+    ax.set_xlabel("Time (s)")
+
+
+def add_freq_annotations(ax, nf_bins, sr, n_fft, skip=50):
+    f_bins = np.arange(nf_bins)
+    f_vals = np.array([frequency_bin_to_value(fb, sr, n_fft) for fb in f_bins])
+    try:
+        ax.set_yticks(f_bins[::skip], f_vals[::skip])
+    except:
+        pass
+    # ax.set_yticks(f_bins[::skip])
+    # ax.set_yticklabels(f_vals[::skip])
+    ax.set_ylabel("Frequency (Hz)")
+
+
+def show_single_spectrogram(
+    spec,
+    sr,
+    n_fft,
+    hop_length,
+    ax=None,
+    fig=None,
+    figsize=(10, 2),
+    cmap="viridis",
+    colorbar=True,
+    show=True,
+    format='%+2.0f dB',
+    xlabel='Time (s)',
+    ylabel="Frequency (Hz)",
+    title=None,
+    show_dom_freq=False,
+):
+
+    if ax is None:
+        fig, ax = plt.subplots(1, 1, figsize=figsize)
+    axim = ax.imshow(spec, origin="lower", cmap=cmap)
+
+    # Show frequency (Hz) values on y-axis
+    nf_bins, nt_bins = spec.shape
+
+    if "frequency" in ylabel.lower():
+        # Add frequency annotation
+        add_freq_annotations(ax, nf_bins, sr, n_fft)
+
+    # Add time annotation
+    add_time_annotations(ax, nt_bins, hop_length, sr)
+
+    ax.set_title(title)
+    ax.set_xlabel(xlabel)
+    ax.set_ylabel(ylabel)
+
+    if colorbar:
+        fig.colorbar(axim, ax=ax, orientation='vertical', fraction=0.01, format=format)
+
+    if show_dom_freq:
+        fmax = spec.argmax(axis=0)
+        ax.scatter(np.arange(spec.shape[1]), fmax, color="white", s=0.2)
+
+    if show:
+        plt.show()
+
+
+def compute_spectrogram(y, n_fft, hop_length, margin, n_mels=None):
+    import librosa
+
+    # STFT
+    D = librosa.stft(y, n_fft=n_fft, hop_length=hop_length)
+
+    # Run HPSS
+    S, _ = librosa.decompose.hpss(D, margin=margin)
+
+    # DB
+    S = librosa.amplitude_to_db(np.abs(S), ref=np.max)
+
+    if n_mels is not None:
+        S = librosa.feature.melspectrogram(S=S, n_mels=n_mels)
+
+    return S
+
+
+def show_spectrogram(S, sr, n_fft=512, hop_length=256, figsize=(10, 3), n_mels=None, ax=None, show=True):
+    import librosa
+    if ax is None:
+        fig, ax = plt.subplots(1, 1, figsize=figsize)
+    y_axis = "mel" if n_mels is not None else "linear"
+    librosa.display.specshow(
+        S,
+        sr=sr,
+        hop_length=hop_length,
+        n_fft=n_fft,
+        y_axis=y_axis,
+        x_axis='time',
+        ax=ax,
+    )
+    ax.set_title("LogSpectrogram" if n_mels is None else "LogMelSpectrogram")
+    if show:
+        plt.show()
+
+
+def show_frame_and_spectrogram(frame, S, sr, figsize=(12, 4), show=True, axes=None, **spec_args):
+    if axes is None:
+        fig, axes = plt.subplots(1, 2, figsize=figsize, gridspec_kw={"width_ratios": [0.2, 0.8]})
+    ax = axes[0]
+    ax.imshow(frame)
+    ax.set_xticks([])
+    ax.set_yticks([])
+
+    ax = axes[1]
+    show_spectrogram(S=S, sr=sr, ax=ax, show=False, **spec_args)
+
+    plt.tight_layout()
+
+    if show:
+        plt.show()

shared/utils/av.py

@@ -0,0 +1,93 @@
+"""Audio-visual helper functions."""
+import cv2
+import numpy as np
+
+
+def save_video_with_audio(video, audio, output_path):
+    """
+    Saves a video file with audio.
+    
+    Args:
+        video (np.ndarray): Video frames.
+        audio (np.ndarray): Audio samples.
+        output_path (str): Output path.
+    """
+
+    # check the correct shape and format for audio
+    assert isinstance(audio, np.ndarray)
+    assert len(audio.shape) == 2
+    assert audio.shape[1] in [1, 2]
+
+    # create video writer
+    video_writer = cv2.VideoWriter(output_path, cv2.VideoWriter_fourcc(*'mp4v'), 30, (video.shape[2], video.shape[1]))
+    # write the image frames to the video
+    for frame in video:
+        video_writer.write(frame)
+    # add the audio data to the video
+    video_writer.write(audio)
+    # release the VideoWriter object
+    video_writer.release()
+
+
+def save_video_from_image_sequence_and_audio(sequence, audio, save_path, video_fps=15, audio_fps=22100):
+    import torch
+    from moviepy.editor import VideoClip, AudioClip, ImageSequenceClip
+    from moviepy.audio.AudioClip import AudioArrayClip
+    
+    assert isinstance(sequence, list) and isinstance(audio, (np.ndarray, torch.Tensor))
+    assert len(audio.shape) == 2 and audio.shape[1] in [1, 2]
+    
+    video_duration = len(sequence) / video_fps
+    audio_duration = len(audio) / audio_fps
+    # # print(f"Video duration: {video_duration:.2f}s, audio duration: {audio_duration:.2f}s")
+    # assert video_duration == audio_duration, \
+    #     f"Video duration ({video_duration}) and audio duration ({audio_duration}) do not match."
+
+    video_clip = ImageSequenceClip(sequence, fps=video_fps)
+    audio_clip = AudioArrayClip(audio, fps=audio_fps)
+    video_clip = video_clip.set_audio(audio_clip)
+    # video_clip.write_videofile(save_path, verbose=True, logger=None, fps=video_fps, audio_fps=audio_fps)
+    video_clip.write_videofile(save_path, verbose=False, logger=None)
+
+
+import cv2, os
+import argparse
+import numpy as np
+from glob import glob
+import librosa
+import subprocess
+
+
+def generate_video(args):
+
+	frames = glob('{}/*.png'.format(args.input_dir))
+	print("Total frames = ", len(frames))
+
+	frames.sort(key = lambda x: int(x.split("/")[-1].split(".")[0]))
+
+	img = cv2.imread(frames[0])
+	print(img.shape)
+	fname = 'inference.avi'
+	video = cv2.VideoWriter(
+        fname, cv2.VideoWriter_fourcc(*'DIVX'), args.fps, (img.shape[1], img.shape[0]),
+    )
+ 
+	for i in range(len(frames)):
+		img = cv2.imread(frames[i])
+		video.write(img)
+	
+	video.release()
+
+	output_file_name = args.output_video
+
+	no_sound_video = output_file_name + '_nosound.mp4'
+	subprocess.call('ffmpeg -hide_banner -loglevel panic -i %s -c copy -an -strict -2 %s' % (fname, no_sound_video), shell=True)
+
+	if args.audio_file is not None:
+		video_output = output_file_name + '.mp4'
+		subprocess.call('ffmpeg -hide_banner -loglevel panic -y -i %s -i %s -strict -2 -q:v 1 %s' % 
+						(args.audio_file, no_sound_video, video_output), shell=True)
+
+		os.remove(no_sound_video)
+	
+	os.remove(fname)

shared/utils/classification.py

@@ -0,0 +1,47 @@
+"""Helper functions for classification tasks."""
+import matplotlib.pyplot as plt
+import pandas as pd
+import numpy as np
+
+
+def plot_metric_curve(
+        xvalues, yvalues, thresholds, title=None,
+        figsize=(8, 7), show_thresholds=True, show_legend=True,
+        ylabel='X', xlabel='Y', ax=None, text_delta=0.01,
+        label="Metric Curve", color="royalblue", show=False,
+        fill=None,
+    ):
+    """Plot a metric curve, e.g., PR curve or ROC curve."""
+
+    if ax is None:
+        fig, ax = plt.subplots(1, 1, figsize=figsize)
+
+    ax.grid(alpha=0.3)
+    ax.set_title(title)
+    ax.set_ylabel(ylabel)
+    ax.set_xlabel(xlabel)
+    
+    ax.plot(xvalues, yvalues, marker='o', label=label, color=color)
+    ax.set_xlim(-0.08, 1.08)
+    ax.set_ylim(-0.08, 1.08)
+    
+    if fill is not None:
+        yticks = ax.get_yticks()
+        ax.fill_between(xvalues, yvalues, "", alpha=0.08, color=color)
+        # Add `fill` inside the curve
+        # Find a single (x, y) s.t. it is inside the curve
+        ax.text(0.4, 0.5, fill, color=color)
+        ax.set_yticks(yticks)
+        ax.set_yticklabels([f"{y:.1f}" for y in yticks])
+        ax.set_ylim(-0.08, 1.08)
+    
+    # Show thresholds
+    if show_thresholds:
+        for x, y, t in zip(xvalues, yvalues, thresholds):
+            ax.text(x + text_delta, y + text_delta, np.round(t, 2), color=color, alpha=0.5)
+    
+    if show_legend:
+        ax.legend()
+    
+    if show:
+        plt.show()

shared/utils/epic.py

@@ -0,0 +1,15 @@
+"""Utils specific for EPIC data."""
+import datetime
+
+
+def timestamp_to_seconds(timestamp: str):
+    # Parse the timestamp string into a datetime object
+    time_obj = datetime.datetime.strptime(timestamp, '%H:%M:%S.%f')
+
+    # Calculate the total number of seconds using the timedelta object
+    total_seconds = time_obj.time().second \
+        + time_obj.time().minute * 60 \
+        + time_obj.time().hour * 3600 \
+        + time_obj.time().microsecond / 1000000
+    
+    return total_seconds

shared/utils/gif.py

@@ -0,0 +1,609 @@
+import subprocess
+import os
+from pathlib import Path
+
+def create_side_by_side_gif(video_paths, output_gif, gap_width=20, fps=10, scale_height=240, gap_color="white", verbose=False):
+    """
+    Create a single GIF with multiple videos placed side by side.
+    
+    Args:
+        video_paths (list): List of paths to input MP4 files
+        output_gif (str): Path for output GIF file
+        gap_width (int): Width of gap between videos in pixels
+        fps (int): Frame rate for output GIF
+        scale_height (int): Height to scale all videos to (maintains aspect ratio)
+        gap_color (str): Color for gaps between videos (e.g., "white", "black", "red", "#FF0000")
+        verbose (bool): Whether to print FFmpeg commands and processing messages
+    """
+    
+    if not video_paths:
+        raise ValueError("No video paths provided")
+    
+    # Verify all input files exist
+    for path in video_paths:
+        if not os.path.exists(path):
+            raise FileNotFoundError(f"Video file not found: {path}")
+    
+    # Create filter complex string for FFmpeg
+    num_videos = len(video_paths)
+    
+    # Input mapping and scaling
+    filter_parts = []
+    scaled_inputs = []
+    
+    for i, _ in enumerate(video_paths):
+        # Scale each video to same height while maintaining aspect ratio
+        filter_parts.append(f"[{i}:v]scale=-1:{scale_height}[v{i}]")
+        scaled_inputs.append(f"[v{i}]")
+    
+    # Create horizontal stack with gaps
+    if num_videos == 1:
+        hstack_filter = f"{scaled_inputs[0]}copy[stacked]"
+    else:
+        # Create colored gap between videos
+        gap_filters = []
+        for i in range(num_videos - 1):
+            gap_filters.append(f"color={gap_color}:{gap_width}x{scale_height}:d=1[gap{i}]")
+        
+        if gap_filters:
+            filter_parts.extend(gap_filters)
+        
+        # Build hstack input list with gaps
+        hstack_inputs = []
+        for i in range(num_videos):
+            hstack_inputs.append(scaled_inputs[i])
+            if i < num_videos - 1:  # Add gap after each video except the last
+                hstack_inputs.append(f"[gap{i}]")
+        
+        hstack_filter = f"{''.join(hstack_inputs)}hstack=inputs={len(hstack_inputs)}[stacked]"
+    
+    filter_parts.append(hstack_filter)
+    
+    # Complete the filter complex for stacked video
+    stacked_filter = ";".join(filter_parts)
+    
+    # Build FFmpeg command with two-pass palette approach
+    cmd = ["ffmpeg", "-y"]  # -y to overwrite output file
+    
+    # Add input files
+    for video_path in video_paths:
+        cmd.extend(["-i", video_path])
+    
+    # Add filter complex and output options
+    cmd.extend([
+        "-filter_complex", f"{stacked_filter};[stacked]split[s0][s1];[s0]palettegen[p];[s1][p]paletteuse=dither=bayer:bayer_scale=3",
+        "-r", str(fps),  # Set frame rate
+        "-loop", "0",    # Infinite loop
+        output_gif
+    ])
+    
+    if verbose:
+        print("Running FFmpeg command:")
+        print(" ".join(cmd))
+        print("\nProcessing...")
+    
+    try:
+        result = subprocess.run(cmd, capture_output=True, text=True, check=True)
+        if verbose:
+            print(f"✓ Successfully created GIF: {output_gif}")
+        return True
+    except subprocess.CalledProcessError as e:
+        if verbose:
+            print(f"✗ FFmpeg error: {e.stderr}")
+        return False
+    except FileNotFoundError:
+        if verbose:
+            print("✗ FFmpeg not found. Please install FFmpeg first.")
+            print("  - Windows: Download from https://ffmpeg.org/download.html")
+            print("  - macOS: brew install ffmpeg")
+            print("  - Linux: sudo apt install ffmpeg (Ubuntu/Debian)")
+        return False
+
+def create_top_to_bottom_gif(video_paths, output_gif, gap_height=20, fps=10, scale_width=320, gap_color="white", verbose=False):
+    """
+    Create a single GIF with multiple videos stacked vertically (top to bottom).
+    
+    Args:
+        video_paths (list): List of paths to input MP4 files
+        output_gif (str): Path for output GIF file
+        gap_height (int): Height of gap between videos in pixels
+        fps (int): Frame rate for output GIF
+        scale_width (int): Width to scale all videos to (maintains aspect ratio)
+        gap_color (str): Color for gaps between videos (e.g., "white", "black", "red", "#FF0000")
+        verbose (bool): Whether to print FFmpeg commands and processing messages
+    """
+    
+    if not video_paths:
+        raise ValueError("No video paths provided")
+    
+    # Verify all input files exist
+    for path in video_paths:
+        if not os.path.exists(path):
+            raise FileNotFoundError(f"Video file not found: {path}")
+    
+    # Create filter complex string for FFmpeg
+    num_videos = len(video_paths)
+    
+    # Input mapping and scaling
+    filter_parts = []
+    scaled_inputs = []
+    
+    for i, _ in enumerate(video_paths):
+        # Scale each video to same width while maintaining aspect ratio
+        filter_parts.append(f"[{i}:v]scale={scale_width}:-1[v{i}]")
+        scaled_inputs.append(f"[v{i}]")
+    
+    # Create vertical stack with gaps
+    if num_videos == 1:
+        vstack_filter = f"{scaled_inputs[0]}copy[stacked]"
+    else:
+        # Create colored gap between videos
+        gap_filters = []
+        for i in range(num_videos - 1):
+            gap_filters.append(f"color={gap_color}:{scale_width}x{gap_height}:d=1[gap{i}]")
+        
+        if gap_filters:
+            filter_parts.extend(gap_filters)
+        
+        # Build vstack input list with gaps
+        vstack_inputs = []
+        for i in range(num_videos):
+            vstack_inputs.append(scaled_inputs[i])
+            if i < num_videos - 1:  # Add gap after each video except the last
+                vstack_inputs.append(f"[gap{i}]")
+        
+        vstack_filter = f"{''.join(vstack_inputs)}vstack=inputs={len(vstack_inputs)}[stacked]"
+    
+    filter_parts.append(vstack_filter)
+    
+    # Complete the filter complex for stacked video
+    stacked_filter = ";".join(filter_parts)
+    
+    # Build FFmpeg command with two-pass palette approach
+    cmd = ["ffmpeg", "-y"]  # -y to overwrite output file
+    
+    # Add input files
+    for video_path in video_paths:
+        cmd.extend(["-i", video_path])
+    
+    # Add filter complex and output options
+    cmd.extend([
+        "-filter_complex", f"{stacked_filter};[stacked]split[s0][s1];[s0]palettegen[p];[s1][p]paletteuse=dither=bayer:bayer_scale=3",
+        "-r", str(fps),  # Set frame rate
+        "-loop", "0",    # Infinite loop
+        output_gif
+    ])
+    
+    if verbose:
+        print("Running FFmpeg command:")
+        print(" ".join(cmd))
+        print("\nProcessing...")
+    
+    try:
+        result = subprocess.run(cmd, capture_output=True, text=True, check=True)
+        if verbose:
+            print(f"✓ Successfully created GIF: {output_gif}")
+        return True
+    except subprocess.CalledProcessError as e:
+        if verbose:
+            print(f"✗ FFmpeg error: {e.stderr}")
+        return False
+    except FileNotFoundError:
+        if verbose:
+            print("✗ FFmpeg not found. Please install FFmpeg first.")
+            print("  - Windows: Download from https://ffmpeg.org/download.html")
+            print("  - macOS: brew install ffmpeg")
+            print("  - Linux: sudo apt install ffmpeg (Ubuntu/Debian)")
+        return False
+
+def reverse_video(input_video_path, output_filename=None, verbose=False):
+    """
+    Reverse a video file and save it to /tmp directory.
+    
+    Args:
+        input_video_path (str): Path to input MP4 file
+        output_filename (str, optional): Name for output file. If None, generates from input filename
+        verbose (bool): Whether to print FFmpeg commands and processing messages
+        
+    Returns:
+        str: Path to the reversed video file in /tmp, or None if failed
+    """
+    
+    if not os.path.exists(input_video_path):
+        raise FileNotFoundError(f"Input video file not found: {input_video_path}")
+    
+    # Generate output filename if not provided
+    if output_filename is None:
+        input_name = Path(input_video_path).stem
+        output_filename = f"{input_name}_reversed.mp4"
+    
+    # Ensure output filename has .mp4 extension
+    if not output_filename.endswith('.mp4'):
+        output_filename += '.mp4'
+    
+    # Create output path in /tmp
+    output_path = os.path.join('/tmp', output_filename)
+    
+    # Build FFmpeg command to reverse video
+    cmd = [
+        "ffmpeg", "-y",  # -y to overwrite output file
+        "-i", input_video_path,
+        "-vf", "reverse",  # Video filter to reverse frames
+        "-af", "areverse",  # Audio filter to reverse audio
+        "-c:v", "libx264",  # Video codec
+        "-c:a", "aac",      # Audio codec
+        output_path
+    ]
+    
+    if verbose:
+        print("Running FFmpeg command to reverse video:")
+        print(" ".join(cmd))
+        print("\nProcessing...")
+    
+    try:
+        result = subprocess.run(cmd, capture_output=True, text=True, check=True)
+        if verbose:
+            print(f"✓ Successfully created reversed video: {output_path}")
+        return output_path
+    except subprocess.CalledProcessError as e:
+        if verbose:
+            print(f"✗ FFmpeg error: {e.stderr}")
+        return None
+    except FileNotFoundError:
+        if verbose:
+            print("✗ FFmpeg not found. Please install FFmpeg first.")
+            print("  - Windows: Download from https://ffmpeg.org/download.html")
+            print("  - macOS: brew install ffmpeg")
+            print("  - Linux: sudo apt install ffmpeg (Ubuntu/Debian)")
+        return None
+
+def add_text_overlay(input_video_path, text, output_filename=None, font_size=12, font_color="white", 
+                    background_color="black", position="top", margin=10, duration=None, verbose=False):
+    """
+    Add a text overlay to a video with a background title bar.
+    
+    Args:
+        input_video_path (str): Path to input MP4 file
+        text (str): Text to display
+        output_filename (str, optional): Name for output file. If None, generates from input filename
+        font_size (int): Font size for the text (default: 24)
+        font_color (str): Color of the text (default: "white")
+        background_color (str): Color of the background bar (default: "black")
+        position (str): Position of text bar - "top", "bottom", "center" (default: "top")
+        margin (int): Margin from edge in pixels (default: 10)
+        duration (float, optional): Duration to show text in seconds. If None, shows for entire video
+        verbose (bool): Whether to print FFmpeg commands and processing messages
+        
+    Returns:
+        str: Path to the video with text overlay in /tmp, or None if failed
+    """
+    
+    if not os.path.exists(input_video_path):
+        raise FileNotFoundError(f"Input video file not found: {input_video_path}")
+    
+    # Generate output filename if not provided
+    if output_filename is None:
+        input_name = Path(input_video_path).stem
+        output_filename = f"{input_name}_with_text.mp4"
+    
+    # Ensure output filename has .mp4 extension
+    if not output_filename.endswith('.mp4'):
+        output_filename += '.mp4'
+    
+    # Create output path in /tmp
+    output_path = os.path.join('/tmp', output_filename)
+    
+    # Determine text position based on position parameter
+    if position == "top":
+        text_position = f"x={margin}:y={margin+5}"  # Add small offset for background box
+    elif position == "bottom":
+        text_position = f"x={margin}:y=h-th-{margin+5}"  # Add small offset for background box
+    elif position == "center":
+        text_position = f"x={margin}:y=(h-th)/2"
+    else:
+        text_position = f"x={margin}:y={margin+5}"  # Default to top with offset
+    
+    # Build the drawtext filter
+    drawtext_filter = f"drawtext=text='{text}':fontsize={font_size}:fontcolor={font_color}:{text_position}"
+    
+    # Add background box if needed
+    if background_color != "transparent":
+        # Create a semi-transparent background box with estimated height based on font size
+        # Estimate text height as approximately 1.2 * font_size
+        estimated_text_height = int(font_size * 1.2)
+        box_height = estimated_text_height + 10  # Add padding
+        
+        # Position box based on text position
+        if position == "top":
+            box_y = margin
+        elif position == "bottom":
+            box_y = f"h-{box_height}-{margin}"
+        elif position == "center":
+            box_y = f"(h-{box_height})/2"
+        else:
+            box_y = margin
+            
+        box_filter = f"drawbox=x={margin-5}:y={box_y}:w=iw-{2*(margin-5)}:h={box_height}:color={background_color}@0.7:t=fill"
+        drawtext_filter = f"{box_filter},{drawtext_filter}"
+    
+    # Add duration constraint if specified
+    if duration is not None:
+        drawtext_filter += f":enable='between(t,0,{duration})'"
+    
+    # Build FFmpeg command
+    cmd = [
+        "ffmpeg", "-y",  # -y to overwrite output file
+        "-i", input_video_path,
+        "-vf", drawtext_filter,
+        "-c:v", "libx264",  # Video codec
+        "-c:a", "copy",     # Copy audio without re-encoding
+        output_path
+    ]
+    
+    if verbose:
+        print("Running FFmpeg command to add text overlay:")
+        print(" ".join(cmd))
+        print("\nProcessing...")
+    
+    try:
+        result = subprocess.run(cmd, capture_output=True, text=True, check=True)
+        if verbose:
+            print(f"✓ Successfully created video with text overlay: {output_path}")
+        return output_path
+    except subprocess.CalledProcessError as e:
+        if verbose:
+            print(f"✗ FFmpeg error: {e.stderr}")
+        return None
+    except FileNotFoundError:
+        if verbose:
+            print("✗ FFmpeg not found. Please install FFmpeg first.")
+            print("  - Windows: Download from https://ffmpeg.org/download.html")
+            print("  - macOS: brew install ffmpeg")
+            print("  - Linux: sudo apt install ffmpeg (Ubuntu/Debian)")
+        return None
+
+def add_text_strip(input_video_path, text, output_filename=None, font_size=16, font_color="white", 
+                  background_color="black", position="top", text_padding=20, max_width_ratio=0.9, verbose=False):
+    """
+    Add a text strip/bar to a video (increases video height) rather than overlaying text.
+    
+    Args:
+        input_video_path (str): Path to input MP4 file
+        text (str): Text to display in the strip
+        output_filename (str, optional): Name for output file. If None, generates from input filename
+        font_size (int): Font size for the text (default: 16)
+        font_color (str): Color of the text (default: "white")
+        background_color (str): Color of the background strip (default: "black")
+        position (str): Position of text strip - "top" or "bottom" (default: "top")
+        text_padding (int): Padding around text in pixels (default: 20)
+        max_width_ratio (float): Maximum width of text as ratio of video width (default: 0.9)
+        verbose (bool): Whether to print FFmpeg commands and processing messages
+        
+    Returns:
+        str: Path to the video with text strip in /tmp, or None if failed
+    """
+    
+    if not os.path.exists(input_video_path):
+        raise FileNotFoundError(f"Input video file not found: {input_video_path}")
+    
+    # Generate output filename if not provided
+    if output_filename is None:
+        input_name = Path(input_video_path).stem
+        output_filename = f"{input_name}_with_strip.mp4"
+    
+    # Ensure output filename has .mp4 extension
+    if not output_filename.endswith('.mp4'):
+        output_filename += '.mp4'
+    
+    # Create output path in /tmp
+    output_path = os.path.join('/tmp', output_filename)
+    
+    # Calculate text strip height based on font size, padding, and estimated line count
+    # Estimate characters per line based on font size (roughly 2 characters per font size pixel)
+    estimated_chars_per_line = int(font_size * 2)
+    text_lines = text.split('\n') if '\n' in text else [text]
+    
+    # If text is too long, wrap it
+    wrapped_lines = []
+    for line in text_lines:
+        if len(line) <= estimated_chars_per_line:
+            wrapped_lines.append(line)
+        else:
+            # Simple word wrapping
+            words = line.split(' ')
+            current_line = ""
+            for word in words:
+                if len(current_line + " " + word) <= estimated_chars_per_line:
+                    current_line += (" " + word) if current_line else word
+                else:
+                    if current_line:
+                        wrapped_lines.append(current_line)
+                    current_line = word
+            if current_line:
+                wrapped_lines.append(current_line)
+    
+    # Calculate strip height based on number of lines
+    line_height = font_size + 5  # Add some line spacing
+    strip_height = (len(wrapped_lines) * line_height) + (2 * text_padding)
+    
+    # Create text strip using a different approach - pad the video and add text
+    # This will add padding above the video and put text in that padded area
+    if position == "top":
+        # Add padding to top of video and put text in the padded area
+        text_strip_filter = f"[0:v]pad=iw:ih+{strip_height}:0:{strip_height}:{background_color}[padded];[padded]drawtext=text='{chr(10).join(wrapped_lines)}':fontsize={font_size}:fontcolor={font_color}:x=(w-tw)/2:y={text_padding}:line_spacing={line_height}[stacked]"
+    else:  # bottom
+        # Add padding to bottom of video and put text in the padded area
+        text_strip_filter = f"[0:v]pad=iw:ih+{strip_height}:0:0:{background_color}[padded];[padded]drawtext=text='{chr(10).join(wrapped_lines)}':fontsize={font_size}:fontcolor={font_color}:x=(w-tw)/2:y=h-th-{text_padding}:line_spacing={line_height}[stacked]"
+    
+    # Build FFmpeg command
+    cmd = [
+        "ffmpeg", "-y",  # -y to overwrite output file
+        "-i", input_video_path,
+        "-filter_complex", text_strip_filter,
+        "-map", "[stacked]",  # Map the processed video
+        "-map", "0:a",        # Map the original audio
+        "-c:v", "libx264",    # Video codec
+        "-c:a", "copy",       # Copy audio without re-encoding
+        output_path
+    ]
+    
+    if verbose:
+        print("Running FFmpeg command to add text strip:")
+        print(" ".join(cmd))
+        print("\nProcessing...")
+    
+    try:
+        result = subprocess.run(cmd, capture_output=True, text=True, check=True)
+        if verbose:
+            print(f"✓ Successfully created video with text strip: {output_path}")
+        return output_path
+    except subprocess.CalledProcessError as e:
+        if verbose:
+            print(f"✗ FFmpeg error: {e.stderr}")
+        return None
+    except FileNotFoundError:
+        if verbose:
+            print("✗ FFmpeg not found. Please install FFmpeg first.")
+            print("  - Windows: Download from https://ffmpeg.org/download.html")
+            print("  - macOS: brew install ffmpeg")
+            print("  - Linux: sudo apt install ffmpeg (Ubuntu/Debian)")
+        return None
+
+def get_video_info(video_path):
+    """Get basic info about a video file."""
+    cmd = [
+        "ffprobe", "-v", "quiet", "-print_format", "json", 
+        "-show_format", "-show_streams", video_path
+    ]
+    
+    try:
+        result = subprocess.run(cmd, capture_output=True, text=True, check=True)
+        import json
+        data = json.loads(result.stdout)
+        
+        # Find video stream
+        for stream in data['streams']:
+            if stream['codec_type'] == 'video':
+                return {
+                    'width': stream['width'],
+                    'height': stream['height'],
+                    'duration': float(stream.get('duration', 0)),
+                    'fps': eval(stream.get('r_frame_rate', '0/1'))
+                }
+    except:
+        pass
+    return None
+
+
+
+
+
+
+# Example usage
+if __name__ == "__main__":
+    # Example video paths - replace with your actual video files
+    video_files = [
+        "examples/folding_paper.mp4",
+        "examples/S008C002P032R002A051.mp4", 
+    ]
+    
+    output_file = "combined_videos.gif"
+    
+    # Check if example files exist
+    existing_files = [f for f in video_files if os.path.exists(f)]
+    
+    if existing_files:
+        print(f"Found {len(existing_files)} video files:")
+        for video in existing_files:
+            info = get_video_info(video)
+            if info:
+                print(f"  {video}: {info['width']}x{info['height']}, {info['duration']:.1f}s")
+            else:
+                print(f"  {video}: (info unavailable)")
+        
+        # Create the horizontal GIF
+        success = create_side_by_side_gif(
+            video_paths=existing_files,
+            output_gif=output_file,
+            gap_width=30,      # 30px gap between videos
+            fps=12,            # 12 frames per second
+            scale_height=300,  # Scale all videos to 300px height
+            gap_color="white"  # White gap between videos
+        )
+        
+        # Also create a vertical GIF
+        vertical_output_file = "combined_videos_vertical.gif"
+        success_vertical = create_top_to_bottom_gif(
+            video_paths=existing_files,
+            output_gif=vertical_output_file,
+            gap_height=20,     # 20px gap between videos
+            fps=12,            # 12 frames per second
+            scale_width=320,   # Scale all videos to 320px width
+            gap_color="white"  # White gap between videos
+        )
+        
+        if success:
+            file_size = os.path.getsize(output_file) / (1024 * 1024)  # MB
+            print(f"\nHorizontal GIF size: {file_size:.1f} MB")
+        
+        if success_vertical:
+            vertical_file_size = os.path.getsize(vertical_output_file) / (1024 * 1024)  # MB
+            print(f"Vertical GIF size: {vertical_file_size:.1f} MB")
+        
+        # Example of reversing a video
+        if existing_files:
+            print(f"\nReversing first video: {existing_files[0]}")
+            reversed_path = reverse_video(existing_files[0])
+            if reversed_path:
+                print(f"Reversed video saved to: {reversed_path}")
+        
+        # Example of adding text overlay
+        if existing_files:
+            print(f"\nAdding text overlay to first video: {existing_files[0]}")
+            text_video_path = add_text_overlay(
+                input_video_path=existing_files[0],
+                text="Sample Title Text",
+                font_size=30,
+                font_color="white",
+                background_color="black",
+                position="top",
+                margin=15
+            )
+            if text_video_path:
+                print(f"Video with text overlay saved to: {text_video_path}")
+        
+        # Example of adding text strip
+        if existing_files:
+            print(f"\nAdding text strip to first video: {existing_files[0]}")
+            strip_video_path = add_text_strip(
+                input_video_path=existing_files[0],
+                text="Video Title Strip",
+                font_size=16,
+                font_color="white",
+                background_color="darkblue",
+                position="top",
+                text_padding=15
+            )
+            if strip_video_path:
+                print(f"Video with text strip saved to: {strip_video_path}")
+    else:
+        print("No video files found. Please update the video_files list with your actual MP4 file paths.")
+        print("\nExample usage:")
+        print("video_files = [")
+        print('    "/path/to/your/video1.mp4",')
+        print('    "/path/to/your/video2.mp4",')
+        print('    "/path/to/your/video3.mp4"')
+        print("]")
+        print("\n# Create horizontal GIF")
+        print("create_side_by_side_gif(video_files, 'horizontal.gif')")
+        print("\n# Create vertical GIF")
+        print("create_top_to_bottom_gif(video_files, 'vertical.gif')")
+        print("\n# Reverse a video")
+        print("reversed_path = reverse_video('/path/to/your/video1.mp4')")
+        print("print(f'Reversed video: {reversed_path}')")
+        print("\n# Add text overlay to video")
+        print("text_video = add_text_overlay('/path/to/your/video1.mp4', 'My Title', font_size=30)")
+        print("print(f'Video with text: {text_video}')")
+        print("\n# Add text strip to video (increases video height)")
+        print("strip_video = add_text_strip('/path/to/your/video1.mp4', 'Title Strip', font_size=16)")
+        print("print(f'Video with strip: {strip_video}')")

shared/utils/hardware.py

@@ -0,0 +1,68 @@
+import subprocess
+import platform
+
+
+def get_cpu_info_linux():
+    try:
+        result = subprocess.run(['lscpu'], capture_output=True, text=True, check=True)
+        print("---- CPU Info (Linux) ----")
+        print(result.stdout)
+    except FileNotFoundError:
+        print("lscpu command not found.")
+    except subprocess.CalledProcessError as e:
+        print(f"Error running lscpu: {e}")
+
+
+def get_gpu_info_nvidia():
+    try:
+        result = subprocess.run(['nvidia-smi'], capture_output=True, text=True, check=True)
+        print("---- GPU Info (NVIDIA) ----")
+        print(result.stdout)
+    except FileNotFoundError:
+        print("nvidia-smi command not found. NVIDIA drivers might not be installed or not in PATH.")
+    except subprocess.CalledProcessError as e:
+        print(f"Error running nvidia-smi: {e}")
+
+
+if __name__ == "__main__":
+    os_type = platform.system()
+    print(f"Operating System: {os_type}")
+
+    if os_type == "Linux":
+        get_cpu_info_linux()
+        get_gpu_info_nvidia() # Also try rocm-smi if you have AMD
+    elif os_type == "Darwin": # macOS
+        print("---- CPU Info (macOS) ----")
+        subprocess.run(['sysctl', '-n', 'machdep.cpu.brand_string'])
+        subprocess.run(['sysctl', '-n', 'hw.ncpu'])
+        # For GPU, check System Information manually or use more specific tools if available
+    elif os_type == "Windows":
+        print("---- CPU Info (Windows) ----")
+        subprocess.run(['wmic', 'cpu', 'get', 'Name,NumberOfCores,NumberOfLogicalProcessors'], shell=True)
+        print("---- GPU Info (Windows - NVIDIA Example) ----")
+        try:
+            subprocess.run(['nvidia-smi'], shell=True) # May need to ensure nvidia-smi is in PATH
+        except FileNotFoundError:
+            print("nvidia-smi not found. For GPU info, check Task Manager or DxDiag.")
+    else:
+        print(f"Unsupported OS for this script: {os_type}")
+
+    # For PyTorch to check CUDA availability and GPU details:
+    try:
+        import torch
+        if torch.cuda.is_available():
+            print("\n---- PyTorch CUDA Info ----")
+            print(f"CUDA Available: {torch.cuda.is_available()}")
+            print(f"CUDA Version (PyTorch compiled with): {torch.version.cuda}")
+            print(f"Number of GPUs: {torch.cuda.device_count()}")
+            for i in range(torch.cuda.device_count()):
+                print(f"  GPU {i}: {torch.cuda.get_device_name(i)}")
+                print(f"    Memory Allocated: {torch.cuda.memory_allocated(i)/1024**2:.2f} MB")
+                print(f"    Memory Cached:    {torch.cuda.memory_reserved(i)/1024**2:.2f} MB") # formerly memory_cached
+                props = torch.cuda.get_device_properties(i)
+                print(f"    Total Memory:     {props.total_memory/1024**2:.2f} MB")
+                print(f"    Compute Capability: {props.major}.{props.minor}")
+        else:
+            print("\nPyTorch: CUDA is not available.")
+    except ImportError:
+        print("\nPyTorch is not installed.")

shared/utils/image.py

@@ -0,0 +1,81 @@
+"""Image operations."""
+from copy import deepcopy
+from PIL import Image
+import numpy as np
+
+
+def center_crop(im: Image):
+    width, height = im.size
+    new_width = width if width < height else height
+    new_height = height if height < width else width 
+
+    left = (width - new_width)/2
+    top = (height - new_height)/2
+    right = (width + new_width)/2
+    bottom = (height + new_height)/2
+
+    # Crop the center of the image
+    im = im.crop((left, top, right, bottom))
+    
+    return im
+
+
+def pad_to_square(im: Image, color=(0, 0, 0)):
+    im = deepcopy(im)
+    width, height = im.size
+
+    vert_pad = (max(width, height) - height) // 2
+    hor_pad = (max(width, height) - width) // 2
+    
+    if len(im.mode) == 3:
+        color = (0, 0, 0)
+    elif len(im.mode) == 1:
+        color = 0
+    else:
+        raise ValueError(f"Image mode not supported. Image has {im.mode} channels.")
+    
+    return add_margin(im, vert_pad, hor_pad, vert_pad, hor_pad, color=color)
+
+
+def add_margin(pil_img, top, right, bottom, left, color=(0, 0, 0)):
+    """Ref: https://note.nkmk.me/en/python-pillow-add-margin-expand-canvas/"""
+    width, height = pil_img.size
+    new_width = width + right + left
+    new_height = height + top + bottom
+    result = Image.new(pil_img.mode, (new_width, new_height), color)
+    result.paste(pil_img, (left, top))
+    return result
+
+
+def resize_image(image, new_height, new_width):
+    # Convert the numpy array image to PIL Image
+    pil_image = Image.fromarray(image)
+
+    # Resize the PIL Image
+    resized_image = pil_image.resize((new_width, new_height))
+
+    # Convert the resized PIL Image back to numpy array
+    resized_image_np = np.array(resized_image)
+
+    return resized_image_np
+
+
+def pad_to_width(pil_image, new_width, color=(0, 0, 0)):
+    """Pad the image to the specified width."""
+    # Convert the numpy array image to PIL Image
+    # pil_image = Image.fromarray(image)
+
+    # Get the current width and height of the image
+    width, height = pil_image.size
+    assert new_width > width, f"New width {new_width} is less than the current width {width}."
+
+    # Calculate the padding required
+    hor_pad = new_width - width
+
+    # Add padding to the image
+    padded_image = add_margin(pil_image, 0, hor_pad, 0, 0, color=color)
+
+    # Convert the padded PIL Image back to numpy array
+    # padded_image_np = np.array(padded_image)
+
+    return padded_image

shared/utils/io.py

@@ -0,0 +1,194 @@
+"""
+Utilities for input-output loading/saving.
+"""
+
+from typing import Any, List
+import yaml
+import pickle
+import json
+import pandas as pd
+
+
+class PrettySafeLoader(yaml.SafeLoader):
+    """Custom loader for reading YAML files"""
+    def construct_python_tuple(self, node):
+        return tuple(self.construct_sequence(node))
+
+
+PrettySafeLoader.add_constructor(
+    u'tag:yaml.org,2002:python/tuple',
+    PrettySafeLoader.construct_python_tuple
+)
+
+
+def load_yml(path: str, loader_type: str = 'default'):
+    """Read params from a yml file.
+
+    Args:
+        path (str): path to the .yml file
+        loader_type (str, optional): type of loader used to load yml files. Defaults to 'default'.
+
+    Returns:
+        Any: object (typically dict) loaded from .yml file
+    """
+    assert loader_type in ['default', 'safe']
+
+    loader = yaml.Loader if (loader_type == "default") else PrettySafeLoader
+
+    with open(path, 'r') as f:
+        data = yaml.load(f, Loader=loader)
+
+    return data
+
+
+def save_yml(data: dict, path: str):
+    """Save params in the given yml file path.
+
+    Args:
+        data (dict): data object to save
+        path (str): path to .yml file to be saved
+    """
+    with open(path, 'w') as f:
+        yaml.dump(data, f, default_flow_style=False)
+
+
+def load_pkl(path: str, encoding: str = "ascii"):
+    """Loads a .pkl file.
+
+    Args:
+        path (str): path to the .pkl file
+        encoding (str, optional): encoding to use for loading. Defaults to "ascii".
+
+    Returns:
+        Any: unpickled object
+    """
+    return pickle.load(open(path, "rb"), encoding=encoding)
+
+
+def save_pkl(data: Any, path: str) -> None:
+    """Saves given object into .pkl file
+
+    Args:
+        data (Any): object to be saved
+        path (str): path to the location to be saved at
+    """
+    with open(path, 'wb') as f:
+        pickle.dump(data, f)
+
+
+def load_json(path: str) -> dict:
+    """Helper to load json file"""
+    with open(path, 'rb') as f:
+        data = json.load(f)
+    return data
+
+
+def save_json(data: dict, path: str):
+    """Helper to save `dict` as .json file."""
+    with open(path, 'w') as f:
+        json.dump(data, f, indent=2)
+
+
+def load_txt(path: str):
+    """Loads lines of a .txt file.
+
+    Args:
+        path (str): path to the .txt file
+
+    Returns:
+        List: lines of .txt file
+    """
+    with open(path) as f:
+        lines = f.read().splitlines()
+    return lines
+
+
+def save_txt(data: dict, path: str):
+    """Writes data (lines) to a txt file.
+
+    Args:
+        data (dict): List of strings
+        path (str): path to .txt file
+    """
+    assert isinstance(data, list)
+
+    lines = "\n".join(data)
+    with open(path, "w") as f:
+        f.write(str(lines))
+
+
+def read_spreadsheet(sheet_id, gid, url=None, drop_na=True, **kwargs):
+    if url is None:
+        BASE_URL = 'https://docs.google.com/spreadsheets/d/'
+        url = BASE_URL + sheet_id + f'/export?gid={gid}&format=csv'
+    df = pd.read_csv(url, **kwargs)
+    
+    if drop_na:
+        # drop all rows which have atleast 1 NaN value
+        df = df.dropna(axis=0)
+
+    return df
+
+
+def load_midi(file, rate=16000):
+    import pretty_midi
+    assert file.endswith('.mid')
+    pm = pretty_midi.PrettyMIDI(file)
+    y = pm.synthesize(fs=rate)
+    return y, rate
+
+
+def load_ptz(path):
+    import gzip
+    import torch
+    with gzip.open(path, 'rb') as f:
+        data = torch.load(f)
+    return data
+
+
+def save_video(frames, path, fps=30):
+    import imageio
+    imageio.mimwrite(path, frames, fps=fps)
+
+
+def read_spreadsheet(sheet_id, gid, gid_key="granularity", **kwargs):
+    BASE_URL = 'https://docs.google.com/spreadsheets/d/'
+    df = df = pd.read_csv(BASE_URL + sheet_id + f'/export?gid={gid}&format=csv', **kwargs)
+    return df
+
+
+def load_jsonl(file_path: str) -> list:
+    """Load data from a JSONL file.
+    
+    Args:
+        file_path (str): Path to the JSONL file
+        
+    Returns:
+        list: List of dictionaries, where each dictionary is a JSON object from the file
+        
+    Example:
+        >>> data = load_jsonl("path/to/file.jsonl")
+        >>> print(data[0])  # Print first JSON object
+    """
+    data = []
+    with open(file_path, 'r', encoding='utf-8') as f:
+        for line in f:
+            if line.strip():  # Skip empty lines
+                data.append(json.loads(line))
+    return data
+
+
+def save_jsonl(data: list, file_path: str) -> None:
+    """Save data to a JSONL file.
+    
+    Args:
+        data (list): List of dictionaries to save
+        file_path (str): Path where to save the JSONL file
+        
+    Example:
+        >>> data = [{"text": "hello"}, {"text": "world"}]
+        >>> save_jsonl(data, "output.jsonl")
+    """
+    with open(file_path, 'w', encoding='utf-8') as f:
+        for item in data:
+            f.write(json.dumps(item) + '\n')

shared/utils/keypoint_matching.py

@@ -0,0 +1,330 @@
+"""Implements keypoint matching for a pair of images."""
+import os
+import numpy as np
+import PIL
+import cv2
+import matplotlib.pyplot as plt
+
+
+def show_single_image(img, figsize=(7, 5), title="Single image"):
+    """Displays a single image."""
+    fig = plt.figure(figsize=figsize)
+    plt.axis("off")
+    plt.imshow(img)
+    plt.title(title)
+    plt.show()
+
+
+def show_two_images(img1, img2, title="Two images"):
+    """Displays a pair of images."""
+    fig, ax = plt.subplots(1, 2, figsize=(10, 5), constrained_layout=True)
+
+    ax[0].axis("off")
+    ax[0].imshow(img1)
+
+    ax[1].axis("off")
+    ax[1].imshow(img2)
+
+    plt.suptitle(title)
+    plt.show()
+
+
+def show_three_images(img1, img2, img3, ax1_title="", ax2_title="", ax3_title="", title="Three images"):
+    """Displays a triplet of images."""
+    fig, ax = plt.subplots(1, 3, figsize=(15, 5), constrained_layout=True)
+
+    ax[0].axis("off")
+    ax[0].imshow(img1)
+    ax[0].set_title(ax1_title)
+
+    ax[1].axis("off")
+    ax[1].imshow(img2)
+    ax[1].set_title(ax2_title)
+
+    ax[2].axis("off")
+    ax[2].imshow(img3)
+    ax[2].set_title(ax3_title)
+
+    plt.suptitle(title)
+    plt.show()
+
+
+class KeypointMatcher:
+    """Class for Keypoint matching for a pair of images."""
+
+    def __init__(self, **sift_args) -> None:
+        self.SIFT = cv2.SIFT_create(**sift_args)
+        self.BFMatcher = cv2.BFMatcher()
+    
+    @staticmethod
+    def _check_images(img1: np.ndarray, img2: np.ndarray):
+        assert isinstance(img1, np.ndarray)
+        assert len(img1.shape) == 2
+
+        assert isinstance(img2, np.ndarray)
+        assert len(img2.shape) == 2
+
+        # assert img1.shape == img2.shape
+    
+    @staticmethod
+    def _show_matches(img1, kp1, img2, kp2, matches, K=10, figsize=(10, 5), drawMatches_args=dict(matchesThickness=3, singlePointColor=(0, 0, 0))):
+        """Displays matches found in the image"""
+        selected_matches = np.random.choice(matches, K)
+        img3 = cv2.drawMatches(img1, kp1, img2, kp2, selected_matches, outImg=None, **drawMatches_args)
+        show_single_image(img3, figsize=figsize, title=f"Randomly selected K = {K} matches between the pair of images.")
+        return img3
+
+    def match(self, img1: PIL.Image, img2: PIL.Image, show_matches: bool = True):
+        """Finds, describes and matches keypoints in given pair of images."""
+        
+        img1 = np.array(img1)
+        img1 = cv2.cvtColor(img1, cv2.COLOR_RGB2GRAY)
+        
+        img2 = np.array(img2)
+        img2 = cv2.cvtColor(img2, cv2.COLOR_RGB2GRAY)
+        
+        # check input images
+        self._check_images(img1, img2)
+
+        # find kps and descriptors in each image
+        kp1, des1 = self.SIFT.detectAndCompute(img1, None)
+        kp2, des2 = self.SIFT.detectAndCompute(img2, None)
+
+        # compute matches via Brute-force matching
+        matches = self.BFMatcher.match(des1, des2)
+
+        # sort them in the order of their distance
+        matches = sorted(matches, key = lambda x:x.distance)
+
+        if show_matches:
+            self._show_matches(img1, kp1, img2, kp2, matches)
+
+        return matches, kp1, des1, kp2, des2
+
+
+def warp(im, M, output_shape):
+    out = np.zeros((output_shape[0], output_shape[1]))
+    for i in range(output_shape[0]):
+        for j in range(output_shape[1]):
+            u, v = np.array([[i, j, 0, 0, 1, 0], [0, 0, i, j, 0, 1]]) @ M
+            u = int(round(u))
+            v = int(round(v))
+            if im.shape[0] > u >= 0 and im.shape[1] > v >= 0:
+                out[i, j] = im[u, v]
+
+    return out
+
+
+def project_2d_to_6d(X: np.ndarray):
+    """Projects X (N x 2) to Z (2N x 6) space."""
+    N = len(X)
+    assert X.shape == (N, 2)
+
+    Z = np.zeros((2 * N, 6))
+    # in columns 0 to 2, fill even indexed rows of Z with X, and fill 5th column with 1
+    Z[::2, 0:2] = X
+    Z[::2, 4] = 1.0
+    # in columns 2 to 4, fill odd indexed rows of Z with X
+    Z[1::2, 2:4] = X
+    Z[1::2, 5] = 1.0
+
+    return Z
+
+
+def project_6d_to_2d(Z: np.ndarray):
+    """Projects Z (2N x 6) to X (N x 2) space."""
+    N = len(Z) // 2
+    assert Z.shape == (2 * N, 6)
+
+    X_from_even_rows = Z[::2, 0:2]
+    X_from_odd_rows = Z[1::2, 2:4]
+    assert (X_from_even_rows == X_from_odd_rows).all()
+
+    return X_from_even_rows
+
+
+
+def project_2d_to_1d(X: np.ndarray):
+    """Returns X (N x 2) from Z (2N, 1)"""
+    N = len(X)
+    X_stretched = np.zeros(2 * N)
+    X_stretched[::2] = X[:, 0]
+    X_stretched[1::2] = X[:, 1]
+    return X_stretched
+
+
+def project_1d_to_2d(Z: np.ndarray):
+    """Returns X (N x 2) from Z (2N, 1)"""
+    N = len(Z) // 2
+    assert Z.shape == (2 * N,)
+
+    X = np.zeros((N, 2))
+    X[:, 0] = Z[::2]
+    X[:, 1] = Z[1::2]
+
+    return X
+
+
+def rigid_body_transform(X: np.ndarray, params: np.ndarray):
+    """Performs rigid body transformation of points X (N x 2) using params (6 x 1 flattened)"""
+    N = len(X)
+    assert X.shape == (N, 2)
+
+    X = project_2d_to_6d(X)
+
+    X_transformed = np.matmul(X, params)
+    X_transformed = project_1d_to_2d(X_transformed)
+    assert X_transformed.shape == (N, 2)
+
+    return X_transformed
+
+
+def rigid_body_transform_params(X1: np.ndarray, X2: np.ndarray):
+    """Returns rigid-body transform parameters RT (6 x 1) assuming transformation between X1 and X2"""
+    N = len(X1)
+    assert X1.shape == X2.shape
+    assert X1.shape == (N, 2)
+
+    # X2 = X1 * params => params = psuedoinverse(X1) * X2
+    X1_expanded = project_2d_to_6d(X1)
+    assert X1_expanded.shape == (2 * N, 6)
+
+    X2_stretched = project_2d_to_1d(X2)
+    assert X2_stretched.shape == (2 * N,)
+
+    params = np.dot(np.linalg.pinv(X1_expanded), X2_stretched)
+    return params
+
+
+class ImageAlignment:
+    """Class to perform alignment of a pair of images given keypoints."""
+
+    def __init__(self) -> None:
+        pass
+    
+    @staticmethod
+    def show_transformed_points(img1, img2, X1, kp1, kp2, matches, params, num_inliers, num_to_show=20):
+        import matplotlib.cm as cm
+
+        H1, W1 = img1.shape
+        H2, W2 = img2.shape
+        img = np.hstack([img1, img2])
+
+        random_matches = np.random.choice(matches, num_to_show)
+
+        fig, ax = plt.subplots(1, 1, figsize=(15, 6))
+        colors = cm.rainbow(np.linspace(0, 1, num_to_show))
+
+        for i, match in enumerate(random_matches):
+
+            # select a single match to visualize
+            x1, y1 = kp1[match.queryIdx].pt
+            x2, y2 = kp2[match.trainIdx].pt
+
+            # get (x1, y1) transformed to (x1_transformed, y1_transformed)
+            A = project_2d_to_6d(np.array([[x1, y1]]))
+            (x1_transformed, y1_transformed) = np.dot(A, params)
+
+            ax.imshow(img, cmap="gray")
+            ax.axis("off")
+            ax.scatter(x1_transformed + W1, y1_transformed, s=200, marker="x", color=colors[i])
+            ax.plot(
+                (x1, x1_transformed + W1), (y1, y1_transformed),
+                linestyle="--", color=colors[i], marker="o",
+            )
+
+        ax.set_title(
+            f"Points in image 1 mapped to transformed points estimated by {num_inliers} points.",
+            fontsize=18,
+        )
+
+        os.makedirs("./results/", exist_ok=True)
+        plt.savefig(f"./results/match_transformed_inliers_{num_inliers}.png", bbox_inches="tight")
+        plt.show()
+
+    def ransac(
+            self, img1, kp1, img2, kp2, matches, num_matches=6, max_iter=500,
+            radius_in_px=10, show_transformed=True, inlier_th_for_show=1000
+        ):
+        """Performs RANSAC to find best matches."""
+
+        best_inlier_count = 0
+        best_params = None
+
+        # get coordinates of all points in image 1
+        X1 = np.array([kp1[matches[i].queryIdx].pt for i in range(len(matches))])
+
+        # get coordinates of all points in image 2
+        X2 = np.array([kp2[matches[i].trainIdx].pt for i in range(len(matches))])
+
+        for i in range(max_iter):
+            # choose matches randomly
+            selected_matches = np.random.choice(matches, num_matches)
+
+            # get matched keypoints in img1
+            X1_selected = np.array([kp1[selected_matches[i].queryIdx].pt for i in range(len(selected_matches))])
+
+            # get matched keypoints in img2
+            X2_selected = np.array([kp2[selected_matches[i].trainIdx].pt for i in range(len(selected_matches))])
+
+            # get transformation parameters
+            params = rigid_body_transform_params(X1_selected, X2_selected)
+            
+            # transform X1 to get X2_transformed
+            X2_transformed = rigid_body_transform(X1, params)
+
+            # find inliers
+            diff = np.linalg.norm(X2_transformed - X2, axis=1)
+            indices = diff < radius_in_px
+            num_inliers = sum(indices)
+            if num_inliers > best_inlier_count:
+                print(f"Found {num_inliers} inliers!")
+                best_params = params
+                best_inlier_count = num_inliers
+
+                if show_transformed and num_inliers > inlier_th_for_show:
+                    self.show_transformed_points(img1, img2, X1, kp1, kp2, matches, best_params, num_inliers)
+
+        return best_params
+    
+    def align(
+            self, img1, kp1, img2, kp2, matches, num_matches=6,
+            max_iter=500, show_warped_image=True,
+            save_warped=False, path="results/sample.png",
+            method="custom"
+        ):
+        best_params = self.ransac(img1, kp1, img2, kp2, matches, max_iter=max_iter, num_matches=num_matches)
+
+        # apply the affine transformation using cv2.warpAffine()
+        rows, cols = img1.shape[:2]
+
+        if method == 'custom':
+            img1_warped = warp(img1, best_params, (rows, cols))
+        else:
+            M = np.zeros((2, 3))
+            M[0, :2] = best_params[:2]
+            M[1, :2] = best_params[2:4]
+            M[0, 2] = best_params[4]
+            M[1, 2] = best_params[5]
+            img1_warped = cv2.warpAffine(img1, M, (cols, rows))
+
+        if show_warped_image:
+            show_three_images(
+                img1, img2, img1_warped, title="",
+                ax1_title="Image 1", ax2_title="Image 2", ax3_title="Transformation: Image 1 to Image 2",
+            )
+
+        if save_warped:
+            plt.imsave(path, img1_warped)
+
+        return best_params
+
+
+if __name__ == "__main__":
+    # read & show images
+    boat1 = cv2.imread('boat1.pgm', cv2.IMREAD_GRAYSCALE)
+    boat2 = cv2.imread('boat2.pgm', cv2.IMREAD_GRAYSCALE)
+    show_two_images(boat1, boat2, title="Given pair of images.")
+
+    kp_matcher = KeypointMatcher(contrastThreshold=0.1, edgeThreshold=5)
+    matches, kp1, des1, kp2, des2 = kp_matcher.match(boat1, boat2, show_matches=True)