(
self,
x: torch.Tensor,
timesteps: torch.Tensor,
context: torch.Tensor,
*args,
**kwargs,
)
| 46 | self.engine = None |
| 47 | |
| 48 | def forward( |
| 49 | self, |
| 50 | x: torch.Tensor, |
| 51 | timesteps: torch.Tensor, |
| 52 | context: torch.Tensor, |
| 53 | *args, |
| 54 | **kwargs, |
| 55 | ) -> torch.Tensor: |
| 56 | nvtx.range_push("forward") |
| 57 | feed_dict = { |
| 58 | "sample": x.float(), |
| 59 | "timesteps": timesteps.float(), |
| 60 | "encoder_hidden_states": context.float(), |
| 61 | } |
| 62 | if "y" in kwargs: |
| 63 | feed_dict["y"] = kwargs["y"].float() |
| 64 | |
| 65 | tmp = torch.empty( |
| 66 | self.engine_vram_req, dtype=torch.uint8, device=devices.device |
| 67 | ) |
| 68 | self.engine.context.device_memory = tmp.data_ptr() |
| 69 | self.cudaStream = torch.cuda.current_stream().cuda_stream |
| 70 | self.engine.allocate_buffers(feed_dict) |
| 71 | |
| 72 | out = self.engine.infer(feed_dict, self.cudaStream)["latent"] |
| 73 | |
| 74 | nvtx.range_pop() |
| 75 | return out |
| 76 | |
| 77 | def apply_loras(self, refit_dict: dict): |
| 78 | if not self.refitted_keys.issubset(set(refit_dict.keys())): |
no test coverage detected