add files

2026-04-23 10:28:20 -04:00
parent 529fe41273
commit dfbf6377f4
10 changed files with 1700 additions and 3 deletions
@@ -0,0 +1,44 @@
+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)