This commit is contained in:
discountchubbs
2026-09-05 10:44:03 -07:00
parent ee881434c1
commit 377e0d5156
@@ -157,7 +157,10 @@ def make_warp(nv12: NV12Frame, model_w: int, model_h: int):
def warp(tfm, big_tfm, frame, big_frame):
tfm = tfm.to(Device.DEFAULT)
big_tfm = big_tfm.to(Device.DEFAULT)
Tensor.realize(tfm, big_tfm)
if Device.DEFAULT == 'AMD':
frame = frame.to(Device.DEFAULT)
big_frame = big_frame.to(Device.DEFAULT)
Tensor.realize(tfm, big_tfm, frame, big_frame)
warped_frame = frame_prepare(frame, tfm).unsqueeze(0)
warped_big_frame = frame_prepare(big_frame, big_tfm).unsqueeze(0)
@@ -372,7 +375,8 @@ if __name__ == "__main__":
for cam_w, cam_h in args.camera_resolutions:
print(f"Compiling warp JIT for {cam_w}x{cam_h}...")
nv12 = NV12Frame(cam_w, cam_h, *get_nv12_info(cam_w, cam_h))
make_random_warp_inputs = partial(make_random_images, keys=['frame', 'big_frame'], shape=nv12.size, device=Device.DEFAULT)
warp_input_dev = 'NPY' if Device.DEFAULT == 'AMD' else Device.DEFAULT
make_random_warp_inputs = partial(make_random_images, keys=['frame', 'big_frame'], shape=nv12.size, device=warp_input_dev)
warp = TinyJit(make_warp(nv12, model_w, model_h), prune=True)
output_data[(cam_w, cam_h)] = compile_jit(warp, make_random_warp_inputs, WARP_INPUTS, make_warp_queues)