multimodal-llm/train.py at main · Open-Superintelligence-Lab/multimodal-llm · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32

import torch
from configs.dit_config import DiTConfig
from models.dit import DiT
from data.video_loader import get_dataloaders
from training.video_trainer import train_video_model

def main():
    config = DiTConfig()
    # Faster training for testing
    config.train_steps = 500
    config.log_every = 10
    config.save_every = 500
    config.batch_size = 2 # Small batch size


    config.hidden_size = 384 # Reduce model size (DiT-S/2 equivalent-ish)
    config.depth = 12
    config.num_heads = 6


    print("Initializing Model...")
    model = DiT(config)

    print("Getting Dataloaders...")
    train_loader, val_loader = get_dataloaders(config)

    print("Starting Training...")
    train_video_model(model, train_loader, val_loader, config)

if __name__ == "__main__":
    main()