PID flux 2 error while decoding final image

[Fqlox]

Custom Node Testing

• I have tried disabling custom nodes and the issue persists (see how to disable custom nodes if you need help)

Expected Behavior

I wanted to use pid on flux2kllein 9b output but the vae decoding seems to give me an error.

Actual Behavior

[ERROR] !!! Exception during processing !!! The size of tensor a (3) must match the size of tensor b (128) at non-singleton dimension 1
[ERROR] Traceback (most recent call last):
  File "dir\ComfyUI\execution.py", line 536, in execute
    output_data, output_ui, has_subgraph, has_pending_tasks = await get_output_data(prompt_id, unique_id, obj, input_data_all, execution_block_cb=execution_block_cb, pre_execute_cb=pre_execute_cb, v3_data=v3_data)
                                                              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "dir\ComfyUI\execution.py", line 336, in get_output_data
    return_values = await _async_map_node_over_list(prompt_id, unique_id, obj, input_data_all, obj.FUNCTION, allow_interrupt=True, execution_block_cb=execution_block_cb, pre_execute_cb=pre_execute_cb, v3_data=v3_data)
                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "dir\ComfyUI\execution.py", line 310, in _async_map_node_over_list
    await process_inputs(input_dict, i)
  File "dir\ComfyUI\execution.py", line 298, in process_inputs
    result = f(**inputs)
             ^^^^^^^^^^^
  File "dir\ComfyUI\nodes.py", line 315, in decode
    images = vae.decode(latent)
             ^^^^^^^^^^^^^^^^^^
  File "dir\ComfyUI\comfy\sd.py", line 1066, in decode
    model_management.raise_non_oom(e)
  File "dir\ComfyUI\comfy\model_management.py", line 394, in raise_non_oom
    raise e
  File "dir\ComfyUI\comfy\sd.py", line 1059, in decode
    out = self.first_stage_model.decode(samples, **vae_options).to(device=self.output_device, dtype=self.vae_output_dtype(), copy=True)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "dir\ComfyUI\comfy\ldm\models\autoencoder.py", line 244, in decode
    z = z * s + m
        ~~^~~
RuntimeError: The size of tensor a (3) must match the size of tensor b (128) at non-singleton dimension 1

Steps to Reproduce

pidonly.json

Debug Logs

[ERROR] !!! Exception during processing !!! The size of tensor a (3) must match the size of tensor b (128) at non-singleton dimension 1
[ERROR] Traceback (most recent call last):
  File "dir\ComfyUI\execution.py", line 536, in execute
    output_data, output_ui, has_subgraph, has_pending_tasks = await get_output_data(prompt_id, unique_id, obj, input_data_all, execution_block_cb=execution_block_cb, pre_execute_cb=pre_execute_cb, v3_data=v3_data)
                                                              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "dir\ComfyUI\execution.py", line 336, in get_output_data
    return_values = await _async_map_node_over_list(prompt_id, unique_id, obj, input_data_all, obj.FUNCTION, allow_interrupt=True, execution_block_cb=execution_block_cb, pre_execute_cb=pre_execute_cb, v3_data=v3_data)
                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "dir\ComfyUI\execution.py", line 310, in _async_map_node_over_list
    await process_inputs(input_dict, i)
  File "dir\ComfyUI\execution.py", line 298, in process_inputs
    result = f(**inputs)
             ^^^^^^^^^^^
  File "dir\ComfyUI\nodes.py", line 315, in decode
    images = vae.decode(latent)
             ^^^^^^^^^^^^^^^^^^
  File "dir\ComfyUI\comfy\sd.py", line 1066, in decode
    model_management.raise_non_oom(e)
  File "dir\ComfyUI\comfy\model_management.py", line 394, in raise_non_oom
    raise e
  File "dir\ComfyUI\comfy\sd.py", line 1059, in decode
    out = self.first_stage_model.decode(samples, **vae_options).to(device=self.output_device, dtype=self.vae_output_dtype(), copy=True)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "dir\ComfyUI\comfy\ldm\models\autoencoder.py", line 244, in decode
    z = z * s + m
        ~~^~~
RuntimeError: The size of tensor a (3) must match the size of tensor b (128) at non-singleton dimension 1

Other

Using comfyui 0.22.0
frontend 1.44.19

[lin-bot23]

Hi, thanks for the detailed report!

The PiD (Pixel Difference) VAE support for Flux 2 was added in commit 28f4ef2, which is after the 0.22.0 release. The version you're currently using (0.22.0) does not include the necessary changes for handling the Flux 2 VAE's channel structure.

To fix this, please update to the latest commit by pulling the latest changes via git:

git pull origin master

The latest commit includes the PiD VAE decoding fix for Flux 2 models and should resolve the tensor size mismatch error you're seeing.

Let us know if the issue persists after updating!

[Fqlox]

Hi, thanks for the detailed report!

The PiD (Pixel Difference) VAE support for Flux 2 was added in commit 28f4ef2, which is after the 0.22.0 release. The version you're currently using (0.22.0) does not include the necessary changes for handling the Flux 2 VAE's channel structure.

To fix this, please update to the latest commit by pulling the latest changes via git:

git pull origin master

The latest commit includes the PiD VAE decoding fix for Flux 2 models and should resolve the tensor size mismatch error you're seeing.

Let us know if the issue persists after updating!

I did update via git I did update the comfyui kitchen too.
I still have the same issue:

[INFO] Python version: 3.12.10 (tags/v3.12.10:0cc8128, Apr  8 2025, 12:21:36) [MSC v.1943 64 bit (AMD64)]
[INFO] ComfyUI version: 0.22.0
[INFO] comfy-aimdo version: 0.4.5
[INFO] comfy-kitchen version: 0.2.9
[INFO] comfyui-frontend-package version: 1.44.19
[INFO] comfyui-workflow-templates version: 0.9.91
[INFO] comfyui-embedded-docs version: 0.5.1
[INFO] comfy-kitchen version: 0.2.9
[INFO] comfy-aimdo version: 0.4.5

[braintm]

You need to decode using the vae that comes with PiD. It's an error to use flux's vae for this.

[Fqlox]

You need to decode using the vae that comes with PiD. It's an error to use flux's vae for this.

Thanks it doesn't output an error. It sill shows a noisy output while using flux2. I've seen this issue still open I think I will wait till the new implementation