| 361 | |
| 362 | |
| 363 | def model_fn_wan_video( |
| 364 | dit: WanModel, |
| 365 | x: torch.Tensor, |
| 366 | timestep: torch.Tensor, |
| 367 | context: torch.Tensor, |
| 368 | clip_feature: Optional[torch.Tensor] = None, |
| 369 | y: Optional[torch.Tensor] = None, |
| 370 | tea_cache: TeaCache = None, |
| 371 | **kwargs, |
| 372 | ): |
| 373 | t = dit.time_embedding(sinusoidal_embedding_1d(dit.freq_dim, timestep)) |
| 374 | t_mod = dit.time_projection(t).unflatten(1, (6, dit.dim)) |
| 375 | context = dit.text_embedding(context) |
| 376 | |
| 377 | if dit.has_image_input: |
| 378 | x = torch.cat([x, y], dim=1) # (b, c_x + c_y, f, h, w) |
| 379 | clip_embdding = dit.img_emb(clip_feature) |
| 380 | context = torch.cat([clip_embdding, context], dim=1) |
| 381 | |
| 382 | x, (f, h, w) = dit.patchify(x) |
| 383 | |
| 384 | freqs = torch.cat([ |
| 385 | dit.freqs[0][:f].view(f, 1, 1, -1).expand(f, h, w, -1), |
| 386 | dit.freqs[1][:h].view(1, h, 1, -1).expand(f, h, w, -1), |
| 387 | dit.freqs[2][:w].view(1, 1, w, -1).expand(f, h, w, -1) |
| 388 | ], dim=-1).reshape(f * h * w, 1, -1).to(x.device) |
| 389 | |
| 390 | # TeaCache |
| 391 | if tea_cache is not None: |
| 392 | tea_cache_update = tea_cache.check(dit, x, t_mod) |
| 393 | else: |
| 394 | tea_cache_update = False |
| 395 | |
| 396 | if tea_cache_update: |
| 397 | x = tea_cache.update(x) |
| 398 | else: |
| 399 | # blocks |
| 400 | for block in dit.blocks: |
| 401 | x = block(x, context, t_mod, freqs) |
| 402 | if tea_cache is not None: |
| 403 | tea_cache.store(x) |
| 404 | |
| 405 | x = dit.head(x, t) |
| 406 | x = dit.unpatchify(x, (f, h, w)) |
| 407 | return x |