Error 'layers.0.weight' at vae encode node in mochi edit wf

[yogotatara3]

Hi, I get this error at the load vae node to encode the vid when executing the mochiedit inversion wf:

got prompt
Initializing model...
Loading model state_dict from C:\tools\StabilityMatrix\Packages\ComfyUI\models\diffusion_models\mochi\mochi_preview_dit_bf16.safetensors...
Using accelerate to load and assign model weights to device...
!!! Exception during processing !!! 'layers.0.weight'
Traceback (most recent call last):
File "C:\tools\StabilityMatrix\Packages\ComfyUI\execution.py", line 323, 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 "C:\tools\StabilityMatrix\Packages\ComfyUI\execution.py", line 198, 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 "C:\tools\StabilityMatrix\Packages\ComfyUI\execution.py", line 169, in _map_node_over_list
process_inputs(input_dict, i)
File "C:\tools\StabilityMatrix\Packages\ComfyUI\execution.py", line 158, in process_inputs
results.append(getattr(obj, func)(**inputs))
File "C:\tools\StabilityMatrix\Packages\ComfyUI\custom_nodes\ComfyUI-MochiWrapper\nodes.py", line 408, in loadmodel
set_module_tensor_to_device(encoder, name, dtype=dtype, device=offload_device, value=encoder_sd[name])
KeyError: 'layers.0.weight'

Would be amazing if you could help me soplve this please.
Thank you in Advance

Greets,
Yogo

[kijai]

When did you last update the nodes?

[yogotatara3]

Hi, I updated them and reinstalled them but eventually it doenst help.

[kijai]

Can't see the full name of the VAE you're trying to load, perhaps you're trying to load the decoder model in the encoder node.

[yogotatara3]

Ofc..that solved the problem, thanks a lot for your time