This repository was archived by the owner on Dec 26, 2025. It is now read-only.
forked from pschroedl/StreamDiffusion
-
Notifications
You must be signed in to change notification settings - Fork 10
Expand file tree
/
Copy pathcompile_raft_tensorrt.py
More file actions
302 lines (249 loc) · 10.9 KB
/
compile_raft_tensorrt.py
File metadata and controls
302 lines (249 loc) · 10.9 KB
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
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
import torch
import logging
from pathlib import Path
from typing import Optional
import fire
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)
try:
import tensorrt as trt
TENSORRT_AVAILABLE = True
except ImportError:
TENSORRT_AVAILABLE = False
logger.error("TensorRT not available. Please install it first.")
try:
from torchvision.models.optical_flow import raft_small, Raft_Small_Weights
TORCHVISION_AVAILABLE = True
except ImportError:
TORCHVISION_AVAILABLE = False
logger.error("torchvision not available. Please install it first.")
def export_raft_to_onnx(
onnx_path: Path,
min_height: int = 512,
min_width: int = 512,
max_height: int = 512,
max_width: int = 512,
device: str = "cuda"
) -> bool:
"""
Export RAFT model to ONNX format
Args:
onnx_path: Path to save the ONNX model
min_height: Minimum input height for the model
min_width: Minimum input width for the model
max_height: Maximum input height for the model
max_width: Maximum input width for the model
device: Device to use for export
Returns:
True if successful, False otherwise
"""
if not TORCHVISION_AVAILABLE:
logger.error("torchvision is required but not installed")
return False
logger.info(f"Exporting RAFT model to ONNX: {onnx_path}")
logger.info(f"Resolution range: {min_height}x{min_width} - {max_height}x{max_width}")
try:
# Load RAFT model
logger.info("Loading RAFT Small model...")
raft_model = raft_small(weights=Raft_Small_Weights.DEFAULT, progress=True)
raft_model = raft_model.to(device=device)
raft_model.eval()
# Create dummy inputs using max resolution for export
dummy_frame1 = torch.randn(1, 3, max_height, max_width).to(device)
dummy_frame2 = torch.randn(1, 3, max_height, max_width).to(device)
# Apply RAFT preprocessing if available
weights = Raft_Small_Weights.DEFAULT
if hasattr(weights, 'transforms') and weights.transforms is not None:
transforms = weights.transforms()
dummy_frame1, dummy_frame2 = transforms(dummy_frame1, dummy_frame2)
# Make batch, height, and width dimensions dynamic
dynamic_axes = {
"frame1": {0: "batch_size", 2: "height", 3: "width"},
"frame2": {0: "batch_size", 2: "height", 3: "width"},
"flow": {0: "batch_size", 2: "height", 3: "width"},
}
logger.info("Exporting to ONNX...")
with torch.no_grad():
torch.onnx.export(
raft_model,
(dummy_frame1, dummy_frame2),
str(onnx_path),
verbose=False,
input_names=['frame1', 'frame2'],
output_names=['flow'],
opset_version=17,
export_params=True,
dynamic_axes=dynamic_axes,
dynamo=False,
)
del raft_model
torch.cuda.empty_cache()
logger.info(f"Successfully exported ONNX model to {onnx_path}")
return True
except Exception as e:
logger.error(f"Failed to export ONNX model: {e}")
import traceback
traceback.print_exc()
return False
def build_tensorrt_engine(
onnx_path: Path,
engine_path: Path,
min_height: int = 512,
min_width: int = 512,
max_height: int = 512,
max_width: int = 512,
fp16: bool = True,
workspace_size_gb: int = 4
) -> bool:
"""
Build TensorRT engine from ONNX model
Args:
onnx_path: Path to the ONNX model
engine_path: Path to save the TensorRT engine
min_height: Minimum input height for optimization
min_width: Minimum input width for optimization
max_height: Maximum input height for optimization
max_width: Maximum input width for optimization
fp16: Enable FP16 precision mode
workspace_size_gb: Maximum workspace size in GB
Returns:
True if successful, False otherwise
"""
if not TENSORRT_AVAILABLE:
logger.error("TensorRT is required but not installed")
return False
if not onnx_path.exists():
logger.error(f"ONNX model not found: {onnx_path}")
return False
logger.info(f"Building TensorRT engine from ONNX model: {onnx_path}")
logger.info(f"Output path: {engine_path}")
logger.info(f"Resolution range: {min_height}x{min_width} - {max_height}x{max_width}")
logger.info(f"FP16 mode: {fp16}")
logger.info("This may take several minutes...")
try:
builder = trt.Builder(trt.Logger(trt.Logger.INFO))
network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
parser = trt.OnnxParser(network, trt.Logger(trt.Logger.WARNING))
logger.info("Parsing ONNX model...")
with open(onnx_path, 'rb') as model:
if not parser.parse(model.read()):
logger.error("Failed to parse ONNX model")
for error in range(parser.num_errors):
logger.error(f"Parser error: {parser.get_error(error)}")
return False
logger.info("Configuring TensorRT builder...")
config = builder.create_builder_config()
config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, workspace_size_gb * (1 << 30))
if fp16:
config.set_flag(trt.BuilderFlag.FP16)
logger.info("FP16 mode enabled")
# Calculate optimal resolution (middle point)
opt_height = (min_height + max_height) // 2
opt_width = (min_width + max_width) // 2
profile = builder.create_optimization_profile()
min_shape = (1, 3, min_height, min_width)
opt_shape = (1, 3, opt_height, opt_width)
max_shape = (1, 3, max_height, max_width)
profile.set_shape("frame1", min_shape, opt_shape, max_shape)
profile.set_shape("frame2", min_shape, opt_shape, max_shape)
config.add_optimization_profile(profile)
logger.info("Building TensorRT engine... (this will take a while)")
engine = builder.build_serialized_network(network, config)
if engine is None:
logger.error("Failed to build TensorRT engine")
return False
logger.info(f"Saving engine to {engine_path}")
engine_path.parent.mkdir(parents=True, exist_ok=True)
with open(engine_path, 'wb') as f:
f.write(engine)
logger.info(f"Successfully built and saved TensorRT engine: {engine_path}")
logger.info(f"Engine size: {engine_path.stat().st_size / (1024*1024):.2f} MB")
# Delete ONNX file after successful engine creation
try:
if onnx_path.exists():
onnx_path.unlink()
logger.info(f"Deleted ONNX file: {onnx_path}")
except Exception as e:
logger.warning(f"Failed to delete ONNX file: {e}")
return True
except Exception as e:
logger.error(f"Failed to build TensorRT engine: {e}")
import traceback
traceback.print_exc()
return False
def compile_raft(
min_resolution: str = "512x512",
max_resolution: str = "512x512",
output_dir: str = "./models/temporal_net",
device: str = "cuda",
fp16: bool = True,
workspace_size_gb: int = 4,
force_rebuild: bool = False
):
"""
Main function to compile RAFT model to TensorRT engine
Args:
min_resolution: Minimum input resolution as "HxW" (e.g., "512x512") (default: "512x512")
max_resolution: Maximum input resolution as "HxW" (e.g., "1024x1024") (default: "512x512")
output_dir: Directory to save the models (default: ./models/temporal_net)
device: Device to use for export (default: cuda)
fp16: Enable FP16 precision mode (default: True)
workspace_size_gb: Maximum workspace size in GB (default: 4)
force_rebuild: Force rebuild even if engine exists (default: False)
"""
if not TENSORRT_AVAILABLE:
logger.error("TensorRT is not available. Please install it first using:")
logger.error(" python -m streamdiffusion.tools.install-tensorrt")
return
if not TORCHVISION_AVAILABLE:
logger.error("torchvision is not available. Please install it first using:")
logger.error(" pip install torchvision")
return
# Parse resolution strings
try:
min_height, min_width = map(int, min_resolution.split('x'))
except:
logger.error(f"Invalid min_resolution format: {min_resolution}. Expected format: HxW (e.g., 512x512)")
return
try:
max_height, max_width = map(int, max_resolution.split('x'))
except:
logger.error(f"Invalid max_resolution format: {max_resolution}. Expected format: HxW (e.g., 1024x1024)")
return
output_path = Path(output_dir)
output_path.mkdir(parents=True, exist_ok=True)
# Add resolution suffix to filenames
onnx_path = output_path / f"raft_small_min_{min_resolution}_max_{max_resolution}.onnx"
engine_path = output_path / f"raft_small_min_{min_resolution}_max_{max_resolution}.engine"
logger.info("="*80)
logger.info("RAFT TensorRT Compilation")
logger.info("="*80)
logger.info(f"Output directory: {output_path.absolute()}")
logger.info(f"Resolution range: {min_resolution} - {max_resolution}")
logger.info(f"ONNX path: {onnx_path}")
logger.info(f"Engine path: {engine_path}")
logger.info("="*80)
if engine_path.exists() and not force_rebuild:
logger.info(f"TensorRT engine already exists: {engine_path}")
logger.info("Use --force_rebuild to rebuild it")
return
if not onnx_path.exists() or force_rebuild:
logger.info("\n[Step 1/2] Exporting RAFT to ONNX...")
if not export_raft_to_onnx(onnx_path, min_height, min_width, max_height, max_width, device):
logger.error("Failed to export ONNX model")
return
else:
logger.info(f"\n[Step 1/2] ONNX model already exists: {onnx_path}")
logger.info("\n[Step 2/2] Building TensorRT engine...")
if not build_tensorrt_engine(onnx_path, engine_path, min_height, min_width, max_height, max_width, fp16, workspace_size_gb):
logger.error("Failed to build TensorRT engine")
return
logger.info("\n" + "="*80)
logger.info("✓ Compilation completed successfully!")
logger.info("="*80)
logger.info(f"Engine path: {engine_path.absolute()}")
logger.info("\nYou can now use this engine in TemporalNetTensorRTPreprocessor:")
logger.info(f' engine_path="{engine_path.absolute()}"')
logger.info("="*80)
if __name__ == "__main__":
fire.Fire(compile_raft)