Mismatch shape when applying and merging a locon model to a SanaTransformer2DModel module

[frutiemax92]

I have trained a locon model based on Efficient-Large-Model/Sana_1600M_1024px_diffusers that I want to apply and merge.
https://huggingface.co/docs/diffusers/en/api/models/sana_transformer2d

Here is a zero initialized rank 4 locon model (untrained)
0.zip

I'm getting a mismatch size error (expecting [11200, 11200, 3, 3] but got [11200, 1, 3, 3]:

import torch
from diffusers import SanaTransformer2
from lycoris import create_lycoris_from_weights
step = '0'
transformer = SanaTransformer2DModel.from_pretrained("Efficient-Large-Model/Sana_1600M_1024px_diffusers", subfolder='transformer').to(dtype=torch.float16)
transformer.train(False)
with torch.no_grad():
    lycoris_net, weights = create_lycoris_from_weights(1.0, '0.safetensors', transformer)
    lycoris_net.train(False)
    for lora in lycoris_net.loras:
        lora = lora.to(dtype=torch.float16, device='cuda:1')
    lycoris_net.apply_to()
    lycoris_net.merge_to()

pipe = SanaPipeline.from_pretrained(
  "Efficient-Large-Model/Sana_1600M_1024px_diffusers",
  torch_dtype=torch.float16,
  variant='fp16',
  transformer=transformer,
  vae = None
)