prernaaa12/abot-world-interactive
0
1import types2from pathlib import Path3from typing import List, Optional4import torch5from torch import nn6from safetensors.torch import load_file as load_safetensors_file7 8from utils.scheduler import SchedulerInterface, FlowMatchScheduler9import os10 11 12def _wan_models_path(*parts) -> str:13 """Resolve wan_models path relative to project root (works with symlink and any cwd)."""14 root = Path(__file__).resolve().parent.parent15 return str((root / "wan_models").joinpath(Path(*parts)))16 17 18def _resolve_wan_path(path: str) -> str:19 """If path starts with wan_models/, resolve to absolute path (project root); else return as-is."""20 if path and path.startswith("wan_models/"):21 return _wan_models_path(path[len("wan_models/"):])22 return path23 24 25def _resolve_wan_path_with_dir(path: str, wan_models_dir: Optional[str] = None) -> str:26 """Resolve path: if wan_models_dir is set and path starts with wan_models/, use wan_models_dir as base; else _resolve_wan_path."""27 if not path:28 return path29 if wan_models_dir and path.startswith("wan_models/"):30 return os.path.join(wan_models_dir, path[len("wan_models/"):])31 return _resolve_wan_path(path)32 33 34def model_kwargs_with_relative_rope(args, default: bool = False) -> dict:35 """Merge top-level use_relative_rope into model_kwargs with a stable default."""36 raw_model_kwargs = getattr(args, "model_kwargs", {}) or {}37 model_kwargs = dict(raw_model_kwargs)38 if "use_relative_rope" not in model_kwargs:39 try:40 model_kwargs["use_relative_rope"] = bool(getattr(args, "use_relative_rope"))41 except Exception:42 model_kwargs["use_relative_rope"] = bool(default)43 return model_kwargs44 45from wan.modules.tokenizers import HuggingfaceTokenizer46from wan.modules.model import WanModel47from wan.modules.t5 import umt5_xxl48from wan.modules.causal_model import CausalWanModel49 50class WanTextEncoder(torch.nn.Module):51 def __init__(52 self,53 tokenizer_path="wan_models/Wan2.1-T2V-1.3B/google/umt5-xxl/",54 encoder_pth_path="wan_models/Wan2.1-T2V-1.3B/models_t5_umt5-xxl-enc-bf16.pth",55 ) -> None:56 super().__init__()57 # tokenizer_path = _resolve_wan_path_with_dir(tokenizer_path, wan_models_dir)58 # encoder_pth_path = _resolve_wan_path_with_dir(encoder_pth_path, wan_models_dir)59 60 self.text_encoder = umt5_xxl(61 encoder_only=True,62 return_tokenizer=False,63 dtype=torch.bfloat16,64 device=torch.device('cpu')65 ).eval().requires_grad_(False)66 state_dict = torch.load(encoder_pth_path,67 map_location='cpu', weights_only=False)68 self.text_encoder.load_state_dict(state_dict)69 del state_dict70 71 self.tokenizer = HuggingfaceTokenizer(72 name=tokenizer_path, seq_len=512, clean='whitespace')73 74 @property75 def device(self):76 return next(self.text_encoder.parameters()).device77 78 def forward(self, text_prompts: List[str], device: torch.device = None) -> dict:79 ids, mask = self.tokenizer(80 text_prompts, return_mask=True, add_special_tokens=True)81 # When DynamicSwapInstaller is active, self.device returns cpu because82 # parameters are swapped to GPU only during forward. Use the explicitly83 # passed device (the intended execution device) when available.84 target_device = device if device is not None else self.device85 ids = ids.to(target_device)86 mask = mask.to(target_device)87 seq_lens = mask.gt(0).sum(dim=1).long()88 context = self.text_encoder(ids, mask)89 90 for u, v in zip(context, seq_lens):91 u[v:] = 0.0 # set padding to 0.092 93 return {94 "prompt_embeds": context95 }96 97 98class WanVAEWrapper(torch.nn.Module):99 def __init__(100 self,101 pretrained_path=None,102 z_dim=48,103 vae_type="Wan2.2_VAE",104 wan_models_dir=None,105 ):106 super().__init__()107 if vae_type != "Wan2.2_VAE":108 raise ValueError(f"Unsupported vae_type={vae_type!r}; only 'Wan2.2_VAE' is supported.")109 from wan.modules.vae2_2 import _video_vae110 self.mean = torch.tensor([111 -0.2289, -0.0052, -0.1323, -0.2339, -0.2799, 0.0174, 0.1838, 0.1557,112 -0.1382, 0.0542, 0.2813, 0.0891, 0.1570, -0.0098, 0.0375, -0.1825,113 -0.2246, -0.1207, -0.0698, 0.5109, 0.2665, -0.2108, -0.2158, 0.2502,114 -0.2055, -0.0322, 0.1109, 0.1567, -0.0729, 0.0899, -0.2799, -0.1230,115 -0.0313, -0.1649, 0.0117, 0.0723, -0.2839, -0.2083, -0.0520, 0.3748,116 0.0152, 0.1957, 0.1433, -0.2944, 0.3573, -0.0548, -0.1681, -0.0667,117 ], dtype=torch.float32)118 119 self.std = torch.tensor([120 0.4765, 1.0364, 0.4514, 1.1677, 0.5313, 0.4990, 0.4818, 0.5013,121 0.8158, 1.0344, 0.5894, 1.0901, 0.6885, 0.6165, 0.8454, 0.4978,122 0.5759, 0.3523, 0.7135, 0.6804, 0.5833, 1.4146, 0.8986, 0.5659,123 0.7069, 0.5338, 0.4889, 0.4917, 0.4069, 0.4999, 0.6866, 0.4093,124 0.5709, 0.6065, 0.6415, 0.4944, 0.5726, 1.2042, 0.5458, 1.6887,125 0.3971, 1.0600, 0.3943, 0.5537, 0.5444, 0.4089, 0.7468, 0.7744,126 ], dtype=torch.float32)127 self.scale = [self.mean, 1.0 / self.std]128 self.upsampling_factor = 16129 130 z_dim = 48131 self.z_dim = z_dim132 self.model = _video_vae(pretrained_path=pretrained_path,133 z_dim=z_dim,).eval().requires_grad_(False)134 135 def generate_noise(self, shape, seed=None, rand_device="cpu", rand_torch_dtype=torch.float32, device=None, torch_dtype=None):136 # Initialize Gaussian noise137 generator = None if seed is None else torch.Generator(rand_device).manual_seed(seed)138 noise = torch.randn(shape, generator=generator, device=rand_device, dtype=rand_torch_dtype)139 noise = noise.to(dtype=torch_dtype, device=device)140 return noise141 142 def encode_to_latent(self, pixel: torch.Tensor) -> torch.Tensor:143 # pixel: [batch_size, num_channels, num_frames, height, width]144 device, dtype = pixel.device, pixel.dtype145 scale = [self.mean.to(device=device, dtype=dtype),146 1.0 / self.std.to(device=device, dtype=dtype)]147 148 output = [149 self.model.encode(u.unsqueeze(0), scale).float().squeeze(0)150 for u in pixel151 ]152 output = torch.stack(output, dim=0)153 # from [batch_size, num_channels, num_frames, height, width]154 # to [batch_size, num_frames, num_channels, height, width]155 output = output.permute(0, 2, 1, 3, 4)156 return output157 158 def decode_to_pixel(self, latent: torch.Tensor, use_cache: bool = False, return_in_cpu: bool = False) -> torch.Tensor:159 # from [batch_size, num_frames, num_channels, height, width]160 # to [batch_size, num_channels, num_frames, height, width]161 zs = latent.permute(0, 2, 1, 3, 4)162 if use_cache:163 assert latent.shape[0] == 1, "Batch size must be 1 when using cache"164 165 device, dtype = latent.device, latent.dtype166 scale = [self.mean.to(device=device, dtype=dtype),167 1.0 / self.std.to(device=device, dtype=dtype)]168 169 if use_cache:170 decode_function = self.model.cached_decode171 else:172 decode_function = self.model.decode173 174 output = []175 for u in zs:176 decoded = decode_function(u.unsqueeze(0), scale).float().clamp_(-1, 1).squeeze(0)177 if return_in_cpu:178 decoded = decoded.cpu()179 output.append(decoded)180 output = torch.stack(output, dim=0)181 # from [batch_size, num_channels, num_frames, height, width]182 # to [batch_size, num_frames, num_channels, height, width]183 output = output.permute(0, 2, 1, 3, 4)184 return output185 186 187class TAEW2_2VAEWrapper(torch.nn.Module):188 """189 VAE wrapper using TAEHV (TAEW2.2) for faster decoding.190 Requires: pip install taehv (or install from https://github.com/madebyollin/taehv)191 Checkpoint: taew2_2.pth (download from taehv releases)192 """193 def __init__(self, checkpoint_path: str = "taew2_2.pth", dtype=torch.float16):194 super().__init__()195 try:196 from wan.modules.taehv import TAEHV, StreamingTAEHV197 except ImportError as e:198 raise ImportError(199 "taehv is required for TAEW2.2 VAE. Install with: pip install taehv"200 ) from e201 self.taehv = TAEHV(checkpoint_path).to(dtype).eval().requires_grad_(False)202 self.taehv = StreamingTAEHV(self.taehv)203 self.dtype = dtype204 # For compatibility with pipeline.vae.model.clear_cache()205 self.model = _TAEW2_2ModelRef(self)206 207 def warmup_first_frame(self, first_frame_latent: torch.Tensor):208 """Warm up the streaming decoder's MemBlock memory with the first-frame latent.209 210 The TAeW2.2 MemBlocks use zero-initialized past context for the first frame,211 causing blur. By feeding the first-frame latent as a warmup pass (output212 discarded), subsequent decodes benefit from real temporal context.213 214 Args:215 first_frame_latent: [B, 1, C, H, W] latent of the first frame.216 """217 if first_frame_latent is None:218 return219 # Reset decoder state, then feed first frame as warmup220 self.taehv.reset()221 with torch.no_grad(), torch.autocast(device_type="cuda", dtype=self.dtype):222 # Feed the first-frame latent to populate MemBlock memory;223 # the output (startup frames) is discarded.224 _ = self.taehv.decode(first_frame_latent)225 226 def decode_to_pixel(227 self,228 latent: torch.Tensor,229 use_cache: bool = False,230 return_in_cpu: bool = False231 ) -> torch.Tensor:232 # latent: [B, F, C, H, W] = [B, T, C, H, W] (same as TAEHV's NTCHW)233 # use_cache=True -> parallel=False for lower memory (streaming)234 parallel = not use_cache235 with torch.autocast(device_type="cuda", dtype=self.dtype):236 # out = self.taehv.decode_video(237 # latent, parallel=parallel, show_progress_bar=False238 # )239 out = self.taehv.decode(latent)240 # TAEHV returns [0, 1], convert to [-1, 1] to match WanVAEWrapper241 out = out.mul(2).sub(1).clamp(-1, 1).float()242 if return_in_cpu:243 out = out.cpu()244 return out245 246 247class _TAEW2_2ModelRef:248 """Dummy ref for clear_cache compatibility; delegates to StreamingTAEHV.reset()."""249 250 def __init__(self, parent):251 self._parent = parent252 253 def clear_cache(self):254 self._parent.taehv.reset()255 256 257class MGLightVAEWrapper(torch.nn.Module):258 """VAE wrapper using MG-LightVAE (pruned Wan2.2 VAE) for faster decoding.259 260 Wraps the ``Wan2_2_VAE`` class which supports different pruning rates.261 The encoder uses the full (unpruned) Wan2.2 VAE teacher, while the decoder262 uses the pruned student model.263 264 Args:265 vae_pth: Path to the pruned LightVAE checkpoint (student decoder).266 lightvae_pruning_rate: Pruning rate for the decoder (e.g. 0.5, 0.75).267 lightvae_encoder_vae_pth: Path to the full Wan2.2 VAE checkpoint268 (teacher encoder). Required for mg_lightvae.269 dtype: Data type for the VAE model.270 device: Device to load the VAE on.271 """272 273 def __init__(274 self,275 vae_pth: str,276 lightvae_pruning_rate: float = 0.75,277 lightvae_encoder_vae_pth: str | None = None,278 dtype=torch.float,279 device="cpu",280 ):281 super().__init__()282 from wan.modules.vae2_2 import Wan2_2_VAE283 284 self._vae = Wan2_2_VAE(285 z_dim=48,286 c_dim=160,287 vae_pth=vae_pth,288 dtype=dtype,289 device=device,290 vae_type="mg_lightvae",291 lightvae_pruning_rate=lightvae_pruning_rate,292 lightvae_encoder_vae_pth=lightvae_encoder_vae_pth,293 )294 # Register model (pruned decoder) and encoder_model (teacher encoder)295 # as submodules so that .to(), .eval(), .requires_grad_() propagate.296 self.model = self._vae.model297 if self._vae.encoder_model is not None:298 self.encoder_model = self._vae.encoder_model299 else:300 self.encoder_model = None301 302 self.z_dim = 48303 self.upsampling_factor = 16304 self.mean = self._vae.scale[0]305 self.std = 1.0 / self._vae.scale[1]306 307 # Initialize streaming cache attributes (_feat_map, _conv_idx, etc.)308 # so cached_decode() can be called before any explicit clear_cache().309 self.model.clear_cache()310 if self.encoder_model is not None:311 self.encoder_model.clear_cache()312 313 def generate_noise(self, shape, seed=None, rand_device="cpu",314 rand_torch_dtype=torch.float32, device=None, torch_dtype=None):315 generator = None if seed is None else torch.Generator(rand_device).manual_seed(seed)316 noise = torch.randn(shape, generator=generator, device=rand_device, dtype=rand_torch_dtype)317 noise = noise.to(dtype=torch_dtype, device=device)318 return noise319 320 def encode_to_latent(self, pixel: torch.Tensor) -> torch.Tensor:321 # pixel: [batch_size, num_channels, num_frames, height, width]322 device, dtype = pixel.device, pixel.dtype323 scale = [self.mean.to(device=device, dtype=dtype),324 1.0 / self.std.to(device=device, dtype=dtype)]325 326 encode_model = self.encoder_model if self.encoder_model is not None else self.model327 output = [328 encode_model.encode(u.unsqueeze(0), scale).float().squeeze(0)329 for u in pixel330 ]331 output = torch.stack(output, dim=0)332 # from [B, C, F, H, W] to [B, F, C, H, W]333 output = output.permute(0, 2, 1, 3, 4)334 return output335 336 def decode_to_pixel(self, latent: torch.Tensor, use_cache: bool = False,337 return_in_cpu: bool = False) -> torch.Tensor:338 # from [B, F, C, H, W] to [B, C, F, H, W]339 zs = latent.permute(0, 2, 1, 3, 4)340 if use_cache:341 assert latent.shape[0] == 1, "Batch size must be 1 when using cache"342 343 device, dtype = latent.device, latent.dtype344 scale = [self.mean.to(device=device, dtype=dtype),345 1.0 / self.std.to(device=device, dtype=dtype)]346 347 if use_cache:348 decode_function = self.model.cached_decode349 else:350 decode_function = self.model.decode351 352 output = []353 for u in zs:354 decoded = decode_function(u.unsqueeze(0), scale).float().clamp_(-1, 1).squeeze(0)355 if return_in_cpu:356 decoded = decoded.cpu()357 output.append(decoded)358 output = torch.stack(output, dim=0)359 # from [B, C, F, H, W] to [B, F, C, H, W]360 output = output.permute(0, 2, 1, 3, 4)361 return output362 363 364def create_vae_from_config(config) -> Optional[torch.nn.Module]:365 """Create a VAE wrapper based on the unified ``vae_type`` config field.366 367 Supported vae_type values:368 - "wan2.2": Standard Wan2.2 VAE (returns None, pipeline creates WanVAEWrapper)369 - "taew2_2": TAeW2.2/taehv fast streaming decoder370 - "mg_lightvae": MG-LightVAE pruned decoder (pruning rate 0.5)371 - "mg_lightvae_v2": MG-LightVAE v2 pruned decoder (pruning rate 0.75)372 373 If ``vae_type`` is not set, defaults to "taew2_2".374 375 Returns:376 A VAE wrapper instance, or None for wan2.2 (let pipeline create377 WanVAEWrapper from vae_kwargs).378 """379 vae_type = getattr(config, "vae_type", None)380 381 if vae_type is None:382 vae_type = "taew2_2"383 else:384 vae_type = str(vae_type).strip().lower()385 386 if vae_type == "wan2.2":387 return None # pipeline creates WanVAEWrapper(**vae_kwargs) internally388 389 if vae_type == "taew2_2":390 ckpt = os.environ.get("TAEW2_2_CHECKPOINT") or getattr(391 config, "taew2_2_checkpoint", "taew2_2.pth"392 )393 return TAEW2_2VAEWrapper(checkpoint_path=ckpt).eval()394 395 if vae_type in ("mg_lightvae", "mg_lightvae_v2"):396 pruning_map = {"mg_lightvae": 0.5, "mg_lightvae_v2": 0.75}397 # Explicit pruning rate overrides the default mapping398 explicit_rate = getattr(config, "lightvae_pruning_rate", None)399 if explicit_rate is not None:400 pruning_rate = float(explicit_rate)401 else:402 pruning_rate = pruning_map[vae_type]403 404 # Select checkpoint based on vae_type405 ckpt_map = {406 "mg_lightvae": "lightvae_checkpoint",407 "mg_lightvae_v2": "lightvae_v2_checkpoint",408 }409 vae_ckpt = getattr(config, ckpt_map[vae_type], None)410 if vae_ckpt is None:411 raise ValueError(412 f"vae_type={vae_type!r} requires '{ckpt_map[vae_type]}' config field "413 f"(path to MG-LightVAE .pth file)."414 )415 416 # Encoder checkpoint: explicit config, or fall back to vae_kwargs.pretrained_path417 encoder_ckpt = getattr(config, "lightvae_encoder_checkpoint", None)418 if encoder_ckpt is None:419 vae_kwargs = getattr(config, "vae_kwargs", {}) or {}420 if isinstance(vae_kwargs, dict):421 encoder_ckpt = vae_kwargs.get("pretrained_path")422 else:423 encoder_ckpt = getattr(vae_kwargs, "pretrained_path", None)424 if encoder_ckpt is None:425 raise ValueError(426 f"vae_type={vae_type!r} requires 'lightvae_encoder_checkpoint' config field "427 f"(path to full Wan2.2_VAE.pth for teacher encoder), "428 f"or 'vae_kwargs.pretrained_path' must be set."429 )430 431 return MGLightVAEWrapper(432 vae_pth=vae_ckpt,433 lightvae_pruning_rate=pruning_rate,434 lightvae_encoder_vae_pth=encoder_ckpt,435 )436 437 raise ValueError(438 f"Unsupported vae_type={vae_type!r}. "439 f"Choose from: wan2.2, taew2_2, mg_lightvae, mg_lightvae_v2."440 )441 442 443class WanDiffusionWrapper(torch.nn.Module):444 @staticmethod445 def _materialize_meta_tensors(module: torch.nn.Module, device: torch.device = torch.device("cpu")):446 materialized_names = []447 448 def _materialize_recursive(mod: torch.nn.Module, prefix: str = ""):449 for name, param in list(mod.named_parameters(recurse=False)):450 if getattr(param, "is_meta", False):451 new_param = torch.nn.Parameter(452 torch.empty(tuple(param.shape), dtype=param.dtype, device=device),453 requires_grad=param.requires_grad,454 )455 setattr(mod, name, new_param)456 materialized_names.append(prefix + name)457 458 for name, buf in list(mod.named_buffers(recurse=False)):459 if getattr(buf, "is_meta", False):460 setattr(mod, name, torch.empty(tuple(buf.shape), dtype=buf.dtype, device=device))461 materialized_names.append(prefix + name)462 463 for child_name, child in mod.named_children():464 _materialize_recursive(child, prefix + child_name + ".")465 466 _materialize_recursive(module)467 return materialized_names468 469 @staticmethod470 def _normalize_model_state_dict_keys(state_dict: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:471 key_prefixes = (472 "generator.model._fsdp_wrapped_module.",473 "generator.model.",474 "model._fsdp_wrapped_module.",475 "model.",476 "_fsdp_wrapped_module.",477 "module.",478 )479 keys = list(state_dict.keys())480 for prefix in key_prefixes:481 if keys and all(k.startswith(prefix) for k in keys):482 return {k[len(prefix):]: v for k, v in state_dict.items()}483 return state_dict484 485 def _load_model_safetensors(self, model_safetensors_path: str) -> None:486 if model_safetensors_path.startswith("oss://"):487 raise ValueError(488 "model_safetensors_path must be a local mounted path on AI-Hub, "489 f"got {model_safetensors_path}"490 )491 492 model_safetensors_path = _resolve_wan_path(model_safetensors_path)493 print(f"[WanDiffusionWrapper] Loading model safetensors from {model_safetensors_path}")494 state_dict = load_safetensors_file(model_safetensors_path, device="cpu")495 state_dict = self._normalize_model_state_dict_keys(state_dict)496 497 model_keys = set(self.model.state_dict().keys())498 matched_keys = model_keys.intersection(state_dict.keys())499 if not matched_keys:500 sample_keys = list(state_dict.keys())[:10]501 raise ValueError(502 "No safetensors keys matched the Wan model state_dict. "503 f"First loaded keys: {sample_keys}"504 )505 506 match_ratio = len(matched_keys) / max(1, len(state_dict))507 if match_ratio < 0.5:508 sample_unexpected = [k for k in state_dict.keys() if k not in model_keys][:10]509 raise ValueError(510 f"Only {len(matched_keys)}/{len(state_dict)} safetensors keys match the Wan model "511 f"state_dict after prefix normalization. Sample unexpected keys: {sample_unexpected}"512 )513 514 missing, unexpected = self.model.load_state_dict(state_dict, strict=False)515 if missing:516 print(517 f"[WanDiffusionWrapper] model_safetensors missing {len(missing)} keys "518 f"(showing first 20): {missing[:20]}"519 )520 if unexpected:521 print(522 f"[WanDiffusionWrapper] model_safetensors unexpected {len(unexpected)} keys "523 f"(showing first 20): {unexpected[:20]}"524 )525 print(526 f"[WanDiffusionWrapper] Loaded model safetensors with {len(matched_keys)} "527 f"matched keys from {model_safetensors_path}"528 )529 530 def __init__(531 self,532 model_name="Wan2.1-T2V-1.3B",533 timestep_shift=8.0,534 is_causal=False,535 local_attn_size=-1,536 sink_size=0,537 subfolder=None,538 model_type='t2v',539 num_frame_per_block=3,540 model_safetensors_path: Optional[str] = None,541 **model_init_kwargs,542 ):543 super().__init__()544 545 self.model_type = model_type546 use_relative_rope = bool(model_init_kwargs.pop("use_relative_rope", False))547 if is_causal:548 model_init_kwargs["use_relative_rope"] = use_relative_rope549 self.model = CausalWanModel.from_pretrained(550 model_name, local_attn_size=local_attn_size, sink_size=sink_size, model_type=model_type, num_frame_per_block=num_frame_per_block,551 **model_init_kwargs)552 else:553 if use_relative_rope:554 print("[WanDiffusionWrapper] use_relative_rope is ignored for non-causal WanModel.")555 self.model = WanModel.from_pretrained(model_name, model_type=model_type, **model_init_kwargs)556 materialized = self._materialize_meta_tensors(self.model, device=torch.device("cpu"))557 if materialized:558 print(f"[WanDiffusionWrapper] Materialized {len(materialized)} meta tensors on CPU.")559 if model_safetensors_path:560 self._load_model_safetensors(model_safetensors_path)561 self.model.eval()562 563 # For non-causal diffusion, all frames share the same timestep564 self.uniform_timestep = not is_causal565 566 self.scheduler = FlowMatchScheduler(567 shift=timestep_shift, sigma_min=0.0, extra_one_step=True568 )569 self.scheduler.set_timesteps(1000, training=True)570 571 self.seq_len = None # [1, 21, 16, 60, 104]572 self.post_init()573 574 def enable_gradient_checkpointing(self) -> None:575 self.model.enable_gradient_checkpointing()576 577 def adding_cls_branch(self, atten_dim=1536, num_class=4, time_embed_dim=0) -> None:578 # NOTE: This is hard coded for WAN2.1-T2V-1.3B for now!!!!!!!!!!!!!!!!!!!!579 self._cls_pred_branch = nn.Sequential(580 # Input: [B, 384, 21, 60, 104]581 nn.LayerNorm(atten_dim * 3 + time_embed_dim),582 nn.Linear(atten_dim * 3 + time_embed_dim, 1536),583 nn.SiLU(),584 nn.Linear(atten_dim, num_class)585 )586 self._cls_pred_branch.requires_grad_(True)587 num_registers = 3588 self._register_tokens = RegisterTokens(num_registers=num_registers, dim=atten_dim)589 self._register_tokens.requires_grad_(True)590 591 gan_ca_blocks = []592 for _ in range(num_registers):593 block = GanAttentionBlock()594 gan_ca_blocks.append(block)595 self._gan_ca_blocks = nn.ModuleList(gan_ca_blocks)596 self._gan_ca_blocks.requires_grad_(True)597 # self.has_cls_branch = True598 599 def _convert_flow_pred_to_x0(self, flow_pred: torch.Tensor, xt: torch.Tensor, timestep: torch.Tensor) -> torch.Tensor:600 """601 Convert flow matching's prediction to x0 prediction.602 flow_pred: the prediction with shape [B, C, H, W]603 xt: the input noisy data with shape [B, C, H, W]604 timestep: the timestep with shape [B]605 606 pred = noise - x0607 x_t = (1-sigma_t) * x0 + sigma_t * noise608 we have x0 = x_t - sigma_t * pred609 see derivations https://chatgpt.com/share/67bf8589-3d04-8008-bc6e-4cf1a24e2d0e610 """611 # use higher precision for calculations612 original_dtype = flow_pred.dtype613 flow_pred, xt, sigmas, timesteps = map(614 lambda x: x.double().to(flow_pred.device), [flow_pred, xt,615 self.scheduler.sigmas,616 self.scheduler.timesteps]617 )618 619 timestep_id = torch.argmin(620 (timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1)621 sigma_t = sigmas[timestep_id].reshape(-1, 1, 1, 1)622 x0_pred = xt - sigma_t * flow_pred623 return x0_pred.to(original_dtype)624 625 @staticmethod626 def _convert_x0_to_flow_pred(scheduler, x0_pred: torch.Tensor, xt: torch.Tensor, timestep: torch.Tensor) -> torch.Tensor:627 """628 Convert x0 prediction to flow matching's prediction.629 x0_pred: the x0 prediction with shape [B, C, H, W]630 xt: the input noisy data with shape [B, C, H, W]631 timestep: the timestep with shape [B]632 633 pred = (x_t - x_0) / sigma_t634 """635 # use higher precision for calculations636 original_dtype = x0_pred.dtype637 x0_pred, xt, sigmas, timesteps = map(638 lambda x: x.double().to(x0_pred.device), [x0_pred, xt,639 scheduler.sigmas,640 scheduler.timesteps]641 )642 timestep_id = torch.argmin(643 (timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1)644 sigma_t = sigmas[timestep_id].reshape(-1, 1, 1, 1)645 flow_pred = (xt - x0_pred) / sigma_t646 return flow_pred.to(original_dtype)647 648 @staticmethod649 def _history_x_to_model_format(history_x):650 if history_x is None:651 return None652 653 if torch.is_tensor(history_x):654 if history_x.ndim != 5:655 raise ValueError(656 f"history_x must be [B,F,C,H,W] when passed as a tensor, got {history_x.shape}"657 )658 return [u.permute(1, 0, 2, 3).contiguous() for u in history_x]659 660 return history_x661 662 @staticmethod663 def _history_condition_to_model_format(value):664 if value is None:665 return None666 667 if torch.is_tensor(value):668 if value.ndim < 3:669 return value670 return [u.contiguous() for u in value]671 672 return value673 674 def forward(675 self,676 noisy_image_or_video: torch.Tensor, conditional_dict: dict,677 timestep: torch.Tensor, kv_cache: Optional[List[dict]] = None,678 crossattn_cache: Optional[List[dict]] = None,679 current_start: Optional[int] = None,680 681 classify_mode: Optional[bool] = False, # DF682 concat_time_embeddings: Optional[bool] = False, #DF683 clean_x: Optional[torch.Tensor] = None, # TF684 aug_t: Optional[torch.Tensor] = None, # for TF clean GT, if it's also noisy and needs denoising by the model, aug_t is its timestep685 686 cache_start: Optional[int] = None,687 updating_cache: Optional[bool] = False,688 replace_first_timestep_and_noise_latents: Optional[bool] = False,689 history_x: Optional[torch.Tensor] = None,690 history_y: Optional[torch.Tensor] = None,691 history_act_context: Optional[torch.Tensor] = None,692 history_y_action: Optional[torch.Tensor] = None,693 noisy_start_frame: int = 0,694 695 ) -> torch.Tensor:696 prompt_embeds = conditional_dict["prompt_embeds"]697 act_context = conditional_dict.get("act_context", None)698 act_context_scale = conditional_dict.get("act_context_scale", 1.0)699 clip_fea = conditional_dict.get("clip_fea", None)700 y = conditional_dict.get("y", None)701 y_action = conditional_dict.get("y_action", None)702 ref_latents = conditional_dict.get("ref_latents", None)703 ref_mask = conditional_dict.get("ref_mask", None)704 # first_frame_latents = conditional_dict.get("first_frame_latents", None)705 706 raw_timestep = timestep707 b, f, c, h, w = noisy_image_or_video.shape708 709 if replace_first_timestep_and_noise_latents:710 # Wan2.2 5B uses the first latent frame as a clean condition. Keep711 # per-frame timesteps for score models so only frame 0 is forced to t=0.712 if raw_timestep.dim() == 2:713 input_timestep = raw_timestep.clone()714 input_timestep[:, 0] = 0715 elif raw_timestep.dim() == 1 and raw_timestep.shape[0] == f:716 input_timestep = raw_timestep.unsqueeze(0).repeat(b, 1)717 input_timestep[:, 0] = 0718 elif raw_timestep.dim() == 1 and raw_timestep.shape[0] == b:719 input_timestep = raw_timestep[:, None].repeat(1, f)720 input_timestep[:, 0] = 0721 else:722 input_timestep = raw_timestep.reshape(-1)[0].view(1, 1).repeat(b, f)723 input_timestep[:, 0] = 0724 elif self.uniform_timestep:725 # [B, F] -> [B] for legacy non-causal uniform score models.726 input_timestep = raw_timestep[:, 0]727 else:728 input_timestep = raw_timestep729 730 logits = None731 if history_x is not None:732 history_kwargs = {733 "history_x": self._history_x_to_model_format(history_x),734 "noisy_start_frame": int(noisy_start_frame),735 }736 history_y = self._history_condition_to_model_format(history_y)737 history_act_context = self._history_condition_to_model_format(history_act_context)738 history_y_action = self._history_condition_to_model_format(history_y_action)739 740 if history_y is not None:741 history_kwargs["history_y"] = history_y742 if history_act_context is not None:743 history_kwargs["history_act_context"] = history_act_context744 if history_y_action is not None:745 history_kwargs["history_y_action"] = history_y_action746 747 model_out = self.model(748 noisy_image_or_video.permute(0, 2, 1, 3, 4),749 t=input_timestep,750 context=prompt_embeds,751 seq_len=self.seq_len,752 kv_cache=None,753 crossattn_cache=crossattn_cache,754 current_start=0 if current_start is None else current_start,755 cache_start=0 if cache_start is None else cache_start,756 act_context=act_context,757 y_action=y_action,758 act_context_scale=act_context_scale,759 clip_fea=clip_fea,760 y=y,761 ref_latents=ref_latents,762 ref_mask=ref_mask,763 **history_kwargs,764 )765 if isinstance(model_out, tuple):766 flow_pred = model_out[0]767 else:768 flow_pred = model_out769 flow_pred = flow_pred.permute(0, 2, 1, 3, 4)770 # X0 prediction771 elif kv_cache is not None:772 kwargs = {}773 if updating_cache:774 kwargs["updating_cache"] = updating_cache775 flow_pred = self.model(776 noisy_image_or_video.permute(0, 2, 1, 3, 4), # => [B, C, F, H, W],777 t=input_timestep, context=prompt_embeds,778 seq_len=self.seq_len,779 kv_cache=kv_cache,780 crossattn_cache=crossattn_cache,781 current_start=current_start,782 cache_start=cache_start,783 act_context=act_context,784 act_context_scale=act_context_scale,785 clip_fea=clip_fea,786 y=y,787 ref_latents=ref_latents,788 ref_mask=ref_mask,789 **kwargs,790 ).permute(0, 2, 1, 3, 4)791 else:792 if clean_x is not None:793 # teacher forcing794 flow_pred = self.model(795 noisy_image_or_video.permute(0, 2, 1, 3, 4), # => [B, C, F, H, W]796 t=input_timestep, context=prompt_embeds,797 seq_len=self.seq_len,798 clean_x=clean_x.permute(0, 2, 1, 3, 4), # => [B, C, F, H, W]799 aug_t=aug_t,800 act_context=act_context,801 act_context_scale=act_context_scale,802 clip_fea=clip_fea,803 y=y,804 ref_latents=ref_latents,805 ref_mask=ref_mask,806 ).permute(0, 2, 1, 3, 4)807 else:808 # diffusion forcing or bidirectional809 if classify_mode:810 flow_pred, logits = self.model(811 noisy_image_or_video.permute(0, 2, 1, 3, 4),812 t=input_timestep, context=prompt_embeds,813 seq_len=self.seq_len,814 classify_mode=True,815 register_tokens=self._register_tokens,816 cls_pred_branch=self._cls_pred_branch,817 gan_ca_blocks=self._gan_ca_blocks,818 concat_time_embeddings=concat_time_embeddings,819 act_context=act_context,820 act_context_scale=act_context_scale,821 clip_fea=clip_fea,822 y=y,823 ref_latents=ref_latents,824 ref_mask=ref_mask,825 )826 flow_pred = flow_pred.permute(0, 2, 1, 3, 4)827 else:828 flow_pred = self.model(829 noisy_image_or_video.permute(0, 2, 1, 3, 4),830 t=input_timestep, context=prompt_embeds,831 seq_len=self.seq_len,832 act_context=act_context,833 act_context_scale=act_context_scale,834 clip_fea=clip_fea,835 y=y,836 ref_latents=ref_latents,837 ref_mask=ref_mask,838 ).permute(0, 2, 1, 3, 4)839 840 pred_x0 = self._convert_flow_pred_to_x0(841 flow_pred=flow_pred.flatten(0, 1),842 xt=noisy_image_or_video.flatten(0, 1),843 timestep=timestep.flatten(0, 1)844 ).unflatten(0, flow_pred.shape[:2])845 846 if logits is not None:847 return flow_pred, pred_x0, logits848 849 return flow_pred, pred_x0850 851 def get_scheduler(self) -> SchedulerInterface:852 """853 Update the current scheduler with the interface's static method854 """855 scheduler = self.scheduler856 scheduler.convert_x0_to_noise = types.MethodType(857 SchedulerInterface.convert_x0_to_noise, scheduler)858 scheduler.convert_noise_to_x0 = types.MethodType(859 SchedulerInterface.convert_noise_to_x0, scheduler)860 scheduler.convert_velocity_to_x0 = types.MethodType(861 SchedulerInterface.convert_velocity_to_x0, scheduler)862 self.scheduler = scheduler863 return scheduler864 865 def post_init(self):866 """867 A few custom initialization steps that should be called after the object is created.868 Currently, the only one we have is to bind a few methods to scheduler.869 We can gradually add more methods here if needed.870 """871 self.get_scheduler()872