(
self,
x,
timestep,
context,
kv_cache: list[torch.Tensor],
y=None,
clip_feature=None,
action=None,
timestep_action=None,
state=None,
)
| 678 | self.engine = Engine(eng_path) |
| 679 | |
| 680 | def forward( |
| 681 | self, |
| 682 | x, |
| 683 | timestep, |
| 684 | context, |
| 685 | kv_cache: list[torch.Tensor], |
| 686 | y=None, |
| 687 | clip_feature=None, |
| 688 | action=None, |
| 689 | timestep_action=None, |
| 690 | state=None, |
| 691 | ): |
| 692 | |
| 693 | kv_cache_packed = torch.stack(kv_cache, dim=0) |
| 694 | |
| 695 | self.engine.set_runtime_tensor_shape("x", x.shape) |
| 696 | self.engine.set_runtime_tensor_shape("timestep", timestep.shape) |
| 697 | self.engine.set_runtime_tensor_shape("context", context.shape) |
| 698 | self.engine.set_runtime_tensor_shape("kv_cache_packed", kv_cache_packed.shape) |
| 699 | self.engine.set_runtime_tensor_shape("y", y.shape) |
| 700 | self.engine.set_runtime_tensor_shape("clip_feature", clip_feature.shape) |
| 701 | self.engine.set_runtime_tensor_shape("action", action.shape) |
| 702 | self.engine.set_runtime_tensor_shape("timestep_action", timestep_action.shape) |
| 703 | self.engine.set_runtime_tensor_shape("state", state.shape) |
| 704 | |
| 705 | |
| 706 | output = self.engine( |
| 707 | x.to(torch.float16), |
| 708 | timestep.to(torch.float16), |
| 709 | context.to(torch.float16), |
| 710 | kv_cache_packed.to(torch.float16), |
| 711 | y.to(torch.float16), |
| 712 | clip_feature.to(torch.float16), |
| 713 | action.to(torch.float16), |
| 714 | timestep_action.to(torch.float16), |
| 715 | state.to(torch.float16), |
| 716 | ) |
| 717 | |
| 718 | if "out.0" in output: # for nvfp4 model export through modelopt |
| 719 | return output["out.0"].to(torch.bfloat16).contiguous(), output["out.1"].to(torch.bfloat16).contiguous() |
| 720 | else: |
| 721 | return output["video_noise_pred"].to(torch.bfloat16).contiguous(), output["action_noise_pred"].to(torch.bfloat16).contiguous() |
| 722 | |
| 723 | def load_tensorrt_engine(engine_path="tensorrt/wan_model.trt", model_type="5B"): |
| 724 | """Load TensorRT engine""" |
nothing calls this directly
no test coverage detected