content: # Copyright 2025 Bytedance Ltd. and/or its affiliates.
  # SPDX-License-Identifier: Apache-2.0
  
  
  import random
  import json
  
  import numpy as np
  import torch
  
  from .data_utils import (
      get_flattened_position_ids_interpolate,
      get_flattened_position_ids_extrapolate, 
      len2weight,
      patchify, 
      prepare_attention_mask_per_sample, 
  )
  from .dataset_info import DATASET_INFO, DATASET_REGISTRY
  from .transforms import ImageTransform
  from .video_utils import FrameSampler
  
  
  class DataConfig:
      def __init__(
          self, 
          grouped_datasets, 
          text_cond_dropout_prob=0.1,
          vit_cond_dropout_prob=0.4,
          vae_cond_dropout_prob=0.1,
          vae_image_downsample=16,
          max_latent_size=32,
          vit_patch_size=14,
          max_num_patch_per_side=70,
      ):
          self.grouped_datasets = grouped_datasets
          self.text_cond_dropout_prob = text_cond_dropout_prob
          self.vit_cond_dropout_prob = vit_cond_dropout_prob
          self.vit_patch_size = vit_patch_size
          self.max_num_patch_per_side = max_num_patch_per_side
          self.vae_cond_dropout_prob = vae_cond_dropout_prob
          self.vae_image_downsample = vae_image_downsample
          self.max_latent_size = max_latent_size
  
  
  class PackedDataset(torch.utils.data.IterableDataset):
      def __init__(
          self, 
          data_config, 
          tokenizer, 
          special_tokens,
          local_rank, 
          world_size, 
          num_workers,
          expected_num_tokens=32768, 
          max_num_tokens_per_sample=16384,
          max_num_tokens=36864,
          prefer_buffer_before=16384,
          max_buffer_size=50,
          interpolate_pos=False,
          use_flex=False,
          data_status=None,
      ):
          super().__init__()
          self.expected_num_tokens = expected_num_tokens
          self.max_num_tokens_per_sample = max_num_tokens_per_sample
          self.prefer_buffer_before = prefer_buffer_before
          self.max_num_tokens = max_num_tokens
          self.max_buffer_size = max_buffer_size
          self.tokenizer = tokenizer
          self.local_rank = local_rank
          self.world_size = world_size
          self.num_workers = num_workers
          self.use_flex = use_flex
          for k, v in special_tokens.items():
              setattr(self, k, v)
  
          grouped_datasets, is_mandatory, grouped_weights = self.build_datasets(
              data_config.grouped_datasets, data_status
          )
          self.grouped_datasets = grouped_datasets
          self.dataset_iters = [iter(dataset) for dataset in grouped_datasets]
          self.is_mandatory = is_mandatory
          self.grouped_weights = grouped_weights
          self.data_config = data_config
          self.interpolate_pos = interpolate_pos
          if self.interpolate_pos:
              self.get_flattened_position_ids = get_flattened_position_ids_interpolate
          else:
              self.get_flattened_position_ids = get_flattened_position_ids_extrapolate
  
      def build_datasets(self, datasets_metainfo, data_status):
          datasets = []
          is_mandatory = []
          grouped_weights = []
          for grouped_dataset_name, dataset_args in datasets_metainfo.items():
              is_mandatory.append(dataset_args.pop('is_mandatory', False))
              grouped_weights.append(dataset_args.pop('weight', 0.0))
  
              if 'frame_sampler_args' in dataset_args.keys():
                  frame_sampler = FrameSampler(**dataset_args.pop('frame_sampler_args'))
                  dataset_args['frame_sampler'] = frame_sampler
              if 'image_transform_args' in dataset_args.keys():
                  transform = ImageTransform(**dataset_args.pop('image_transform_args'))
                  dataset_args['transform'] = transform
              if 'vit_image_transform_args' in dataset_args.keys():
                  vit_transform = ImageTransform(**dataset_args.pop('vit_image_transform_args'))
                  dataset_args['vit_transform'] = vit_transform
  
              assert 'dataset_names' in dataset_args.keys()
              dataset_names = dataset_args.pop('dataset_names')
              dataset_args['data_dir_list'] = []
              for item in dataset_names:
                  if self.local_rank == 0:
                      print(f'Preparing Dataset {grouped_dataset_name}/{item}')
                  meta_info = DATASET_INFO[grouped_dataset_name][item]
                  dataset_args['data_dir_list'].append(meta_info['data_dir'])
  
                  if "parquet_info_path" in meta_info.keys():
                      if 'parquet_info' not in dataset_args.keys():
                          dataset_args['parquet_info'] = {}
                      with open(meta_info['parquet_info_path'], 'r') as f:
                          parquet_info = json.load(f)
                      dataset_args['parquet_info'].update(parquet_info)
  
                  if 'json_dir' in meta_info.keys():
                      # parquet/tar with json
                      if 'json_dir_list' not in dataset_args.keys():
                          dataset_args['json_dir_list'] = [meta_info['json_dir']]
                      else:
                          dataset_args['json_dir_list'].append(meta_info['json_dir'])
  
                  if 'jsonl_path' in meta_info.keys():
                      # jsonl with jpeg
                      if 'jsonl_path_list' not in dataset_args.keys():
                          dataset_args['jsonl_path_list'] = [meta_info['jsonl_path']]
                      else:
                          dataset_args['jsonl_path_list'].append(meta_info['jsonl_path'])
  
              resume_data_status = dataset_args.pop('resume_data_status', True)
              if data_status is not None and grouped_dataset_name in data_status.keys() and resume_data_status:
                  data_status_per_group = data_status[grouped_dataset_name]
              else:
                  data_status_per_group = None
              dataset = DATASET_REGISTRY[grouped_dataset_name](
                  dataset_name=grouped_dataset_name,
                  tokenizer=self.tokenizer,
                  local_rank=self.local_rank,
                  world_size=self.world_size,
                  num_workers=self.num_workers,
                  data_status=data_status_per_group,
                  **dataset_args
              )
              datasets.append(dataset)
  
          return datasets, is_mandatory, grouped_weights
  
      def set_epoch(self, seed):
          for dataset in self.grouped_datasets:
              dataset.set_epoch(seed)
  
      def set_sequence_status(self):
          sequence_status = dict(
              curr                        = 0,
              sample_lens                 = list(),
              packed_position_ids         = list(),
              nested_attention_masks      = list(),
              split_lens                  = list(),
              attn_modes                  = list(),
              packed_text_ids             = list(), 
              packed_text_indexes         = list(),
              packed_label_ids            = list(),
              ce_loss_indexes             = list(),
              ce_loss_weights             = list(),
              vae_image_tensors           = list(), 
              packed_latent_position_ids  = list(),
              vae_latent_shapes           = list(), 
              packed_vae_token_indexes    = list(), 
              packed_timesteps            = list(), 
              mse_loss_indexes            = list(),
              packed_vit_tokens           = list(), 
              vit_token_seqlens           = list(),
              packed_vit_position_ids     = list(),
              packed_vit_token_indexes    = list(), 
          )
          return sequence_status
  
      def to_tensor(self, sequence_status):
          data = dict(
              sequence_length=sum(sequence_status['sample_lens']),
              sample_lens=sequence_status['sample_lens'],
              packed_text_ids=torch.tensor(sequence_status['packed_text_ids']),
              packed_text_indexes=torch.tensor(sequence_status['packed_text_indexes']),
              packed_position_ids=torch.tensor(sequence_status['packed_position_ids']),
          )
          if not self.use_flex:
              data['nested_attention_masks'] = sequence_status['nested_attention_masks']
          else:
              sequence_len = data['sequence_length']
              pad_len = self.max_num_tokens - sequence_len
              data['split_lens'] = sequence_status['split_lens'] + [pad_len]
              data['attn_modes'] = sequence_status['attn_modes'] + ['causal']
              data['sample_lens'] += [pad_len]
  
          # if the model has a convnet vae (e.g., as visual tokenizer)
          if len(sequence_status['vae_image_tensors']) > 0:
              image_tensors = sequence_status.pop('vae_image_tensors')
              image_sizes = [item.shape for item in image_tensors]
              max_image_size = [max(item) for item in list(zip(*image_sizes))]
              padded_images = torch.zeros(size=(len(image_tensors), *max_image_size))
              for i, image_tensor in enumerate(image_tensors):
                  padded_images[i, :, :image_tensor.shape[1], :image_tensor.shape[2]] = image_tensor
  
              data['padded_images'] = padded_images
              data['patchified_vae_latent_shapes'] = sequence_status['vae_latent_shapes']
              data['packed_latent_position_ids'] = torch.cat(sequence_status['packed_latent_position_ids'], dim=0)
              data['packed_vae_token_indexes'] = torch.tensor(sequence_status['packed_vae_token_indexes'])
  
          # if the model has a vit (e.g., as visual tokenizer)
          if len(sequence_status['packed_vit_tokens']) > 0:
              data['packed_vit_tokens'] = torch.cat(sequence_status['packed_vit_tokens'], dim=0)
              data['packed_vit_position_ids'] = torch.cat(sequence_status['packed_vit_position_ids'], dim=0)
              data['packed_vit_token_indexes'] = torch.tensor(sequence_status['packed_vit_token_indexes'])
              data['vit_token_seqlens'] = torch.tensor(sequence_status['vit_token_seqlens'])
  
          # if the model is required to perform visual generation
          if len(sequence_status['packed_timesteps']) > 0:
              data['packed_timesteps'] = torch.tensor(sequence_status['packed_timesteps'])
              data['mse_loss_indexes'] = torch.tensor(sequence_status['mse_loss_indexes'])
  
          # if the model is required to perform text generation
          if len(sequence_status['packed_label_ids']) > 0:
              data['packed_label_ids'] = torch.tensor(sequence_status['packed_label_ids'])
              data['ce_loss_indexes'] = torch.tensor(sequence_status['ce_loss_indexes'])
              data['ce_loss_weights'] = torch.tensor(sequence_status['ce_loss_weights'])
  
          return data
  
      def __iter__(self):
          total_weights = sum(self.grouped_weights)
          assert total_weights > 0.0
          group_cumprobs = [sum(self.grouped_weights[:i + 1]) / total_weights 
                            for i in range(len(self.grouped_weights))]
          sequence_status = self.set_sequence_status()
          batch_data_indexes = []
  
          buffer = []
          while True:
              # Ensure at least one sample from each group
              if sequence_status['curr'] == 0:
                  for group_index, group_iter in enumerate(self.dataset_iters):
                      if self.is_mandatory[group_index]:
                          while True:
                              sample = next(group_iter)
                              # if a sample is too long, skip it
                              num_tokens = sample['num_tokens'] + 2 * len(sample['sequence_plan'])
                              if num_tokens < self.max_num_tokens_per_sample:
                                  sequence_status = self.pack_sequence(sample, sequence_status)
                                  batch_data_indexes.append(sample['data_indexes'])
                                  break
                              else:
                                  print(f"skip a sample with length {num_tokens}")
                                  continue
  
              if sequence_status['curr'] < self.prefer_buffer_before and len(buffer) > 0:
                  sample = buffer.pop(0)
                  sample_from_buffer = True
              else:
                  # sample normally across all groups
                  n = random.random()
                  group_index = 0
                  for i, cumprob in enumerate(group_cumprobs):
                      if n < cumprob:
                          group_index = i
                          break
                  sample = next(self.dataset_iters[group_index])
                  sample_from_buffer = False
  
              # if a sample is too long, skip it
              num_tokens = sample['num_tokens'] + 2 * len(sample['sequence_plan'])
              if num_tokens > self.max_num_tokens_per_sample:
                  print(f"skip a sample with length {num_tokens}")
                  continue
  
              if sequence_status['curr'] + num_tokens > self.max_num_tokens:
                  if len(buffer) < self.max_buffer_size and not sample_from_buffer:
                      buffer.append(sample)
                  else:
                      print(f"Yielding data with length {sum(sequence_status['sample_lens'])}")
                      data = self.to_tensor(sequence_status)
                      data['batch_data_indexes'] = batch_data_indexes
                      yield data
                      sequence_status = self.set_sequence_status()
                      batch_data_indexes = []
                  continue
  
              sequence_status = self.pack_sequence(sample, sequence_status)
              batch_data_indexes.append(sample['data_indexes'])
  
              if sequence_status['curr'] >= self.expected_num_tokens:
                  data = self.to_tensor(sequence_status)
                  data['batch_data_indexes'] = batch_data_indexes
                  yield data
                  sequence_status = self.set_sequence_status()
                  batch_data_indexes = []
  
      def pack_sequence(self, sample, sequence_status):
          image_tensor_list = sample['image_tensor_list']
          text_ids_list = sample['text_ids_list']
          sequence_plan = sample['sequence_plan']
  
          split_lens, attn_modes = list(), list()
          curr = sequence_status['curr']
          curr_rope_id = 0
          sample_lens = 0
  
          for item in sequence_plan:
              split_start = item.get('split_start', True)
              if split_start:
                  curr_split_len = 0
  
              if item['type'] == 'text':
                  text_ids = text_ids_list.pop(0)
                  if item['enable_cfg'] == 1 and random.random() < self.data_config.text_cond_dropout_prob:
                      continue
  
                  shifted_text_ids = [self.bos_token_id] + text_ids
                  sequence_status['packed_text_ids'].extend(shifted_text_ids)
                  sequence_status['packed_text_indexes'].extend(range(curr, curr + len(shifted_text_ids)))
                  if item['loss'] == 1:
                      sequence_status['ce_loss_indexes'].extend(range(curr, curr + len(shifted_text_ids)))
                      sequence_status['ce_loss_weights'].extend(
                          [len2weight(len(shifted_text_ids))] * len(shifted_text_ids)
                      )
                      sequence_status['packed_label_ids'].extend(text_ids + [self.eos_token_id])
                  curr += len(shifted_text_ids)
                  curr_split_len += len(shifted_text_ids)
  
                  # add a <|im_end|> token
                  sequence_status['packed_text_ids'].append(self.eos_token_id)
                  sequence_status['packed_text_indexes'].append(curr)
                  if item['special_token_loss'] == 1: # <|im_end|> may have loss
                      sequence_status['ce_loss_indexes'].append(curr)
                      sequence_status['ce_loss_weights'].append(1.0)
                      sequence_status['packed_label_ids'].append(item['special_token_label'])
                  curr += 1
                  curr_split_len += 1
  
                  # update sequence status
                  attn_modes.append("causal")
                  sequence_status['packed_position_ids'].extend(range(curr_rope_id, curr_rope_id + curr_split_len))
                  curr_rope_id += curr_split_len
  
              elif item['type'] == 'vit_image':
                  image_tensor = image_tensor_list.pop(0)
                  if item['enable_cfg'] == 1 and random.random() < self.data_config.vit_cond_dropout_prob:
                      curr_rope_id += 1
                      continue
  
                  # add a <|startofimage|> token
                  sequence_status['packed_text_ids'].append(self.start_of_image)
                  sequence_status['packed_text_indexes'].append(curr)
                  curr += 1
                  curr_split_len += 1
  
                  # preprocess image
                  vit_tokens = patchify(image_tensor, self.data_config.vit_patch_size)
                  num_img_tokens = vit_tokens.shape[0]
                  sequence_status['packed_vit_token_indexes'].extend(range(curr, curr + num_img_tokens))
                  curr += num_img_tokens
                  curr_split_len += num_img_tokens
  
                  sequence_status['packed_vit_tokens'].append(vit_tokens)
                  sequence_status['vit_token_seqlens'].append(num_img_tokens)
                  sequence_status['packed_vit_position_ids'].append(
                      self.get_flattened_position_ids(
                          image_tensor.size(1), image_tensor.size(2),
                          self.data_config.vit_patch_size, 
                          max_num_patches_per_side=self.data_config.max_num_patch_per_side
                      )
                  )
  
                  # add a <|endofimage|> token
                  sequence_status['packed_text_ids'].append(self.end_of_image)
                  sequence_status['packed_text_indexes'].append(curr)
                  if item['special_token_loss'] == 1: # <|endofimage|> may have loss
                      sequence_status['ce_loss_indexes'].append(curr)
                      sequence_status['ce_loss_weights'].append(1.0)
                      sequence_status['packed_label_ids'].append(item['special_token_label'])
                  curr += 1
                  curr_split_len += 1
  
                  # update sequence status
                  attn_modes.append("full")
                  sequence_status['packed_position_ids'].extend([curr_rope_id] * curr_split_len)
                  curr_rope_id += 1
  
              elif item['type'] == 'vae_image':
                  image_tensor = image_tensor_list.pop(0)
                  if item['enable_cfg'] == 1 and random.random() < self.data_config.vae_cond_dropout_prob:
                      # FIXME fix vae dropout in video2video setting.
                      curr_rope_id += 1
                      continue
  
                  # add a <|startofimage|> token
                  sequence_status['packed_text_ids'].append(self.start_of_image)
                  sequence_status['packed_text_indexes'].append(curr)
                  curr += 1
                  curr_split_len += 1
  
                  # preprocess image
                  sequence_status['vae_image_tensors'].append(image_tensor)
                  sequence_status['packed_latent_position_ids'].append(
                      self.get_flattened_position_ids(
                          image_tensor.size(1), image_tensor.size(2),
                          self.data_config.vae_image_downsample, 
                          max_num_patches_per_side=self.data_config.max_latent_size
                      )
                  )
                  H, W = image_tensor.shape[1:]
                  h = H // self.data_config.vae_image_downsample
                  w = W // self.data_config.vae_image_downsample
                  sequence_status['vae_latent_shapes'].append((h, w))
  
                  num_img_tokens = w * h
                  sequence_status['packed_vae_token_indexes'].extend(range(curr, curr + num_img_tokens))
                  if item['loss'] == 1:
                      sequence_status['mse_loss_indexes'].extend(range(curr, curr + num_img_tokens))
                      if split_start:
                          timestep = np.random.randn()
                  else:
                      timestep = float('-inf')
  
                  sequence_status['packed_timesteps'].extend([timestep] * num_img_tokens)
                  curr += num_img_tokens
                  curr_split_len += num_img_tokens
  
                  # add a <|endofimage|> token
                  sequence_status['packed_text_ids'].append(self.end_of_image)
                  sequence_status['packed_text_indexes'].append(curr)
                  # <|endofimage|> may have loss
                  if item['special_token_loss'] == 1:
                      sequence_status['ce_loss_indexes'].append(curr)
                      sequence_status['ce_loss_weights'].append(1.0)
                      sequence_status['packed_label_ids'].append(item['special_token_label'])
                  curr += 1
                  curr_split_len += 1
  
                  # update sequence status
                  if split_start:
                      if item['loss'] == 1 and 'frame_delta' not in item.keys():
                          attn_modes.append("noise")
                      else:
                          attn_modes.append("full")
                  sequence_status['packed_position_ids'].extend([curr_rope_id] * (num_img_tokens + 2))
                  if 'frame_delta' in item.keys():
                      curr_rope_id += item['frame_delta']
                  elif item['loss'] == 0:
                      curr_rope_id += 1
  
              if item.get('split_end', True):
                  split_lens.append(curr_split_len)
                  sample_lens += curr_split_len
  
          sequence_status['curr'] = curr
          sequence_status['sample_lens'].append(sample_lens)
          # prepare attention mask
          if not self.use_flex:
              sequence_status['nested_attention_masks'].append(
                  prepare_attention_mask_per_sample(split_lens, attn_modes)
              )
          else:
              sequence_status['split_lens'].extend(split_lens)
              sequence_status['attn_modes'].extend(attn_modes)
  
          return sequence_status
  
  
  class SimpleCustomBatch:
      def __init__(self, batch):
          data = batch[0]
          self.batch_data_indexes = data['batch_data_indexes']
          self.sequence_length = data["sequence_length"]
          self.sample_lens = data["sample_lens"]
          self.packed_text_ids = data["packed_text_ids"]
          self.packed_text_indexes = data["packed_text_indexes"]
          self.packed_position_ids = data["packed_position_ids"]
  
          self.use_flex = "nested_attention_masks" not in data.keys()
  
          if self.use_flex:
              self.split_lens = data["split_lens"]
              self.attn_modes = data["attn_modes"]
          else:
              self.nested_attention_masks = data["nested_attention_masks"]
  
          if "padded_images" in data.keys():
              self.padded_images = data["padded_images"]
              self.patchified_vae_latent_shapes = data["patchified_vae_latent_shapes"]
              self.packed_latent_position_ids = data["packed_latent_position_ids"]
              self.packed_vae_token_indexes = data["packed_vae_token_indexes"]
  
          if "packed_vit_tokens" in data.keys():
              self.packed_vit_tokens = data["packed_vit_tokens"]
              self.packed_vit_position_ids = data["packed_vit_position_ids"]
              self.packed_vit_token_indexes = data["packed_vit_token_indexes"]
              self.vit_token_seqlens = data["vit_token_seqlens"]
  
          if "packed_timesteps" in data.keys():
              self.packed_timesteps = data["packed_timesteps"]
              self.mse_loss_indexes = data["mse_loss_indexes"]
  
          if "packed_label_ids" in data.keys():
              self.packed_label_ids = data["packed_label_ids"]
              self.ce_loss_indexes = data["ce_loss_indexes"]
              self.ce_loss_weights = data["ce_loss_weights"]
  
      def pin_memory(self):
          self.packed_text_ids = self.packed_text_ids.pin_memory()
          self.packed_text_indexes = self.packed_text_indexes.pin_memory()
          self.packed_position_ids = self.packed_position_ids.pin_memory()
  
          if not self.use_flex:
              self.nested_attention_masks = [item.pin_memory() for item in self.nested_attention_masks]
  
          if hasattr(self, 'padded_images'):
              self.padded_images = self.padded_images.pin_memory()
              self.packed_vae_token_indexes = self.packed_vae_token_indexes.pin_memory()
              self.packed_latent_position_ids = self.packed_latent_position_ids.pin_memory()
  
          if hasattr(self, 'packed_timesteps'):
              self.packed_timesteps = self.packed_timesteps.pin_memory()
              self.mse_loss_indexes = self.mse_loss_indexes.pin_memory()
  
          if hasattr(self, 'packed_vit_tokens'):
              self.packed_vit_tokens = self.packed_vit_tokens.pin_memory()
              self.packed_vit_position_ids = self.packed_vit_position_ids.pin_memory()
              self.packed_vit_token_indexes = self.packed_vit_token_indexes.pin_memory()
              self.vit_token_seqlens = self.vit_token_seqlens.pin_memory()
  
          if hasattr(self, 'packed_label_ids'):
              self.packed_label_ids = self.packed_label_ids.pin_memory()
              self.ce_loss_indexes = self.ce_loss_indexes.pin_memory()
              self.ce_loss_weights = self.ce_loss_weights.pin_memory()
  
          return self
  
      def cuda(self, device):
          self.packed_text_ids = self.packed_text_ids.to(device)
          self.packed_text_indexes = self.packed_text_indexes.to(device)
          self.packed_position_ids = self.packed_position_ids.to(device)
  
          if not self.use_flex:
              self.nested_attention_masks = [item.to(device) for item in self.nested_attention_masks]
  
          if hasattr(self, 'padded_images'):
              self.padded_images = self.padded_images.to(device)
              self.packed_vae_token_indexes = self.packed_vae_token_indexes.to(device)
              self.packed_latent_position_ids = self.packed_latent_position_ids.to(device)
  
          if hasattr(self, 'packed_timesteps'):
              self.packed_timesteps = self.packed_timesteps.to(device)
              self.mse_loss_indexes = self.mse_loss_indexes.to(device)
  
          if hasattr(self, 'packed_vit_tokens'):
              self.packed_vit_tokens = self.packed_vit_tokens.to(device)
              self.packed_vit_position_ids = self.packed_vit_position_ids.to(device)
              self.packed_vit_token_indexes = self.packed_vit_token_indexes.to(device)
              self.vit_token_seqlens = self.vit_token_seqlens.to(device)
  
          if hasattr(self, 'packed_label_ids'):
              self.packed_label_ids = self.packed_label_ids.to(device)
              self.ce_loss_indexes = self.ce_loss_indexes.to(device)
              self.ce_loss_weights = self.ce_loss_weights.to(device)
  
          return self
  
      def to_dict(self):
          data = dict(
              sequence_length = self.sequence_length,
              sample_lens = self.sample_lens,
              packed_text_ids = self.packed_text_ids,
              packed_text_indexes = self.packed_text_indexes,
              packed_position_ids = self.packed_position_ids,
              batch_data_indexes = self.batch_data_indexes,
          )
  
          if not self.use_flex:
              data['nested_attention_masks'] = self.nested_attention_masks
          else:
              data['split_lens'] = self.split_lens
              data['attn_modes'] = self.attn_modes
  
          if hasattr(self, 'padded_images'):
              data['padded_images'] = self.padded_images
              data['patchified_vae_latent_shapes'] = self.patchified_vae_latent_shapes
              data['packed_latent_position_ids'] = self.packed_latent_position_ids
              data['packed_vae_token_indexes'] = self.packed_vae_token_indexes
  
          if hasattr(self, 'packed_vit_tokens'):
              data['packed_vit_tokens'] = self.packed_vit_tokens
              data['packed_vit_position_ids'] = self.packed_vit_position_ids
              data['packed_vit_token_indexes'] = self.packed_vit_token_indexes
              data['vit_token_seqlens'] = self.vit_token_seqlens
  
          if hasattr(self, 'packed_timesteps'):
              data['packed_timesteps'] = self.packed_timesteps
              data['mse_loss_indexes'] = self.mse_loss_indexes
  
          if hasattr(self, 'packed_label_ids'):
              data['packed_label_ids'] = self.packed_label_ids
              data['ce_loss_indexes'] = self.ce_loss_indexes
              data['ce_loss_weights'] = self.ce_loss_weights
  
          return data
  
  
  def collate_wrapper():
      def collate_fn(batch):
          return SimpleCustomBatch(batch)
      return collate_fn
  