iyedjb/self-forcing
0
1import types2from typing import List, Optional3import torch4from torch import nn5 6from utils.scheduler import SchedulerInterface, FlowMatchScheduler7from wan.modules.tokenizers import HuggingfaceTokenizer8from wan.modules.model import WanModel, RegisterTokens, GanAttentionBlock9from wan.modules.vae import _video_vae10from wan.modules.t5 import umt5_xxl11from wan.modules.causal_model import CausalWanModel12 13 14class WanTextEncoder(torch.nn.Module):15 def __init__(self) -> None:16 super().__init__()17 18 self.text_encoder = umt5_xxl(19 encoder_only=True,20 return_tokenizer=False,21 dtype=torch.float32,22 device=torch.device('cpu')23 ).eval().requires_grad_(False)24 self.text_encoder.load_state_dict(25 torch.load("wan_models/Wan2.1-T2V-1.3B/models_t5_umt5-xxl-enc-bf16.pth",26 map_location='cpu', weights_only=False)27 )28 29 self.tokenizer = HuggingfaceTokenizer(30 name="wan_models/Wan2.1-T2V-1.3B/google/umt5-xxl/", seq_len=512, clean='whitespace')31 32 @property33 def device(self):34 # Assume we are always on GPU35 return torch.cuda.current_device()36 37 def forward(self, text_prompts: List[str]) -> dict:38 ids, mask = self.tokenizer(39 text_prompts, return_mask=True, add_special_tokens=True)40 ids = ids.to(self.device)41 mask = mask.to(self.device)42 seq_lens = mask.gt(0).sum(dim=1).long()43 context = self.text_encoder(ids, mask)44 45 for u, v in zip(context, seq_lens):46 u[v:] = 0.0 # set padding to 0.047 48 return {49 "prompt_embeds": context50 }51 52 53class WanVAEWrapper(torch.nn.Module):54 def __init__(self):55 super().__init__()56 mean = [57 -0.7571, -0.7089, -0.9113, 0.1075, -0.1745, 0.9653, -0.1517, 1.5508,58 0.4134, -0.0715, 0.5517, -0.3632, -0.1922, -0.9497, 0.2503, -0.292159 ]60 std = [61 2.8184, 1.4541, 2.3275, 2.6558, 1.2196, 1.7708, 2.6052, 2.0743,62 3.2687, 2.1526, 2.8652, 1.5579, 1.6382, 1.1253, 2.8251, 1.916063 ]64 self.mean = torch.tensor(mean, dtype=torch.float32)65 self.std = torch.tensor(std, dtype=torch.float32)66 67 # init model68 self.model = _video_vae(69 pretrained_path="wan_models/Wan2.1-T2V-1.3B/Wan2.1_VAE.pth",70 z_dim=16,71 ).eval().requires_grad_(False)72 73 def encode_to_latent(self, pixel: torch.Tensor) -> torch.Tensor:74 # pixel: [batch_size, num_channels, num_frames, height, width]75 device, dtype = pixel.device, pixel.dtype76 scale = [self.mean.to(device=device, dtype=dtype),77 1.0 / self.std.to(device=device, dtype=dtype)]78 79 output = [80 self.model.encode(u.unsqueeze(0), scale).float().squeeze(0)81 for u in pixel82 ]83 output = torch.stack(output, dim=0)84 # from [batch_size, num_channels, num_frames, height, width]85 # to [batch_size, num_frames, num_channels, height, width]86 output = output.permute(0, 2, 1, 3, 4)87 return output88 89 def decode_to_pixel(self, latent: torch.Tensor, use_cache: bool = False) -> torch.Tensor:90 # from [batch_size, num_frames, num_channels, height, width]91 # to [batch_size, num_channels, num_frames, height, width]92 zs = latent.permute(0, 2, 1, 3, 4)93 if use_cache:94 assert latent.shape[0] == 1, "Batch size must be 1 when using cache"95 96 device, dtype = latent.device, latent.dtype97 scale = [self.mean.to(device=device, dtype=dtype),98 1.0 / self.std.to(device=device, dtype=dtype)]99 100 if use_cache:101 decode_function = self.model.cached_decode102 else:103 decode_function = self.model.decode104 105 output = []106 for u in zs:107 output.append(decode_function(u.unsqueeze(0), scale).float().clamp_(-1, 1).squeeze(0))108 output = torch.stack(output, dim=0)109 # from [batch_size, num_channels, num_frames, height, width]110 # to [batch_size, num_frames, num_channels, height, width]111 output = output.permute(0, 2, 1, 3, 4)112 return output113 114 115class WanDiffusionWrapper(torch.nn.Module):116 def __init__(117 self,118 model_name="Wan2.1-T2V-1.3B",119 timestep_shift=8.0,120 is_causal=False,121 local_attn_size=-1,122 sink_size=0123 ):124 super().__init__()125 126 if is_causal:127 self.model = CausalWanModel.from_pretrained(128 f"wan_models/{model_name}/", local_attn_size=local_attn_size, sink_size=sink_size)129 else:130 self.model = WanModel.from_pretrained(f"wan_models/{model_name}/")131 self.model.eval()132 133 # For non-causal diffusion, all frames share the same timestep134 self.uniform_timestep = not is_causal135 136 self.scheduler = FlowMatchScheduler(137 shift=timestep_shift, sigma_min=0.0, extra_one_step=True138 )139 self.scheduler.set_timesteps(1000, training=True)140 141 self.seq_len = 32760 # [1, 21, 16, 60, 104]142 self.post_init()143 144 def enable_gradient_checkpointing(self) -> None:145 self.model.enable_gradient_checkpointing()146 147 def adding_cls_branch(self, atten_dim=1536, num_class=4, time_embed_dim=0) -> None:148 # NOTE: This is hard coded for WAN2.1-T2V-1.3B for now!!!!!!!!!!!!!!!!!!!!149 self._cls_pred_branch = nn.Sequential(150 # Input: [B, 384, 21, 60, 104]151 nn.LayerNorm(atten_dim * 3 + time_embed_dim),152 nn.Linear(atten_dim * 3 + time_embed_dim, 1536),153 nn.SiLU(),154 nn.Linear(atten_dim, num_class)155 )156 self._cls_pred_branch.requires_grad_(True)157 num_registers = 3158 self._register_tokens = RegisterTokens(num_registers=num_registers, dim=atten_dim)159 self._register_tokens.requires_grad_(True)160 161 gan_ca_blocks = []162 for _ in range(num_registers):163 block = GanAttentionBlock()164 gan_ca_blocks.append(block)165 self._gan_ca_blocks = nn.ModuleList(gan_ca_blocks)166 self._gan_ca_blocks.requires_grad_(True)167 # self.has_cls_branch = True168 169 def _convert_flow_pred_to_x0(self, flow_pred: torch.Tensor, xt: torch.Tensor, timestep: torch.Tensor) -> torch.Tensor:170 """171 Convert flow matching's prediction to x0 prediction.172 flow_pred: the prediction with shape [B, C, H, W]173 xt: the input noisy data with shape [B, C, H, W]174 timestep: the timestep with shape [B]175 176 pred = noise - x0177 x_t = (1-sigma_t) * x0 + sigma_t * noise178 we have x0 = x_t - sigma_t * pred179 see derivations https://chatgpt.com/share/67bf8589-3d04-8008-bc6e-4cf1a24e2d0e180 """181 # use higher precision for calculations182 original_dtype = flow_pred.dtype183 flow_pred, xt, sigmas, timesteps = map(184 lambda x: x.double().to(flow_pred.device), [flow_pred, xt,185 self.scheduler.sigmas,186 self.scheduler.timesteps]187 )188 189 timestep_id = torch.argmin(190 (timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1)191 sigma_t = sigmas[timestep_id].reshape(-1, 1, 1, 1)192 x0_pred = xt - sigma_t * flow_pred193 return x0_pred.to(original_dtype)194 195 @staticmethod196 def _convert_x0_to_flow_pred(scheduler, x0_pred: torch.Tensor, xt: torch.Tensor, timestep: torch.Tensor) -> torch.Tensor:197 """198 Convert x0 prediction to flow matching's prediction.199 x0_pred: the x0 prediction with shape [B, C, H, W]200 xt: the input noisy data with shape [B, C, H, W]201 timestep: the timestep with shape [B]202 203 pred = (x_t - x_0) / sigma_t204 """205 # use higher precision for calculations206 original_dtype = x0_pred.dtype207 x0_pred, xt, sigmas, timesteps = map(208 lambda x: x.double().to(x0_pred.device), [x0_pred, xt,209 scheduler.sigmas,210 scheduler.timesteps]211 )212 timestep_id = torch.argmin(213 (timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1)214 sigma_t = sigmas[timestep_id].reshape(-1, 1, 1, 1)215 flow_pred = (xt - x0_pred) / sigma_t216 return flow_pred.to(original_dtype)217 218 def forward(219 self,220 noisy_image_or_video: torch.Tensor, conditional_dict: dict,221 timestep: torch.Tensor, kv_cache: Optional[List[dict]] = None,222 crossattn_cache: Optional[List[dict]] = None,223 current_start: Optional[int] = None,224 classify_mode: Optional[bool] = False,225 concat_time_embeddings: Optional[bool] = False,226 clean_x: Optional[torch.Tensor] = None,227 aug_t: Optional[torch.Tensor] = None,228 cache_start: Optional[int] = None229 ) -> torch.Tensor:230 prompt_embeds = conditional_dict["prompt_embeds"]231 232 # [B, F] -> [B]233 if self.uniform_timestep:234 input_timestep = timestep[:, 0]235 else:236 input_timestep = timestep237 238 logits = None239 # X0 prediction240 if kv_cache is not None:241 flow_pred = self.model(242 noisy_image_or_video.permute(0, 2, 1, 3, 4),243 t=input_timestep, context=prompt_embeds,244 seq_len=self.seq_len,245 kv_cache=kv_cache,246 crossattn_cache=crossattn_cache,247 current_start=current_start,248 cache_start=cache_start249 ).permute(0, 2, 1, 3, 4)250 else:251 if clean_x is not None:252 # teacher forcing253 flow_pred = self.model(254 noisy_image_or_video.permute(0, 2, 1, 3, 4),255 t=input_timestep, context=prompt_embeds,256 seq_len=self.seq_len,257 clean_x=clean_x.permute(0, 2, 1, 3, 4),258 aug_t=aug_t,259 ).permute(0, 2, 1, 3, 4)260 else:261 if classify_mode:262 flow_pred, logits = self.model(263 noisy_image_or_video.permute(0, 2, 1, 3, 4),264 t=input_timestep, context=prompt_embeds,265 seq_len=self.seq_len,266 classify_mode=True,267 register_tokens=self._register_tokens,268 cls_pred_branch=self._cls_pred_branch,269 gan_ca_blocks=self._gan_ca_blocks,270 concat_time_embeddings=concat_time_embeddings271 )272 flow_pred = flow_pred.permute(0, 2, 1, 3, 4)273 else:274 flow_pred = self.model(275 noisy_image_or_video.permute(0, 2, 1, 3, 4),276 t=input_timestep, context=prompt_embeds,277 seq_len=self.seq_len278 ).permute(0, 2, 1, 3, 4)279 280 pred_x0 = self._convert_flow_pred_to_x0(281 flow_pred=flow_pred.flatten(0, 1),282 xt=noisy_image_or_video.flatten(0, 1),283 timestep=timestep.flatten(0, 1)284 ).unflatten(0, flow_pred.shape[:2])285 286 if logits is not None:287 return flow_pred, pred_x0, logits288 289 return flow_pred, pred_x0290 291 def get_scheduler(self) -> SchedulerInterface:292 """293 Update the current scheduler with the interface's static method294 """295 scheduler = self.scheduler296 scheduler.convert_x0_to_noise = types.MethodType(297 SchedulerInterface.convert_x0_to_noise, scheduler)298 scheduler.convert_noise_to_x0 = types.MethodType(299 SchedulerInterface.convert_noise_to_x0, scheduler)300 scheduler.convert_velocity_to_x0 = types.MethodType(301 SchedulerInterface.convert_velocity_to_x0, scheduler)302 self.scheduler = scheduler303 return scheduler304 305 def post_init(self):306 """307 A few custom initialization steps that should be called after the object is created.308 Currently, the only one we have is to bind a few methods to scheduler.309 We can gradually add more methods here if needed.310 """311 self.get_scheduler()312