(
self,
x: torch.Tensor,
action: torch.Tensor,
timestep: torch.Tensor,
context: torch.Tensor,
state: torch.Tensor,
embodiment_id: torch.Tensor,
clip_feature: torch.Tensor,
y: torch.Tensor,
)
| 591 | self.engine = Engine(eng_path) |
| 592 | |
| 593 | def forward( |
| 594 | self, |
| 595 | x: torch.Tensor, |
| 596 | action: torch.Tensor, |
| 597 | timestep: torch.Tensor, |
| 598 | context: torch.Tensor, |
| 599 | state: torch.Tensor, |
| 600 | embodiment_id: torch.Tensor, |
| 601 | clip_feature: torch.Tensor, |
| 602 | y: torch.Tensor, |
| 603 | ): |
| 604 | |
| 605 | self.engine.set_runtime_tensor_shape("x", x.shape) |
| 606 | self.engine.set_runtime_tensor_shape("action", action.shape) |
| 607 | self.engine.set_runtime_tensor_shape("context", context.shape) |
| 608 | self.engine.set_runtime_tensor_shape("state", state.shape) |
| 609 | self.engine.set_runtime_tensor_shape("clip_feature", clip_feature.shape) |
| 610 | self.engine.set_runtime_tensor_shape("y", y.shape) |
| 611 | |
| 612 | output = self.engine( |
| 613 | x=x.to(torch.float16), |
| 614 | action=action.to(torch.float16), |
| 615 | timestep=timestep.to(torch.float16), |
| 616 | context=context.to(torch.float16), |
| 617 | state=state.to(torch.float16), |
| 618 | embodiment_id=embodiment_id.to(torch.int32), |
| 619 | clip_feature=clip_feature.to(torch.float16), |
| 620 | y=y.to(torch.float16), |
| 621 | ) |
| 622 | if "out.0" in output: # for nvfp4 model export through modelopt |
| 623 | return output["out.0"].to(torch.bfloat16).contiguous(), output["out.1"].to(torch.bfloat16).contiguous() |
| 624 | else: |
| 625 | return output["video_noise_pred"].to(torch.bfloat16).contiguous(), output["action_noise_pred"].to(torch.bfloat16).contiguous() |
| 626 | |
| 627 | |
| 628 | class WanTrtModelAr5B(torch.nn.Module): |
nothing calls this directly
no test coverage detected