A minimal viable prototype of a multimodal language model demonstrating unified discrete representation for text and video modalities.
Create a simple multimodal model that can:
- Video Captioning: Video → Text (generate captions describing the video)
- Video Question Answering (VQA): Video + Question → Answer
This is an MVP designed to demonstrate multimodal capability with minimal resources.
This is an MVP prototype created to demonstrate multimodal capability. Please note the following:
-
Training Data: The demo uses samples from the training dataset itself. Due to the limited amount of training data, the model has overfit to the caption task, resulting in generated captions that match the ground truth exactly.
-
Video Question Answering (VQA): The VQA task is included to showcase multimodal capability beyond simple captioning. However, the model has not been trained on VQA data, so many outputs contain gibberish text or empty responses. This is expected behavior for an MVP that hasn't been trained on question-answer pairs.
-
Demo Samples: Five cherry-picked samples (input videos + generated text) are available in the
demo_outputs/folder for reference. These samples were selected to best demonstrate the model's multimodal capabilities despite the limitations mentioned above.
Video -> Cosmos-Tokenizer (frozen) -> Discrete Tokens -> DistilGPT-2 -> Text Tokens -> Caption
┌─────────────────────────────────────────────────────────────────────────────┐
│ TRAINING PIPELINE │
└─────────────────────────────────────────────────────────────────────────────┘
INPUT: Video (frames @ 128×128)
│
▼
┌─────────────────────────────────────────────────────────────────────────┐
│ COSMOS TOKENIZER (FROZEN - Pretrained) │
│ ───────────────────────────────────────────────────────────────────────│
│ • Encoder: 3D CNN → Latent Features │
│ • Quantizer: FSQ with levels [8,8,8,5,5,5] → Codebook size: 64,000 │
│ • Output: Discrete token indices (0-63,999) │
│ • Compression: 8× temporal, 8×8 spatial → 512 tokens per video │
└─────────────────────────────────────────────────────────────────────────┘
│
▼
┌─────────────────────────────────────────────────────────────────────────┐
│ SEQUENCE ASSEMBLY │
│ ───────────────────────────────────────────────────────────────────────│
│ <bos> <vid> [video_token_1] ... [video_token_512] </vid> │
│ <txt> [text_token_1] ... [text_token_N] </txt> <eos> │
│ │
│ Video tokens offset: +50,264 (to avoid collision with text vocab) │
└─────────────────────────────────────────────────────────────────────────┘
│
▼
┌─────────────────────────────────────────────────────────────────────────┐
│ MULTIMODAL GPT MODEL │
│ ───────────────────────────────────────────────────────────────────────│
│ │
│ ┌───────────────────────────────────────────────────────────────────┐ │
│ │ VOCABULARY EXTENSION (TRAINED) │ │
│ │ ─────────────────────────────────────────────────────────────────│ │
│ │ • Original DistilGPT-2 vocab: 50,257 tokens │ │
│ │ • Special tokens: +7 (<bos>, <eos>, <vid>, </vid>, etc.) │ │
│ │ • Video tokens: +64,000 │ │
│ │ • Total vocabulary: 114,264 tokens │ │
│ │ • New embeddings initialized from old embeddings' distribution │ │
│ └───────────────────────────────────────────────────────────────────┘ │
│ │
│ ┌───────────────────────────────────────────────────────────────────┐ │
│ │ DISTILGPT-2 BASE (FROZEN - Pretrained) │ │
│ │ ─────────────────────────────────────────────────────────────────│ │
│ │ • Transformer layers: 6 layers (frozen) │ │
│ │ • Attention heads: 12 per layer (frozen) │ │
│ │ • Hidden size: 768 (frozen) │ │
│ │ • Max positions: 1,024 (frozen) │ │
│ └───────────────────────────────────────────────────────────────────┘ │
│ │
│ ┌───────────────────────────────────────────────────────────────────┐ │
│ │ OUTPUT HEAD (TRAINED) │ │
│ │ ─────────────────────────────────────────────────────────────────│ │
│ │ • Language modeling head: Extended to 114,264 vocab (trained) │ │
│ │ • Loss computed only on text tokens (video tokens masked) │ │
│ └───────────────────────────────────────────────────────────────────┘ │
│ │
└─────────────────────────────────────────────────────────────────────────┘
│
▼
OUTPUT: Generated Text Caption
┌─────────────────────────────────────────────────────────────────────────┐
│ KEY: │
│ ───────────────────────────────────────────────────────────────────────│
│ [FROZEN] = Pretrained weights, not updated during training │
│ [TRAINED] = Weights updated during training │
└─────────────────────────────────────────────────────────────────────────┘
- Core LLM: DistilGPT-2 (82M parameters, base frozen, embeddings trained)
- Video Tokenizer: NVIDIA Cosmos-Tokenizer DV8x8x8 (frozen, pretrained)
- Dataset: UCF101 training split from HuggingFace (all videos by default)
- Task: Video Captioning (video -> text)
- Training: 10,000 steps by default
MM_gen_prototype/
+-- configs/
| +-- config.py # Dataclass configurations
+-- data/
| +-- dataset.py # PyTorch Dataset
| +-- raw/ # Custom videos (optional)
+-- models/
| +-- video_tokenizer.py # Cosmos tokenizer wrapper
| +-- multimodal_gpt.py # Extended DistilGPT-2
+-- training/
| +-- trainer.py # Training loop
+-- inference/
| +-- generate.py # Generation utilities
+-- scripts/
| +-- setup_cosmos.py # Setup Cosmos
| +-- preprocess_data.py # Data preprocessing
| +-- train.py # Training script
| +-- demo.py # Demo script
+-- utils/
| +-- helpers.py # Utility functions
+-- requirements.txt
+-- README.md
# Create virtual environment
python -m venv venv
venv\Scripts\activate # Windows
# source venv/bin/activate # Linux/Mac
# Install dependencies
pip install -r requirements.txtRun the setup script to automatically clone the repo, install dependencies, and download the checkpoint:
python scripts/setup_cosmos.pyOptions:
# Use a specific model
python scripts/setup_cosmos.py --model Cosmos-0.1-Tokenizer-DV8x8x8
# Skip dependency installation
python scripts/setup_cosmos.py --skip-deps
# List available models
python scripts/setup_cosmos.py --list-modelsUCF101 from HuggingFace
This automatically downloads videos and generates captions from action labels:
# Process all videos in training split (default)
python scripts/preprocess_data.py --source ucf101-hf
# Or specify number of videos
python scripts/preprocess_data.py --source ucf101-hf --num-videos 500UCF101 from local directory:
# Process all videos (default)
python scripts/preprocess_data.py --source ucf101-local
# Or specify number
python scripts/preprocess_data.py --source ucf101-local --num-videos 500Custom videos:
- Place videos in
data/raw/videos/ - Create
data/raw/captions.json:
[
{"video_id": "vid001", "video_file": "vid001.mp4", "caption": "A person walking"},
{"video_id": "vid002", "video_file": "vid002.mp4", "caption": "A dog running"}
]- Run:
# Process all videos in captions.json (default)
python scripts/preprocess_data.py --source custom
# Or specify number
python scripts/preprocess_data.py --source custom --num-videos 50# Train with default settings (10,000 steps)
python scripts/train.py
# Or customize
python scripts/train.py --max-steps 5000 --batch-size 4Options:
--max-steps: Number of training steps (default: 10,000)--batch-size: Batch size (default: 2)--learning-rate: Learning rate (default: 5e-5)--checkpoint-dir: Where to save checkpoints--no-amp: Disable mixed precision
Default Demo (shows both Captioning and VQA):
# Demo with pre-tokenized data
python scripts/demo.py --num-samples 5
# Save outputs for showcasing
python scripts/demo.py --num-samples 5 --save-outputsSpecific Modes:
# Captioning only
python scripts/demo.py --mode captioning --num-samples 5
# VQA only
python scripts/demo.py --mode vqa --num-samples 5
# VQA with custom questions
python scripts/demo.py --mode vqa --vqa-questions "What is happening?" "What is the person doing?"
# Demo with a specific video
python scripts/demo.py --video path/to/video.mp4
# Interactive mode
python scripts/demo.py --interactiveDemo Outputs: Use --save-outputs to save results to demo_outputs/ folder for showcasing.
pip install -r requirements.txt
python scripts/setup_cosmos.py
python scripts/preprocess_data.py --source ucf101-hf
python scripts/train.py
python scripts/demo.py --num-samples 5Note: By default, preprocessing uses all videos in the training split and training runs for 10,000 steps. Use command-line flags to override these defaults.
Edit configs/config.py to modify:
- Model settings (num_frames, frame_size, etc.)
- Training hyperparameters
- Data paths
pip install datasetspython scripts/setup_cosmos.py- Reduce
batch_size - Reduce
num_framesin config - Use
--no-ampflag
MIT License
UCF101 dataset is licensed under CC BY-NC 4.0 (non-commercial use).