File size: 1,789 Bytes
4fcca49 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 | from __future__ import annotations
import argparse, statistics, time
from pathlib import Path
import torch
from huggingface_hub import hf_hub_download
import sys
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from src.ort_dit import OrtDitModule
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--model", default=None, help="Local dit_fp16.onnx path")
ap.add_argument("--runs", type=int, default=50)
ap.add_argument("--warmup", type=int, default=10)
ap.add_argument("--no-trt", action="store_true")
args = ap.parse_args()
path = args.model or hf_hub_download("patdev/NitroGen-RTX2060-ONNX", "onnx/dit_fp16.onnx")
device = "cuda" if torch.cuda.is_available() else "cpu"
mod = OrtDitModule(path, prefer_tensorrt=not args.no_trt).to(device)
print(f"device={device} provider={mod.provider}")
if torch.cuda.is_available():
print("gpu=", torch.cuda.get_device_name(0))
h = torch.randn(1, 18, 1024, device=device, dtype=torch.float16)
e = torch.randn(1, 256, 1024, device=device, dtype=torch.float16)
t = torch.tensor([500], device=device, dtype=torch.int64)
for _ in range(args.warmup):
mod(h, e, t)
if torch.cuda.is_available(): torch.cuda.synchronize()
times=[]
for _ in range(args.runs):
t0=time.perf_counter(); mod(h,e,t)
if torch.cuda.is_available(): torch.cuda.synchronize()
times.append((time.perf_counter()-t0)*1000)
med=statistics.median(times); p95=sorted(times)[max(0,int(len(times)*.95)-1)]
print(f"DiT median={med:.2f} ms p95={p95:.2f} ms")
for steps in (4,8,16):
print(f"steps={steps:2d}: DiT-only estimate {1000/(med*steps):.2f} action-chunks/s (vision/head overhead excluded)")
if __name__ == "__main__": main()
|