refactor with Sajad's PR - DiT data to common dir

Author: Huy Vu2

…tron/data/dit/base_energon_datamodule.py …n/data/common/base_energon_datamodule.pydfm/src/megatron/data/dit/base_energon_datamodule.py renamed to dfm/src/megatron/data/common/base_energon_datamodule.py dfm/src/megatron/data/dit/base_energon_datamodule.py renamed to dfm/src/megatron/data/common/base_energon_datamodule.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
 
 

…data/dit/diffusion_energon_datamodule.py …a/common/diffusion_energon_datamodule.pydfm/src/megatron/data/dit/diffusion_energon_datamodule.py renamed to dfm/src/megatron/data/common/diffusion_energon_datamodule.py dfm/src/megatron/data/dit/diffusion_energon_datamodule.py renamed to dfm/src/megatron/data/common/diffusion_energon_datamodule.py

@@ -7,8 +7,8 @@
 from megatron.energon.task_encoder.base import stateless
 from megatron.energon.task_encoder.cooking import Cooker, basic_sample_keys
 
-from dfm.src.megatron.data.dit.diffusion_sample import DiffusionSample
-from dfm.src.megatron.data.dit.sequence_packing_utils import first_fit_decreasing
+from dfm.src.megatron.data.common.diffusion_sample import DiffusionSample
+from dfm.src.megatron.data.common.sequence_packing_utils import first_fit_decreasing
 
 
 def cook(sample: dict) -> dict:

…rc/megatron/data/dit/diffusion_sample.py …megatron/data/common/diffusion_sample.pydfm/src/megatron/data/dit/diffusion_sample.py renamed to dfm/src/megatron/data/common/diffusion_sample.py dfm/src/megatron/data/dit/diffusion_sample.py renamed to dfm/src/megatron/data/common/diffusion_sample.py

@@ -21,8 +21,8 @@
 from megatron.core import parallel_state
 from megatron.energon import SkipSample, stateless
 
-from dfm.src.megatron.data.dit.diffusion_sample import DiffusionSample
-from dfm.src.megatron.data.dit.diffusion_task_encoder_with_sp import DiffusionTaskEncoderWithSequencePacking
+from dfm.src.megatron.data.common.diffusion_sample import DiffusionSample
+from dfm.src.megatron.data.common.diffusion_task_encoder_with_sp import DiffusionTaskEncoderWithSequencePacking
 
 
 class DiTTaskEncoder(DiffusionTaskEncoderWithSequencePacking):

…ta/dit/diffusion_task_encoder_with_sp.py …common/diffusion_task_encoder_with_sp.pydfm/src/megatron/data/dit/diffusion_task_encoder_with_sp.py renamed to dfm/src/megatron/data/common/diffusion_task_encoder_with_sp.py dfm/src/megatron/data/dit/diffusion_task_encoder_with_sp.py renamed to dfm/src/megatron/data/common/diffusion_task_encoder_with_sp.py

@@ -19,7 +19,7 @@
 from megatron.bridge.data.utils import DatasetBuildContext, DatasetProvider
 from torch import int_repr
 
-from dfm.src.megatron.data.dit.diffusion_energon_datamodule import DiffusionDataModuleConfig, DiffusionDataModule
+from dfm.src.megatron.data.common.diffusion_energon_datamodule import DiffusionDataModuleConfig, DiffusionDataModule
 from dfm.src.megatron.data.wan.wan_taskencoder import WanTaskEncoder
 
 

…atron/data/dit/sequence_packing_utils.py …on/data/common/sequence_packing_utils.pydfm/src/megatron/data/dit/sequence_packing_utils.py renamed to dfm/src/megatron/data/common/sequence_packing_utils.py dfm/src/megatron/data/dit/sequence_packing_utils.py renamed to dfm/src/megatron/data/common/sequence_packing_utils.py

@@ -19,7 +19,7 @@
 import torch
 from megatron.bridge.data.utils import DatasetBuildContext, DatasetProvider
 
-from dfm.src.megatron.data.dit.diffusion_energon_datamodule import DiffusionDataModule
+from dfm.src.megatron.data.common.diffusion_energon_datamodule import DiffusionDataModule
 from dfm.src.megatron.data.wan.wan_taskencoder import WanTaskEncoder
 
 

dfm/src/megatron/data/dit/dit_taskencoder.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 
 from dataclasses import dataclass
-from dfm.src.megatron.data.dit.diffusion_sample import DiffusionSample
+from dfm.src.megatron.data.common.diffusion_sample import DiffusionSample
 
 
 @dataclass

dfm/src/megatron/data/wan/wan_energon_datamodule.py

@@ -21,7 +21,7 @@
 from megatron.core import parallel_state
 from typing import List
 from megatron.energon import SkipSample
-from dfm.src.megatron.data.dit.diffusion_task_encoder_with_sp import DiffusionTaskEncoderWithSequencePacking
+from dfm.src.megatron.data.common.diffusion_task_encoder_with_sp import DiffusionTaskEncoderWithSequencePacking
 from megatron.energon.task_encoder.cooking import Cooker, basic_sample_keys
 from dfm.src.megatron.data.wan.wan_sample import WanSample
 from dfm.src.megatron.model.wan.utils import grid_sizes_calculation, patchify