content: #!/usr/bin/env python3
  
  # Hunyuan 3D is licensed under the TENCENT HUNYUAN NON-COMMERCIAL LICENSE AGREEMENT
  # except for the third-party components listed below.
  # Hunyuan 3D does not impose any additional limitations beyond what is outlined
  # in the repsective licenses of these third-party components.
  # Users must comply with all terms and conditions of original licenses of these third-party
  # components and must ensure that the usage of the third party components adheres to
  # all relevant laws and regulations.
  
  # For avoidance of doubts, Hunyuan 3D means the large language models and
  # their software and algorithms, including trained model weights, parameters (including
  # optimizer states), machine-learning model code, inference-enabling code, training-enabling code,
  # fine-tuning enabling code and other elements of the foregoing made publicly available
  # by Tencent in accordance with TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT.
  
  import pytorch_lightning as pl
  from torch.utils.data import Dataset, ConcatDataset, DataLoader
  from torch.utils.data.distributed import DistributedSampler
  
  
  class DataModuleFromConfig(pl.LightningDataModule):
      def __init__(
          self,
          batch_size=8,
          num_workers=4,
          train=None,
          validation=None,
          test=None,
          **kwargs,
      ):
          super().__init__()
  
          self.batch_size = batch_size
          self.num_workers = num_workers
  
          self.dataset_configs = dict()
          if train is not None:
              self.dataset_configs["train"] = train
          if validation is not None:
              self.dataset_configs["validation"] = validation
          if test is not None:
              self.dataset_configs["test"] = test
  
      def setup(self, stage):
          from src.utils.train_util import instantiate_from_config
  
          if stage in ["fit"]:
              dataset_dict = {}
              for k in self.dataset_configs:
                  dataset_dict[k] = []
                  for loader in self.dataset_configs[k]:
                      dataset_dict[k].append(instantiate_from_config(loader))
              self.datasets = dataset_dict
              print(self.datasets)
          else:
              raise NotImplementedError
  
      def train_dataloader(self):
          datasets = ConcatDataset(self.datasets["train"])
          sampler = DistributedSampler(datasets)
          return DataLoader(
              datasets,
              batch_size=self.batch_size,
              num_workers=self.num_workers,
              shuffle=False,
              sampler=sampler,
              prefetch_factor=2,
              pin_memory=True,
          )
  
      def val_dataloader(self):
          datasets = ConcatDataset(self.datasets["validation"])
          sampler = DistributedSampler(datasets)
          return DataLoader(datasets, batch_size=4, num_workers=self.num_workers, shuffle=False, sampler=sampler)
  
      def test_dataloader(self):
          datasets = ConcatDataset(self.datasets["test"])
          return DataLoader(datasets, batch_size=self.batch_size, num_workers=self.num_workers, shuffle=False)
  