AstroAUmin/self-forcing

1from tqdm import tqdm2from typing import List3import torch4 5from wan.utils.fm_solvers import FlowDPMSolverMultistepScheduler, get_sampling_sigmas, retrieve_timesteps6from wan.utils.fm_solvers_unipc import FlowUniPCMultistepScheduler7from utils.wan_wrapper import WanDiffusionWrapper, WanTextEncoder, WanVAEWrapper8 9 10class BidirectionalDiffusionInferencePipeline(torch.nn.Module):11    def __init__(12            self,13            args,14            device,15            generator=None,16            text_encoder=None,17            vae=None18    ):19        super().__init__()20        # Step 1: Initialize all models21        self.generator = WanDiffusionWrapper(22            **getattr(args, "model_kwargs", {}), is_causal=False) if generator is None else generator23        self.text_encoder = WanTextEncoder() if text_encoder is None else text_encoder24        self.vae = WanVAEWrapper() if vae is None else vae25 26        # Step 2: Initialize scheduler27        self.num_train_timesteps = args.num_train_timestep28        self.sampling_steps = 5029        self.sample_solver = 'unipc'30        self.shift = 8.031 32        self.args = args33 34    def inference(35        self,36        noise: torch.Tensor,37        text_prompts: List[str],38        return_latents=False39    ) -> torch.Tensor:40        """41        Perform inference on the given noise and text prompts.42        Inputs:43            noise (torch.Tensor): The input noise tensor of shape44                (batch_size, num_frames, num_channels, height, width).45            text_prompts (List[str]): The list of text prompts.46        Outputs:47            video (torch.Tensor): The generated video tensor of shape48                (batch_size, num_frames, num_channels, height, width). It is normalized to be in the range [0, 1].49        """50 51        conditional_dict = self.text_encoder(52            text_prompts=text_prompts53        )54        unconditional_dict = self.text_encoder(55            text_prompts=[self.args.negative_prompt] * len(text_prompts)56        )57 58        latents = noise59 60        sample_scheduler = self._initialize_sample_scheduler(noise)61        for _, t in enumerate(tqdm(sample_scheduler.timesteps)):62            latent_model_input = latents63            timestep = t * torch.ones([latents.shape[0], 21], device=noise.device, dtype=torch.float32)64 65            flow_pred_cond, _ = self.generator(latent_model_input, conditional_dict, timestep)66            flow_pred_uncond, _ = self.generator(latent_model_input, unconditional_dict, timestep)67 68            flow_pred = flow_pred_uncond + self.args.guidance_scale * (69                flow_pred_cond - flow_pred_uncond)70 71            temp_x0 = sample_scheduler.step(72                flow_pred.unsqueeze(0),73                t,74                latents.unsqueeze(0),75                return_dict=False)[0]76            latents = temp_x0.squeeze(0)77 78        x0 = latents79        video = self.vae.decode_to_pixel(x0)80        video = (video * 0.5 + 0.5).clamp(0, 1)81 82        del sample_scheduler83 84        if return_latents:85            return video, latents86        else:87            return video88 89    def _initialize_sample_scheduler(self, noise):90        if self.sample_solver == 'unipc':91            sample_scheduler = FlowUniPCMultistepScheduler(92                num_train_timesteps=self.num_train_timesteps,93                shift=1,94                use_dynamic_shifting=False)95            sample_scheduler.set_timesteps(96                self.sampling_steps, device=noise.device, shift=self.shift)97            self.timesteps = sample_scheduler.timesteps98        elif self.sample_solver == 'dpm++':99            sample_scheduler = FlowDPMSolverMultistepScheduler(100                num_train_timesteps=self.num_train_timesteps,101                shift=1,102                use_dynamic_shifting=False)103            sampling_sigmas = get_sampling_sigmas(self.sampling_steps, self.shift)104            self.timesteps, _ = retrieve_timesteps(105                sample_scheduler,106                device=noise.device,107                sigmas=sampling_sigmas)108        else:109            raise NotImplementedError("Unsupported solver.")110        return sample_scheduler111