| """Modified from https://github.com/THUDM/CogVideo/blob/3710a612d8760f5cdb1741befeebb65b9e0f2fe0/sat/sgm/modules/diffusionmodules/sigma_sampling.py
|
| """
|
| import torch
|
|
|
| class DiscreteSampling:
|
| def __init__(self, num_idx, uniform_sampling=False, start_num_idx=0, sp_size=1):
|
| self.num_idx = num_idx
|
| self.start_num_idx = start_num_idx
|
| self.uniform_sampling = uniform_sampling
|
| self.is_distributed = torch.distributed.is_available() and torch.distributed.is_initialized()
|
|
|
| if self.is_distributed and self.uniform_sampling:
|
| world_size = torch.distributed.get_world_size()
|
| self.rank = torch.distributed.get_rank()
|
|
|
| i = 1
|
| while True:
|
| if world_size % i != 0 or num_idx % (world_size // i) != 0:
|
| i += 1
|
| else:
|
| if i >= sp_size:
|
| self.group_num = world_size // i
|
| elif sp_size > world_size:
|
| self.group_num = 1
|
| else:
|
| self.group_num = world_size // sp_size
|
| break
|
| assert self.group_num > 0
|
| assert world_size % self.group_num == 0
|
|
|
| self.group_width = world_size // self.group_num
|
| self.sigma_interval = self.num_idx // self.group_num
|
| print('rank=%d world_size=%d group_num=%d group_width=%d sigma_interval=%s' % (
|
| self.rank, world_size, self.group_num,
|
| self.group_width, self.sigma_interval))
|
|
|
| def __call__(self, n_samples, generator=None, device=None):
|
| if self.is_distributed and self.uniform_sampling:
|
| group_index = self.rank // self.group_width
|
| idx = torch.randint(
|
| self.start_num_idx + group_index * self.sigma_interval,
|
| self.start_num_idx + (group_index + 1) * self.sigma_interval,
|
| (n_samples,),
|
| generator=generator, device=device,
|
| )
|
| print('proc[%d] idx=%s' % (self.rank, idx))
|
| else:
|
| idx = torch.randint(
|
| self.start_num_idx, self.start_num_idx + self.num_idx, (n_samples,),
|
| generator=generator, device=device,
|
| )
|
| return idx |
