(
self,
x,
timestep,
context,
kv_cache: list[torch.Tensor],
y=None,
action=None,
timestep_action=None,
state=None,
)
| 632 | self.engine = Engine(eng_path) |
| 633 | |
| 634 | def forward( |
| 635 | self, |
| 636 | x, |
| 637 | timestep, |
| 638 | context, |
| 639 | kv_cache: list[torch.Tensor], |
| 640 | y=None, |
| 641 | action=None, |
| 642 | timestep_action=None, |
| 643 | state=None, |
| 644 | ): |
| 645 | |
| 646 | kv_cache_packed = torch.stack(kv_cache, dim=0) |
| 647 | |
| 648 | self.engine.set_runtime_tensor_shape("x", x.shape) |
| 649 | self.engine.set_runtime_tensor_shape("timestep", timestep.shape) |
| 650 | self.engine.set_runtime_tensor_shape("context", context.shape) |
| 651 | self.engine.set_runtime_tensor_shape("kv_cache_packed", kv_cache_packed.shape) |
| 652 | # self.engine.set_runtime_tensor_shape("y", y.shape) |
| 653 | self.engine.set_runtime_tensor_shape("action", action.shape) |
| 654 | self.engine.set_runtime_tensor_shape("timestep_action", timestep_action.shape) |
| 655 | self.engine.set_runtime_tensor_shape("state", state.shape) |
| 656 | |
| 657 | |
| 658 | output = self.engine( |
| 659 | x.to(torch.float16), |
| 660 | timestep.to(torch.float16), |
| 661 | context.to(torch.float16), |
| 662 | kv_cache_packed.to(torch.float16), |
| 663 | # y.to(torch.float16), |
| 664 | action.to(torch.float16), |
| 665 | timestep_action.to(torch.float16), |
| 666 | state.to(torch.float16), |
| 667 | ) |
| 668 | |
| 669 | if "out.0" in output: # for nvfp4 model export through modelopt |
| 670 | return output["out.0"].to(torch.bfloat16).contiguous(), output["out.1"].to(torch.bfloat16).contiguous() |
| 671 | else: |
| 672 | return output["video_noise_pred"].to(torch.bfloat16).contiguous(), output["action_noise_pred"].to(torch.bfloat16).contiguous() |
| 673 | |
| 674 | |
| 675 | class WanTrtModelAr14B(torch.nn.Module): |
nothing calls this directly
no test coverage detected