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
74 lines
2.9 KiB
Python
74 lines
2.9 KiB
Python
"""Time repeated official inference requests, without precision or RNG overrides."""
|
|
|
|
import argparse
|
|
import json
|
|
import os
|
|
import sys
|
|
import time
|
|
from pathlib import Path
|
|
|
|
import torch
|
|
from accelerate.utils import set_seed
|
|
from omegaconf import OmegaConf
|
|
|
|
|
|
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("--video", type=Path, required=True)
|
|
parser.add_argument("--audio", type=Path, required=True)
|
|
parser.add_argument("--output", type=Path, required=True)
|
|
parser.add_argument("--deepcache", action="store_true")
|
|
parser.add_argument("--repeats", type=int, default=2)
|
|
args = parser.parse_args()
|
|
for name in ("reference", "checkpoint", "video", "audio", "output"):
|
|
setattr(args, name, getattr(args, name).resolve())
|
|
args.output.mkdir(parents=True, exist_ok=True)
|
|
sys.path.insert(0, str(args.reference))
|
|
os.chdir(args.reference)
|
|
from scripts import inference
|
|
|
|
original_call = inference.LipsyncPipeline.__call__
|
|
|
|
def timed_call(pipeline, **kwargs):
|
|
totals = {}
|
|
|
|
def timed(name, function):
|
|
def call(*positional, **keywords):
|
|
torch.cuda.synchronize()
|
|
start = time.perf_counter()
|
|
result = function(*positional, **keywords)
|
|
torch.cuda.synchronize()
|
|
totals[name] = totals.get(name, 0) + (time.perf_counter() - start) * 1000
|
|
return result
|
|
return call
|
|
|
|
pipeline.unet.forward = timed("unet.wall_ms", pipeline.unet.forward)
|
|
pipeline.vae.encode = timed("vae.encode.wall_ms", pipeline.vae.encode)
|
|
pipeline.vae.decode = timed("vae.decode.wall_ms", pipeline.vae.decode)
|
|
for index in range(args.repeats):
|
|
totals.clear()
|
|
set_seed(1247)
|
|
torch.cuda.reset_peak_memory_stats()
|
|
torch.cuda.synchronize()
|
|
start = time.perf_counter()
|
|
kwargs["video_out_path"] = str(args.output / f"request-{index}.mp4")
|
|
result = original_call(pipeline, **kwargs)
|
|
torch.cuda.synchronize()
|
|
totals["session.wall_ms"] = (time.perf_counter() - start) * 1000
|
|
totals["peak_allocated_mib"] = torch.cuda.max_memory_allocated() / 1048576
|
|
print(json.dumps({"request": index, "deepcache": args.deepcache, **totals}), flush=True)
|
|
return result
|
|
|
|
inference.LipsyncPipeline.__call__ = timed_call
|
|
inference.main(OmegaConf.load("configs/unet/stage2_512.yaml"), argparse.Namespace(
|
|
video_path=str(args.video), audio_path=str(args.audio),
|
|
video_out_path=str(args.output / "output.mp4"),
|
|
inference_ckpt_path=str(args.checkpoint), enable_deepcache=args.deepcache,
|
|
seed=1247, inference_steps=20, guidance_scale=1.5,
|
|
temp_dir=str(args.output / "temp")))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|