ALSv/self-forcing

1from abc import abstractmethod, ABC2import torch3 4 5class SchedulerInterface(ABC):6    """7    Base class for diffusion noise schedule.8    """9    alphas_cumprod: torch.Tensor  # [T], alphas for defining the noise schedule10 11    @abstractmethod12    def add_noise(13        self, clean_latent: torch.Tensor,14        noise: torch.Tensor, timestep: torch.Tensor15    ):16        """17        Diffusion forward corruption process.18        Input:19            - clean_latent: the clean latent with shape [B, C, H, W]20            - noise: the noise with shape [B, C, H, W]21            - timestep: the timestep with shape [B]22        Output: the corrupted latent with shape [B, C, H, W]23        """24        pass25 26    def convert_x0_to_noise(27        self, x0: torch.Tensor, xt: torch.Tensor,28        timestep: torch.Tensor29    ) -> torch.Tensor:30        """31        Convert the diffusion network's x0 prediction to noise predidction.32        x0: the predicted clean data with shape [B, C, H, W]33        xt: the input noisy data with shape [B, C, H, W]34        timestep: the timestep with shape [B]35 36        noise = (xt-sqrt(alpha_t)*x0) / sqrt(beta_t) (eq 11 in https://arxiv.org/abs/2311.18828)37        """38        # use higher precision for calculations39        original_dtype = x0.dtype40        x0, xt, alphas_cumprod = map(41            lambda x: x.double().to(x0.device), [x0, xt,42                                                 self.alphas_cumprod]43        )44 45        alpha_prod_t = alphas_cumprod[timestep].reshape(-1, 1, 1, 1)46        beta_prod_t = 1 - alpha_prod_t47 48        noise_pred = (xt - alpha_prod_t **49                      (0.5) * x0) / beta_prod_t ** (0.5)50        return noise_pred.to(original_dtype)51 52    def convert_noise_to_x0(53        self, noise: torch.Tensor, xt: torch.Tensor,54        timestep: torch.Tensor55    ) -> torch.Tensor:56        """57        Convert the diffusion network's noise prediction to x0 predidction.58        noise: the predicted noise with shape [B, C, H, W]59        xt: the input noisy data with shape [B, C, H, W]60        timestep: the timestep with shape [B]61 62        x0 = (x_t - sqrt(beta_t) * noise) / sqrt(alpha_t) (eq 11 in https://arxiv.org/abs/2311.18828)63        """64        # use higher precision for calculations65        original_dtype = noise.dtype66        noise, xt, alphas_cumprod = map(67            lambda x: x.double().to(noise.device), [noise, xt,68                                                    self.alphas_cumprod]69        )70        alpha_prod_t = alphas_cumprod[timestep].reshape(-1, 1, 1, 1)71        beta_prod_t = 1 - alpha_prod_t72 73        x0_pred = (xt - beta_prod_t **74                   (0.5) * noise) / alpha_prod_t ** (0.5)75        return x0_pred.to(original_dtype)76 77    def convert_velocity_to_x0(78        self, velocity: torch.Tensor, xt: torch.Tensor,79        timestep: torch.Tensor80    ) -> torch.Tensor:81        """82        Convert the diffusion network's velocity prediction to x0 predidction.83        velocity: the predicted noise with shape [B, C, H, W]84        xt: the input noisy data with shape [B, C, H, W]85        timestep: the timestep with shape [B]86 87        v = sqrt(alpha_t) * noise - sqrt(beta_t) x088        noise = (xt-sqrt(alpha_t)*x0) / sqrt(beta_t)89        given v, x_t, we have90        x0 = sqrt(alpha_t) * x_t - sqrt(beta_t) * v91        see derivations https://chatgpt.com/share/679fb6c8-3a30-8008-9b0e-d1ae892dac5692        """93        # use higher precision for calculations94        original_dtype = velocity.dtype95        velocity, xt, alphas_cumprod = map(96            lambda x: x.double().to(velocity.device), [velocity, xt,97                                                       self.alphas_cumprod]98        )99        alpha_prod_t = alphas_cumprod[timestep].reshape(-1, 1, 1, 1)100        beta_prod_t = 1 - alpha_prod_t101 102        x0_pred = (alpha_prod_t ** 0.5) * xt - (beta_prod_t ** 0.5) * velocity103        return x0_pred.to(original_dtype)104 105 106class FlowMatchScheduler():107 108    def __init__(self, num_inference_steps=100, num_train_timesteps=1000, shift=3.0, sigma_max=1.0, sigma_min=0.003 / 1.002, inverse_timesteps=False, extra_one_step=False, reverse_sigmas=False):109        self.num_train_timesteps = num_train_timesteps110        self.shift = shift111        self.sigma_max = sigma_max112        self.sigma_min = sigma_min113        self.inverse_timesteps = inverse_timesteps114        self.extra_one_step = extra_one_step115        self.reverse_sigmas = reverse_sigmas116        self.set_timesteps(num_inference_steps)117 118    def set_timesteps(self, num_inference_steps=100, denoising_strength=1.0, training=False):119        sigma_start = self.sigma_min + \120            (self.sigma_max - self.sigma_min) * denoising_strength121        if self.extra_one_step:122            self.sigmas = torch.linspace(123                sigma_start, self.sigma_min, num_inference_steps + 1)[:-1]124        else:125            self.sigmas = torch.linspace(126                sigma_start, self.sigma_min, num_inference_steps)127        if self.inverse_timesteps:128            self.sigmas = torch.flip(self.sigmas, dims=[0])129        self.sigmas = self.shift * self.sigmas / \130            (1 + (self.shift - 1) * self.sigmas)131        if self.reverse_sigmas:132            self.sigmas = 1 - self.sigmas133        self.timesteps = self.sigmas * self.num_train_timesteps134        if training:135            x = self.timesteps136            y = torch.exp(-2 * ((x - num_inference_steps / 2) /137                          num_inference_steps) ** 2)138            y_shifted = y - y.min()139            bsmntw_weighing = y_shifted * \140                (num_inference_steps / y_shifted.sum())141            self.linear_timesteps_weights = bsmntw_weighing142 143    def step(self, model_output, timestep, sample, to_final=False):144        if timestep.ndim == 2:145            timestep = timestep.flatten(0, 1)146        self.sigmas = self.sigmas.to(model_output.device)147        self.timesteps = self.timesteps.to(model_output.device)148        timestep_id = torch.argmin(149            (self.timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1)150        sigma = self.sigmas[timestep_id].reshape(-1, 1, 1, 1)151        if to_final or (timestep_id + 1 >= len(self.timesteps)).any():152            sigma_ = 1 if (153                self.inverse_timesteps or self.reverse_sigmas) else 0154        else:155            sigma_ = self.sigmas[timestep_id + 1].reshape(-1, 1, 1, 1)156        prev_sample = sample + model_output * (sigma_ - sigma)157        return prev_sample158 159    def add_noise(self, original_samples, noise, timestep):160        """161        Diffusion forward corruption process.162        Input:163            - clean_latent: the clean latent with shape [B*T, C, H, W]164            - noise: the noise with shape [B*T, C, H, W]165            - timestep: the timestep with shape [B*T]166        Output: the corrupted latent with shape [B*T, C, H, W]167        """168        if timestep.ndim == 2:169            timestep = timestep.flatten(0, 1)170        self.sigmas = self.sigmas.to(noise.device)171        self.timesteps = self.timesteps.to(noise.device)172        timestep_id = torch.argmin(173            (self.timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1)174        sigma = self.sigmas[timestep_id].reshape(-1, 1, 1, 1)175        sample = (1 - sigma) * original_samples + sigma * noise176        return sample.type_as(noise)177 178    def training_target(self, sample, noise, timestep):179        target = noise - sample180        return target181 182    def training_weight(self, timestep):183        """184        Input:185            - timestep: the timestep with shape [B*T]186        Output: the corresponding weighting [B*T]187        """188        if timestep.ndim == 2:189            timestep = timestep.flatten(0, 1)190        self.linear_timesteps_weights = self.linear_timesteps_weights.to(timestep.device)191        timestep_id = torch.argmin(192            (self.timesteps.unsqueeze(1) - timestep.unsqueeze(0)).abs(), dim=0)193        weights = self.linear_timesteps_weights[timestep_id]194        return weights195