(self,
x: torch.Tensor,
timestep: torch.Tensor,
context: torch.Tensor,
use_gradient_checkpointing: bool = False,
use_gradient_checkpointing_offload: bool = False,
LQ_latents: Optional[List[torch.Tensor]] = None,
train_img: bool = False,
topk_ratio: Optional[float] = None,
kv_ratio: Optional[float] = None,
local_num: Optional[int] = None,
is_full_block: bool = False,
causal_idx: Optional[int] = None,
**kwargs,
)
| 591 | ) |
| 592 | |
| 593 | def forward(self, |
| 594 | x: torch.Tensor, |
| 595 | timestep: torch.Tensor, |
| 596 | context: torch.Tensor, |
| 597 | use_gradient_checkpointing: bool = False, |
| 598 | use_gradient_checkpointing_offload: bool = False, |
| 599 | LQ_latents: Optional[List[torch.Tensor]] = None, |
| 600 | train_img: bool = False, |
| 601 | topk_ratio: Optional[float] = None, |
| 602 | kv_ratio: Optional[float] = None, |
| 603 | local_num: Optional[int] = None, |
| 604 | is_full_block: bool = False, |
| 605 | causal_idx: Optional[int] = None, |
| 606 | **kwargs, |
| 607 | ): |
| 608 | # time / text embeds |
| 609 | t = self.time_embedding( |
| 610 | sinusoidal_embedding_1d(self.freq_dim, timestep)) |
| 611 | t_mod = self.time_projection(t).unflatten(1, (6, self.dim)) |
| 612 | |
| 613 | # 这里仍会嵌入 text（CrossAttention 若已有缓存会忽略它） |
| 614 | # context = self.text_embedding(context) |
| 615 | |
| 616 | # 输入打补丁 |
| 617 | x, (f, h, w) = self.patchify(x) |
| 618 | B = x.shape[0] |
| 619 | |
| 620 | # window / masks 超参 |
| 621 | win = (2, 8, 8) |
| 622 | seqlen = f//win[0] |
| 623 | if local_num is None: |
| 624 | local_random = random.random() |
| 625 | if local_random < 0.3: |
| 626 | local_num = seqlen - 3 |
| 627 | elif local_random < 0.4: |
| 628 | local_num = seqlen - 4 |
| 629 | elif local_random < 0.5: |
| 630 | local_num = seqlen - 2 |
| 631 | else: |
| 632 | local_num = seqlen |
| 633 | |
| 634 | window_size = win[0]*h*w//128 |
| 635 | square_num = window_size*window_size |
| 636 | topk_ratio = 2.0 |
| 637 | topk = min(max(int(square_num*topk_ratio), 1), int(square_num*seqlen)-1) |
| 638 | |
| 639 | if kv_ratio is None: |
| 640 | kv_ratio = (random.uniform(0., 1.0)**2)*(local_num-2-2)+2 |
| 641 | kv_len = min(max(int(window_size*kv_ratio), 1), int(window_size*seqlen)-1) |
| 642 | |
| 643 | decay_ratio = random.uniform(0.7, 1.0) |
| 644 | |
| 645 | # RoPE 3D |
| 646 | freqs = torch.cat([ |
| 647 | self.freqs[0][:f].view(f, 1, 1, -1).expand(f, h, w, -1), |
| 648 | self.freqs[1][:h].view(1, h, 1, -1).expand(f, h, w, -1), |
| 649 | self.freqs[2][:w].view(1, 1, w, -1).expand(f, h, w, -1) |
| 650 | ], dim=-1).reshape(f * h * w, 1, -1).to(x.device) |
nothing calls this directly
no test coverage detected