flux nf4 model loading error

[Simon12062313]

When I load model flux1-dev-bnb-nf4.safetensors and flux1-dev-bnb-nf4-v2.safetensors with nf4 checkpoint loader node, following error will raise:
got prompt
!!! Exception during processing !!! Got [16, 56, 56] but expected positional dim 262144
Traceback (most recent call last):
File "E:\ComfyUI\ComfyUI-aki-v2\ComfyUI\execution.py", line 510, 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, hidden_inputs=hidden_inputs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "E:\ComfyUI\ComfyUI-aki-v2\ComfyUI\execution.py", line 324, 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, hidden_inputs=hidden_inputs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "E:\ComfyUI\ComfyUI-aki-v2\ComfyUI\execution.py", line 298, in async_map_node_over_list
await process_inputs(input_dict, i)
File "E:\ComfyUI\ComfyUI-aki-v2\ComfyUI\execution.py", line 286, in process_inputs
result = f(**inputs)
^^^^^^^^^^^
File "E:\ComfyUI\ComfyUI-aki-v2\ComfyUI\custom_nodes\ComfyUI_bnb_nf4_fp4_Loaders_init.py", line 205, in load_checkpoint
out = comfy.sd.load_checkpoint_guess_config(ckpt_path, output_vae=True, output_clip=True, embedding_directory=folder_paths.get_folder_paths("embeddings"), model_options={"custom_operations": ops})
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "E:\ComfyUI\ComfyUI-aki-v2\ComfyUI\comfy\sd.py", line 1268, in load_checkpoint_guess_config
out = load_state_dict_guess_config(sd, output_vae, output_clip, output_clipvision, embedding_directory, output_model, model_options, te_model_options=te_model_options, metadata=metadata)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "E:\ComfyUI\ComfyUI-aki-v2\ComfyUI\comfy\sd.py", line 1313, in load_state_dict_guess_config
model = model_config.get_model(sd, diffusion_model_prefix, device=inital_load_device)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "E:\ComfyUI\ComfyUI-aki-v2\ComfyUI\comfy\supported_models.py", line 713, in get_model
out = model_base.Flux(self, device=device)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "E:\ComfyUI\ComfyUI-aki-v2\ComfyUI\comfy\model_base.py", line 849, in init
super().init(model_config, model_type, device=device, unet_model=unet_model)
File "E:\ComfyUI\ComfyUI-aki-v2\ComfyUI\comfy\model_base.py", line 140, in init
self.diffusion_model = unet_model(**unet_config, device=device, operations=operations)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "E:\ComfyUI\ComfyUI-aki-v2\ComfyUI\comfy\ldm\flux\model.py", line 63, in init
raise ValueError(f"Got {params.axes_dim} but expected positional dim {pe_dim}")
ValueError: Got [16, 56, 56] but expected positional dim 262144

end_vram - start_vram: 0 - 0 = 0

#32 [CheckpointLoaderNF4]: 0.06s - vram 0b