45 lines
1.3 KiB
Python
45 lines
1.3 KiB
Python
import sys
|
|
import os
|
|
sys.path.append(os.path.dirname(os.path.abspath(__file__)))
|
|
import torch
|
|
import time
|
|
from trt_utilities import Engine
|
|
import tensorrt
|
|
|
|
print("Tensorrt version: ", tensorrt.__version__)
|
|
|
|
def export_trt(trt_path=None, onnx_path=None, use_fp16=True):
|
|
if trt_path is None:
|
|
trt_path = input(
|
|
"Enter the path to save the TensorRT engine (e.g ./realesrgan.engine): ")
|
|
if onnx_path is None:
|
|
onnx_path = input(
|
|
"Enter the path to the ONNX model (e.g ./realesrgan.onnx): ")
|
|
|
|
engine = Engine(trt_path)
|
|
|
|
torch.cuda.empty_cache()
|
|
|
|
s = time.time()
|
|
ret = engine.build(
|
|
onnx_path,
|
|
use_fp16,
|
|
enable_preview=True,
|
|
input_profile=[
|
|
# any sizes from 256x256 to 3840x3840, batch size 1
|
|
{
|
|
"img0": [(1, 3, 256, 256), (1, 3, 512, 512), (1, 3, 3840, 3840)],
|
|
"img1": [(1, 3, 256, 256), (1, 3, 512, 512), (1, 3, 3840, 3840)],
|
|
},
|
|
],
|
|
)
|
|
e = time.time()
|
|
print(f"Time taken to build: {(e-s)} seconds")
|
|
print(f"Tensorrt engine saved at: {trt_path}")
|
|
|
|
return ret
|
|
|
|
|
|
export_trt(trt_path="./models/rife49_ensemble_True_scale_1_sim.engine",
|
|
onnx_path="./models/rife49_ensemble_True_scale_1_sim.onnx", use_fp16=True)
|