Adapter/T2I-Adapter

1# Modified from https://github.com/open-mmlab/mmcv/blob/master/mmcv/runner/dist_utils.py  # noqa: E5012import functools3import os4import subprocess5import torch6import torch.distributed as dist7import torch.multiprocessing as mp8from torch.nn.parallel import DataParallel, DistributedDataParallel9 10 11def init_dist(launcher, backend='nccl', **kwargs):12    if mp.get_start_method(allow_none=True) is None:13        mp.set_start_method('spawn')14    if launcher == 'pytorch':15        _init_dist_pytorch(backend, **kwargs)16    elif launcher == 'slurm':17        _init_dist_slurm(backend, **kwargs)18    else:19        raise ValueError(f'Invalid launcher type: {launcher}')20 21 22def _init_dist_pytorch(backend, **kwargs):23    rank = int(os.environ['RANK'])24    num_gpus = torch.cuda.device_count()25    torch.cuda.set_device(rank % num_gpus)26    dist.init_process_group(backend=backend, **kwargs)27 28 29def _init_dist_slurm(backend, port=None):30    """Initialize slurm distributed training environment.31 32    If argument ``port`` is not specified, then the master port will be system33    environment variable ``MASTER_PORT``. If ``MASTER_PORT`` is not in system34    environment variable, then a default port ``29500`` will be used.35 36    Args:37        backend (str): Backend of torch.distributed.38        port (int, optional): Master port. Defaults to None.39    """40    proc_id = int(os.environ['SLURM_PROCID'])41    ntasks = int(os.environ['SLURM_NTASKS'])42    node_list = os.environ['SLURM_NODELIST']43    num_gpus = torch.cuda.device_count()44    torch.cuda.set_device(proc_id % num_gpus)45    addr = subprocess.getoutput(f'scontrol show hostname {node_list} | head -n1')46    # specify master port47    if port is not None:48        os.environ['MASTER_PORT'] = str(port)49    elif 'MASTER_PORT' in os.environ:50        pass  # use MASTER_PORT in the environment variable51    else:52        # 29500 is torch.distributed default port53        os.environ['MASTER_PORT'] = '29500'54    os.environ['MASTER_ADDR'] = addr55    os.environ['WORLD_SIZE'] = str(ntasks)56    os.environ['LOCAL_RANK'] = str(proc_id % num_gpus)57    os.environ['RANK'] = str(proc_id)58    dist.init_process_group(backend=backend)59 60 61def get_dist_info():62    if dist.is_available():63        initialized = dist.is_initialized()64    else:65        initialized = False66    if initialized:67        rank = dist.get_rank()68        world_size = dist.get_world_size()69    else:70        rank = 071        world_size = 172    return rank, world_size73 74 75def master_only(func):76 77    @functools.wraps(func)78    def wrapper(*args, **kwargs):79        rank, _ = get_dist_info()80        if rank == 0:81            return func(*args, **kwargs)82 83    return wrapper84 85def get_bare_model(net):86    """Get bare model, especially under wrapping with87    DistributedDataParallel or DataParallel.88    """89    if isinstance(net, (DataParallel, DistributedDataParallel)):90        net = net.module91    return net92