Anonymous-123/ImageNet-Editing

1import numpy as np2import torch as th3 4from .gaussian_diffusion import GaussianDiffusion5 6 7def space_timesteps(num_timesteps, section_counts):8    """9    Create a list of timesteps to use from an original diffusion process,10    given the number of timesteps we want to take from equally-sized portions11    of the original process.12 13    For example, if there's 300 timesteps and the section counts are [10,15,20]14    then the first 100 timesteps are strided to be 10 timesteps, the second 10015    are strided to be 15 timesteps, and the final 100 are strided to be 20.16 17    If the stride is a string starting with "ddim", then the fixed striding18    from the DDIM paper is used, and only one section is allowed.19 20    :param num_timesteps: the number of diffusion steps in the original21                          process to divide up.22    :param section_counts: either a list of numbers, or a string containing23                           comma-separated numbers, indicating the step count24                           per section. As a special case, use "ddimN" where N25                           is a number of steps to use the striding from the26                           DDIM paper.27    :return: a set of diffusion steps from the original process to use.28    """29    if isinstance(section_counts, str):30        if section_counts.startswith("ddim"):31            desired_count = int(section_counts[len("ddim") :])32            for i in range(1, num_timesteps):33                if len(range(0, num_timesteps, i)) == desired_count:34                    return set(range(0, num_timesteps, i))35            raise ValueError(36                f"cannot create exactly {num_timesteps} steps with an integer stride"37            )38        section_counts = [int(x) for x in section_counts.split(",")]39    size_per = num_timesteps // len(section_counts)40    extra = num_timesteps % len(section_counts)41    start_idx = 042    all_steps = []43    for i, section_count in enumerate(section_counts):44        size = size_per + (1 if i < extra else 0)45        if size < section_count:46            raise ValueError(47                f"cannot divide section of {size} steps into {section_count}"48            )49        if section_count <= 1:50            frac_stride = 151        else:52            frac_stride = (size - 1) / (section_count - 1)53        cur_idx = 0.054        taken_steps = []55        for _ in range(section_count):56            taken_steps.append(start_idx + round(cur_idx))57            cur_idx += frac_stride58        all_steps += taken_steps59        start_idx += size60    return set(all_steps)61 62 63class SpacedDiffusion(GaussianDiffusion):64    """65    A diffusion process which can skip steps in a base diffusion process.66 67    :param use_timesteps: a collection (sequence or set) of timesteps from the68                          original diffusion process to retain.69    :param kwargs: the kwargs to create the base diffusion process.70    """71 72    def __init__(self, use_timesteps, **kwargs):73        self.use_timesteps = set(use_timesteps)74        self.timestep_map = []75        self.original_num_steps = len(kwargs["betas"])76 77        base_diffusion = GaussianDiffusion(**kwargs)  # pylint: disable=missing-kwoa78        last_alpha_cumprod = 1.079        new_betas = []80        for i, alpha_cumprod in enumerate(base_diffusion.alphas_cumprod):81            if i in self.use_timesteps:82                new_betas.append(1 - alpha_cumprod / last_alpha_cumprod)83                last_alpha_cumprod = alpha_cumprod84                self.timestep_map.append(i)85        kwargs["betas"] = np.array(new_betas)86        super().__init__(**kwargs)87 88    def p_mean_variance(89        self, model, *args, **kwargs90    ):  # pylint: disable=signature-differs91        return super().p_mean_variance(self._wrap_model(model), *args, **kwargs)92 93    def training_losses(94        self, model, *args, **kwargs95    ):  # pylint: disable=signature-differs96        return super().training_losses(self._wrap_model(model), *args, **kwargs)97 98    def condition_mean(self, cond_fn, *args, **kwargs):99        return super().condition_mean(self._wrap_model(cond_fn), *args, **kwargs)100 101    def condition_score(self, cond_fn, *args, **kwargs):102        return super().condition_score(self._wrap_model(cond_fn), *args, **kwargs)103 104    def _wrap_model(self, model):105        if isinstance(model, _WrappedModel):106            return model107        return _WrappedModel(108            model, self.timestep_map, self.rescale_timesteps, self.original_num_steps109        )110 111    def _scale_timesteps(self, t):112        # Scaling is done by the wrapped model.113        return t114 115 116class _WrappedModel:117    def __init__(self, model, timestep_map, rescale_timesteps, original_num_steps):118        self.model = model119        self.timestep_map = timestep_map120        self.rescale_timesteps = rescale_timesteps121        self.original_num_steps = original_num_steps122 123    def __call__(self, x, ts, **kwargs):124        map_tensor = th.tensor(self.timestep_map, device=ts.device, dtype=ts.dtype)125        new_ts = map_tensor[ts]126        if self.rescale_timesteps:127            new_ts = new_ts.float() * (1000.0 / self.original_num_steps)128        return self.model(x, new_ts, **kwargs)129