[Python] How can we make torch.frombuffer support torch.compile?

I have a python function like below:

def custom_rgb2tensor(binary):
t_image = torch.frombuffer(binary,dtype=torch.uint8)
img = cv_imdecode(t_image, cv2.IMREAD_COLOR)
img_tensor = torch.from_numpy(img[:, :, ::-1].copy()).unsqueeze(0) / 255.0
return img_tensor


I tried many ways to make it compilable with torch.compile, but got error below:

File "/root/butterfly/ComfyUI/env_wl/lib/python3.11/site-packages/torch/fx/proxy.py", line 298, in create_arg
raise NotImplementedError(f"argument of type: {type(a)}")
torch._dynamo.exc.InternalTorchDynamoError: argument of type: <class 'bytes'>

from user code:
File "/root/butterfly/ComfyUI/./custom_nodes/ComfyUI-Optimus/compilable/op_compilable.py", line 73, in custom_rgb2tensor
t_image = torch.frombuffer(binary,dtype=torch.uint8)


Does anyone have a good idea?

If I remove the torch.compile annotation, it works fine.

Continue reading...