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
47 lines
1.9 KiB
Python
47 lines
1.9 KiB
Python
"""Export deterministic SD VAE encoder/decoder references from the official demo."""
|
|
|
|
import argparse
|
|
import shutil
|
|
from pathlib import Path
|
|
|
|
import cv2
|
|
import torch
|
|
from diffusers import AutoencoderKL
|
|
from safetensors.torch import save_file
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument("--vae-dir", type=Path, required=True)
|
|
parser.add_argument("--native-weights", type=Path, required=True)
|
|
parser.add_argument("--video", type=Path, required=True)
|
|
parser.add_argument("--output", type=Path, required=True)
|
|
args = parser.parse_args()
|
|
capture = cv2.VideoCapture(str(args.video))
|
|
success, frame = capture.read()
|
|
capture.release()
|
|
if not success:
|
|
raise RuntimeError("Could not read the reference video")
|
|
frame = cv2.resize(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB), (512, 512))
|
|
image = torch.from_numpy(frame).permute(2, 0, 1).unsqueeze(0).float().cuda() / 127.5 - 1
|
|
vae = AutoencoderKL.from_pretrained(str(args.vae_dir), torch_dtype=torch.float32).cuda().eval()
|
|
with torch.inference_mode():
|
|
moments = vae.encode(image).latent_dist.parameters
|
|
latent = moments[:, :4].contiguous()
|
|
decoded = vae.decode(latent).sample
|
|
for component, input_value, output_value in (
|
|
("vae-encoder", image, moments), ("vae-decoder", latent, decoded)
|
|
):
|
|
directory = args.output / component
|
|
directory.mkdir(parents=True, exist_ok=True)
|
|
shutil.copyfile(args.vae_dir / "config.json", directory / "config.json")
|
|
link = directory / "weights.safetensors"
|
|
if not link.exists():
|
|
link.symlink_to(args.native_weights.resolve())
|
|
save_file({"input": input_value.cpu().contiguous(), "output": output_value.cpu().contiguous()},
|
|
directory / "reference.safetensors")
|
|
print(f"{component}: {tuple(input_value.shape)} -> {tuple(output_value.shape)}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|