import sys, torch
text, decoder, out = sys.argv[1:4]
merged = torch.load(decoder, map_location="cpu", weights_only=False)
merged = {k: merged[k] for k in ("model", "mpd", "msd")}
merged["model"].update({k: v for k, v in torch.load(text, map_location="cpu", weights_only=True).items() if not k.startswith("decoder.")})
torch.save(merged, out)
print("assembled", round(sum(v.numel() for v in merged["model"].values()) / 1e6, 2), "M tensors")