move data modules to common, move cosmos to common/tokenizer instad of common/models. Make inference compatible with the new data loading.

Signed-off-by: Sajad Norouzi <snorouzi@nvidia.com>

Author: sajadn

dfm/src/common/models/__init__ .py dfm/src/common/tokenizers/__init__ .pydfm/src/common/models/__init__ .py renamed to dfm/src/common/tokenizers/__init__ .py

@@ -25,7 +25,7 @@
 from huggingface_hub import hf_hub_download
 from tqdm import tqdm
 
-from dfm.src.common.models.cosmos.cosmos1.tokenizer_utils import (
+from dfm.src.common.tokenizers.cosmos.cosmos1.video_tokenizer_utils import (
     load_jit_model,
     numpy2tensor,
     pad_video_batch,

dfm/src/common/models/cosmos/__init__.py …src/common/tokenizers/cosmos/__init__.pydfm/src/common/models/cosmos/__init__.py renamed to dfm/src/common/tokenizers/cosmos/__init__.py dfm/src/common/models/cosmos/__init__.py renamed to dfm/src/common/tokenizers/cosmos/__init__.py

@@ -14,7 +14,6 @@
 
 import imageio
 import numpy as np
-import torch
 
 
 def save_video(
@@ -35,12 +34,3 @@ def save_video(
 
     print("video_save_path", video_save_path)
     imageio.mimsave(video_save_path, grid, "mp4", **kwargs)
-
-
-def print_dict(dict):
-    for key, value in dict.items():
-        if isinstance(value, torch.Tensor):
-            print(key, value.shape)
-        else:
-            print(key, value)
-    print("-" * 40)

…common/models/cosmos/cosmos1/__init__.py …on/tokenizers/cosmos/cosmos1/__init__.pydfm/src/common/models/cosmos/cosmos1/__init__.py renamed to dfm/src/common/tokenizers/cosmos/cosmos1/__init__.py dfm/src/common/models/cosmos/cosmos1/__init__.py renamed to dfm/src/common/tokenizers/cosmos/cosmos1/__init__.py

@@ -24,7 +24,7 @@
 
 class EnergonMultiModalDataModule:
     """
-    A PyTorch Lightning DataModule for handling multimodal datasets with images and text.
+    A DataModule for handling multimodal datasets with images and text.
 
     This data module is designed to work with multimodal datasets that involve both images and text.
     It provides a seamless interface to load training and validation data, manage batching, and handle

…cosmos/cosmos1/causal_video_tokenizer.py …cosmos/cosmos1/causal_video_tokenizer.pydfm/src/common/models/cosmos/cosmos1/causal_video_tokenizer.py renamed to dfm/src/common/tokenizers/cosmos/cosmos1/causal_video_tokenizer.py dfm/src/common/models/cosmos/cosmos1/causal_video_tokenizer.py renamed to dfm/src/common/tokenizers/cosmos/cosmos1/causal_video_tokenizer.py

@@ -23,7 +23,7 @@
 from megatron.energon import DefaultTaskEncoder, get_train_dataset
 from torch import int_repr
 
-from dfm.src.megatron.data.dit.base_energon_datamodule import EnergonMultiModalDataModule
+from dfm.src.megatron.data.common.base_energon_datamodule import EnergonMultiModalDataModule
 from dfm.src.megatron.data.dit.dit_taskencoder import DiTTaskEncoder