Issue loading the 600M 512px SANA model

[frutiemax92]

I've tried to load the SANA 512px model but it doesn't seem to work.
Here is the workflow:
SANA512px.json

The error log:
model_type FLOW !!! Exception during processing !!! Error(s) in loading state_dict for SanaMS: size mismatch for pos_embed: copying a param with shape torch.Size([1, 256, 1152]) from checkpoint, the shape in current model is torch.Size([1, 1024, 1152]). Traceback (most recent call last): File "/home/lucas/Documents/ComfyUI/execution.py", line 328, 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 "/home/lucas/Documents/ComfyUI/execution.py", line 203, 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 "/home/lucas/Documents/ComfyUI/execution.py", line 174, in _map_node_over_list process_inputs(input_dict, i) File "/home/lucas/Documents/ComfyUI/execution.py", line 163, in process_inputs results.append(getattr(obj, func)(**inputs)) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/home/lucas/Documents/ComfyUI/custom_nodes/ComfyUI_ExtraModels/Sana/nodes.py", line 33, in load_checkpoint model = load_sana( ^^^^^^^^^^ File "/home/lucas/Documents/ComfyUI/custom_nodes/ComfyUI_ExtraModels/Sana/loader.py", line 88, in load_sana m, u = model.diffusion_model.load_state_dict(state_dict, strict=False) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/home/lucas/Documents/ComfyUI/myenv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 2584, in load_state_dict raise RuntimeError( RuntimeError: Error(s) in loading state_dict for SanaMS: size mismatch for pos_embed: copying a param with shape torch.Size([1, 256, 1152]) from checkpoint, the shape in current model is torch.Size([1, 1024, 1152]).