`JoinImageWithAlpha` node: all input tensors must be on the same device. Received cpu and mps:0

[bigcat88]

Expected Behavior

"Join Image With Alpha" node checks whether the mask and the images is on a different devices and move them to the one device.

ComfyUI itself is running on MPS, no specific flags specified. I expect it can move image to mps in this case in the "Join Image With Alpha" to succeed or move mask to the cpu.

Actual Behavior

Steps to Reproduce

reproduce-workflow.json

Only this custom node is required: https://github.com/viperyl/ComfyUI-BiRefNet

Debug Logs

!!! Exception during processing !!! torch.cat(): all input tensors must be on the same device. Received cpu and mps:0
Traceback (most recent call last):
  File "/Users/shurik/PycharmProjects/ComfyUI/execution.py", line 317, in execute
    output_data, output_ui, has_subgraph = get_output_data(obj, input_data_all, execution_block_cb=execution_block_cb, pre_execute_cb=pre_execute_cb)
                                           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/shurik/PycharmProjects/ComfyUI/execution.py", line 192, in get_output_data
    return_values = _map_node_over_list(obj, input_data_all, obj.FUNCTION, allow_interrupt=True, execution_block_cb=execution_block_cb, pre_execute_cb=pre_execute_cb)
                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/shurik/PycharmProjects/ComfyUI/execution.py", line 169, in _map_node_over_list
    process_inputs(input_dict, i)
  File "/Users/shurik/PycharmProjects/ComfyUI/execution.py", line 158, in process_inputs
    results.append(getattr(obj, func)(**inputs))
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/shurik/PycharmProjects/ComfyUI/comfy_extras/nodes_compositing.py", line 198, in join_image_with_alpha
    out_images.append(torch.cat((image[i][:,:,:3], alpha[i].unsqueeze(2)), dim=2))
                      ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
RuntimeError: torch.cat(): all input tensors must be on the same device. Received cpu and mps:0

Other

No response

[bigcat88]

As a workaround I found way to use "Image To Device" node from ComfyUI_essentials to pass mask to cpu device before passing it to Join Image with Alpha

The downside of it that it requires ComfyUI_essentials and 3 additional nodes, but at least it is working.

note: we are deploying ComfyUI remotely and using it through API