iyedjb/self-forcing

1from typing import Tuple2import torch3 4from model.base import BaseModel5from utils.wan_wrapper import WanDiffusionWrapper, WanTextEncoder, WanVAEWrapper6 7 8class CausalDiffusion(BaseModel):9    def __init__(self, args, device):10        """11        Initialize the Diffusion loss module.12        """13        super().__init__(args, device)14        self.num_frame_per_block = getattr(args, "num_frame_per_block", 1)15        if self.num_frame_per_block > 1:16            self.generator.model.num_frame_per_block = self.num_frame_per_block17        self.independent_first_frame = getattr(args, "independent_first_frame", False)18        if self.independent_first_frame:19            self.generator.model.independent_first_frame = True20 21        if args.gradient_checkpointing:22            self.generator.enable_gradient_checkpointing()23 24        # Step 2: Initialize all hyperparameters25        self.num_train_timestep = args.num_train_timestep26        self.min_step = int(0.02 * self.num_train_timestep)27        self.max_step = int(0.98 * self.num_train_timestep)28        self.guidance_scale = args.guidance_scale29        self.timestep_shift = getattr(args, "timestep_shift", 1.0)30        self.teacher_forcing = getattr(args, "teacher_forcing", False)31        # Noise augmentation in teacher forcing, we add small noise to clean context latents32        self.noise_augmentation_max_timestep = getattr(args, "noise_augmentation_max_timestep", 0)33 34    def _initialize_models(self, args):35        self.generator = WanDiffusionWrapper(**getattr(args, "model_kwargs", {}), is_causal=True)36        self.generator.model.requires_grad_(True)37 38        self.text_encoder = WanTextEncoder()39        self.text_encoder.requires_grad_(False)40 41        self.vae = WanVAEWrapper()42        self.vae.requires_grad_(False)43 44    def generator_loss(45        self,46        image_or_video_shape,47        conditional_dict: dict,48        unconditional_dict: dict,49        clean_latent: torch.Tensor,50        initial_latent: torch.Tensor = None51    ) -> Tuple[torch.Tensor, dict]:52        """53        Generate image/videos from noise and compute the DMD loss.54        The noisy input to the generator is backward simulated.55        This removes the need of any datasets during distillation.56        See Sec 4.5 of the DMD2 paper (https://arxiv.org/abs/2405.14867) for details.57        Input:58            - image_or_video_shape: a list containing the shape of the image or video [B, F, C, H, W].59            - conditional_dict: a dictionary containing the conditional information (e.g. text embeddings, image embeddings).60            - unconditional_dict: a dictionary containing the unconditional information (e.g. null/negative text embeddings, null/negative image embeddings).61            - clean_latent: a tensor containing the clean latents [B, F, C, H, W]. Need to be passed when no backward simulation is used.62        Output:63            - loss: a scalar tensor representing the generator loss.64            - generator_log_dict: a dictionary containing the intermediate tensors for logging.65        """66        noise = torch.randn_like(clean_latent)67        batch_size, num_frame = image_or_video_shape[:2]68 69        # Step 2: Randomly sample a timestep and add noise to denoiser inputs70        index = self._get_timestep(71            0,72            self.scheduler.num_train_timesteps,73            image_or_video_shape[0],74            image_or_video_shape[1],75            self.num_frame_per_block,76            uniform_timestep=False77        )78        timestep = self.scheduler.timesteps[index].to(dtype=self.dtype, device=self.device)79        noisy_latents = self.scheduler.add_noise(80            clean_latent.flatten(0, 1),81            noise.flatten(0, 1),82            timestep.flatten(0, 1)83        ).unflatten(0, (batch_size, num_frame))84        training_target = self.scheduler.training_target(clean_latent, noise, timestep)85 86        # Step 3: Noise augmentation, also add small noise to clean context latents87        if self.noise_augmentation_max_timestep > 0:88            index_clean_aug = self._get_timestep(89                0,90                self.noise_augmentation_max_timestep,91                image_or_video_shape[0],92                image_or_video_shape[1],93                self.num_frame_per_block,94                uniform_timestep=False95            )96            timestep_clean_aug = self.scheduler.timesteps[index_clean_aug].to(dtype=self.dtype, device=self.device)97            clean_latent_aug = self.scheduler.add_noise(98                clean_latent.flatten(0, 1),99                noise.flatten(0, 1),100                timestep_clean_aug.flatten(0, 1)101            ).unflatten(0, (batch_size, num_frame))102        else:103            clean_latent_aug = clean_latent104            timestep_clean_aug = None105 106        # Compute loss107        flow_pred, x0_pred = self.generator(108            noisy_image_or_video=noisy_latents,109            conditional_dict=conditional_dict,110            timestep=timestep,111            clean_x=clean_latent_aug if self.teacher_forcing else None,112            aug_t=timestep_clean_aug if self.teacher_forcing else None113        )114        # loss = torch.nn.functional.mse_loss(flow_pred.float(), training_target.float())115        loss = torch.nn.functional.mse_loss(116            flow_pred.float(), training_target.float(), reduction='none'117        ).mean(dim=(2, 3, 4))118        loss = loss * self.scheduler.training_weight(timestep).unflatten(0, (batch_size, num_frame))119        loss = loss.mean()120 121        log_dict = {122            "x0": clean_latent.detach(),123            "x0_pred": x0_pred.detach()124        }125        return loss, log_dict126