models/datasets/precompute_tokens.py

import os
import shutil
import signal
import sys
import time
from contextlib import ExitStack
from functools import partial
from pathlib import Path

from accelerate.utils import gather_object, gather
from torchinfo import summary

from unidisc.tokenizers.chameleon_tokenizers import tokenize_chameleon, tokenize_chameleon_fast, get_chameleon_images, decode_ids, decode_ids_batched, tokenize_chameleon_mmc4, tokenize_regular_cambrian_mmc4
from utils import _print_config, set_numa_affinity, set_omega_conf_resolvers

sys.path.append(str(Path(__file__).parent.parent.parent / "unidisc/misc/hydra_submitit_launcher"))

import json
import os
import random
import sys
from contextlib import nullcontext
from pathlib import Path

import fsspec
import hydra
import numpy as np
import omegaconf
import rich.syntax
import rich.tree
import torch
from accelerate import Accelerator
from PIL import Image
from tensordict import TensorDict
from tqdm import tqdm
try:
    from viztracer import VizTracer
except ImportError:
    print("VizTracer not installed, skipping")

from dataloader import get_dataloaders, get_tokenizer, tokenize_text
from decoupled_utils import (barrier, breakpoint_on_error, get_local_rank, get_rank, get_world_size,
                             is_local_main_process, is_main_process,
                             rank_zero_fn, rprint, set_global_breakpoint,
                             set_global_exists, gprint)
from model import decode_latents, get_image_batch, get_vae
from models.datasets.combine_token_dicts import main as combine_token_dicts
from models.datasets.vggface_v2_attributes import (get_inference_func,
                                                   get_output)
from utils import (_print_config, set_numa_affinity, set_omega_conf_resolvers,
                   set_torch_defaults)
from omegaconf import DictConfig, OmegaConf, open_dict, read_write

os.environ["HYDRA_FULL_ERROR"] = "1"

set_global_breakpoint()  # Overrides breakpoint() to use ipdb.set_trace() instead and handle distributed training
set_global_exists()
set_omega_conf_resolvers()
set_torch_defaults()

def get_batch_size(config):
    with open_dict(config):
        if any(x.lower() in torch.cuda.get_device_name().lower() for x in ["v100", "1080", "2080", "quadro", "titan"]) or torch.cuda.get_device_capability()[0] <= 7:
            config.trainer.precision = "no"
            config.model.force_optimized_native_attn = False
            config.trainer.compile = False
            config.loader.batch_size = config.loader.batch_size // 3
            print(f"Found {torch.cuda.get_device_name().lower()}, set batch size to {config.loader.batch_size}")
    return config

def enc(data, idx, encode_images, config, vae, batch, accelerator, mixed_precision, tokenizer, vgg_data, existing_ids=None, device=None, mapping=None):
    
    if isinstance(batch, list):
        bs = len(batch)
    elif "img" in batch:
        bs = batch["img"].shape[0]
    else:
        bs = batch["attention_mask"].shape[0]
    
    sl = slice(idx * bs, (idx + 1) * bs)
    if not isinstance(batch, list) and "idx" in batch:
        if set(data[sl]["idx"].flatten().tolist()) == set(batch["idx"].tolist()):
            rprint(f"Skipping {idx} as all samples have already been processed 1")
            return
        if existing_ids is not None:
            set_inter = set(batch["idx"].tolist()) & existing_ids
            if len(set_inter) == bs:
                rprint(f"Skipping {idx} as all samples have already been processed 2")
                return
            elif len(set_inter) > 0:
                rprint(f"Running {idx} as some samples have already been processed: {len(set_inter)}")
    else:
        if (data[sl]["idx"] != -1).all():
            rprint(f"Skipping {idx} as all samples have already been processed")
            return
    
    if not isinstance(batch, list) and "img" in batch:
        batch["img"] = batch["img"].to(device=device, dtype=torch.bfloat16 if mixed_precision else None)

    with torch.no_grad():
        with torch.autocast(device_type="cuda", dtype=torch.bfloat16, enabled=mixed_precision):
            use_chameleon = getattr(config.data, "use_chameleon", False)
            use_mmc4 = config.data.train == "mmc4"
            use_cambrian = config.data.train == "cambrian"
            if not use_chameleon and not use_mmc4 and not use_cambrian:
                if tokenizer is not None and getattr(config.model, "unified_model", False):
                    if "input_ids" in batch and "attention_mask" in batch:
                        tokens = batch
                    else:
                        tokens = tokenize_text(tokenizer, config.data.block_size, batch[".txt"])

                    batch["txt_input_ids"] = tokens["input_ids"]
                    batch["txt_attention_mask"] = tokens["attention_mask"].float()
                elif getattr(config.data, "add_vggface_v2_attributes", False) and "vggface" not in config.data.train:
                    txt_input_ids, txt_attention_mask = get_output(batch, **vgg_data)
                    batch["txt_input_ids"] = txt_input_ids
                    batch["txt_attention_mask"] = txt_attention_mask
                elif getattr(config.data, "txt_only", False):
                    batch["txt_input_ids"] = batch["input_ids"]
                    batch["txt_attention_mask"] = batch["attention_mask"]

            if getattr(config.model, "unified_model", False) is False:
                if getattr(config.data, "txt_only", False):
                    batch["modality"] = torch.full((bs, 1), fill_value=0, dtype=torch.int16)
                else:
                    batch["modality"] = torch.full((bs, 1), fill_value=1, dtype=torch.int16)

            if isinstance(batch, list) and batch[0].get("idx", None) is not None:
                _idx = torch.tensor([x["idx"] for x in batch], dtype=torch.int32).unsqueeze(-1)
            elif "idx" in batch:
                _idx = batch["idx"].to(torch.int32).unsqueeze(-1)
            else:
                _idx = torch.full((bs, 1), fill_value=0, dtype=torch.int32)

            if "is_valid" in batch:
                _idx[~batch["is_valid"]] = -1
                if (_idx == -1).all():
                    gprint(f"WARNING: All samples are invalid")

            sl = slice(idx * bs, (idx + 1) * bs)
            assert (idx + 1) * bs <= len(data), f"Index {idx} + batch size {bs} is greater than the data length {len(data)}"

            if encode_images:
                if use_chameleon:
                    if isinstance(batch, list):
                        all_input_ids, all_attention_masks = tokenize_chameleon_mmc4(config, tokenizer, vae, batch, device, mapping)
                    else:
                        all_input_ids, all_attention_masks = tokenize_chameleon_fast(config, tokenizer, vae, batch)

                    # all_input_ids_gt, all_attention_masks_gt = tokenize_chameleon(config, tokenizer, vae, batch)
                    # txt_tokens, img_tokens = decode_ids_batched(_vae, all_input_ids[:4], return_tokens=True)
                    # img = decode_latents(config, _vae, img_tokens)
                    # from image_utils import Im;  Im(img).save()

                elif use_mmc4 or use_cambrian:
                    all_input_ids, all_attention_masks, all_modality = tokenize_regular_cambrian_mmc4(config, tokenizer, vae, batch, device, mapping)
                    if all_input_ids is None:
                        return
                else:
                    image_ids = get_image_batch(config, vae, batch, device)

                if use_chameleon or use_mmc4 or use_cambrian:
                    if not use_chameleon:
                        assert (all_input_ids < torch.iinfo(torch.int16).max).all()
                    
                    _kwargs = {}
                    if use_mmc4 or use_cambrian:
                        _kwargs["modality"] = all_modality.to(torch.int8)

                    data[sl] = TensorDict(
                        {
                            "input_ids": all_input_ids.to(torch.int32 if use_chameleon else torch.int16),
                            "attention_mask": all_attention_masks.to(torch.bool),
                            "idx": _idx,
                            "write_flag": torch.ones((bs, 1), dtype=torch.bool),
                            **_kwargs,
                        },
                        batch_size=[bs],
                    )
                elif getattr(config.model, "cond_label", False):
                    data[sl] = TensorDict(
                        {
                            "img_input_ids": image_ids.to(torch.int16),
                            "img_label": batch["label"].to(torch.int32).unsqueeze(-1),
                            "idx": _idx,
                            "write_flag": torch.ones((bs, 1), dtype=torch.bool),
                        },
                        batch_size=[bs],
                    )
                elif getattr(config.model, "unified_model", False) or getattr(config.data, "add_vggface_v2_attributes", False):
                    data[sl] = TensorDict(
                        {
                            "img_input_ids": image_ids.to(torch.int16),
                            "txt_input_ids": (batch.get("txt_input_ids") if batch.get("txt_input_ids") is not None else batch["input_ids"]).to(
                                torch.int32
                            ),
                            "txt_attention_mask": (
                                batch.get("txt_attention_mask") if batch.get("txt_attention_mask") is not None else batch["attention_mask"]
                            ).to(torch.bool),
                            "idx": _idx,
                            "write_flag": torch.ones((bs, 1), dtype=torch.bool),
                        },
                        batch_size=[bs],
                    )
                else:
                    data[sl] = TensorDict(
                        {"input_ids": image_ids.to(torch.int32), "attention_mask": torch.ones((image_ids.shape[0], image_ids.shape[1]), dtype=torch.bool), "idx": _idx, "write_flag": torch.ones((bs, 1), dtype=torch.bool), "modality": batch["modality"].to(torch.int16)},
                        batch_size=[bs],
                    )

            elif getattr(config.data, "txt_only", False):
                data[sl] = TensorDict(
                    {"input_ids": batch['input_ids'].to(torch.int32), "attention_mask": batch['attention_mask'].to(torch.bool), "idx": _idx, "write_flag": torch.ones((bs, 1), dtype=torch.bool), "modality": batch["modality"].to(torch.int16)},
                    batch_size=[bs],
                )
            else:
                real_image = batch["img"]
                if (config.data.resolution == 512 and batch["img"].shape[0] > 16) or (config.model.downscale_ratio <= 8):
                    chunk_size = 8 if (config.model.image_vocab_size > 64000 or config.model.downscale_ratio <= 8) else 16
                    chunks = [batch["img"][i : i + chunk_size] for i in range(0, batch["img"].shape[0], chunk_size)]
                    rec_img_list = []
                    for chunk in chunks:
                        batch_chunk = {"img": chunk}
                        image_ids = get_image_batch(config, vae, batch_chunk, device)
                        rec_img = decode_latents(config, vae, image_ids)
                        rec_img_list.append(rec_img)
                    rec_img = torch.cat(rec_img_list, dim=0)
                else:
                    image_ids = get_image_batch(config, vae, batch, device)
                    rec_img = decode_latents(config, vae, image_ids)

                viz_img = torch.cat([real_image, rec_img], dim=-1)
                from image_utils import Im
                
                if getattr(config.model, 'custom_vae_name', None) is not None:
                    custom_str = getattr(config.model, 'custom_vae_name')
                else:
                    custom_str = f"{'_custom' if getattr(config.model, 'use_custom_vae_ckpt', False) else ''}"
                (Path(__file__).parent.parent.parent / "output").mkdir(parents=True, exist_ok=True)
                Im(viz_img).save(
                    Path(__file__).parent.parent.parent / f"output/{config.data.train.replace('/', '')}_seq{image_ids.shape[1]}_res{config.data.resolution}_{config.model.vae_type}{custom_str}_voc{config.model.image_vocab_size}.png"
                )
                
                # Create directories for saving images
                dataset_name = config.data.train.replace('/', '')
                vae_name = f"seq{image_ids.shape[1]}_res{config.data.resolution}_{config.model.vae_type}{custom_str}_voc{config.model.image_vocab_size}"
                output_dir = Path(__file__).parent.parent.parent / "output" / dataset_name / vae_name
                gt_output_dir = Path(__file__).parent.parent.parent / "output" / dataset_name / f"GT_{config.data.resolution}"
                output_dir.mkdir(parents=True, exist_ok=True)
                gt_output_dir.mkdir(parents=True, exist_ok=True)

                # Save each image separately
                for i, (real, rec) in enumerate(zip(real_image, rec_img)):
                    print(Im(rec).save(output_dir / f"{i}.png"))
                    if (gt_output_dir / f"{i}.png").exists() is False:
                        print(Im(real).save(gt_output_dir / f"{i}.png"))

                gprint(f"Exiting")
                exit()


def get_dict(config, dataset_size):
    if getattr(config.data, "use_chameleon", False) or config.data.train == "cambrian" or config.data.train == "mmc4":
        input_ids_dtype = torch.int32 if getattr(config.data, "use_chameleon", False) else torch.int16
        data = TensorDict(
            {
                "input_ids": torch.zeros(dataset_size, config.model.length, dtype=input_ids_dtype),
                "attention_mask": torch.zeros(dataset_size, config.model.length, dtype=torch.bool),
                "modality": torch.full((dataset_size, config.model.length), fill_value=-1, dtype=torch.int8),
                "idx": torch.full((dataset_size, 1), fill_value=-1, dtype=torch.int32),
                "write_flag": torch.zeros(dataset_size, 1, dtype=torch.bool),
            },
            batch_size=[dataset_size],
        )
    elif getattr(config.model, "cond_label", False):
        data = TensorDict(
            {
                "img_input_ids": torch.zeros(dataset_size, config.model.img_length, dtype=torch.int16),
                "img_label": torch.zeros(dataset_size, 1, dtype=torch.int32),
                "idx": torch.full((dataset_size,), fill_value=-1, dtype=torch.int32),
                "write_flag": torch.zeros(dataset_size, 1, dtype=torch.bool),
            },
            batch_size=[dataset_size],
        )
    elif getattr(config.model, "unified_model", False) or getattr(config.data, "add_vggface_v2_attributes", False):
        data = TensorDict(
            {
                "img_input_ids": torch.zeros(dataset_size, config.model.img_length, dtype=torch.int16),
                "txt_input_ids": torch.zeros(dataset_size, config.model.txt_length, dtype=torch.int32),
                "txt_attention_mask": torch.zeros(dataset_size, config.model.txt_length, dtype=torch.bool),
                "idx": torch.full((dataset_size, 1), fill_value=-1, dtype=torch.int32),
                "write_flag": torch.zeros(dataset_size, 1, dtype=torch.bool),
            },
            batch_size=[dataset_size],
        )
    else:
        data = TensorDict(
            {
                "input_ids": torch.zeros(dataset_size, config.model.txt_length if config.data.txt_only else config.model.img_length, dtype=torch.int16),
                "idx": torch.full((dataset_size, 1), fill_value=-1, dtype=torch.int32),
                "write_flag": torch.zeros(dataset_size, 1, dtype=torch.bool),
                "modality": torch.full((dataset_size, 1), fill_value=-1, dtype=torch.int16),
            },
            batch_size=[dataset_size],
        )
    return data

def signal_handler(signum, frame, train_data, tmp_path):
    """Handle signals to save temporary train data."""
    rprint(f"Received signal {signum}, saving temporary train data.")
    print(f"[PRINT] Received signal {signum}, saving temporary train data.")
    save_tmp_data(train_data, tmp_path)
    sys.exit

def save_tmp_data(data, tmp_path):
    """Save data to a temporary path."""
    if tmp_path.exists() and tmp_path.is_dir():
        rprint(f"Deleting {tmp_path}")
        shutil.rmtree(tmp_path)  # Delete old tmp directory if it exists
    rprint(f"Saving tmp data to {tmp_path}")
    data.memmap(tmp_path, copy_existing=True)

def periodic_save(data, tmp_path, start_time, interval=2 * 60 * 60):
    """Periodically save data to a temporary path."""
    current_time = time.time()
    if current_time - start_time >= interval:
        rprint(f"Hit periodic save interval, saving tmp data to {tmp_path}")
        save_tmp_data(data, tmp_path)
        return current_time  # Reset start time
    return start_time

@hydra.main(version_base=None, config_path="../../configs", config_name="config")
def main(config):
    """Main entry point for training."""
    
    try:
        import resource
        soft_limit, hard_limit = resource.getrlimit(resource.RLIMIT_NOFILE)
        resource.setrlimit(resource.RLIMIT_NOFILE, (hard_limit, hard_limit)) # Set the soft limit to the hard limit
        rprint(f"Successfully set RLIMIT_NOFILE to {hard_limit}")
    except Exception as e:
        rprint(f"Failed to set RLIMIT_NOFILE: {e}")

    mixed_precision = False
    train_start_time = time.time()

    from datetime import timedelta
    from accelerate import Accelerator, DataLoaderConfiguration
    from accelerate.utils import InitProcessGroupKwargs
    kwargs = InitProcessGroupKwargs(timeout=timedelta(seconds=3600))
    prepare_kwargs = {}
    if config.data.train == "mmc4":
        prepare_kwargs["dispatch_batches"] = False

    accelerator = Accelerator(mixed_precision="bf16" if mixed_precision else None, kwargs_handlers=[kwargs], dataloader_config=DataLoaderConfiguration(**prepare_kwargs))
    device = torch.device(f"cuda:{accelerator.local_process_index}")

    import socket
    hostname = socket.gethostname()
    print(f"Hostname: {hostname}, Process index: {accelerator.process_index}, {device}, local_process_index: {accelerator.local_process_index}, get_local_process_index: {get_local_rank()}, device: {device}")
    _print_config(config, resolve=True, save_cfg=True)

    config = get_batch_size(config)

    # with omegaconf.open_dict(config):
    #     batch_sizes = gather_object([config.loader.batch_size])
    #     rprint(f"Batch sizes: {batch_sizes}")
    #     smallest_batch_size = min(batch_sizes)
    #     config.loader.batch_size = smallest_batch_size
    #     rprint(f"New config batch size: {config.loader.batch_size}")

    prefix = f"[Rank {accelerator.process_index}/{accelerator.num_processes}, Node: {os.environ.get('SLURM_NODEID', 'N/A')}, Hostname: {os.environ.get('SLURM_JOB_NODELIST', 'N/A')}, {config.data.train}]"
    print(f"{prefix} Starting precomputing tokens")
    save_validation_dataloader = getattr(config.data, "save_validation_dataloader", False)
    save_train_dataloader = getattr(config.data, "save_train_dataloader", False)

    tokenizer = get_tokenizer(config)
    train_dataloader, val_dataloader = get_dataloaders(
        config, tokenizer=tokenizer, allow_aug=False, force_aug=getattr(config.data, "force_aug", False), skip_valid=not save_validation_dataloader
    )

    train_dataloader = accelerator.prepare(train_dataloader)
    if save_validation_dataloader:
        val_dataloader = accelerator.prepare(val_dataloader)
    encode_images = getattr(config.model, "encode_images", False)

    use_chameleon = getattr(config.data, "use_chameleon", False)
    use_mmc4 = config.data.train == "mmc4"
    use_cambrian = config.data.train == "cambrian"
    mapping = None

    if use_chameleon:
        from unidisc.tokenizers.chameleon_tokenizers import ItemProcessor
        vae = ItemProcessor(target_size=config.data.resolution)
    else:
        vae = get_vae(config, device)

    if use_mmc4:
        import pandas as pd
        mapping = pd.read_parquet(config.data.mmc4_mapping_parquet)
        # Keep tar_filepath if it exists, otherwise use shard_path or map img2dataset_shard_id
        if "tar_filepath" in mapping.columns:
            pass
        elif "shard_path" in mapping.columns:
            mapping = mapping.rename(columns={"shard_path": "tar_filepath"})
            mapping["tar_filepath"] = mapping["tar_filepath"].str.replace(".parquet", ".tar")
        else:
            tar_path = Path(config.data.mmc4_tar_path)
            mapping["tar_filepath"] = mapping["img2dataset_shard_id"].apply(lambda x: tar_path / f"{x}.tar")
        
        mapping = mapping[['url', 'tar_filepath', 'key']]
        mapping = mapping.set_index("url").sort_index()

    if use_mmc4 or use_cambrian:
        assert config.data.use_slow_tokenizer and config.data.add_image_token

    if config.data.iterable:
        train_dataset_size = getattr(config.data, "train_dataset_size", None)
    else:
        print(f"{prefix} Train dataloader: {len(train_dataloader)} batches")
        print(f"{prefix} Train underlying dataset: {len(train_dataloader.dataset)} samples")
        train_dataset_size = (len(train_dataloader.dataset) // accelerator.num_processes) + config.loader.batch_size
        if save_validation_dataloader:
            print(f"{prefix} Val dataloader: {len(val_dataloader)} batches")
            print(f"Val underlying dataset: {len(val_dataloader.dataset)} samples")
            val_dataset_size = (len(val_dataloader.dataset) // accelerator.num_processes) + config.loader.batch_size

    print(f"{prefix} Train dataset size: {train_dataset_size} for 1 GPU")
    if save_validation_dataloader:
        print(f"{prefix} Val dataset size: {val_dataset_size} for 1 GPU")

    rank = accelerator.process_index
    output_dir = config.data.token_output_dir
    output_dir = Path(f"{output_dir}")
    output_dir.mkdir(parents=True, exist_ok=True)
    assert config.data.force_disable_shuffle

    debug = getattr(config.data, "debug", False)
    print(f"{prefix} Output dir: {output_dir}")

    vgg_data = None
    if getattr(config.data, "add_vggface_v2_attributes", False):
        print(f"{prefix} Adding VGGFace V2 attributes")
        vgg_data = get_inference_func()
        vgg_data["model"] = accelerator.prepare(vgg_data["model"])
    
    if not config.data.split_dataset and is_main_process() and any(output_dir.iterdir()):
        rprint(f"Found temporary directories in output dir, combining them")
        combine_token_dicts(output_dir, use_tmp=False, use_timestamp=True, delete_after_combining=True)
        for item in output_dir.iterdir():
            if item.is_dir() and "tmp" in item.name:
                rprint(f"Removing temporary directory: {item}")
                shutil.rmtree(item)

    # barrier() # TODO: Should be a barrier here
    if not config.data.split_dataset:
        existing_folders = sorted([folder for folder in output_dir.iterdir() if folder.is_dir() and "existing" in folder.name])
        if existing_folders:
            rprint(f"Found existing folders: {existing_folders}")
            existing_data = torch.cat([TensorDict.load_memmap(folder) for folder in existing_folders], dim=0)
            rprint(f"Concatenated existing data with shape: {existing_data.shape}")
            existing_ids = set(existing_data["idx"].to(torch.int32).flatten().tolist())
        else:
            rprint("No existing folders found")
            existing_ids = None
    else:
        existing_ids = None
    
    if save_train_dataloader:
        if not config.data.split_dataset and getattr(config.data, "allow_load_from_tmp", True) and Path(output_dir / f"tmp_train_{rank}").exists():
            rprint("Found tmp_train_{rank} in output dir, loading from it")
            train_data = TensorDict.load_memmap(output_dir / f"tmp_train_{rank}")
            train_data = train_data.clone()
        else:
            train_data = get_dict(config, train_dataset_size)

        print(f"{prefix} Starting train dataloader")
        if config.data.split_dataset:
            rank = int(os.getenv("SLURM_ARRAY_TASK_ID"))
            print(f"Using task id: {rank}")

        split_path = output_dir / f"train_{rank}"
        tmp_train_path = output_dir / f"tmp_train_{rank}"

        signal.signal(signal.SIGUSR1, partial(signal_handler, train_data=train_data, tmp_path=tmp_train_path))
        signal.signal(signal.SIGUSR2, partial(signal_handler, train_data=train_data, tmp_path=tmp_train_path))

        try:
            signal.signal(signal.SIGKILL, partial(signal_handler, train_data=train_data, tmp_path=tmp_train_path))
        except:
            rprint(f"Failed to set SIGKILL handler")

        start_time = time.time()
        with VizTracer(output_file="optional.json", tracer_entries=5000000) if debug else nullcontext():
            for i, batch in tqdm(enumerate(train_dataloader), leave=False, disable=not is_local_main_process()):
                if i == 0 and "img" in batch:
                    print(f"Batch shape: {batch['img'].shape}")
                if debug and i >= 1:
                    break
                enc(train_data, i, encode_images, config, vae, batch, accelerator, mixed_precision, tokenizer, vgg_data=vgg_data, existing_ids=existing_ids, device=device, mapping=mapping)
                try:
                    if not config.data.split_dataset or True:
                        start_time = periodic_save(train_data, tmp_train_path, start_time, getattr(config.data, "periodic_save", 2 * 60 * 60))
                except Exception as e:
                    gprint(f"Failed to save train data: {e}")
                    start_time = time.time()

        if debug:
            exit()

        del train_dataloader
        print(f"{prefix} Saving train data")
        if split_path.exists() and split_path.is_dir():
            rprint(f"Removing {split_path}")
            shutil.rmtree(split_path)

        split_path.mkdir(parents=True, exist_ok=True)
        gprint(f"Saving train data to {split_path}: {train_data.shape}")
        train_data.memmap(split_path, copy_existing=True)

        if tmp_train_path.exists() and tmp_train_path.is_dir():
            rprint(f"Removing {tmp_train_path}")
            shutil.rmtree(tmp_train_path)

        if not config.data.split_dataset:
            with open(output_dir / f"train_{rank}.completed", 'w') as f:
                f.write(f"Processing done for rank {rank}\n")

        print(f"{prefix} Finished train dataloader")

    if save_validation_dataloader:
        val_data = get_dict(config, val_dataset_size)
        split_path = output_dir / f"val_{rank}"
        split_path.mkdir(parents=True, exist_ok=True)
        tmp_val_path = output_dir / f"tmp_val_{rank}"
        print(f"Starting val dataloader")
        start_time = time.time()  # Track start time for periodic saving
        for i, batch in tqdm(enumerate(val_dataloader), leave=False):
            if debug and i >= 10:
                break
            enc(val_data, i, encode_images, config, vae, batch, accelerator, mixed_precision, tokenizer, vgg_data=vgg_data, device=device)
            
            # Periodically save data
            start_time = periodic_save(val_data, tmp_val_path, start_time)

        print(f"{prefix} Saving val data")
        if split_path.exists() and split_path.is_dir():
            rprint(f"Removing {split_path}")
            shutil.rmtree(split_path)
        split_path.mkdir(parents=True, exist_ok=True)
        rprint(f"Saving val data to {split_path}")
        val_data.memmap(split_path, copy_existing=True)
        if tmp_val_path.exists() and tmp_val_path.is_dir():
            shutil.rmtree(tmp_val_path)  # Delete tmp directory after final save
        print(f"{prefix} Finished val dataloader")

    rprint(f"{prefix} Finished precomputing tokens")

    if config.data.split_dataset:
        rprint(f"We are splitting the dataset and thus exiting.")
        exit()

    if get_world_size() > 1 and (time.time() - train_start_time) > 60 * 60:
        time.sleep(60 * 60)
        barrier()

    rprint('after barrier')
    if is_main_process():
        combine_token_dicts(data_dir=output_dir, allow_zero_idx=True, move_files=True, delete_after_combining=True)

    barrier()
    rprint(f"Finished concating tokens")


if __name__ == "__main__":
    with breakpoint_on_error():
        main()