vs-undistort 2.2.0__tar.gz

vs_undistort-2.2.0/.github/workflows/publish.yml

@@ -0,0 +1,76 @@
+name: Publish to PyPI
+
+on:
+  release:
+    types: [published]
+
+permissions:
+  contents: read
+
+jobs:
+  build:
+    name: Build distributions
+    if: ${{ !github.event.release.prerelease }}
+    runs-on: ubuntu-latest
+
+    steps:
+      - name: Check out released tag
+        uses: actions/checkout@v6
+        with:
+          ref: ${{ github.event.release.tag_name }}
+          persist-credentials: false
+
+      - name: Set up Python
+        uses: actions/setup-python@v6
+        with:
+          python-version: "3.12"
+
+      - name: Check tag matches package version
+        shell: bash
+        run: |
+          VERSION="$(python -c 'import tomllib; print(tomllib.load(open("pyproject.toml", "rb"))["project"]["version"])')"
+          TAG="${{ github.event.release.tag_name }}"
+          TAG_VERSION="${TAG#v}"
+
+          if [ "$TAG_VERSION" != "$VERSION" ]; then
+            echo "Release tag $TAG does not match pyproject.toml version $VERSION"
+            exit 1
+          fi
+
+      - name: Install build tools
+        run: python -m pip install --upgrade build twine
+
+      - name: Build wheel and source distribution
+        run: python -m build
+
+      - name: Validate distributions
+        run: python -m twine check dist/*
+
+      - name: Upload distributions
+        uses: actions/upload-artifact@v5
+        with:
+          name: python-package-distributions
+          path: dist/
+          if-no-files-found: error
+
+  publish:
+    name: Publish distributions to PyPI
+    needs: build
+    runs-on: ubuntu-latest
+
+    environment:
+      name: pypi
+      url: https://pypi.org/p/vs_undistort
+
+    permissions:
+      id-token: write
+
+    steps:
+      - name: Download distributions
+        uses: actions/download-artifact@v6
+        with:
+          name: python-package-distributions
+          path: dist/
+
+      - name: Publish to PyPI
+        uses: pypa/gh-action-pypi-publish@release/v1

vs_undistort-2.2.0/PKG-INFO

@@ -0,0 +1,92 @@
+Metadata-Version: 2.4
+Name: vs_undistort
+Version: 2.2.0
+Summary: Vapoursynth function to remove video distortions, turbulence, wobble, warp, heat haze, or similar.
+Project-URL: Homepage, https://github.com/pifroggi/vs_undistort
+Project-URL: Repository, https://github.com/pifroggi/vs_undistort
+Project-URL: Issues, https://github.com/pifroggi/vs_undistort/issues
+Author: pifroggi
+Keywords: distortion,filtering,heat haze,rectification,restoration,turbulence,vapoursynth,video,wobble
+Classifier: Programming Language :: Python :: 3
+Classifier: Topic :: Multimedia :: Video
+Requires-Python: >=3.12
+Requires-Dist: numpy
+Requires-Dist: tensorrt-cu13>=11.1.0.106
+Requires-Dist: vapoursynth-mlrt-trt>=16.0
+Requires-Dist: vapoursynth>=74
+Description-Content-Type: text/markdown
+
+# Video Distortion Removal for VapourSynth
+Also known as dewobble, warp stabilization, film or VHS distortion fix, rectification, atmospheric turbulence mitigation, or heat haze removal.
+
+This is a partial implementation of the [Turbulence Mitigation Transformer](https://github.com/xg416/TMT) (only distortion removal, no deblurring). It does not do general video stabilization for shaky footage, only removes distortions within the frames. It is recommented to stabilize first if needed.
+
+Check out stinkybread's comparisons [here](https://kaizoku.pw/c/tezc1iieeHympQGQeWKl) and [here](https://kaizoku.pw/c/4HsCLWDxuFMNpOIdSNI8).
+
+<p align="center">
+    <img src="https://github.com/xg416/TMT/blob/main/figs/video_22.gif"/>
+</p>
+
+<br />
+
+## Installation
+
+```
+pip install -U vs_undistort
+```
+* To enable the CPU/CUDA backends, install [PyTorch with CUDA](https://pytorch.org/). *(optional)*
+<br />
+
+For older vapoursynth versions below R74, follow the manual installation steps [here](https://github.com/pifroggi/vs_temporalfix/wiki/Manual-Installation).
+
+<br />
+
+## Usage
+
+```python
+from vs_undistort import vs_undistort
+clip = vs_undistort(clip, temp_window=10, tiles=1, overlap=8, interpolation="bicubic", backend="tensorrt", num_streams=1, engine_folder=None)
+```
+
+__*`clip`*__  
+Distorted clip. Must be in RGBH format.
+
+__*`temp_window`*__  
+Temporal window length. How many frames are grouped together and processed as a single chunk. Larger means higher VRAM requirements, but better temporal averaging and slower distortions can be removed. If this is too small, some distortions may not get removed, small jumps/hitches may be visible between windows and seams from tiling may become more obvious.
+
+__*`tiles`* (optional)__  
+Amount of tiles to split the frames into. A higher amount reduces VRAM requirements, but also worsens spatial averaging. Default tiles=1 uses the full frame.
+
+__*`overlap`* (optional)__  
+Overlap from one tile to the next. Increase if seams between tiles are visible.
+
+__*`interpolation`* (optional)__  
+The interpolation mode used to warp the frames:
+* `bilinear` More blurry.
+* `bicubic` No blur, but may oversharpen slightly.
+
+__*`backend`* (optional)__  
+The backend used to run the model:
+* `cpu` CPU mode using PyTorch *(very slow)*.
+* `cuda` GPU mode using PyTorch with CUDA support. Requires any Nvidia GPU *(fast)*.
+* `tensorrt` GPU mode using vs-mlrt with TensorRT support. Requires an Nvidia RTX GPU. On the first run, this mode will automatically build an engine, which may take a few minutes. Changing interpolation, temp_window, or input dimensions will trigger rebuilding, but previously build engines are stored *(very fast/low vram)*.
+
+__*`num_streams`* (optional)__  
+Number of parallel TensorRT streams. For high end GPUs higher can be a bit faster, but requires more VRAM. Only affects the TensorRT backend.
+
+__*`engine_folder`* (optional)__  
+Optional path to the TensorRT engine storage location. By default engines are stored in `vs_undistort/engines`. Only affects the TensorRT backend.
+
+> [!TIP]
+> * If you see jumps/hitches between temporal windows, you can crossfade them with [vs_tiletools](https://github.com/pifroggi/vs_tiletools?tab=readme-ov-file#fix-jumpshitches-on-chunktemporal-window-based-filters-via-crossfading).
+> * If you have an undistorted reference clip, you can also try to align to it with [vs_align](https://github.com/pifroggi/vs_align).
+
+<br />
+
+## Benchmarks
+
+| Hardware | Resolution | TensorRT | CUDA
+|----------|------------|----------|---------
+| RTX 4090 | 720x480    | ~35 fps  | ~6.5 fps
+| RTX 4090 | 1440x1080  | ~7.5 fps | ~1.5 fps
+| RTX 4090 | 2880x2160  | ~2 fps   | ~0.5 fps

vs_undistort-2.2.0/README.md

@@ -0,0 +1,74 @@
+# Video Distortion Removal for VapourSynth
+Also known as dewobble, warp stabilization, film or VHS distortion fix, rectification, atmospheric turbulence mitigation, or heat haze removal.
+
+This is a partial implementation of the [Turbulence Mitigation Transformer](https://github.com/xg416/TMT) (only distortion removal, no deblurring). It does not do general video stabilization for shaky footage, only removes distortions within the frames. It is recommented to stabilize first if needed.
+
+Check out stinkybread's comparisons [here](https://kaizoku.pw/c/tezc1iieeHympQGQeWKl) and [here](https://kaizoku.pw/c/4HsCLWDxuFMNpOIdSNI8).
+
+<p align="center">
+    <img src="https://github.com/xg416/TMT/blob/main/figs/video_22.gif"/>
+</p>
+
+<br />
+
+## Installation
+
+```
+pip install -U vs_undistort
+```
+* To enable the CPU/CUDA backends, install [PyTorch with CUDA](https://pytorch.org/). *(optional)*
+<br />
+
+For older vapoursynth versions below R74, follow the manual installation steps [here](https://github.com/pifroggi/vs_temporalfix/wiki/Manual-Installation).
+
+<br />
+
+## Usage
+
+```python
+from vs_undistort import vs_undistort
+clip = vs_undistort(clip, temp_window=10, tiles=1, overlap=8, interpolation="bicubic", backend="tensorrt", num_streams=1, engine_folder=None)
+```
+
+__*`clip`*__  
+Distorted clip. Must be in RGBH format.
+
+__*`temp_window`*__  
+Temporal window length. How many frames are grouped together and processed as a single chunk. Larger means higher VRAM requirements, but better temporal averaging and slower distortions can be removed. If this is too small, some distortions may not get removed, small jumps/hitches may be visible between windows and seams from tiling may become more obvious.
+
+__*`tiles`* (optional)__  
+Amount of tiles to split the frames into. A higher amount reduces VRAM requirements, but also worsens spatial averaging. Default tiles=1 uses the full frame.
+
+__*`overlap`* (optional)__  
+Overlap from one tile to the next. Increase if seams between tiles are visible.
+
+__*`interpolation`* (optional)__  
+The interpolation mode used to warp the frames:
+* `bilinear` More blurry.
+* `bicubic` No blur, but may oversharpen slightly.
+
+__*`backend`* (optional)__  
+The backend used to run the model:
+* `cpu` CPU mode using PyTorch *(very slow)*.
+* `cuda` GPU mode using PyTorch with CUDA support. Requires any Nvidia GPU *(fast)*.
+* `tensorrt` GPU mode using vs-mlrt with TensorRT support. Requires an Nvidia RTX GPU. On the first run, this mode will automatically build an engine, which may take a few minutes. Changing interpolation, temp_window, or input dimensions will trigger rebuilding, but previously build engines are stored *(very fast/low vram)*.
+
+__*`num_streams`* (optional)__  
+Number of parallel TensorRT streams. For high end GPUs higher can be a bit faster, but requires more VRAM. Only affects the TensorRT backend.
+
+__*`engine_folder`* (optional)__  
+Optional path to the TensorRT engine storage location. By default engines are stored in `vs_undistort/engines`. Only affects the TensorRT backend.
+
+> [!TIP]
+> * If you see jumps/hitches between temporal windows, you can crossfade them with [vs_tiletools](https://github.com/pifroggi/vs_tiletools?tab=readme-ov-file#fix-jumpshitches-on-chunktemporal-window-based-filters-via-crossfading).
+> * If you have an undistorted reference clip, you can also try to align to it with [vs_align](https://github.com/pifroggi/vs_align).
+
+<br />
+
+## Benchmarks
+
+| Hardware | Resolution | TensorRT | CUDA
+|----------|------------|----------|---------
+| RTX 4090 | 720x480    | ~35 fps  | ~6.5 fps
+| RTX 4090 | 1440x1080  | ~7.5 fps | ~1.5 fps
+| RTX 4090 | 2880x2160  | ~2 fps   | ~0.5 fps

vs_undistort-2.2.0/pyproject.toml

@@ -0,0 +1,40 @@
+[build-system]
+requires = ["hatchling>=1.27.0"]
+build-backend = "hatchling.build"
+
+[project]
+name = "vs_undistort"
+version = "2.2.0"
+description = "Vapoursynth function to remove video distortions, turbulence, wobble, warp, heat haze, or similar."
+readme = "README.md"
+requires-python = ">=3.12"
+authors = [{ name = "pifroggi" }]
+
+keywords = [
+	"vapoursynth",
+	"video",
+	"distortion",
+	"turbulence",
+	"heat haze",
+	"wobble",
+	"rectification",
+	"filtering",
+	"restoration"
+]
+
+classifiers = [
+	"Programming Language :: Python :: 3",
+	"Topic :: Multimedia :: Video"
+]
+
+dependencies = [
+	"vapoursynth>=74",
+    "vapoursynth-mlrt-trt>=16.0",
+    "tensorrt-cu13>=11.1.0.106",
+    "numpy"
+]
+
+[project.urls]
+Homepage = "https://github.com/pifroggi/vs_undistort"
+Repository = "https://github.com/pifroggi/vs_undistort"
+Issues = "https://github.com/pifroggi/vs_undistort/issues"

vs_undistort-2.2.0/vs_undistort/__init__.py

@@ -0,0 +1,9 @@
+import sys
+from types import ModuleType
+from .vs_undistort import vs_undistort
+
+class _CallableModule(ModuleType):
+    __call__ = staticmethod(vs_undistort)
+
+sys.modules[__name__].__class__ = _CallableModule
+__all__ = ["vs_undistort"]

vs_undistort-2.2.0/vs_undistort/models/UNet3d_TMT_arch.py

@@ -0,0 +1,319 @@
+# Code from "Turbulence Mitigation Transformer" https://github.com/xg416/TMT
+
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+class LayerNorm3D(nn.Module):
+    def __init__(self, dim, bias=True):
+        super().__init__()
+        self.LN = nn.LayerNorm(dim, elementwise_affine=bias)
+
+    def forward(self, x):
+        # x: (B, C, T, H, W)
+        x = x.permute(0, 2, 3, 4, 1)  # (B, T, H, W, C)
+        x = self.LN(x)                # normalize over last dim (C)
+        x = x.permute(0, 4, 1, 2, 3)  # (B, C, T, H, W)
+        return x
+        
+class DetiltUNet3DS(nn.Module):
+    # __                            __
+    #  1|__   ________________   __|1
+    #     2|__  ____________  __|2
+    #        3|__  ______  __|3
+    #           4|__ __ __|4
+    # The convolution operations on either side are residual subject to 1*1 Convolution for channel homogeneity
+    def __init__(self, num_channels=3, feat_channels=[64, 256, 256, 512], norm='BN', conv_type='normal', residual='conv', interpolation='bilinear'):
+        # residual: conv for residual input x through 1*1 conv across every layer for downsampling, None for removal of residuals
+        super(DetiltUNet3DS, self).__init__()
+
+        # Encoder downsamplers
+        self.pool1 = nn.MaxPool3d((1, 2, 2))
+        self.pool2 = nn.MaxPool3d((1, 2, 2))
+        self.pool3 = nn.MaxPool3d((1, 2, 2))
+        self.pool4 = nn.MaxPool3d((1, 2, 2))
+
+        # Encoder convolutions
+        if norm =='BN': 
+            norm3d = nn.BatchNorm3d 
+        elif norm == 'LN':
+            norm3d = LayerNorm3D
+            
+        self.conv_blk1 = Conv3D_Block(num_channels, feat_channels[0], kernel=7, stride=1, padding=3, norm=norm3d, conv_type='normal', residual=residual, first=True)
+        self.conv_blk2 = Conv3D_Block(feat_channels[0], feat_channels[1], kernel=7, stride=1, padding=3, norm=norm3d, conv_type=conv_type, residual=residual)
+        self.conv_blk3 = Conv3D_Block(feat_channels[1], feat_channels[2], kernel=5, stride=1, padding=2, norm=norm3d, conv_type=conv_type, residual=residual)
+        self.conv_blk4 = Conv3D_Block(feat_channels[2], feat_channels[3], kernel=3, stride=1, padding=1, norm=norm3d, conv_type=conv_type, residual=residual)
+
+        # Decoder convolutions
+        self.dec_conv_blk3 = Conv3D_Block(2 * feat_channels[2], feat_channels[2], norm=norm3d, conv_type=conv_type, residual=residual)
+        self.dec_conv_blk2 = Conv3D_Block(2 * feat_channels[1], feat_channels[1], norm=norm3d, conv_type=conv_type, residual=residual)
+        self.dec_conv_blk1 = Conv3D_Block(2 * feat_channels[0], feat_channels[0], norm=norm3d, conv_type=conv_type, residual=residual)
+
+        # Decoder upsamplers
+        self.deconv_blk3 = Deconv3D_Block(feat_channels[3], feat_channels[2], conv_type=conv_type)
+        self.deconv_blk2 = Deconv3D_Block(feat_channels[2], feat_channels[1], conv_type=conv_type)
+        self.deconv_blk1 = Deconv3D_Block(feat_channels[1], feat_channels[0], conv_type=conv_type)
+
+        # Final 1*1 Conv Segmentation map
+        self.out_conv3 = nn.Conv3d(feat_channels[2], 2, kernel_size=3, stride=1, padding=1, bias=True)
+        self.out_conv2 = nn.Conv3d(feat_channels[1], 2, kernel_size=3, stride=1, padding=1, bias=True)
+        self.out_conv1 = nn.Conv3d(feat_channels[0], 2, kernel_size=3, stride=1, padding=1, bias=True)
+        
+        self.upsample3 = nn.Upsample(size=None, scale_factor=(1,4,4), mode='trilinear', align_corners=True)
+        self.upsample2 = nn.Upsample(size=None, scale_factor=(1,2,2), mode='trilinear', align_corners=True)
+
+        # TiltWarp settings
+        self.interpolation = interpolation
+        if interpolation == 'bilinear':
+            self.interp_mode = 'bilinear'
+            self.padding_mode = 'zeros'
+        elif interpolation == 'bicubic':
+            self.interp_mode = 'bicubic'
+            self.padding_mode = 'reflection'
+
+    def forward(self, x, scales=[True, True, True]):
+        # Encoder part
+        xin = x.permute(0,2,1,3,4)
+        x1 = self.conv_blk1(xin)
+
+        x_low1 = self.pool1(x1)
+        x2 = self.conv_blk2(x_low1)
+
+        x_low2 = self.pool2(x2)
+        x3 = self.conv_blk3(x_low2)
+
+        x_low3 = self.pool3(x3)
+        base = self.conv_blk4(x_low3)
+
+        # Decoder part
+
+        d3 = torch.cat([self.deconv_blk3(base), x3], dim=1)
+        d_high3 = self.dec_conv_blk3(d3)
+        d2 = torch.cat([self.deconv_blk2(d_high3), x2], dim=1)
+        d_high2 = self.dec_conv_blk2(d2)
+        d1 = torch.cat([self.deconv_blk1(d_high2), x1], dim=1)
+        d_high1 = self.dec_conv_blk1(d1)
+        
+        s3, s2, s1 = scales
+        flow3 = self.out_conv3(d_high3)
+        flow2 = self.out_conv2(d_high2)
+        flow1 = self.out_conv1(d_high1)
+
+        # keep grid_sample in fp32 for stability
+        orig_dtype = x.dtype
+        x_warp = x.float()
+
+        # only do warp if needed
+        if s3:
+            flow3_up = self.upsample3(flow3).float()
+            out_3 = TiltWarp(x_warp, flow3_up, interp_mode=self.interp_mode, padding_mode=self.padding_mode)
+        else:
+            out_3 = x_warp
+
+        if s2:
+            flow2_up = self.upsample2(flow2).float()
+            out_2 = TiltWarp(out_3, flow2_up, interp_mode=self.interp_mode, padding_mode=self.padding_mode)
+        else:
+            out_2 = out_3
+
+        if s1:
+            flow1 = flow1.float()
+            out = TiltWarp(out_2, flow1, interp_mode=self.interp_mode, padding_mode=self.padding_mode)
+        else:
+            out = out_2
+
+        # cast back
+        return out.to(orig_dtype)
+
+class UNet3D(nn.Module):
+    # __                            __
+    #  1|__   ________________   __|1
+    #     2|__  ____________  __|2
+    #        3|__  ______  __|3
+    #           4|__ __ __|4
+    # The convolution operations on either side are residual subject to 1*1 Convolution for channel homogeneity
+
+    def __init__(self, num_channels=3, feat_channels=[64, 256, 256, 512, 1024], norm='BN', out_channels=2, conv_type='normal', residual='conv'):
+        # residual: conv for residual input x through 1*1 conv across every layer for downsampling, None for removal of residuals
+
+        super(UNet3D, self).__init__()
+
+        # Encoder downsamplers
+        self.pool1 = nn.MaxPool3d((1, 2, 2))
+        self.pool2 = nn.MaxPool3d((1, 2, 2))
+        self.pool3 = nn.MaxPool3d((1, 2, 2))
+        self.pool4 = nn.MaxPool3d((1, 2, 2))
+
+        # Encoder convolutions
+        if norm =='BN': 
+            norm3d = nn.BatchNorm3d 
+        elif norm == 'LN':
+            norm3d = LayerNorm3D
+            
+        self.conv_blk1 = Conv3D_Block(num_channels, feat_channels[0], norm=norm3d, conv_type='normal', residual=residual)
+        self.conv_blk2 = Conv3D_Block(feat_channels[0], feat_channels[1], norm=norm3d, conv_type=conv_type, residual=residual)
+        self.conv_blk3 = Conv3D_Block(feat_channels[1], feat_channels[2], norm=norm3d, conv_type=conv_type, residual=residual)
+        self.conv_blk4 = Conv3D_Block(feat_channels[2], feat_channels[3], norm=norm3d, conv_type=conv_type, residual=residual)
+        self.conv_blk5 = Conv3D_Block(feat_channels[3], feat_channels[4], norm=norm3d,conv_type=conv_type, residual=residual)
+
+        # Decoder convolutions
+        self.dec_conv_blk4 = Conv3D_Block(2 * feat_channels[3], feat_channels[3], norm=norm3d, conv_type=conv_type, residual=residual)
+        self.dec_conv_blk3 = Conv3D_Block(2 * feat_channels[2], feat_channels[2], norm=norm3d, conv_type=conv_type, residual=residual)
+        self.dec_conv_blk2 = Conv3D_Block(2 * feat_channels[1], feat_channels[1], norm=norm3d, conv_type=conv_type, residual=residual)
+        self.dec_conv_blk1 = Conv3D_Block(2 * feat_channels[0], feat_channels[0], norm=norm3d, conv_type=conv_type, residual=residual)
+
+        # Decoder upsamplers
+        self.deconv_blk4 = Deconv3D_Block(feat_channels[4], feat_channels[3], conv_type=conv_type)
+        self.deconv_blk3 = Deconv3D_Block(feat_channels[3], feat_channels[2], conv_type=conv_type)
+        self.deconv_blk2 = Deconv3D_Block(feat_channels[2], feat_channels[1], conv_type=conv_type)
+        self.deconv_blk1 = Deconv3D_Block(feat_channels[1], feat_channels[0], conv_type=conv_type)
+
+        # Final output stage
+        self.out_conv = nn.Conv3d(feat_channels[0], out_channels, kernel_size=3, stride=1, padding=1, bias=False)
+        
+    def forward(self, x):
+        # Encoder part
+        x1 = self.conv_blk1(x)
+
+        x_low1 = self.pool1(x1)
+        x2 = self.conv_blk2(x_low1)
+
+        x_low2 = self.pool2(x2)
+        x3 = self.conv_blk3(x_low2)
+
+        x_low3 = self.pool3(x3)
+        x4 = self.conv_blk4(x_low3)
+
+        x_low4 = self.pool4(x4)
+        base = self.conv_blk5(x_low4)
+
+        # Decoder part
+        d4 = torch.cat([self.deconv_blk4(base), x4], dim=1)
+        d_high4 = self.dec_conv_blk4(d4)
+        d3 = torch.cat([self.deconv_blk3(d_high4), x3], dim=1)
+        d_high3 = self.dec_conv_blk3(d3)
+        d2 = torch.cat([self.deconv_blk2(d_high3), x2], dim=1)
+        d_high2 = self.dec_conv_blk2(d2)
+        d1 = torch.cat([self.deconv_blk1(d_high2), x1], dim=1)
+        d_high1 = self.dec_conv_blk1(d1)
+
+        # seg = self.tanh(self.one_conv(d_high1))
+        seg = self.out_conv(d_high1)
+        return seg
+
+class DWconv3D(nn.Module):
+    def __init__(self, in_channels, out_channels, kernel_size, stride=1, padding=0, dilation=1, bias=True):
+        super(DWconv3D, self).__init__()
+
+        self.conv = nn.Sequential(
+            nn.Conv3d(in_channels=in_channels, out_channels=in_channels, kernel_size=kernel_size, stride=stride, padding=padding, dilation=dilation, 
+                                groups=in_channels, bias=bias),
+            nn.Conv3d(in_channels=in_channels, out_channels=out_channels, kernel_size=1))
+
+    def forward(self, x):
+        return self.conv(x)
+
+class DWconvTFL(nn.Module):
+    def __init__(self, in_channels, out_channels, kernel_size, stride=1, padding=0, dilation=1, bias=True, n_frames=12):
+        super(DWconvTFL, self).__init__()
+
+        self.conv = nn.Sequential(
+                nn.Conv3d(in_channels=in_channels, out_channels=in_channels, kernel_size=(1,kernel_size,kernel_size), 
+                            stride=stride, padding=(0,padding,padding), dilation=dilation, groups=in_channels, bias=bias),
+                Rearrange('b c t h w -> b c h w t'),
+                nn.Linear(n_frames, n_frames),
+                Rearrange('b c h w t -> b c t h w'),
+                nn.Conv3d(in_channels=in_channels, out_channels=out_channels, kernel_size=1))
+                
+    def forward(self, x):
+        return self.conv(x)
+            
+class Conv3D_Block(nn.Module):
+    def __init__(self, inp_feat, out_feat, kernel=3, stride=1, padding=1, norm=LayerNorm3D, conv_type='dw', residual=None, first=False):
+        super(Conv3D_Block, self).__init__()
+        if conv_type == 'normal' or first:
+            conv3d = nn.Conv3d
+        elif conv_type == 'dw':
+            conv3d = DWconv3D
+
+        self.conv1 = nn.Sequential(
+            conv3d(inp_feat, out_feat, kernel_size=kernel, stride=stride, padding=padding, bias=True),
+            norm(out_feat),
+            nn.GELU())
+        if first:
+            self.conv2 = nn.Sequential(
+                DWconv3D(out_feat, out_feat, kernel_size=kernel, stride=stride, padding=padding, bias=True),
+                norm(out_feat),
+                nn.GELU())
+        else:      
+            self.conv2 = nn.Sequential(
+                conv3d(out_feat, out_feat, kernel_size=kernel, stride=stride, padding=padding, bias=True),
+                norm(out_feat),
+                nn.GELU())
+        self.residual = residual
+        if self.residual is not None:
+            self.residual_upsampler = conv3d(inp_feat, out_feat, kernel_size=1, bias=False)
+
+    def forward(self, x):
+        res = x
+        if not self.residual:
+            return self.conv2(self.conv1(x))
+        else:
+            return self.conv2(self.conv1(x)) + self.residual_upsampler(res)
+
+
+class Deconv3D_Block(nn.Module):
+    def __init__(self, inp_feat, out_feat, kernel=3, stride=2, padding=1, norm=LayerNorm3D, conv_type='dw'):
+        super(Deconv3D_Block, self).__init__()
+        if conv_type == 'normal':
+            self.deconv = nn.Sequential(
+                norm(inp_feat),
+                nn.ConvTranspose3d(inp_feat, out_feat, kernel_size=(kernel, kernel, kernel),
+                                stride=(1, stride, stride), padding=(padding, padding, padding), 
+                                output_padding=(0,1,1), bias=True),
+                nn.GELU())
+        if conv_type == 'dw':
+            self.deconv = nn.Sequential(
+                norm(inp_feat),
+                nn.ConvTranspose3d(inp_feat, inp_feat, kernel_size=(kernel, kernel, kernel),
+                                stride=(1, stride, stride), padding=(padding, padding, padding), 
+                                output_padding=(0,1,1), groups=inp_feat, bias=True),
+                nn.Conv3d(in_channels=inp_feat, out_channels=out_feat, kernel_size=1),
+                nn.GELU())
+
+    def forward(self, x):
+        return self.deconv(x)
+
+def TiltWarp(x, flow, interp_mode='bilinear', padding_mode='zeros', align_corners=True, use_pad_mask=False):
+    """
+    Args:
+        x (Tensor): Tensor with size (b, n, c, h, w) -> (b*n, c, h, w).
+        flow (Tensor): Tensor with size (b, 2, n, h, w) -> (b*n, h, w, 2), normal value.
+        interp_mode (str): 'nearest' or 'bilinear' or 'nearest4'. Default: 'bilinear'.
+        padding_mode (str): 'zeros' or 'border' or 'reflection'.
+            Default: 'zeros'.
+        align_corners (bool): Before pytorch 1.3, the default value is
+            align_corners=True. After pytorch 1.3, the default value is
+            align_corners=False. Here, we use the True as default.
+        use_pad_mask (bool): only used for PWCNet, x is first padded with ones along the channel dimension.
+            The mask is generated according to the grid_sample results of the padded dimension.
+    Returns:
+        Tensor: Warped image or feature map.
+    """
+    B, N, C, H, W = x.shape
+    x_flat = x.reshape(B * N, C, H, W)  # merge batch and time so grid_sample can work as (BN, C, H, W)
+    flow_flat = (flow.permute(0,2,3,4,1).reshape(B * N, H, W, 2))  # (BN, H, W, 2)
+    # create mesh grid
+    grid_y, grid_x = torch.meshgrid(torch.arange(H, dtype=x_flat.dtype, device=x_flat.device),
+                                    torch.arange(W, dtype=x_flat.dtype, device=x_flat.device), indexing='ij')  # H, W are taken from the tensor's shape, so ONNX can track them.
+    grid = torch.stack((grid_x, grid_y), dim=2)                    # (H, W, 2)
+    vgrid = grid.unsqueeze(0) + flow_flat                          # (BN, H, W, 2)
+
+    vgrid_x = 2.0 * vgrid[..., 0] / (W - 1) - 1.0
+    vgrid_y = 2.0 * vgrid[..., 1] / (H - 1) - 1.0
+    vgrid_scaled = torch.stack((vgrid_x, vgrid_y), dim=-1)         # (BN, H, W, 2)
+    output_flat = F.grid_sample(x_flat, vgrid_scaled, mode=interp_mode, padding_mode=padding_mode, align_corners=align_corners)
+    output = output_flat.reshape(B, N, C, H, W)
+    return output
+

vs_undistort-2.2.0/vs_undistort/utils.py

@@ -0,0 +1,146 @@
+
+# Script by pifroggi https://github.com/pifroggi/vs_undistort
+# or tepete and pifroggi on Discord
+
+import math
+import vapoursynth as vs
+
+core = vs.core
+
+
+def expression(clips, expr, format=None):
+    # backend for basic exprs supported by std.Expr
+    if hasattr(core, "akarin"):
+        return core.akarin.Expr(clips, expr, format=format)
+    else:
+        return core.std.Expr(clips, expr, format=format)
+
+
+def get_window(clip, temp_window):
+    # how many pad frames to reach a multiple of temp_window
+    num_frames = clip.num_frames
+    pad = (-num_frames) % temp_window
+
+    if pad:
+        # pad black frames if too short
+        pad_clip = core.std.BlankClip(clip=clip, length=pad)
+        clip     = core.std.Splice([clip, pad_clip])
+
+    # offset_clips[i] contains frames i, i+temp_window, i+2*temp_window, ...
+    return [core.std.SelectEvery(clip[i:], cycle=temp_window, offsets=[0]) for i in range(temp_window)]
+
+
+def get_tiles(clip_w, clip_h, tiles, overlap=0):
+    # calculate tile size and choose the most square layout
+    if tiles not in (1, 2, 4, 6, 8, 12, 16, 24, 32):
+        raise ValueError("vs_undistort: Tiles must be 1, 2, 4, 6, 8, 12, 16, 24, or 32.")
+
+    layouts = {
+        1: [(1, 1)],
+        2: [(2, 1), (1, 2)],
+        4: [(4, 1), (2, 2), (1, 4)],
+        6: [(6, 1), (3, 2), (2, 3), (1, 6)],
+        8: [(8, 1), (4, 2), (2, 4), (1, 8)],
+        12: [(12, 1), (6, 2), (4, 3), (3, 4), (2, 6), (1, 12)],
+        16: [(16, 1), (8, 2), (4, 4), (2, 8), (1, 16)],
+        24: [(24, 1), (12, 2), (8, 3), (6, 4), (4, 6), (3, 8), (2, 12), (1, 24)],
+        32: [(32, 1), (16, 2), (8, 4), (4, 8), (2, 16), (1, 32)],
+    }[tiles]
+
+    def _tile_size(layout):
+        cols, rows = layout
+        tile_w = math.ceil(math.ceil((clip_w + 2 * overlap * (cols - 1)) / cols) / 16) * 16
+        tile_h = math.ceil(math.ceil((clip_h + 2 * overlap * (rows - 1)) / rows) / 16) * 16
+        return tile_w, tile_h
+
+    def _layout_valid(layout):
+        # tiles must have a positive non overlapped stride
+        cols, rows = layout
+        tile_w, tile_h = _tile_size(layout)
+        if cols > 1 and tile_w <= 2 * overlap:
+            return False
+        if rows > 1 and tile_h <= 2 * overlap:
+            return False
+        return True
+
+    def _score(layout):
+        tile_w, tile_h = _tile_size(layout)
+        cols, rows = layout
+        tile_aspect = tile_w / tile_h
+        square_error = abs(math.log(tile_aspect))
+        orientation_penalty = rows if clip_w >= clip_h else cols
+        balance_penalty = abs(cols - rows)
+        return (square_error, orientation_penalty, balance_penalty)
+
+    valid_layouts = [layout for layout in layouts if _layout_valid(layout)]
+    if not valid_layouts:
+        raise ValueError("vs_undistort: Clip dimensions are too small for current tile amount. Reduce tiles or overlap.")
+
+    cols, rows = min(valid_layouts, key=_score)
+    tile_w, tile_h = _tile_size((cols, rows))
+    if tile_w > 2 * tile_h or tile_h > 2 * tile_w:
+        raise ValueError("vs_undistort: The current tile amount produces tiles that are too elongated. Try a different tile amount.")
+
+    return tile_w, tile_h
+
+
+def get_tile_positions(h, w, s_h, s_w, overlap=0):
+    if overlap is None:
+        overlap = 0
+    
+    def _axis_positions(size, patch, overlap):
+        
+        # single patch or smaller image = no tiling / overlap needed
+        if patch >= size:
+            return [0]
+
+        # overlap is removed from both sides of the tile stride to match vsmlrt
+        step      = patch - 2 * overlap
+        max_start = size - patch
+
+        if step <= 0 or max_start <= 0:
+            return [0]
+
+        positions = [0]
+        while positions[-1] < max_start:
+            positions.append(min(positions[-1] + step, max_start))
+
+        return positions
+
+    hpos = _axis_positions(h, s_h, overlap)
+    wpos = _axis_positions(w, s_w, overlap)
+    return hpos, wpos
+
+
+def inference_tiled(input_blk, model_tilt, patch_height, patch_width, overlap=0, scales=[True, True, True], tile_device=None):
+    import torch
+    tile_device   = torch.device("cpu") if tile_device is None else tile_device
+    model_device  = next(model_tilt.parameters()).device
+    b, l, c, h, w = input_blk.shape
+    
+    hpos, wpos = get_tile_positions(h, w, patch_height, patch_width, overlap)
+    out_spaces = torch.empty_like(input_blk, device=tile_device)                           # (B, L, C, H, W)
+
+    for hi in hpos:
+        y_crop_start = 0 if hi == 0 else overlap
+        y_crop_end   = 0 if hi == h - patch_height else overlap
+        for wi in wpos:
+            x_crop_start = 0 if wi == 0 else overlap
+            x_crop_end   = 0 if wi == w - patch_width else overlap
+            inp = input_blk[..., hi:hi + patch_height, wi:wi + patch_width]                # (B, L, C, ph, pw)
+            
+            # move only the tile to the model device
+            if inp.device != model_device:
+                non_blocking = (inp.device.type == "cpu" and model_device.type == "cuda" and inp.is_pinned())
+                inp = inp.to(model_device, non_blocking=non_blocking)
+            
+            rectified = model_tilt(inp, scales=scales)
+            hs, ws = rectified.shape[-2:]
+            
+            # move tile result back to tile device
+            if rectified.device != tile_device:
+                rectified = rectified.to(tile_device)
+            
+            out_spaces[..., hi + y_crop_start:hi + hs - y_crop_end, wi + x_crop_start:wi + ws - x_crop_end].copy_(rectified[..., y_crop_start:hs - y_crop_end, x_crop_start:ws - x_crop_end])
+
+    return out_spaces

vs_undistort-2.2.0/vs_undistort/vs_undistort.py

@@ -0,0 +1,479 @@
+# Script to run "Turbulence Mitigation Transformer" https://github.com/xg416/TMT
+
+# Vapoursynth Implementation by pifroggi https://github.com/pifroggi/vs_undistort
+# or tepete and pifroggi on Discord
+
+import os
+import re
+import shutil
+import logging
+import subprocess
+import vapoursynth as vs
+from pathlib import Path
+from .utils import get_window, get_tiles, expression
+
+core = vs.core
+
+
+def _pytorch(clip, temp_window=10, tiles=1, overlap=8, interpolation="bicubic", device="cuda"):
+    import threading
+    import numpy as np
+    from collections import OrderedDict
+
+    if device == "cpu":
+        try:
+            import torch
+        except ImportError:
+            raise RuntimeError("vs_undistort: The CPU/CUDA backends require PyTorch. Please install it from https://pytorch.org/ or choose a different backend. For the CUDA backend specifically, install a version of PyTorch with CUDA support.") from None
+
+    if device == "cuda":
+        try:
+            import torch
+        except ImportError:
+            raise RuntimeError("vs_undistort: The CUDA backend requires PyTorch with CUDA. Please install a version of PyTorch with CUDA support from https://pytorch.org/ or choose a different backend.") from None
+        if not torch.cuda.is_available():
+            raise RuntimeError("vs_undistort: The CUDA backend requires PyTorch with CUDA, but the installed version has no CUDA support. Please upgrade to a version with CUDA support from https://pytorch.org/ or choose a different backend.")
+
+    from .utils import inference_tiled
+    from .models.UNet3d_TMT_arch import DetiltUNet3DS
+    os.environ["CUDA_MODULE_LOADING"] = "LAZY"
+
+    def _frames_to_tensor(frames, device, tile_device, fp16=False):
+        temp_window = len(frames)
+        h, w = frames[0].height, frames[0].width
+        num_planes = frames[0].format.num_planes
+        dtype = np.float16 if fp16 == True else np.float32
+        arr = np.empty((temp_window, num_planes, h, w), dtype=dtype)
+        for i, fr in enumerate(frames):
+            for p in range(num_planes):
+                arr[i, p, :, :] = np.asarray(fr[p])
+        tensor = torch.from_numpy(arr)
+        if tile_device.type != "cpu":
+            return tensor.to(tile_device).clamp_(0, 1).unsqueeze(0)
+        if device.type == "cuda":
+            tensor = tensor.pin_memory()
+        return tensor.clamp_(0, 1).unsqueeze(0)          
+    
+    def _tensor_to_frame(tensor, frame):
+        tensor_np = tensor.detach().cpu().numpy()
+        for p in range(frame.format.num_planes):
+            frame_arr = np.asarray(frame[p])
+            np.copyto(frame_arr, tensor_np[:, :, p])
+    
+    def _load_model(device, interpolation, fp16=False):
+        model_tilt = DetiltUNet3DS(norm="LN", residual="pool", conv_type="dw", interpolation=interpolation)
+        current_folder = os.path.dirname(os.path.abspath(__file__))
+        path_tilt = os.path.join(current_folder, "models", "UNet3d_TMT.pth")
+        ckpt_tilt = torch.load(path_tilt, map_location=device, weights_only=True)
+        model_tilt.load_state_dict(ckpt_tilt["state_dict"] if "state_dict" in ckpt_tilt else ckpt_tilt)
+        model_tilt.to(device)
+        if fp16:
+            model_tilt.half()
+        model_tilt.eval()
+        return model_tilt
+    
+    def _pytorch_inference(n, f):
+        with torch.inference_mode():
+            out = f[0].copy()
+            tensor = _frames_to_tensor(f[1:], device, tile_device, fp16=fp16)          # (1, T, C, H, W)
+            tensor = inference_tiled(tensor, model_tilt, tile_h, tile_w, overlap=overlap, scales=[True, True, True], tile_device=tile_device)
+            stacked_tensor = tensor[0].permute(2, 0, 3, 1).contiguous().flatten(1, 2)  # (H, T * W, C), output one large frame to let vs handle the caching
+            _tensor_to_frame(stacked_tensor, out)
+            return out
+    
+    # checks
+    if not isinstance(clip, vs.VideoNode):
+        raise TypeError("vs_undistort: Clip must be a vapoursynth clip.")
+    if clip.format.id == vs.PresetVideoFormat.NONE or clip.width == 0 or clip.height == 0:
+        raise TypeError("vs_undistort: Clip must have constant format and dimensions.")
+    if clip.format.color_family != vs.RGB:
+        raise ValueError("vs_undistort: Clip must be in RGB format.")
+    if not isinstance(temp_window, int) or isinstance(temp_window, bool):
+        raise TypeError("vs_undistort: Temporal window length must be an integer.")
+    if temp_window < 2:
+        raise ValueError("vs_undistort: Temporal window length must be at least 2.")
+    if not isinstance(tiles, int) or isinstance(tiles, bool):
+        raise TypeError("vs_undistort: Tiles must be an integer.")
+    if not isinstance(overlap, int) or isinstance(overlap, bool):
+        raise TypeError("vs_undistort: Overlap must be an integer.")
+    if overlap < 0:
+        raise ValueError("vs_undistort: Overlap can not be negative.")
+    if interpolation not in ["bilinear", "bicubic"]:
+        raise ValueError("vs_undistort: Warp interpolation mode must be 'bilinear' or 'bicubic'.")
+    
+    orig_format = clip.format.id
+    device      = torch.device(device)
+    fp16        = device.type == "cuda" and torch.cuda.get_device_capability()[0] >= 7
+    req_format  = vs.RGBH if fp16 else vs.RGBS
+    model_tilt  = _load_model(device, interpolation, fp16=fp16)
+    
+    # decide tile size
+    tile_w, tile_h = get_tiles(clip_w=clip.width, clip_h=clip.height, tiles=tiles, overlap=overlap)
+
+    # pad if tile is larger than clip
+    pad_w = max(0, tile_w - clip.width)
+    pad_h = max(0, tile_h - clip.height)
+    if pad_w or pad_h:
+        clip = core.std.AddBorders(clip, right=pad_w, bottom=pad_h)
+    
+    width       = clip.width
+    height      = clip.height
+    tile_device = device if tiles == 1 else torch.device("cpu")
+
+    # convert inputs to needed precision
+    if clip.format.id != req_format:
+        clip = core.resize.Point(clip, format=req_format)
+    if tile_w < 128 or tile_h < 128:
+        raise ValueError("vs_undistort: Tile size must be at least 128 x 128. Reduce tiles.")
+    if tiles > 1 and overlap > min(tile_w, tile_h) // 2:
+        raise ValueError("vs_undistort: Overlap can not be larger than half of tile size. Reduce overlap.")
+    
+    # inference
+    offset_clips   = get_window(clip, temp_window)
+    out_shape      = core.std.BlankClip(clip=offset_clips[0], width=width * temp_window, height=height, keep=True)
+    stacked_clip   = core.std.ModifyFrame(out_shape, clips=[out_shape, *offset_clips], selector=_pytorch_inference)
+    offset_clips   = [core.std.Crop(stacked_clip, left=i * width, right=(temp_window - 1 - i) * width) for i in range(temp_window)]
+    unstacked_clip = core.std.Interleave(offset_clips)
+    
+    if unstacked_clip.format.id != orig_format:
+        unstacked_clip = core.resize.Point(unstacked_clip, format=orig_format)
+    if pad_w or pad_h:
+        unstacked_clip = core.std.Crop(unstacked_clip, right=pad_w, bottom=pad_h)
+    if unstacked_clip.num_frames != clip.num_frames:
+        unstacked_clip = core.std.Trim(unstacked_clip, last=clip.num_frames - 1)
+    
+    return core.std.CopyFrameProps(unstacked_clip, clip)
+
+
+def _get_builder(plugin_path, trt_version, cuda_major):
+    # finds compatible tensorrt engine builders
+    exe_name = "trtexec.exe" if os.name == "nt" else "trtexec"
+    builders = []
+    errors   = []
+    
+    # check for python tensorrt
+    try:
+        import tensorrt
+        package_version = list(map(int, tensorrt.__version__.split(".")[:3]))
+        if package_version == trt_version:
+            builders.append(["python", tensorrt])
+        else:
+            errors.append(f"Python TensorRT: Wrong version {'.'.join(map(str, package_version))}")
+    except ImportError:
+        errors.append("Python TensorRT: Not found.")
+    except Exception:
+        errors.append("Python TensorRT: Found but failed to check version.")
+    
+    # check for bundled trtexec
+    bundled_trtexec = Path(plugin_path) / "vsmlrt-cuda" / exe_name
+    if bundled_trtexec.is_file() and os.access(str(bundled_trtexec), os.X_OK):
+        builders.append(["trtexec", bundled_trtexec])
+    else:
+        errors.append(f"Bundled trtexec: Not found.")
+
+    # check for system trtexec
+    system_trtexec = shutil.which("trtexec")
+    if system_trtexec is not None:
+        try:
+            trtexec_path = Path(system_trtexec)
+            help_output  = subprocess.run([str(trtexec_path), "--help"], check=False, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, encoding="locale", errors="replace")
+            help_output  = f"{help_output.stdout}\n{help_output.stderr}"
+            
+            trtexec_version = None
+            trtexec_version = re.search(r"\[TensorRT v(\d+)\]", help_output)
+            if trtexec_version is None:
+                raise RuntimeError("vs_undistort: Internal Error: Regex failed to find the version.")
+
+            trtexec_version = int(trtexec_version.group(1))
+            trtexec_version = [trtexec_version // 10000, (trtexec_version % 10000) // 100, trtexec_version % 100]
+            if trtexec_version == trt_version:
+                builders.append(["trtexec", trtexec_path])
+            else:
+                errors.append(f"System trtexec: Wrong version {'.'.join(map(str, trtexec_version))}")
+        except Exception:
+            errors.append("System trtexec: Found but failed to check version.")
+    else:
+        errors.append("System trtexec: Not found.")
+    
+    # return first compatible builder
+    if builders:
+        return builders[0]
+    
+    errors = "\n".join(f"{builder}" for builder in errors)
+    raise FileNotFoundError(f"vs_undistort: No compatible TensorRT engine builder found. Please install the python package 'tensorrt' or install trtexec. The required TensorRT version is {'.'.join(map(str, trt_version))}. The required CUDA version is {cuda_major}.\n{errors}")
+
+
+def _build_engine_trtexec(onnx_path, engine_path, temp_window, engine_w, engine_h, interpolation, trt_version, trtexec_path):
+    # build engine using trtexec, supports trt 10 and 11
+
+    # settings
+    opt_shapes = f"input:1x{temp_window * 3}x{engine_h}x{engine_w}"
+    io_formats = f"fp16:chw" if trt_version[0] < 11 else "chw"
+    cmd = [
+        str(trtexec_path),
+        *(["--stronglyTyped"] if trt_version[0] < 11 else []),
+        *(["--markDebug=grid_sampler,grid_sampler_1"] if interpolation == "bicubic" else []),  # part of gridsample bicubic workaround
+        "--skipInference",
+        "--memPoolSize=workspace:6144",
+        "--builderOptimizationLevel=3",
+        f"--inputIOFormats={io_formats}",
+        f"--outputIOFormats={io_formats}",
+        f"--onnx={onnx_path}",
+        f"--saveEngine={engine_path}",
+        f"--optShapes={opt_shapes}",
+    ]
+
+    # build
+    try:
+        result = subprocess.run(cmd, check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, encoding="locale", errors="replace")
+    except subprocess.CalledProcessError as e:
+        msg = (
+            "vs_undistort: Internal Error: trtexec failed while building the TensorRT engine.\n"
+            f"  Command: {' '.join(cmd)}\n"
+            f"  Return code: {e.returncode}\n"
+        )
+        if e.stdout:
+            msg += f"\n=== trtexec stdout ===\n{e.stdout}"
+        if e.stderr:
+            msg += f"\n=== trtexec stderr ===\n{e.stderr}"
+        raise RuntimeError(msg) from e
+
+
+def _build_engine_python(onnx_path, engine_path, temp_window, engine_w, engine_h, interpolation, trt_package):
+    # build engine using tensorrt python package, supports only trt 11 because of vapoursynth-mlrt-trt
+    trt = trt_package
+
+    # custom logger for errors
+    class _TrtLogger(trt.ILogger):
+        def __init__(self):
+            trt.ILogger.__init__(self)
+            self.messages = []
+            self.fatal    = False
+        def log(self, severity, msg):
+            if severity <= trt.Logger.WARNING:
+                self.messages.append((severity, msg))
+                if self.fatal:
+                    logging.critical(f"  [{severity}] {msg}")
+                elif severity == trt.Logger.INTERNAL_ERROR:  # print fatal errors immediately because python may not get control back
+                    self.fatal = True
+                    log = "\n".join(f"  [{log_severity}] {log_msg}" for log_severity, log_msg in self.messages)
+                    logging.critical(f"vs_undistort: Internal Error: TensorRT failed while building the TensorRT engine.\n=== TensorRT log ===\n{log}")
+        def get_log(self):
+            return "\n".join(f"  [{severity}] {msg}" for severity, msg in self.messages)
+
+    # initialize trt and load model
+    logger  = _TrtLogger()
+    builder = trt.Builder(logger)
+    network = builder.create_network()
+    config  = builder.create_builder_config()
+    parser  = trt.OnnxParser(network, logger)
+    if not parser.parse_from_file(str(onnx_path)):
+        errors = "\n".join(f"  {parser.get_error(i)}" for i in range(parser.num_errors))
+        raise RuntimeError(f"vs_undistort: Internal Error: TensorRT failed while parsing the ONNX model.\n{errors}")
+    
+    # mark debug
+    if interpolation == "bicubic":
+        for layer_index in range(network.num_layers):
+            layer = network.get_layer(layer_index)
+            for output_index in range(layer.num_outputs):
+                tensor = layer.get_output(output_index)
+                if tensor is not None and tensor.name in ("grid_sampler", "grid_sampler_1"):  # part of gridsample bicubic workaround
+                    network.mark_debug(tensor)                         
+    
+    # settings
+    opt_shapes = (1, temp_window * 3, engine_h, engine_w)                                                             # optShapes
+    network.get_input(0).allowed_formats = network.get_output(0).allowed_formats = 1 << int(trt.TensorFormat.LINEAR)  # IOFormats:chw
+    config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 6144 << 20)                                            # workspace
+    config.builder_optimization_level = 3                                                                             # builderOptimizationLevel
+
+    # build
+    profile = builder.create_optimization_profile()
+    profile.set_shape(network.get_input(0).name, opt_shapes, opt_shapes, opt_shapes)
+    config.add_optimization_profile(profile)
+    engine  = builder.build_serialized_network(network, config)
+    if engine is None:
+        log = logger.get_log()
+        msg = "vs_undistort: Internal Error: TensorRT failed while building the TensorRT engine."
+        if log:
+            msg += f"\n=== TensorRT log ===\n{log}"
+        raise RuntimeError(msg)
+    
+    # save engine
+    with open(engine_path, "wb") as f:
+        f.write(engine)
+
+
+def _get_engine(onnx_path, engine_dir, temp_window, engine_w, engine_h, interpolation, force_rebuild=False) -> str:
+    # get path to tensorrt engine
+    os.makedirs(engine_dir, exist_ok=True)  # create engine folder if needed
+    engine_name = f"undistort_{interpolation}_t{temp_window}_h{engine_h}_w{engine_w}_fp16.engine"
+    engine_path = os.path.join(engine_dir, engine_name)
+    
+    # check plugin version
+    try:
+        info = core.trt.Version()
+    except Exception as e:
+        raise RuntimeError("vs_undistort: Please install a version of vs-mlrt with TensorRT support or choose a different backend.") from e
+    
+    # if engine file exist, return it
+    if not force_rebuild and os.path.isfile(engine_path) and os.path.getsize(engine_path) >= 512:
+        return engine_path
+    
+    # get plugin info
+    plugin_path = os.path.dirname(info["path"].decode(errors="ignore"))
+    trt_version = int(info["tensorrt_version"].decode(errors="ignore"))
+    trt_version = [trt_version // 10000, (trt_version % 10000) // 100, trt_version % 100]
+    cuda_major  = int(info["cuda_runtime_version"].decode(errors="ignore")) // 1000
+    
+    # build new engine
+    logging.warning("vs_undistort: Building new TensorRT engine for interpolation='%s' with temp_window=%d, width=%d, and height=%d. This may take a few minutes.", interpolation, temp_window, engine_w, engine_h)
+    builder_info = _get_builder(plugin_path=plugin_path, trt_version=trt_version, cuda_major=cuda_major)
+    if builder_info[0] == "python":
+        _build_engine_python(onnx_path=onnx_path, engine_path=engine_path, temp_window=temp_window, engine_w=engine_w, engine_h=engine_h, interpolation=interpolation, trt_package=builder_info[1])
+    elif builder_info[0] == "trtexec":
+        _build_engine_trtexec(onnx_path=onnx_path, engine_path=engine_path, temp_window=temp_window, engine_w=engine_w, engine_h=engine_h, interpolation=interpolation, trt_version=trt_version, trtexec_path=builder_info[1])
+    else:
+        raise RuntimeError(f"vs_undistort: Internal Error: Unknown TensorRT engine builder: {builder_info[0]}")
+    logging.warning("vs_undistort: Engine building complete.")
+    return engine_path
+
+
+def _tensorrt_inference(input_clips, onnx_path, engine_dir, temp_window, tile_w, tile_h, overlap, tiles, interpolation, num_streams, flex_out_prop, force_rebuild=False):
+    engine_path = _get_engine(onnx_path=onnx_path, engine_dir=engine_dir, temp_window=temp_window, engine_w=tile_w, engine_h=tile_h, interpolation=interpolation, force_rebuild=force_rebuild)
+    model_args  = dict(engine_path=engine_path, num_streams=num_streams, flexible_output_prop=flex_out_prop, **(dict(tilesize=(tile_w, tile_h), overlap=(overlap, overlap)) if tiles > 1 else {}))
+
+    # try inference, rebuild engine if it fails
+    try:
+        out = core.trt.Model(input_clips, **model_args)
+    except vs.Error as e:
+        err_msg = str(e).lower()
+        serialization_keywords = ("serialize", "serialization", "deserialize", "deserialization")
+        if any(k in err_msg for k in serialization_keywords) and not force_rebuild:
+            logging.warning("vs_undistort: Engine loading failed. This may be due to a TensorRT or driver update. Rebuilding...")
+            model_args["engine_path"] = _get_engine(onnx_path=onnx_path, engine_dir=engine_dir, temp_window=temp_window, engine_w=tile_w, engine_h=tile_h, interpolation=interpolation, force_rebuild=True)
+            out = core.trt.Model(input_clips, **model_args)
+        else:
+            raise
+    return out
+
+
+def _tensorrt(clip, temp_window=10, tiles=1, overlap=8, interpolation="bicubic", num_streams=1, engine_folder=None):
+    
+    # checks
+    if not isinstance(clip, vs.VideoNode):
+        raise TypeError("vs_undistort: Clip must be a vapoursynth clip.")
+    if clip.format.id == vs.PresetVideoFormat.NONE or clip.width == 0 or clip.height == 0:
+        raise TypeError("vs_undistort: Clip must have constant format and dimensions.")
+    if clip.format.id not in [vs.RGBH]:
+        raise ValueError("vs_undistort: Clip must be in RGBH format for the TensorRT backend.")
+    if not isinstance(temp_window, int) or isinstance(temp_window, bool):
+        raise TypeError("vs_undistort: Temporal window length must be an integer.")
+    if temp_window < 2:
+        raise ValueError("vs_undistort: Temporal window length must be at least 2.")
+    if not isinstance(tiles, int) or isinstance(tiles, bool):
+        raise TypeError("vs_undistort: Tiles must be an integer.")
+    if not isinstance(overlap, int) or isinstance(overlap, bool):
+        raise TypeError("vs_undistort: Overlap must be an integer.")
+    if overlap < 0:
+        raise ValueError("vs_undistort: Overlap can not be negative.")
+    if not isinstance(num_streams, int) or isinstance(num_streams, bool):
+        raise TypeError("vs_undistort: Number of parallel TensorRT streams (num_streams) must be an integer.")
+    if num_streams < 1:
+        raise ValueError("vs_undistort: Number of parallel TensorRT streams (num_streams) must be at least 1.")
+    
+    # clamp
+    clip = expression(clip, expr=["x 0 max 1 min"])
+    
+    # decide which dimensions to build the engine for
+    tile_w, tile_h = get_tiles(clip_w=clip.width, clip_h=clip.height, tiles=tiles, overlap=overlap)
+
+    # pad if tile is larger than clip
+    pad_w = max(0, tile_w - clip.width)
+    pad_h = max(0, tile_h - clip.height)
+    if pad_w or pad_h:
+        clip = core.std.AddBorders(clip, right=pad_w, bottom=pad_h)
+
+    if tile_w < 128 or tile_h < 128:
+        raise ValueError("vs_undistort: Tile size must be at least 128 x 128. Reduce tiles.")
+    if tiles > 1 and overlap > min(tile_w, tile_h) // 2:
+        raise ValueError("vs_undistort: Overlap can not be larger than half of tile size. Reduce overlap.")
+
+    # make sure the extremely slow engine isn't build due to some tensorrt tactic limitations
+    cur_pixels = temp_window * tile_w * tile_h
+    max_pixels = 4194304  #2^22
+    if cur_pixels >= max_pixels:
+        raise ValueError(f"vs_undistort: temp_window * tile width * tile height must be smaller than {max_pixels} (currently {cur_pixels}). Increase tiles or reduce temp_window or overlap.")
+    
+    if interpolation == "bilinear":
+        model_name = "UNet3d_TMT_lin_op19_fp16.onnx"
+    elif interpolation == "bicubic":
+        model_name = "UNet3d_TMT_cub_op19_fp16.onnx"
+    else:
+        raise ValueError("vs_undistort: Warp interpolation mode must be 'bilinear' or 'bicubic'.")
+    
+    current_dir   = os.path.dirname(os.path.abspath(__file__))
+    engine_dir    = os.path.join(current_dir, "engines") if engine_folder is None else os.path.abspath(engine_folder)
+    onnx_path     = os.path.join(current_dir, "models", model_name)
+    flex_out_prop = "vs_undistort_output"
+    force_rebuild = False
+
+    # get inference window and do inference
+    input_clips   = get_window(clip, temp_window)
+    stacked_clips = _tensorrt_inference(
+        input_clips=input_clips,
+        onnx_path=onnx_path,
+        engine_dir=engine_dir,
+        temp_window=temp_window,
+        tile_w=tile_w,
+        tile_h=tile_h,
+        overlap=overlap,
+        tiles=tiles,
+        interpolation=interpolation,
+        num_streams=num_streams,
+        flex_out_prop=flex_out_prop,
+        force_rebuild=force_rebuild,
+    )
+    
+    # turn stacked output into normal clip
+    carrier_clip   = stacked_clips["clip"]
+    num_planes     = stacked_clips["num_planes"]
+    planes         = [carrier_clip.std.PropToClip(prop=f"{flex_out_prop}{i}")  for i in range(num_planes)]        # turn planes back into normal clips
+    grouped_clips  = [core.std.ShufflePlanes(planes[i:i+3], [0, 0, 0], vs.RGB) for i in range(0, num_planes, 3)]  # group every 3 planes back into RGB clips
+    unstacked_clip = core.std.Interleave(grouped_clips)                                                           # interleave to restore chronological order
+    
+    if pad_w or pad_h:
+        unstacked_clip = core.std.Crop(unstacked_clip, right=pad_w, bottom=pad_h)
+    if unstacked_clip.num_frames != clip.num_frames:
+        unstacked_clip = core.std.Trim(unstacked_clip, last=clip.num_frames - 1)
+    
+    return core.std.CopyFrameProps(unstacked_clip, clip)
+
+
+def vs_undistort(clip, temp_window=10, tiles=1, overlap=8, interpolation="bicubic", backend="tensorrt", num_streams=1, engine_folder=None):
+    """Removes distortions. Also known as atmospheric turbulence mitigation, warp stabilization, film shrink or VHS distortion fix, heat haze removal.
+
+    Args:
+        clip: Distorted clip. Must be in RGB format.
+        temp_window: Temporal window length. This is how many frames are grouped together and processed as a single chunk. Larger means
+            higher VRAM requirements, but better temporal averaging and slower distortions can be removed. If this is too small,
+            some distortions may not get removed, small jumps/hitches may be visible between windows and seams from tiling
+            may become more obvious.
+        tiles: Amount of tiles to split the frames into. A higher amount reduces VRAM requirements, but also worsens spatial averaging.
+            Default `tiles=1` uses the full frame.
+        overlap: Overlap from one tile to the next. Use if seams between tiles are visible.
+        interpolation: Interpolation mode used to warp the frames.  
+            - `bilinear` = More blurry.
+            - `bicubic` = No blur, but may oversharpen slightly.
+        backend: The backend used to run the model.
+            - `cpu` = CPU mode using PyTorch (very slow).
+            - `cuda` = GPU mode using PyTorch with CUDA support. Requires any Nvidia GPU (fast).
+            - `tensorrt` = GPU mode using vs-mlrt with TensorRT support. Requires an Nvidia RTX GPU (very fast and lower vram usage).
+        num_streams: Number of parallel TensorRT streams. For high end GPUs higher can be a bit faster, but requires more VRAM. Only affects the TensorRT backend.
+        engine_folder: Optional path to the TensorRT engine storage location. By default engines are stored in `vs_undistort/engines`. Only affects the TensorRT backend.
+    """
+    
+    if backend in ["cpu", "cuda"]:
+        return _pytorch(clip, temp_window=temp_window, tiles=tiles, overlap=overlap, interpolation=interpolation, device=backend)
+    if backend in ["tensorrt", "trt"]:
+        return _tensorrt(clip, temp_window=temp_window, tiles=tiles, overlap=overlap, interpolation=interpolation, num_streams=num_streams, engine_folder=engine_folder)
+    raise ValueError("vs_undistort: Backend must be 'cpu', 'cuda', or 'tensorrt'.")