| 456 | self.gate = GateModule() |
| 457 | |
| 458 | def forward(self, x, context, t_mod, freqs, f, h, w, local_num=None, topk=None, |
| 459 | train_img=False, block_id=None, kv_len=None, is_full_block=False, |
| 460 | is_stream=False, pre_cache_k=None, pre_cache_v=None, local_range = 9): |
| 461 | shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = ( |
| 462 | self.modulation.to(dtype=t_mod.dtype, device=t_mod.device) + t_mod).chunk(6, dim=1) |
| 463 | input_x = modulate(self.norm1(x), shift_msa, scale_msa) |
| 464 | self_attn_output, self_attn_cache_k, self_attn_cache_v = self.self_attn( |
| 465 | input_x, freqs, f, h, w, local_num, topk, train_img, block_id, |
| 466 | kv_len=kv_len, is_full_block=is_full_block, is_stream=is_stream, |
| 467 | pre_cache_k=pre_cache_k, pre_cache_v=pre_cache_v, local_range = local_range) |
| 468 | |
| 469 | x = self.gate(x, gate_msa, self_attn_output) |
| 470 | x = x + self.cross_attn(self.norm3(x), context, is_stream=is_stream) |
| 471 | input_x = modulate(self.norm2(x), shift_mlp, scale_mlp) |
| 472 | x = self.gate(x, gate_mlp, self.ffn(input_x)) |
| 473 | if is_stream: |
| 474 | return x, self_attn_cache_k, self_attn_cache_v |
| 475 | return x |
| 476 | |
| 477 | |
| 478 | class MLP(torch.nn.Module): |