iyedjb/self-forcing

1from abc import ABC, abstractmethod2import torch3 4 5class DenoisingLoss(ABC):6    @abstractmethod7    def __call__(8        self, x: torch.Tensor, x_pred: torch.Tensor,9        noise: torch.Tensor, noise_pred: torch.Tensor,10        alphas_cumprod: torch.Tensor,11        timestep: torch.Tensor,12        **kwargs13    ) -> torch.Tensor:14        """15        Base class for denoising loss.16        Input:17            - x: the clean data with shape [B, F, C, H, W]18            - x_pred: the predicted clean data with shape [B, F, C, H, W]19            - noise: the noise with shape [B, F, C, H, W]20            - noise_pred: the predicted noise with shape [B, F, C, H, W]21            - alphas_cumprod: the cumulative product of alphas (defining the noise schedule) with shape [T]22            - timestep: the current timestep with shape [B, F]23        """24        pass25 26 27class X0PredLoss(DenoisingLoss):28    def __call__(29        self, x: torch.Tensor, x_pred: torch.Tensor,30        noise: torch.Tensor, noise_pred: torch.Tensor,31        alphas_cumprod: torch.Tensor,32        timestep: torch.Tensor,33        **kwargs34    ) -> torch.Tensor:35        return torch.mean((x - x_pred) ** 2)36 37 38class VPredLoss(DenoisingLoss):39    def __call__(40        self, x: torch.Tensor, x_pred: torch.Tensor,41        noise: torch.Tensor, noise_pred: torch.Tensor,42        alphas_cumprod: torch.Tensor,43        timestep: torch.Tensor,44        **kwargs45    ) -> torch.Tensor:46        weights = 1 / (1 - alphas_cumprod[timestep].reshape(*timestep.shape, 1, 1, 1))47        return torch.mean(weights * (x - x_pred) ** 2)48 49 50class NoisePredLoss(DenoisingLoss):51    def __call__(52        self, x: torch.Tensor, x_pred: torch.Tensor,53        noise: torch.Tensor, noise_pred: torch.Tensor,54        alphas_cumprod: torch.Tensor,55        timestep: torch.Tensor,56        **kwargs57    ) -> torch.Tensor:58        return torch.mean((noise - noise_pred) ** 2)59 60 61class FlowPredLoss(DenoisingLoss):62    def __call__(63        self, x: torch.Tensor, x_pred: torch.Tensor,64        noise: torch.Tensor, noise_pred: torch.Tensor,65        alphas_cumprod: torch.Tensor,66        timestep: torch.Tensor,67        **kwargs68    ) -> torch.Tensor:69        return torch.mean((kwargs["flow_pred"] - (noise - x)) ** 2)70 71 72NAME_TO_CLASS = {73    "x0": X0PredLoss,74    "v": VPredLoss,75    "noise": NoisePredLoss,76    "flow": FlowPredLoss77}78 79 80def get_denoising_loss(loss_type: str) -> DenoisingLoss:81    return NAME_TO_CLASS[loss_type]82