| 488 | # 简化版模型前向封装（无 vace / 无 motion_controller） |
| 489 | # ----------------------------- |
| 490 | def model_fn_wan_video( |
| 491 | dit: WanModel, |
| 492 | x: torch.Tensor, |
| 493 | timestep: torch.Tensor, |
| 494 | context: torch.Tensor, |
| 495 | tea_cache: Optional[TeaCache] = None, |
| 496 | use_unified_sequence_parallel: bool = False, |
| 497 | LQ_latents: Optional[torch.Tensor] = None, |
| 498 | is_full_block: bool = False, |
| 499 | is_stream: bool = False, |
| 500 | pre_cache_k: Optional[list[torch.Tensor]] = None, |
| 501 | pre_cache_v: Optional[list[torch.Tensor]] = None, |
| 502 | topk_ratio: float = 2.0, |
| 503 | kv_ratio: float = 3.0, |
| 504 | cur_process_idx: int = 0, |
| 505 | t_mod : torch.Tensor = None, |
| 506 | t : torch.Tensor = None, |
| 507 | local_range: int = 9, |
| 508 | **kwargs, |
| 509 | ): |
| 510 | # patchify |
| 511 | x, (f, h, w) = dit.patchify(x) |
| 512 | |
| 513 | win = (2, 8, 8) |
| 514 | seqlen = f // win[0] |
| 515 | local_num = seqlen |
| 516 | window_size = win[0] * h * w // 128 |
| 517 | square_num = window_size * window_size |
| 518 | topk = int(square_num * topk_ratio) - 1 |
| 519 | kv_len = int(kv_ratio) |
| 520 | |
| 521 | # RoPE 位置（分段） |
| 522 | if cur_process_idx == 0: |
| 523 | freqs = torch.cat([ |
| 524 | dit.freqs[0][:f].view(f, 1, 1, -1).expand(f, h, w, -1), |
| 525 | dit.freqs[1][:h].view(1, h, 1, -1).expand(f, h, w, -1), |
| 526 | dit.freqs[2][:w].view(1, 1, w, -1).expand(f, h, w, -1) |
| 527 | ], dim=-1).reshape(f * h * w, 1, -1).to(x.device) |
| 528 | else: |
| 529 | freqs = torch.cat([ |
| 530 | dit.freqs[0][4 + cur_process_idx*2:4 + cur_process_idx*2 + f].view(f, 1, 1, -1).expand(f, h, w, -1), |
| 531 | dit.freqs[1][:h].view(1, h, 1, -1).expand(f, h, w, -1), |
| 532 | dit.freqs[2][:w].view(1, 1, w, -1).expand(f, h, w, -1) |
| 533 | ], dim=-1).reshape(f * h * w, 1, -1).to(x.device) |
| 534 | |
| 535 | # TeaCache（默认不启用） |
| 536 | tea_cache_update = tea_cache.check(dit, x, t_mod) if tea_cache is not None else False |
| 537 | |
| 538 | # 统一序列并行（此处默认关闭） |
| 539 | if use_unified_sequence_parallel: |
| 540 | import torch.distributed as dist |
| 541 | from xfuser.core.distributed import (get_sequence_parallel_rank, |
| 542 | get_sequence_parallel_world_size, |
| 543 | get_sp_group) |
| 544 | if dist.is_initialized() and dist.get_world_size() > 1: |
| 545 | x = torch.chunk(x, get_sequence_parallel_world_size(), dim=1)[get_sequence_parallel_rank()] |
| 546 | |
| 547 | # Block 堆叠 |