Добавлены пропсы конвейера и стереодвижки, задействованные в прогоне

assets/conveyors (274 МБ) - ленты и угловая секция NVIDIA, на которые ссылается сцена
относительным путём. Раньше исключались как перекачиваемые, но без них сцена не
композится из коробки.

cv/ - код стереодвижков, которые вызывает control_test, без весов:
* defom-stereo - рабочий бейзлайн (DEFOM vitl, вход 480, iters 24)
* crestereo - второй движок, точнее по габаритам (MAE 23.5 против 32.8 мм)
* fast-foundationstereo - проверялся, в бейзлайн не вошёл
* circular_section.py - показатель кругового сечения, перенесён в measure_plane.py:
  выравнивает облако по СОБСТВЕННЫМ главным осям и режет на пяти высотах вдоль каждой.
  Три самодельные версии (мировые оси, одно сечение) давали хуже; результаты проверки
  на эталонной геометрии - в circular_section_results.json

Веса по-прежнему не в репозитории - источники в MODELS.md. Наборы кадров прежних
прогонов (cv/flow_*, 1.26 ГБ) исключены: это выход, а не исходники.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

cv/fast-foundationstereo/.gitignore

@@ -0,0 +1,14 @@
+__pycache__
+*.vscode
+cfg.yaml
+*.pth
+*.pkl
+*.pyc
+*.so
+*.egg-info
+# Ignore weight files but keep the folder structure
+weights/*
+!weights/.gitkeep
+output/
+.claude/
+cpp/build/

cv/fast-foundationstereo/LICENSE.txt

@@ -0,0 +1,94 @@
+Copyright (c) 2026-Present, NVIDIA Corporation & affiliates. All rights reserved.
+
+
+=======================================================================
+
+1. Definitions
+
+"Licensor" means any person or entity that distributes its Work.
+
+"Software" means the original work of authorship made available under
+this License.
+
+"Work" means the Software and any additions to or derivative works of
+the Software that are made available under this License.
+
+The terms "reproduce," "reproduction," "derivative works," and
+"distribution" have the meaning as provided under U.S. copyright law;
+provided, however, that for the purposes of this License, derivative
+works shall not include works that remain separable from, or merely
+link (or bind by name) to the interfaces of, the Work.
+
+Works, including the Software, are "made available" under this License
+by including in or with the Work either (a) a copyright notice
+referencing the applicability of this License to the Work, or (b) a
+copy of this License.
+
+2. License Grants
+
+    2.1 Copyright Grant. Subject to the terms and conditions of this
+    License, each Licensor grants to you a perpetual, worldwide,
+    non-exclusive, royalty-free, copyright license to reproduce,
+    prepare derivative works of, publicly display, publicly perform,
+    sublicense and distribute its Work and any resulting derivative
+    works in any form.
+
+3. Limitations
+
+    3.1 Redistribution. You may reproduce or distribute the Work only
+    if (a) you do so under this License, (b) you include a complete
+    copy of this License with your distribution, and (c) you retain
+    without modification any copyright, patent, trademark, or
+    attribution notices that are present in the Work.
+
+    3.2 Derivative Works. You may specify that additional or different
+    terms apply to the use, reproduction, and distribution of your
+    derivative works of the Work ("Your Terms") only if (a) Your Terms
+    provide that the use limitation in Section 3.3 applies to your
+    derivative works, and (b) you identify the specific derivative
+    works that are subject to Your Terms. Notwithstanding Your Terms,
+    this License (including the redistribution requirements in Section
+    3.1) will continue to apply to the Work itself.
+
+    3.3 Use Limitation. The Work and any derivative works thereof only
+    may be used or intended for use non-commercially. Notwithstanding
+    the foregoing, NVIDIA and its affiliates may use the Work and any
+    derivative works commercially. As used herein, "non-commercially"
+    means for research purposes only.
+
+    3.4 Patent Claims. If you bring or threaten to bring a patent claim
+    against any Licensor (including any claim, cross-claim or
+    counterclaim in a lawsuit) to enforce any patents that you allege
+    are infringed by any Work, then your rights under this License from
+    such Licensor (including the grant in Section 2.1) will terminate
+    immediately.
+
+    3.5 Trademarks. This License does not grant any rights to use any
+    Licensor�s or its affiliates� names, logos, or trademarks, except
+    as necessary to reproduce the notices described in this License.
+
+    3.6 Termination. If you violate any term of this License, then your
+    rights under this License (including the grant in Section 2.1) will
+    terminate immediately.
+
+4. Disclaimer of Warranty.
+
+THE WORK IS PROVIDED "AS IS" WITHOUT WARRANTIES OR CONDITIONS OF ANY
+KIND, EITHER EXPRESS OR IMPLIED, INCLUDING WARRANTIES OR CONDITIONS OF
+MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE, TITLE OR
+NON-INFRINGEMENT. YOU BEAR THE RISK OF UNDERTAKING ANY ACTIVITIES UNDER
+THIS LICENSE.
+
+5. Limitation of Liability.
+
+EXCEPT AS PROHIBITED BY APPLICABLE LAW, IN NO EVENT AND UNDER NO LEGAL
+THEORY, WHETHER IN TORT (INCLUDING NEGLIGENCE), CONTRACT, OR OTHERWISE
+SHALL ANY LICENSOR BE LIABLE TO YOU FOR DAMAGES, INCLUDING ANY DIRECT,
+INDIRECT, SPECIAL, INCIDENTAL, OR CONSEQUENTIAL DAMAGES ARISING OUT OF
+OR RELATED TO THIS LICENSE, THE USE OR INABILITY TO USE THE WORK
+(INCLUDING BUT NOT LIMITED TO LOSS OF GOODWILL, BUSINESS INTERRUPTION,
+LOST PROFITS OR DATA, COMPUTER FAILURE OR MALFUNCTION, OR ANY OTHER
+COMMERCIAL DAMAGES OR LOSSES), EVEN IF THE LICENSOR HAS BEEN ADVISED OF
+THE POSSIBILITY OF SUCH DAMAGES.
+
+=======================================================================

cv/fast-foundationstereo/Utils.py

@@ -0,0 +1,88 @@
+import os, sys, torch, imageio, logging, importlib, argparse
+import cv2
+import numpy as np
+import yaml
+try:
+  import open3d as o3d
+except:
+  o3d = None
+
+AMP_DTYPE = torch.float16
+
+def set_logging_format(level=logging.INFO):
+  importlib.reload(logging)
+  FORMAT = '%(message)s'
+  logging.basicConfig(level=level, format=FORMAT, datefmt='%m-%d|%H:%M:%S')
+
+
+def set_seed(random_seed):
+  import torch,random
+  np.random.seed(random_seed)
+  random.seed(random_seed)
+  torch.manual_seed(random_seed)
+  torch.cuda.manual_seed_all(random_seed)
+  torch.backends.cudnn.deterministic = True
+  torch.backends.cudnn.benchmark = False
+
+
+def toOpen3dCloud(points,colors=None,normals=None):
+  cloud = o3d.geometry.PointCloud()
+  cloud.points = o3d.utility.Vector3dVector(points.astype(np.float64))
+  if colors is not None:
+    if colors.max()>1:
+      colors = colors/255.0
+    cloud.colors = o3d.utility.Vector3dVector(colors.astype(np.float64))
+  if normals is not None:
+    cloud.normals = o3d.utility.Vector3dVector(normals.astype(np.float64))
+  return cloud
+
+
+
+def depth2xyzmap(depth:np.ndarray, K, uvs:np.ndarray=None, zmin=0.1):
+  invalid_mask = (depth<zmin)
+  H,W = depth.shape[:2]
+  if uvs is None:
+    vs,us = np.meshgrid(np.arange(0,H),np.arange(0,W), sparse=False, indexing='ij')
+    vs = vs.reshape(-1)
+    us = us.reshape(-1)
+  else:
+    uvs = uvs.round().astype(int)
+    us = uvs[:,0]
+    vs = uvs[:,1]
+  zs = depth[vs,us]
+  xs = (us-K[0,2])*zs/K[0,0]
+  ys = (vs-K[1,2])*zs/K[1,1]
+  pts = np.stack((xs.reshape(-1),ys.reshape(-1),zs.reshape(-1)), 1)  #(N,3)
+  xyz_map = np.zeros((H,W,3), dtype=np.float32)
+  xyz_map[vs,us] = pts
+  if invalid_mask.any():
+    xyz_map[invalid_mask] = 0
+  return xyz_map
+
+
+def vis_disparity(disp, min_val=None, max_val=None, invalid_thres=np.inf, color_map=cv2.COLORMAP_TURBO, cmap=None, other_output={}):
+  """
+  @disp: np array (H,W)
+  @invalid_thres: > thres is invalid
+  """
+  disp = disp.copy()
+  H,W = disp.shape[:2]
+  invalid_mask = disp>=invalid_thres
+  if (invalid_mask==0).sum()==0:
+    other_output['min_val'] = None
+    other_output['max_val'] = None
+    return np.zeros((H,W,3))
+  if min_val is None:
+    min_val = disp[invalid_mask==0].min()
+  if max_val is None:
+    max_val = disp[invalid_mask==0].max()
+  other_output['min_val'] = min_val
+  other_output['max_val'] = max_val
+  vis = ((disp-min_val)/(max_val-min_val)).clip(0,1) * 255
+  if cmap is None:
+    vis = cv2.applyColorMap(vis.clip(0, 255).astype(np.uint8), color_map)[...,::-1]
+  else:
+    vis = cmap(vis.astype(np.uint8))[...,:3]*255
+  if invalid_mask.any():
+    vis[invalid_mask] = 0
+  return vis.astype(np.uint8)

cv/fast-foundationstereo/core/distill_block.py

@@ -0,0 +1,50 @@
+import os,sys
+code_dir = os.path.dirname(os.path.abspath(__file__))
+sys.path.append(code_dir+'/../')
+from foundation_stereo_ori.submodule import FeatureAtt
+import torch
+import torch.nn as nn
+import Utils as U
+import pickle
+
+
+class ForwardHelper(nn.Module):
+  def __init__(self, layers:list):
+    super().__init__()
+    self.layers = nn.ModuleList(layers)
+
+  def forward(self, x, left_feat=None):
+    for layer in self.layers:
+      if isinstance(layer, FeatureAtt):
+        x = layer(x, left_feat)
+      else:
+        x = layer(x)
+    return x
+
+
+class PostForwardHelper(nn.Module):
+  def __init__(self, layers:list):
+    super().__init__()
+    for pos in range(len(layers)):
+      if layers[pos] in ['sum', 'concat']:
+        self.op = layers[pos]
+        break
+    self.upsample = nn.Sequential(*layers[:pos])
+    self.out = nn.ModuleList(layers[pos+1:])
+
+  def forward(self, conv2, conv3, left_feat=None):
+    conv3_up = self.upsample(conv3)
+    if self.op == 'sum':
+      x = conv3_up + conv2
+    elif self.op == 'concat':
+      x = torch.cat((conv3_up, conv2), dim=1)
+    else:
+      raise ValueError(f"Unknown operation: {self.op}")
+
+    for layer in self.out:
+      if isinstance(layer, FeatureAtt):
+        x = layer(x, left_feat)
+      else:
+        x = layer(x)
+    return x
+

cv/fast-foundationstereo/core/extractor.py

@@ -0,0 +1,78 @@
+import torch,os,sys
+import torch.nn as nn
+import torch.nn.functional as F
+code_dir = os.path.dirname(os.path.realpath(__file__))
+sys.path.append(f'{code_dir}/../')
+from core.submodule import Conv2x_IN
+import timm
+
+
+
+class ContextNetSharedBackbone(nn.Module):
+  def __init__(self, args, c04, c08, c16, output_dim=[(128,128,128), (128,128,128)], norm_fn='batch', downsample=3):
+    super().__init__()
+    self.args = args
+    self.conv04 = nn.ModuleList([
+      nn.Conv2d(c04, output_dim[0][0], kernel_size=3, padding=1),
+      nn.Conv2d(c04, output_dim[1][0], kernel_size=3, padding=1),
+    ])
+
+  def forward(self, x4, x8, x16):
+    outputs04 = []
+    for i in range(len(self.conv04)):
+      outputs04.append(self.conv04[i](x4))
+    return (outputs04,)
+
+
+
+class DepthAnythingFeature:
+    model_configs = {
+        'vitl': {'encoder': 'vitl', 'features': 256, 'out_channels': [256, 512, 1024, 1024]},
+        'vitb': {'encoder': 'vitb', 'features': 128, 'out_channels': [96, 192, 384, 768]},
+        'vits': {'encoder': 'vits', 'features': 64, 'out_channels': [48, 96, 192, 384]}
+    }
+
+
+
+class Feature(nn.Module):
+    def __init__(self, args):
+        super(Feature, self).__init__()
+        self.args = args
+        model = timm.create_model('edgenext_small', pretrained=True, features_only=False)
+        self.stem = model.stem
+        self.stages = model.stages
+        chans = [48, 96, 160, 304]
+        self.chans = chans
+        vit_feat_dim = DepthAnythingFeature.model_configs[self.args.vit_size]['features']//2
+
+        self.deconv32_16 = Conv2x_IN(chans[3], chans[2], deconv=True, concat=True)
+        self.deconv16_8 = Conv2x_IN(chans[2]*2, chans[1], deconv=True, concat=True)
+        self.deconv8_4 = Conv2x_IN(chans[1]*2, chans[0], deconv=True, concat=True)
+
+        self.conv4 = nn.Conv2d(chans[0]*2, self.chans[0]*2+vit_feat_dim, kernel_size=1, stride=1, padding=0)
+
+        self.d_out = [self.chans[0]*2+vit_feat_dim, self.chans[1]*2, self.chans[2]*2, self.chans[3]]
+
+
+    def forward(self, x):
+        B,C,H,W = x.shape
+        if hasattr(self, 'stem'):
+          x = self.stem(x)
+          x4 = self.stages[0](x)
+          x8 = self.stages[1](x4)
+          x16 = self.stages[2](x8)
+          x32 = self.stages[3](x16)
+        else:
+          intermediates = self.model.forward_intermediates(x, intermediates_only=True)
+          x4, x8, x16, x32 = intermediates[-4:]
+
+        with torch.profiler.record_function("feature_deconv"):
+          x16 = self.deconv32_16(x32, x16)
+          x8 = self.deconv16_8(x16, x8)
+          x4 = self.deconv8_4(x8, x4)
+          x4 = self.conv4(x4)
+          if hasattr(self, 'conv8'):
+            x8 = self.conv8(x8)
+            x16 = self.conv16(x16)
+            x32 = self.conv32(x32)
+        return [x4, x8, x16, x32]

cv/fast-foundationstereo/core/foundation_stereo.py

@@ -0,0 +1,445 @@
+import torch,pdb,logging,timm
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+import sys,os
+code_dir = os.path.dirname(os.path.realpath(__file__))
+sys.path.append(f'{code_dir}/../')
+from core.update import BasicSelectiveMultiUpdateBlock
+from core.extractor import ContextNetSharedBackbone, Feature
+from core.geometry import Combined_Geo_Encoding_Volume
+from core.submodule import (
+    BasicConv, Conv3dNormActReduced, ResnetBasicBlock3D, BasicConv_IN, Conv2x,
+    FeatureAtt, CostVolumeDisparityAttention, SpatialAttentionExtractor,
+    ChannelAttentionEnhancement, disparity_regression, context_upsample,
+    build_gwc_volume_optimized_pytorch1, build_gwc_volume_triton,
+    build_concat_volume_optimized_pytorch1, build_concat_volume_optimized_pytorch,
+)
+from core.utils.utils import InputPadder
+import Utils as U
+import time
+
+sys.modules['foundation_stereo_ori'] = sys.modules['core']
+sys.modules['foundation_stereo_ori.submodule'] = sys.modules['core.submodule']
+sys.modules['foundation_stereo_ori.extractor'] = sys.modules['core.extractor']
+sys.modules['foundation_stereo_ori.update'] = sys.modules['core.update']
+sys.modules['foundation_stereo_ori.foundation_stereo'] = sys.modules['core.foundation_stereo']
+class FoundationStereo(nn.Module):
+  pass
+
+
+
+def normalize_image(img):
+    '''
+    @img: (B,C,H,W) in range 0-255, RGB order
+    '''
+    mean = img.new_tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1)
+    std = img.new_tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1)
+    return (img/255.0 - mean) / std
+
+
+class hourglass(nn.Module):
+    def __init__(self, cfg, in_channels, feat_dims=None):
+        super().__init__()
+        self.cfg = cfg
+        self.conv1 = nn.Sequential(BasicConv(in_channels, in_channels*2, is_3d=True, bn=True, relu=True, kernel_size=3,
+                                             padding=1, stride=2, dilation=1),
+                                   Conv3dNormActReduced(in_channels*2, in_channels*2, kernel_size=3, kernel_disp=17))
+
+        self.conv2 = nn.Sequential(BasicConv(in_channels*2, in_channels*4, is_3d=True, bn=True, relu=True, kernel_size=3,
+                                             padding=1, stride=2, dilation=1),
+                                   Conv3dNormActReduced(in_channels*4, in_channels*4, kernel_size=3, kernel_disp=17))
+
+        self.conv3 = nn.Sequential(BasicConv(in_channels*4, in_channels*6, is_3d=True, bn=True, relu=True, kernel_size=3,
+                                             padding=1, stride=2, dilation=1),
+                                   Conv3dNormActReduced(in_channels*6, in_channels*6, kernel_size=3, kernel_disp=17))
+
+
+        self.conv3_up = BasicConv(in_channels*6, in_channels*4, deconv=True, is_3d=True, bn=True,
+                                  relu=True, kernel_size=(4, 4, 4), padding=(1, 1, 1), stride=(2, 2, 2))
+
+        self.conv2_up = BasicConv(in_channels*4, in_channels*2, deconv=True, is_3d=True, bn=True,
+                                  relu=True, kernel_size=(4, 4, 4), padding=(1, 1, 1), stride=(2, 2, 2))
+
+        self.conv1_up = BasicConv(in_channels*2, in_channels, deconv=True, is_3d=True, bn=True,
+                                  relu=True, kernel_size=(4, 4, 4), padding=(1, 1, 1), stride=(2, 2, 2))
+        self.conv_out = nn.Sequential(
+          Conv3dNormActReduced(in_channels, in_channels, kernel_size=3, kernel_disp=17),
+          Conv3dNormActReduced(in_channels, in_channels, kernel_size=3, kernel_disp=17),
+        )
+
+        self.agg_0 = nn.Sequential(BasicConv(in_channels*8, in_channels*4, is_3d=True, kernel_size=1, padding=0, stride=1),
+                                   Conv3dNormActReduced(in_channels*4, in_channels*4, kernel_size=3, kernel_disp=17),
+                                   Conv3dNormActReduced(in_channels*4, in_channels*4, kernel_size=3, kernel_disp=17),)
+
+        self.agg_1 = nn.Sequential(BasicConv(in_channels*4, in_channels*2, is_3d=True, kernel_size=1, padding=0, stride=1),
+                                   Conv3dNormActReduced(in_channels*2, in_channels*2, kernel_size=3, kernel_disp=17),
+                                   Conv3dNormActReduced(in_channels*2, in_channels*2, kernel_size=3, kernel_disp=17))
+
+        self.atts = nn.ModuleDict({
+          "4": CostVolumeDisparityAttention(d_model=in_channels, nhead=4, dim_feedforward=in_channels, norm_first=False, num_transformer=4, max_len=self.cfg['max_disp']//16),
+        })
+        self.conv_patch = nn.Sequential(
+          nn.Conv3d(in_channels, in_channels, kernel_size=4, stride=4, padding=0, groups=in_channels),
+          nn.BatchNorm3d(in_channels),
+        )
+        self.feature_att_8 = FeatureAtt(in_channels*2, feat_dims[1])
+        self.feature_att_16 = FeatureAtt(in_channels*4, feat_dims[2])
+        self.feature_att_32 = FeatureAtt(in_channels*6, feat_dims[3])
+        self.feature_att_up_16 = FeatureAtt(in_channels*4, feat_dims[2])
+        self.feature_att_up_8 = FeatureAtt(in_channels*2, feat_dims[1])
+
+        self.post32_to_16 = None
+        self.post16_to_8 = None
+        self.post8_to_4 = None
+
+    def forward(self, x, features):
+        conv1 = self.conv1(x)
+        conv1 = self.feature_att_8(conv1, features[1])
+
+        conv2 = self.conv2(conv1)
+        conv2 = self.feature_att_16(conv2, features[2])
+
+        conv3 = self.conv3(conv2)
+        conv3 = self.feature_att_32(conv3, features[3])
+        if not hasattr(self, 'post32_to_16') or self.post32_to_16 is None:
+          conv3_up = self.conv3_up(conv3)
+          conv2 = torch.cat((conv3_up, conv2), dim=1)
+          conv2 = self.agg_0(conv2)
+          conv2 = self.feature_att_up_16(conv2, features[2])
+        else:
+          conv2 = self.post32_to_16(conv2, conv3, features[2])
+
+        if not hasattr(self, 'post16_to_8') or self.post16_to_8 is None:
+          conv2_up = self.conv2_up(conv2)
+          conv1 = torch.cat((conv2_up, conv1), dim=1)
+          conv1 = self.agg_1(conv1)
+          conv1 = self.feature_att_up_8(conv1, features[1])
+        else:
+          conv1 = self.post16_to_8(conv1, conv2, features[1])
+
+        conv = self.conv1_up(conv1)
+        if not hasattr(self, 'post8_to_4') or self.post8_to_4 is None:
+          x = self.conv_patch(x)
+          x = self.atts["4"](x)
+          x = F.interpolate(x, scale_factor=4, mode='trilinear', align_corners=False)
+          conv = conv + x
+          conv = self.conv_out(conv)
+        else:
+          conv = self.post8_to_4(x, conv)
+
+        return conv
+
+
+class FastFoundationStereo(nn.Module):
+  def __init__(self, args):
+    super().__init__()
+    self.args = args
+    self.dtype = torch.float32
+
+    context_dims = args.hidden_dims
+    self.cv_group = args.get('cv_group', 8)
+    self.concat_channel = 24
+    volume_dim = args.get('volume_dim', 28)
+    self.volume_dim = volume_dim
+    self.update_block = BasicSelectiveMultiUpdateBlock(self.args, self.args.hidden_dims[0], volume_dim=volume_dim)
+    self.sam = SpatialAttentionExtractor()
+    self.cam = ChannelAttentionEnhancement(self.args.hidden_dims[0])
+    self.context_zqr_convs = nn.ModuleList([nn.Conv2d(context_dims[i], args.hidden_dims[i]*3, kernel_size=3, padding=3//2) for i in range(self.args.n_gru_layers)])
+    self.feature = Feature(args)
+    self.proj_cmb = nn.Conv2d(self.feature.d_out[0], self.concat_channel//2, kernel_size=1, padding=0)
+    self.cnet = ContextNetSharedBackbone(args, c04=self.feature.d_out[0], c08=self.feature.d_out[1], c16=self.feature.d_out[2], output_dim=[args.hidden_dims, context_dims])
+
+    self.stem_2 = nn.Sequential(
+      BasicConv_IN(3, 32, kernel_size=3, stride=2, padding=1),
+      nn.Conv2d(32, 32, 3, 1, 1, bias=False),
+      nn.InstanceNorm2d(32), nn.ReLU()
+    )
+    self.spx_2_gru = Conv2x(32, 32, deconv=True, bn=False, concat=True)
+    self.spx_gru = nn.Sequential(
+          nn.ConvTranspose2d(2*32, 9, kernel_size=4, stride=2, padding=1),
+          )
+
+    self.corr_stem = nn.Sequential(
+      nn.Conv3d(self.proj_cmb.out_channels*2+self.cv_group, volume_dim, kernel_size=1),
+      BasicConv(volume_dim, volume_dim, kernel_size=3, padding=1, is_3d=True),
+      ResnetBasicBlock3D(volume_dim, volume_dim, kernel_size=3, stride=1, padding=1),
+      ResnetBasicBlock3D(volume_dim, volume_dim, kernel_size=3, stride=1, padding=1),
+    )
+    self.corr_feature_att = FeatureAtt(volume_dim, self.feature.d_out[0])
+    self.cost_agg = hourglass(cfg=self.args, in_channels=volume_dim, feat_dims=self.feature.d_out)
+    self.classifier = nn.Sequential(
+      BasicConv(volume_dim, volume_dim//2, kernel_size=3, padding=1, is_3d=True),
+      ResnetBasicBlock3D(volume_dim//2, volume_dim//2, kernel_size=3, stride=1, padding=1),
+      nn.Conv3d(volume_dim//2, 1, kernel_size=7, padding=3),
+    )
+
+    r = self.args.corr_radius
+    dx = torch.arange(-r, r+1, requires_grad=False, dtype=torch.int8).reshape(1, 1, 2*r+1, 1)
+    self.register_buffer("dx", dx)
+
+
+  def upsample_disp(self, disp, mask_feat_4, stem_2x):
+    with torch.amp.autocast('cuda', enabled=self.args.mixed_precision, dtype=U.AMP_DTYPE):
+      xspx = self.spx_2_gru(mask_feat_4, stem_2x)   # 1/2 resolution
+      spx_pred = self.spx_gru(xspx)
+      spx_pred = F.softmax(spx_pred, 1)
+      up_disp = context_upsample(disp*4., spx_pred).unsqueeze(1)
+    return up_disp.to(self.dtype)
+
+
+  def forward(self, image1, image2, iters=12, test_mode=False, low_memory=False, init_disp=None, profile=False, optimize_build_volume='pytorch1'):
+    """ Estimate disparity between pair of frames """
+    B,C,H,W = image1.shape
+    low_memory = low_memory or (self.args.get('low_memory', False))
+    image1 = normalize_image(image1)
+    image2 = normalize_image(image2)
+    with torch.amp.autocast('cuda', enabled=self.args.mixed_precision, dtype=U.AMP_DTYPE):
+      out = self.feature(torch.cat([image1, image2], dim=0))
+      features_left = [o[:B] for o in out]
+      features_right = [o[B:] for o in out]
+      stem_2x = self.stem_2(image1)
+
+      if optimize_build_volume=='pytorch1':
+        gwc_volume = build_gwc_volume_optimized_pytorch1(features_left[0], features_right[0], self.args.max_disp//4, self.cv_group, normalize=self.args.normalize)
+      elif optimize_build_volume=='triton':
+        gwc_volume = build_gwc_volume_triton(features_left[0], features_right[0], self.args.max_disp//4, self.cv_group, normalize=self.args.normalize)
+      else:
+        raise RuntimeError(f"Invalid optimize_build_volume: {optimize_build_volume}")
+
+      left_tmp = self.proj_cmb(features_left[0])
+      right_tmp = self.proj_cmb(features_right[0])
+      concat_volume = build_concat_volume_optimized_pytorch1(left_tmp, right_tmp, maxdisp=self.args.max_disp//4)
+      del left_tmp, right_tmp
+      comb_volume = torch.cat([gwc_volume, concat_volume], dim=1)
+      del concat_volume, gwc_volume
+
+      comb_volume = self.corr_stem(comb_volume)
+      comb_volume = self.corr_feature_att(comb_volume, features_left[0])
+      comb_volume = self.cost_agg(comb_volume, features_left)
+
+      # Init disp from geometry encoding volume
+      logits = self.classifier(comb_volume).squeeze(1)
+      prob = F.softmax(logits, dim=1)
+      if init_disp is None:
+        init_disp = disparity_regression(prob, self.args.max_disp//4)
+
+      cnet_list = self.cnet(features_left[0], features_left[1], features_left[2])
+      cnet_list = list(cnet_list)
+      net_list = [torch.tanh(x[0]) for x in cnet_list]
+      inp_list = [torch.relu(x[1]) for x in cnet_list]
+      inp_list = [self.cam(x) * x for x in inp_list]
+      att = [self.sam(x) for x in inp_list]
+
+    geo_fn = Combined_Geo_Encoding_Volume(features_left[0].to(self.dtype), features_right[0].to(self.dtype), comb_volume.to(self.dtype), num_levels=self.args.corr_levels)
+    b, c, h, w = features_left[0].shape
+    coords = torch.arange(w, dtype=torch.float, device=init_disp.device).reshape(1,1,w,1).repeat(b, h, 1, 1)
+    disp = init_disp.to(self.dtype)
+    disp_preds = []
+
+    del comb_volume, features_left, features_right, cnet_list
+
+    # GRUs iterations to update disparity (1/4 resolution)
+    for itr in range(iters):
+      disp = disp.detach()
+      geo_feat = geo_fn(disp, coords, dx=self.dx, low_memory=low_memory)
+      with torch.amp.autocast('cuda', enabled=self.args.mixed_precision, dtype=U.AMP_DTYPE):
+        net_list, mask_feat_4, delta_disp = self.update_block(net_list, inp_list, geo_feat.to(self.dtype), disp, att)
+
+      disp = disp + delta_disp.to(self.dtype)
+      if test_mode and itr < iters-1:
+        continue
+
+      # upsample predictions
+      disp_up = self.upsample_disp(disp.to(self.dtype), mask_feat_4.to(self.dtype), stem_2x.to(self.dtype))
+      disp_preds.append(disp_up)
+
+
+    if test_mode:
+      return disp_up
+
+    return init_disp, disp_preds
+
+
+  def run_hierachical(self, image1, image2, iters=12, test_mode=False, low_memory=False, small_ratio=0.5):
+      B,_,H,W = image1.shape
+      img1_small = F.interpolate(image1, scale_factor=small_ratio, align_corners=False, mode='bilinear')
+      img2_small = F.interpolate(image2, scale_factor=small_ratio, align_corners=False, mode='bilinear')
+      padder = InputPadder(img1_small.shape[-2:], divis_by=32, force_square=False)
+      img1_small, img2_small = padder.pad(img1_small, img2_small)
+      disp_small = self.forward(img1_small, img2_small, test_mode=True, iters=iters, low_memory=low_memory)
+      disp_small = padder.unpad(disp_small)
+      disp_small_up = F.interpolate(disp_small, size=(H,W), mode='bilinear', align_corners=True) * 1/small_ratio
+      disp_small_up = disp_small_up.clip(0, None)
+
+      padder = InputPadder(image1.shape[-2:], divis_by=32, force_square=False)
+      image1, image2, disp_small_up = padder.pad(image1, image2, disp_small_up)
+      disp_small_up += padder._pad[0]
+      init_disp = F.interpolate(disp_small_up, scale_factor=0.25, mode='bilinear', align_corners=True) * 0.25   # Init disp will be 1/4
+      disp = self.forward(image1, image2, iters=iters, test_mode=test_mode, low_memory=low_memory, init_disp=init_disp)
+      disp = padder.unpad(disp)
+      return disp
+
+FoundationStereoLite = FastFoundationStereo
+
+
+class TrtFeatureRunner(nn.Module):
+  def __init__(self, model):
+    super().__init__()
+    self.feature = model.feature
+    self.stem_2 = model.stem_2
+
+  def forward(self, image1, image2):
+    image1 = normalize_image(image1)
+    image2 = normalize_image(image2)
+    B = len(image1)
+    out = self.feature(torch.cat([image1, image2], dim=0))
+    features_left = [o[:B] for o in out]
+    features_right = [o[B:] for o in out]
+    stem_2x = self.stem_2(image1)
+    return *features_left, features_right[0], stem_2x
+
+
+class TrtPostRunner(nn.Module):
+  def __init__(self, model):
+    super().__init__()
+    self.args = model.args
+    self.dtype = model.dtype
+    self.register_buffer("dx", model.dx)
+    self.proj_cmb = model.proj_cmb
+    self.corr_stem = model.corr_stem
+    self.corr_feature_att = model.corr_feature_att
+    self.cost_agg = model.cost_agg
+    self.classifier = model.classifier
+    self.update_block = model.update_block
+    self.sam = model.sam
+    self.cam = model.cam
+    self.feature = model.feature
+    self.spx_2_gru = model.spx_2_gru
+    self.spx_gru = model.spx_gru
+    self.cnet = model.cnet
+
+  def upsample_disp(self, disp, mask_feat_4, stem_2x):
+    with torch.amp.autocast('cuda', enabled=self.args.mixed_precision, dtype=U.AMP_DTYPE):
+      xspx = self.spx_2_gru(mask_feat_4, stem_2x)   # 1/2 resolution
+      spx_pred = self.spx_gru(xspx)
+      spx_pred = F.softmax(spx_pred, 1)
+      up_disp = context_upsample(disp*4., spx_pred).unsqueeze(1)
+    return up_disp.to(self.dtype)
+
+
+  def forward(self, features_left_04, features_left_08, features_left_16, features_left_32, features_right_04, stem_2x, gwc_volume):
+    features_left = [features_left_04, features_left_08, features_left_16, features_left_32]
+    left_tmp = self.proj_cmb(features_left_04)
+    right_tmp = self.proj_cmb(features_right_04)
+    concat_volume = build_concat_volume_optimized_pytorch(left_tmp, right_tmp, maxdisp=self.args.max_disp//4)
+    del left_tmp, right_tmp
+    comb_volume = torch.cat([gwc_volume, concat_volume], dim=1)
+    del concat_volume, gwc_volume
+    comb_volume = self.corr_stem(comb_volume)
+    comb_volume = self.corr_feature_att(comb_volume, features_left_04)
+    comb_volume = self.cost_agg(comb_volume, features_left)
+
+    # Init disp from geometry encoding volume
+    logits = self.classifier(comb_volume).squeeze(1)
+    prob = F.softmax(logits, dim=1)
+    init_disp = disparity_regression(prob, self.args.max_disp//4)
+
+    cnet_list = self.cnet(features_left[0], features_left[1], features_left[2])
+    cnet_list = list(cnet_list)
+    net_list = [torch.tanh(x[0]) for x in cnet_list]
+    inp_list = [torch.relu(x[1]) for x in cnet_list]
+    inp_list = [self.cam(x) * x for x in inp_list]
+    att = [self.sam(x) for x in inp_list]
+
+    geo_fn = Combined_Geo_Encoding_Volume(features_left_04.to(self.dtype), features_right_04.to(self.dtype), comb_volume.to(self.dtype), num_levels=self.args.corr_levels)
+    b, c, h, w = features_left[0].shape
+    coords = torch.arange(w, dtype=torch.float, device=init_disp.device).reshape(1,1,w,1).repeat(b, h, 1, 1)
+    disp = init_disp.to(self.dtype)
+
+    # GRUs iterations to update disparity (1/4 resolution)
+    for itr in range(self.args.valid_iters):
+      disp = disp.detach()
+      geo_feat = geo_fn(disp, coords, dx=self.dx, low_memory=True)
+      net_list, mask_feat_4, delta_disp = self.update_block(net_list, inp_list, geo_feat.to(self.dtype), disp, att)
+
+      disp = disp + delta_disp.to(self.dtype)
+      if itr < self.args.valid_iters-1:
+        continue
+
+      disp_up = self.upsample_disp(disp.to(self.dtype), mask_feat_4.to(self.dtype), stem_2x.to(self.dtype))
+
+    return disp_up
+
+
+class TrtRunner(nn.Module):
+  def __init__(self, args, feature_runner_engine_path, post_runner_engine_path):
+    super().__init__()
+    import tensorrt as trt
+    self.args = args
+    with open(feature_runner_engine_path, 'rb') as file:
+      engine_data = file.read()
+    self.trt_logger = trt.Logger(trt.Logger.WARNING)
+    self.feature_engine = trt.Runtime(self.trt_logger).deserialize_cuda_engine(engine_data)
+    self.feature_context = self.feature_engine.create_execution_context()
+
+    with open(post_runner_engine_path, 'rb') as file:
+      engine_data = file.read()
+    self.post_engine = trt.Runtime(self.trt_logger).deserialize_cuda_engine(engine_data)
+    self.post_context = self.post_engine.create_execution_context()
+    self.max_disp = args.max_disp
+    self.cv_group = args.get('cv_group', 8)
+
+  def trt_dtype_to_torch(self, dt):
+    import tensorrt as trt
+    if dt==trt.DataType.FLOAT: return torch.float32
+    if dt==trt.DataType.HALF: return torch.float16
+    if dt==trt.DataType.BF16: return torch.bfloat16
+    if dt==trt.DataType.INT32: return torch.int32
+    if dt==trt.DataType.INT8: return torch.int8
+    if dt==trt.DataType.BOOL: return torch.bool
+    raise RuntimeError(f"Unsupported TRT dtype: {dt}")
+
+  def get_io_tensor_names(self, engine, mode):
+    names = []
+    n = engine.num_io_tensors
+    for i in range(n):
+      name = engine.get_tensor_name(i)
+      if engine.get_tensor_mode(name)==mode:
+        names.append(name)
+    return names
+
+  def run_trt(self, engine, context, inputs_by_name:dict):
+    import tensorrt as trt
+    for name, tensor in list(inputs_by_name.items()):
+      expected_dtype = self.trt_dtype_to_torch(engine.get_tensor_dtype(name))
+      if tensor.dtype != expected_dtype: inputs_by_name[name] = tensor.to(expected_dtype)
+      if not inputs_by_name[name].is_contiguous(): inputs_by_name[name] = inputs_by_name[name].contiguous()
+      context.set_input_shape(name, tuple(inputs_by_name[name].shape))
+    outputs = {}
+    out_names = [n for n in self.get_io_tensor_names(engine, trt.TensorIOMode.OUTPUT)]
+    for name in out_names:
+      shp = tuple(context.get_tensor_shape(name))
+      dtype = self.trt_dtype_to_torch(engine.get_tensor_dtype(name))
+      outputs[name] = torch.empty(shp, device='cuda', dtype=dtype)
+    for name, tensor in inputs_by_name.items(): context.set_tensor_address(name, int(tensor.data_ptr()))
+    for name, tensor in outputs.items(): context.set_tensor_address(name, int(tensor.data_ptr()))
+    stream = torch.cuda.current_stream().cuda_stream
+    ok = context.execute_async_v3(stream)
+    assert ok
+    return outputs
+
+  def forward(self, image1, image2):
+    import tensorrt as trt
+    feat_out = self.run_trt(self.feature_engine, self.feature_context, {'left': image1, 'right': image2})
+    gwc_volume = build_gwc_volume_triton(feat_out['features_left_04'].half(), feat_out['features_right_04'].half(), self.args.max_disp//4, self.cv_group, normalize=self.args.normalize)
+    post_inputs = feat_out
+    post_inputs['gwc_volume'] = gwc_volume
+    in_names = self.get_io_tensor_names(self.post_engine, trt.TensorIOMode.INPUT)
+    tmp_keys = list(post_inputs.keys())
+    for k in tmp_keys:
+      if k not in in_names:
+        del post_inputs[k]
+    out = self.run_trt(self.post_engine, self.post_context, post_inputs)
+    disp = out['disp']
+    return disp

cv/fast-foundationstereo/core/geometry.py

@@ -0,0 +1,81 @@
+import torch,os,sys
+import torch.nn.functional as F
+from core.utils.utils import bilinear_sampler, bilinear_sampler1d
+code_dir = os.path.dirname(os.path.realpath(__file__))
+sys.path.append(f'{code_dir}/../')
+
+class Combined_Geo_Encoding_Volume:
+    def __init__(self, init_fmap1, init_fmap2, geo_volume, num_levels=2):
+        self.num_levels = num_levels
+        self.geo_volume_pyramid = []
+        self.init_corr_pyramid = []
+
+        # all pairs correlation
+        init_corr = Combined_Geo_Encoding_Volume.corr(init_fmap1, init_fmap2)
+
+        b, h, w, _, w2 = init_corr.shape
+        b, c, d, h, w = geo_volume.shape
+        geo_volume = geo_volume.permute(0, 3, 4, 1, 2).reshape(b*h*w, c, 1, d)
+
+        init_corr = init_corr.view(b*h*w, 1, 1, w2)
+        self.geo_volume_pyramid.append(geo_volume)
+        self.init_corr_pyramid.append(init_corr)
+        for _ in range(self.num_levels-1):
+            geo_volume = F.avg_pool2d(geo_volume, [1,2], stride=[1,2])
+            self.geo_volume_pyramid.append(geo_volume)
+
+        for _ in range(self.num_levels-1):
+            init_corr = F.avg_pool2d(init_corr, [1,2], stride=[1,2])
+            self.init_corr_pyramid.append(init_corr)
+
+
+
+    def __call__(self, disp, coords, dx, low_memory=True):
+        b, _, h, w = disp.shape
+        out_pyramid = []
+        for i in range(self.num_levels):
+            with torch.profiler.record_function(f"make disp_lvl {i}"):
+              geo_volume = self.geo_volume_pyramid[i]
+              x0 = dx + disp.view(b*h*w, 1, 1, 1) / 2**i
+            with torch.profiler.record_function(f"bilinear_sampler geo_volume {i}"):
+              if low_memory:
+                geo_volume = bilinear_sampler1d(geo_volume, x0, mode='bilinear', align_corners=True)
+              else:
+                y0 = torch.zeros_like(x0)
+                disp_lvl = torch.cat([x0,y0], dim=-1)
+                geo_volume = bilinear_sampler(geo_volume, disp_lvl, low_memory=low_memory)
+              geo_volume = geo_volume.view(b, h, w, -1)   #(b, h, h, 3x3xC)
+
+            with torch.profiler.record_function(f"make init_coords_lvl {i}"):
+              init_corr = self.init_corr_pyramid[i]   # (B*H*W, 1, 1, W2)
+              init_x0 = coords.view(b*h*w, 1, 1, 1)/2**i - disp.view(b*h*w, 1, 1, 1) / 2**i + dx   # X on right image
+            with torch.profiler.record_function(f"bilinear_sampler init_corr {i}"):
+              if low_memory:
+                init_corr = bilinear_sampler1d(init_corr, init_x0, mode='bilinear', align_corners=True)
+              else:
+                init_coords_lvl = torch.cat([init_x0,y0], dim=-1)
+                init_corr = bilinear_sampler(init_corr, init_coords_lvl, low_memory=low_memory)
+              init_corr = init_corr.view(b, h, w, -1)
+
+            out_pyramid.append(geo_volume)
+            out_pyramid.append(init_corr)
+
+        with torch.profiler.record_function(f"make out_pyramid"):
+          out_pyramid = torch.cat(out_pyramid, dim=-1)
+          return out_pyramid.permute(0, 3, 1, 2)   #(B,C,H,W)
+
+
+    @staticmethod
+    def corr(fmap1, fmap2, normalize=True):
+        with torch.profiler.record_function("build corr"):
+          B, D, H, W1 = fmap1.shape
+          _, _, _, W2 = fmap2.shape
+          fmap1 = fmap1.view(B, D, H, W1)
+          fmap2 = fmap2.view(B, D, H, W2)
+          if normalize:
+            with torch.cuda.amp.autocast(enabled=False):
+              corr = torch.einsum('aijk,aijh->ajkh', F.normalize(fmap1.float(), dim=1), F.normalize(fmap2.float(), dim=1))
+          else:
+            corr = corr.view(B, H, W1, 1, W2).to(fmap1.dtype)
+          corr = corr.view(B, H, W1, 1, W2).to(fmap1.dtype)
+        return corr

cv/fast-foundationstereo/core/submodule.py

@@ -0,0 +1,675 @@
+import torch,pdb,os,sys
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+code_dir = os.path.dirname(os.path.realpath(__file__))
+sys.path.append(f'{code_dir}/../')
+from Utils import AMP_DTYPE
+import Utils as U
+try:
+  import triton
+  import triton.language as tl
+except Exception:
+  triton = None
+  tl = None
+
+def _is_contiguous(tensor: torch.Tensor) -> bool:
+    if torch.jit.is_scripting():
+        return tensor.is_contiguous()
+    else:
+        return tensor.is_contiguous(memory_format=torch.contiguous_format)
+
+
+
+class LayerNorm2d(nn.LayerNorm):
+    r""" https://huggingface.co/spaces/Roll20/pet_score/blob/b258ef28152ab0d5b377d9142a23346f863c1526/lib/timm/models/convnext.py#L85
+    LayerNorm for channels_first tensors with 2d spatial dimensions (ie N, C, H, W).
+    """
+
+    def __init__(self, normalized_shape, eps=1e-6):
+        """
+        @normalized_shape: channel dim
+        """
+        super().__init__(normalized_shape, eps=eps)
+
+    def forward(self, x) -> torch.Tensor:
+        """
+        @x: (B,C,H,W)
+        """
+        if _is_contiguous(x):
+            return F.layer_norm(x.permute(0, 2, 3, 1), self.normalized_shape, self.weight, self.bias, self.eps).permute(0, 3, 1, 2).contiguous()
+        else:
+            s, u = torch.var_mean(x, dim=1, keepdim=True)
+            x = (x - u) * torch.rsqrt(s + self.eps)
+            x = x * self.weight[:, None, None] + self.bias[:, None, None]
+            return x
+
+
+class BasicConv(nn.Module):
+
+    def __init__(self, in_channels, out_channels, deconv=False, is_3d=False, bn=True, relu=True, norm='batch', **kwargs):
+        super(BasicConv, self).__init__()
+
+        self.relu = nn.LeakyReLU(inplace=True) if relu else nn.Identity()
+        self.use_bn = bn
+        self.bn = nn.Identity()
+        if is_3d:
+            if deconv:
+                self.conv = nn.ConvTranspose3d(in_channels, out_channels, bias=False, **kwargs)
+            else:
+                self.conv = nn.Conv3d(in_channels, out_channels, bias=False, **kwargs)
+            if self.use_bn:
+              if norm=='batch':
+                self.bn = nn.BatchNorm3d(out_channels)
+              elif norm=='instance':
+                self.bn = nn.InstanceNorm3d(out_channels)
+        else:
+            if deconv:
+                self.conv = nn.ConvTranspose2d(in_channels, out_channels, bias=False, **kwargs)
+            else:
+                self.conv = nn.Conv2d(in_channels, out_channels, bias=False, **kwargs)
+            if self.use_bn:
+              if norm=='batch':
+                self.bn = nn.BatchNorm2d(out_channels)
+              elif norm=='instance':
+                self.bn = nn.InstanceNorm2d(out_channels)
+
+    def forward(self, x):
+        x = self.conv(x)
+        if self.use_bn:
+            x = self.bn(x)
+        if isinstance(self.relu, bool):
+          if self.relu:
+            self.relu = nn.LeakyReLU(inplace=True)
+          else:
+            self.relu = nn.Identity()
+        x = self.relu(x)
+        return x
+
+
+class Conv3dNormActReduced(nn.Module):
+    def __init__(self, C_in, C_out, hidden=None, kernel_size=3, kernel_disp=None, stride=1, norm=nn.BatchNorm3d):
+        super().__init__()
+        if kernel_disp is None:
+          kernel_disp = kernel_size
+        if hidden is None:
+            hidden = C_out
+        self.conv1 = nn.Sequential(
+            nn.Conv3d(C_in, hidden, kernel_size=(1,kernel_size,kernel_size), padding=(0, kernel_size//2, kernel_size//2), stride=(1, stride, stride)),
+            norm(hidden),
+            nn.ReLU(),
+        )
+        self.conv2 = nn.Sequential(
+            nn.Conv3d(hidden, C_out, kernel_size=(kernel_disp, 1, 1), padding=(kernel_disp//2, 0, 0), stride=(stride, 1, 1)),
+            norm(C_out),
+            nn.ReLU(),
+        )
+
+
+    def forward(self, x):
+        """
+        @x: (B,C,D,H,W)
+        """
+        x = self.conv1(x)
+        x = self.conv2(x)
+        return x
+
+
+class ResnetBasicBlock(nn.Module):
+  def __init__(self, inplanes, planes, kernel_size=3, stride=1, padding=1, downsample=None, groups=1, base_width=64, dilation=1, norm_layer=nn.BatchNorm2d, bias=False):
+    super().__init__()
+    self.norm_layer = norm_layer
+    if groups != 1 or base_width != 64:
+            raise ValueError('BasicBlock only supports groups=1 and base_width=64')
+    if dilation > 1:
+            raise NotImplementedError("Dilation > 1 not supported in BasicBlock")
+    # Both self.conv1 and self.downsample layers downsample the input when stride != 1
+    self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=kernel_size, stride=stride, bias=bias, padding=padding)
+    if self.norm_layer is not None:
+      self.bn1 = norm_layer(planes)
+    self.relu = nn.ReLU(inplace=True)
+    self.conv2 = nn.Conv2d(planes, planes, kernel_size=kernel_size, bias=bias, padding=padding)
+    if self.norm_layer is not None:
+      self.bn2 = norm_layer(planes)
+    self.downsample = downsample
+    self.stride = stride
+
+
+  def forward(self, x):
+    identity = x
+
+    out = self.conv1(x)
+    if self.norm_layer is not None:
+      out = self.bn1(out)
+    out = self.relu(out)
+
+    out = self.conv2(out)
+    if self.norm_layer is not None:
+      out = self.bn2(out)
+
+    if self.downsample is not None:
+      identity = self.downsample(x)
+    out += identity
+    out = self.relu(out)
+
+    return out
+
+
+class ResnetBasicBlock3D(nn.Module):
+  def __init__(self, inplanes, planes, kernel_size=3, stride=1, padding=1, downsample=None, groups=1, base_width=64, dilation=1, norm_layer=nn.BatchNorm3d, bias=False):
+    super().__init__()
+    self.norm_layer = norm_layer
+    if groups != 1 or base_width != 64:
+            raise ValueError('BasicBlock only supports groups=1 and base_width=64')
+    if dilation > 1:
+            raise NotImplementedError("Dilation > 1 not supported in BasicBlock")
+    # Both self.conv1 and self.downsample layers downsample the input when stride != 1
+    self.conv1 = nn.Conv3d(inplanes, planes, kernel_size=kernel_size, stride=stride, bias=bias, padding=padding)
+    if self.norm_layer is not None:
+      self.bn1 = norm_layer(planes)
+    self.relu = nn.ReLU(inplace=True)
+    self.conv2 = nn.Conv3d(planes, planes, kernel_size=kernel_size, bias=bias, padding=padding)
+    if self.norm_layer is not None:
+      self.bn2 = norm_layer(planes)
+    self.downsample = downsample
+    self.stride = stride
+
+
+  def forward(self, x):
+    identity = x
+
+    out = self.conv1(x)
+    if self.norm_layer is not None:
+      out = self.bn1(out)
+    out = self.relu(out)
+
+    out = self.conv2(out)
+    if self.norm_layer is not None:
+      out = self.bn2(out)
+
+    if self.downsample is not None:
+      identity = self.downsample(x)
+    out += identity
+    out = self.relu(out)
+
+    return out
+
+
+class FlashMultiheadAttention(nn.Module):
+    def __init__(self, embed_dim, num_heads):
+        super().__init__()
+        self.num_heads = num_heads
+        self.embed_dim = embed_dim
+        self.head_dim = embed_dim // num_heads
+        assert self.head_dim * num_heads == self.embed_dim, "embed_dim must be divisible by num_heads"
+
+        self.q_proj = nn.Linear(embed_dim, embed_dim)
+        self.k_proj = nn.Linear(embed_dim, embed_dim)
+        self.v_proj = nn.Linear(embed_dim, embed_dim)
+        self.out_proj = nn.Linear(embed_dim, embed_dim)
+
+    def forward(self, query, key, value, attn_mask=None, window_size=(-1,-1)):
+        """
+        @query: (B,L,C)
+        """
+        B,L,C = query.shape
+        Q = self.q_proj(query)
+        K = self.k_proj(key)
+        V = self.v_proj(value)
+
+        Q = Q.view(Q.size(0), Q.size(1), self.num_heads, self.head_dim)
+        K = K.view(K.size(0), K.size(1), self.num_heads, self.head_dim)
+        V = V.view(V.size(0), V.size(1), self.num_heads, self.head_dim)
+
+        attn_output = F.scaled_dot_product_attention(Q, K, V)
+
+        attn_output = attn_output.reshape(B,L,-1)
+        output = self.out_proj(attn_output)
+
+        return output
+
+
+
+class FlashAttentionTransformerEncoderLayer(nn.Module):
+    def __init__(self, embed_dim, num_heads, dim_feedforward, dropout=0.1, act=nn.GELU, norm=nn.LayerNorm):
+        super().__init__()
+        self.self_attn = FlashMultiheadAttention(embed_dim, num_heads)
+        self.act = act()
+
+        self.linear1 = nn.Linear(embed_dim, dim_feedforward)
+        self.dropout = nn.Dropout(dropout)
+        self.linear2 = nn.Linear(dim_feedforward, embed_dim)
+
+        self.norm1 = norm(embed_dim)
+        self.norm2 = norm(embed_dim)
+        self.dropout1 = nn.Dropout(dropout)
+        self.dropout2 = nn.Dropout(dropout)
+
+    def forward(self, src, src_mask=None, window_size=(-1, -1)):
+        dtype = src.dtype
+        src2 = self.self_attn(src, src, src, src_mask, window_size=window_size)
+        src = src + self.dropout1(src2)
+        src = self.norm1(src).to(dtype)
+
+        src2 = self.linear2(self.dropout(self.act(self.linear1(src))))
+        src = src + self.dropout2(src2)
+        src = self.norm2(src).to(dtype)
+
+        return src
+
+
+class Conv2x(nn.Module):
+
+    def __init__(self, in_channels, out_channels, deconv=False, is_3d=False, concat=True, keep_concat=True, bn=True, relu=True, keep_dispc=False):
+        super(Conv2x, self).__init__()
+        self.concat = concat
+        self.is_3d = is_3d
+        if deconv and is_3d:
+            kernel = (4, 4, 4)
+        elif deconv:
+            kernel = 4
+        else:
+            kernel = 3
+
+        if deconv and is_3d and keep_dispc:
+            kernel = (1, 4, 4)
+            stride = (1, 2, 2)
+            padding = (0, 1, 1)
+            self.conv1 = BasicConv(in_channels, out_channels, deconv, is_3d, bn=bn, relu=True, kernel_size=kernel, stride=stride, padding=padding)
+        else:
+            self.conv1 = BasicConv(in_channels, out_channels, deconv, is_3d, bn=bn, relu=True, kernel_size=kernel, stride=2, padding=1)
+
+        if self.concat:
+            mul = 2 if keep_concat else 1
+            self.conv2 = BasicConv(out_channels*2, out_channels*mul, False, is_3d, bn, relu, kernel_size=3, stride=1, padding=1)
+        else:
+            self.conv2 = BasicConv(out_channels, out_channels, False, is_3d, bn, relu, kernel_size=3, stride=1, padding=1)
+
+    def forward(self, x, rem):
+        x = self.conv1(x)
+        if x.shape != rem.shape:
+            x = F.interpolate(x, size=(rem.shape[-2], rem.shape[-1]), mode='bilinear')
+        if self.concat:
+            x = torch.cat((x, rem), 1)
+        else:
+            x = x + rem
+        x = self.conv2(x)
+        return x
+
+
+class BasicConv_IN(nn.Module):
+
+    def __init__(self, in_channels, out_channels, deconv=False, is_3d=False, IN=True, relu=True, **kwargs):
+        super(BasicConv_IN, self).__init__()
+        if relu:
+          self.relu = nn.LeakyReLU(inplace=True)
+        else:
+          self.relu = nn.Identity()
+        self.use_in = IN
+        if is_3d:
+            if deconv:
+                self.conv = nn.ConvTranspose3d(in_channels, out_channels, bias=False, **kwargs)
+            else:
+                self.conv = nn.Conv3d(in_channels, out_channels, bias=False, **kwargs)
+            self.IN = nn.InstanceNorm3d(out_channels)
+        else:
+            if deconv:
+                self.conv = nn.ConvTranspose2d(in_channels, out_channels, bias=False, **kwargs)
+            else:
+                self.conv = nn.Conv2d(in_channels, out_channels, bias=False, **kwargs)
+            self.IN = nn.InstanceNorm2d(out_channels)
+
+    def forward(self, x):
+        x = self.conv(x)
+        if self.use_in:
+          x = self.IN(x)
+        if isinstance(self.relu, bool):
+          if self.relu:
+            self.relu = nn.LeakyReLU(inplace=True)
+          else:
+            self.relu = nn.Identity()
+        x = self.relu(x)
+        return x
+
+
+class Conv2x_IN(nn.Module):
+    def __init__(self, in_channels, out_channels, c_middle=None, deconv=False, is_3d=False, concat=True, keep_concat=True, IN=True, relu=True, keep_dispc=False):
+        super(Conv2x_IN, self).__init__()
+        self.concat = concat
+        self.is_3d = is_3d
+        if deconv and is_3d:
+            kernel = (4, 4, 4)
+        elif deconv:
+            kernel = 4
+        else:
+            kernel = 3
+        if c_middle is None:
+          c_middle = out_channels
+
+        if deconv and is_3d and keep_dispc:
+            kernel = (1, 4, 4)
+            stride = (1, 2, 2)
+            padding = (0, 1, 1)
+            self.conv1 = BasicConv_IN(in_channels, c_middle, deconv, is_3d, IN=True, relu=True, kernel_size=kernel, stride=stride, padding=padding)
+        else:
+            self.conv1 = BasicConv_IN(in_channels, c_middle, deconv, is_3d, IN=True, relu=True, kernel_size=kernel, stride=2, padding=1)
+
+        if self.concat:
+            mul = 2 if keep_concat else 1
+            self.conv2 = ResnetBasicBlock(out_channels*2, out_channels*mul, kernel_size=3, stride=1, padding=1, norm_layer=nn.InstanceNorm2d)
+        else:
+            self.conv2 = BasicConv_IN(c_middle, out_channels, False, is_3d, IN, relu, kernel_size=3, stride=1, padding=1)
+
+    def forward(self, x, rem):
+        x = self.conv1(x)
+        if x.shape != rem.shape:
+            x = F.interpolate(x, size=(rem.shape[-2], rem.shape[-1]), mode='bilinear')
+        if self.concat:
+            x = torch.cat((x, rem), 1)
+        else:
+            x = x + rem
+        x = self.conv2(x)
+        return x
+
+
+
+@torch.compile
+def build_gwc_volume_optimized_pytorch1(refimg_fea: torch.Tensor, targetimg_fea: torch.Tensor, maxdisp: int, num_groups: int, normalize=True):
+  dtype = refimg_fea.dtype
+  B, C, H, W = refimg_fea.shape
+  channels_per_group = C // num_groups
+
+  ref_volume = refimg_fea.unsqueeze(2).expand(B, C, maxdisp, H, W)
+  padded_target = F.pad(targetimg_fea, (maxdisp - 1, 0, 0, 0))
+  unfolded_target = padded_target.unfold(3, W, 1)
+  target_volume = torch.flip(unfolded_target, [3]).permute(0, 1, 3, 2, 4)
+  ref_volume = ref_volume.view(B, num_groups, channels_per_group, maxdisp, H, W)
+  target_volume = target_volume.view(B, num_groups, channels_per_group, maxdisp, H, W)
+  if normalize:
+    ref_volume = F.normalize(ref_volume.float(), dim=2).to(dtype)
+    target_volume = F.normalize(target_volume.float(), dim=2).to(dtype)
+
+  cost_volume = (ref_volume * target_volume).sum(dim=2)
+
+  return cost_volume.contiguous()
+
+
+if triton is not None and torch.cuda.is_available():
+  @triton.autotune(configs=[
+    triton.Config({'BLOCK_C':4,'BLOCK_W':128,'BLOCK_D':8}, num_warps=4, num_stages=2),
+    triton.Config({'BLOCK_C':8,'BLOCK_W':128,'BLOCK_D':8}, num_warps=4, num_stages=2),
+    triton.Config({'BLOCK_C':16,'BLOCK_W':128,'BLOCK_D':8}, num_warps=4, num_stages=2),
+    triton.Config({'BLOCK_C':64,'BLOCK_W':128,'BLOCK_D':8}, num_warps=8, num_stages=2),
+    triton.Config({'BLOCK_C':128,'BLOCK_W':64,'BLOCK_D':8}, num_warps=8, num_stages=2),
+    triton.Config({'BLOCK_C':128,'BLOCK_W':128,'BLOCK_D':8}, num_warps=8, num_stages=2),
+  ], key=['C','W','D','G','K','NORMALIZE'])
+  @triton.jit
+  def _gwc_triton_kernel(ref_ptr, tar_ptr, ref_norm_ptr, tar_norm_ptr, out_ptr, BH, C, W, D: tl.constexpr, G: tl.constexpr, K: tl.constexpr,
+                         stride_rn, stride_rw, stride_rc, stride_tn, stride_tw, stride_tc,
+                         stride_nn, stride_ng, stride_nw,
+                         stride_on, stride_og, stride_od, stride_ow,
+                         NORMALIZE: tl.constexpr,
+                         BLOCK_C: tl.constexpr, BLOCK_W: tl.constexpr, BLOCK_D: tl.constexpr):
+    pid0 = tl.program_id(0)
+    db = tl.program_id(1)
+    wb = tl.program_id(2)
+    bh = pid0 // G
+    g = pid0 % G
+    w_off = wb*BLOCK_W + tl.arange(0, BLOCK_W)
+    d_off = db*BLOCK_D + tl.arange(0, BLOCK_D)
+    w_mask = w_off < W
+    w_src = w_off[None, :] - d_off[:, None]
+    td_mask = (w_src >= 0) & w_mask[None, :]
+    acc = tl.zeros((BLOCK_D, BLOCK_W), dtype=tl.float32)
+    for k0 in tl.static_range(0, K, BLOCK_C):
+      k_off = k0 + tl.arange(0, BLOCK_C)
+      k_mask = k_off < K
+      c_idx = g*K + k_off
+      ref_ptrs = ref_ptr + bh*stride_rn + w_off[None, :]*stride_rw + c_idx[:, None]*stride_rc
+      ref_vals = tl.load(ref_ptrs, mask=k_mask[:, None] & w_mask[None, :], other=0.).to(tl.float32)
+      tar_ptrs = tar_ptr + bh*stride_tn + w_src[None, :, :]*stride_tw + c_idx[:, None, None]*stride_tc
+      tar_vals = tl.load(tar_ptrs, mask=k_mask[:, None, None] & td_mask[None, :, :], other=0.).to(tl.float32)
+      acc += tl.sum(tar_vals * ref_vals[:, None, :], axis=0)
+
+    if NORMALIZE:
+      norm_offset = bh*stride_nn + g*stride_ng
+      ref_norm = tl.load(ref_norm_ptr + norm_offset + w_off*stride_nw, mask=w_mask, other=1.0).to(tl.float32)
+      tar_norm = tl.load(tar_norm_ptr + norm_offset + w_src*stride_nw, mask=td_mask, other=1.0).to(tl.float32)
+      denom = (ref_norm[None, :] * tar_norm) + 1e-5
+      acc = acc / denom
+    out_ptrs = out_ptr + bh*stride_on + g*stride_og + d_off[:, None]*stride_od + w_off[None, :]*stride_ow
+    tl.store(out_ptrs, acc, mask=w_mask[None, :])
+
+@torch.no_grad()
+def build_gwc_volume_triton(refimg_fea: torch.Tensor, targetimg_fea: torch.Tensor, maxdisp: int, num_groups: int, normalize=True):
+  if triton is None:
+    raise RuntimeError('Triton is not available. Please install triton to use build_gwc_volume_triton.')
+  B, C, H, W = refimg_fea.shape
+  assert maxdisp > 0 and C % num_groups == 0
+  K = C // num_groups
+  in_dtype = refimg_fea.dtype if refimg_fea.dtype in (torch.float16, torch.bfloat16, torch.float32) else torch.float32
+
+  if normalize:
+    ref_norm = refimg_fea.float().view(B, num_groups, K, H, W).norm(dim=2)
+    tar_norm = targetimg_fea.float().view(B, num_groups, K, H, W).norm(dim=2)
+    ref_norm = ref_norm.permute(0, 2, 1, 3).reshape(B*H, num_groups, W).to(in_dtype).contiguous()
+    tar_norm = tar_norm.permute(0, 2, 1, 3).reshape(B*H, num_groups, W).to(in_dtype).contiguous()
+  else:
+    # Dummy tensors; kernel won't read them when NORMALIZE=False
+    ref_norm = refimg_fea.new_empty((1, 1, 1), dtype=in_dtype)
+    tar_norm = refimg_fea.new_empty((1, 1, 1), dtype=in_dtype)
+
+  ref = refimg_fea.to(in_dtype)
+  tar = targetimg_fea.to(in_dtype)
+  ref_bhwc = ref.permute(0, 2, 3, 1).view(B * H, W, C).contiguous()
+  tar_bhwc = tar.permute(0, 2, 3, 1).view(B * H, W, C).contiguous()
+  out_bhw = torch.empty((B * H, num_groups, maxdisp, W), device=ref.device, dtype=in_dtype)
+  BH = B * H
+  D_eff = min(maxdisp, W)
+  grid = lambda META: (BH * num_groups, triton.cdiv(D_eff, META['BLOCK_D']), triton.cdiv(W, META['BLOCK_W']))
+  _gwc_triton_kernel[grid](ref_bhwc, tar_bhwc, ref_norm, tar_norm, out_bhw, BH, C, W, D_eff, num_groups, K,
+                           ref_bhwc.stride(0), ref_bhwc.stride(1), ref_bhwc.stride(2),
+                           tar_bhwc.stride(0), tar_bhwc.stride(1), tar_bhwc.stride(2),
+                           ref_norm.stride(0), ref_norm.stride(1), ref_norm.stride(2),
+                           out_bhw.stride(0), out_bhw.stride(1), out_bhw.stride(2), out_bhw.stride(3),
+                           NORMALIZE=normalize)
+  if D_eff < maxdisp: out_bhw[:, :, D_eff:, :] = 0
+  volume = out_bhw.view(B, H, num_groups, maxdisp, W).permute(0, 2, 3, 1, 4).contiguous()
+  return volume
+
+
+
+@torch.compile
+def build_concat_volume_optimized_pytorch(refimg_fea, targetimg_fea, maxdisp:int):
+  B, C, H, W = refimg_fea.shape
+  ref_volume = refimg_fea.unsqueeze(2).expand(B, C, maxdisp, H, W)
+  shifted_target_list = [F.pad(targetimg_fea, (int(d), 0, 0, 0), "constant", 0.0)[:, :, :, :W] for d in range(maxdisp)]
+  target_volume = torch.stack(shifted_target_list, dim=2)
+  volume = torch.cat((ref_volume, target_volume), dim=1)
+  return volume.contiguous()
+
+
+@torch.compile
+def build_concat_volume_optimized_pytorch1(refimg_fea, targetimg_fea, maxdisp:int):
+    B, C, H, W = refimg_fea.shape
+
+    ref_volume = refimg_fea.unsqueeze(2).expand(B, C, maxdisp, H, W)
+    padded_target = F.pad(targetimg_fea, (maxdisp - 1, 0, 0, 0))  # (B, C, H, W + maxdisp - 1)
+    unfolded_target = padded_target.unfold(dimension=3, size=W, step=1)  # (B, C, H, maxdisp, W)
+    target_volume = torch.flip(unfolded_target, [3]).permute(0, 1, 3, 2, 4)
+    volume = torch.cat((ref_volume, target_volume), dim=1)
+    return volume.contiguous()
+
+
+
+
+def disparity_regression(x, maxdisp):
+    assert len(x.shape) == 4
+    disp_values = torch.arange(0, maxdisp, dtype=x.dtype, device=x.device)
+    disp_values = disp_values.reshape(1, maxdisp, 1, 1)
+    return torch.sum(x * disp_values, 1, keepdim=True)   #(B,1,H,W)
+
+
+class FeatureAtt(nn.Module):
+    def __init__(self, cv_chan, feat_chan):
+        super(FeatureAtt, self).__init__()
+
+        self.feat_att = nn.Sequential(
+            BasicConv(feat_chan, feat_chan//2, kernel_size=1, stride=1, padding=0),
+            nn.Conv2d(feat_chan//2, cv_chan, 1)
+            )
+
+    def forward(self, cv, feat):
+        '''
+        @cv: cost volume (B,C,D,H,W)
+        @feat: (B,C,H,W)
+        '''
+        feat_att = self.feat_att(feat).unsqueeze(2)   #(B,C,1,H,W)
+        cv = torch.sigmoid(feat_att)*cv
+        return cv
+
+def context_upsample(disp_low, up_weights):
+    """
+    @disp_low: (b,1,h,w)  1/4 resolution
+    @up_weights: (b,9,4*h,4*w)  Image resolution
+    """
+    b, c, h, w = disp_low.shape
+
+    disp_unfold = F.unfold(disp_low.reshape(b,c,h,w),3,1,1).reshape(b,-1,h,w)
+    disp_unfold = F.interpolate(disp_unfold,(h*4,w*4),mode='nearest').reshape(b,9,h*4,w*4)
+
+    disp = (disp_unfold*up_weights).sum(1)
+
+    return disp
+
+
+
+class PositionalEmbedding(nn.Module):
+  def __init__(self, d_model, max_len=512):
+    super().__init__()
+
+    # Compute the positional encodings once in log space.
+    pe = torch.zeros(max_len, d_model, dtype=torch.float)
+    pe.require_grad = False
+
+    position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)  #(N,1)
+    div_term = (torch.arange(0, d_model, 2, dtype=torch.float) * -(np.log(10000.0) / d_model)).exp()[None]
+
+    pe[:, 0::2] = torch.sin(position * div_term)  #(N, d_model/2)
+    pe[:, 1::2] = torch.cos(position * div_term)
+
+    pe = pe.unsqueeze(0)
+    self.pe = pe
+
+
+  def forward(self, x, resize_embed=False):
+    '''
+    @x: (B,N,D)
+    '''
+    dtype = x.dtype
+    self.pe = self.pe.to(x.device).to(x.dtype)
+    pe = self.pe
+    if pe.shape[1]<x.shape[1]:
+      if resize_embed:
+        pe = F.interpolate(pe.permute(0,2,1), size=x.shape[1], mode='linear', align_corners=True).permute(0,2,1)
+      else:
+        raise RuntimeError(f'x:{x.shape}, pe:{pe.shape}')
+    return (x + pe[:, :x.size(1)]).to(dtype)
+
+
+
+class CostVolumeDisparityAttention(nn.Module):
+  def __init__(self, d_model, nhead, dim_feedforward, dropout=0.1, act=nn.GELU, norm_first=False, num_transformer=6, max_len=512, resize_embed=False):
+    super().__init__()
+    self.resize_embed = resize_embed
+    self.sa = nn.ModuleList([])
+    for _ in range(num_transformer):
+      self.sa.append(FlashAttentionTransformerEncoderLayer(embed_dim=d_model, num_heads=nhead, dim_feedforward=dim_feedforward, act=act, dropout=dropout))
+    self.pos_embed0 = PositionalEmbedding(d_model, max_len=max_len)
+
+
+  def forward(self, cv, window_size=(-1,-1)):
+    """
+    @cv: (B,C,D,H,W) where D is max disparity
+    """
+    x = cv
+    B,C,D,H,W = x.shape
+    x = x.permute(0,3,4,2,1).reshape(B*H*W, D, C)
+    x = self.pos_embed0(x, resize_embed=self.resize_embed)  #!NOTE No resize since disparity is pre-determined
+    for i in range(len(self.sa)):
+        x = self.sa[i](x, window_size=window_size)
+    x = x.reshape(B,H,W,D,C).permute(0,4,3,1,2)
+
+    return x
+
+
+
+class ChannelAttentionEnhancement(nn.Module):
+    def __init__(self, in_planes, ratio=16):
+        """From selective-IGEV
+        """
+        super(ChannelAttentionEnhancement, self).__init__()
+        self.avg_pool = nn.AdaptiveAvgPool2d(1)
+        self.max_pool = nn.AdaptiveMaxPool2d(1)
+
+        self.fc = nn.Sequential(nn.Conv2d(in_planes, in_planes // 16, 1, bias=False),
+                               nn.ReLU(),
+                               nn.Conv2d(in_planes // 16, in_planes, 1, bias=False))
+        self.sigmoid = nn.Sigmoid()
+
+    def forward(self, x):
+        avg_out = self.fc(self.avg_pool(x))
+        max_out = self.fc(self.max_pool(x))
+        out = avg_out + max_out
+        return self.sigmoid(out)
+
+class SpatialAttentionExtractor(nn.Module):
+    def __init__(self, kernel_size=7):
+        """From selective-IGEV
+        """
+        super(SpatialAttentionExtractor, self).__init__()
+
+        self.samconv = nn.Conv2d(2, 1, kernel_size, padding=kernel_size//2, bias=False)
+        self.sigmoid = nn.Sigmoid()
+
+    def forward(self, x):
+        avg_out = torch.mean(x, dim=1, keepdim=True)
+        max_out, _ = torch.max(x, dim=1, keepdim=True)
+        x = torch.cat([avg_out, max_out], dim=1)
+        x = self.samconv(x)
+        return self.sigmoid(x)
+
+
+
+class EdgeNextConvEncoder(nn.Module):
+    def __init__(self, dim, layer_scale_init_value=1e-6, expan_ratio=4, kernel_size=7, norm='layer'):
+        """https://github.com/mmaaz60/EdgeNeXt/blob/main/models/conv_encoder.py#L7
+        """
+        super().__init__()
+        self.dwconv = nn.Conv2d(dim, dim, kernel_size=kernel_size, padding=kernel_size // 2, groups=dim)
+        if norm=='layer':
+          self.norm = LayerNorm2d(dim, eps=1e-6)
+        elif norm=='batch':
+          self.norm = nn.BatchNorm2d(dim)
+        else:
+          self.norm = nn.Identity()
+        self.pwconv1 = nn.Linear(dim, expan_ratio * dim)
+        self.act = nn.GELU()
+        self.pwconv2 = nn.Linear(expan_ratio * dim, dim)
+        self.gamma = nn.Parameter(layer_scale_init_value * torch.ones(dim), requires_grad=True) if layer_scale_init_value > 0 else None
+
+    def forward(self, x):
+        input = x
+        x = self.dwconv(x)
+        x = self.norm(x)
+        x = x.permute(0, 2, 3, 1)  # (N, C, H, W) -> (N, H, W, C)
+        x = self.pwconv1(x)
+        x = self.act(x)
+        x = self.pwconv2(x)
+        if self.gamma is not None:
+            x = self.gamma * x
+        x = x.permute(0, 3, 1, 2)  # (N, H, W, C) -> (N, C, H, W)
+
+        x = input + x
+        return x
+

cv/fast-foundationstereo/core/update.py

@@ -0,0 +1,108 @@
+import torch,os,sys
+import torch.nn as nn
+import torch.nn.functional as F
+code_dir = os.path.dirname(os.path.realpath(__file__))
+sys.path.append(f'{code_dir}/../')
+from core.submodule import EdgeNextConvEncoder
+
+
+class DispHead(nn.Module):
+    def __init__(self, input_dim=128, hidden_dim=256, output_dim=1):
+        super(DispHead, self).__init__()
+        self.conv = nn.Sequential(
+          nn.Conv2d(input_dim, input_dim, kernel_size=3, padding=1),
+          nn.ReLU(),
+          EdgeNextConvEncoder(input_dim, expan_ratio=4, kernel_size=7, norm=None),
+          EdgeNextConvEncoder(input_dim, expan_ratio=4, kernel_size=7, norm=None),
+          nn.Conv2d(input_dim, output_dim, 3, padding=1),
+        )
+
+    def forward(self, x):
+        return self.conv(x)
+
+
+class BasicMotionEncoder(nn.Module):
+    def __init__(self, args, ngroup=8):
+        super(BasicMotionEncoder, self).__init__()
+        self.args = args
+        cor_planes = args.corr_levels * (2*args.corr_radius + 1) * (ngroup+1)
+        self.convc1 = nn.Conv2d(cor_planes, 256, kernel_size=1, padding=0)
+        self.convc2 = nn.Conv2d(256, 256, kernel_size=3, padding=1)
+        self.convd1 = nn.Conv2d(1, 64, kernel_size=7, padding=3)
+        self.convd2 = nn.Conv2d(64, 64, kernel_size=3, padding=1)
+        self.conv = nn.Conv2d(64+256, args.hidden_dims[0]-1, kernel_size=1, padding=0)
+
+    def forward(self, disp, corr):
+        cor = F.relu(self.convc1(corr))
+        cor = F.relu(self.convc2(cor))
+        disp_ = F.relu(self.convd1(disp))
+        disp_ = F.relu(self.convd2(disp_))
+
+        cor_disp = torch.cat([cor, disp_], dim=1)
+        out = F.relu(self.conv(cor_disp))
+        return torch.cat([out, disp], dim=1)
+
+class RaftConvGRU(nn.Module):
+    def __init__(self, hidden_dim=128, input_dim=256, kernel_size=3):
+        super().__init__()
+        self.convz = nn.Conv2d(hidden_dim+input_dim, hidden_dim, kernel_size, padding=kernel_size // 2)
+        self.convr = nn.Conv2d(hidden_dim+input_dim, hidden_dim, kernel_size, padding=kernel_size // 2)
+        self.convq = nn.Conv2d(hidden_dim+input_dim, hidden_dim, kernel_size, padding=kernel_size // 2)
+
+    def forward(self, h, x, hx):
+        z = torch.sigmoid(self.convz(hx))
+        r = torch.sigmoid(self.convr(hx))
+        q = torch.tanh(self.convq(torch.cat([r*h, x], dim=1)))
+        h = (1-z) * h + z * q
+        return h
+
+
+class SelectiveConvGRU(nn.Module):
+    def __init__(self, hidden_dim=128, input_dim=256, small_kernel_size=1, large_kernel_size=3, patch_size=None):
+        super(SelectiveConvGRU, self).__init__()
+        self.conv0 = nn.Sequential(
+            nn.Conv2d(input_dim, input_dim, kernel_size=3, padding=1),
+            nn.ReLU(),
+        )
+        self.conv1 = nn.Sequential(
+            nn.Conv2d(input_dim+hidden_dim, input_dim+hidden_dim, kernel_size=3, padding=1),
+            nn.ReLU(),
+        )
+        self.small_gru = RaftConvGRU(hidden_dim, input_dim, small_kernel_size)
+        self.large_gru = RaftConvGRU(hidden_dim, input_dim, large_kernel_size)
+
+    def forward(self, att, h, *x):
+        x = torch.cat(x, dim=1)
+        x = self.conv0(x)
+        hx = torch.cat([x, h], dim=1)
+        hx = self.conv1(hx)
+        h = self.small_gru(h, x, hx) * att + self.large_gru(h, x, hx) * (1 - att)
+
+        return h
+
+
+class BasicSelectiveMultiUpdateBlock(nn.Module):
+    def __init__(self, args, hidden_dim=128, volume_dim=8):
+        super().__init__()
+        self.args = args
+        self.encoder = BasicMotionEncoder(args, volume_dim)
+
+        self.gru04 = SelectiveConvGRU(hidden_dim, hidden_dim*2)
+        self.disp_head = DispHead(hidden_dim, 256)
+        self.mask = nn.Sequential(
+            nn.Conv2d(hidden_dim, 64, 3, padding=1),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(64, 32, 3, padding=1),
+            nn.ReLU(inplace=True),
+            )
+
+    def forward(self, net, inp, corr, disp, att):
+        motion_features = self.encoder(disp, corr)
+        motion_features = torch.cat([inp[0], motion_features], dim=1)
+        net[0] = self.gru04(att[0], net[0], motion_features)
+
+        delta_disp = self.disp_head(net[0])
+
+        mask = .25 * self.mask(net[0])
+        return net, mask, delta_disp
+

cv/fast-foundationstereo/core/utils/frame_utils.py

@@ -0,0 +1,202 @@
+import numpy as np
+from PIL import Image
+from os.path import basename, exists, splitext
+import os,sys
+code_dir = os.path.dirname(os.path.realpath(__file__))
+sys.path.append(f'{code_dir}/../../')
+import re
+import json
+import imageio
+import cv2
+from turbojpeg import TurboJPEG, TJPF_GRAY, TJSAMP_GRAY, TJFLAG_PROGRESSIVE, TJFLAG_FASTUPSAMPLE, TJFLAG_FASTDCT
+jpeg = TurboJPEG()
+cv2.setNumThreads(0)
+cv2.ocl.setUseOpenCL(False)
+
+TAG_CHAR = np.array([202021.25], np.float32)
+
+def readFlow(fn):
+    """ Read .flo file in Middlebury format"""
+    # Code adapted from:
+    # http://stackoverflow.com/questions/28013200/reading-middlebury-flow-files-with-python-bytes-array-numpy
+
+    # WARNING: this will work on little-endian architectures (eg Intel x86) only!
+    # print 'fn = %s'%(fn)
+    with open(fn, 'rb') as f:
+        magic = np.fromfile(f, np.float32, count=1)
+        if 202021.25 != magic:
+            print('Magic number incorrect. Invalid .flo file')
+            return None
+        else:
+            w = np.fromfile(f, np.int32, count=1)
+            h = np.fromfile(f, np.int32, count=1)
+            # print 'Reading %d x %d flo file\n' % (w, h)
+            data = np.fromfile(f, np.float32, count=2*int(w)*int(h))
+            # Reshape data into 3D array (columns, rows, bands)
+            # The reshape here is for visualization, the original code is (w,h,2)
+            return np.resize(data, (int(h), int(w), 2))
+
+def readPFM(file):
+    file = open(file, 'rb')
+
+    color = None
+    width = None
+    height = None
+    scale = None
+    endian = None
+
+    header = file.readline().rstrip()
+    if header == b'PF':
+        color = True
+    elif header == b'Pf':
+        color = False
+    else:
+        raise Exception('Not a PFM file.')
+
+    dim_match = re.match(rb'^(\d+)\s(\d+)\s$', file.readline())
+    if dim_match:
+        width, height = map(int, dim_match.groups())
+    else:
+        raise Exception('Malformed PFM header.')
+
+    scale = float(file.readline().rstrip())
+    if scale < 0: # little-endian
+        endian = '<'
+        scale = -scale
+    else:
+        endian = '>' # big-endian
+
+    data = np.fromfile(file, endian + 'f')
+    shape = (height, width, 3) if color else (height, width)
+
+    data = np.reshape(data, shape)
+    data = np.flipud(data)
+    return data
+
+def writePFM(file, array):
+    import os
+    assert type(file) is str and type(array) is np.ndarray and \
+           os.path.splitext(file)[1] == ".pfm"
+    with open(file, 'wb') as f:
+        H, W = array.shape
+        headers = ["Pf\n", f"{W} {H}\n", "-1\n"]
+        for header in headers:
+            f.write(str.encode(header))
+        array = np.flip(array, axis=0).astype(np.float32)
+        f.write(array.tobytes())
+
+
+
+def writeFlow(filename,uv,v=None):
+    """ Write optical flow to file.
+
+    If v is None, uv is assumed to contain both u and v channels,
+    stacked in depth.
+    Original code by Deqing Sun, adapted from Daniel Scharstein.
+    """
+    nBands = 2
+
+    if v is None:
+        assert(uv.ndim == 3)
+        assert(uv.shape[2] == 2)
+        u = uv[:,:,0]
+        v = uv[:,:,1]
+    else:
+        u = uv
+
+    assert(u.shape == v.shape)
+    height,width = u.shape
+    f = open(filename,'wb')
+    # write the header
+    f.write(TAG_CHAR)
+    np.array(width).astype(np.int32).tofile(f)
+    np.array(height).astype(np.int32).tofile(f)
+    # arrange into matrix form
+    tmp = np.zeros((height, width*nBands))
+    tmp[:,np.arange(width)*2] = u
+    tmp[:,np.arange(width)*2 + 1] = v
+    tmp.astype(np.float32).tofile(f)
+    f.close()
+
+
+def readFlowKITTI(filename):
+    flow = cv2.imread(filename, cv2.IMREAD_ANYDEPTH|cv2.IMREAD_COLOR)
+    flow = flow[:,:,::-1].astype(np.float32)
+    flow, valid = flow[:, :, :2], flow[:, :, 2]
+    flow = (flow - 2**15) / 64.0
+    return flow, valid
+
+def readDispKITTI(filename):
+    disp = cv2.imread(filename, cv2.IMREAD_ANYDEPTH) / 256.0
+    valid = disp > 0.0
+    return disp, valid
+
+# Method taken from /n/fs/raft-depth/RAFT-Stereo/datasets/SintelStereo/sdk/python/sintel_io.py
+def readDispSintelStereo(file_name):
+    a = np.array(Image.open(file_name))
+    d_r, d_g, d_b = np.split(a, axis=2, indices_or_sections=3)
+    disp = (d_r * 4 + d_g / (2**6) + d_b / (2**14))[..., 0]
+    mask = np.array(Image.open(file_name.replace('disparities', 'occlusions')))
+    valid = ((mask == 0) & (disp > 0))
+    return disp, valid
+
+# Method taken from https://research.nvidia.com/sites/default/files/pubs/2018-06_Falling-Things/readme_0.txt
+def readDispFallingThings(file_name):
+    a = np.array(Image.open(file_name))
+    with open('/'.join(file_name.split('/')[:-1] + ['_camera_settings.json']), 'r') as f:
+        intrinsics = json.load(f)
+    fx = intrinsics['camera_settings'][0]['intrinsic_settings']['fx']
+    disp = (fx * 6.0 * 100) / a.astype(np.float32)
+    valid = disp > 0
+    return disp, valid
+
+# Method taken from https://github.com/castacks/tartanair_tools/blob/master/data_type.md
+def readDispTartanAir(file_name):
+    depth = np.load(file_name)
+    disp = 80.0 / depth
+    valid = disp > 0
+    return disp, valid
+
+
+def readDispMiddlebury(file_name):
+    assert basename(file_name) == 'disp0GT.pfm'
+    disp = readPFM(file_name).astype(np.float32)
+    assert len(disp.shape) == 2
+    nocc_pix = file_name.replace('disp0GT.pfm', 'mask0nocc.png')
+    assert exists(nocc_pix)
+    nocc_pix = imageio.imread(nocc_pix) == 255
+    assert np.any(nocc_pix)
+    return disp, nocc_pix
+
+def writeFlowKITTI(filename, uv):
+    uv = 64.0 * uv + 2**15
+    valid = np.ones([uv.shape[0], uv.shape[1], 1])
+    uv = np.concatenate([uv, valid], axis=-1).astype(np.uint16)
+    cv2.imwrite(filename, uv[..., ::-1])
+
+
+def read_gen(file_name, pil=False):
+    ext = splitext(file_name)[-1]
+    if ext in ['.jpeg','.jpg']:
+      with open(file_name, 'rb') as ff:
+        bgr_array = jpeg.decode(ff.read())
+      img = bgr_array[...,::-1]
+      return img
+    elif ext == '.png' or ext == '.ppm':
+        img = cv2.imread(file_name)[...,:3]
+        if len(img.shape)==3:
+          img = img[...,::-1]
+        elif len(img.shape)==2:
+          img = np.tile(img[...,None], (1,1,3))
+        return img
+    elif ext == '.bin' or ext == '.raw':
+        return np.load(file_name)
+    elif ext == '.flo':
+        return readFlow(file_name).astype(np.float32)
+    elif ext == '.pfm':
+        flow = readPFM(file_name).astype(np.float32)
+        if len(flow.shape) == 2:
+            return flow
+        else:
+            return flow[:, :, :-1]
+    return []

cv/fast-foundationstereo/core/utils/utils.py

@@ -0,0 +1,122 @@
+import torch,os,sys
+code_dir = os.path.dirname(os.path.realpath(__file__))
+sys.path.append(f'{code_dir}/../../')
+import torch.nn.functional as F
+import numpy as np
+
+
+class InputPadder:
+    """ Pads images such that dimensions are divisible by 8 """
+    def __init__(self, dims, mode='sintel', divis_by=8, force_square=False):
+        self.ht, self.wd = dims[-2:]
+        if force_square:
+          max_side = max(self.ht, self.wd)
+          pad_ht = ((max_side // divis_by) + 1) * divis_by - self.ht
+          pad_wd = ((max_side // divis_by) + 1) * divis_by - self.wd
+        else:
+          pad_ht = (((self.ht // divis_by) + 1) * divis_by - self.ht) % divis_by
+          pad_wd = (((self.wd // divis_by) + 1) * divis_by - self.wd) % divis_by
+        if mode == 'sintel':
+            self._pad = [pad_wd//2, pad_wd - pad_wd//2, pad_ht//2, pad_ht - pad_ht//2]
+        else:
+            self._pad = [pad_wd//2, pad_wd - pad_wd//2, 0, pad_ht]
+
+    def pad(self, *inputs):
+        assert all((x.ndim == 4) for x in inputs)
+        return [F.pad(x, self._pad, mode='replicate') for x in inputs]
+
+    def unpad(self, x):
+        assert x.ndim == 4
+        ht, wd = x.shape[-2:]
+        c = [self._pad[2], ht-self._pad[3], self._pad[0], wd-self._pad[1]]
+        return x[..., c[0]:c[1], c[2]:c[3]]
+
+
+@torch.compile
+def bilinear_sampler1d(img, x_coords, mode='bilinear', align_corners=True):
+    """
+    1D bilinear sampling along width dimension only (for stereo applications)
+    Much faster than grid_sample for stereo where y is constant
+
+    Args:
+        img: (B, C, 1, W) input tensor
+        x_coords: (B, 1, W_out, 1) x coordinates in pixel space [0, W-1]
+        mode: interpolation mode ('bilinear' or 'nearest')
+        align_corners: if True, corner pixels are aligned (like grid_sample)
+
+    Returns:
+        sampled: (B, C, 1, W_coords) sampled tensor
+        mask: (B, 1, H, W) validity mask (if mask=True)
+    """
+    B, C, H_img, W = img.shape
+    x = x_coords.reshape(B,-1) # (B, W_out)
+
+    if align_corners:
+        # align_corners=True: coordinate range [0, W-1] maps to pixel centers
+        # This matches grid_sample with align_corners=True behavior
+        x_normalized = x
+    else:
+        # align_corners=False: coordinate range [0, W-1] maps to pixel edges
+        # Need to adjust coordinates to match grid_sample with align_corners=False
+        # grid_sample maps [-1, 1] to [0, W-1] when align_corners=False
+        # So our [0, W-1] input should be treated as [0.5, W-0.5] in pixel space
+        x_normalized = x + 0.5
+
+    if mode == 'nearest':
+        # Nearest neighbor sampling with zero padding outside [0, W-1]
+        if align_corners:
+            x_nearest = torch.round(x_normalized).long()
+        else:
+            x_nearest = torch.floor(x_normalized).long()
+        valid = (x_nearest >= 0) & (x_nearest < W)  # (B, W_out)
+        x_index = torch.clamp(x_nearest, 0, W-1)
+        sampled = torch.gather(img, 3, x_index.view(B,1,1,-1).expand(B,C,1,-1))
+        sampled = sampled * valid.view(B,1,1,-1).to(img.dtype)
+
+    else:  # bilinear
+        # Get integer and fractional parts
+        x_floor = torch.floor(x_normalized)
+        x_ceil = x_floor + 1
+        x_frac = x_normalized - x_floor  # (B, W_out)
+
+        # Zero padding behavior: mark validity and zero-out invalid contributions
+        valid_floor = (x_floor >= 0) & (x_floor < W)
+        valid_ceil = (x_ceil >= 0) & (x_ceil < W)
+        x_floor_clamped = torch.clamp(x_floor, 0, W-1)
+        x_ceil_clamped = torch.clamp(x_ceil, 0, W-1)
+
+        # Create index tensors
+        batch_idx = torch.arange(B, device=img.device).view(B, 1)
+        img_floor = torch.gather(img, 3, x_floor_clamped.view(B,1,1,-1).expand(B,C,1,-1).long())
+        img_ceil = torch.gather(img, 3, x_ceil_clamped.view(B,1,1,-1).expand(B,C,1,-1).long())
+
+        # Apply validity masks (zero out-of-bounds samples)
+        img_floor = img_floor * valid_floor.view(B,1,1,-1).to(img.dtype)
+        img_ceil = img_ceil * valid_ceil.view(B,1,1,-1).to(img.dtype)
+
+        # Linear interpolation
+        x_frac = x_frac.view(B,1,1,-1)
+        sampled = img_floor * (1 - x_frac) + img_ceil * x_frac
+
+    return sampled
+
+
+def bilinear_sampler(img, coords, mode='bilinear', mask=False, low_memory=False, use1d=False):
+    """ Wrapper for grid_sample, uses pixel coordinates """
+    H, W = img.shape[-2:]
+    coords[...,0] = 2*coords[...,0]/(W-1) - 1
+    if low_memory:
+      B = img.shape[0]
+      out = []
+      bs = 102400
+      for b in np.arange(0,B,bs):
+        tmp = F.grid_sample(img[b:b+bs], coords[b:b+bs], align_corners=True)
+        out.append(tmp)
+      img = torch.cat(out, dim=0)
+    else:
+      img = F.grid_sample(img, coords, align_corners=True)
+    if mask:
+        mask = (xgrid > -1) & (ygrid > -1) & (xgrid < 1) & (ygrid < 1)
+        return img, mask.float()
+    return img
+

cv/fast-foundationstereo/cpp/CMakeLists.txt

@@ -0,0 +1,156 @@
+cmake_minimum_required(VERSION 3.18)
+
+project(fast_foundation_stereo_cpp LANGUAGES CXX CUDA)
+
+set(CMAKE_CXX_STANDARD 17)
+set(CMAKE_CXX_STANDARD_REQUIRED ON)
+set(CMAKE_CUDA_STANDARD 17)
+set(CMAKE_CUDA_STANDARD_REQUIRED ON)
+
+if(NOT CMAKE_BUILD_TYPE)
+  set(CMAKE_BUILD_TYPE "Release" CACHE STRING
+      "Build type: Debug, Release, RelWithDebInfo" FORCE)
+endif()
+
+if(NOT DEFINED CMAKE_CUDA_ARCHITECTURES)
+  set(CMAKE_CUDA_ARCHITECTURES "80;86;89;90")
+endif()
+
+add_compile_options(
+  $<$<COMPILE_LANGUAGE:CUDA>:--expt-relaxed-constexpr>
+  $<$<COMPILE_LANGUAGE:CUDA>:--diag-suppress=20015>
+  $<$<COMPILE_LANGUAGE:CUDA>:--diag-suppress=20013>
+  $<$<COMPILE_LANGUAGE:CUDA>:--diag-suppress=20011>
+  $<$<COMPILE_LANGUAGE:CUDA>:--diag-suppress=20091>
+)
+
+find_package(CUDAToolkit REQUIRED)
+
+# ---- TensorRT ----
+set(TENSORRT_ROOT "/usr" CACHE PATH "TensorRT installation path")
+set(TENSORRT_HINT_DIRS
+  ${TENSORRT_ROOT}
+  $ENV{CONDA_PREFIX}
+  /usr
+  /usr/local
+  /usr/local/tensorrt
+  ${CUDAToolkit_TARGET_DIR}
+  /usr/local/cuda
+  /usr/local/cuda/targets/x86_64-linux
+  /usr/local/cuda-13.0
+  /usr/local/cuda-13.0/targets/x86_64-linux
+  /usr/src/tensorrt
+)
+if(DEFINED ENV{CONDA_PREFIX})
+  file(GLOB TENSORRT_PY_LIB_DIRS
+    "$ENV{CONDA_PREFIX}/lib/python*/site-packages/tensorrt_libs"
+  )
+  file(GLOB TENSORRT_PY_INCLUDE_DIRS
+    "$ENV{CONDA_PREFIX}/lib/python*/site-packages/tensorrt/include"
+  )
+endif()
+
+find_library(NVINFER_LIB
+  NAMES nvinfer nvinfer.so.10
+  HINTS
+    ${TENSORRT_HINT_DIRS}
+    ${TENSORRT_PY_LIB_DIRS}
+  PATH_SUFFIXES
+    lib
+    lib64
+    lib/x86_64-linux-gnu
+)
+if(NOT NVINFER_LIB)
+  message(FATAL_ERROR "TensorRT nvinfer library not found. Set -DTENSORRT_ROOT=/path/to/tensorrt")
+endif()
+message(STATUS "TensorRT nvinfer: ${NVINFER_LIB}")
+
+find_library(NVONNXPARSER_LIB
+  NAMES nvonnxparser nvonnxparser.so.10
+  HINTS
+    ${TENSORRT_HINT_DIRS}
+    ${TENSORRT_PY_LIB_DIRS}
+  PATH_SUFFIXES
+    lib
+    lib64
+    lib/x86_64-linux-gnu
+)
+if(NOT NVONNXPARSER_LIB)
+  message(FATAL_ERROR "TensorRT nvonnxparser library not found. Set -DTENSORRT_ROOT=/path/to/tensorrt")
+endif()
+message(STATUS "TensorRT nvonnxparser: ${NVONNXPARSER_LIB}")
+
+find_path(NVINFER_INCLUDE_DIR NvInfer.h
+  HINTS
+    ${TENSORRT_HINT_DIRS}
+    ${TENSORRT_PY_INCLUDE_DIRS}
+  PATH_SUFFIXES
+    include
+    include/x86_64-linux-gnu
+)
+if(NOT NVINFER_INCLUDE_DIR)
+  message(FATAL_ERROR "TensorRT headers (NvInfer.h) not found. Set -DTENSORRT_ROOT=/path/to/tensorrt")
+endif()
+
+find_package(OpenCV QUIET COMPONENTS core imgcodecs imgproc)
+
+add_library(ffs_gwc_plugin SHARED
+  src/gwc_volume_plugin.cpp
+  src/depth_kernels.cu
+)
+set_target_properties(ffs_gwc_plugin PROPERTIES
+  CUDA_SEPARABLE_COMPILATION ON
+  CUDA_RESOLVE_DEVICE_SYMBOLS ON
+  OUTPUT_NAME "ffs_gwc_plugin"
+)
+target_include_directories(ffs_gwc_plugin PUBLIC
+  ${CMAKE_CURRENT_SOURCE_DIR}/include
+  ${NVINFER_INCLUDE_DIR}
+)
+target_link_libraries(ffs_gwc_plugin PUBLIC
+  ${NVINFER_LIB}
+  CUDA::cudart
+)
+
+add_library(ffs_depth_inference STATIC
+  src/ffs_depth_tensorrt.cpp
+  src/ffs_depth_single_tensorrt.cpp
+  src/gwc_volume_plugin.cpp
+  src/depth_kernels.cu
+)
+set_target_properties(ffs_depth_inference PROPERTIES
+  CUDA_SEPARABLE_COMPILATION ON
+  CUDA_RESOLVE_DEVICE_SYMBOLS ON
+)
+target_include_directories(ffs_depth_inference PUBLIC
+  ${CMAKE_CURRENT_SOURCE_DIR}/include
+  ${NVINFER_INCLUDE_DIR}
+)
+target_link_libraries(ffs_depth_inference PUBLIC
+  ${NVINFER_LIB}
+  CUDA::cudart
+)
+
+if(OpenCV_FOUND)
+  add_executable(ffs_depth_main app/main.cpp)
+  target_link_libraries(ffs_depth_main PRIVATE
+    ffs_depth_inference
+    ${OpenCV_LIBS}
+  )
+else()
+  message(WARNING "OpenCV C++ development package not found; skipping ffs_depth_main and ffs_profile_speed")
+endif()
+
+add_executable(ffs_build_single_engine app/build_single_engine.cpp)
+target_link_libraries(ffs_build_single_engine PRIVATE
+  ffs_depth_inference
+  ${NVONNXPARSER_LIB}
+)
+
+if(OpenCV_FOUND)
+  add_executable(ffs_profile_speed app/profile_speed.cpp)
+  target_link_libraries(ffs_profile_speed PRIVATE
+    ffs_depth_inference
+    ${OpenCV_LIBS}
+  )
+endif()

cv/fast-foundationstereo/cpp/README.md

@@ -0,0 +1,249 @@
+# Fast-FoundationStereo C++ Inference
+
+C++ runtime for Fast-FoundationStereo stereo depth inference on TensorRT.
+
+## Run from inside the C++ Docker container
+
+**All commands in this README are intended to be run from inside the container built from [`docker/dockerfile_cpp`](../docker/dockerfile_cpp).** That image carries the CUDA toolkit, TensorRT 10 runtime + ONNX parser + C++ headers, OpenCV development package, `trtexec`, and a Python environment with PyTorch / ONNX export tooling. The plain `docker/dockerfile` image does **not** ship the C++ TRT headers and will fail to build the C++ targets.
+
+Environment setup:
+
+```bash
+docker build --network host -t ffs -f docker/dockerfile_cpp .
+bash docker/run_container.sh
+```
+
+
+`run_container.sh` mounts the parent of `docker/` (the repo root) into `/workspace`, so `cd /workspace/<repo-name>/cpp` lands you in this folder. Every `cd cpp`, `cpp/build/...`, `python3 scripts/...`, and `trtexec ...` example below assumes that working directory.
+
+```text
+cpp/
++-- CMakeLists.txt
++-- README.md
++-- app/
+|   +-- main.cpp                  -- ffs_depth_main: single image inference + visualization
+|   +-- build_single_engine.cpp   -- ffs_build_single_engine: ONNX (with FFSGWCVolume plugin) -> .engine
+|   +-- profile_speed.cpp         -- ffs_profile_speed: latency profiler for either route
++-- include/
+|   +-- ffs_depth_tensorrt.hpp
+|   +-- ffs_depth_single_tensorrt.hpp
+|   +-- ffs_gwc_plugin.hpp
++-- src/
+    +-- ffs_depth_tensorrt.cpp
+    +-- ffs_depth_single_tensorrt.cpp
+    +-- gwc_volume_plugin.cpp
+    +-- depth_kernels.cu
+```
+
+Two interchangeable inference routes are supported:
+
+- `FFSSingleEngineInference` loads **one** TensorRT engine that contains the `FFSGWCVolume` plugin node.
+- `FFSDepthInference` loads **two** TensorRT engines (feature_runner + post_runner) and computes the GWC cost volume between them with a hand-written CUDA kernel.
+
+`ffs_depth_main` auto-detects the route based on the engine directory contents:
+
+- If the directory contains `fast_foundationstereo.engine`, it uses the single-engine plugin path.
+- Otherwise it uses the two-engine reference path.
+
+Inputs and outputs are the same in both cases:
+
+- Input: a stereo image pair plus an intrinsic file (`demo_data/K.txt` format: 9 floats on line 1 for the 3x3 camera matrix, one float on line 2 for the stereo baseline in meters).
+- Output: float32 disparity (input-pixel units), float32 depth in meters, and PNG visualizations.
+
+## Dependencies
+
+All C++ build dependencies are provided by the [`docker/dockerfile_cpp`](../docker/dockerfile_cpp) image (see [Run from inside the C++ Docker container](#run-from-inside-the-c-docker-container) for how to build and enter it):
+
+- CUDA Toolkit (for `nvcc` and the CUDA runtime).
+- TensorRT 10 runtime, ONNX parser, and C++ headers.
+- OpenCV development package (image I/O and depth visualization in `app/main.cpp`).
+- `trtexec`, used by Route B to build the two-engine TensorRT engines from ONNX.
+
+## Build
+
+```bash
+cd cpp
+cmake -B build
+cmake --build build -j
+```
+
+This produces:
+
+- `build/libffs_gwc_plugin.so` -- the FFSGWCVolume plugin as a shared library (loadable by `trtexec --staticPlugins=...`).
+- `build/libffs_depth_inference.a` -- the inference static library.
+- `build/ffs_build_single_engine` -- single-engine builder.
+- `build/ffs_depth_main` -- demo / inference CLI.
+- `build/ffs_profile_speed` -- latency profiler.
+
+## Route A: Single Engine with FFSGWCVolume Plugin
+
+### A.1 Export ONNX with the plugin node
+
+`scripts/make_plugin_onnx.py` exports one ONNX graph in which the GWC cost volume is represented by an `FFSGWCVolume` custom plugin node (resolved at engine-build time by `libffs_gwc_plugin.so`):
+
+```bash
+python3 scripts/make_plugin_onnx.py \
+  --model_dir weights/23-36-37/model_best_bp2_serialize.pth \
+  --save_path output_plugin_onnx \
+  --height 480 \
+  --width 640 \
+  --valid_iters 8 \
+  --max_disp 192
+```
+
+This writes:
+
+```text
+output_plugin_onnx/
++-- fast_foundationstereo_plugin.onnx
++-- onnx.yaml
+```
+
+### A.2 Build the single TensorRT engine
+
+```bash
+cpp/build/ffs_build_single_engine \
+  output_plugin_onnx/fast_foundationstereo_plugin.onnx \
+  output_plugin_onnx/fast_foundationstereo.engine
+```
+
+By default the engine is built with FP16 enabled. Pass `--fp32` to disable FP16. Pass `--workspace-mb N` to override the workspace (default 4096 MB).
+
+After this step `output_plugin_onnx/` contains everything the runtime needs:
+
+```text
+output_plugin_onnx/
++-- fast_foundationstereo.engine
++-- onnx.yaml
+```
+
+### A.3 Run inference
+
+Run from the repository root:
+
+```bash
+cpp/build/ffs_depth_main \
+  output_plugin_onnx \
+  demo_data/left.png \
+  demo_data/right.png \
+  demo_data/K.txt \
+  output_plugin_onnx
+```
+
+The last argument is the output directory (default `ffs_output`). `ffs_depth_main` sees `fast_foundationstereo.engine` inside `output_plugin_onnx/` and uses `FFSSingleEngineInference`, which deserializes the engine and registers the `FFSGWCVolume` plugin before inference.
+
+### A.4 Python alternative (build engine and run inference without the C++ apps)
+
+The same plugin ONNX can be turned into an engine and executed end-to-end from Python. Only the C++ plugin shared library (`libffs_gwc_plugin.so`) is required from the C++ build; the C++ apps (`ffs_build_single_engine`, `ffs_depth_main`) are not.
+
+```bash
+# 1. Export plugin ONNX
+python3 scripts/make_plugin_onnx.py \
+  --model_dir weights/23-36-37/model_best_bp2_serialize.pth \
+  --save_path output_plugin_onnx \
+  --height 480 \
+  --width 640
+
+# 2. Build the C++ plugin shared library
+cmake -S cpp -B cpp/build
+cmake --build cpp/build -j
+
+# 3. Build the TensorRT engine from Python
+python3 scripts/build_plugin_trt.py \
+  output_plugin_onnx/fast_foundationstereo_plugin.onnx \
+  output_plugin_onnx/fast_foundationstereo.engine
+
+# 4. Run inference from Python
+python3 scripts/run_demo_plugin_trt.py \
+  --model_dir output_plugin_onnx \
+  --left_file demo_data/left.png \
+  --right_file demo_data/right.png \
+  --intrinsic_file demo_data/K.txt \
+  --out_dir output_plugin_onnx
+```
+
+Both `build_plugin_trt.py` and `run_demo_plugin_trt.py` auto-discover `libffs_gwc_plugin.so` in `cpp/build/`. Pass `--plugin_lib /path/to/libffs_gwc_plugin.so` to override the location. `build_plugin_trt.py` accepts `--fp32` and `--workspace-mb N` with the same meaning as `ffs_build_single_engine`. `run_demo_plugin_trt.py` writes the same five output files as the C++ demo (see [Outputs](#outputs)).
+
+## Route B: Two TensorRT Engines (feature_runner + post_runner)
+
+### B.1 Build the two engines
+
+See the [Two-stage ONNX section in the top-level README](../readme.md#two-stage-onnx) for how to export `feature_runner.onnx` / `post_runner.onnx` with `scripts/make_onnx.py` and then build `feature_runner.engine` / `post_runner.engine` with `trtexec`. After running those steps your engine directory should contain:
+
+```text
+output_two_onnx/
++-- feature_runner.engine
++-- post_runner.engine
++-- onnx.yaml
+```
+
+The two engines do **not** use the FFSGWCVolume plugin: the GWC cost volume is computed externally on GPU by `cpp/src/depth_kernels.cu` between the two engine calls, so a plain `trtexec --fp16` build with no custom plugin library is enough.
+
+### B.2 Run inference
+
+```bash
+cpp/build/ffs_depth_main \
+  output_two_onnx \
+  demo_data/left.png \
+  demo_data/right.png \
+  demo_data/K.txt \
+  output_two_onnx
+```
+
+`ffs_depth_main` sees no `fast_foundationstereo.engine` in the directory and falls back to `FFSDepthInference`, which executes the feature engine, builds the GWC volume with the CUDA kernel, then executes the post engine.
+
+## Outputs
+
+Both routes write the same files into the output directory:
+
+- `disparity.bin` -- raw float32 disparity (input-pixel units), prefixed by int32 `[height, width]`.
+- `depth_meter.bin` -- raw float32 depth in meters, prefixed by int32 `[height, width]`.
+- `depth_meter.npy` -- NumPy float32 depth in meters, shape `[height, width]`.
+- `disp_vis.png` -- left/right/colorized-disparity side-by-side visualization.
+- `depth_vis.png` -- colorized depth visualization.
+
+## Profile Speed
+
+`ffs_profile_speed` benchmarks either route end-to-end on a single image pair using CUDA events for GPU-side timing and `steady_clock` for host wall time. It auto-detects the route the same way `ffs_depth_main` does (presence of `fast_foundationstereo.engine` in the engine directory selects single-engine), and `--mode` can be set explicitly:
+
+```bash
+cpp/build/ffs_profile_speed <engine_dir> <left_image> <right_image> <intrinsic_file> \
+    [--mode auto|two|single] [--warmup N] [--runs N] [--include-depth]
+```
+
+Defaults: `--mode auto`, `--warmup 10`, `--runs 30`. With `--include-depth`, the disparity-to-depth conversion is included in the timed region (otherwise only the `infer()` call is timed).
+
+### Profile the single engine
+
+```bash
+cpp/build/ffs_profile_speed \
+  output_plugin_onnx \
+  demo_data/left.png demo_data/right.png demo_data/K.txt \
+  --mode single --warmup 20 --runs 100
+```
+
+### Profile the two engines
+
+```bash
+cpp/build/ffs_profile_speed \
+  output_two_onnx \
+  demo_data/left.png demo_data/right.png demo_data/K.txt \
+  --mode two --warmup 20 --runs 100
+```
+
+The output looks like:
+
+```text
+mode=single
+image=960x540
+model=640x480
+warmup=20 runs=100
+timed_region=infer
+gpu  mean_ms=... p50_ms=... p90_ms=... min_ms=... max_ms=... std_ms=...
+host mean_ms=... p50_ms=... p90_ms=... min_ms=... max_ms=... std_ms=...
+```
+
+- `gpu` is `cudaEventElapsedTime` between start/stop events on the inference stream (pure GPU work).
+- `host` is `std::chrono::steady_clock` around the same region, including the cost of `cudaEventSynchronize`. `host >= gpu` always.
+
+For a side-by-side comparison, run the profiler twice with different engine directories and `--mode`, then compare the `gpu mean_ms` columns.

cv/fast-foundationstereo/cpp/app/build_single_engine.cpp

@@ -0,0 +1,120 @@
+#include "ffs_gwc_plugin.hpp"
+
+#include <NvInfer.h>
+#include <NvOnnxParser.h>
+
+#include <filesystem>
+#include <fstream>
+#include <iostream>
+#include <memory>
+#include <stdexcept>
+#include <string>
+
+namespace {
+
+class Logger : public nvinfer1::ILogger {
+public:
+    void log(Severity severity, const char* msg) noexcept override {
+        if (severity <= Severity::kINFO) {
+            std::cerr << "[TRT] " << msg << "\n";
+        }
+    }
+};
+
+struct TrtDestroy {
+    template <typename T>
+    void operator()(T* p) const {
+        delete p;
+    }
+};
+
+template <typename T>
+using TrtPtr = std::unique_ptr<T, TrtDestroy>;
+
+void printUsage(const char* prog) {
+    std::cerr
+        << "Usage: " << prog << " <plugin_onnx> <output_engine> [--fp32] [--workspace-mb N]\n"
+        << "\n"
+        << "Builds a single TensorRT engine from an ONNX graph containing the\n"
+        << "FFSGWCVolume custom plugin node.\n";
+}
+
+}  // namespace
+
+int main(int argc, char** argv) {
+    if (argc < 3) {
+        printUsage(argv[0]);
+        return 1;
+    }
+
+    try {
+        const std::filesystem::path onnx_path = argv[1];
+        const std::filesystem::path engine_path = argv[2];
+        bool fp16 = true;
+        size_t workspace_mb = 4096;
+
+        for (int i = 3; i < argc; ++i) {
+            const std::string arg = argv[i];
+            if (arg == "--fp32") {
+                fp16 = false;
+            } else if (arg == "--workspace-mb" && i + 1 < argc) {
+                workspace_mb = static_cast<size_t>(std::stoull(argv[++i]));
+            } else {
+                throw std::runtime_error("unknown argument: " + arg);
+            }
+        }
+
+        if (!std::filesystem::exists(onnx_path)) {
+            throw std::runtime_error("ONNX file does not exist: " + onnx_path.string());
+        }
+        if (!engine_path.parent_path().empty()) {
+            std::filesystem::create_directories(engine_path.parent_path());
+        }
+
+        if (!ffs_depth::registerFFSGWCPlugin()) {
+            throw std::runtime_error("failed to register FFSGWCVolume plugin");
+        }
+
+        Logger logger;
+        TrtPtr<nvinfer1::IBuilder> builder(nvinfer1::createInferBuilder(logger));
+        if (!builder) throw std::runtime_error("createInferBuilder failed");
+
+        const auto explicit_batch =
+            1U << static_cast<uint32_t>(nvinfer1::NetworkDefinitionCreationFlag::kEXPLICIT_BATCH);
+        TrtPtr<nvinfer1::INetworkDefinition> network(builder->createNetworkV2(explicit_batch));
+        if (!network) throw std::runtime_error("createNetworkV2 failed");
+
+        TrtPtr<nvonnxparser::IParser> parser(nvonnxparser::createParser(*network, logger));
+        if (!parser) throw std::runtime_error("createParser failed");
+
+        if (!parser->parseFromFile(onnx_path.string().c_str(),
+                                   static_cast<int32_t>(nvinfer1::ILogger::Severity::kWARNING))) {
+            for (int32_t i = 0; i < parser->getNbErrors(); ++i) {
+                std::cerr << parser->getError(i)->desc() << "\n";
+            }
+            throw std::runtime_error("failed to parse ONNX: " + onnx_path.string());
+        }
+
+        TrtPtr<nvinfer1::IBuilderConfig> config(builder->createBuilderConfig());
+        if (!config) throw std::runtime_error("createBuilderConfig failed");
+        config->setMemoryPoolLimit(nvinfer1::MemoryPoolType::kWORKSPACE,
+                                   workspace_mb * 1024ULL * 1024ULL);
+        if (fp16 && builder->platformHasFastFp16()) {
+            config->setFlag(nvinfer1::BuilderFlag::kFP16);
+        }
+
+        TrtPtr<nvinfer1::IHostMemory> serialized(builder->buildSerializedNetwork(*network, *config));
+        if (!serialized) throw std::runtime_error("buildSerializedNetwork failed");
+
+        std::ofstream out(engine_path, std::ios::binary);
+        if (!out) throw std::runtime_error("cannot write engine: " + engine_path.string());
+        out.write(static_cast<char const*>(serialized->data()), serialized->size());
+
+        std::cout << "Built engine: " << engine_path << "\n";
+        std::cout << "Precision: " << (fp16 ? "FP16 allowed" : "FP32") << "\n";
+    } catch (const std::exception& e) {
+        std::cerr << "ERROR: " << e.what() << "\n";
+        return 1;
+    }
+    return 0;
+}

cv/fast-foundationstereo/cpp/app/main.cpp

@@ -0,0 +1,346 @@
+/**
+ * Standalone Fast-FoundationStereo TensorRT inference demo.
+ *
+ * Usage:
+ *   ./ffs_depth_main <engine_dir> <left_image> <right_image> <intrinsic_file> [output_dir]
+ *
+ * Outputs:
+ *   <output_dir>/disparity.bin   int32 H, int32 W, then H*W float32 disparity
+ *   <output_dir>/depth_meter.bin int32 H, int32 W, then H*W float32 depth in meters
+ *   <output_dir>/depth_meter.npy NumPy float32 depth in meters, shape (H, W)
+ *   <output_dir>/disp_vis.png    left/right/colorized-disparity visualization
+ *   <output_dir>/depth_vis.png   colorized depth visualization
+ */
+
+#include "ffs_depth_tensorrt.hpp"
+#include "ffs_depth_single_tensorrt.hpp"
+
+#include <cuda_runtime.h>
+#include <opencv2/core.hpp>
+#include <opencv2/imgcodecs.hpp>
+#include <opencv2/imgproc.hpp>
+
+#include <algorithm>
+#include <cstdint>
+#include <cmath>
+#include <filesystem>
+#include <fstream>
+#include <iostream>
+#include <sstream>
+#include <stdexcept>
+#include <string>
+#include <vector>
+
+namespace {
+
+void checkCuda(cudaError_t status, const char* what) {
+    if (status != cudaSuccess) {
+        throw std::runtime_error(std::string(what) + ": " + cudaGetErrorString(status));
+    }
+}
+
+void printUsage(const char* prog) {
+    std::cerr
+        << "Usage: " << prog << " <engine_dir> <left_image> <right_image> <intrinsic_file> [output_dir]\n"
+        << "\n"
+        << "  engine_dir : directory containing feature_runner.engine, post_runner.engine, onnx.yaml\n"
+        << "  left_image : left stereo image readable by OpenCV\n"
+        << "  right_image: right stereo image with the same size as left_image\n"
+        << "  intrinsic_file: text file with 3x3 K on line 1 and baseline in meters on line 2\n"
+        << "  output_dir : output directory (default: ffs_output)\n";
+}
+
+struct CudaBuffer {
+    void* ptr = nullptr;
+
+    ~CudaBuffer() {
+        if (ptr) cudaFree(ptr);
+    }
+
+    CudaBuffer() = default;
+    CudaBuffer(const CudaBuffer&) = delete;
+    CudaBuffer& operator=(const CudaBuffer&) = delete;
+
+    void allocate(size_t bytes) {
+        checkCuda(cudaMalloc(&ptr, bytes), "cudaMalloc");
+    }
+
+    template <typename T>
+    T* as() {
+        return static_cast<T*>(ptr);
+    }
+};
+
+struct Intrinsics {
+    float k[9] = {};
+    float baseline = 0.0f;
+};
+
+Intrinsics loadIntrinsics(const std::string& path) {
+    std::ifstream in(path);
+    if (!in) {
+        throw std::runtime_error("cannot open intrinsic file: " + path);
+    }
+
+    std::string k_line;
+    std::string baseline_line;
+    if (!std::getline(in, k_line) || !std::getline(in, baseline_line)) {
+        throw std::runtime_error("intrinsic file must contain K on line 1 and baseline on line 2");
+    }
+
+    Intrinsics intr;
+    std::istringstream k_stream(k_line);
+    for (float& value : intr.k) {
+        if (!(k_stream >> value)) {
+            throw std::runtime_error("intrinsic file K line must contain 9 float values");
+        }
+    }
+
+    std::istringstream baseline_stream(baseline_line);
+    if (!(baseline_stream >> intr.baseline)) {
+        throw std::runtime_error("intrinsic file baseline line must contain one float value");
+    }
+    if (!(intr.k[0] > 0.0f) || !(intr.baseline > 0.0f)) {
+        throw std::runtime_error("invalid focal length or baseline in intrinsic file");
+    }
+
+    return intr;
+}
+
+void saveFloatMatrix(const std::filesystem::path& path,
+                     const std::vector<float>& values,
+                     int height,
+                     int width) {
+    std::ofstream out(path, std::ios::binary);
+    if (!out) {
+        throw std::runtime_error("cannot write matrix: " + path.string());
+    }
+    const int32_t dims[2] = {height, width};
+    out.write(reinterpret_cast<const char*>(dims), sizeof(dims));
+    out.write(reinterpret_cast<const char*>(values.data()),
+              values.size() * sizeof(float));
+}
+
+void saveNpyFloat32(const std::filesystem::path& path,
+                    const std::vector<float>& values,
+                    int height,
+                    int width) {
+    std::ofstream out(path, std::ios::binary);
+    if (!out) {
+        throw std::runtime_error("cannot write npy matrix: " + path.string());
+    }
+
+    std::ostringstream header_stream;
+    header_stream << "{'descr': '<f4', 'fortran_order': False, 'shape': ("
+                  << height << ", " << width << "), }";
+    std::string header = header_stream.str();
+
+    const size_t prefix_len = 10;  // magic(6) + version(2) + header_len(2)
+    size_t padded_len = header.size() + 1;
+    const size_t rem = (prefix_len + padded_len) % 16;
+    if (rem != 0) {
+        padded_len += 16 - rem;
+    }
+    header.append(padded_len - header.size() - 1, ' ');
+    header.push_back('\n');
+
+    if (header.size() > 65535) {
+        throw std::runtime_error("npy header too large: " + path.string());
+    }
+    const uint16_t header_len = static_cast<uint16_t>(header.size());
+
+    out.write("\x93NUMPY", 6);
+    const char version[2] = {1, 0};
+    out.write(version, 2);
+    out.write(reinterpret_cast<const char*>(&header_len), sizeof(header_len));
+    out.write(header.data(), static_cast<std::streamsize>(header.size()));
+    out.write(reinterpret_cast<const char*>(values.data()),
+              static_cast<std::streamsize>(values.size() * sizeof(float)));
+}
+
+cv::Mat colorizeDisparity(const std::vector<float>& disparity, int height, int width) {
+    std::vector<float> safe_disp(disparity.size(), 0.0f);
+    std::vector<uint8_t> valid_mask(disparity.size(), 0);
+    for (size_t i = 0; i < disparity.size(); ++i) {
+        if (std::isfinite(disparity[i])) {
+            safe_disp[i] = std::max(disparity[i], 0.0f);
+            valid_mask[i] = 255;
+        }
+    }
+
+    cv::Mat disp_mat(height, width, CV_32FC1, safe_disp.data());
+    cv::Mat valid(height, width, CV_8UC1, valid_mask.data());
+
+    if (cv::countNonZero(valid) == 0) {
+        return cv::Mat::zeros(height, width, CV_8UC3);
+    }
+
+    double min_val = 0.0;
+    double max_val = 0.0;
+    cv::minMaxLoc(disp_mat, &min_val, &max_val, nullptr, nullptr, valid);
+    if (max_val <= min_val) {
+        max_val = min_val + 1.0;
+    }
+
+    cv::Mat disp_u8;
+    disp_mat.convertTo(disp_u8, CV_8UC1, 255.0 / (max_val - min_val),
+                       -255.0 * min_val / (max_val - min_val));
+
+    cv::Mat colored;
+    cv::applyColorMap(disp_u8, colored, cv::COLORMAP_TURBO);
+    colored.setTo(cv::Scalar(0, 0, 0), ~valid);
+    return colored;
+}
+
+cv::Mat colorizeDepth(const std::vector<float>& depth, int height, int width) {
+    std::vector<float> safe_depth(depth.size(), 0.0f);
+    std::vector<uint8_t> valid_mask(depth.size(), 0);
+    for (size_t i = 0; i < depth.size(); ++i) {
+        if (std::isfinite(depth[i]) && depth[i] > 0.0f) {
+            safe_depth[i] = depth[i];
+            valid_mask[i] = 255;
+        }
+    }
+
+    cv::Mat depth_mat(height, width, CV_32FC1, safe_depth.data());
+    cv::Mat valid(height, width, CV_8UC1, valid_mask.data());
+
+    if (cv::countNonZero(valid) == 0) {
+        return cv::Mat::zeros(height, width, CV_8UC3);
+    }
+
+    double min_val = 0.0;
+    double max_val = 0.0;
+    cv::minMaxLoc(depth_mat, &min_val, &max_val, nullptr, nullptr, valid);
+    if (max_val <= min_val) {
+        max_val = min_val + 1.0;
+    }
+
+    cv::Mat depth_u8;
+    depth_mat.convertTo(depth_u8, CV_8UC1, 255.0 / (max_val - min_val),
+                        -255.0 * min_val / (max_val - min_val));
+
+    cv::Mat colored;
+    cv::applyColorMap(depth_u8, colored, cv::COLORMAP_TURBO);
+    colored.setTo(cv::Scalar(0, 0, 0), ~valid);
+    return colored;
+}
+
+}  // namespace
+
+int main(int argc, char** argv) {
+    if (argc < 5 || argc > 6) {
+        printUsage(argv[0]);
+        return 1;
+    }
+
+    try {
+        const std::string engine_dir = argv[1];
+        const std::string left_path = argv[2];
+        const std::string right_path = argv[3];
+        const std::string intrinsic_path = argv[4];
+        const std::filesystem::path output_dir = (argc == 6) ? argv[5] : "ffs_output";
+        const Intrinsics intr = loadIntrinsics(intrinsic_path);
+
+        cv::Mat left_bgr = cv::imread(left_path, cv::IMREAD_COLOR);
+        cv::Mat right_bgr = cv::imread(right_path, cv::IMREAD_COLOR);
+
+        if (left_bgr.empty()) {
+            throw std::runtime_error("cannot read left image: " + left_path);
+        }
+        if (right_bgr.empty()) {
+            throw std::runtime_error("cannot read right image: " + right_path);
+        }
+        if (left_bgr.size() != right_bgr.size()) {
+            throw std::runtime_error("left and right images must have identical dimensions");
+        }
+        if (!left_bgr.isContinuous()) left_bgr = left_bgr.clone();
+        if (!right_bgr.isContinuous()) right_bgr = right_bgr.clone();
+
+        std::filesystem::create_directories(output_dir);
+
+        const int height = left_bgr.rows;
+        const int width = left_bgr.cols;
+        const size_t image_bytes = static_cast<size_t>(height) * width * 3 * sizeof(uint8_t);
+        const size_t map_bytes = static_cast<size_t>(height) * width * sizeof(float);
+
+        CudaBuffer d_left;
+        CudaBuffer d_right;
+        CudaBuffer d_disparity;
+        CudaBuffer d_depth;
+        d_left.allocate(image_bytes);
+        d_right.allocate(image_bytes);
+        d_disparity.allocate(map_bytes);
+        d_depth.allocate(map_bytes);
+
+        checkCuda(cudaMemcpy(d_left.ptr, left_bgr.data, image_bytes, cudaMemcpyHostToDevice),
+                  "cudaMemcpy left image");
+        checkCuda(cudaMemcpy(d_right.ptr, right_bgr.data, image_bytes, cudaMemcpyHostToDevice),
+                  "cudaMemcpy right image");
+
+        std::cout << "Input images: " << width << "x" << height << "\n";
+        std::cout << "Depth: fx=" << intr.k[0] << " baseline=" << intr.baseline << " m\n";
+        std::cout << "Loading engines from: " << engine_dir << "\n";
+        const bool use_single_engine =
+            std::filesystem::exists(std::filesystem::path(engine_dir) / "fast_foundationstereo.engine");
+
+        if (use_single_engine) {
+            std::cout << "Runtime: single TensorRT engine with FFSGWCVolume plugin\n";
+            ffs_depth::FFSSingleEngineInference ffs(engine_dir);
+            ffs.infer(d_left.as<uint8_t>(), d_right.as<uint8_t>(), height, width,
+                      d_disparity.as<float>());
+            ffs.dispToDepth(d_disparity.as<float>(), height, width,
+                            intr.k[0], intr.baseline, d_depth.as<float>());
+            ffs.sync();
+        } else {
+            std::cout << "Runtime: two TensorRT engines with external CUDA GWC\n";
+            ffs_depth::FFSDepthInference ffs(engine_dir);
+            ffs.infer(d_left.as<uint8_t>(), d_right.as<uint8_t>(), height, width,
+                      d_disparity.as<float>());
+            ffs.dispToDepth(d_disparity.as<float>(), height, width,
+                            intr.k[0], intr.baseline, d_depth.as<float>());
+            ffs.sync();
+        }
+
+        std::vector<float> disparity(static_cast<size_t>(height) * width);
+        std::vector<float> depth(static_cast<size_t>(height) * width);
+        checkCuda(cudaMemcpy(disparity.data(), d_disparity.ptr, map_bytes, cudaMemcpyDeviceToHost),
+                  "cudaMemcpy disparity");
+        checkCuda(cudaMemcpy(depth.data(), d_depth.ptr, map_bytes, cudaMemcpyDeviceToHost),
+                  "cudaMemcpy depth");
+
+        const auto disparity_path = output_dir / "disparity.bin";
+        saveFloatMatrix(disparity_path, disparity, height, width);
+
+        const auto depth_path = output_dir / "depth_meter.bin";
+        saveFloatMatrix(depth_path, depth, height, width);
+
+        const auto depth_npy_path = output_dir / "depth_meter.npy";
+        saveNpyFloat32(depth_npy_path, depth, height, width);
+
+        const cv::Mat disp_color = colorizeDisparity(disparity, height, width);
+        cv::Mat disp_vis;
+        cv::hconcat(std::vector<cv::Mat>{left_bgr, right_bgr, disp_color}, disp_vis);
+        const auto disp_vis_path = output_dir / "disp_vis.png";
+        if (!cv::imwrite(disp_vis_path.string(), disp_vis)) {
+            throw std::runtime_error("failed to write disparity visualization: " + disp_vis_path.string());
+        }
+
+        const cv::Mat depth_vis = colorizeDepth(depth, height, width);
+        const auto depth_vis_path = output_dir / "depth_vis.png";
+        if (!cv::imwrite(depth_vis_path.string(), depth_vis)) {
+            throw std::runtime_error("failed to write depth visualization: " + depth_vis_path.string());
+        }
+
+        std::cout << "Saved: " << disparity_path << "\n";
+        std::cout << "Saved: " << depth_path << "\n";
+        std::cout << "Saved: " << depth_npy_path << "\n";
+        std::cout << "Saved: " << disp_vis_path << "\n";
+        std::cout << "Saved: " << depth_vis_path << "\n";
+    } catch (const std::exception& e) {
+        std::cerr << "ERROR: " << e.what() << "\n";
+        return 1;
+    }
+
+    return 0;
+}

cv/fast-foundationstereo/cpp/app/profile_speed.cpp

@@ -0,0 +1,312 @@
+#include "ffs_depth_tensorrt.hpp"
+#include "ffs_depth_single_tensorrt.hpp"
+
+#include <cuda_runtime.h>
+#include <opencv2/core.hpp>
+#include <opencv2/imgcodecs.hpp>
+
+#include <algorithm>
+#include <chrono>
+#include <cmath>
+#include <cstdint>
+#include <filesystem>
+#include <fstream>
+#include <iostream>
+#include <numeric>
+#include <sstream>
+#include <stdexcept>
+#include <string>
+#include <vector>
+
+namespace {
+
+void checkCuda(cudaError_t status, const char* what) {
+    if (status != cudaSuccess) {
+        throw std::runtime_error(std::string(what) + ": " + cudaGetErrorString(status));
+    }
+}
+
+struct CudaBuffer {
+    void* ptr = nullptr;
+    ~CudaBuffer() { if (ptr) cudaFree(ptr); }
+    void allocate(size_t bytes) { checkCuda(cudaMalloc(&ptr, bytes), "cudaMalloc"); }
+    template <typename T> T* as() { return static_cast<T*>(ptr); }
+};
+
+struct Intrinsics {
+    float k[9] = {};
+    float baseline = 0.0f;
+};
+
+struct Args {
+    std::string engine_dir;
+    std::string left_path;
+    std::string right_path;
+    std::string intrinsic_path;
+    std::string mode = "auto";
+    int warmup = 10;
+    int runs = 30;
+    bool include_depth = false;
+};
+
+struct Stats {
+    double mean = 0.0;
+    double min = 0.0;
+    double p50 = 0.0;
+    double p90 = 0.0;
+    double max = 0.0;
+    double stddev = 0.0;
+};
+
+void printUsage(const char* prog) {
+    std::cerr
+        << "Usage: " << prog
+        << " <engine_dir> <left_image> <right_image> <intrinsic_file>"
+        << " [--mode auto|two|single] [--warmup N] [--runs N] [--include-depth]\n";
+}
+
+Args parseArgs(int argc, char** argv) {
+    if (argc < 5) {
+        printUsage(argv[0]);
+        throw std::runtime_error("missing required arguments");
+    }
+
+    Args args;
+    args.engine_dir = argv[1];
+    args.left_path = argv[2];
+    args.right_path = argv[3];
+    args.intrinsic_path = argv[4];
+
+    for (int i = 5; i < argc; ++i) {
+        const std::string key = argv[i];
+        auto requireValue = [&](const char* name) -> std::string {
+            if (i + 1 >= argc) {
+                throw std::runtime_error(std::string("missing value for ") + name);
+            }
+            return argv[++i];
+        };
+
+        if (key == "--mode") {
+            args.mode = requireValue("--mode");
+        } else if (key == "--warmup") {
+            args.warmup = std::stoi(requireValue("--warmup"));
+        } else if (key == "--runs") {
+            args.runs = std::stoi(requireValue("--runs"));
+        } else if (key == "--include-depth") {
+            args.include_depth = true;
+        } else {
+            throw std::runtime_error("unknown argument: " + key);
+        }
+    }
+
+    if (args.mode != "auto" && args.mode != "two" && args.mode != "single") {
+        throw std::runtime_error("--mode must be auto, two, or single");
+    }
+    if (args.warmup < 0 || args.runs <= 0) {
+        throw std::runtime_error("--warmup must be >= 0 and --runs must be > 0");
+    }
+    return args;
+}
+
+Intrinsics loadIntrinsics(const std::string& path) {
+    std::ifstream in(path);
+    if (!in) throw std::runtime_error("cannot open intrinsic file: " + path);
+    std::string k_line;
+    std::string baseline_line;
+    if (!std::getline(in, k_line) || !std::getline(in, baseline_line)) {
+        throw std::runtime_error("intrinsic file must contain K on line 1 and baseline on line 2");
+    }
+
+    Intrinsics intr;
+    std::istringstream ks(k_line);
+    for (float& v : intr.k) {
+        if (!(ks >> v)) throw std::runtime_error("K line must contain 9 floats");
+    }
+    std::istringstream bs(baseline_line);
+    if (!(bs >> intr.baseline)) throw std::runtime_error("baseline line must contain one float");
+    return intr;
+}
+
+Stats summarize(std::vector<float> values) {
+    if (values.empty()) throw std::runtime_error("cannot summarize empty timing vector");
+    std::sort(values.begin(), values.end());
+    const double sum = std::accumulate(values.begin(), values.end(), 0.0);
+    const double mean = sum / static_cast<double>(values.size());
+    double var = 0.0;
+    for (float v : values) {
+        const double d = static_cast<double>(v) - mean;
+        var += d * d;
+    }
+    var /= static_cast<double>(values.size());
+
+    auto percentile = [&](double p) {
+        const size_t idx = static_cast<size_t>(
+            std::llround((values.size() - 1) * p / 100.0));
+        return static_cast<double>(values[std::min(idx, values.size() - 1)]);
+    };
+
+    Stats stats;
+    stats.mean = mean;
+    stats.min = values.front();
+    stats.p50 = percentile(50.0);
+    stats.p90 = percentile(90.0);
+    stats.max = values.back();
+    stats.stddev = std::sqrt(var);
+    return stats;
+}
+
+void printStats(const char* label, const Stats& s) {
+    std::cout
+        << label
+        << " mean_ms=" << s.mean
+        << " p50_ms=" << s.p50
+        << " p90_ms=" << s.p90
+        << " min_ms=" << s.min
+        << " max_ms=" << s.max
+        << " std_ms=" << s.stddev
+        << "\n";
+}
+
+template <typename Runner>
+std::vector<float> profileRunner(
+    Runner& runner,
+    uint8_t* d_left,
+    uint8_t* d_right,
+    int height,
+    int width,
+    float* d_disp,
+    float* d_depth,
+    float fx,
+    float baseline,
+    int warmup,
+    int runs,
+    bool include_depth,
+    std::vector<float>& host_ms)
+{
+    cudaEvent_t start = nullptr;
+    cudaEvent_t stop = nullptr;
+    checkCuda(cudaEventCreate(&start), "cudaEventCreate start");
+    checkCuda(cudaEventCreate(&stop), "cudaEventCreate stop");
+
+    std::vector<float> gpu_ms;
+    gpu_ms.reserve(static_cast<size_t>(runs));
+    host_ms.clear();
+    host_ms.reserve(static_cast<size_t>(runs));
+
+    for (int i = 0; i < warmup + runs; ++i) {
+        const auto host_start = std::chrono::steady_clock::now();
+        checkCuda(cudaEventRecord(start, runner.stream()), "cudaEventRecord start");
+        runner.infer(d_left, d_right, height, width, d_disp);
+        if (include_depth) {
+            runner.dispToDepth(d_disp, height, width, fx, baseline, d_depth);
+        }
+        checkCuda(cudaEventRecord(stop, runner.stream()), "cudaEventRecord stop");
+        checkCuda(cudaEventSynchronize(stop), "cudaEventSynchronize stop");
+        const auto host_stop = std::chrono::steady_clock::now();
+
+        if (i >= warmup) {
+            float elapsed = 0.0f;
+            checkCuda(cudaEventElapsedTime(&elapsed, start, stop), "cudaEventElapsedTime");
+            gpu_ms.push_back(elapsed);
+            host_ms.push_back(static_cast<float>(
+                std::chrono::duration<double, std::milli>(host_stop - host_start).count()));
+        }
+    }
+
+    cudaEventDestroy(stop);
+    cudaEventDestroy(start);
+    return gpu_ms;
+}
+
+template <typename Runner>
+void runProfile(const Args& args,
+                Runner& runner,
+                uint8_t* d_left,
+                uint8_t* d_right,
+                int height,
+                int width,
+                float* d_disp,
+                float* d_depth,
+                const Intrinsics& intr,
+                const char* mode_name) {
+    std::vector<float> host_ms;
+    const std::vector<float> gpu_ms = profileRunner(
+        runner,
+        d_left,
+        d_right,
+        height,
+        width,
+        d_disp,
+        d_depth,
+        intr.k[0],
+        intr.baseline,
+        args.warmup,
+        args.runs,
+        args.include_depth,
+        host_ms);
+
+    std::cout << "mode=" << mode_name << "\n";
+    std::cout << "image=" << width << "x" << height << "\n";
+    std::cout << "model=" << runner.modelWidth() << "x" << runner.modelHeight() << "\n";
+    std::cout << "warmup=" << args.warmup << " runs=" << args.runs << "\n";
+    std::cout << "timed_region=" << (args.include_depth ? "infer+dispToDepth" : "infer") << "\n";
+    printStats("gpu", summarize(gpu_ms));
+    printStats("host", summarize(host_ms));
+}
+
+}  // namespace
+
+int main(int argc, char** argv) {
+    try {
+        const Args args = parseArgs(argc, argv);
+        const Intrinsics intr = loadIntrinsics(args.intrinsic_path);
+
+        cv::Mat left = cv::imread(args.left_path, cv::IMREAD_COLOR);
+        cv::Mat right = cv::imread(args.right_path, cv::IMREAD_COLOR);
+        if (left.empty()) throw std::runtime_error("cannot read left image: " + args.left_path);
+        if (right.empty()) throw std::runtime_error("cannot read right image: " + args.right_path);
+        if (left.size() != right.size()) throw std::runtime_error("left/right size mismatch");
+        if (!left.isContinuous()) left = left.clone();
+        if (!right.isContinuous()) right = right.clone();
+
+        const int height = left.rows;
+        const int width = left.cols;
+        const size_t image_bytes = static_cast<size_t>(height) * width * 3;
+        const size_t map_bytes = static_cast<size_t>(height) * width * sizeof(float);
+
+        CudaBuffer d_left;
+        CudaBuffer d_right;
+        CudaBuffer d_disp;
+        CudaBuffer d_depth;
+        d_left.allocate(image_bytes);
+        d_right.allocate(image_bytes);
+        d_disp.allocate(map_bytes);
+        d_depth.allocate(map_bytes);
+        checkCuda(cudaMemcpy(d_left.ptr, left.data, image_bytes, cudaMemcpyHostToDevice), "copy left");
+        checkCuda(cudaMemcpy(d_right.ptr, right.data, image_bytes, cudaMemcpyHostToDevice), "copy right");
+
+        std::string mode = args.mode;
+        if (mode == "auto") {
+            mode = std::filesystem::exists(
+                std::filesystem::path(args.engine_dir) / "fast_foundationstereo.engine")
+                ? "single"
+                : "two";
+        }
+
+        if (mode == "single") {
+            ffs_depth::FFSSingleEngineInference runner(args.engine_dir);
+            runProfile(args, runner, d_left.as<uint8_t>(), d_right.as<uint8_t>(),
+                       height, width, d_disp.as<float>(), d_depth.as<float>(),
+                       intr, "single");
+        } else {
+            ffs_depth::FFSDepthInference runner(args.engine_dir);
+            runProfile(args, runner, d_left.as<uint8_t>(), d_right.as<uint8_t>(),
+                       height, width, d_disp.as<float>(), d_depth.as<float>(),
+                       intr, "two");
+        }
+    } catch (const std::exception& e) {
+        std::cerr << "ERROR: " << e.what() << "\n";
+        return 1;
+    }
+    return 0;
+}

cv/fast-foundationstereo/cpp/include/ffs_depth_single_tensorrt.hpp

@@ -0,0 +1,80 @@
+#pragma once
+
+#include <NvInfer.h>
+#include <cuda_runtime.h>
+#include <cstdint>
+#include <memory>
+#include <string>
+
+#include "ffs_depth_tensorrt.hpp"
+
+namespace ffs_depth {
+
+/**
+ * Single TensorRT engine inference path.
+ *
+ * Expected engine directory layout:
+ *   - fast_foundationstereo.engine
+ *   - onnx.yaml
+ *
+ * The engine is built from the plugin ONNX export path where the GWC cost
+ * volume is represented by an FFSGWCVolume TensorRT plugin node.
+ */
+class FFSSingleEngineInference {
+public:
+    explicit FFSSingleEngineInference(const std::string& engine_dir);
+    ~FFSSingleEngineInference();
+
+    FFSSingleEngineInference(const FFSSingleEngineInference&) = delete;
+    FFSSingleEngineInference& operator=(const FFSSingleEngineInference&) = delete;
+
+    void infer(const uint8_t* d_left_rgb, const uint8_t* d_right_rgb,
+               int input_h, int input_w,
+               float* d_disp_out);
+
+    void dispToDepth(const float* d_disp,
+                     int height, int width,
+                     float fx, float baseline_m,
+                     float* d_depth_out);
+
+    void inferDepth(const uint8_t* d_left_rgb, const uint8_t* d_right_rgb,
+                    int input_h, int input_w,
+                    float fx, float baseline_m,
+                    float* d_depth_out);
+
+    void sync();
+
+    int modelHeight() const { return config_.image_height; }
+    int modelWidth() const { return config_.image_width; }
+    int maxDisp() const { return config_.max_disp; }
+    int cvGroup() const { return config_.cv_group; }
+    const FFSDepthInference::Config& config() const { return config_; }
+    cudaStream_t stream() const { return stream_; }
+
+private:
+    void loadConfig(const std::string& config_path);
+    void loadEngine(const std::string& path);
+    void allocateBuffers();
+    void freeDeviceBuffers();
+    void preprocessRGBGPU(const uint8_t* d_rgb, int src_h, int src_w, float* d_output);
+
+    FFSDepthInference::Config config_;
+
+    std::unique_ptr<nvinfer1::IRuntime> runtime_;
+    std::unique_ptr<nvinfer1::ICudaEngine> engine_;
+    std::unique_ptr<nvinfer1::IExecutionContext> context_;
+
+    cudaStream_t stream_ = nullptr;
+
+    float* d_left_ = nullptr;
+    float* d_right_ = nullptr;
+    float* d_disp_ = nullptr;
+    float* d_disp_cropped_ = nullptr;
+    float* d_disp_for_depth_ = nullptr;
+
+    int64_t depth_alloc_pixels_ = 0;
+    int scaled_w_ = 0;
+    int scaled_h_ = 0;
+};
+
+}  // namespace ffs_depth

cv/fast-foundationstereo/cpp/include/ffs_depth_tensorrt.hpp

@@ -0,0 +1,190 @@
+#pragma once
+
+#include <NvInfer.h>
+#include <cuda_runtime.h>
+#include <cstdint>
+#include <memory>
+#include <string>
+#include <vector>
+
+namespace ffs_depth {
+
+/**
+ * FoundationStereo TensorRT depth inference.
+ *
+ * Uses the two-engine architecture (feature_runner + post_runner).
+ * The GWC (Group-wise Correlation) volume is computed on GPU between the two engines.
+ *
+ * Expected engine directory layout:
+ *   - feature_runner.engine
+ *   - post_runner.engine
+ *   - onnx.yaml  (contains image_size, max_disp, cv_group, valid_iters)
+ */
+class FFSDepthInference {
+public:
+    struct Config {
+        int image_height = 480;   // Model input height (must be divisible by 32)
+        int image_width  = 864;   // Model input width (must be divisible by 32)
+        int max_disp     = 192;   // Maximum disparity search range in pixels
+        int cv_group     = 8;     // Number of groups for group-wise correlation (GWC) volume
+        int valid_iters  = 8;     // Number of GRU refinement iterations during inference
+        bool normalize   = true;  // Whether to L2-normalize features in GWC correlation
+    };
+
+    /**
+     * Load engines and allocate all GPU buffers.
+     * @param engine_dir  Directory containing the two .engine files and onnx.yaml.
+     */
+    explicit FFSDepthInference(const std::string& engine_dir);
+    ~FFSDepthInference();
+
+    FFSDepthInference(const FFSDepthInference&) = delete;
+    FFSDepthInference& operator=(const FFSDepthInference&) = delete;
+
+    /**
+     * Run stereo disparity inference entirely on the GPU.
+     *
+     * This method is asynchronous: all work (preprocessing, TensorRT
+     * inference, postprocessing) is enqueued on an internal CUDA stream
+     * and the call returns immediately. Call sync() to block until the
+     * output buffer is safe to read.
+     *
+     * Input images are expected in HWC uint8 BGR (OpenCV) format on device memory.
+     * If (input_h, input_w) differs from the model size, the images are
+     * uniformly (aspect-preserving) bilinearly resized down to fit the model
+     * input and the remaining right/bottom strip is filled with replicate
+     * padding. After inference the disparity is cropped back to the scaled
+     * region and nearest-neighbour upsampled to (input_h, input_w), with
+     * a horizontal scale correction so the values are in input-pixel units.
+     *
+     * @param d_left_rgb    Left  image on GPU  (input_h x input_w x 3, uint8 BGR)
+     * @param d_right_rgb   Right image on GPU  (input_h x input_w x 3, uint8 BGR)
+     * @param input_h       Height of input images (must be > 0)
+     * @param input_w       Width  of input images (must be > 0)
+     * @param d_disp_out    Output disparity on GPU (input_h x input_w, float32,
+     *                      in input-pixel units, clamped to >= 0).
+     */
+    void infer(const uint8_t* d_left_rgb, const uint8_t* d_right_rgb,
+               int input_h, int input_w,
+               float* d_disp_out);
+
+    /**
+     * Convert an input-resolution disparity map to float32 depth in meters.
+     *
+     * This method is asynchronous: work is enqueued on the internal CUDA
+     * stream and the call returns immediately. Call sync() before reading.
+     *
+     * Conversion semantics match scripts/run_demo.py (see inferDepth docstring
+     * for the +inf / off-image-correspondence handling).
+     *
+     * @param d_disp       Input disparity on GPU (height x width, float32).
+     *                     Must be expressed in INPUT-image pixel units (which
+     *                     is what infer() produces).
+     * @param height       Disparity / depth height
+     * @param width        Disparity / depth width
+     * @param fx           Focal length in pixels at INPUT resolution
+     * @param baseline_m   Stereo baseline in meters
+     * @param d_depth_out  Output depth on GPU (height x width, float32, meters)
+     */
+    void dispToDepth(const float* d_disp,
+                     int height, int width,
+                     float fx, float baseline_m,
+                     float* d_depth_out);
+
+    /**
+     * Run stereo inference and convert disparity to float32 depth (meters) in one call.
+     *
+     * This method is asynchronous: all work is enqueued on an internal
+     * CUDA stream and the call returns immediately. Call sync() to block
+     * until d_depth_out is safe to read.
+     *
+     * The conversion matches scripts/run_demo.py:
+     *   depth_m = fx * baseline_m / disparity
+     * where:
+     *   - disparity is in INPUT-image pixel units (infer() already rescales it),
+     *     so `fx` must come from the INPUT-resolution intrinsics (i.e. K[0,0]
+     *     of the unscaled camera matrix, not the model-resolution intrinsics).
+     *   - pixels with x - disparity < 0 (right-image correspondence off-image)
+     *     are marked invalid: disparity is replaced with +inf, yielding depth 0.
+     *   - disparity == 0 yields depth = +inf (consumers should mask non-finite
+     *     values before using the depth map for downstream geometry).
+     *
+     * @param d_left_rgb   Left  image on GPU (input_h x input_w x 3, uint8 BGR/RGB)
+     * @param d_right_rgb  Right image on GPU (input_h x input_w x 3, uint8 BGR/RGB)
+     * @param input_h      Height of input images
+     * @param input_w      Width of input images
+     * @param fx           Focal length in pixels at INPUT resolution
+     * @param baseline_m   Stereo baseline in meters
+     * @param d_depth_out  Output depth on GPU (input_h x input_w, float32, meters)
+     */
+    void inferDepth(const uint8_t* d_left_rgb, const uint8_t* d_right_rgb,
+                    int input_h, int input_w,
+                    float fx, float baseline_m,
+                    float* d_depth_out);
+
+    /** Block until all async work on the internal CUDA stream has completed. */
+    void sync();
+
+    int modelHeight()  const { return config_.image_height; }
+    int modelWidth()   const { return config_.image_width;  }
+    int maxDisp()      const { return config_.max_disp;     }
+    int cvGroup()      const { return config_.cv_group;     }
+    const Config& config() const { return config_; }
+    cudaStream_t stream() const { return stream_; }
+
+private:
+    void loadConfig(const std::string& config_path);
+    void loadEngine(const std::string& path,
+                    std::unique_ptr<nvinfer1::ICudaEngine>& engine,
+                    std::unique_ptr<nvinfer1::IExecutionContext>& context);
+    void allocateBuffers();
+    void allocateFeatureBuffers();
+    void allocatePostBuffers();
+    void freeDeviceBuffers();  // safe to call repeatedly; nulls every pointer
+    void preprocessRGBGPU(const uint8_t* d_rgb, int src_h, int src_w, float* d_output);
+
+    void runFeatureRunner();
+    void buildGWCVolume();
+    void runPostRunner();
+
+    Config config_;
+
+    std::unique_ptr<nvinfer1::IRuntime>          runtime_;
+    std::unique_ptr<nvinfer1::ICudaEngine>       feature_engine_;
+    std::unique_ptr<nvinfer1::IExecutionContext>  feature_context_;
+    std::unique_ptr<nvinfer1::ICudaEngine>       post_engine_;
+    std::unique_ptr<nvinfer1::IExecutionContext>  post_context_;
+
+    cudaStream_t stream_ = nullptr;
+
+    float* d_left_  = nullptr;
+    float* d_right_ = nullptr;
+
+    float* d_feat_left_04_  = nullptr;
+    float* d_feat_left_08_  = nullptr;
+    float* d_feat_left_16_  = nullptr;
+    float* d_feat_left_32_  = nullptr;
+    float* d_feat_right_04_ = nullptr;
+    float* d_stem_2x_       = nullptr;
+
+    std::vector<int> feat_04_dims_;
+    std::vector<int> feat_08_dims_;
+    std::vector<int> feat_16_dims_;
+    std::vector<int> feat_32_dims_;
+    std::vector<int> stem_2x_dims_;
+
+    float* d_gwc_volume_   = nullptr;
+    int    gwc_disp_levels_ = 0;
+
+    float* d_disp_         = nullptr;
+    float* d_disp_cropped_ = nullptr;
+    float* d_disp_for_depth_ = nullptr;
+    int64_t depth_alloc_pixels_ = 0;
+
+    int    scaled_w_ = 0;  // Uniform-scaled width (before padding)
+    int    scaled_h_ = 0;  // Uniform-scaled height (before padding)
+
+    bool gwc_fp16_  = false;  // GWC volume tensor is FP16 (vs. FP32)
+};
+
+}  // namespace ffs_depth

cv/fast-foundationstereo/cpp/include/ffs_gwc_plugin.hpp

@@ -0,0 +1,12 @@
+#pragma once
+
+namespace ffs_depth {
+
+// Registers the FFSGWCVolume TensorRT plugin creator in the global registry.
+// Safe to call more than once.
+bool registerFFSGWCPlugin();
+
+}  // namespace ffs_depth
+
+// C ABI wrapper for loading/registering the plugin from Python via ctypes.
+extern "C" bool ffs_register_gwc_plugin();

cv/fast-foundationstereo/cpp/src/depth_kernels.cu

@@ -0,0 +1,345 @@
+/**
+ * CUDA kernels for the Fast-FoundationStereo TensorRT depth inference pipeline.
+ *
+ * Active kernels:
+ *   1.  GWC volume (mixed precision) -- FP32 features -> FP16 correlation volume
+ *   2.  HWC uint8 -> CHW float      -- exact-match path (no resize, with zero-padding)
+ *   3. Uniform resize + pad        -- aspect-ratio-preserving bilinear resize with
+ *                                      border-replicate padding to model dimensions
+ *   4.  Disparity crop              -- removes padding from model-resolution disparity
+ *   5.  Disparity upsample          -- nearest-neighbor upsample with scale correction
+ *   6.  Disparity clamp             -- clamps to minimum value
+ *   7.  Disparity to depth          -- depth_m = fx * baseline_m / disparity (float32, meters)
+ */
+
+#include <cuda_runtime.h>
+#include <cuda_fp16.h>
+#include <device_launch_parameters.h>
+#include <math_constants.h>
+#include <cstdint>
+
+namespace ffs_depth {
+namespace cuda {
+
+// =========================================================================
+// 1. GWC Volume
+// =========================================================================
+template <typename T>
+__device__ __forceinline__ float toFloat(T v) {
+    return static_cast<float>(v);
+}
+
+template <>
+__device__ __forceinline__ float toFloat<__half>(__half v) {
+    return __half2float(v);
+}
+
+template <typename T>
+__device__ __forceinline__ T fromFloat(float v) {
+    return static_cast<T>(v);
+}
+
+template <>
+__device__ __forceinline__ __half fromFloat<__half>(float v) {
+    return __float2half(v);
+}
+
+template <typename InputT, typename OutputT>
+__global__ void buildGWCVolumeKernel(
+    const InputT* __restrict__ feat_left,
+    const InputT* __restrict__ feat_right,
+    OutputT* __restrict__ gwc_volume,
+    int B, int C, int H, int W,
+    int max_disp, int num_groups, bool normalize)
+{
+    const int w = blockIdx.x * blockDim.x + threadIdx.x;
+    const int h = blockIdx.y * blockDim.y + threadIdx.y;
+    const int dgb = blockIdx.z;
+
+    if (w >= W || h >= H) return;
+
+    const int d = dgb % max_disp;
+    const int g = (dgb / max_disp) % num_groups;
+    const int b = dgb / (max_disp * num_groups);
+    if (b >= B) return;
+
+    const int K = C / num_groups;
+    const int w_right = w - d;
+
+    const int out_idx = ((b * num_groups + g) * max_disp + d) * H * W + h * W + w;
+
+    if (w_right < 0) { gwc_volume[out_idx] = fromFloat<OutputT>(0.0f); return; }
+
+    float dot = 0.f, nl = 0.f, nr = 0.f;
+    const int left_base  = (b * C + g * K) * H * W + h * W + w;
+    const int right_base = (b * C + g * K) * H * W + h * W + w_right;
+    const int stride = H * W;
+
+    for (int k = 0; k < K; ++k) {
+        float l = toFloat<InputT>(feat_left [left_base  + k * stride]);
+        float r = toFloat<InputT>(feat_right[right_base + k * stride]);
+        dot += l * r;  nl += l * l;  nr += r * r;
+    }
+    if (normalize)
+        gwc_volume[out_idx] = fromFloat<OutputT>(dot / (sqrtf(nl) * sqrtf(nr) + 1e-5f));
+    else
+        gwc_volume[out_idx] = fromFloat<OutputT>(dot);
+}
+
+// =========================================================================
+// 2. Preprocess RGB HWC uint8 -> CHW float  (exact-match path, no resize)
+//    Used when input resolution matches model resolution exactly.
+// =========================================================================
+__global__ void preprocessRGBToCHWKernel(
+    const uint8_t* __restrict__ d_rgb_hwc,
+    float* __restrict__ d_chw,
+    int src_h, int src_w,
+    int dst_h, int dst_w)
+{
+    const int x = blockIdx.x * blockDim.x + threadIdx.x;
+    const int y = blockIdx.y * blockDim.y + threadIdx.y;
+    if (x >= dst_w || y >= dst_h) return;
+
+    const int hw  = dst_h * dst_w;
+    const int idx = y * dst_w + x;
+
+    if (x < src_w && y < src_h) {
+        const int si = (y * src_w + x) * 3;
+        d_chw[idx]          = static_cast<float>(d_rgb_hwc[si + 2]);  // R
+        d_chw[hw + idx]     = static_cast<float>(d_rgb_hwc[si + 1]);  // G
+        d_chw[2 * hw + idx] = static_cast<float>(d_rgb_hwc[si]);      // B
+    } else {
+        d_chw[idx] = 0.f;  d_chw[hw + idx] = 0.f;  d_chw[2 * hw + idx] = 0.f;
+    }
+}
+
+// =========================================================================
+// 3. Uniform resize + border-replicate padding  (aspect-ratio preserving)
+//    Pixels in [0, scaled_w) x [0, scaled_h) are bilinear-sampled from src.
+//    Pixels in the padding region replicate the nearest edge pixel.
+// =========================================================================
+__global__ void resizeUniformAndPadKernel(
+    const uint8_t* __restrict__ d_rgb_hwc,
+    float* __restrict__ d_chw,
+    int src_h, int src_w,
+    int scaled_h, int scaled_w,
+    int dst_h, int dst_w)
+{
+    const int x = blockIdx.x * blockDim.x + threadIdx.x;
+    const int y = blockIdx.y * blockDim.y + threadIdx.y;
+    if (x >= dst_w || y >= dst_h) return;
+
+    const int hw  = dst_h * dst_w;
+    const int idx = y * dst_w + x;
+
+    const int cx = min(x, scaled_w - 1);
+    const int cy = min(y, scaled_h - 1);
+
+    const float sx = static_cast<float>(src_w) / scaled_w;
+    const float sy = static_cast<float>(src_h) / scaled_h;
+
+    const float fx = (cx + 0.5f) * sx - 0.5f;
+    const float fy = (cy + 0.5f) * sy - 0.5f;
+
+    const int x0 = max(0, min(__float2int_rd(fx), src_w - 1));
+    const int y0 = max(0, min(__float2int_rd(fy), src_h - 1));
+    const int x1 = min(x0 + 1, src_w - 1);
+    const int y1 = min(y0 + 1, src_h - 1);
+
+    const float wx = fx - floorf(fx);
+    const float wy = fy - floorf(fy);
+
+    for (int c = 0; c < 3; ++c) {
+        const int sc = 2 - c;  // RGB channel reorder
+        float v00 = static_cast<float>(d_rgb_hwc[(y0 * src_w + x0) * 3 + sc]);
+        float v01 = static_cast<float>(d_rgb_hwc[(y0 * src_w + x1) * 3 + sc]);
+        float v10 = static_cast<float>(d_rgb_hwc[(y1 * src_w + x0) * 3 + sc]);
+        float v11 = static_cast<float>(d_rgb_hwc[(y1 * src_w + x1) * 3 + sc]);
+        float val = (1.f - wy) * ((1.f - wx) * v00 + wx * v01)
+                  +        wy  * ((1.f - wx) * v10 + wx * v11);
+        d_chw[c * hw + idx] = val;
+    }
+}
+
+// =========================================================================
+// 4. Crop disparity  (remove border-replicate padding)
+//    Extracts the valid (scaled_w x scaled_h) region from the
+//    model-resolution (model_w x model_h) disparity output.
+// =========================================================================
+__global__ void cropDisparityKernel(
+    const float* __restrict__ src, float* __restrict__ dst,
+    int src_w, int dst_h, int dst_w)
+{
+    const int x = blockIdx.x * blockDim.x + threadIdx.x;
+    const int y = blockIdx.y * blockDim.y + threadIdx.y;
+    if (x >= dst_w || y >= dst_h) return;
+    dst[y * dst_w + x] = src[y * src_w + x];
+}
+
+// =========================================================================
+// 5. Upsample float disparity (nearest neighbor) with scale correction
+//    Upsamples from cropped model resolution to input resolution.
+//    disp_scale converts disparity from scaled-pixel units to input-pixel units.
+// =========================================================================
+__global__ void upsampleDisparityKernel(
+    const float* __restrict__ src, float* __restrict__ dst,
+    int src_w, int src_h,
+    int dst_w, int dst_h,
+    float disp_scale)
+{
+    const int x = blockIdx.x * blockDim.x + threadIdx.x;
+    const int y = blockIdx.y * blockDim.y + threadIdx.y;
+    if (x >= dst_w || y >= dst_h) return;
+
+    const float sx = static_cast<float>(src_w) / dst_w;
+    const float sy = static_cast<float>(src_h) / dst_h;
+
+    const float fx = (x + 0.5f) * sx - 0.5f;
+    const float fy = (y + 0.5f) * sy - 0.5f;
+
+    const int x0 = max(0, min(__float2int_rn(fx), src_w - 1));
+    const int y0 = max(0, min(__float2int_rn(fy), src_h - 1));
+    float val = src[y0 * src_w + x0];
+
+    dst[y * dst_w + x] = val * disp_scale;
+}
+
+// =========================================================================
+// 6. Clamp disparity to minimum value (removes negative/zero artifacts)
+// =========================================================================
+__global__ void clampDisparityKernel(
+    float* __restrict__ disp, int count, float min_val)
+{
+    const int idx = blockIdx.x * blockDim.x + threadIdx.x;
+    if (idx >= count) return;
+    disp[idx] = fmaxf(disp[idx], min_val);
+}
+
+// =========================================================================
+// 7. Disparity to depth: depth_m = fx * baseline_m / disparity (float32 meters)
+//
+//    Matches scripts/run_demo.py behaviour:
+//      - disparity is assumed already clamped to >= 0 (see clampDisparityKernel)
+//      - pixels where the right-image correspondence would fall off-image
+//        (x - disp < 0) are marked invalid by setting disparity to +inf,
+//        which yields depth = 0 after division.
+//      - disp == 0 yields depth = +inf (matches numpy's float divide-by-zero).
+//      - no clip on the upper end; output is float32 meters.
+// =========================================================================
+__global__ void dispToDepthKernel(
+    const float* __restrict__ disp,
+    float* __restrict__ depth_m,
+    int height, int width, float fx, float baseline_m)
+{
+    const int x = blockIdx.x * blockDim.x + threadIdx.x;
+    const int y = blockIdx.y * blockDim.y + threadIdx.y;
+    if (x >= width || y >= height) return;
+    const int i = y * width + x;
+
+    float d = disp[i];
+    if (static_cast<float>(x) - d < 0.0f) {
+        d = CUDART_INF_F;  // remove_invisible: drop pixel via inf -> depth 0
+    }
+    depth_m[i] = fx * baseline_m / d;
+}
+
+// =========================================================================
+// Host wrappers  (extern "C" for linkage)
+// =========================================================================
+extern "C" {
+
+void ffsCudaBuildGWCVolumeFloat(
+    const float* d_fl, const float* d_fr, float* d_gwc,
+    int B, int C, int H, int W, int max_disp, int ngroups, bool normalize, cudaStream_t s)
+{
+    dim3 blk(16, 16);
+    dim3 grd((W + 15) / 16, (H + 15) / 16, B * ngroups * max_disp);
+    buildGWCVolumeKernel<float, float><<<grd, blk, 0, s>>>(d_fl, d_fr, d_gwc, B, C, H, W, max_disp, ngroups, normalize);
+}
+
+void ffsCudaBuildGWCVolumeHalf(
+    const __half* d_fl, const __half* d_fr, __half* d_gwc,
+    int B, int C, int H, int W, int max_disp, int ngroups, bool normalize, cudaStream_t s)
+{
+    dim3 blk(16, 16);
+    dim3 grd((W + 15) / 16, (H + 15) / 16, B * ngroups * max_disp);
+    buildGWCVolumeKernel<__half, __half><<<grd, blk, 0, s>>>(d_fl, d_fr, d_gwc, B, C, H, W, max_disp, ngroups, normalize);
+}
+
+void ffsCudaBuildGWCVolumeHalfToFloat(
+    const __half* d_fl, const __half* d_fr, float* d_gwc,
+    int B, int C, int H, int W, int max_disp, int ngroups, bool normalize, cudaStream_t s)
+{
+    dim3 blk(16, 16);
+    dim3 grd((W + 15) / 16, (H + 15) / 16, B * ngroups * max_disp);
+    buildGWCVolumeKernel<__half, float><<<grd, blk, 0, s>>>(d_fl, d_fr, d_gwc, B, C, H, W, max_disp, ngroups, normalize);
+}
+
+void ffsCudaBuildGWCVolumeMixed(
+    const float* d_fl, const float* d_fr, __half* d_gwc,
+    int B, int C, int H, int W, int max_disp, int ngroups, bool normalize, cudaStream_t s)
+{
+    dim3 blk(16, 16);
+    dim3 grd((W + 15) / 16, (H + 15) / 16, B * ngroups * max_disp);
+    buildGWCVolumeKernel<float, __half><<<grd, blk, 0, s>>>(d_fl, d_fr, d_gwc, B, C, H, W, max_disp, ngroups, normalize);
+}
+
+void ffsCudaPreprocessRGBToCHW(
+    const uint8_t* d_rgb, float* d_chw,
+    int src_h, int src_w, int dst_h, int dst_w, cudaStream_t s)
+{
+    dim3 blk(32, 16);
+    dim3 grd((dst_w + 31) / 32, (dst_h + 15) / 16);
+    preprocessRGBToCHWKernel<<<grd, blk, 0, s>>>(d_rgb, d_chw, src_h, src_w, dst_h, dst_w);
+}
+
+void ffsCudaResizeUniformAndPad(
+    const uint8_t* d_rgb, float* d_chw,
+    int src_h, int src_w, int scaled_h, int scaled_w,
+    int dst_h, int dst_w, cudaStream_t s)
+{
+    dim3 blk(32, 16);
+    dim3 grd((dst_w + 31) / 32, (dst_h + 15) / 16);
+    resizeUniformAndPadKernel<<<grd, blk, 0, s>>>(d_rgb, d_chw, src_h, src_w, scaled_h, scaled_w, dst_h, dst_w);
+}
+
+void ffsCudaCropDisparity(
+    const float* d_src, float* d_dst,
+    int src_h, int src_w, int dst_h, int dst_w, cudaStream_t s)
+{
+    dim3 blk(32, 16);
+    dim3 grd((dst_w + 31) / 32, (dst_h + 15) / 16);
+    cropDisparityKernel<<<grd, blk, 0, s>>>(d_src, d_dst, src_w, dst_h, dst_w);
+}
+
+void ffsCudaUpsampleDisparity(
+    const float* d_src, float* d_dst,
+    int src_w, int src_h, int dst_w, int dst_h,
+    float disp_scale, cudaStream_t s)
+{
+    dim3 blk(32, 16);
+    dim3 grd((dst_w + 31) / 32, (dst_h + 15) / 16);
+    upsampleDisparityKernel<<<grd, blk, 0, s>>>(d_src, d_dst, src_w, src_h, dst_w, dst_h, disp_scale);
+}
+
+void ffsCudaClampDisparity(
+    float* d_disp, int count, float min_val, cudaStream_t s)
+{
+    int threads = 256;
+    int blocks = (count + threads - 1) / threads;
+    clampDisparityKernel<<<blocks, threads, 0, s>>>(d_disp, count, min_val);
+}
+
+void ffsCudaDispToDepth(
+    const float* d_disp, float* d_depth_m,
+    int height, int width, float fx, float baseline_m, cudaStream_t s)
+{
+    dim3 blk(32, 16);
+    dim3 grd((width + 31) / 32, (height + 15) / 16);
+    dispToDepthKernel<<<grd, blk, 0, s>>>(d_disp, d_depth_m, height, width, fx, baseline_m);
+}
+
+}  // extern "C"
+
+}  // namespace cuda
+}  // namespace ffs_depth

cv/fast-foundationstereo/cpp/src/ffs_depth_single_tensorrt.cpp

@@ -0,0 +1,392 @@
+#include "ffs_depth_single_tensorrt.hpp"
+
+#include "ffs_gwc_plugin.hpp"
+
+#include <algorithm>
+#include <cctype>
+#include <cmath>
+#include <cstdint>
+#include <filesystem>
+#include <fstream>
+#include <iostream>
+#include <sstream>
+#include <stdexcept>
+#include <string>
+#include <vector>
+
+namespace ffs_depth {
+
+namespace cuda {
+extern "C" {
+
+void ffsCudaPreprocessRGBToCHW(
+    const uint8_t* d_rgb, float* d_chw,
+    int src_h, int src_w, int dst_h, int dst_w, cudaStream_t s);
+
+void ffsCudaResizeUniformAndPad(
+    const uint8_t* d_rgb, float* d_chw,
+    int src_h, int src_w, int scaled_h, int scaled_w,
+    int dst_h, int dst_w, cudaStream_t s);
+
+void ffsCudaCropDisparity(
+    const float* d_src, float* d_dst,
+    int src_h, int src_w, int dst_h, int dst_w, cudaStream_t s);
+
+void ffsCudaUpsampleDisparity(
+    const float* d_src, float* d_dst,
+    int src_w, int src_h, int dst_w, int dst_h,
+    float disp_scale, cudaStream_t s);
+
+void ffsCudaClampDisparity(
+    float* d_disp, int count, float min_val, cudaStream_t s);
+
+void ffsCudaDispToDepth(
+    const float* d_disp, float* d_depth_m,
+    int height, int width, float fx, float baseline_m, cudaStream_t s);
+
+}  // extern "C"
+}  // namespace cuda
+
+namespace {
+
+class TrtLogger : public nvinfer1::ILogger {
+public:
+    void log(Severity severity, const char* msg) noexcept override {
+        if (severity <= Severity::kWARNING) {
+            std::cerr << "[TRT] " << msg << std::endl;
+        }
+    }
+};
+
+TrtLogger g_trt_logger;
+
+size_t elementSize(nvinfer1::DataType dt) {
+    switch (dt) {
+        case nvinfer1::DataType::kFLOAT: return 4;
+        case nvinfer1::DataType::kHALF:  return 2;
+        case nvinfer1::DataType::kINT8:  return 1;
+        case nvinfer1::DataType::kINT32: return 4;
+        default: return 4;
+    }
+}
+
+void cudaMallocChecked(void** ptr, size_t bytes, const char* what) {
+    cudaError_t err = cudaMalloc(ptr, bytes);
+    if (err != cudaSuccess) {
+        throw std::runtime_error(std::string("[FFS single] cudaMalloc failed for ") +
+                                 what + " (" + std::to_string(bytes) + " bytes): " +
+                                 cudaGetErrorString(err));
+    }
+}
+
+bool hasTensor(nvinfer1::ICudaEngine& engine, const char* name) {
+    for (int32_t i = 0; i < engine.getNbIOTensors(); ++i) {
+        if (std::string(engine.getIOTensorName(i)) == name) return true;
+    }
+    return false;
+}
+
+std::string findSingleYamlConfig(const std::string& engine_dir) {
+    namespace fs = std::filesystem;
+
+    std::vector<fs::path> yaml_files;
+    std::error_code ec;
+    fs::directory_iterator it(engine_dir, ec);
+    if (ec) {
+        throw std::runtime_error("[FFS single] Cannot list engine directory: " +
+                                 engine_dir + " (" + ec.message() + ")");
+    }
+
+    for (const auto& entry : it) {
+        if (ec) {
+            throw std::runtime_error("[FFS single] Cannot list engine directory: " +
+                                     engine_dir + " (" + ec.message() + ")");
+        }
+        if (!entry.is_regular_file()) continue;
+
+        std::string ext = entry.path().extension().string();
+        std::transform(ext.begin(), ext.end(), ext.begin(), [](unsigned char c) {
+            return static_cast<char>(std::tolower(c));
+        });
+        if (ext == ".yaml" || ext == ".yml") {
+            yaml_files.push_back(entry.path());
+        }
+    }
+
+    if (yaml_files.empty()) {
+        throw std::runtime_error("[FFS single] No YAML config found in: " + engine_dir);
+    }
+    if (yaml_files.size() > 1) {
+        std::ostringstream oss;
+        oss << "[FFS single] Expected exactly one YAML config in " << engine_dir
+            << ", found " << yaml_files.size() << ":";
+        for (const auto& path : yaml_files) {
+            oss << " " << path.string();
+        }
+        throw std::runtime_error(oss.str());
+    }
+
+    return yaml_files.front().string();
+}
+
+std::string trim(std::string s) {
+    const char* ws = " \t\r\n";
+    const size_t start = s.find_first_not_of(ws);
+    if (start == std::string::npos) return "";
+    const size_t end = s.find_last_not_of(ws);
+    return s.substr(start, end - start + 1);
+}
+
+std::vector<int> parseInts(std::string s) {
+    for (char& c : s) {
+        if (!(c >= '0' && c <= '9')) c = ' ';
+    }
+    std::stringstream ss(s);
+    std::vector<int> values;
+    int v = 0;
+    while (ss >> v) values.push_back(v);
+    return values;
+}
+
+bool parseBool(std::string s) {
+    s = trim(s);
+    std::transform(s.begin(), s.end(), s.begin(), [](unsigned char c) {
+        return static_cast<char>(std::tolower(c));
+    });
+    return s == "true" || s == "1" || s == "yes" || s == "on";
+}
+
+}  // namespace
+
+FFSSingleEngineInference::FFSSingleEngineInference(const std::string& engine_dir) {
+    cudaError_t err = cudaStreamCreate(&stream_);
+    if (err != cudaSuccess) {
+        throw std::runtime_error(std::string("[FFS single] cudaStreamCreate failed: ") +
+                                 cudaGetErrorString(err));
+    }
+    try {
+        loadConfig(findSingleYamlConfig(engine_dir));
+        if (!registerFFSGWCPlugin()) {
+            throw std::runtime_error("[FFS single] failed to register FFSGWCVolume plugin");
+        }
+        runtime_.reset(nvinfer1::createInferRuntime(g_trt_logger));
+        if (!runtime_) {
+            throw std::runtime_error("[FFS single] createInferRuntime returned null");
+        }
+        loadEngine(engine_dir + "/fast_foundationstereo.engine");
+        allocateBuffers();
+    } catch (...) {
+        freeDeviceBuffers();
+        if (stream_) {
+            cudaStreamDestroy(stream_);
+            stream_ = nullptr;
+        }
+        throw;
+    }
+}
+
+FFSSingleEngineInference::~FFSSingleEngineInference() {
+    freeDeviceBuffers();
+    if (stream_) {
+        cudaStreamDestroy(stream_);
+        stream_ = nullptr;
+    }
+}
+
+void FFSSingleEngineInference::freeDeviceBuffers() {
+    auto free = [](void*& p) { if (p) { cudaFree(p); p = nullptr; } };
+    free(reinterpret_cast<void*&>(d_left_));
+    free(reinterpret_cast<void*&>(d_right_));
+    free(reinterpret_cast<void*&>(d_disp_));
+    free(reinterpret_cast<void*&>(d_disp_cropped_));
+    free(reinterpret_cast<void*&>(d_disp_for_depth_));
+    depth_alloc_pixels_ = 0;
+}
+
+void FFSSingleEngineInference::loadConfig(const std::string& path) {
+    std::ifstream f(path);
+    if (!f.good()) throw std::runtime_error("[FFS single] Cannot open config: " + path);
+
+    std::string line;
+    int image_index = -1;
+    while (std::getline(f, line)) {
+        const size_t comment = line.find('#');
+        if (comment != std::string::npos) line = line.substr(0, comment);
+        line = trim(line);
+        if (line.empty()) continue;
+
+        if (image_index >= 0) {
+            if (line.rfind("-", 0) == 0) {
+                const auto values = parseInts(line);
+                if (!values.empty()) {
+                    if (image_index == 0) config_.image_height = values[0];
+                    if (image_index == 1) config_.image_width = values[0];
+                    ++image_index;
+                    if (image_index >= 2) image_index = -1;
+                    continue;
+                }
+            }
+            image_index = -1;
+        }
+
+        const size_t colon = line.find(':');
+        if (colon == std::string::npos) continue;
+        const std::string key = trim(line.substr(0, colon));
+        const std::string value = trim(line.substr(colon + 1));
+
+        if (key == "image_size") {
+            const auto values = parseInts(value);
+            if (values.size() >= 2) {
+                config_.image_height = values[0];
+                config_.image_width = values[1];
+            } else {
+                image_index = 0;
+            }
+        } else if (key == "max_disp") {
+            const auto values = parseInts(value);
+            if (!values.empty()) config_.max_disp = values[0];
+        } else if (key == "cv_group") {
+            const auto values = parseInts(value);
+            if (!values.empty()) config_.cv_group = values[0];
+        } else if (key == "valid_iters") {
+            const auto values = parseInts(value);
+            if (!values.empty()) config_.valid_iters = values[0];
+        } else if (key == "normalize") {
+            config_.normalize = parseBool(value);
+        }
+    }
+}
+
+void FFSSingleEngineInference::loadEngine(const std::string& path) {
+    std::ifstream f(path, std::ios::binary);
+    if (!f.good()) throw std::runtime_error("[FFS single] Cannot open engine: " + path);
+
+    f.seekg(0, std::ios::end);
+    size_t sz = f.tellg();
+    f.seekg(0, std::ios::beg);
+    std::vector<char> buf(sz);
+    f.read(buf.data(), sz);
+
+    engine_.reset(runtime_->deserializeCudaEngine(buf.data(), sz));
+    if (!engine_) throw std::runtime_error("[FFS single] Deserialize failed: " + path);
+
+    context_.reset(engine_->createExecutionContext());
+    if (!context_) throw std::runtime_error("[FFS single] Context creation failed: " + path);
+
+    if (!hasTensor(*engine_, "left") || !hasTensor(*engine_, "right") || !hasTensor(*engine_, "disp")) {
+        throw std::runtime_error("[FFS single] Engine must expose tensors named left, right, and disp");
+    }
+}
+
+void FFSSingleEngineInference::allocateBuffers() {
+    const size_t H = static_cast<size_t>(config_.image_height);
+    const size_t W = static_cast<size_t>(config_.image_width);
+    const size_t input_bytes = 3 * H * W * sizeof(float);
+
+    cudaMallocChecked(reinterpret_cast<void**>(&d_left_), input_bytes, "d_left_");
+    cudaMallocChecked(reinterpret_cast<void**>(&d_right_), input_bytes, "d_right_");
+
+    const size_t disp_bytes = H * W * elementSize(engine_->getTensorDataType("disp"));
+    cudaMallocChecked(reinterpret_cast<void**>(&d_disp_), disp_bytes, "d_disp_");
+    cudaMallocChecked(reinterpret_cast<void**>(&d_disp_cropped_),
+                      H * W * sizeof(float), "d_disp_cropped_");
+
+    context_->setTensorAddress("left", d_left_);
+    context_->setTensorAddress("right", d_right_);
+    context_->setTensorAddress("disp", d_disp_);
+}
+
+void FFSSingleEngineInference::preprocessRGBGPU(
+    const uint8_t* d_rgb, int src_h, int src_w, float* d_output) {
+    const int mH = config_.image_height;
+    const int mW = config_.image_width;
+    if (src_h == mH && src_w == mW) {
+        scaled_h_ = mH;
+        scaled_w_ = mW;
+        cuda::ffsCudaPreprocessRGBToCHW(d_rgb, d_output, src_h, src_w, mH, mW, stream_);
+    } else {
+        const float scale = std::min(static_cast<float>(mW) / src_w,
+                                     static_cast<float>(mH) / src_h);
+        scaled_w_ = std::max(1, static_cast<int>(std::round(src_w * scale)));
+        scaled_h_ = std::max(1, static_cast<int>(std::round(src_h * scale)));
+        cuda::ffsCudaResizeUniformAndPad(
+            d_rgb, d_output, src_h, src_w, scaled_h_, scaled_w_, mH, mW, stream_);
+    }
+}
+
+void FFSSingleEngineInference::infer(
+    const uint8_t* d_left_rgb, const uint8_t* d_right_rgb,
+    int input_h, int input_w,
+    float* d_disp_out) {
+    if (!d_left_rgb || !d_right_rgb || !d_disp_out) {
+        throw std::runtime_error("[FFS single] infer: null device pointer");
+    }
+    if (input_h <= 0 || input_w <= 0) {
+        throw std::runtime_error("[FFS single] infer: input dimensions must be positive");
+    }
+
+    const int mH = config_.image_height;
+    const int mW = config_.image_width;
+    const bool needs_resize = (input_h != mH || input_w != mW);
+
+    preprocessRGBGPU(d_left_rgb, input_h, input_w, d_left_);
+    preprocessRGBGPU(d_right_rgb, input_h, input_w, d_right_);
+
+    if (!needs_resize) {
+        context_->setTensorAddress("disp", d_disp_out);
+    }
+    if (!context_->enqueueV3(stream_)) {
+        throw std::runtime_error("[FFS single] enqueue failed");
+    }
+
+    if (needs_resize) {
+        cuda::ffsCudaClampDisparity(d_disp_, mH * mW, 0.0f, stream_);
+        cuda::ffsCudaCropDisparity(d_disp_, d_disp_cropped_,
+                                   mH, mW, scaled_h_, scaled_w_, stream_);
+        const float disp_scale = static_cast<float>(input_w) / scaled_w_;
+        cuda::ffsCudaUpsampleDisparity(d_disp_cropped_, d_disp_out,
+                                       scaled_w_, scaled_h_, input_w, input_h,
+                                       disp_scale, stream_);
+    } else {
+        cuda::ffsCudaClampDisparity(d_disp_out, mH * mW, 0.0f, stream_);
+        context_->setTensorAddress("disp", d_disp_);
+    }
+}
+
+void FFSSingleEngineInference::dispToDepth(
+    const float* d_disp,
+    int height, int width,
+    float fx, float baseline_m,
+    float* d_depth_out) {
+    cuda::ffsCudaDispToDepth(d_disp, d_depth_out,
+                             height, width, fx, baseline_m, stream_);
+}
+
+void FFSSingleEngineInference::inferDepth(
+    const uint8_t* d_left_rgb, const uint8_t* d_right_rgb,
+    int input_h, int input_w,
+    float fx, float baseline_m,
+    float* d_depth_out) {
+    if (input_h <= 0 || input_w <= 0) {
+        throw std::runtime_error("[FFS single] inferDepth: input dimensions must be positive");
+    }
+    const size_t num_pixels = static_cast<size_t>(input_h) * static_cast<size_t>(input_w);
+    if (static_cast<int64_t>(num_pixels) > depth_alloc_pixels_) {
+        if (d_disp_for_depth_) {
+            cudaFree(d_disp_for_depth_);
+            d_disp_for_depth_ = nullptr;
+        }
+        cudaMallocChecked(reinterpret_cast<void**>(&d_disp_for_depth_),
+                          num_pixels * sizeof(float), "d_disp_for_depth_");
+        depth_alloc_pixels_ = static_cast<int64_t>(num_pixels);
+    }
+    infer(d_left_rgb, d_right_rgb, input_h, input_w, d_disp_for_depth_);
+    cuda::ffsCudaDispToDepth(d_disp_for_depth_, d_depth_out,
+                             input_h, input_w, fx, baseline_m, stream_);
+}
+
+void FFSSingleEngineInference::sync() {
+    cudaStreamSynchronize(stream_);
+}
+
+}  // namespace ffs_depth

cv/fast-foundationstereo/cpp/src/ffs_depth_tensorrt.cpp

@@ -0,0 +1,488 @@
+#include "ffs_depth_tensorrt.hpp"
+
+#include <cuda_fp16.h>
+
+#include <algorithm>
+#include <cctype>
+#include <cmath>
+#include <cstdint>
+#include <fstream>
+#include <iostream>
+#include <sstream>
+#include <stdexcept>
+#include <string>
+#include <vector>
+
+namespace ffs_depth {
+
+// Forward declarations of CUDA kernel wrappers (defined in depth_kernels.cu)
+namespace cuda {
+extern "C" {
+
+void ffsCudaBuildGWCVolumeMixed(
+    const float* d_fl, const float* d_fr, __half* d_gwc,
+    int B, int C, int H, int W, int max_disp, int ngroups, bool normalize, cudaStream_t s);
+
+void ffsCudaPreprocessRGBToCHW(
+    const uint8_t* d_rgb, float* d_chw,
+    int src_h, int src_w, int dst_h, int dst_w, cudaStream_t s);
+
+void ffsCudaResizeUniformAndPad(
+    const uint8_t* d_rgb, float* d_chw,
+    int src_h, int src_w, int scaled_h, int scaled_w,
+    int dst_h, int dst_w, cudaStream_t s);
+
+void ffsCudaCropDisparity(
+    const float* d_src, float* d_dst,
+    int src_h, int src_w, int dst_h, int dst_w, cudaStream_t s);
+
+void ffsCudaUpsampleDisparity(
+    const float* d_src, float* d_dst,
+    int src_w, int src_h, int dst_w, int dst_h,
+    float disp_scale, cudaStream_t s);
+
+void ffsCudaClampDisparity(
+    float* d_disp, int count, float min_val, cudaStream_t s);
+
+void ffsCudaDispToDepth(
+    const float* d_disp, float* d_depth_m,
+    int height, int width, float fx, float baseline_m, cudaStream_t s);
+
+}  // extern "C"
+}  // namespace cuda
+
+namespace {
+
+class TrtLogger : public nvinfer1::ILogger {
+public:
+    void log(Severity severity, const char* msg) noexcept override {
+        if (severity <= Severity::kWARNING)
+            std::cerr << "[TRT] " << msg << std::endl;
+    }
+};
+
+TrtLogger g_trt_logger;
+
+size_t elementSize(nvinfer1::DataType dt) {
+    switch (dt) {
+        case nvinfer1::DataType::kFLOAT: return 4;
+        case nvinfer1::DataType::kHALF:  return 2;
+        case nvinfer1::DataType::kINT8:  return 1;
+        case nvinfer1::DataType::kINT32: return 4;
+        default: return 4;
+    }
+}
+
+void cudaMallocChecked(void** ptr, size_t bytes, const char* what) {
+    cudaError_t err = cudaMalloc(ptr, bytes);
+    if (err != cudaSuccess) {
+        throw std::runtime_error(std::string("[FFS] cudaMalloc failed for ") + what +
+                                 " (" + std::to_string(bytes) + " bytes): " +
+                                 cudaGetErrorString(err));
+    }
+}
+
+std::string trim(std::string s) {
+    const char* ws = " \t\r\n";
+    const size_t start = s.find_first_not_of(ws);
+    if (start == std::string::npos) return "";
+    const size_t end = s.find_last_not_of(ws);
+    return s.substr(start, end - start + 1);
+}
+
+std::vector<int> parseInts(std::string s) {
+    for (char& c : s) {
+        if (!(c >= '0' && c <= '9')) c = ' ';
+    }
+    std::stringstream ss(s);
+    std::vector<int> values;
+    int v = 0;
+    while (ss >> v) values.push_back(v);
+    return values;
+}
+
+bool parseBool(std::string s) {
+    s = trim(s);
+    std::transform(s.begin(), s.end(), s.begin(), [](unsigned char c) {
+        return static_cast<char>(std::tolower(c));
+    });
+    return s == "true" || s == "1" || s == "yes" || s == "on";
+}
+
+}  // anonymous namespace
+
+// =========================================================================
+// Construction / destruction
+// =========================================================================
+
+FFSDepthInference::FFSDepthInference(const std::string& engine_dir) {
+    cudaError_t err = cudaStreamCreate(&stream_);
+    if (err != cudaSuccess) {
+        throw std::runtime_error(std::string("[FFS] cudaStreamCreate failed: ") +
+                                 cudaGetErrorString(err));
+    }
+    try {
+        loadConfig(engine_dir + "/onnx.yaml");
+        runtime_.reset(nvinfer1::createInferRuntime(g_trt_logger));
+        if (!runtime_) {
+            throw std::runtime_error("[FFS] createInferRuntime returned null");
+        }
+        loadEngine(engine_dir + "/feature_runner.engine", feature_engine_, feature_context_);
+        loadEngine(engine_dir + "/post_runner.engine", post_engine_, post_context_);
+        allocateBuffers();
+    } catch (...) {
+        // Best-effort cleanup of anything allocated so far. Buffers may be partially
+        // populated; freeDeviceBuffers() is a no-op for null pointers.
+        freeDeviceBuffers();
+        if (stream_) {
+            cudaStreamDestroy(stream_);
+            stream_ = nullptr;
+        }
+        throw;
+    }
+}
+
+
+FFSDepthInference::~FFSDepthInference() {
+    freeDeviceBuffers();
+    if (stream_) {
+        cudaStreamDestroy(stream_);
+        stream_ = nullptr;
+    }
+}
+
+void FFSDepthInference::freeDeviceBuffers() {
+    auto free = [](void*& p) { if (p) { cudaFree(p); p = nullptr; } };
+
+    free(reinterpret_cast<void*&>(d_left_));
+    free(reinterpret_cast<void*&>(d_right_));
+    free(reinterpret_cast<void*&>(d_feat_left_04_));
+    free(reinterpret_cast<void*&>(d_feat_left_08_));
+    free(reinterpret_cast<void*&>(d_feat_left_16_));
+    free(reinterpret_cast<void*&>(d_feat_left_32_));
+    free(reinterpret_cast<void*&>(d_feat_right_04_));
+    free(reinterpret_cast<void*&>(d_stem_2x_));
+    free(reinterpret_cast<void*&>(d_gwc_volume_));
+    free(reinterpret_cast<void*&>(d_disp_));
+    free(reinterpret_cast<void*&>(d_disp_cropped_));
+    free(reinterpret_cast<void*&>(d_disp_for_depth_));
+    depth_alloc_pixels_ = 0;
+}
+
+// =========================================================================
+// Config / engine loading
+// =========================================================================
+
+void FFSDepthInference::loadConfig(const std::string& path) {
+    std::ifstream f(path);
+    if (!f.good()) throw std::runtime_error("[FFS] Cannot open config: " + path);
+
+    std::string line;
+    int image_index = -1;
+    while (std::getline(f, line)) {
+        const size_t comment = line.find('#');
+        if (comment != std::string::npos) line = line.substr(0, comment);
+        line = trim(line);
+        if (line.empty()) continue;
+
+        if (image_index >= 0) {
+            if (line.rfind("-", 0) == 0) {
+                const auto values = parseInts(line);
+                if (!values.empty()) {
+                    if (image_index == 0) config_.image_height = values[0];
+                    if (image_index == 1) config_.image_width = values[0];
+                    ++image_index;
+                    if (image_index >= 2) image_index = -1;
+                    continue;
+                }
+            }
+            image_index = -1;
+        }
+
+        const size_t colon = line.find(':');
+        if (colon == std::string::npos) continue;
+        const std::string key = trim(line.substr(0, colon));
+        const std::string value = trim(line.substr(colon + 1));
+
+        if (key == "image_size") {
+            const auto values = parseInts(value);
+            if (values.size() >= 2) {
+                config_.image_height = values[0];
+                config_.image_width = values[1];
+            } else {
+                image_index = 0;
+            }
+        } else if (key == "max_disp") {
+            const auto values = parseInts(value);
+            if (!values.empty()) config_.max_disp = values[0];
+        } else if (key == "cv_group") {
+            const auto values = parseInts(value);
+            if (!values.empty()) config_.cv_group = values[0];
+        } else if (key == "valid_iters") {
+            const auto values = parseInts(value);
+            if (!values.empty()) config_.valid_iters = values[0];
+        } else if (key == "normalize") {
+            config_.normalize = parseBool(value);
+        }
+    }
+
+    gwc_disp_levels_ = config_.max_disp / 4;
+}
+
+void FFSDepthInference::loadEngine(
+    const std::string& path,
+    std::unique_ptr<nvinfer1::ICudaEngine>& engine,
+    std::unique_ptr<nvinfer1::IExecutionContext>& ctx)
+{
+    std::ifstream f(path, std::ios::binary);
+    if (!f.good()) throw std::runtime_error("[FFS] Cannot open engine: " + path);
+
+    f.seekg(0, std::ios::end);
+    size_t sz = f.tellg();
+    f.seekg(0, std::ios::beg);
+    std::vector<char> buf(sz);
+    f.read(buf.data(), sz);
+
+    engine.reset(runtime_->deserializeCudaEngine(buf.data(), sz));
+    if (!engine) throw std::runtime_error("[FFS] Deserialize failed: " + path);
+
+    ctx.reset(engine->createExecutionContext());
+    if (!ctx) throw std::runtime_error("[FFS] Context creation failed: " + path);
+}
+
+// =========================================================================
+// Buffer allocation
+// =========================================================================
+
+void FFSDepthInference::allocateBuffers() {
+    const size_t H = static_cast<size_t>(config_.image_height);
+    const size_t W = static_cast<size_t>(config_.image_width);
+    const size_t input_bytes = 3 * H * W * sizeof(float);
+
+    cudaMallocChecked(reinterpret_cast<void**>(&d_left_),  input_bytes, "d_left_");
+    cudaMallocChecked(reinterpret_cast<void**>(&d_right_), input_bytes, "d_right_");
+
+    allocateFeatureBuffers();
+    allocatePostBuffers();
+}
+
+void FFSDepthInference::allocateFeatureBuffers() {
+    auto getDims = [&](const char* name) {
+        auto d = feature_engine_->getTensorShape(name);
+        std::vector<int> v(d.nbDims);
+        for (int i = 0; i < d.nbDims; ++i) v[i] = d.d[i];
+        return v;
+    };
+
+    auto allocTensor = [&](const char* name, const std::vector<int>& dims) -> float* {
+        auto dt = feature_engine_->getTensorDataType(name);
+        size_t bytes = elementSize(dt);
+        for (int d : dims) {
+            if (d <= 0) {
+                throw std::runtime_error(std::string("[FFS] tensor '") + name +
+                                         "' has non-positive dimension");
+            }
+            bytes *= static_cast<size_t>(d);
+        }
+        void* p = nullptr;
+        cudaMallocChecked(&p, bytes, name);
+        return static_cast<float*>(p);
+    };
+
+    feat_04_dims_  = getDims("features_left_04");
+    feat_08_dims_  = getDims("features_left_08");
+    feat_16_dims_  = getDims("features_left_16");
+    feat_32_dims_  = getDims("features_left_32");
+    stem_2x_dims_  = getDims("stem_2x");
+
+    d_feat_left_04_  = allocTensor("features_left_04",  feat_04_dims_);
+    d_feat_left_08_  = allocTensor("features_left_08",  feat_08_dims_);
+    d_feat_left_16_  = allocTensor("features_left_16",  feat_16_dims_);
+    d_feat_left_32_  = allocTensor("features_left_32",  feat_32_dims_);
+    d_feat_right_04_ = allocTensor("features_right_04", feat_04_dims_);
+    d_stem_2x_       = allocTensor("stem_2x",           stem_2x_dims_);
+
+    feature_context_->setTensorAddress("left",               d_left_);
+    feature_context_->setTensorAddress("right",              d_right_);
+    feature_context_->setTensorAddress("features_left_04",   d_feat_left_04_);
+    feature_context_->setTensorAddress("features_left_08",   d_feat_left_08_);
+    feature_context_->setTensorAddress("features_left_16",   d_feat_left_16_);
+    feature_context_->setTensorAddress("features_left_32",   d_feat_left_32_);
+    feature_context_->setTensorAddress("features_right_04",  d_feat_right_04_);
+    feature_context_->setTensorAddress("stem_2x",            d_stem_2x_);
+}
+
+void FFSDepthInference::allocatePostBuffers() {
+    const size_t H  = static_cast<size_t>(config_.image_height);
+    const size_t W  = static_cast<size_t>(config_.image_width);
+    const size_t H4 = H / 4;
+    const size_t W4 = W / 4;
+
+    const auto gwc_dt = post_engine_->getTensorDataType("gwc_volume");
+    gwc_fp16_ = (gwc_dt == nvinfer1::DataType::kHALF);
+    const size_t gwc_elem = gwc_fp16_ ? sizeof(__half) : sizeof(float);
+    const size_t gwc_bytes =
+        static_cast<size_t>(config_.cv_group) *
+        static_cast<size_t>(gwc_disp_levels_) * H4 * W4 * gwc_elem;
+    cudaMallocChecked(reinterpret_cast<void**>(&d_gwc_volume_), gwc_bytes, "gwc_volume");
+
+    const auto disp_dt = post_engine_->getTensorDataType("disp");
+    const size_t disp_bytes = H * W * elementSize(disp_dt);
+    cudaMallocChecked(reinterpret_cast<void**>(&d_disp_), disp_bytes, "d_disp_");
+
+    const size_t disp_cropped_bytes = H * W * sizeof(float);
+    cudaMallocChecked(reinterpret_cast<void**>(&d_disp_cropped_),
+                      disp_cropped_bytes, "d_disp_cropped_");
+
+    post_context_->setTensorAddress("features_left_04",  d_feat_left_04_);
+    post_context_->setTensorAddress("features_left_08",  d_feat_left_08_);
+    post_context_->setTensorAddress("features_left_16",  d_feat_left_16_);
+    post_context_->setTensorAddress("features_left_32",  d_feat_left_32_);
+    post_context_->setTensorAddress("features_right_04", d_feat_right_04_);
+    post_context_->setTensorAddress("stem_2x",           d_stem_2x_);
+    post_context_->setTensorAddress("gwc_volume",        d_gwc_volume_);
+    post_context_->setTensorAddress("disp",              d_disp_);
+}
+
+// =========================================================================
+// Pre-processing
+// =========================================================================
+
+void FFSDepthInference::preprocessRGBGPU(
+    const uint8_t* d_rgb, int src_h, int src_w, float* d_output)
+{
+    const int mH = config_.image_height;
+    const int mW = config_.image_width;
+
+    if (src_h == mH && src_w == mW) {
+        scaled_h_ = mH;
+        scaled_w_ = mW;
+        cuda::ffsCudaPreprocessRGBToCHW(
+            d_rgb, d_output, src_h, src_w, mH, mW, stream_);
+    } else {
+        const float scale = std::min(static_cast<float>(mW) / src_w,
+                                     static_cast<float>(mH) / src_h);
+        scaled_w_ = std::max(1, static_cast<int>(std::round(src_w * scale)));
+        scaled_h_ = std::max(1, static_cast<int>(std::round(src_h * scale)));
+        cuda::ffsCudaResizeUniformAndPad(
+            d_rgb, d_output, src_h, src_w, scaled_h_, scaled_w_, mH, mW, stream_);
+    }
+}
+
+// =========================================================================
+// Pipeline stages
+// =========================================================================
+
+void FFSDepthInference::runFeatureRunner() {
+    if (!feature_context_->enqueueV3(stream_))
+        throw std::runtime_error("[FFS] feature_runner failed");
+}
+
+void FFSDepthInference::buildGWCVolume() {
+    int B = feat_04_dims_[0];
+    int C = feat_04_dims_[1];
+    int H = feat_04_dims_[2];
+    int W = feat_04_dims_[3];
+
+    // Only mixed precision path is used (FP32 features -> FP16 GWC volume)
+    cuda::ffsCudaBuildGWCVolumeMixed(
+        d_feat_left_04_, d_feat_right_04_,
+        reinterpret_cast<__half*>(d_gwc_volume_),
+        B, C, H, W, gwc_disp_levels_, config_.cv_group, config_.normalize, stream_);
+
+}
+
+void FFSDepthInference::runPostRunner() {
+    if (!post_context_->enqueueV3(stream_))
+        throw std::runtime_error("[FFS] post_runner failed");
+}
+
+
+// =========================================================================
+// Public inference entry point
+// =========================================================================
+
+void FFSDepthInference::infer(
+    const uint8_t* d_left_rgb, const uint8_t* d_right_rgb,
+    int input_h, int input_w,
+    float* d_disp_out)
+{
+    if (!d_left_rgb || !d_right_rgb || !d_disp_out) {
+        throw std::runtime_error("[FFS] infer: null device pointer");
+    }
+    if (input_h <= 0 || input_w <= 0) {
+        throw std::runtime_error("[FFS] infer: input dimensions must be positive");
+    }
+    const int mH = config_.image_height;
+    const int mW = config_.image_width;
+    const bool needs_resize = (input_h != mH || input_w != mW);
+    preprocessRGBGPU(d_left_rgb,  input_h, input_w, d_left_);
+    preprocessRGBGPU(d_right_rgb, input_h, input_w, d_right_);
+    runFeatureRunner();
+    buildGWCVolume();
+    if (!needs_resize) {
+        post_context_->setTensorAddress("disp", d_disp_out);
+    }
+    runPostRunner();
+    if (needs_resize) {
+        cuda::ffsCudaClampDisparity(d_disp_, mH * mW, 0.0f, stream_);
+        // Crop padding: model (mW x mH) -> scaled (scaled_w_ x scaled_h_)
+        cuda::ffsCudaCropDisparity(
+            d_disp_, d_disp_cropped_,
+            mH, mW, scaled_h_, scaled_w_, stream_);
+        // Upsample to input resolution with disparity scale correction
+        float disp_scale = static_cast<float>(input_w) / scaled_w_;
+        cuda::ffsCudaUpsampleDisparity(
+            d_disp_cropped_, d_disp_out,
+            scaled_w_, scaled_h_, input_w, input_h,
+            disp_scale, stream_);
+    } else {
+        cuda::ffsCudaClampDisparity(d_disp_out, mH * mW, 0.0f, stream_);
+        post_context_->setTensorAddress("disp", d_disp_);
+    }
+}
+
+void FFSDepthInference::sync() {
+    cudaStreamSynchronize(stream_);
+}
+
+void FFSDepthInference::dispToDepth(
+    const float* d_disp,
+    int height, int width,
+    float fx, float baseline_m,
+    float* d_depth_out)
+{
+    cuda::ffsCudaDispToDepth(d_disp, d_depth_out,
+                             height, width, fx, baseline_m, stream_);
+}
+
+void FFSDepthInference::inferDepth(
+    const uint8_t* d_left_rgb, const uint8_t* d_right_rgb,
+    int input_h, int input_w,
+    float fx, float baseline_m,
+    float* d_depth_out)
+{
+    if (input_h <= 0 || input_w <= 0) {
+        throw std::runtime_error("[FFS] inferDepth: input dimensions must be positive");
+    }
+    const size_t num_pixels = static_cast<size_t>(input_h) * static_cast<size_t>(input_w);
+
+    if (static_cast<int64_t>(num_pixels) > depth_alloc_pixels_) {
+        // cudaFree is implicitly stream-ordered: it waits for prior work on the
+        // device before reclaiming the allocation, so this is safe to call here
+        // even though earlier inferDepth() calls may still be in flight.
+        if (d_disp_for_depth_) {
+            cudaFree(d_disp_for_depth_);
+            d_disp_for_depth_ = nullptr;
+        }
+        cudaMallocChecked(reinterpret_cast<void**>(&d_disp_for_depth_),
+                          num_pixels * sizeof(float), "d_disp_for_depth_");
+        depth_alloc_pixels_ = static_cast<int64_t>(num_pixels);
+    }
+
+    infer(d_left_rgb, d_right_rgb, input_h, input_w, d_disp_for_depth_);
+    cuda::ffsCudaDispToDepth(d_disp_for_depth_, d_depth_out,
+                             input_h, input_w, fx, baseline_m, stream_);
+}
+
+}  // namespace ffs_depth

cv/fast-foundationstereo/cpp/src/gwc_volume_plugin.cpp

@@ -0,0 +1,325 @@
+#include "ffs_gwc_plugin.hpp"
+
+#include <NvInfer.h>
+#include <NvInferRuntimePlugin.h>
+#include <cuda_fp16.h>
+#include <cuda_runtime.h>
+
+#include <cstring>
+#include <iostream>
+#include <mutex>
+#include <string>
+#include <vector>
+
+namespace ffs_depth {
+namespace cuda {
+extern "C" {
+
+void ffsCudaBuildGWCVolumeFloat(
+    const float* d_fl, const float* d_fr, float* d_gwc,
+    int B, int C, int H, int W, int max_disp, int ngroups, bool normalize, cudaStream_t s);
+
+void ffsCudaBuildGWCVolumeHalf(
+    const __half* d_fl, const __half* d_fr, __half* d_gwc,
+    int B, int C, int H, int W, int max_disp, int ngroups, bool normalize, cudaStream_t s);
+
+void ffsCudaBuildGWCVolumeHalfToFloat(
+    const __half* d_fl, const __half* d_fr, float* d_gwc,
+    int B, int C, int H, int W, int max_disp, int ngroups, bool normalize, cudaStream_t s);
+
+void ffsCudaBuildGWCVolumeMixed(
+    const float* d_fl, const float* d_fr, __half* d_gwc,
+    int B, int C, int H, int W, int max_disp, int ngroups, bool normalize, cudaStream_t s);
+
+}  // extern "C"
+}  // namespace cuda
+
+namespace {
+
+constexpr char kPluginName[] = "FFSGWCVolume";
+constexpr char kPluginVersion[] = "1";
+
+struct GWCParams {
+    int32_t max_disp = 0;   // Disparity levels at feature resolution, e.g. max_disp / 4.
+    int32_t cv_group = 0;
+    int32_t normalize = 1;
+};
+
+int32_t fieldToInt(nvinfer1::PluginField const& field, int32_t fallback) {
+    if (!field.data || field.length <= 0) return fallback;
+    if (field.type == nvinfer1::PluginFieldType::kINT32) {
+        return *static_cast<int32_t const*>(field.data);
+    }
+    if (field.type == nvinfer1::PluginFieldType::kINT64) {
+        return static_cast<int32_t>(*static_cast<int64_t const*>(field.data));
+    }
+    return fallback;
+}
+
+class FFSGWCVolumePlugin final : public nvinfer1::IPluginV2DynamicExt {
+public:
+    explicit FFSGWCVolumePlugin(GWCParams params) : params_(params) {}
+
+    FFSGWCVolumePlugin(void const* data, size_t length) {
+        if (data && length == sizeof(GWCParams)) {
+            std::memcpy(&params_, data, sizeof(GWCParams));
+        }
+    }
+
+    nvinfer1::IPluginV2DynamicExt* clone() const noexcept override {
+        auto* plugin = new FFSGWCVolumePlugin(params_);
+        plugin->setPluginNamespace(namespace_.c_str());
+        return plugin;
+    }
+
+    char const* getPluginType() const noexcept override { return kPluginName; }
+    char const* getPluginVersion() const noexcept override { return kPluginVersion; }
+    int32_t getNbOutputs() const noexcept override { return 1; }
+
+    nvinfer1::DimsExprs getOutputDimensions(
+        int32_t outputIndex,
+        nvinfer1::DimsExprs const* inputs,
+        int32_t nbInputs,
+        nvinfer1::IExprBuilder& exprBuilder) noexcept override {
+        nvinfer1::DimsExprs out{};
+        if (outputIndex != 0 || nbInputs != 2 || inputs[0].nbDims != 4) {
+            out.nbDims = -1;
+            return out;
+        }
+
+        out.nbDims = 5;
+        out.d[0] = inputs[0].d[0];  // B
+        out.d[1] = exprBuilder.constant(params_.cv_group);
+        out.d[2] = exprBuilder.constant(params_.max_disp);
+        out.d[3] = inputs[0].d[2];  // H
+        out.d[4] = inputs[0].d[3];  // W
+        return out;
+    }
+
+    bool supportsFormatCombination(
+        int32_t pos,
+        nvinfer1::PluginTensorDesc const* inOut,
+        int32_t nbInputs,
+        int32_t nbOutputs) noexcept override {
+        if (nbInputs != 2 || nbOutputs != 1 || pos < 0 || pos >= 3) return false;
+        if (inOut[pos].format != nvinfer1::TensorFormat::kLINEAR) return false;
+
+        if (pos == 0) {
+            return inOut[0].type == nvinfer1::DataType::kFLOAT ||
+                   inOut[0].type == nvinfer1::DataType::kHALF;
+        }
+        if (pos == 1) {
+            return inOut[1].type == inOut[0].type;
+        }
+        return inOut[2].type == nvinfer1::DataType::kFLOAT ||
+               inOut[2].type == nvinfer1::DataType::kHALF;
+    }
+
+    void configurePlugin(
+        nvinfer1::DynamicPluginTensorDesc const*,
+        int32_t,
+        nvinfer1::DynamicPluginTensorDesc const*,
+        int32_t) noexcept override {}
+
+    size_t getWorkspaceSize(
+        nvinfer1::PluginTensorDesc const*,
+        int32_t,
+        nvinfer1::PluginTensorDesc const*,
+        int32_t) const noexcept override {
+        return 0;
+    }
+
+    int32_t enqueue(
+        nvinfer1::PluginTensorDesc const* inputDesc,
+        nvinfer1::PluginTensorDesc const* outputDesc,
+        void const* const* inputs,
+        void* const* outputs,
+        void*,
+        cudaStream_t stream) noexcept override {
+        if (!inputs || !outputs || !inputs[0] || !inputs[1] || !outputs[0]) return 1;
+        if (inputDesc[0].dims.nbDims != 4) return 1;
+
+        const int B = inputDesc[0].dims.d[0];
+        const int C = inputDesc[0].dims.d[1];
+        const int H = inputDesc[0].dims.d[2];
+        const int W = inputDesc[0].dims.d[3];
+        const bool normalize = params_.normalize != 0;
+
+        if (inputDesc[0].type == nvinfer1::DataType::kFLOAT &&
+            outputDesc[0].type == nvinfer1::DataType::kFLOAT) {
+            cuda::ffsCudaBuildGWCVolumeFloat(
+                static_cast<float const*>(inputs[0]),
+                static_cast<float const*>(inputs[1]),
+                static_cast<float*>(outputs[0]),
+                B, C, H, W, params_.max_disp, params_.cv_group, normalize, stream);
+            return 0;
+        }
+        if (inputDesc[0].type == nvinfer1::DataType::kFLOAT &&
+            outputDesc[0].type == nvinfer1::DataType::kHALF) {
+            cuda::ffsCudaBuildGWCVolumeMixed(
+                static_cast<float const*>(inputs[0]),
+                static_cast<float const*>(inputs[1]),
+                static_cast<__half*>(outputs[0]),
+                B, C, H, W, params_.max_disp, params_.cv_group, normalize, stream);
+            return 0;
+        }
+        if (inputDesc[0].type == nvinfer1::DataType::kHALF &&
+            outputDesc[0].type == nvinfer1::DataType::kHALF) {
+            cuda::ffsCudaBuildGWCVolumeHalf(
+                static_cast<__half const*>(inputs[0]),
+                static_cast<__half const*>(inputs[1]),
+                static_cast<__half*>(outputs[0]),
+                B, C, H, W, params_.max_disp, params_.cv_group, normalize, stream);
+            return 0;
+        }
+        if (inputDesc[0].type == nvinfer1::DataType::kHALF &&
+            outputDesc[0].type == nvinfer1::DataType::kFLOAT) {
+            cuda::ffsCudaBuildGWCVolumeHalfToFloat(
+                static_cast<__half const*>(inputs[0]),
+                static_cast<__half const*>(inputs[1]),
+                static_cast<float*>(outputs[0]),
+                B, C, H, W, params_.max_disp, params_.cv_group, normalize, stream);
+            return 0;
+        }
+
+        return 1;
+    }
+
+    nvinfer1::DataType getOutputDataType(
+        int32_t,
+        nvinfer1::DataType const* inputTypes,
+        int32_t nbInputs) const noexcept override {
+        if (nbInputs > 0 && inputTypes[0] == nvinfer1::DataType::kHALF) {
+            return nvinfer1::DataType::kHALF;
+        }
+        return nvinfer1::DataType::kFLOAT;
+    }
+
+    int32_t initialize() noexcept override { return 0; }
+    void terminate() noexcept override {}
+    size_t getSerializationSize() const noexcept override { return sizeof(GWCParams); }
+    void serialize(void* buffer) const noexcept override {
+        std::memcpy(buffer, &params_, sizeof(GWCParams));
+    }
+    void destroy() noexcept override { delete this; }
+    void setPluginNamespace(char const* pluginNamespace) noexcept override {
+        namespace_ = pluginNamespace ? pluginNamespace : "";
+    }
+    char const* getPluginNamespace() const noexcept override { return namespace_.c_str(); }
+
+    void attachToContext(cudnnContext*, cublasContext*, nvinfer1::IGpuAllocator*) noexcept override {}
+    void detachFromContext() noexcept override {}
+
+private:
+    GWCParams params_;
+    std::string namespace_;
+};
+
+class FFSGWCVolumePluginCreator final : public nvinfer1::IPluginCreator {
+public:
+    FFSGWCVolumePluginCreator() {
+        fields_.emplace_back("max_disp", nullptr, nvinfer1::PluginFieldType::kINT32, 1);
+        fields_.emplace_back("cv_group", nullptr, nvinfer1::PluginFieldType::kINT32, 1);
+        fields_.emplace_back("normalize", nullptr, nvinfer1::PluginFieldType::kINT32, 1);
+        field_collection_.nbFields = static_cast<int32_t>(fields_.size());
+        field_collection_.fields = fields_.data();
+    }
+
+    char const* getPluginName() const noexcept override { return kPluginName; }
+    char const* getPluginVersion() const noexcept override { return kPluginVersion; }
+    nvinfer1::PluginFieldCollection const* getFieldNames() noexcept override {
+        return &field_collection_;
+    }
+
+    nvinfer1::IPluginV2* createPlugin(
+        char const*,
+        nvinfer1::PluginFieldCollection const* fc) noexcept override {
+        GWCParams params;
+        if (fc) {
+            for (int32_t i = 0; i < fc->nbFields; ++i) {
+                auto const& field = fc->fields[i];
+                if (!std::strcmp(field.name, "max_disp")) {
+                    params.max_disp = fieldToInt(field, params.max_disp);
+                } else if (!std::strcmp(field.name, "cv_group")) {
+                    params.cv_group = fieldToInt(field, params.cv_group);
+                } else if (!std::strcmp(field.name, "normalize")) {
+                    params.normalize = fieldToInt(field, params.normalize);
+                }
+            }
+        }
+        if (params.max_disp <= 0 || params.cv_group <= 0) return nullptr;
+        auto* plugin = new FFSGWCVolumePlugin(params);
+        plugin->setPluginNamespace(namespace_.c_str());
+        return plugin;
+    }
+
+    nvinfer1::IPluginV2* deserializePlugin(
+        char const*,
+        void const* serialData,
+        size_t serialLength) noexcept override {
+        auto* plugin = new FFSGWCVolumePlugin(serialData, serialLength);
+        plugin->setPluginNamespace(namespace_.c_str());
+        return plugin;
+    }
+
+    void setPluginNamespace(char const* pluginNamespace) noexcept override {
+        namespace_ = pluginNamespace ? pluginNamespace : "";
+    }
+    char const* getPluginNamespace() const noexcept override {
+        return namespace_.c_str();
+    }
+
+private:
+    std::vector<nvinfer1::PluginField> fields_;
+    nvinfer1::PluginFieldCollection field_collection_{};
+    std::string namespace_;
+};
+
+FFSGWCVolumePluginCreator g_creator;
+
+}  // namespace
+
+bool registerFFSGWCPlugin() {
+    static bool result = false;
+    static std::once_flag once;
+    std::call_once(once, []() {
+        auto try_register = [](nvinfer1::IPluginRegistry* registry, const char* tag) -> bool {
+            if (!registry) {
+                std::cerr << "[FFS plugin] " << tag << ": registry unavailable\n";
+                return false;
+            }
+            if (registry->getPluginCreator(kPluginName, kPluginVersion, "")) {
+                std::cerr << "[FFS plugin] " << tag
+                          << ": FFSGWCVolume v1 ns=\"\" found-already\n";
+                return true;
+            }
+            if (registry->registerCreator(g_creator, "")) {
+                std::cerr << "[FFS plugin] " << tag
+                          << ": FFSGWCVolume v1 ns=\"\" newly-registered\n";
+                return true;
+            }
+            std::cerr << "[FFS plugin] " << tag
+                      << ": FFSGWCVolume v1 ns=\"\" register FAILED\n";
+            return false;
+        };
+
+        const bool a = try_register(::getPluginRegistry(), "runtime-registry");
+        const bool b = try_register(
+            nvinfer1::getBuilderPluginRegistry(nvinfer1::EngineCapability::kSTANDARD),
+            "builder-registry");
+        result = a || b;
+    });
+    return result;
+}
+
+// Auto-register the creator with the global TRT plugin registry on shared
+// library load. Lets external tools (trtexec --staticPlugins, polygraphy
+// --plugins) deserialize engines containing FFSGWCVolume without calling
+// registerFFSGWCPlugin() themselves.
+REGISTER_TENSORRT_PLUGIN(FFSGWCVolumePluginCreator);
+
+}  // namespace ffs_depth
+
+extern "C" bool ffs_register_gwc_plugin() {
+    return ffs_depth::registerFFSGWCPlugin();
+}

cv/fast-foundationstereo/demo_data/K.txt

@@ -0,0 +1,2 @@
+754.6680908203125 0.0 489.3794860839844 0.0 754.6680908203125 265.16162109375 0.0 0.0 1.0 
+0.063

cv/fast-foundationstereo/docker/dockerfile

@@ -0,0 +1,43 @@
+FROM nvidia/cuda:12.4.1-cudnn-devel-ubuntu22.04
+
+ENV TZ=US/Pacific
+RUN ln -snf /usr/share/zoneinfo/$TZ /etc/localtime && echo $TZ > /etc/timezone
+
+RUN apt-get update --fix-missing && \
+    apt-get install -y apt-utils build-essential ca-certificates cmake curl ffmpeg git libturbojpeg-dev pkg-config wget libnvinfer10 libnvinfer-plugin10 libnvonnxparsers10 libnvinfer-dispatch10 libnvinfer-bin zstd libx11-xcb1 libxcb-xinerama0 libxcb-icccm4 libxcb-render-util0 libxcb-shape0 libxcb-keysyms1 libxcb-image0 libxkbcommon-x11-0 libxcb-cursor0 libxcb-xkb1 libxcb-render0 libxcb-shm0 libxcb-sync1 libxcb-xfixes0 libxcb-randr0 libxcb-xtest0 libsm6 libxext6 libxkbcommon0 && \
+    rm -rf /var/lib/apt/lists/*
+
+SHELL ["/bin/bash", "--login", "-c"]
+
+RUN cd / && wget --quiet https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh -O /miniconda.sh && \
+    /bin/bash /miniconda.sh -b -p /opt/conda &&\
+    ln -s /opt/conda/etc/profile.d/conda.sh /etc/profile.d/conda.sh &&\
+    echo ". /opt/conda/etc/profile.d/conda.sh" >> ~/.bashrc &&\
+    /bin/bash -c "source ~/.bashrc" && \
+    /opt/conda/bin/conda tos accept --override-channels --channel https://repo.anaconda.com/pkgs/main &&\
+    /opt/conda/bin/conda tos accept --override-channels --channel https://repo.anaconda.com/pkgs/r &&\
+    /opt/conda/bin/conda update -n base -c defaults conda -y &&\
+    /opt/conda/bin/conda create -n my python=3.12
+
+ENV PATH=$PATH:/opt/conda/envs/my/bin
+ENV OPENCV_IO_ENABLE_OPENEXR=1
+
+RUN conda init bash &&\
+    echo "conda activate my" >> ~/.bashrc &&\
+    conda activate my &&\
+    pip install uv &&\
+    uv pip install torch==2.6.0 torchvision==0.21.0 xformers --index-url https://download.pytorch.org/whl/cu124
+
+COPY requirements.txt /tmp/requirements.txt
+RUN conda activate my &&\
+    uv pip install -r /tmp/requirements.txt &&\
+    uv pip install onnxruntime-gpu onnx pycuda cuda-python tensorrt-cu12 tensorrt-lean-cu12 tensorrt-dispatch-cu12 nvidia-modelopt[torch] &&\
+    conda install -y -c anaconda h5py &&\
+    conda install -y -c conda-forge libstdcxx-ng
+
+RUN cd / && wget https://developer.download.nvidia.com/compute/tensorrt/10.11.0/local_installers/nv-tensorrt-local-repo-ubuntu2204-10.11.0-cuda-12.9_1.0-1_amd64.deb &&\
+    apt install -y ./nv-tensorrt-local-repo-ubuntu2204-10.11.0-cuda-12.9_1.0-1_amd64.deb &&\
+    echo 'alias trtexec="/usr/src/tensorrt/bin/trtexec"' >> ~/.bashrc &&\
+    rm -f /nv-tensorrt-local-repo-ubuntu2204-10.11.0-cuda-12.9_1.0-1_amd64.deb
+
+SHELL ["/bin/bash", "-c", "source ~/.bashrc && conda activate my"]

cv/fast-foundationstereo/docker/dockerfile_cpp

@@ -0,0 +1,48 @@
+FROM nvidia/cuda:12.4.1-cudnn-devel-ubuntu22.04
+
+ENV TZ=US/Pacific
+RUN ln -snf /usr/share/zoneinfo/$TZ /etc/localtime && echo $TZ > /etc/timezone
+
+RUN apt-get update --fix-missing && \
+    apt-get install -y apt-utils build-essential ca-certificates cmake curl ffmpeg git libturbojpeg-dev pkg-config wget libnvinfer10 libnvinfer-dev libnvinfer-plugin10 libnvinfer-plugin-dev libnvonnxparsers10 libnvonnxparsers-dev libnvinfer-dispatch10 libnvinfer-bin zstd libyaml-cpp-dev libopencv-dev libx11-xcb1 libxcb-xinerama0 libxcb-icccm4 libxcb-render-util0 libxcb-shape0 libxcb-keysyms1 libxcb-image0 libxkbcommon-x11-0 libxcb-cursor0 libxcb-xkb1 libxcb-render0 libxcb-shm0 libxcb-sync1 libxcb-xfixes0 libxcb-randr0 libxcb-xtest0 libsm6 libxext6 libxkbcommon0 && \
+    rm -rf /var/lib/apt/lists/*
+
+SHELL ["/bin/bash", "--login", "-c"]
+
+RUN cd / && wget --quiet https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh -O /miniconda.sh && \
+    /bin/bash /miniconda.sh -b -p /opt/conda &&\
+    ln -s /opt/conda/etc/profile.d/conda.sh /etc/profile.d/conda.sh &&\
+    echo ". /opt/conda/etc/profile.d/conda.sh" >> ~/.bashrc &&\
+    /bin/bash -c "source ~/.bashrc" && \
+    /opt/conda/bin/conda tos accept --override-channels --channel https://repo.anaconda.com/pkgs/main &&\
+    /opt/conda/bin/conda tos accept --override-channels --channel https://repo.anaconda.com/pkgs/r &&\
+    /opt/conda/bin/conda update -n base -c defaults conda -y &&\
+    /opt/conda/bin/conda create -n my python=3.12
+
+ENV PATH=$PATH:/opt/conda/envs/my/bin
+ENV OPENCV_IO_ENABLE_OPENEXR=1
+
+RUN conda init bash &&\
+    echo "conda activate my" >> ~/.bashrc &&\
+    conda activate my &&\
+    pip install uv &&\
+    uv pip install torch==2.9.1 torchvision==0.24.1 --index-url https://download.pytorch.org/whl/cu128
+
+COPY requirements.txt /tmp/requirements.txt
+# Do not install nvidia-modelopt[torch] here. It has a strong transitive torch
+# requirement and can silently replace the pinned torch/torchvision pair above,
+# causing torchvision ABI mismatches such as:
+#   RuntimeError: operator torchvision::nms does not exist
+# Install modelopt manually only if you need quantization tooling.
+RUN conda activate my &&\
+    uv pip install -r /tmp/requirements.txt &&\
+    uv pip install onnxruntime-gpu onnx onnxscript pycuda cuda-python tensorrt-cu12 tensorrt-lean-cu12 tensorrt-dispatch-cu12 && nvidia-modelopt[torch] &&\
+    conda install -y -c anaconda h5py &&\
+    conda install -y -c conda-forge libstdcxx-ng
+
+RUN cd / && wget https://developer.download.nvidia.com/compute/tensorrt/10.11.0/local_installers/nv-tensorrt-local-repo-ubuntu2204-10.11.0-cuda-12.9_1.0-1_amd64.deb &&\
+    apt install -y ./nv-tensorrt-local-repo-ubuntu2204-10.11.0-cuda-12.9_1.0-1_amd64.deb &&\
+    echo 'alias trtexec="/usr/src/tensorrt/bin/trtexec"' >> ~/.bashrc &&\
+    rm -f /nv-tensorrt-local-repo-ubuntu2204-10.11.0-cuda-12.9_1.0-1_amd64.deb
+
+SHELL ["/bin/bash", "-c", "source ~/.bashrc && conda activate my"]

cv/fast-foundationstereo/docker/run_container.sh

@@ -0,0 +1,4 @@
+docker rm -f ffs || true
+xhost +local:root || true
+DIR=$(pwd)/../
+docker run --gpus all --runtime nvidia --env NVIDIA_DISABLE_REQUIRE=1 -it --network=host --name ffs --cap-add=SYS_PTRACE --security-opt seccomp=unconfined -v $DIR:/workspace --ipc=host -e DISPLAY=${DISPLAY} -v /tmp/.X11-unix:/tmp/.X11-unix -v /tmp:/tmp -v /home:/home -v /mnt:/mnt ffs bash

cv/fast-foundationstereo/model_card.md

@@ -0,0 +1,138 @@
+<p align="center"><b>Model Card - Fast-FoundationStereo</b></p>
+
+#  Overview
+
+## Description:
+The Fast-FoundationStereo model estimates the disparity of each pixel in a rectified binocular stereo pair of images. This is a transformer based foundational model which shows strong generalization running in real-time. This model is for research and evaluation purposes only.
+
+### License/Terms of Use:
+The code is released using the NVIDIA Source Code License. https://github.com/NVlabs/FoundationStereo/blob/master/LICENSE
+
+### Deployment Geography:
+Global
+
+### Use Case:
+Researchers and developers in the field of computer vision, specifically those interested in depth estimation, are expected to use this method for tasks such as three dimensional reconstruction, object detection, object pose estimation, and scene understanding.
+
+### Release Date:
+Github [02/01/2026] via [https://github.com/NVlabs/Fast-FoundationStereo]
+
+## Reference(s):
+[Fast-FoundationStereo: Real-Time Zero-Shot Stereo Matching](https://arxiv.org/abs/2512.11130)
+
+## Model Architecture:
+**Architecture Type:** Transformers and convolutional neural networks (CNNs).
+
+**Network Architecture:** The network contains three parts:  1) EdgeNeXt student module that distills the original FoundationStereo feature extractor.  2) Set of blocks (CNNs and transformers) that performs matching with long-range dependencies.  3) Reduced set of convGRU blocks.
+
+**Number of model parameters:** 14.6M.
+
+
+## Input:
+**Input Type(s):** A pair of two rectified binocular stereo images
+
+**Input Format(s):** Red, Green, Blue (RGB)
+
+**Input Parameters:** The input parameters to this model are rectified stereo images, specifically Two-Dimensional (2D) images from a camera like Zed.  In addition, the baseline is needed to convert disparity to depth.
+
+**Other Properties Related to Input:** Additional input properties:
+- No Alpha Channel or Pre-Processing Needed.  Bit:  24-bit.
+
+## Output:
+**Output Type(s):** Disparity image
+
+**Output Format(s):** 16-bit unsigned integer
+
+**Output Parameters:** The output parameter of this model is the final 2D disparity map.
+
+**Other Properties Related to Output:** No Alpha Channel or Post-Processing Needed Bit:  16 bit.
+
+Our AI models are designed and/or optimized to run on NVIDIA GPU-accelerated systems. By leveraging NVIDIA’s hardware (e.g. GPU cores) and software frameworks (e.g., CUDA libraries), the model achieves faster training and inference times compared to CPU-only solutions.
+
+## Software Integration :
+**Runtime Engine(s):**
+Not Applicable (N/A)
+
+
+**Supported Hardware Microarchitecture Compatibility:**
+NVIDIA Ampere
+
+
+**[Preferred/Supported] Operating System(s):**
+Linux
+
+
+## Model Version(s):
+v1.0: Initial model version with full capabilities, unpruned and trained.
+
+# Training and Evaluation Datasets:
+
+## Training Dataset:
+
+**Data Modality:**
+Image
+
+**Image Training Data Size:**
+1 Million to 1 Billion Images.
+
+
+**Link:** Internal, proprietary dataset, and Stereo4D dataset
+
+**Data Collection Method by dataset**
+[Hybrid: Synthetic, Automatic/Sensors]
+
+**Labeling Method by dataset**
+[Hybrid: Synthetic, Automatic/Sensors]
+
+**Properties:** The training dataset includes: 1) a large-scale synthetic dataset featuring 1.4 million stereo pairs with large diversity of objects and scenes and high photorealism; 2) real dataset from Stereo4D (external)
+
+## Testing Dataset:
+**Link:** Middlebury dataset
+
+**Data Collection Method by dataset**
+[Automatic/Sensors]
+
+**Labeling Method by dataset**
+[Automatic/Sensors]
+
+
+**Properties:** The dataset encompasses a wide range of scenarios, includes diverse three dimensional assets, captures stereo images under diversely randomized camera parameters, and achieves high fidelity in both rendering and spatial layouts.
+
+## Evaluation Dataset:
+- We evaluated the model using public leaderboards well-known to the stereo community: <br>
+  Middlebury: https://vision.middlebury.edu/stereo/
+
+
+- This is the classic leaderboard for dense stereo matching, developed at Middlebury College. <br>
+  ETH3D: https://www.eth3d.net/
+
+
+- This is another popular leaderboard for stereo, developed at ETH. <br>
+  KITTI: https://www.cvlibs.net/datasets/kitti/eval_stereo.php
+
+**Data Collection Method by dataset**
+[Automatic/Sensors]
+
+** Labeling Method by dataset**
+[Automatic/Sensors]
+
+**Properties:**
+* The Middlebury Stereo dataset consists of high-resolution stereo sequences with complex geometry and pixel-accurate ground-truth disparity data. The ground-truth disparities were acquired using a novel technique that employs structured lighting and infrared paint. <br>
+
+* ETHD is a multi-view stereo benchmark / 3D reconstruction benchmark that covers a variety of indoor and outdoor scenes. Ground truth geometry was obtained using a high-precision laser scanner. A DSLR camera as well as a synchronized multi-camera rig with varying field-of-view was used to capture images. <br>
+* KITTI stereo dataset is a cornerstone of autonomous driving research, developed by the Karlsruhe Institute of Technology (KIT) and the Toyota Technological Institute at Chicago (TTIC). It provides real-world, high-resolution stereo imagery paired with precise ground-truth depth data collected from a moving vehicle in diverse urban environments.
+
+
+
+## Inference:
+**Engine:** Tensor(RT)
+
+**Test Hardware :**
+* Zed Stereo Camera, 3090
+
+
+## Ethical Considerations:
+NVIDIA believes Trustworthy AI is a shared responsibility and we have established policies and practices to enable development for a wide array of AI applications.  When downloaded or used in accordance with our terms of service, developers should work with their internal model team to ensure this model meets requirements for the relevant industry and use case and addresses unforeseen product misuse.
+
+Please report security vulnerabilities or NVIDIA AI Concerns [here](https://www.nvidia.com/en-us/support/submit-security-vulnerability/).
+

cv/fast-foundationstereo/readme.md

@@ -0,0 +1,216 @@
+# Fast-FoundationStereo: Real-Time Zero-Shot Stereo Matching
+
+This is the official implementation of our paper accepted to CVPR 2026
+
+[[Website]](https://nvlabs.github.io/Fast-FoundationStereo/) [[Paper]](https://arxiv.org/abs/2512.11130) [[Video]](https://www.youtube.com/watch?v=2BUYZojCzXE)
+
+Authors: [Bowen Wen](https://wenbowen123.github.io/), [Shaurya Dewan](https://www.linkedin.com/in/shaurya-dewan-1b07231a2), [Stan Birchfield](https://research.nvidia.com/person/stan-birchfield)
+
+
+# Abstract
+Stereo foundation models achieve strong zero-shot generalization but remain computationally prohibitive for real-time applications. Efficient stereo architectures, on the other hand, sacrifice robustness for speed and require costly per-domain fine-tuning. To bridge this gap, we present Fast-FoundationStereo, a family of architectures that achieve, for the first time, strong zero-shot generalization at real-time frame rate. We employ a divide-and-conquer acceleration strategy with three components: (1) knowledge distillation to compress the hybrid backbone into a single efficient student; (2) blockwise neural architecture search for automatically discovering optimal cost filtering designs under latency budgets, reducing search complexity exponentially; and (3) structured pruning for eliminating redundancy in the iterative refinement module. Furthermore, we introduce an automatic pseudo-labeling pipeline used to curate 1.4M in-the-wild stereo pairs to supplement synthetic training data and facilitate knowledge distillation. The resulting model can run over 10× faster than FoundationStereo while closely matching its zero-shot accuracy, thus establishing a new state-of-the-art among real-time methods.
+
+ [NOTE] This model is designed for real-time applications. For offline computation for the best accuracy, please checkout our earlier work [FoundationStereo](https://github.com/NVlabs/FoundationStereo).
+
+
+
+<p align="center">
+  <img src="assets/intro.jpg" width="100%"/>
+</p>
+
+<td align="center">
+  <img src="assets/intro_c.webp" width="60%"/>
+</td>
+<td align="center">
+  <img src="assets/bp2_vs_runtime.jpg" width="60%"/>
+</td>
+
+
+# Environment setup
+- Option 1: Docker
+```bash
+docker build --network host -t ffs -f docker/dockerfile .
+bash docker/run_container.sh
+```
+
+- Option 2: pip
+```bash
+conda create -n ffs python=3.12 && conda activate ffs
+pip install torch==2.6.0 torchvision==0.21.0 xformers --index-url https://download.pytorch.org/whl/cu124
+pip install -r requirements.txt
+```
+
+
+# Weights and Trade-off
+download from [here](https://drive.google.com/drive/folders/1HuTt7UIp7gQsMiDvJwVuWmKpvFzIIMap?usp=drive_link) and put under the folder `weights/` (e.g. `./weights/23-36-37`). Below table compares the differences among some representative models of varying sizes from our trained family. They are sorted from slowest to fastest, with accuracy descending, where runtime is profiled on GPU 3090, image size 640x480.
+
+To trade-off speed and accuracy, there are two options:
+1) Try with different checkpoints.
+2) Tune the config flags (see explanations in the "Run demo" section below).
+
+| Checkpoint     | valid_iters | Runtime-Pytorch (ms) | Runtime-TRT (ms) | Peak Memory (MB) |
+|---------------|-------------|-------------|-----------------|-----------------|
+| `23-36-37`    | 8           | 49.4        | 23.4            | 653             |
+| `23-36-37`    | 4           | 41.1        | 18.4            | 653             |
+| `20-26-39`    | 8           | 43.6        | 19.4            | 651             |
+| `20-26-39`    | 4           | 37.5        | 16.4            | 651             |
+| `20-30-48`    | 8           | 38.4        | 16.6            | 646             |
+| `20-30-48`    | 4           | 29.3        | 14.0            | 646             |
+
+# Run demo
+```
+python scripts/run_demo.py --model_dir weights/23-36-37/model_best_bp2_serialize.pth --left_file demo_data/left.png --right_file demo_data/right.png --intrinsic_file demo_data/K.txt --out_dir output/ --remove_invisible 0 --denoise_cloud 1  --scale 1 --get_pc 1 --valid_iters 8 --max_disp 192 --zfar 100
+```
+| Flag                        | Meaning                                                                |
+|-----------------------------|------------------------------------------------------------------------|
+| `--model_dir`               | Path to the trained weights/model file                                 |
+| `--left_file`               | Path to the left image file                                            |
+| `--right_file`              | Path to the right image file                                           |
+| `--intrinsic_file`          | Path to the camera intrinsic matrix and baseline file                  |
+| `--out_dir`                 | Output directory for saving results                                    |
+| `--remove_invisible`        | Whether to ignore non-overlapping region's depth (0: no, 1: yes)      |
+| `--denoise_cloud`           | Whether to apply denoising to the point cloud (0: no, 1: yes)          |
+| `--scale`                   | Image scaling factor                                                   |
+| `--get_pc`                  | Obtain point cloud output (0: no, 1: yes)                              |
+| `--valid_iters`             | Number of refinement updates during forward pass                       |
+| `--max_disp`                | Maximum disparity for volume encoding, 192 should be enough, unless you need to sense very near objects (e.g. <0.1m). Increasing it runs slower and uses more memory. |
+| `--zfar`                    | Maximum depth to include in point cloud                                |
+
+Refer to `scripts/run_demo.py` for comprehensive list of flags.
+
+**Tips:**
+- The input left and right images should be rectified and undistorted, which means there should not be fisheye kind of lens distortion and the epipolar lines are horizontal between the left/right images. If you obtain images from stereo cameras such as Zed, they usually have handled this for you.
+- Do not swap left and right image. The left image should really be obtained from the left-side camera (objects will appear righter in the image).
+- We recommend to use PNG files with no lossy compression
+- Our method works best on stereo RGB images. However, we have also tested it on monochrome or IR stereo images (e.g. from RealSense D4XX series) and it works well too.
+- To get point cloud for your own data, you need to specify the intrinsics. In the intrinsic file in args, 1st line is the flattened 1x9 intrinsic matrix, 2nd line is the baseline (distance) between the left and right camera, unit in meters.
+- The model performs better for image width size <1000. You can run with smaller scale, e.g. `--scale 0.5` to downsize input image, then upsize the output depth to your need with nearest neighbor interpolation.
+- For faster inference, you can reduce the input image resolution by e.g. `--scale 0.5`, and reduce refine iterations by e.g. `--valid_iters 4`.
+- Note that the 1st time running is slower due to compilation, use a while loop after warm up for live running.
+
+Expect to see results like below:
+- Disparity/Depth:
+  <p align="center">
+    <img src="assets/disp_vis.png" alt="Disparity Visualization" width="100%">
+  </p>
+
+- Point cloud:
+  <p align="center">
+    <img src="assets/pcl_vis.png" alt="Point Cloud Visualization" width="100%">
+  </p>
+
+
+# ONNX/TRT
+For TRT, we recommend first setup env in docker.
+
+## Single ONNX
+
+Export the full model as a single ONNX file. This replaces the Triton GWC kernel with ONNX-compatible ops so no intermediate engine split is needed.
+
+```bash
+python scripts/make_single_onnx.py --model_dir weights/23-36-37/model_best_bp2_serialize.pth --save_path output/ --height 480 --width 640 --valid_iters 8 --max_disp 192
+```
+
+| Flag              | Meaning                                                                  |
+|-------------------|--------------------------------------------------------------------------|
+| `--model_dir`     | Path to the trained weights/model file                                   |
+| `--save_path`     | Directory to save the ONNX model and config                              |
+| `--height`        | Input image height, must be divisible by 32. Reduce for faster speed.    |
+| `--width`         | Input image width, must be divisible by 32. Reduce for faster speed.     |
+| `--valid_iters`   | Number of refinement updates during forward pass, reduce for faster speed, but may drop quality |
+| `--max_disp`      | Maximum disparity for volume encoding, 192 should be enough, unless you need to sense very near objects (e.g. <0.1m). Increasing it runs slower and uses more memory. |
+| `--onnx_name`     | Base name for the saved ONNX file (default: `fast_foundationstereo`)     |
+
+Then convert to a single TRT engine:
+```bash
+trtexec --onnx=output/fast_foundationstereo.onnx --saveEngine=output/fast_foundationstereo.engine --fp16
+```
+
+To run inference with the single ONNX or TRT engine:
+```bash
+python scripts/run_demo_single_trt.py --model_dir output/ --left_file demo_data/left.png --right_file demo_data/right.png --intrinsic_file demo_data/K.txt --out_dir output_demo/ --get_pc 1 --remove_invisible 0 --denoise_cloud 1 --zfar 100
+```
+
+The script auto-detects `.engine` or `.onnx` files in `--model_dir`. To use a specific file, pass `--model_file` directly.
+
+| Flag                  | Meaning                                                                  |
+|-----------------------|--------------------------------------------------------------------------|
+| `--model_dir`         | Directory containing the .onnx/.engine file and its .yaml config         |
+| `--model_file`        | Explicit path to .onnx or .engine file (overrides auto-search)           |
+| `--left_file`         | Path to the left image file                                              |
+| `--right_file`        | Path to the right image file                                             |
+| `--intrinsic_file`    | Path to the camera intrinsic matrix and baseline file                    |
+| `--out_dir`           | Output directory for saving results                                      |
+| `--remove_invisible`  | Whether to ignore non-overlapping region's depth (0: no, 1: yes)         |
+| `--denoise_cloud`     | Whether to apply denoising to the point cloud (0: no, 1: yes)            |
+| `--get_pc`            | Obtain point cloud output (0: no, 1: yes)                                |
+| `--zfar`              | Maximum depth (m) to include in point cloud                              |
+
+**Note:** The single ONNX model expects **pre-normalized** float32 inputs (ImageNet normalization stripped). The inference script handles this automatically. If integrating into your own pipeline, apply normalization beforehand:
+```
+normalized = (pixel - mean) / std
+mean = [123.675, 116.28, 103.53]
+std  = [ 58.395, 57.12, 57.375]
+```
+
+## Two-stage ONNX
+
+The original export splits the model into two ONNX files around the Triton GWC kernel, which runs as an intermediate step between the two TRT engines.
+
+```bash
+python scripts/make_onnx.py --model_dir weights/23-36-37/model_best_bp2_serialize.pth --save_path output/ --height 448 --width 640 --valid_iters 8 --max_disp 192
+```
+
+| Flag              | Meaning                                                                  |
+|-------------------|--------------------------------------------------------------------------|
+| `--model_dir`     | Path to the trained weights/model file                                   |
+| `--save_path`     | Directory to save ONNX outputs and zip file                             |
+| `--height`        | Input image height, better to be divisible by 32. Reduce image size can increase speed.                                          |
+| `--width`         | Input image width,  better to be divisible by 32. Reduce image size can increase speed.                                     |
+| `--valid_iters`   | Number of updates during forward pass, reduce it for faster speed, but may drop quality                                    |
+| `--max_disp`      | Maximum disparity for volume encoding, 192 should be enough, unless you need to sense very near objects (e.g. <0.1m). Increasing it runs slower and uses more memory.                                    |
+
+Refer to `scripts/make_onnx.py` for a comprehensive list of available flags.
+
+Then convert from ONNX to TRT:
+```bash
+trtexec --onnx=output/feature_runner.onnx --saveEngine=output/feature_runner.engine --fp16  --useCudaGraph
+trtexec --onnx=output/post_runner.onnx --saveEngine=output/post_runner.engine --fp16  --useCudaGraph
+```
+
+To use the two-stage TRT for inference:
+```bash
+python scripts/run_demo_tensorrt.py --onnx_dir output/ --left_file demo_data/left.png --right_file demo_data/right.png --intrinsic_file demo_data/K.txt --out_dir output/ --remove_invisible 0 --denoise_cloud 1  --get_pc 1 --zfar 100
+```
+
+# Internet-Scale Pseudo-Labeling
+Real-world data offers greater diversity and realism than synthetic data. However, obtaining real stereo images with ground-truth metric depth annotation is notoriously difficult. To address this challenge, we propose an automatic data curation pipeline to generate pseudo-labels on internet-scale stereo images from [Stereo4D](https://stereo4d.github.io/) dataset. **Top:** Pseudo-labeling pipeline on in-the-wild internet stereo data. **Bottom:** Visualization of our generated pseudo-labels.
+
+<p align="center">
+  <img src="assets/stereo4d.jpg" width="50%">
+</p>
+
+Below are visualizations of the intermediate results in our pseudo-labeling process. In the rightmost column, green checkmark or red cross denotes whether samples are kept for training or not, based on the percentage of positive pixels in the consistency mask. Our data curation process can automatically discover failures on noisy internet data such as images containing subtitle (bottom), mosaic (2nd last row) and overly challenging samples that are unsuitable for training (top). The final pseudo-labels can also correct erroneous predictions from FoundationStereo on sky regions (5th row).
+
+<p align="center">
+  <img src="assets/stereo4d_labeling.jpg" width="100%">
+</p>
+
+
+The dataset is available at HuggingFace: https://huggingface.co/datasets/nvidia/ffs_stereo4d
+
+# Citation
+```bibtex
+@article{wen2026fastfoundationstereo,
+  title={{Fast-FoundationStereo}: Real-Time Zero-Shot Stereo Matching},
+  author={Bowen Wen and Shaurya Dewan and Stan Birchfield},
+  journal={CVPR},
+  year={2026}
+}
+```
+
+# Contact
+Please contact [Bowen Wen](https://wenbowen123.github.io/) (bowenw@nvidia.com) for questions and commercial inquiries.
+
+# Acknowledgement
+We would like to thank Xutong Ren, Karsten Patzwaldt, Yonggan Fu, Saurav Muralidharan, Han Cai, Pavlo Molchanov, Yu Wang, Varun Praveen, Joseph Aribido and Jun Gao for their insightful early discussions for this project. We would also like to thank NVIDIA Isaac and TAO teams for their engineering support and valuable discussions. Thanks to the authors of [FoundationStereo](https://github.com/NVlabs/FoundationStereo), [Selective-IGEV](https://github.com/Windsrain/Selective-Stereo), [Stereo4D](https://github.com/Stereo4d/stereo4d-code) and [RAFT-Stereo](https://github.com/princeton-vl/RAFT-Stereo) for their code release. Finally, thanks to CVPR reviewers and AC for their appreciation of this work and constructive feedback.

cv/fast-foundationstereo/requirements.txt

@@ -0,0 +1,19 @@
+# === Model & inference ===
+timm
+einops
+omegaconf
+scipy
+numpy
+scikit-image
+opencv-contrib-python
+imageio
+pyyaml
+open3d
+
+# === TensorRT (optional — uncomment if needed) ===
+# onnx
+# onnxruntime-gpu
+# tensorrt-cu12
+# tensorrt-lean-cu12
+# tensorrt-dispatch-cu12
+# nvidia-modelopt[torch]

cv/fast-foundationstereo/scripts/build_plugin_trt.py

@@ -0,0 +1,150 @@
+#!/usr/bin/env python3
+"""Build a TensorRT engine from the FFSGWCVolume plugin ONNX.
+
+This mirrors cpp/app/build_single_engine.cpp in Python. The ONNX parser still
+needs the custom FFSGWCVolume plugin creator registered before parsing, so the
+shared plugin library is auto-detected from cpp/build or can be passed with
+--plugin_lib.
+"""
+
+import argparse
+import ctypes
+import os
+from pathlib import Path
+
+
+PLUGIN_NAME = "FFSGWCVolume"
+PLUGIN_VERSION = "1"
+_LOADED_PLUGIN_LIBS = []
+
+
+def find_default_plugin_library() -> Path | None:
+    repo_dir = Path(__file__).resolve().parents[1]
+    candidates = [
+        repo_dir / "cpp" / "build" / "libffs_gwc_plugin.so",
+        repo_dir / "cpp" / "build" / "lib" / "libffs_gwc_plugin.so",
+        repo_dir / "cpp" / "build" / "Release" / "libffs_gwc_plugin.so",
+    ]
+    for path in candidates:
+        if path.exists():
+            return path
+    return None
+
+
+def load_plugin_library(path: str) -> None:
+    """Load an optional shared library that registers FFSGWCVolume."""
+    lib = ctypes.CDLL(path, mode=ctypes.RTLD_GLOBAL)
+    _LOADED_PLUGIN_LIBS.append(lib)
+
+    # The current C++ code registers via ffs_depth::registerFFSGWCPlugin().
+    # A loadable Python plugin library should expose an extern "C" wrapper with
+    # one of these names so ctypes can call it without C++ name mangling.
+    for symbol in ("ffs_register_gwc_plugin", "registerFFSGWCPlugin"):
+        try:
+            fn = getattr(lib, symbol)
+        except AttributeError:
+            continue
+        fn.restype = ctypes.c_bool
+        if not fn():
+            raise RuntimeError(f"{symbol}() returned false for {path}")
+        return
+
+    # Some TensorRT plugin libraries register creators during library load. That
+    # is not true for this repo's current static C++ helper, but allow it here.
+
+
+def find_plugin_creator(trt) -> bool:
+    registry = trt.get_plugin_registry()
+    creator = registry.get_plugin_creator(PLUGIN_NAME, PLUGIN_VERSION, "")
+    return creator is not None
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(
+        description="Build a TensorRT engine from an ONNX graph containing FFSGWCVolume."
+    )
+    parser.add_argument("plugin_onnx", type=Path, help="Path to plugin ONNX file")
+    parser.add_argument("output_engine", type=Path, help="Path to write TensorRT engine")
+    parser.add_argument(
+        "--plugin_lib",
+        type=Path,
+        default=None,
+        help=(
+            "Shared library that registers FFSGWCVolume. Defaults to "
+            "cpp/build/libffs_gwc_plugin.so when present. Pure Python cannot "
+            "provide this repo's CUDA plugin implementation."
+        ),
+    )
+    parser.add_argument(
+        "--fp32",
+        action="store_true",
+        help="Disable FP16 builder flag. Default allows FP16 when supported.",
+    )
+    parser.add_argument(
+        "--workspace-mb",
+        type=int,
+        default=4096,
+        help="TensorRT workspace memory limit in MiB.",
+    )
+    return parser.parse_args()
+
+
+def main() -> int:
+    args = parse_args()
+    if not args.plugin_onnx.exists():
+        raise FileNotFoundError(f"ONNX file does not exist: {args.plugin_onnx}")
+    args.output_engine.parent.mkdir(parents=True, exist_ok=True)
+
+    import tensorrt as trt
+
+    logger = trt.Logger(trt.Logger.INFO)
+    trt.init_libnvinfer_plugins(logger, "")
+
+    plugin_lib = args.plugin_lib or find_default_plugin_library()
+    if plugin_lib is not None:
+        if not plugin_lib.exists():
+            raise FileNotFoundError(f"Plugin library does not exist: {plugin_lib}")
+        load_plugin_library(str(plugin_lib))
+
+    if not find_plugin_creator(trt):
+        raise RuntimeError(
+            f"{PLUGIN_NAME} plugin creator is not registered. "
+            "Build/load a shared library for cpp/src/gwc_volume_plugin.cpp and pass "
+            "--plugin_lib, or use cpp/build/ffs_build_single_engine."
+        )
+
+    explicit_batch = 1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)
+    builder = trt.Builder(logger)
+    network = builder.create_network(explicit_batch)
+    parser = trt.OnnxParser(network, logger)
+
+    parsed = False
+    if hasattr(parser, "parse_from_file"):
+        parsed = parser.parse_from_file(str(args.plugin_onnx))
+    else:
+        parsed = parser.parse(args.plugin_onnx.read_bytes())
+    if not parsed:
+        for i in range(parser.num_errors):
+            print(parser.get_error(i))
+        raise RuntimeError(f"failed to parse ONNX: {args.plugin_onnx}")
+
+    config = builder.create_builder_config()
+    config.set_memory_pool_limit(
+        trt.MemoryPoolType.WORKSPACE, int(args.workspace_mb) * 1024 * 1024
+    )
+    if not args.fp32 and builder.platform_has_fast_fp16:
+        config.set_flag(trt.BuilderFlag.FP16)
+
+    serialized = builder.build_serialized_network(network, config)
+    if serialized is None:
+        raise RuntimeError("build_serialized_network failed")
+
+    args.output_engine.write_bytes(bytes(serialized))
+    precision = "FP32" if args.fp32 else "FP16 allowed"
+    print(f"Built engine: {args.output_engine}")
+    print(f"Precision: {precision}")
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

cv/fast-foundationstereo/scripts/make_onnx.py

@@ -0,0 +1,88 @@
+import warnings, argparse, logging, os, sys,zipfile
+os.environ['TORCH_COMPILE_DISABLE'] = '1'
+os.environ['TORCHDYNAMO_DISABLE'] = '1'
+code_dir = os.path.dirname(os.path.abspath(__file__))
+sys.path.append(f'{code_dir}/../')
+import omegaconf, yaml, torch,pdb
+from omegaconf import OmegaConf
+from core.foundation_stereo import FastFoundationStereo, TrtFeatureRunner, TrtPostRunner, build_gwc_volume_triton
+import Utils as U
+
+
+class FoundationStereoOnnx(FastFoundationStereo):
+    def __init__(self, args):
+        super().__init__(args)
+
+    @torch.no_grad()
+    def forward(self, left, right):
+        """ Removes extra outputs and hyper-parameters """
+        with torch.amp.autocast('cuda', enabled=True, dtype=U.AMP_DTYPE):
+            disp = FastFoundationStereo.forward(self, left, right, iters=self.args.valid_iters, test_mode=True, optimize_build_volume=False)
+        return disp
+
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    code_dir = os.path.dirname(os.path.abspath(__file__))
+    parser.add_argument('--model_dir', type=str, default=f'{code_dir}/../weights/model_best_bp2_serialize.pth')
+    parser.add_argument('--save_path', type=str, default=f'/home/bowen/debug/', help='Path to save results.')
+    parser.add_argument('--height', type=int, default=448)
+    parser.add_argument('--width', type=int, default=640)
+    parser.add_argument('--valid_iters', type=int, default=8, help='number of flow-field updates during forward pass')
+    parser.add_argument('--corr_levels', type=int, default=2, help="number of levels in the correlation pyramid")
+    parser.add_argument('--mixed_precision', default=True, action='store_true', help='use mixed precision')
+    parser.add_argument('--corr_radius', type=int, default=4, help="width of the correlation pyramid")
+    parser.add_argument('--n_downsample', type=int, default=2, help="resolution of the disparity field (1/2^K)")
+    parser.add_argument('--n_gru_layers', type=int, default=1, help="number of hidden GRU levels")
+    parser.add_argument('--max_disp', type=int, default=192, help="max disp of geometry encoding volume")
+    parser.add_argument('--low_memory', type=int, default=1, help='reduce memory usage')
+    args = parser.parse_args()
+    os.makedirs(os.path.dirname(args.save_path), exist_ok=True)
+
+    torch.autograd.set_grad_enabled(False)
+
+    model = torch.load(args.model_dir, map_location='cpu', weights_only=False)
+    model.args.max_disp = args.max_disp
+    model.args.valid_iters = args.valid_iters
+    model.cuda().eval()
+
+    feature_runner = TrtFeatureRunner(model)
+    post_runner = TrtPostRunner(model)
+
+    feature_runner.cuda().eval()
+    post_runner.cuda().eval()
+    assert args.height % 32 == 0 and args.width % 32 == 0, "height and width must be divisible by 32"
+    left_img = torch.randn(1, 3, args.height, args.width).cuda().float()*255
+    right_img = torch.randn(1, 3, args.height, args.width).cuda().float()*255
+
+    torch.onnx.export(
+        feature_runner,
+        (left_img, right_img),
+        args.save_path+'/feature_runner.onnx',
+        opset_version=17,
+        input_names = ['left', 'right'],
+        output_names = ['features_left_04', 'features_left_08', 'features_left_16', 'features_left_32', 'features_right_04', 'stem_2x'],
+        do_constant_folding=True,
+        dynamo=False,
+    )
+
+    features_left_04, features_left_08, features_left_16, features_left_32, features_right_04, stem_2x = feature_runner(left_img, right_img)
+    gwc_volume = build_gwc_volume_triton(features_left_04.half(), features_right_04.half(), args.max_disp//4, model.cv_group)
+    disp = post_runner(features_left_04.float(), features_left_08.float(), features_left_16.float(), features_left_32.float(), features_right_04.float(), stem_2x.float(), gwc_volume.float())
+
+    torch.onnx.export(
+        post_runner,
+        (features_left_04, features_left_08, features_left_16, features_left_32, features_right_04, stem_2x, gwc_volume),
+        args.save_path+'/post_runner.onnx',
+        opset_version=17,
+        input_names = ['features_left_04', 'features_left_08', 'features_left_16', 'features_left_32', 'features_right_04', 'stem_2x', 'gwc_volume'],
+        output_names = ['disp'],
+        do_constant_folding=True,
+        dynamo=False,
+    )
+
+    with open(f'{args.save_path}/onnx.yaml', 'w') as f:
+      cfg = OmegaConf.to_container(model.args)
+      cfg['image_size'] = [args.height, args.width]
+      yaml.safe_dump(cfg, f)

cv/fast-foundationstereo/scripts/make_plugin_onnx.py

@@ -0,0 +1,160 @@
+import argparse
+import logging
+import os
+import sys
+
+os.environ['TORCH_COMPILE_DISABLE'] = '1'
+os.environ['TORCHDYNAMO_DISABLE'] = '1'
+
+code_dir = os.path.dirname(os.path.abspath(__file__))
+sys.path.append(f'{code_dir}/../')
+
+
+def build_parser():
+    parser = argparse.ArgumentParser(
+        description='Export Fast-FoundationStereo as one ONNX with an FFSGWCVolume TensorRT plugin node')
+    parser.add_argument('--model_dir', type=str,
+                        default=f'{code_dir}/../weights/23-36-37/model_best_bp2_serialize.pth')
+    parser.add_argument('--save_path', type=str, default=f'{code_dir}/../output_plugin_onnx')
+    parser.add_argument('--height', type=int, default=608)
+    parser.add_argument('--width', type=int, default=960)
+    parser.add_argument('--valid_iters', type=int, default=8)
+    parser.add_argument('--max_disp', type=int, default=192)
+    parser.add_argument('--onnx_name', type=str, default='fast_foundationstereo_plugin.onnx')
+    return parser
+
+
+if any(arg in ('-h', '--help') for arg in sys.argv[1:]):
+    build_parser().print_help()
+    sys.exit(0)
+
+import torch
+import torch.nn as nn
+import yaml
+from omegaconf import OmegaConf
+from torch.onnx import symbolic_helper
+
+from core.foundation_stereo import TrtFeatureRunner, TrtPostRunner
+
+
+class FFSGWCVolumeOp(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, features_left_04, features_right_04, max_disp, cv_group, normalize):
+        # ONNX export only needs a tensor with the correct static shape here.
+        # symbolic() emits the TensorRT plugin node that computes the real volume.
+        batch, _, height, width = features_left_04.shape
+        return features_left_04.new_zeros(
+            (batch, int(cv_group), int(max_disp), height, width)
+        )
+
+    @staticmethod
+    def symbolic(g, features_left_04, features_right_04, max_disp, cv_group, normalize):
+        def as_int(value):
+            if isinstance(value, int):
+                return value
+            return symbolic_helper._parse_arg(value, 'i')
+
+        max_disp = as_int(max_disp)
+        cv_group = as_int(cv_group)
+        normalize = as_int(normalize)
+        out = g.op(
+            'FFSGWCVolume',
+            features_left_04,
+            features_right_04,
+            max_disp_i=int(max_disp),
+            cv_group_i=int(cv_group),
+            normalize_i=int(normalize),
+        )
+        sizes = features_left_04.type().sizes()
+        if sizes is not None and len(sizes) == 4:
+            out.setType(features_left_04.type().with_sizes(
+                [sizes[0], int(cv_group), int(max_disp), sizes[2], sizes[3]]))
+        return out
+
+
+class FastFoundationStereoPluginOnnx(nn.Module):
+    def __init__(self, model, max_disp_levels, cv_group, normalize):
+        super().__init__()
+        self.feature_runner = TrtFeatureRunner(model)
+        self.post_runner = TrtPostRunner(model)
+        self.max_disp_levels = int(max_disp_levels)
+        self.cv_group = int(cv_group)
+        self.normalize = int(bool(normalize))
+
+    @torch.no_grad()
+    def forward(self, left, right):
+        features_left_04, features_left_08, features_left_16, features_left_32, features_right_04, stem_2x = (
+            self.feature_runner(left, right)
+        )
+        gwc_volume = FFSGWCVolumeOp.apply(
+            features_left_04,
+            features_right_04,
+            self.max_disp_levels,
+            self.cv_group,
+            self.normalize,
+        )
+        disp = self.post_runner(
+            features_left_04.float(),
+            features_left_08.float(),
+            features_left_16.float(),
+            features_left_32.float(),
+            features_right_04.float(),
+            stem_2x.float(),
+            gwc_volume.float(),
+        )
+        return disp
+
+
+if __name__ == '__main__':
+    args = build_parser().parse_args()
+
+    logging.basicConfig(level=logging.INFO, format='%(asctime)s %(levelname)s %(message)s')
+    assert args.height % 32 == 0 and args.width % 32 == 0, 'height and width must be divisible by 32'
+    os.makedirs(args.save_path, exist_ok=True)
+    torch.autograd.set_grad_enabled(False)
+
+    logging.info('Loading model: %s', args.model_dir)
+    model = torch.load(args.model_dir, map_location='cpu', weights_only=False)
+    model.args.max_disp = args.max_disp
+    model.args.valid_iters = args.valid_iters
+    model.cuda().eval()
+
+    cv_group = int(getattr(model, 'cv_group', getattr(model.args, 'cv_group', 8)))
+    normalize = bool(getattr(model.args, 'normalize', True))
+    wrapper = FastFoundationStereoPluginOnnx(
+        model,
+        max_disp_levels=args.max_disp // 4,
+        cv_group=cv_group,
+        normalize=normalize,
+    ).cuda().eval()
+
+    left = torch.randn(1, 3, args.height, args.width, device='cuda').float() * 255
+    right = torch.randn(1, 3, args.height, args.width, device='cuda').float() * 255
+
+    onnx_name = args.onnx_name if args.onnx_name.endswith('.onnx') else f'{args.onnx_name}.onnx'
+    onnx_path = os.path.join(args.save_path, onnx_name)
+    logging.info('Exporting plugin ONNX: %s', onnx_path)
+
+    torch.onnx.export(
+        wrapper,
+        (left, right),
+        onnx_path,
+        opset_version=17,
+        input_names=['left', 'right'],
+        output_names=['disp'],
+        do_constant_folding=True,
+        dynamo=False,
+    )
+
+    cfg = OmegaConf.to_container(model.args)
+    cfg['image_size'] = [args.height, args.width]
+    cfg['cv_group'] = cv_group
+    cfg['normalize'] = normalize
+    with open(os.path.join(args.save_path, 'onnx.yaml'), 'w') as f:
+        yaml.safe_dump(cfg, f)
+
+    logging.info('ONNX model: %s', onnx_path)
+    logging.info('Config    : %s', os.path.join(args.save_path, 'onnx.yaml'))
+    logging.info('Build with:')
+    logging.info('  cpp/build/ffs_build_single_engine %s %s',
+                 onnx_path, os.path.join(args.save_path, 'fast_foundationstereo.engine'))

cv/fast-foundationstereo/scripts/make_single_onnx.py

@@ -0,0 +1,222 @@
+# SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES
+# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+Export Fast FoundationStereo as a **single** ONNX model.
+
+Unlike make_onnx.py (which splits into feature_runner + post_runner with a
+Triton GWC kernel in between), this script produces one self-contained ONNX
+that can be converted to a single TensorRT engine via trtexec.
+
+Key design choices:
+  - The GWC and concat cost volumes are built with ONNX-compatible ops
+    (pad + slice + stack).  The upstream pytorch1 variants use
+    Tensor.unfold / torch.flip which the ONNX exporter cannot handle.
+  - ImageNet normalization is STRIPPED from the model so that it can be
+    applied externally (e.g. via Isaac ROS ImageNormalizeNode).  The ONNX
+    model expects **pre-normalized** float inputs:
+        pixel = (pixel_0_255 - mean) / std
+        mean  = [123.675, 116.28, 103.53]   (ImageNet, in 0-255 scale)
+        std   = [ 58.395, 57.12, 57.375]
+  - Inputs:  left_image  (1, 3, H, W)  float32, ImageNet-normalised
+             right_image (1, 3, H, W)  float32, ImageNet-normalised
+  - Output:  disparity   (1, 1, H, W)  float32
+
+Usage:
+  python make_single_onnx.py \\
+      --model_dir ../weights/model_best_bp2_serialize.pth \\
+      --save_path ./output_single_onnx --height 480 --width 640
+
+  # Then build a TensorRT engine:
+  trtexec --onnx=./output_single_onnx/fast_foundationstereo.onnx \\
+      --saveEngine=fast_foundationstereo.engine --fp16
+"""
+
+import argparse
+import logging
+import os
+import sys
+
+os.environ['TORCH_COMPILE_DISABLE'] = '1'
+os.environ['TORCHDYNAMO_DISABLE'] = '1'
+
+code_dir = os.path.dirname(os.path.abspath(__file__))
+sys.path.append(f'{code_dir}/../')
+
+import yaml
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from omegaconf import OmegaConf
+import core.foundation_stereo as _fs_module
+from core.foundation_stereo import FastFoundationStereo
+
+
+# ---------------------------------------------------------------------------
+# ONNX-compatible cost-volume builders
+#
+# The upstream *_optimized_pytorch1 variants use Tensor.unfold + torch.flip
+# which the ONNX tracer cannot export.  These replacements build the
+# disparity-shifted target volume with an explicit loop over disparities
+# using only F.pad, slicing, and torch.stack — all fully ONNX-exportable.
+# ---------------------------------------------------------------------------
+
+def _build_gwc_volume_onnx(refimg_fea, targetimg_fea, maxdisp,
+                           num_groups, normalize=True):
+    dtype = refimg_fea.dtype
+    B, C, H, W = refimg_fea.shape
+    channels_per_group = C // num_groups
+
+    ref_volume = refimg_fea.unsqueeze(2).expand(B, C, maxdisp, H, W)
+
+    shifted = [
+        F.pad(targetimg_fea, (d, 0, 0, 0), 'constant', 0.0)[:, :, :, :W]
+        for d in range(maxdisp)
+    ]
+    target_volume = torch.stack(shifted, dim=2)
+
+    ref_volume = ref_volume.view(B, num_groups, channels_per_group,
+                                 maxdisp, H, W)
+    target_volume = target_volume.view(B, num_groups, channels_per_group,
+                                       maxdisp, H, W)
+
+    if normalize:
+        ref_volume = F.normalize(ref_volume.float(), dim=2).to(dtype)
+        target_volume = F.normalize(target_volume.float(), dim=2).to(dtype)
+
+    return (ref_volume * target_volume).sum(dim=2).contiguous()
+
+
+def _build_concat_volume_onnx(refimg_fea, targetimg_fea, maxdisp):
+    B, C, H, W = refimg_fea.shape
+
+    ref_volume = refimg_fea.unsqueeze(2).expand(B, C, maxdisp, H, W)
+
+    shifted = [
+        F.pad(targetimg_fea, (d, 0, 0, 0), 'constant', 0.0)[:, :, :, :W]
+        for d in range(maxdisp)
+    ]
+    target_volume = torch.stack(shifted, dim=2)
+
+    return torch.cat((ref_volume, target_volume), dim=1).contiguous()
+
+
+# ---------------------------------------------------------------------------
+
+class FastFoundationStereoSingleOnnx(nn.Module):
+    """Thin wrapper that calls the full model with ONNX-compatible settings.
+
+    Before ONNX tracing the caller must monkey-patch:
+      - normalize_image        → identity  (normalization done externally)
+      - build_gwc_volume_*     → _build_gwc_volume_onnx
+      - build_concat_volume_*  → _build_concat_volume_onnx
+    """
+
+    def __init__(self, model: FastFoundationStereo):
+        super().__init__()
+        self.model = model
+
+    @torch.no_grad()
+    def forward(self, left_image, right_image):
+        return self.model.forward(
+            left_image, right_image,
+            iters=self.model.args.valid_iters,
+            test_mode=True,
+            optimize_build_volume='pytorch1',
+        )
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(
+        description='Export Fast FoundationStereo as a single ONNX model')
+    parser.add_argument(
+        '--model_dir', type=str,
+        default=f'{code_dir}/../weights/model_best_bp2_serialize.pth',
+        help='Path to the serialized .pth model')
+    parser.add_argument(
+        '--save_path', type=str,
+        default=f'{code_dir}/output_single_onnx',
+        help='Directory to save the ONNX model and config')
+    parser.add_argument('--height', type=int, default=480)
+    parser.add_argument('--width', type=int, default=640)
+    parser.add_argument('--valid_iters', type=int, default=8,
+                        help='GRU refinement iterations')
+    parser.add_argument('--max_disp', type=int, default=192,
+                        help='Maximum disparity (in pixels at full resolution)')
+    parser.add_argument('--onnx_name', type=str, default='fast_foundationstereo',
+                        help='Base name for the saved ONNX file (without .onnx extension)')
+    args = parser.parse_args()
+
+    logging.basicConfig(
+        level=logging.INFO,
+        format='%(asctime)s %(levelname)s %(message)s')
+
+    assert args.height % 32 == 0 and args.width % 32 == 0, \
+        'height and width must be divisible by 32'
+
+    os.makedirs(args.save_path, exist_ok=True)
+    torch.autograd.set_grad_enabled(False)
+
+    if not os.path.isfile(args.model_dir):
+        raise FileNotFoundError(f'Model file not found: {args.model_dir}')
+
+    logging.info(f'Loading model from {args.model_dir}')
+    model = torch.load(args.model_dir, map_location='cpu', weights_only=False)
+    model.args.max_disp = args.max_disp
+    model.args.valid_iters = args.valid_iters
+    model.args.mixed_precision = False
+    model.cuda().eval()
+
+    wrapper = FastFoundationStereoSingleOnnx(model)
+    wrapper.cuda().eval()
+
+    left_img = torch.randn(1, 3, args.height, args.width, device='cuda')
+    right_img = torch.randn(1, 3, args.height, args.width, device='cuda')
+
+    onnx_name = args.onnx_name if args.onnx_name.endswith('.onnx') else f'{args.onnx_name}.onnx'
+    onnx_path = os.path.join(args.save_path, onnx_name)
+    logging.info(f'Exporting ONNX ({args.height}x{args.width}) → {onnx_path}')
+
+    # Monkey-patch non-ONNX-exportable functions before tracing
+    _fs_module.normalize_image = lambda img: img
+    _fs_module.build_gwc_volume_optimized_pytorch1 = _build_gwc_volume_onnx
+    _fs_module.build_concat_volume_optimized_pytorch1 = _build_concat_volume_onnx
+
+    torch.onnx.export(
+        wrapper,
+        (left_img, right_img),
+        onnx_path,
+        opset_version=17,
+        input_names=['left_image', 'right_image'],
+        output_names=['disparity'],
+        do_constant_folding=True,
+    )
+
+    cfg = OmegaConf.to_container(model.args)
+    cfg['image_size'] = [args.height, args.width]
+    config_name = os.path.splitext(onnx_name)[0] + '.yaml'
+    config_path = os.path.join(args.save_path, config_name)
+    with open(config_path, 'w') as f:
+        yaml.safe_dump(cfg, f)
+
+    logging.info(f'ONNX model  : {onnx_path}')
+    logging.info(f'Config      : {config_path}')
+    logging.info(f'Resolution  : {args.height} x {args.width}')
+    logging.info(
+        f'Build TRT engine:\n'
+        f'  trtexec --onnx={onnx_path} '
+        f'--saveEngine={args.save_path}/{os.path.splitext(onnx_name)[0]}.engine --fp16')

cv/fast-foundationstereo/scripts/profile_memory.py

@@ -0,0 +1,75 @@
+# Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
+#
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+import os,sys
+code_dir = os.path.dirname(os.path.realpath(__file__))
+sys.path.append(f'{code_dir}/../')
+from omegaconf import OmegaConf
+from core.utils.utils import InputPadder
+import argparse, torch, logging, yaml, time
+import numpy as np
+from Utils import AMP_DTYPE, set_logging_format, set_seed
+
+
+if __name__=="__main__":
+  code_dir = os.path.dirname(os.path.realpath(__file__))
+  parser = argparse.ArgumentParser()
+  parser.add_argument('--model_dir', default=f'{code_dir}/../weights/23-36-37/model_best_bp2_serialize.pth', type=str)
+  parser.add_argument('--hiera', default=0, type=int)
+  parser.add_argument('--valid_iters', type=int, default=8, help='number of flow-field updates during forward pass')
+  parser.add_argument('--max_disp', type=int, default=192, help='maximum disparity')
+  parser.add_argument('--warmup', type=int, default=15, help='number of warmup iterations')
+  parser.add_argument('--total', type=int, default=30, help='total number of iterations')
+  args = parser.parse_args()
+
+  set_logging_format()
+  set_seed(0)
+  torch.backends.cudnn.benchmark = True
+  torch.autograd.set_grad_enabled(False)
+
+  with open(f'{os.path.dirname(args.model_dir)}/cfg.yaml', 'r') as ff:
+    cfg:dict = yaml.safe_load(ff)
+  for k in args.__dict__:
+    if args.__dict__[k] is not None:
+      cfg[k] = args.__dict__[k]
+  args = OmegaConf.create(cfg)
+  model = torch.load(args.model_dir, map_location='cpu', weights_only=False)
+  model.args.valid_iters = args.valid_iters
+  model.args.max_disp = args.max_disp
+  model.cuda().eval()
+
+  H, W = 480, 640
+  img0 = torch.randint(0, 256, (1, 3, H, W), dtype=torch.float32).cuda()
+  img1 = torch.randint(0, 256, (1, 3, H, W), dtype=torch.float32).cuda()
+  padder = InputPadder(img0.shape, divis_by=32, force_square=False)
+  img0, img1 = padder.pad(img0, img1)
+
+  logging.info(f"Image size: {H}x{W}, warmup: {args.warmup}, total: {args.total}")
+
+  times = []
+  peak_memories = []
+  with torch.amp.autocast('cuda', enabled=True, dtype=AMP_DTYPE):
+    for i in range(args.total):
+      torch.cuda.reset_peak_memory_stats()
+      torch.cuda.synchronize()
+      t0 = time.perf_counter()
+      disp = model.forward(img0, img1, iters=args.valid_iters, test_mode=True, optimize_build_volume='triton')
+      torch.cuda.synchronize()
+      elapsed = time.perf_counter() - t0
+      peak_mem = torch.cuda.max_memory_allocated() / (1024 ** 2)
+      times.append(elapsed)
+      peak_memories.append(peak_mem)
+      logging.info(f"Iter {i:2d}: {elapsed*1000:.1f} ms, peak mem: {peak_mem:.1f} MB {'(warmup)' if i < args.warmup else ''}")
+
+  measure_times = times[args.warmup:]
+  measure_mems = peak_memories[args.warmup:]
+  avg = np.mean(measure_times) * 1000
+  avg_mem = np.mean(measure_mems)
+  max_mem = np.max(measure_mems)
+  logging.info(f"Vanilla Pytorch speed average (after warmup): {avg:.1f}[ms] over {len(measure_times)} iters")
+  logging.info(f"Peak GPU memory (after warmup): avg {avg_mem:.1f} MB, max {max_mem:.1f} MB")

cv/fast-foundationstereo/scripts/profile_speed.py

@@ -0,0 +1,67 @@
+# Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
+#
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+import os,sys
+code_dir = os.path.dirname(os.path.realpath(__file__))
+sys.path.append(f'{code_dir}/../')
+from omegaconf import OmegaConf
+from core.utils.utils import InputPadder
+import argparse, torch, logging, yaml, time
+import numpy as np
+from Utils import AMP_DTYPE, set_logging_format, set_seed
+
+
+if __name__=="__main__":
+  code_dir = os.path.dirname(os.path.realpath(__file__))
+  parser = argparse.ArgumentParser()
+  parser.add_argument('--model_dir', default=f'{code_dir}/../weights/23-36-37/model_best_bp2_serialize.pth', type=str)
+  parser.add_argument('--hiera', default=0, type=int)
+  parser.add_argument('--valid_iters', type=int, default=8, help='number of flow-field updates during forward pass')
+  parser.add_argument('--max_disp', type=int, default=192, help='maximum disparity')
+  parser.add_argument('--warmup', type=int, default=15, help='number of warmup iterations')
+  parser.add_argument('--total', type=int, default=30, help='total number of iterations')
+  args = parser.parse_args()
+
+  set_logging_format()
+  set_seed(0)
+  torch.backends.cudnn.benchmark = True
+  torch.autograd.set_grad_enabled(False)
+
+  with open(f'{os.path.dirname(args.model_dir)}/cfg.yaml', 'r') as ff:
+    cfg:dict = yaml.safe_load(ff)
+  for k in args.__dict__:
+    if args.__dict__[k] is not None:
+      cfg[k] = args.__dict__[k]
+  args = OmegaConf.create(cfg)
+  model = torch.load(args.model_dir, map_location='cpu', weights_only=False)
+  model.args.valid_iters = args.valid_iters
+  model.args.max_disp = args.max_disp
+  model.cuda().eval()
+
+  H, W = 480, 640
+  img0 = torch.randint(0, 256, (1, 3, H, W), dtype=torch.float32).cuda()
+  img1 = torch.randint(0, 256, (1, 3, H, W), dtype=torch.float32).cuda()
+  padder = InputPadder(img0.shape, divis_by=32, force_square=False)
+  img0, img1 = padder.pad(img0, img1)
+
+  logging.info(f"Image size: {H}x{W}, warmup: {args.warmup}, total: {args.total}")
+
+  times = []
+  with torch.amp.autocast('cuda', enabled=True, dtype=AMP_DTYPE):
+    for i in range(args.total):
+      torch.cuda.synchronize()
+      t0 = time.perf_counter()
+      disp = model.forward(img0, img1, iters=args.valid_iters, test_mode=True, optimize_build_volume='triton')
+      torch.cuda.synchronize()
+      elapsed = time.perf_counter() - t0
+      times.append(elapsed)
+      logging.info(f"Iter {i:2d}: {elapsed*1000:.1f} ms {'(warmup)' if i < args.warmup else ''}")
+
+  measure_times = times[args.warmup:]
+  avg = np.mean(measure_times) * 1000
+  logging.info(f"Vanilla Pytorch speed average (after warmup): {avg:.1f}[ms] over {len(measure_times)} iters")

cv/fast-foundationstereo/scripts/profile_speed_tensorrt.py

@@ -0,0 +1,66 @@
+import os,sys
+code_dir = os.path.dirname(os.path.realpath(__file__))
+sys.path.append(f'{code_dir}/../')
+from omegaconf import OmegaConf
+import argparse, torch, logging, yaml, time
+import numpy as np
+from Utils import set_logging_format, set_seed
+from core.foundation_stereo import TrtRunner
+
+
+def resolve_onnx_cfg_path(onnx_dir: str) -> str:
+  direct = os.path.join(onnx_dir, 'onnx.yaml')
+  if os.path.exists(direct):
+    return direct
+  parent = os.path.join(os.path.dirname(onnx_dir), 'onnx.yaml')
+  if os.path.exists(parent):
+    return parent
+  raise FileNotFoundError(f"onnx.yaml not found in {onnx_dir} or its parent directory")
+
+
+if __name__=="__main__":
+  parser = argparse.ArgumentParser()
+  parser.add_argument('--model_dir', default=None, type=str)
+  parser.add_argument('--hiera', default=None, type=int)
+  parser.add_argument('--valid_iters', type=int, default=None, help='number of flow-field updates during forward pass (default: from onnx.yaml)')
+  parser.add_argument('--max_disp', type=int, default=None, help='maximum disparity (default: from onnx.yaml)')
+  parser.add_argument('--warmup', type=int, default=15, help='number of warmup iterations')
+  parser.add_argument('--total', type=int, default=30, help='total number of iterations')
+  parser.add_argument('--build_volume_backend', default=None, choices=['pytorch1', 'triton'], help='backend for cost-volume build (default: from onnx.yaml)')
+  parser.add_argument('--onnx_dir', default=f'{code_dir}/../output', type=str, help='directory containing TensorRT engines and onnx.yaml')
+  args = parser.parse_args()
+
+  set_logging_format()
+  set_seed(0)
+  torch.backends.cudnn.benchmark = True
+  torch.autograd.set_grad_enabled(False)
+
+  cfg_path = resolve_onnx_cfg_path(args.onnx_dir)
+  with open(cfg_path, 'r') as ff:
+    cfg:dict = yaml.safe_load(ff)
+  for k in args.__dict__:
+    if args.__dict__[k] is not None:
+      cfg[k] = args.__dict__[k]
+  args = OmegaConf.create(cfg)
+
+  model = TrtRunner(args, args.onnx_dir+'/feature_runner.engine', args.onnx_dir+'/post_runner.engine')
+
+  H, W = int(args.image_size[0]), int(args.image_size[1])
+  img0 = torch.randint(0, 256, (1, 3, H, W), dtype=torch.float32).cuda()
+  img1 = torch.randint(0, 256, (1, 3, H, W), dtype=torch.float32).cuda()
+
+  logging.info(f"TensorRT image size: {H}x{W}, warmup: {args.warmup}, total: {args.total}")
+
+  times = []
+  for i in range(args.total):
+    torch.cuda.synchronize()
+    t0 = time.perf_counter()
+    _ = model.forward(img0, img1)
+    torch.cuda.synchronize()
+    elapsed = time.perf_counter() - t0
+    times.append(elapsed)
+    logging.info(f"Iter {i:2d}: {elapsed*1000:.1f} ms {'(warmup)' if i < args.warmup else ''}")
+
+  measure_times = times[args.warmup:]
+  avg = np.mean(measure_times) * 1000
+  logging.info(f"TensorRT speed average (after warmup): {avg:.1f}[ms] over {len(measure_times)} iters")

cv/fast-foundationstereo/scripts/run_demo.py

@@ -0,0 +1,150 @@
+# Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
+#
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+import os,sys
+code_dir = os.path.dirname(os.path.realpath(__file__))
+sys.path.append(f'{code_dir}/../')
+from omegaconf import OmegaConf
+from core.utils.utils import InputPadder
+import argparse, torch, imageio, logging, yaml
+import numpy as np
+from Utils import (
+    AMP_DTYPE, set_logging_format, set_seed, vis_disparity,
+    depth2xyzmap, toOpen3dCloud, o3d,
+)
+import cv2
+
+
+if __name__=="__main__":
+  code_dir = os.path.dirname(os.path.realpath(__file__))
+  parser = argparse.ArgumentParser()
+  parser.add_argument('--model_dir', default=f'{code_dir}/../weights/23-36-37/model_best_bp2_serialize.pth', type=str)
+  parser.add_argument('--left_file', default=f'{code_dir}/../demo_data/left.png', type=str)
+  parser.add_argument('--right_file', default=f'{code_dir}/../demo_data/right.png', type=str)
+  parser.add_argument('--intrinsic_file', default=f'{code_dir}/../demo_data/K.txt', type=str, help='camera intrinsic matrix and baseline file')
+  parser.add_argument('--out_dir', default='/home/bowen/debug/stereo_output', type=str)
+  parser.add_argument('--remove_invisible', default=1, type=int)
+  parser.add_argument('--denoise_cloud', default=0, type=int)
+  parser.add_argument('--denoise_nb_points', type=int, default=30, help='number of points to consider for radius outlier removal')
+  parser.add_argument('--denoise_radius', type=float, default=0.03, help='radius to use for outlier removal')
+  parser.add_argument('--scale', default=1, type=float)
+  parser.add_argument('--hiera', default=0, type=int)
+  parser.add_argument('--get_pc', type=int, default=1, help='save point cloud output')
+  parser.add_argument('--valid_iters', type=int, default=8, help='number of flow-field updates during forward pass')
+  parser.add_argument('--max_disp', type=int, default=192, help='maximum disparity')
+  parser.add_argument('--zfar', type=float, default=100, help="max depth to include in point cloud")
+  args = parser.parse_args()
+
+  set_logging_format()
+  set_seed(0)
+  torch.autograd.set_grad_enabled(False)
+
+  os.system(f'rm -rf {args.out_dir} && mkdir -p {args.out_dir}')
+
+  with open(f'{os.path.dirname(args.model_dir)}/cfg.yaml', 'r') as ff:
+    cfg:dict = yaml.safe_load(ff)
+  for k in args.__dict__:
+    if args.__dict__[k] is not None:
+      cfg[k] = args.__dict__[k]
+  args = OmegaConf.create(cfg)
+  logging.info(f"args:\n{args}")
+  model = torch.load(args.model_dir, map_location='cpu', weights_only=False)
+  model.args.valid_iters = args.valid_iters
+  model.args.max_disp = args.max_disp
+
+  model.cuda().eval()
+
+  scale = args.scale
+
+  img0 = imageio.imread(args.left_file)
+  img1 = imageio.imread(args.right_file)
+  if len(img0.shape)==2:
+    img0 = np.tile(img0[...,None], (1,1,3))
+    img1 = np.tile(img1[...,None], (1,1,3))
+  img0 = img0[...,:3]
+  img1 = img1[...,:3]
+  H,W = img0.shape[:2]
+
+  img0 = cv2.resize(img0, fx=scale, fy=scale, dsize=None)
+  img1 = cv2.resize(img1, dsize=(img0.shape[1], img0.shape[0]))
+  H,W = img0.shape[:2]
+  img0_ori = img0.copy()
+  img1_ori = img1.copy()
+  logging.info(f"img0: {img0.shape}")
+  imageio.imwrite(f'{args.out_dir}/left.png', img0)
+  imageio.imwrite(f'{args.out_dir}/right.png', img1)
+
+  img0 = torch.as_tensor(img0).cuda().float()[None].permute(0,3,1,2)
+  img1 = torch.as_tensor(img1).cuda().float()[None].permute(0,3,1,2)
+  padder = InputPadder(img0.shape, divis_by=32, force_square=False)
+  img0, img1 = padder.pad(img0, img1)
+
+  logging.info(f"Start forward, 1st time run can be slow due to compilation")
+  with torch.amp.autocast('cuda', enabled=True, dtype=AMP_DTYPE):
+    if not args.hiera:
+      disp = model.forward(img0, img1, iters=args.valid_iters, test_mode=True, optimize_build_volume='pytorch1')
+    else:
+      disp = model.run_hierachical(img0, img1, iters=args.valid_iters, test_mode=True, small_ratio=0.5)
+  logging.info("forward done")
+  disp = padder.unpad(disp.float())
+  disp = disp.data.cpu().numpy().reshape(H,W).clip(0, None)
+
+  cmap = None
+  min_val = None
+  max_val = None
+  vis = vis_disparity(disp, min_val=min_val, max_val=max_val, cmap=cmap, color_map=cv2.COLORMAP_TURBO)
+  vis = np.concatenate([img0_ori, img1_ori, vis], axis=1)
+  imageio.imwrite(f'{args.out_dir}/disp_vis.png', vis)
+  s = 1280/vis.shape[1]
+  resized_vis = cv2.resize(vis, (int(vis.shape[1]*s), int(vis.shape[0]*s)))
+  cv2.imshow('disp', resized_vis[:,:,::-1])
+  cv2.waitKey(0)
+
+  if args.remove_invisible:
+    yy,xx = np.meshgrid(np.arange(disp.shape[0]), np.arange(disp.shape[1]), indexing='ij')
+    us_right = xx-disp
+    invalid = us_right<0
+    disp[invalid] = np.inf
+
+  if args.get_pc:
+    with open(args.intrinsic_file, 'r') as f:
+      lines = f.readlines()
+      K = np.array(list(map(float, lines[0].rstrip().split()))).astype(np.float32).reshape(3,3)
+      baseline = float(lines[1])
+    K[:2] *= scale
+    depth = K[0,0]*baseline/disp
+    np.save(f'{args.out_dir}/depth_meter.npy', depth)
+    xyz_map = depth2xyzmap(depth, K)
+    pcd = toOpen3dCloud(xyz_map.reshape(-1,3), img0_ori.reshape(-1,3))
+    keep_mask = (np.asarray(pcd.points)[:,2]>0) & (np.asarray(pcd.points)[:,2]<=args.zfar)
+    keep_ids = np.arange(len(np.asarray(pcd.points)))[keep_mask]
+    pcd = pcd.select_by_index(keep_ids)
+    o3d.io.write_point_cloud(f'{args.out_dir}/cloud.ply', pcd)
+    logging.info(f"PCL saved to {args.out_dir}")
+
+    if args.denoise_cloud:
+      logging.info("[Optional step] denoise point cloud...")
+      pcd = pcd.voxel_down_sample(voxel_size=0.001)
+      cl, ind = pcd.remove_radius_outlier(nb_points=args.denoise_nb_points, radius=args.denoise_radius)
+      inlier_cloud = pcd.select_by_index(ind)
+      o3d.io.write_point_cloud(f'{args.out_dir}/cloud_denoise.ply', inlier_cloud)
+      pcd = inlier_cloud
+
+    logging.info("Visualizing point cloud. Press ESC to exit.")
+    vis = o3d.visualization.Visualizer()
+    vis.create_window()
+    vis.add_geometry(pcd)
+    vis.get_render_option().point_size = 1.0
+    vis.get_render_option().background_color = np.array([0.5, 0.5, 0.5])
+    ctr = vis.get_view_control()
+    ctr.set_front([0, 0, -1])
+    id = np.asarray(pcd.points)[:,2].argmin()
+    ctr.set_lookat(np.asarray(pcd.points)[id])
+    ctr.set_up([0, -1, 0])
+    vis.run()
+    vis.destroy_window()

cv/fast-foundationstereo/scripts/run_demo_plugin_trt.py

@@ -0,0 +1,380 @@
+#!/usr/bin/env python3
+"""Run the FFSGWCVolume-plugin TensorRT engine from Python.
+
+This is the Python equivalent of cpp/src/ffs_single_depth_inference.cpp plus
+cpp/app/main.cpp for the single-engine plugin path:
+
+  - engine inputs: left, right
+  - engine output: disp
+  - input tensors are RGB, CHW, float32, raw 0-255 values
+  - resize is aspect-ratio-preserving with right/bottom replicate padding
+  - disparity is cropped, nearest-neighbor upsampled, and scaled back to the
+    original input-image pixel units
+"""
+
+import argparse
+import logging
+import os
+import sys
+from pathlib import Path
+
+import cv2
+import imageio.v2 as imageio
+import numpy as np
+import torch
+import yaml
+
+code_dir = Path(__file__).resolve().parent
+repo_dir = code_dir.parent
+sys.path.append(str(repo_dir))
+
+from Utils import set_logging_format, set_seed, vis_disparity
+from build_plugin_trt import (
+    PLUGIN_NAME,
+    find_default_plugin_library,
+    find_plugin_creator,
+    load_plugin_library,
+)
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(
+        description="Run Fast-FoundationStereo single TensorRT plugin engine."
+    )
+    parser.add_argument(
+        "--model_dir",
+        type=Path,
+        default=repo_dir / "engine1_plug",
+        help="Directory containing fast_foundationstereo.engine and one YAML config.",
+    )
+    parser.add_argument(
+        "--engine_file",
+        "--model_file",
+        dest="engine_file",
+        type=Path,
+        default=None,
+        help="Explicit TensorRT engine path. Defaults to <model_dir>/fast_foundationstereo.engine.",
+    )
+    parser.add_argument(
+        "--config_file",
+        type=Path,
+        default=None,
+        help="Explicit YAML config path. Defaults to the single YAML in the engine directory.",
+    )
+    parser.add_argument(
+        "--plugin_lib",
+        type=Path,
+        default=None,
+        help="Path to libffs_gwc_plugin.so. Defaults to cpp/build/libffs_gwc_plugin.so.",
+    )
+    parser.add_argument("--left_file", type=Path, default=repo_dir / "demo_data" / "left.png")
+    parser.add_argument("--right_file", type=Path, default=repo_dir / "demo_data" / "right.png")
+    parser.add_argument(
+        "--intrinsic_file",
+        type=Path,
+        default=repo_dir / "demo_data" / "K.txt",
+        help="Text file with 3x3 K on line 1 and baseline in meters on line 2.",
+    )
+    parser.add_argument("--out_dir", type=Path, default=repo_dir / "output_plugin_trt")
+    return parser.parse_args()
+
+
+def resolve_engine_path(model_dir: Path, engine_file: Path | None) -> Path:
+    path = engine_file if engine_file is not None else model_dir / "fast_foundationstereo.engine"
+    if not path.exists():
+        raise FileNotFoundError(f"TensorRT engine does not exist: {path}")
+    return path
+
+
+def resolve_config_path(engine_path: Path, model_dir: Path, config_file: Path | None) -> Path:
+    if config_file is not None:
+        if not config_file.exists():
+            raise FileNotFoundError(f"Config file does not exist: {config_file}")
+        return config_file
+
+    search_dir = model_dir if model_dir.exists() else engine_path.parent
+    yaml_files = sorted(
+        p for p in search_dir.iterdir()
+        if p.is_file() and p.suffix.lower() in (".yaml", ".yml")
+    )
+    if not yaml_files:
+        raise FileNotFoundError(f"No YAML config found in: {search_dir}")
+    if len(yaml_files) > 1:
+        names = " ".join(str(p) for p in yaml_files)
+        raise RuntimeError(
+            f"Expected exactly one YAML config in {search_dir}, found {len(yaml_files)}: {names}"
+        )
+    return yaml_files[0]
+
+
+def load_intrinsics(path: Path) -> tuple[np.ndarray, float]:
+    with path.open("r") as f:
+        lines = f.readlines()
+    if len(lines) < 2:
+        raise RuntimeError("intrinsic file must contain K on line 1 and baseline on line 2")
+    k = np.array(list(map(float, lines[0].strip().split())), dtype=np.float32).reshape(3, 3)
+    baseline = float(lines[1])
+    if not (k[0, 0] > 0.0 and baseline > 0.0):
+        raise RuntimeError("invalid focal length or baseline in intrinsic file")
+    return k, baseline
+
+
+def load_rgb(path: Path) -> np.ndarray:
+    img = imageio.imread(path)
+    if img.ndim == 2:
+        img = np.tile(img[..., None], (1, 1, 3))
+    img = img[..., :3]
+    if img.dtype != np.uint8:
+        img = np.clip(img, 0, 255).astype(np.uint8)
+    return np.ascontiguousarray(img)
+
+
+def resize_uniform_and_pad_rgb(img: np.ndarray, target_h: int, target_w: int):
+    src_h, src_w = img.shape[:2]
+    if src_h == target_h and src_w == target_w:
+        return img.copy(), target_h, target_w
+
+    scale = min(float(target_w) / src_w, float(target_h) / src_h)
+    scaled_w = max(1, int(round(src_w * scale)))
+    scaled_h = max(1, int(round(src_h * scale)))
+
+    scaled = cv2.resize(img, (scaled_w, scaled_h), interpolation=cv2.INTER_LINEAR)
+    padded = cv2.copyMakeBorder(
+        scaled,
+        0,
+        target_h - scaled_h,
+        0,
+        target_w - scaled_w,
+        cv2.BORDER_REPLICATE,
+    )
+    return np.ascontiguousarray(padded), scaled_h, scaled_w
+
+
+def tensor_from_rgb_255(img: np.ndarray) -> torch.Tensor:
+    arr = img.astype(np.float32, copy=False)
+    return torch.as_tensor(arr, device="cuda").permute(2, 0, 1).unsqueeze(0).contiguous()
+
+
+def cpp_nearest_resize(src: np.ndarray, dst_h: int, dst_w: int) -> np.ndarray:
+    src_h, src_w = src.shape
+    xs = np.rint((np.arange(dst_w, dtype=np.float32) + 0.5) * src_w / dst_w - 0.5)
+    ys = np.rint((np.arange(dst_h, dtype=np.float32) + 0.5) * src_h / dst_h - 0.5)
+    xs = np.clip(xs.astype(np.int64), 0, src_w - 1)
+    ys = np.clip(ys.astype(np.int64), 0, src_h - 1)
+    return src[ys[:, None], xs[None, :]]
+
+
+def postprocess_disparity(
+    raw_disp: np.ndarray,
+    input_h: int,
+    input_w: int,
+    model_h: int,
+    model_w: int,
+    scaled_h: int,
+    scaled_w: int,
+) -> np.ndarray:
+    disp = raw_disp.reshape(model_h, model_w).astype(np.float32, copy=False)
+    disp = np.maximum(disp, 0.0)
+
+    if input_h == model_h and input_w == model_w:
+        return np.ascontiguousarray(disp)
+
+    cropped = disp[:scaled_h, :scaled_w]
+    upsampled = cpp_nearest_resize(cropped, input_h, input_w)
+    return np.ascontiguousarray(upsampled * (float(input_w) / scaled_w), dtype=np.float32)
+
+
+def disparity_to_depth_cpp(disp: np.ndarray, fx: float, baseline_m: float) -> np.ndarray:
+    d = disp.astype(np.float32, copy=True)
+    xs = np.arange(d.shape[1], dtype=np.float32)[None, :]
+    d[(xs - d) < 0.0] = np.inf
+    with np.errstate(divide="ignore", invalid="ignore"):
+        depth = np.float32(fx * baseline_m) / d
+    return np.ascontiguousarray(depth, dtype=np.float32)
+
+
+def save_float_matrix_bin(path: Path, values: np.ndarray) -> None:
+    arr = np.ascontiguousarray(values, dtype=np.float32)
+    dims = np.array(arr.shape[:2], dtype=np.int32)
+    with path.open("wb") as f:
+        dims.tofile(f)
+        arr.tofile(f)
+
+
+def colorize_depth(depth: np.ndarray) -> np.ndarray:
+    valid = np.isfinite(depth) & (depth > 0.0)
+    if not np.any(valid):
+        return np.zeros((*depth.shape, 3), dtype=np.uint8)
+
+    safe = np.zeros_like(depth, dtype=np.float32)
+    safe[valid] = depth[valid]
+    min_val = float(safe[valid].min())
+    max_val = float(safe[valid].max())
+    if max_val <= min_val:
+        max_val = min_val + 1.0
+
+    depth_u8 = np.clip((safe - min_val) * (255.0 / (max_val - min_val)), 0, 255).astype(np.uint8)
+    colored_bgr = cv2.applyColorMap(depth_u8, cv2.COLORMAP_TURBO)
+    colored = colored_bgr[..., ::-1]
+    colored[~valid] = 0
+    return colored
+
+
+class PluginTensorRTRunner:
+    def __init__(self, engine_path: Path, plugin_lib: Path):
+        import tensorrt as trt
+
+        self.trt = trt
+        self.logger = trt.Logger(trt.Logger.WARNING)
+        trt.init_libnvinfer_plugins(self.logger, "")
+        load_plugin_library(str(plugin_lib))
+        if not find_plugin_creator(trt):
+            raise RuntimeError(f"{PLUGIN_NAME} plugin creator is not registered after loading {plugin_lib}")
+
+        self.runtime = trt.Runtime(self.logger)
+        self.engine = self.runtime.deserialize_cuda_engine(engine_path.read_bytes())
+        if self.engine is None:
+            raise RuntimeError(f"Failed to deserialize TensorRT engine: {engine_path}")
+        self.context = self.engine.create_execution_context()
+        if self.context is None:
+            raise RuntimeError(f"Failed to create execution context: {engine_path}")
+
+        names = [self.engine.get_tensor_name(i) for i in range(self.engine.num_io_tensors)]
+        missing = [name for name in ("left", "right", "disp") if name not in names]
+        if missing:
+            raise RuntimeError(
+                f"Plugin engine must expose tensors named left, right, and disp. "
+                f"Missing {missing}; found {names}"
+            )
+
+    def _torch_dtype(self, trt_dtype):
+        trt = self.trt
+        mapping = {
+            trt.DataType.FLOAT: torch.float32,
+            trt.DataType.HALF: torch.float16,
+            trt.DataType.INT32: torch.int32,
+            trt.DataType.INT8: torch.int8,
+            trt.DataType.BOOL: torch.bool,
+        }
+        if hasattr(trt.DataType, "BF16"):
+            mapping[trt.DataType.BF16] = torch.bfloat16
+        if trt_dtype not in mapping:
+            raise RuntimeError(f"Unsupported TensorRT dtype: {trt_dtype}")
+        return mapping[trt_dtype]
+
+    def infer(self, left: torch.Tensor, right: torch.Tensor) -> torch.Tensor:
+        inputs = {"left": left, "right": right}
+        for name, tensor in list(inputs.items()):
+            expected = self._torch_dtype(self.engine.get_tensor_dtype(name))
+            if tensor.dtype != expected:
+                tensor = tensor.to(expected)
+            inputs[name] = tensor.contiguous()
+            self.context.set_input_shape(name, tuple(inputs[name].shape))
+
+        out_shape = tuple(self.context.get_tensor_shape("disp"))
+        out_dtype = self._torch_dtype(self.engine.get_tensor_dtype("disp"))
+        output = torch.empty(out_shape, device="cuda", dtype=out_dtype)
+
+        self.context.set_tensor_address("left", int(inputs["left"].data_ptr()))
+        self.context.set_tensor_address("right", int(inputs["right"].data_ptr()))
+        self.context.set_tensor_address("disp", int(output.data_ptr()))
+
+        stream = torch.cuda.current_stream().cuda_stream
+        if not self.context.execute_async_v3(stream):
+            raise RuntimeError("TensorRT enqueue failed")
+        return output
+
+
+def main() -> int:
+    args = parse_args()
+    set_logging_format()
+    set_seed(0)
+    torch.autograd.set_grad_enabled(False)
+    args.out_dir.mkdir(parents=True, exist_ok=True)
+
+    engine_path = resolve_engine_path(args.model_dir, args.engine_file)
+    config_path = resolve_config_path(engine_path, args.model_dir, args.config_file)
+    plugin_lib = args.plugin_lib or find_default_plugin_library()
+    if plugin_lib is None or not plugin_lib.exists():
+        raise FileNotFoundError(
+            "Could not find libffs_gwc_plugin.so. Build cpp first or pass "
+            "--plugin_lib /path/to/libffs_gwc_plugin.so"
+        )
+
+    with config_path.open("r") as f:
+        cfg = yaml.safe_load(f)
+    model_h, model_w = [int(v) for v in cfg["image_size"]]
+
+    left = load_rgb(args.left_file)
+    right = load_rgb(args.right_file)
+    if left.shape != right.shape:
+        raise RuntimeError(f"left/right image size mismatch: {left.shape} vs {right.shape}")
+    input_h, input_w = left.shape[:2]
+
+    logging.info(f"Engine: {engine_path}")
+    logging.info(f"Plugin: {plugin_lib}")
+    logging.info(f"Config: {config_path}")
+    logging.info(f"Input images: {input_w}x{input_h}")
+    logging.info(f"Model target resolution: {model_w}x{model_h}")
+
+    left_model, scaled_h, scaled_w = resize_uniform_and_pad_rgb(left, model_h, model_w)
+    right_model, right_scaled_h, right_scaled_w = resize_uniform_and_pad_rgb(right, model_h, model_w)
+    if (scaled_h, scaled_w) != (right_scaled_h, right_scaled_w):
+        raise RuntimeError("left/right images produced different resize scales")
+    if (input_h, input_w) != (model_h, model_w):
+        logging.info(
+            f"Uniform resize+pad: {input_w}x{input_h} -> "
+            f"{scaled_w}x{scaled_h} inside {model_w}x{model_h}"
+        )
+
+    runner = PluginTensorRTRunner(engine_path, plugin_lib)
+    t_left = tensor_from_rgb_255(left_model)
+    t_right = tensor_from_rgb_255(right_model)
+
+    logging.info("Running TensorRT inference")
+    disp_raw_t = runner.infer(t_left, t_right)
+    torch.cuda.current_stream().synchronize()
+    logging.info("Inference done")
+
+    disp_raw = disp_raw_t.float().detach().cpu().numpy()
+    disp = postprocess_disparity(
+        disp_raw,
+        input_h=input_h,
+        input_w=input_w,
+        model_h=model_h,
+        model_w=model_w,
+        scaled_h=scaled_h,
+        scaled_w=scaled_w,
+    )
+
+    if not np.isfinite(disp).any():
+        raise RuntimeError("Model produced no finite disparity values")
+
+    k, baseline = load_intrinsics(args.intrinsic_file)
+    depth = disparity_to_depth_cpp(disp, float(k[0, 0]), baseline)
+
+    disparity_path = args.out_dir / "disparity.bin"
+    depth_bin_path = args.out_dir / "depth_meter.bin"
+    depth_npy_path = args.out_dir / "depth_meter.npy"
+    disp_vis_path = args.out_dir / "disp_vis.png"
+    depth_vis_path = args.out_dir / "depth_vis.png"
+
+    save_float_matrix_bin(disparity_path, disp)
+    save_float_matrix_bin(depth_bin_path, depth)
+    np.save(depth_npy_path, depth)
+
+    disp_color = vis_disparity(disp, color_map=cv2.COLORMAP_TURBO)
+    disp_vis = np.concatenate([left, right, disp_color], axis=1)
+    imageio.imwrite(disp_vis_path, disp_vis)
+    imageio.imwrite(depth_vis_path, colorize_depth(depth))
+
+    logging.info(f"Saved: {disparity_path}")
+    logging.info(f"Saved: {depth_bin_path}")
+    logging.info(f"Saved: {depth_npy_path}")
+    logging.info(f"Saved: {disp_vis_path}")
+    logging.info(f"Saved: {depth_vis_path}")
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

cv/fast-foundationstereo/scripts/run_demo_single_trt.py

@@ -0,0 +1,348 @@
+# SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES
+# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+Run Fast FoundationStereo inference with the single ONNX model (or TRT engine)
+produced by make_single_onnx.py.
+
+Supports two backends:
+  - ONNX Runtime (default if --model_file points to an .onnx, or auto-detected)
+  - TensorRT     (if --model_file points to an .engine)
+
+The model expects ImageNet-normalised inputs, so this script applies
+normalisation during preprocessing.
+
+Usage:
+  # Run directly with ONNX (no trtexec step needed):
+  python run_demo_single_trt.py \
+      --model_dir ./output_single_onnx \
+      --left_file  ../demo_data/left.png \
+      --right_file ../demo_data/right.png
+
+  # Or with an explicit model file:
+  python run_demo_single_trt.py \
+      --model_dir  ./output_single_onnx \
+      --model_file ./output_single_onnx/fast_foundationstereo.onnx \
+      --left_file  ../demo_data/left.png \
+      --right_file ../demo_data/right.png
+"""
+
+import argparse
+import logging
+import os
+import sys
+
+import cv2
+import imageio
+import numpy as np
+import torch
+import yaml
+
+code_dir = os.path.dirname(os.path.realpath(__file__))
+sys.path.append(f'{code_dir}/../')
+
+from Utils import (
+    set_logging_format, set_seed, vis_disparity,
+    depth2xyzmap, toOpen3dCloud, o3d,
+)
+
+IMAGENET_MEAN = np.array([0.485, 0.456, 0.406], dtype=np.float32)
+IMAGENET_STD  = np.array([0.229, 0.224, 0.225], dtype=np.float32)
+
+
+class SingleEngineTrtRunner:
+    """Minimal TensorRT runner for a single engine with named I/O."""
+
+    def __init__(self, engine_path):
+        import tensorrt as trt
+        self.trt = trt
+        self.logger = trt.Logger(trt.Logger.WARNING)
+
+        with open(engine_path, 'rb') as f:
+            self.engine = trt.Runtime(self.logger).deserialize_cuda_engine(f.read())
+        if self.engine is None:
+            raise RuntimeError(
+                f'Failed to deserialize TRT engine from {engine_path}. '
+                f'This usually means the engine was built with a different '
+                f'TensorRT version (yours: {trt.__version__}). '
+                f'Rebuild with:  trtexec --onnx=<your .onnx> '
+                f'--saveEngine={engine_path} --fp16')
+        self.context = self.engine.create_execution_context()
+
+    def _trt_to_torch_dtype(self, dt):
+        trt = self.trt
+        mapping = {
+            trt.DataType.FLOAT:  torch.float32,
+            trt.DataType.HALF:   torch.float16,
+            trt.DataType.BF16:   torch.bfloat16,
+            trt.DataType.INT32:  torch.int32,
+            trt.DataType.INT8:   torch.int8,
+            trt.DataType.BOOL:   torch.bool,
+        }
+        if dt not in mapping:
+            raise RuntimeError(f'Unsupported TRT dtype: {dt}')
+        return mapping[dt]
+
+    def __call__(self, inputs: dict) -> dict:
+        """Run inference.
+
+        Args:
+            inputs: {binding_name: torch.Tensor} for every input tensor.
+        Returns:
+            {binding_name: torch.Tensor} for every output tensor.
+        """
+        trt = self.trt
+
+        for name, tensor in inputs.items():
+            expected = self._trt_to_torch_dtype(self.engine.get_tensor_dtype(name))
+            if tensor.dtype != expected:
+                inputs[name] = tensor.to(expected)
+            if not inputs[name].is_contiguous():
+                inputs[name] = inputs[name].contiguous()
+            self.context.set_input_shape(name, tuple(inputs[name].shape))
+
+        out_names = [
+            self.engine.get_tensor_name(i)
+            for i in range(self.engine.num_io_tensors)
+            if self.engine.get_tensor_mode(self.engine.get_tensor_name(i))
+               == trt.TensorIOMode.OUTPUT
+        ]
+
+        outputs = {}
+        for name in out_names:
+            shape = tuple(self.context.get_tensor_shape(name))
+            dtype = self._trt_to_torch_dtype(self.engine.get_tensor_dtype(name))
+            outputs[name] = torch.empty(shape, device='cuda', dtype=dtype)
+
+        for name, tensor in inputs.items():
+            self.context.set_tensor_address(name, int(tensor.data_ptr()))
+        for name, tensor in outputs.items():
+            self.context.set_tensor_address(name, int(tensor.data_ptr()))
+
+        stream = torch.cuda.current_stream().cuda_stream
+        assert self.context.execute_async_v3(stream)
+
+        return outputs
+
+
+class OnnxRuntimeRunner:
+    """Run inference via ONNX Runtime (GPU if available, else CPU)."""
+
+    def __init__(self, onnx_path):
+        import onnxruntime as ort
+        providers = []
+        if 'CUDAExecutionProvider' in ort.get_available_providers():
+            providers.append('CUDAExecutionProvider')
+        providers.append('CPUExecutionProvider')
+        logging.info(f'ONNX Runtime providers: {providers}')
+        self.session = ort.InferenceSession(onnx_path, providers=providers)
+        self.input_names = [inp.name for inp in self.session.get_inputs()]
+        self.output_names = [out.name for out in self.session.get_outputs()]
+
+    def __call__(self, inputs: dict) -> dict:
+        feed = {}
+        for name in self.input_names:
+            tensor = inputs[name]
+            if isinstance(tensor, torch.Tensor):
+                tensor = tensor.cpu().float().numpy()
+            feed[name] = tensor
+        raw_outputs = self.session.run(self.output_names, feed)
+        outputs = {}
+        for name, arr in zip(self.output_names, raw_outputs):
+            outputs[name] = torch.as_tensor(arr).cuda()
+        return outputs
+
+
+def normalize_imagenet(img_uint8: np.ndarray) -> np.ndarray:
+    """Apply ImageNet normalization: (img/255 - mean) / std."""
+    return ((img_uint8.astype(np.float32) / 255.0) - IMAGENET_MEAN) / IMAGENET_STD
+
+
+def resolve_config(model_path: str) -> str:
+    """Find the YAML config matching the model file, falling back to defaults."""
+    model_dir = os.path.dirname(model_path)
+    base = os.path.splitext(os.path.basename(model_path))[0]
+    candidates = [
+        os.path.join(model_dir, f'{base}.yaml'),
+        os.path.join(model_dir, 'config.yaml'),
+        os.path.join(model_dir, 'onnx.yaml'),
+    ]
+    for p in candidates:
+        if os.path.exists(p):
+            return p
+    raise FileNotFoundError(
+        f'No .yaml config found for {model_path}. '
+        'Run make_single_onnx.py first.')
+
+
+def find_model(model_dir: str) -> str:
+    """Find an .engine or .onnx file in the directory (prefer .engine)."""
+    for ext in ('.engine', '.onnx'):
+        for f in os.listdir(model_dir):
+            if f.endswith(ext):
+                return os.path.join(model_dir, f)
+    raise FileNotFoundError(
+        f'No .engine or .onnx file found in {model_dir}. '
+        'Run make_single_onnx.py first.')
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(
+        description='Run Fast FoundationStereo with ONNX Runtime or TensorRT')
+    parser.add_argument('--model_dir', type=str,
+                        default=f'{code_dir}/output_single_onnx',
+                        help='Directory containing .onnx/.engine + config.yaml')
+    parser.add_argument('--model_file', type=str, default='',
+                        help='Explicit path to .onnx or .engine file (overrides auto-search)')
+    parser.add_argument('--left_file', type=str,
+                        default=f'{code_dir}/../demo_data/left.png')
+    parser.add_argument('--right_file', type=str,
+                        default=f'{code_dir}/../demo_data/right.png')
+    parser.add_argument('--intrinsic_file', type=str,
+                        default=f'{code_dir}/../demo_data/K.txt',
+                        help='Camera intrinsic matrix and baseline file')
+    parser.add_argument('--out_dir', type=str,
+                        default=f'{code_dir}/../output_demo')
+    parser.add_argument('--remove_invisible', type=int, default=1)
+    parser.add_argument('--denoise_cloud', type=int, default=1)
+    parser.add_argument('--denoise_nb_points', type=int, default=30)
+    parser.add_argument('--denoise_radius', type=float, default=0.03)
+    parser.add_argument('--get_pc', type=int, default=1,
+                        help='Generate and save point cloud')
+    parser.add_argument('--zfar', type=float, default=100,
+                        help='Max depth (m) to include in point cloud')
+    args = parser.parse_args()
+
+    set_logging_format()
+    set_seed(0)
+    torch.autograd.set_grad_enabled(False)
+    os.makedirs(args.out_dir, exist_ok=True)
+
+    # ── Find model and config ─────────────────────────────────────────────
+    model_path = args.model_file if args.model_file else find_model(args.model_dir)
+    cfg_path = resolve_config(model_path)
+    with open(cfg_path, 'r') as f:
+        cfg = yaml.safe_load(f)
+    target_h, target_w = cfg['image_size']
+    logging.info(f'Model target resolution: {target_h} x {target_w}')
+
+    # ── Load model (ONNX Runtime or TensorRT) ────────────────────────────
+    logging.info(f'Loading model: {model_path}')
+    if model_path.endswith('.onnx'):
+        runner = OnnxRuntimeRunner(model_path)
+    else:
+        runner = SingleEngineTrtRunner(model_path)
+
+    # ── Read images ───────────────────────────────────────────────────────
+    img0 = imageio.imread(args.left_file)
+    img1 = imageio.imread(args.right_file)
+
+    if img0.ndim == 2:
+        img0 = np.tile(img0[..., None], (1, 1, 3))
+        img1 = np.tile(img1[..., None], (1, 1, 3))
+    img0 = img0[..., :3]
+    img1 = img1[..., :3]
+
+    # ── Resize to model resolution (direct stretch) ────────────────────────
+    orig_h, orig_w = img0.shape[:2]
+    fx = target_w / orig_w
+    fy = target_h / orig_h
+
+    if fx != 1 or fy != 1:
+        logging.info(
+            f'Resizing images: {orig_h}x{orig_w} → {target_h}x{target_w} '
+            f'(fx={fx:.4f}, fy={fy:.4f})')
+        img0 = cv2.resize(img0, (target_w, target_h), interpolation=cv2.INTER_LINEAR)
+        img1 = cv2.resize(img1, (target_w, target_h), interpolation=cv2.INTER_LINEAR)
+    H, W = img0.shape[:2]
+
+    img0_ori = img0.copy()
+    img1_ori = img1.copy()
+    logging.info(f'Image size after resize: {img0.shape}')
+    imageio.imwrite(f'{args.out_dir}/left.png', img0)
+    imageio.imwrite(f'{args.out_dir}/right.png', img1)
+
+    # ── Preprocess: ImageNet normalize → NCHW float tensor ────────────────
+    img0_norm = normalize_imagenet(img0)
+    img1_norm = normalize_imagenet(img1)
+
+    t_left  = torch.as_tensor(img0_norm).cuda().float()[None].permute(0, 3, 1, 2)
+    t_right = torch.as_tensor(img1_norm).cuda().float()[None].permute(0, 3, 1, 2)
+
+    # ── Inference ─────────────────────────────────────────────────────────
+    logging.info('Running inference (first run may be slow due to TRT warmup)')
+    outputs = runner({'left_image': t_left, 'right_image': t_right})
+    disp = outputs['disparity']
+    logging.info('Inference done')
+
+    disp = disp.float().cpu().numpy().reshape(H, W).clip(0, None) * (1.0 / fx)
+
+    # ── Visualise disparity ──────────────────────────────────────────────
+    vis = vis_disparity(disp, color_map=cv2.COLORMAP_TURBO)
+    vis = np.concatenate([img0_ori, img1_ori, vis], axis=1)
+    imageio.imwrite(f'{args.out_dir}/disp_vis.png', vis)
+    s = 1280 / vis.shape[1]
+    resized_vis = cv2.resize(vis, (int(vis.shape[1] * s), int(vis.shape[0] * s)))
+    cv2.imshow('disp', resized_vis[:, :, ::-1])
+    cv2.waitKey(0)
+
+    # ── Remove invisible pixels ──────────────────────────────────────────
+    if args.remove_invisible:
+        _, xx = np.meshgrid(np.arange(H), np.arange(W), indexing='ij')
+        invalid = (xx - disp) < 0
+        disp[invalid] = np.inf
+
+    # ── Point cloud generation ───────────────────────────────────────────
+    if args.get_pc:
+        with open(args.intrinsic_file, 'r') as f:
+            lines = f.readlines()
+            K = (np.array(list(map(float, lines[0].rstrip().split())))
+                 .astype(np.float32).reshape(3, 3))
+            baseline = float(lines[1])
+        K[:2] *= np.array([fx, fy], dtype=np.float32)[:, np.newaxis]
+        depth = K[0, 0] * baseline / disp
+        np.save(f'{args.out_dir}/depth_meter.npy', depth)
+
+        xyz_map = depth2xyzmap(depth, K)
+        pcd = toOpen3dCloud(xyz_map.reshape(-1, 3), img0_ori.reshape(-1, 3))
+        pts = np.asarray(pcd.points)
+        keep = (pts[:, 2] > 0) & (pts[:, 2] <= args.zfar)
+        pcd = pcd.select_by_index(np.where(keep)[0])
+        o3d.io.write_point_cloud(f'{args.out_dir}/cloud.ply', pcd)
+        logging.info(f'Point cloud saved to {args.out_dir}')
+
+        if args.denoise_cloud:
+            logging.info('Denoising point cloud...')
+            _, ind = pcd.remove_radius_outlier(
+                nb_points=args.denoise_nb_points,
+                radius=args.denoise_radius)
+            pcd = pcd.select_by_index(ind)
+            o3d.io.write_point_cloud(f'{args.out_dir}/cloud_denoise.ply', pcd)
+
+        logging.info('Visualizing point cloud. Press ESC to exit.')
+        vis = o3d.visualization.Visualizer()
+        vis.create_window()
+        vis.add_geometry(pcd)
+        vis.get_render_option().point_size = 1.0
+        vis.get_render_option().background_color = np.array([0.5, 0.5, 0.5])
+        ctr = vis.get_view_control()
+        ctr.set_front([0, 0, -1])
+        closest = np.asarray(pcd.points)[:, 2].argmin()
+        ctr.set_lookat(np.asarray(pcd.points)[closest])
+        ctr.set_up([0, -1, 0])
+        vis.run()
+        vis.destroy_window()

cv/fast-foundationstereo/scripts/run_demo_tensorrt.py

@@ -0,0 +1,145 @@
+import os,sys
+code_dir = os.path.dirname(os.path.realpath(__file__))
+sys.path.append(f'{code_dir}/../')
+from omegaconf import OmegaConf
+from core.utils.utils import InputPadder
+import argparse, torch, logging, yaml
+import imageio
+import numpy as np
+from Utils import (
+    set_logging_format, set_seed, vis_disparity,
+    depth2xyzmap, toOpen3dCloud, o3d,
+)
+from core.foundation_stereo import TrtRunner
+import cv2
+
+
+def resolve_onnx_cfg_path(onnx_dir: str) -> str:
+  onnx_dir = os.path.normpath(onnx_dir)
+  candidates = [
+    os.path.join(onnx_dir, 'onnx.yaml'),
+    os.path.join(os.path.dirname(onnx_dir), 'onnx.yaml'),
+  ]
+  for p in candidates:
+    if os.path.exists(p):
+      return p
+  raise FileNotFoundError(
+    f"onnx.yaml not found. Looked in: {candidates}. "
+    "Please run scripts/make_onnx.py first to generate ONNX metadata."
+  )
+
+
+if __name__=="__main__":
+  code_dir = os.path.dirname(os.path.realpath(__file__))
+  parser = argparse.ArgumentParser()
+  parser.add_argument('--onnx_dir', default=f'{code_dir}/output', type=str)
+  parser.add_argument('--left_file', default=f'{code_dir}/../assets/left.png', type=str)
+  parser.add_argument('--right_file', default=f'{code_dir}/../assets/right.png', type=str)
+  parser.add_argument('--intrinsic_file', default=f'{code_dir}/../assets/K.txt', type=str, help='camera intrinsic matrix and baseline file')
+  parser.add_argument('--out_dir', default='/home/bowen/debug/stereo_output', type=str)
+  parser.add_argument('--remove_invisible', default=1, type=int)
+  parser.add_argument('--denoise_cloud', default=1, type=int)
+  parser.add_argument('--denoise_nb_points', type=int, default=30, help='number of points to consider for radius outlier removal')
+  parser.add_argument('--denoise_radius', type=float, default=0.03, help='radius to use for outlier removal')
+  parser.add_argument('--get_pc', type=int, default=1, help='save point cloud output')
+  parser.add_argument('--zfar', type=float, default=100, help="max depth to include in point cloud")
+  args = parser.parse_args()
+
+  set_logging_format()
+  set_seed(0)
+  torch.autograd.set_grad_enabled(False)
+  os.makedirs(args.out_dir, exist_ok=True)
+
+  onnx_cfg_path = resolve_onnx_cfg_path(args.onnx_dir)
+  with open(onnx_cfg_path, 'r') as ff:
+    cfg:dict = yaml.safe_load(ff)
+  for k in args.__dict__:
+    if args.__dict__[k] is not None:
+      cfg[k] = args.__dict__[k]
+  args = OmegaConf.create(cfg)
+  logging.info(f"args:\n{args}")
+  model = TrtRunner(args, args.onnx_dir+'/feature_runner.engine', args.onnx_dir+'/post_runner.engine')
+
+  img0 = imageio.imread(args.left_file)
+  img1 = imageio.imread(args.right_file)
+  if len(img0.shape)==2:
+    img0 = np.tile(img0[...,None], (1,1,3))
+    img1 = np.tile(img1[...,None], (1,1,3))
+  img0 = img0[...,:3]
+  img1 = img1[...,:3]
+  H,W = img0.shape[:2]
+
+  fx = args.image_size[1] / img0.shape[1]
+  fy = args.image_size[0] / img0.shape[0]
+  if fx != 1 or fy != 1:
+    logging.info(f">>>>>>>>>>>>>>>WARNING: resizing image to {args.image_size}, fx: {fx}, fy: {fy}, this is not recommended. It's best to make tensorrt engine with the same image size as the input image.")
+  img0 = cv2.resize(img0, fx=fx, fy=fy, dsize=None)
+  img1 = cv2.resize(img1, fx=fx, fy=fy, dsize=None)
+  H,W = img0.shape[:2]
+  img0_ori = img0.copy()
+  img1_ori = img1.copy()
+  logging.info(f"img0: {img0.shape}")
+  imageio.imwrite(f'{args.out_dir}/left.png', img0)
+  imageio.imwrite(f'{args.out_dir}/right.png', img1)
+
+  img0 = torch.as_tensor(img0).cuda().float()[None].permute(0,3,1,2)
+  img1 = torch.as_tensor(img1).cuda().float()[None].permute(0,3,1,2)
+
+  logging.info(f"Start forward, 1st time run can be slow due to compilation")
+  disp = model.forward(img0, img1)
+  logging.info("forward done")
+  disp = disp.data.cpu().numpy().reshape(H,W).clip(0, None) * 1/fx
+
+  cmap = None
+  min_val = None
+  max_val = None
+  vis = vis_disparity(disp, min_val=min_val, max_val=max_val, cmap=cmap, color_map=cv2.COLORMAP_TURBO)
+  vis = np.concatenate([img0_ori, img1_ori, vis], axis=1)
+  imageio.imwrite(f'{args.out_dir}/disp_vis.png', vis)
+  s = 1280/vis.shape[1]
+  resized_vis = cv2.resize(vis, (int(vis.shape[1]*s), int(vis.shape[0]*s)))
+  cv2.imshow('disp', resized_vis[:,:,::-1])
+  cv2.waitKey(0)
+
+  if args.remove_invisible:
+    yy,xx = np.meshgrid(np.arange(disp.shape[0]), np.arange(disp.shape[1]), indexing='ij')
+    us_right = xx-disp
+    invalid = us_right<0
+    disp[invalid] = np.inf
+
+  if args.get_pc:
+    with open(args.intrinsic_file, 'r') as f:
+      lines = f.readlines()
+      K = np.array(list(map(float, lines[0].rstrip().split()))).astype(np.float32).reshape(3,3)
+      baseline = float(lines[1])
+    K[:2] *= np.array([fx, fy], dtype=np.float32)[:, np.newaxis]
+    depth = K[0,0]*baseline/disp
+    np.save(f'{args.out_dir}/depth_meter.npy', depth)
+    xyz_map = depth2xyzmap(depth, K)
+    pcd = toOpen3dCloud(xyz_map.reshape(-1,3), img0_ori.reshape(-1,3))
+    keep_mask = (np.asarray(pcd.points)[:,2]>0) & (np.asarray(pcd.points)[:,2]<=args.zfar)
+    keep_ids = np.arange(len(np.asarray(pcd.points)))[keep_mask]
+    pcd = pcd.select_by_index(keep_ids)
+    o3d.io.write_point_cloud(f'{args.out_dir}/cloud.ply', pcd)
+    logging.info(f"PCL saved to {args.out_dir}")
+
+    if args.denoise_cloud:
+      logging.info("[Optional step] denoise point cloud...")
+      cl, ind = pcd.remove_radius_outlier(nb_points=args.denoise_nb_points, radius=args.denoise_radius)
+      inlier_cloud = pcd.select_by_index(ind)
+      o3d.io.write_point_cloud(f'{args.out_dir}/cloud_denoise.ply', inlier_cloud)
+      pcd = inlier_cloud
+
+    logging.info("Visualizing point cloud. Press ESC to exit.")
+    vis = o3d.visualization.Visualizer()
+    vis.create_window()
+    vis.add_geometry(pcd)
+    vis.get_render_option().point_size = 1.0
+    vis.get_render_option().background_color = np.array([0.5, 0.5, 0.5])
+    ctr = vis.get_view_control()
+    ctr.set_front([0, 0, -1])
+    id = np.asarray(pcd.points)[:,2].argmin()
+    ctr.set_lookat(np.asarray(pcd.points)[id])
+    ctr.set_up([0, -1, 0])
+    vis.run()
+    vis.destroy_window()