TSPANv2

TSPANv2 (Temporal SPAN) is a custom VSR architecture, and is a continuation of TSPAN. This version focuses on improving the temporal component of the architecture. It provides a drastic improvement in overall temporal stability and quality! This arch has support for PyTorch, ONNX, and TensorRT.

This is the inference and ONNX conversion code. To train a model, you'll want to use traiNNer-redux with the TSPANv2 config and a video dataset (WIP!). For easier inference, try out Vapourkit. To make a video dataset, try my other tool, video destroyer.

Getting Started

1. Clone the repository:

   git clone https://github.com/Kim2091/TSPANv2

2. Install PyTorch with CUDA: Follow the instructions at pytorch.org.

3. Install required packages:

   pip install -r requirements.txt

Usage

You can use TSPANv2 through the included graphical user interface (GUI), Vapourkit, or the command line.

TensorRT

For high-performance inference, refer to the TensorRT guide.

GUI Usage

For an easy-to-use experience with PyTorch or ONNX models, launch the GUI:

python vsr_gui.py

Command-Line Usage

For more advanced control, you can use the command-line scripts.

Video upscaling (PyTorch):

python test_vsr.py --model_path pretrained_models/tspan.pth --input path/to/video.mp4 --output path/to/output.mp4

ONNX Video upscaling:

python test_onnx.py --model_path model.onnx --input path/to/video.mp4 --output path/to/output.mp4

Key arguments for test_vsr.py and test_onnx.py:

• --video_codec: Specify the video codec (e.g., libx264, libx265).
• --crf: Set the Constant Rate Factor for quality (for libx264/libx265).
• --providers: (ONNX only) Set ONNX Runtime execution providers.

Tools

Utility scripts are located in the tools/ directory.

Convert PyTorch model to ONNX:

python tools/convert_to_onnx.py --model pretrained_models/model.pth --output model.onnx

• --dynamic: Create a model that supports various input sizes.
• --fp16: Convert the model to FP16 for a speed boost.

Credits (thanks all!)

Thank you to leobby, Hermes, and Bendel for testing the arch!

• Uses SCUNet's respository as a base
• Uses a modified version of SPAN's architecture
• Uses the additions made in both the SCUNet fork and TSCUNet repositories
• Is a continuation of the original TSPAN