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
51 lines
2.4 KiB
Python
51 lines
2.4 KiB
Python
"""Export LatentSync Whisper conditioning tensors for a controlled component check."""
|
|
|
|
import argparse
|
|
import json
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
import torch
|
|
import numpy as np
|
|
from safetensors.torch import save_file
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument("--reference", type=Path, required=True)
|
|
parser.add_argument("--checkpoint", type=Path, required=True)
|
|
parser.add_argument("--audio", type=Path, required=True)
|
|
parser.add_argument("--output", type=Path, required=True)
|
|
parser.add_argument("--audio-repeat", type=int, default=1)
|
|
args = parser.parse_args()
|
|
sys.path.insert(0, str(args.reference.resolve()))
|
|
from latentsync.whisper.whisper import load_model
|
|
from latentsync.whisper.whisper.audio import load_audio, log_mel_spectrogram, pad_or_trim
|
|
from latentsync.whisper.audio2feature import Audio2Feature
|
|
|
|
args.output.mkdir(parents=True, exist_ok=True)
|
|
model = load_model(str(args.checkpoint), device="cuda").float().eval()
|
|
audio = np.tile(load_audio(str(args.audio)), args.audio_repeat)
|
|
mel = log_mel_spectrogram(pad_or_trim(audio)).unsqueeze(0).cuda()
|
|
with torch.inference_mode():
|
|
_, embeddings = model.encoder(mel, include_embeddings=True)
|
|
segments = model.transcribe(audio, fp16=False)["segments"]
|
|
audio_features = torch.cat([torch.from_numpy(segment["encoder_embeddings"]).permute(0, 2, 1, 3)[0,
|
|
:int((segment["end"] - segment["start"]) / 2)] for segment in segments])
|
|
save_file({"whisper.encoder." + name: value.detach().cpu().contiguous()
|
|
for name, value in model.encoder.state_dict().items()}, args.output / "weights.safetensors")
|
|
features = torch.from_numpy(embeddings).permute(0, 2, 1, 3).contiguous()
|
|
slicer = Audio2Feature.__new__(Audio2Feature)
|
|
slicer.embedding_dim = 384
|
|
slicer.audio_feat_length = [2, 2]
|
|
windows = torch.stack([slicer.get_sliced_feature(features[0], i, 25)[0] for i in range(751)])
|
|
save_file({"audio": torch.from_numpy(audio), "audio_features": audio_features.contiguous(),
|
|
"mel": mel.cpu().contiguous(), "windows": windows.contiguous(),
|
|
"embeddings": features},
|
|
args.output / "reference.safetensors")
|
|
(args.output / "config.json").write_text(json.dumps(vars(model.dims), indent=2) + "\n")
|
|
print(f"conditioning_shape={embeddings.shape}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|