import gc

import json

import math

import os

import os.path as osp

import random

import sys

import time

from functools import partial

from typing import List, Optional, Tuple

os.environ["TOKENIZERS_PARALLELISM"] = "false"

os.environ['XFORMERS_FORCE_DISABLE_TRITON'] = '1'

import numpy as np

import torch

torch._dynamo.config.cache_size_limit = 64

from torch.nn import functional as F

from torch.profiler import record_function

from torch.utils.data import DataLoader

from transformers import T5EncoderModel, T5TokenizerFast

import torch.distributed as tdist

import grn.utils_t2iv.dist as dist

from grn.dataset.build import build_joint_dataset

from grn.models.ema import get_ema_model

from grn.utils_t2iv import arg_util, misc

from grn.utils import wandb_utils

from grn.trainer import get_trainer

def build_everything_from_args(args: arg_util.Args, saver):

return text_tokenizer, text_encoder, trainer, global_it

def build_model_optimizer(args):

return vae_local, gpt_uncompiled, gpt_wo_ddp, gpt_ddp, gpt_wo_ddp_ema, gpt_ddp_ema, gpt_optim

def build_dataset(args):

return train_dataset

def main_train(args: arg_util.Args):

return

def train_one_ep(

return {k: meter.global_avg for k, meter in me.meters.items()}, me.iter_time.time_preds(max_it - (g_it + 1) + (args.ep - ep) * 15) # +15: other cost

def main():

time.sleep(120)

if __name__ == '__main__':

main()