Add modality info and transforms for raw RGB, tokens, video descriptions, video transcripts, and video bounding boxes.

cfgs/default/4m/alphas_mixture/main/mix_mod7_all2all_rgb2dense_a0.5.yaml

@@ -1,7 +1,7 @@
 # Half of all samples from this mixture are rgb2dense, half are all2all
-rgb@224:
-  input_alphas: [1000.0, 0.5]
-  target_alphas: [0., 0.]  # RGB is not a target
+rgb@224: # num cols is the number of sampling strategies to try.
+  input_alphas: [1000.0, 0.5] # NOTE: input alphas: how much weight to put on the input modalities for encoding.
+  target_alphas: [0., 0.]  # RGB is not a target # NOTE: target alphas - how much weight to put on the modalities to be unmasked/sampled during decoding.
 tok_rgb@224:
   input_alphas: [0., 0.5]
   target_alphas: [0.5, 0.5]

cfgs/default/4m/alphas_mixture/video/mix_mod3_rgb_tok_a0.5.yaml

@@ -0,0 +1,7 @@
+# What's up with the naming configs here?
+video_rgb@224:
+  input_alphas: [1000.0, 0.5]
+  target_alphas: [0., 0.]  # RGB is not a target
+video_tok_rgb@224:
+  input_alphas: [0., 0.5]
+  target_alphas: [0.5, 0.5]

cfgs/default/4m/alphas_mixture/video/mix_mod3_rgb_tok_det_a0.5.yaml

@@ -0,0 +1,11 @@
+# What's up with the naming configs here?
+video_rgb@224:
+  input_alphas: [1000.0, 0.5]
+  target_alphas: [0., 0.]  # RGB is not a target
+video_tok_rgb@224:
+  input_alphas: [0., 0.5]
+  target_alphas: [0.5, 0.5]
+det:
+  input_alphas: [0., 0.5]
+  target_alphas: [0.5, 0.5]
+  keep: ['random', 'random']

cfgs/default/4m/data/video/mix_mod3_rgb_tok_det_to_all_a0.5.yaml

@@ -0,0 +1,44 @@
+train:
+  datasets:
+    my_video_dataset:
+      type: multimodal
+
+      # Input and output domain names, separated by hyphen
+      in_domains: video_rgb@224-video_tok_rgb@224-video_det
+      out_domains: video_rgb@224-video_tok_rgb@224-video_det
+
+      # Dirichlet alphas concentration parameter for input and output. 
+      # Can be either one value, or one value per input modality separated by hyphen.
+      input_alphas: null
+      target_alphas: null
+      # Path to specific alphas configuration to enable mixture of Dirichlets. 
+      # If provided, overrides input_alphas and target_alphas
+      alphas_config: "cfgs/default/4m/alphas_mixture/video/mix_mod3_rgb_tok_det_a0.5.yaml"
+
+      # Optionally, min_input_tokens, min_target_tokens, num_input_tokens, num_target_tokens can be specified here
+      # If so, they will override the values provided in the main config
+
+      # Data can either be local or on cloud storage (e.g. S3), see data docs for more info
+      # Use braceexpand notation to indicate shard range (e.g. shard-{0000..9999}.tar)
+      # Use brackets to indicate multiple modalities (e.g. [modality1,modality2,modality3]) 
+      data_path: '/store/swissai/a08/data/4m/train/[video_rgb,video_rgb_tok,video_det]/shard-{00000..00100}.tar' # TODO: need to reformat the data correctly here. 
+      use_wds: True # Use webdataset
+      wds_n_repeats: 4 # Number of repeats for webdataset loader to improve efficiency
+      wds_shuffle_buffer_tar: 1_000 # Webdatasets shuffle buffer after loading tar files
+      wds_shuffle_buffer_repeat: 1_000 # Webdatasets shuffle buffer after repeating samples
+
+      main_augment_domain: video_rgb@224 # Select from which modality to get the original full image size (mostly important for resizing bounding boxes)
+      aligned_captions: True # Align captions to crop_settings # TODO: tbd?
+      tok_train_aug: True # Apply data augmentation to tokens (if multiple crop settings are available) # TODO: tbd?
+
+      # modality_name_map: # Use modality_name_map to define a mapping from a folder name to a modality name
+      #   tok_rgb_folder_name: tok_rgb@224
+      #   tok_depth_folder_nme: tok_depth@224
+      #   ...
+
+  weights: [1.0] # Sampling weights for the training datasets
+
+val:
+  datasets:
+    my_video_dataset: 
+      data_path: '/store/swissai/a08/data/4m/val/[video_rgb,video_rgb_tok,video_det]/shard-{00000..00100}.tar'

cfgs/default/4m/data/video/mix_mod3_rgb_tok_to_all_a0.5.yaml

@@ -0,0 +1,44 @@
+train:
+  datasets:
+    my_video_dataset:
+      type: multimodal
+
+      # Input and output domain names, separated by hyphen
+      in_domains: video_rgb@224-video_tok_rgb@224
+      out_domains: video_rgb@224-video_tok_rgb@224
+
+      # Dirichlet alphas concentration parameter for input and output. 
+      # Can be either one value, or one value per input modality separated by hyphen.
+      input_alphas: null
+      target_alphas: null
+      # Path to specific alphas configuration to enable mixture of Dirichlets. 
+      # If provided, overrides input_alphas and target_alphas
+      alphas_config: "cfgs/default/4m/alphas_mixture/video/mix_mod3_rgb_tok_a0.5.yaml"
+
+      # Optionally, min_input_tokens, min_target_tokens, num_input_tokens, num_target_tokens can be specified here
+      # If so, they will override the values provided in the main config
+
+      # Data can either be local or on cloud storage (e.g. S3), see data docs for more info
+      # Use braceexpand notation to indicate shard range (e.g. shard-{0000..9999}.tar)
+      # Use brackets to indicate multiple modalities (e.g. [modality1,modality2,modality3]) 
+      data_path: '/store/swissai/a08/data/4m/cleaned/train/[video_rgb,video_tok_rgb]/0000000000.tar' # TODO: need to reformat the data correctly here. 
+      use_wds: True # Use webdataset
+      wds_n_repeats: 4 # Number of repeats for webdataset loader to improve efficiency
+      wds_shuffle_buffer_tar: 1_000 # Webdatasets shuffle buffer after loading tar files
+      wds_shuffle_buffer_repeat: 1_000 # Webdatasets shuffle buffer after repeating samples
+
+      main_augment_domain: video_rgb@224 # Select from which modality to get the original full image size (mostly important for resizing bounding boxes)
+      aligned_captions: True # Align captions to crop_settings # TODO: tbd?
+      tok_train_aug: True # Apply data augmentation to tokens (if multiple crop settings are available) # TODO: tbd?
+
+      # modality_name_map: # Use modality_name_map to define a mapping from a folder name to a modality name
+      #   tok_rgb_folder_name: tok_rgb@224
+      #   tok_depth_folder_nme: tok_depth@224
+      #   ...
+
+  weights: [1.0] # Sampling weights for the training datasets
+
+# val:
+#   datasets:
+#     my_video_dataset: 
+#       data_path: '/store/swissai/a08/data/4m/val/[video_rgb,video_tok_rgb]/00000{00175..00199}.tar'

cfgs/default/4m/models/video/4m-b_mod3.yaml

@@ -0,0 +1,46 @@
+# Config for DDP
+
+# Arch: SwiGLU No Bias
+# Modalities: Mix of rgb2all and all2all, with alphas=0.5
+# To be run on 64 GPUs for batch size = 8192
+run_name: auto
+
+# Input & output
+num_input_tokens: 128
+num_target_tokens: 128
+loss_type: mod
+
+# Architecture
+model: fm_base_12e_12d_swiglu_nobias
+patch_size: 16
+input_size: 224
+dtype: bfloat16
+tokenizer_path: "fourm/utils/tokenizer/trained/text_tokenizer_4m_wordpiece_30k.json"
+
+# Train
+epochs: -1
+total_tokens: 500 # in billions
+opt: adamw
+blr: 0.0001 # this is base_lr = 1e-4, lr = base_lr * batch_size / 256
+min_blr: 0.
+warmup_epochs: -1
+warmup_tokens: 10 # in billions
+batch_size: 128 # 128 x 64 = 8192
+
+# Data
+
+data_config: "cfgs/default/4m/data/video/mix_mod3_rgb_tok_to_all_a0.5.yaml"
+s3_data_endpoint: null # Change me
+eval_freq: 1
+fixed_eval: True
+epoch_size: 10_000_000 # Number of samples per "epoch"
+
+# Saving
+save_ckpt_freq: 1
+output_dir: 'output/auto'
+
+# Wandb
+log_wandb: False # Set to True to log to Weights & Biases
+wandb_project: '4m-train'
+wandb_entity: null # Change if needed
+wandb_run_name: auto

fourm/data/masking.py

@@ -128,7 +128,7 @@ def chunk_span_masking(sequence_chunks: List[List[int]], sentinel_to_id: Dict[in
 
 
 
-class UnifiedMasking(object):
+class UnifiedMasking(object): # this defines the masking logic
     def __init__(self,
                  modality_info: Dict,
                  text_tokenizer: Optional[Tokenizer],

fourm/data/modality_info.py

@@ -29,6 +29,11 @@
     ColorPaletteTransform,
     SAMInstanceTokTransform,
     SAMInstanceTransform,
+    VideoDescriptionTransform,
+    VideoDetectionTransform,
+    VideoRGBTransform,
+    VideoTokTransform,
+    VideoTranscriptTransform,
 )
 from fourm.models.decoder_embeddings import ImageTokenDecoderEmbedding, SequenceDecoderEmbedding
 from fourm.models.encoder_embeddings import (
@@ -39,6 +44,7 @@
 )
 from fourm.utils import generate_uint15_hash
 
+# Specifications about different modalities
 MODALITY_INFO = {
     # 4M-7 modalities
     "rgb@224": {
@@ -406,6 +412,38 @@
     },
 }
 
+VIDEO_MODALITY_INFO = {
+    ### Video modalities
+    # TODO: do we need to keep the image versions? These probably should generalize over those, right?
+    "video_rgb@224": {
+        **MODALITY_INFO["rgb@224"],
+        "id": generate_uint15_hash("video_rgb@224"),
+        "path": "video_rgb",  # TODO: video_rgb or keep aas rgb? (probably only keep rgb if this generalizes over single images too)
+    },
+    "video_description": {
+        **MODALITY_INFO["caption"],  # TODO: do we want to increase the default 'max_tokens/max_length' from 256?
+        "id": generate_uint15_hash("video_description"),
+    },
+    "video_transcript": {
+        **MODALITY_INFO["caption"],  # TODO: do we want to increase the default 'max_tokens/max_length' from 256?
+        "id": generate_uint15_hash("video_transcript"),
+    },
+    "video_det": {
+        **MODALITY_INFO["det"],
+        "id": generate_uint15_hash("video_det"),
+    },
+    "video_tok_rgb@224": {
+        **MODALITY_INFO["tok_rgb@224"],
+        "id": generate_uint15_hash("video_tok_rgb@224"),
+    },
+    "video_tok_clip@224": {
+        **MODALITY_INFO["tok_clip@224"],
+        "id": generate_uint15_hash("video_tok_clip@224"),
+    },
+}
+
+MODALITY_INFO = {**MODALITY_INFO, **VIDEO_MODALITY_INFO}
+
 # Note: @res suffix is ignored for modality transforms
 MODALITY_TRANSFORMS = {
     # 4M-7 modalities
@@ -414,7 +452,7 @@
     "det": DetectionTransform(
         det_threshold=0.6, det_max_instances=None, bbox_order="dist_to_orig", coord_bins=1000, min_visibility=0.0
     ),
-    "tok_rgb": TokTransform(),
+    "tok_rgb": TokTransform(),  # tok_ indicates its a token representation
     "tok_depth": TokTransform(),
     "tok_normal": TokTransform(),
     "tok_semseg": TokTransform(),
@@ -435,6 +473,15 @@
     "tok_imagebind_global": TokTransform(),
     # Other
     "mask_valid": MaskTransform(mask_pool_size=1),
+    # Video
+    "video_rgb": VideoRGBTransform(imagenet_default_mean_and_std=True),  # TODO: check parameters
+    "video_tok_rgb": VideoTokTransform(),  # tok_ indicates its a token representation
+    "video_tok_clip": VideoTokTransform(),  # TODO: check parameters
+    "video_description": VideoDescriptionTransform(aligned_captions=True),  # TODO: check parameters
+    "video_transcript": VideoTranscriptTransform(aligned_captions=True),  # TODO: check parameters
+    "video_det": VideoDetectionTransform(
+        det_threshold=0.6, det_max_instances=None, bbox_order="dist_to_orig", coord_bins=1000, min_visibility=0.0
+    ),  # TODO: check parameters
 }
 
 MODALITY_TRANSFORMS_DIVAE = {