iyedjb/self-forcing

1from typing import Tuple2from einops import rearrange3from torch import nn4import torch.distributed as dist5import torch6 7from pipeline import SelfForcingTrainingPipeline8from utils.loss import get_denoising_loss9from utils.wan_wrapper import WanDiffusionWrapper, WanTextEncoder, WanVAEWrapper10 11 12class BaseModel(nn.Module):13    def __init__(self, args, device):14        super().__init__()15        self._initialize_models(args, device)16 17        self.device = device18        self.args = args19        self.dtype = torch.bfloat16 if args.mixed_precision else torch.float3220        if hasattr(args, "denoising_step_list"):21            self.denoising_step_list = torch.tensor(args.denoising_step_list, dtype=torch.long)22            if args.warp_denoising_step:23                timesteps = torch.cat((self.scheduler.timesteps.cpu(), torch.tensor([0], dtype=torch.float32)))24                self.denoising_step_list = timesteps[1000 - self.denoising_step_list]25 26    def _initialize_models(self, args, device):27        self.real_model_name = getattr(args, "real_name", "Wan2.1-T2V-1.3B")28        self.fake_model_name = getattr(args, "fake_name", "Wan2.1-T2V-1.3B")29 30        self.generator = WanDiffusionWrapper(**getattr(args, "model_kwargs", {}), is_causal=True)31        self.generator.model.requires_grad_(True)32 33        self.real_score = WanDiffusionWrapper(model_name=self.real_model_name, is_causal=False)34        self.real_score.model.requires_grad_(False)35 36        self.fake_score = WanDiffusionWrapper(model_name=self.fake_model_name, is_causal=False)37        self.fake_score.model.requires_grad_(True)38 39        self.text_encoder = WanTextEncoder()40        self.text_encoder.requires_grad_(False)41 42        self.vae = WanVAEWrapper()43        self.vae.requires_grad_(False)44 45        self.scheduler = self.generator.get_scheduler()46        self.scheduler.timesteps = self.scheduler.timesteps.to(device)47 48    def _get_timestep(49            self,50            min_timestep: int,51            max_timestep: int,52            batch_size: int,53            num_frame: int,54            num_frame_per_block: int,55            uniform_timestep: bool = False56    ) -> torch.Tensor:57        """58        Randomly generate a timestep tensor based on the generator's task type. It uniformly samples a timestep59        from the range [min_timestep, max_timestep], and returns a tensor of shape [batch_size, num_frame].60        - If uniform_timestep, it will use the same timestep for all frames.61        - If not uniform_timestep, it will use a different timestep for each block.62        """63        if uniform_timestep:64            timestep = torch.randint(65                min_timestep,66                max_timestep,67                [batch_size, 1],68                device=self.device,69                dtype=torch.long70            ).repeat(1, num_frame)71            return timestep72        else:73            timestep = torch.randint(74                min_timestep,75                max_timestep,76                [batch_size, num_frame],77                device=self.device,78                dtype=torch.long79            )80            # make the noise level the same within every block81            if self.independent_first_frame:82                # the first frame is always kept the same83                timestep_from_second = timestep[:, 1:]84                timestep_from_second = timestep_from_second.reshape(85                    timestep_from_second.shape[0], -1, num_frame_per_block)86                timestep_from_second[:, :, 1:] = timestep_from_second[:, :, 0:1]87                timestep_from_second = timestep_from_second.reshape(88                    timestep_from_second.shape[0], -1)89                timestep = torch.cat([timestep[:, 0:1], timestep_from_second], dim=1)90            else:91                timestep = timestep.reshape(92                    timestep.shape[0], -1, num_frame_per_block)93                timestep[:, :, 1:] = timestep[:, :, 0:1]94                timestep = timestep.reshape(timestep.shape[0], -1)95            return timestep96 97 98class SelfForcingModel(BaseModel):99    def __init__(self, args, device):100        super().__init__(args, device)101        self.denoising_loss_func = get_denoising_loss(args.denoising_loss_type)()102 103    def _run_generator(104        self,105        image_or_video_shape,106        conditional_dict: dict,107        initial_latent: torch.tensor = None108    ) -> Tuple[torch.Tensor, torch.Tensor]:109        """110        Optionally simulate the generator's input from noise using backward simulation111        and then run the generator for one-step.112        Input:113            - image_or_video_shape: a list containing the shape of the image or video [B, F, C, H, W].114            - conditional_dict: a dictionary containing the conditional information (e.g. text embeddings, image embeddings).115            - unconditional_dict: a dictionary containing the unconditional information (e.g. null/negative text embeddings, null/negative image embeddings).116            - clean_latent: a tensor containing the clean latents [B, F, C, H, W]. Need to be passed when no backward simulation is used.117            - initial_latent: a tensor containing the initial latents [B, F, C, H, W].118        Output:119            - pred_image: a tensor with shape [B, F, C, H, W].120            - denoised_timestep: an integer121        """122        # Step 1: Sample noise and backward simulate the generator's input123        assert getattr(self.args, "backward_simulation", True), "Backward simulation needs to be enabled"124        if initial_latent is not None:125            conditional_dict["initial_latent"] = initial_latent126        if self.args.i2v:127            noise_shape = [image_or_video_shape[0], image_or_video_shape[1] - 1, *image_or_video_shape[2:]]128        else:129            noise_shape = image_or_video_shape.copy()130 131        # During training, the number of generated frames should be uniformly sampled from132        # [21, self.num_training_frames], but still being a multiple of self.num_frame_per_block133        min_num_frames = 20 if self.args.independent_first_frame else 21134        max_num_frames = self.num_training_frames - 1 if self.args.independent_first_frame else self.num_training_frames135        assert max_num_frames % self.num_frame_per_block == 0136        assert min_num_frames % self.num_frame_per_block == 0137        max_num_blocks = max_num_frames // self.num_frame_per_block138        min_num_blocks = min_num_frames // self.num_frame_per_block139        num_generated_blocks = torch.randint(min_num_blocks, max_num_blocks + 1, (1,), device=self.device)140        dist.broadcast(num_generated_blocks, src=0)141        num_generated_blocks = num_generated_blocks.item()142        num_generated_frames = num_generated_blocks * self.num_frame_per_block143        if self.args.independent_first_frame and initial_latent is None:144            num_generated_frames += 1145            min_num_frames += 1146        # Sync num_generated_frames across all processes147        noise_shape[1] = num_generated_frames148 149        pred_image_or_video, denoised_timestep_from, denoised_timestep_to = self._consistency_backward_simulation(150            noise=torch.randn(noise_shape,151                              device=self.device, dtype=self.dtype),152            **conditional_dict,153        )154        # Slice last 21 frames155        if pred_image_or_video.shape[1] > 21:156            with torch.no_grad():157                # Reencode to get image latent158                latent_to_decode = pred_image_or_video[:, :-20, ...]159                # Deccode to video160                pixels = self.vae.decode_to_pixel(latent_to_decode)161                frame = pixels[:, -1:, ...].to(self.dtype)162                frame = rearrange(frame, "b t c h w -> b c t h w")163                # Encode frame to get image latent164                image_latent = self.vae.encode_to_latent(frame).to(self.dtype)165            pred_image_or_video_last_21 = torch.cat([image_latent, pred_image_or_video[:, -20:, ...]], dim=1)166        else:167            pred_image_or_video_last_21 = pred_image_or_video168 169        if num_generated_frames != min_num_frames:170            # Currently, we do not use gradient for the first chunk, since it contains image latents171            gradient_mask = torch.ones_like(pred_image_or_video_last_21, dtype=torch.bool)172            if self.args.independent_first_frame:173                gradient_mask[:, :1] = False174            else:175                gradient_mask[:, :self.num_frame_per_block] = False176        else:177            gradient_mask = None178 179        pred_image_or_video_last_21 = pred_image_or_video_last_21.to(self.dtype)180        return pred_image_or_video_last_21, gradient_mask, denoised_timestep_from, denoised_timestep_to181 182    def _consistency_backward_simulation(183        self,184        noise: torch.Tensor,185        **conditional_dict: dict186    ) -> torch.Tensor:187        """188        Simulate the generator's input from noise to avoid training/inference mismatch.189        See Sec 4.5 of the DMD2 paper (https://arxiv.org/abs/2405.14867) for details.190        Here we use the consistency sampler (https://arxiv.org/abs/2303.01469)191        Input:192            - noise: a tensor sampled from N(0, 1) with shape [B, F, C, H, W] where the number of frame is 1 for images.193            - conditional_dict: a dictionary containing the conditional information (e.g. text embeddings, image embeddings).194        Output:195            - output: a tensor with shape [B, T, F, C, H, W].196            T is the total number of timesteps. output[0] is a pure noise and output[i] and i>0197            represents the x0 prediction at each timestep.198        """199        if self.inference_pipeline is None:200            self._initialize_inference_pipeline()201 202        return self.inference_pipeline.inference_with_trajectory(203            noise=noise, **conditional_dict204        )205 206    def _initialize_inference_pipeline(self):207        """208        Lazy initialize the inference pipeline during the first backward simulation run.209        Here we encapsulate the inference code with a model-dependent outside function.210        We pass our FSDP-wrapped modules into the pipeline to save memory.211        """212        self.inference_pipeline = SelfForcingTrainingPipeline(213            denoising_step_list=self.denoising_step_list,214            scheduler=self.scheduler,215            generator=self.generator,216            num_frame_per_block=self.num_frame_per_block,217            independent_first_frame=self.args.independent_first_frame,218            same_step_across_blocks=self.args.same_step_across_blocks,219            last_step_only=self.args.last_step_only,220            num_max_frames=self.num_training_frames,221            context_noise=self.args.context_noise222        )223