iyedjb/self-forcing

1import torch.nn.functional as F2from typing import Tuple3import torch4 5from model.base import BaseModel6from utils.wan_wrapper import WanDiffusionWrapper, WanTextEncoder, WanVAEWrapper7 8 9class ODERegression(BaseModel):10    def __init__(self, args, device):11        """12        Initialize the ODERegression module.13        This class is self-contained and compute generator losses14        in the forward pass given precomputed ode solution pairs.15        This class supports the ode regression loss for both causal and bidirectional models.16        See Sec 4.3 of CausVid https://arxiv.org/abs/2412.07772 for details17        """18        super().__init__(args, device)19 20        # Step 1: Initialize all models21 22        self.generator = WanDiffusionWrapper(**getattr(args, "model_kwargs", {}), is_causal=True)23        self.generator.model.requires_grad_(True)24        if getattr(args, "generator_ckpt", False):25            print(f"Loading pretrained generator from {args.generator_ckpt}")26            state_dict = torch.load(args.generator_ckpt, map_location="cpu")[27                'generator']28            self.generator.load_state_dict(29                state_dict, strict=True30            )31 32        self.num_frame_per_block = getattr(args, "num_frame_per_block", 1)33 34        if self.num_frame_per_block > 1:35            self.generator.model.num_frame_per_block = self.num_frame_per_block36 37        self.independent_first_frame = getattr(args, "independent_first_frame", False)38        if self.independent_first_frame:39            self.generator.model.independent_first_frame = True40        if args.gradient_checkpointing:41            self.generator.enable_gradient_checkpointing()42 43        # Step 2: Initialize all hyperparameters44        self.timestep_shift = getattr(args, "timestep_shift", 1.0)45 46    def _initialize_models(self, args):47        self.generator = WanDiffusionWrapper(**getattr(args, "model_kwargs", {}), is_causal=True)48        self.generator.model.requires_grad_(True)49 50        self.text_encoder = WanTextEncoder()51        self.text_encoder.requires_grad_(False)52 53        self.vae = WanVAEWrapper()54        self.vae.requires_grad_(False)55 56    @torch.no_grad()57    def _prepare_generator_input(self, ode_latent: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:58        """59        Given a tensor containing the whole ODE sampling trajectories,60        randomly choose an intermediate timestep and return the latent as well as the corresponding timestep.61        Input:62            - ode_latent: a tensor containing the whole ODE sampling trajectories [batch_size, num_denoising_steps, num_frames, num_channels, height, width].63        Output:64            - noisy_input: a tensor containing the selected latent [batch_size, num_frames, num_channels, height, width].65            - timestep: a tensor containing the corresponding timestep [batch_size].66        """67        batch_size, num_denoising_steps, num_frames, num_channels, height, width = ode_latent.shape68 69        # Step 1: Randomly choose a timestep for each frame70        index = self._get_timestep(71            0,72            len(self.denoising_step_list),73            batch_size,74            num_frames,75            self.num_frame_per_block,76            uniform_timestep=False77        )78        if self.args.i2v:79            index[:, 0] = len(self.denoising_step_list) - 180 81        noisy_input = torch.gather(82            ode_latent, dim=1,83            index=index.reshape(batch_size, 1, num_frames, 1, 1, 1).expand(84                -1, -1, -1, num_channels, height, width).to(self.device)85        ).squeeze(1)86 87        timestep = self.denoising_step_list[index].to(self.device)88 89        # if self.extra_noise_step > 0:90        #     random_timestep = torch.randint(0, self.extra_noise_step, [91        #                                     batch_size, num_frames], device=self.device, dtype=torch.long)92        #     perturbed_noisy_input = self.scheduler.add_noise(93        #         noisy_input.flatten(0, 1),94        #         torch.randn_like(noisy_input.flatten(0, 1)),95        #         random_timestep.flatten(0, 1)96        #     ).detach().unflatten(0, (batch_size, num_frames)).type_as(noisy_input)97 98        #     noisy_input[timestep == 0] = perturbed_noisy_input[timestep == 0]99 100        return noisy_input, timestep101 102    def generator_loss(self, ode_latent: torch.Tensor, conditional_dict: dict) -> Tuple[torch.Tensor, dict]:103        """104        Generate image/videos from noisy latents and compute the ODE regression loss.105        Input:106            - ode_latent: a tensor containing the ODE latents [batch_size, num_denoising_steps, num_frames, num_channels, height, width].107            They are ordered from most noisy to clean latents.108            - conditional_dict: a dictionary containing the conditional information (e.g. text embeddings, image embeddings).109        Output:110            - loss: a scalar tensor representing the generator loss.111            - log_dict: a dictionary containing additional information for loss timestep breakdown.112        """113        # Step 1: Run generator on noisy latents114        target_latent = ode_latent[:, -1]115 116        noisy_input, timestep = self._prepare_generator_input(117            ode_latent=ode_latent)118 119        _, pred_image_or_video = self.generator(120            noisy_image_or_video=noisy_input,121            conditional_dict=conditional_dict,122            timestep=timestep123        )124 125        # Step 2: Compute the regression loss126        mask = timestep != 0127 128        loss = F.mse_loss(129            pred_image_or_video[mask], target_latent[mask], reduction="mean")130 131        log_dict = {132            "unnormalized_loss": F.mse_loss(pred_image_or_video, target_latent, reduction='none').mean(dim=[1, 2, 3, 4]).detach(),133            "timestep": timestep.float().mean(dim=1).detach(),134            "input": noisy_input.detach(),135            "output": pred_image_or_video.detach(),136        }137 138        return loss, log_dict139