Fabio
c36510c21e
Some checks failed
Linux Build / Linux cpu (push) Has been cancelled
Linux Build / Linux cpu-clang (push) Has been cancelled
Linux Build / Linux vulkan (push) Has been cancelled
mac-build / macOS CPU (push) Has been cancelled
Nix Build / Nix (cpu) (push) Has been cancelled
Nix Build / Nix (vulkan) (push) Has been cancelled
windows-build / Windows CPU (push) Has been cancelled
87 lines
4 KiB
Python
87 lines
4 KiB
Python
"""Prepare LatentSync's original checkpoints for the native GGUF converter."""
|
|
|
|
import argparse
|
|
import json
|
|
import subprocess
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
import torch
|
|
import cv2
|
|
import numpy as np
|
|
from omegaconf import OmegaConf
|
|
from safetensors import safe_open
|
|
from safetensors.torch import load_file, save_file
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument("--checkpoint-dir", type=Path, required=True)
|
|
parser.add_argument("--vae-dir", type=Path, required=True)
|
|
parser.add_argument("--reference-dir", type=Path, required=True)
|
|
parser.add_argument("--face-model-dir", type=Path, required=True)
|
|
parser.add_argument("--output-dir", type=Path, required=True)
|
|
args = parser.parse_args()
|
|
args.output_dir.mkdir(parents=True, exist_ok=True)
|
|
index = {}
|
|
|
|
checkpoint = torch.load(args.checkpoint_dir / "latentsync_unet.pt",
|
|
map_location="cpu", weights_only=True, mmap=True)
|
|
unet = {"unet." + key: value.contiguous() for key, value in checkpoint["state_dict"].items()}
|
|
save_file(unet, args.output_dir / "unet.safetensors")
|
|
index.update({key: "unet.safetensors" for key in unet})
|
|
del unet, checkpoint
|
|
|
|
checkpoint = torch.load(args.checkpoint_dir / "whisper/tiny.pt",
|
|
map_location="cpu", weights_only=True)
|
|
whisper_config = checkpoint["dims"]
|
|
whisper = {"whisper." + key: value.contiguous()
|
|
for key, value in checkpoint["model_state_dict"].items()
|
|
if key.startswith("encoder.")}
|
|
save_file(whisper, args.output_dir / "whisper.safetensors")
|
|
index.update({key: "whisper.safetensors" for key in whisper})
|
|
del whisper, checkpoint
|
|
|
|
vae = {"vae." + key: value.contiguous()
|
|
for key, value in load_file(args.vae_dir / "diffusion_pytorch_model.safetensors").items()}
|
|
save_file(vae, args.output_dir / "vae.safetensors")
|
|
index.update({key: "vae.safetensors" for key in vae})
|
|
del vae
|
|
|
|
for script, checkpoint, directory, config_name in [
|
|
("prepare_scrfd_weights.py", "det_10g.onnx", "detector", "detector_config.json"),
|
|
("prepare_landmark_weights.py", "2d106det.onnx", "landmarks", "landmark_config.json"),
|
|
]:
|
|
destination = args.output_dir / directory
|
|
subprocess.run([sys.executable, str(Path(__file__).with_name(script)),
|
|
"--checkpoint", str(args.face_model_dir / checkpoint),
|
|
"--output", str(destination)], check=True)
|
|
(args.output_dir / config_name).write_bytes((destination / "config.json").read_bytes())
|
|
shard = f"{directory}.safetensors"
|
|
(destination / "weights.safetensors").replace(args.output_dir / shard)
|
|
with safe_open(args.output_dir / shard, framework="pt") as weights:
|
|
index.update({key: shard for key in weights.keys()})
|
|
|
|
mask = cv2.imread(str(args.reference_dir / "latentsync/utils/mask.png"))
|
|
if mask is None:
|
|
raise ValueError("Cannot read LatentSync's face mask")
|
|
mask = cv2.cvtColor(cv2.resize(mask, (512, 512), interpolation=cv2.INTER_LANCZOS4), cv2.COLOR_BGR2RGB)
|
|
mask = np.ascontiguousarray(mask.transpose(2, 0, 1), dtype=np.float32) / 255
|
|
save_file({"face.mask": torch.from_numpy(mask)}, args.output_dir / "mask.safetensors")
|
|
index["face.mask"] = "mask.safetensors"
|
|
|
|
config = {
|
|
"model_type": "latentsync",
|
|
"unet": OmegaConf.to_container(OmegaConf.load(args.reference_dir / "configs/unet/stage2_512.yaml"))["model"],
|
|
"whisper": whisper_config,
|
|
"vae": json.loads((args.vae_dir / "config.json").read_text()),
|
|
"scheduler": json.loads((args.reference_dir / "configs/scheduler_config.json").read_text()),
|
|
}
|
|
(args.output_dir / "config.json").write_text(json.dumps(config, indent=2) + "\n")
|
|
(args.output_dir / "model.safetensors.index.json").write_text(
|
|
json.dumps({"weight_map": index}, indent=2) + "\n")
|
|
print(f"Prepared {len(index)} tensors without dtype conversion")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|