Spaces:
Running
on
Zero
Running
on
Zero
File size: 3,109 Bytes
0fd2f06 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 |
from typing import List
import torch
from utils.wan_wrapper import WanDiffusionWrapper, WanTextEncoder, WanVAEWrapper
class BidirectionalInferencePipeline(torch.nn.Module):
def __init__(
self,
args,
device,
generator=None,
text_encoder=None,
vae=None
):
super().__init__()
# Step 1: Initialize all models
self.generator = WanDiffusionWrapper(
**getattr(args, "model_kwargs", {}), is_causal=False) if generator is None else generator
self.text_encoder = WanTextEncoder() if text_encoder is None else text_encoder
self.vae = WanVAEWrapper() if vae is None else vae
# Step 2: Initialize all bidirectional wan hyperparmeters
self.scheduler = self.generator.get_scheduler()
self.denoising_step_list = torch.tensor(
args.denoising_step_list, dtype=torch.long, device=device)
if self.denoising_step_list[-1] == 0:
self.denoising_step_list = self.denoising_step_list[:-1] # remove the zero timestep for inference
if args.warp_denoising_step:
timesteps = torch.cat((self.scheduler.timesteps.cpu(), torch.tensor([0], dtype=torch.float32)))
self.denoising_step_list = timesteps[1000 - self.denoising_step_list]
def inference(self, noise: torch.Tensor, text_prompts: List[str]) -> torch.Tensor:
"""
Perform inference on the given noise and text prompts.
Inputs:
noise (torch.Tensor): The input noise tensor of shape
(batch_size, num_frames, num_channels, height, width).
text_prompts (List[str]): The list of text prompts.
Outputs:
video (torch.Tensor): The generated video tensor of shape
(batch_size, num_frames, num_channels, height, width). It is normalized to be in the range [0, 1].
"""
conditional_dict = self.text_encoder(
text_prompts=text_prompts
)
# initial point
noisy_image_or_video = noise
# use the last n-1 timesteps to simulate the generator's input
for index, current_timestep in enumerate(self.denoising_step_list[:-1]):
_, pred_image_or_video = self.generator(
noisy_image_or_video=noisy_image_or_video,
conditional_dict=conditional_dict,
timestep=torch.ones(
noise.shape[:2], dtype=torch.long, device=noise.device) * current_timestep
) # [B, F, C, H, W]
next_timestep = self.denoising_step_list[index + 1] * torch.ones(
noise.shape[:2], dtype=torch.long, device=noise.device)
noisy_image_or_video = self.scheduler.add_noise(
pred_image_or_video.flatten(0, 1),
torch.randn_like(pred_image_or_video.flatten(0, 1)),
next_timestep.flatten(0, 1)
).unflatten(0, noise.shape[:2])
video = self.vae.decode_to_pixel(pred_image_or_video)
video = (video * 0.5 + 0.5).clamp(0, 1)
return video
|