import datetime import argparse, importlib from pytorch_lightning import seed_everything import torch import torch.distributed as dist def setup_dist(local_rank): if dist.is_initialized(): return torch.cuda.set_device(local_rank) torch.distributed.init_process_group('nccl', init_method='env://') def get_dist_info(): if dist.is_available(): initialized = dist.is_initialized() else: initialized = False if initialized: rank = dist.get_rank() world_size = dist.get_world_size() else: rank = 0 world_size = 1 return rank, world_size if __name__ == '__main__': now = datetime.datetime.now().strftime("%Y-%m-%d-%H-%M-%S") parser = argparse.ArgumentParser() parser.add_argument("--module", type=str, help="module name", default="inference") parser.add_argument("--local_rank", type=int, nargs="?", help="for ddp", default=0) args, unknown = parser.parse_known_args() inference_api = importlib.import_module(args.module, package=None) inference_parser = inference_api.get_parser() inference_args, unknown = inference_parser.parse_known_args() seed_everything(inference_args.seed) setup_dist(args.local_rank) torch.backends.cudnn.benchmark = True rank, gpu_num = get_dist_info() print("@CoVideoGen Inference [rank%d]: %s"%(rank, now)) inference_api.run_inference(inference_args, rank)