# SPDX-License-Identifier: Apache-2.0
"""Repack a Viggle Remake adapter from diffusers PEFT layout into ComfyUI's LoRA layout.

    python to_comfyui.py lora/swap/pytorch_lora_weights.safetensors viggle_remake_swap_comfyui.safetensors

Same delta, different packing (names, fused qkv, SwiGLU half order). Tensors keep their dtype.
"""
import sys

import torch
from safetensors.torch import load_file, save_file

BLOCKS, QKV = 50, 7168


def convert(src):
    w = load_file(src)
    r = w["proj_in.lora_A.weight"].shape[0]
    out = {}

    def put(name, A, B, rank):
        out[f"diffusion_model.{name}.lora_A.weight"] = A.contiguous()
        out[f"diffusion_model.{name}.lora_B.weight"] = B.contiguous()
        out[f"diffusion_model.{name}.alpha"] = torch.tensor(float(rank))

    put("video_patch_proj", w["proj_in.lora_A.weight"], w["proj_in.lora_B.weight"], r)
    put("final_layer.video_out", w["proj_out.lora_A.weight"], w["proj_out.lora_B.weight"], r)
    for i in range(BLOCKS):
        p, q = f"transformer_blocks.{i}", f"blocks.{i}"
        A = torch.cat([w[f"{p}.attn.to_{x}.lora_A.weight"] for x in "qkv"])
        B = A.new_zeros(3 * QKV, 3 * r)
        for j, x in enumerate("qkv"):
            B[j * QKV:(j + 1) * QKV, j * r:(j + 1) * r] = w[f"{p}.attn.to_{x}.lora_B.weight"]
        put(f"{q}.attn.qkv_proj", A, B, 3 * r)
        put(f"{q}.attn.out_proj", w[f"{p}.attn.to_out.0.lora_A.weight"], w[f"{p}.attn.to_out.0.lora_B.weight"], r)
        value, gate = w[f"{p}.ff.net.0.proj.lora_B.weight"].chunk(2)
        put(f"{q}.mlp.fc1", w[f"{p}.ff.net.0.proj.lora_A.weight"], torch.cat([gate, value]), r)
        put(f"{q}.mlp.fc2", w[f"{p}.ff.net.2.lora_A.weight"], w[f"{p}.ff.net.2.lora_B.weight"], r)
    assert len(out) == 3 * (2 + 4 * BLOCKS)
    return out


if __name__ == "__main__":
    out = convert(sys.argv[1])
    save_file(out, sys.argv[2], metadata={"format": "pt"})
    print(f"{len(out)} tensors -> {sys.argv[2]}")