add new real-esrgan model real-esrgan-general-x4v3 (#1566)

* add real-esrgan-general-x4v3 for new real-esrgan model real-esrgan-general-x4v3

* Reorganize the files, change tensorrtx/real-esrgan/ and tensorrtx/real-esrgan-general-x4v3/ to tensorrtx/real-esrgan/x4plus/ and tensorrtx/real-esrgan/general-x4v3/.

---------

Co-authored-by: xuejiejie <xuejiejie@xiaoyingai.com>

real-esrgan/general-x4v3/CMakeLists.txt

@@ -0,0 +1,42 @@
+cmake_minimum_required(VERSION 3.16)
+project(real-esrgan)
+
+set(CMAKE_MODULE_PATH ${CMAKE_MODULE_PATH} "${CMAKE_SOURCE_DIR}/cmake/")
+
+add_definitions(-std=c++17)
+add_definitions(-DAPI_EXPORTS)
+option(CUDA_USE_STATIC_CUDA_RUNTIME OFF)
+#set(CMAKE_CXX_STANDARD 11)
+set(CMAKE_BUILD_TYPE Debug)
+
+#find_package(CUDA REQUIRED)
+
+INCLUDE_DIRECTORIES(${PROJECT_SOURCE_DIR}/src/include)
+
+# cuda
+FIND_PACKAGE(CUDA REQUIRED)
+#INCLUDE_DIRECTORIES(${CUDA_INCLUDE_DIRS})
+include_directories(/usr/local/cuda/include)
+link_directories(/usr/local/cuda/lib64)
+
+# <------------------------TensorRT Related------------------------->
+include_directories(YOUR_TENSORRT_INCLUDE_DIR) # TensorRT-8.6.1.6/include
+link_directories(YOUR_TENSORRT_LIB_DIR) # TensorRT-8.6.1.6/lib
+
+# <------------------------OpenCV Related------------------------->
+# opencv
+FIND_PACKAGE(OpenCV REQUIRED)
+INCLUDE_DIRECTORIES(${OpenCV_INCLUDE_DIRS})
+
+set(CMAKE_CXX_STANDARD 17)
+
+add_executable(${PROJECT_NAME} main.cpp)
+
+cuda_add_library(myplugins SHARED ${PROJECT_SOURCE_DIR}/src/pixel_shuffle/pixel_shuffle.cu)
+target_link_libraries(myplugins nvinfer cudart)
+
+
+TARGET_LINK_LIBRARIES(${PROJECT_NAME} nvinfer)
+TARGET_LINK_LIBRARIES(${PROJECT_NAME} cudart)
+TARGET_LINK_LIBRARIES(${PROJECT_NAME} ${OpenCV_LIBS})
+TARGET_LINK_LIBRARIES(${PROJECT_NAME} myplugins)

real-esrgan/general-x4v3/README.md

@@ -0,0 +1,41 @@
+# Real-ESRGAN realesr-general-x4v3 model
+
+## How to Run
+0. Replace YOUR_TENSORRT_INCLUDE_DIR and YOUR_TENSORRT_LIB_DIR in CMakeLists.txt with your TensorRT include and lib directories.
+1. generate .wts from pytorch with .pt
+```
+git clone https://github.com/xinntao/Real-ESRGAN.git
+cd Real-ESRGAN
+
+# Install basicsr - https://github.com/xinntao/BasicSR
+# We use BasicSR for both training and inference
+pip install basicsr
+# facexlib and gfpgan are for face enhancement
+pip install facexlib
+pip install gfpgan
+pip install -r requirements.txt
+python setup.py develop
+```
+download realesr-general-x4v3.pth (and realesr-general-wdn-x4v3.pth if needed) from
+https://github.com/xinntao/Real-ESRGAN/releases
+
+```
+cp {tensorrtx}/real-esrgan-general-x4v3/gen_wts.py {xinntao}/Real-ESRGAN
+cd {xinntao}/Real-ESRGAN
+python gen_wts.py
+// a file 'real-esrgan.wts' will be generated.
+```
+
+**Be aware that if you need both realesr-general-x4v3.pth and realesr-general-wdn-x4v3.pth, please write a Python script to average all weights of realesr-general-x4v3.pth and realesr-general-wdn-x4v3.pth (from {xinntao}/Real-ESRGAN), then save it as a .pth file, and use this new file to generate a .wts file.**
+
+2. build tensorrtx/real-esrgan-general-x4v3 and run
+
+```
+cd {tensorrtx}/real-esrgan-general-x4v3/
+mkdir build
+cd build
+cp {xinntao}/Real-ESRGAN/real-esrgan.wts {tensorrtx}/real-esrgan/weights/
+cmake ..
+make
+./real-esrgan your_images_dir
+```

real-esrgan/general-x4v3/cmake/FindTensorRT.cmake

@@ -0,0 +1,87 @@
+# source:
+# https://github.com/NVIDIA/tensorrt-laboratory/blob/master/cmake/FindTensorRT.cmake
+
+# This module defines the following variables:
+#
+# ::
+#
+#   TensorRT_INCLUDE_DIRS
+#   TensorRT_LIBRARIES
+#   TensorRT_FOUND
+#
+# ::
+#
+#   TensorRT_VERSION_STRING - version (x.y.z)
+#   TensorRT_VERSION_MAJOR  - major version (x)
+#   TensorRT_VERSION_MINOR  - minor version (y)
+#   TensorRT_VERSION_PATCH  - patch version (z)
+#
+# Hints
+# ^^^^^
+# A user may set ``TensorRT_DIR`` to an installation root to tell this module where to look.
+#
+set(_TensorRT_SEARCHES)
+
+if(TensorRT_DIR)
+    set(_TensorRT_SEARCH_ROOT PATHS ${TensorRT_DIR} NO_DEFAULT_PATH)
+    list(APPEND _TensorRT_SEARCHES _TensorRT_SEARCH_ROOT)
+endif()
+
+# appends some common paths
+set(_TensorRT_SEARCH_NORMAL
+        PATHS "/usr"
+        )
+list(APPEND _TensorRT_SEARCHES _TensorRT_SEARCH_NORMAL)
+
+# Include dir
+foreach(search ${_TensorRT_SEARCHES})
+    find_path(TensorRT_INCLUDE_DIR NAMES NvInfer.h ${${search}} PATH_SUFFIXES include)
+endforeach()
+
+if(NOT TensorRT_LIBRARY)
+    foreach(search ${_TensorRT_SEARCHES})
+        find_library(TensorRT_LIBRARY NAMES nvinfer ${${search}} PATH_SUFFIXES lib)
+    endforeach()
+endif()
+
+if(NOT TensorRT_PARSERS_LIBRARY)
+    foreach(search ${_TensorRT_SEARCHES})
+        find_library(TensorRT_NVPARSERS_LIBRARY NAMES nvparsers ${${search}} PATH_SUFFIXES lib)
+    endforeach()
+endif()
+
+if(NOT TensorRT_NVONNXPARSER_LIBRARY)
+    foreach(search ${_TensorRT_SEARCHES})
+        find_library(TensorRT_NVONNXPARSER_LIBRARY NAMES nvonnxparser ${${search}} PATH_SUFFIXES lib)
+    endforeach()
+endif()
+
+mark_as_advanced(TensorRT_INCLUDE_DIR)
+
+if(TensorRT_INCLUDE_DIR AND EXISTS "${TensorRT_INCLUDE_DIR}/NvInfer.h")
+    file(STRINGS "${TensorRT_INCLUDE_DIR}/NvInfer.h" TensorRT_MAJOR REGEX "^#define NV_TENSORRT_MAJOR [0-9]+.*$")
+    file(STRINGS "${TensorRT_INCLUDE_DIR}/NvInfer.h" TensorRT_MINOR REGEX "^#define NV_TENSORRT_MINOR [0-9]+.*$")
+    file(STRINGS "${TensorRT_INCLUDE_DIR}/NvInfer.h" TensorRT_PATCH REGEX "^#define NV_TENSORRT_PATCH [0-9]+.*$")
+
+    string(REGEX REPLACE "^#define NV_TENSORRT_MAJOR ([0-9]+).*$" "\\1" TensorRT_VERSION_MAJOR "${TensorRT_MAJOR}")
+    string(REGEX REPLACE "^#define NV_TENSORRT_MINOR ([0-9]+).*$" "\\1" TensorRT_VERSION_MINOR "${TensorRT_MINOR}")
+    string(REGEX REPLACE "^#define NV_TENSORRT_PATCH ([0-9]+).*$" "\\1" TensorRT_VERSION_PATCH "${TensorRT_PATCH}")
+    set(TensorRT_VERSION_STRING "${TensorRT_VERSION_MAJOR}.${TensorRT_VERSION_MINOR}.${TensorRT_VERSION_PATCH}")
+endif()
+
+include(FindPackageHandleStandardArgs)
+FIND_PACKAGE_HANDLE_STANDARD_ARGS(TensorRT REQUIRED_VARS TensorRT_LIBRARY TensorRT_INCLUDE_DIR VERSION_VAR TensorRT_VERSION_STRING)
+
+if(TensorRT_FOUND)
+    set(TensorRT_INCLUDE_DIRS ${TensorRT_INCLUDE_DIR})
+
+    if(NOT TensorRT_LIBRARIES)
+        set(TensorRT_LIBRARIES ${TensorRT_LIBRARY} ${TensorRT_NVONNXPARSER_LIBRARY} ${TensorRT_NVPARSERS_LIBRARY})
+    endif()
+
+    if(NOT TARGET TensorRT::TensorRT)
+        add_library(TensorRT::TensorRT UNKNOWN IMPORTED)
+        set_target_properties(TensorRT::TensorRT PROPERTIES INTERFACE_INCLUDE_DIRECTORIES "${TensorRT_INCLUDE_DIRS}")
+        set_property(TARGET TensorRT::TensorRT APPEND PROPERTY IMPORTED_LOCATION "${TensorRT_LIBRARY}")
+    endif()
+endif()

real-esrgan/general-x4v3/gen_wts.py

@@ -0,0 +1,136 @@
+import argparse
+import os
+import struct
+from realesrgan import RealESRGANer
+from realesrgan.archs.srvgg_arch import SRVGGNetCompact
+
+from basicsr.archs.rrdbnet_arch import RRDBNet
+from basicsr.utils.download_util import load_file_from_url
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('-i', '--input', type=str, help='Input image or folder')
+    parser.add_argument(
+        '-n',
+        '--model_name',
+        type=str,
+        default='realesr-general-x4v3',
+        help=('RealESRGAN_x2plus Model names: '
+              'realesr-animevideov3 | realesr-general-x4v3'))
+    parser.add_argument('-o', '--output', type=str, help='Output folder')
+    parser.add_argument(
+        '-dn',
+        '--denoise_strength',
+        type=float,
+        default=0.5,
+        help=('Denoise strength. 0 for weak denoise (keep noise), 1 for strong denoise ability. '
+              'Only used for the realesr-general-x4v3 model'))
+    parser.add_argument('-s', '--outscale', type=float, default=4, help='The final upsampling scale of the image')
+    parser.add_argument(
+        '--model_path', type=str, default=None, help='[Option] Model path. Usually, you do not need to specify it')
+    parser.add_argument('--suffix', type=str, default='out', help='Suffix of the restored image')
+    parser.add_argument('-t', '--tile', type=int, default=0, help='Tile size, 0 for no tile during testing')
+    parser.add_argument('--tile_pad', type=int, default=10, help='Tile padding')
+    parser.add_argument('--pre_pad', type=int, default=0, help='Pre padding size at each border')
+    parser.add_argument('--face_enhance', action='store_true', help='Use GFPGAN to enhance face')
+    parser.add_argument(
+        '--fp32', action='store_true', help='Use fp32 precision during inference. Default: fp16 (half precision).')
+    parser.add_argument(
+        '--alpha_upsampler',
+        type=str,
+        default='realesrgan',
+        help='The upsampler for the alpha channels. Options: realesrgan | bicubic')
+    parser.add_argument(
+        '--ext',
+        type=str,
+        default='auto',
+        help='Image extension. Options: auto | jpg | png, auto means using the same extension as inputs')
+    parser.add_argument(
+        '-g', '--gpu-id', type=int, default=None, help='gpu device to use (default=None) can be 0,1,2 for multi-gpu')
+
+    args = parser.parse_args()
+
+    # determine models according to model names
+    args.model_name = args.model_name.split('.')[0]
+    if args.model_name == 'RealESRGAN_x4plus':  # x4 RRDBNet model
+        model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=4)
+        netscale = 4
+        file_url = ['https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth']
+    elif args.model_name == 'RealESRNet_x4plus':  # x4 RRDBNet model
+        model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=4)
+        netscale = 4
+        file_url = ['https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.1/RealESRNet_x4plus.pth']
+    elif args.model_name == 'RealESRGAN_x4plus_anime_6B':  # x4 RRDBNet model with 6 blocks
+        model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=6, num_grow_ch=32, scale=4)
+        netscale = 4
+        file_url = ['https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.2.4/RealESRGAN_x4plus_anime_6B.pth']
+    elif args.model_name == 'RealESRGAN_x2plus':  # x2 RRDBNet model
+        model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=2)
+        netscale = 2
+        file_url = ['https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.1/RealESRGAN_x2plus.pth']
+    elif args.model_name == 'realesr-animevideov3':  # x4 VGG-style model (XS size)
+        model = SRVGGNetCompact(num_in_ch=3, num_out_ch=3, num_feat=64, num_conv=16, upscale=4, act_type='prelu')
+        netscale = 4
+        file_url = ['https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-animevideov3.pth']
+    elif args.model_name == 'realesr-general-x4v3':  # x4 VGG-style model (S size)
+        model = SRVGGNetCompact(num_in_ch=3, num_out_ch=3, num_feat=64, num_conv=32, upscale=4, act_type='prelu')
+        netscale = 4
+        file_url = [
+            'https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-general-wdn-x4v3.pth',
+            'https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-general-x4v3.pth'
+        ]
+
+    # determine model paths
+    if args.model_path is not None:
+        model_path = args.model_path
+    else:
+        model_path = os.path.join('weights', args.model_name + '.pth')
+        if not os.path.isfile(model_path):
+            ROOT_DIR = os.path.dirname(os.path.abspath(__file__))
+            for url in file_url:
+                # model_path will be updated
+                model_path = load_file_from_url(
+                    url=url, model_dir=os.path.join(ROOT_DIR, 'weights'), progress=True, file_name=None)
+
+    # use dni to control the denoise strength
+    dni_weight = None
+    if args.model_name == 'realesr-general-x4v3' and args.denoise_strength != 1:
+        # wdn_model_path = model_path.replace('realesr-general-x4v3', 'realesr-general-wdn-x4v3')
+        # model_path = [model_path, wdn_model_path]
+        # dni_weight = [args.denoise_strength, 1 - args.denoise_strength]
+        model_path = model_path.replace('realesr-general-x4v3', 'realesr-general-x4v3-cat')
+        dni_weight = None
+
+    # restorer
+    upsampler = RealESRGANer(
+        scale=netscale,
+        model_path=model_path,
+        dni_weight=dni_weight,
+        model=model,
+        tile=args.tile,
+        tile_pad=args.tile_pad,
+        pre_pad=args.pre_pad,
+        half=not args.fp32,
+        gpu_id=args.gpu_id)
+
+    if os.path.isfile('real-esrgan.wts'):
+        print('Already, real-esrgan.wts file exists.')
+    else:
+        print('making real-esrgan.wts file ...')
+        f = open("real-esrgan.wts", 'w')
+        f.write("{}\n".format(len(upsampler.model.state_dict().keys())))
+        for k, v in upsampler.model.state_dict().items():
+            print('key: ', k)
+            print('value: ', v.shape)
+            vr = v.reshape(-1).cpu().numpy()
+            f.write("{} {}".format(k, len(vr)))
+            for vv in vr:
+                f.write(" ")
+                f.write(struct.pack(">f", float(vv)).hex())
+            f.write("\n")
+        print('Completed real-esrgan.wts file!')
+
+
+if __name__ == '__main__':
+    main()

real-esrgan/general-x4v3/main.cpp

@@ -0,0 +1,342 @@
+#include <NvInfer.h>
+#include <dirent.h>
+#include <fstream>
+#include <iostream>
+#include <memory>
+#include <opencv4/opencv2/opencv.hpp>
+#include <vector>
+
+#include "config/config.hpp"
+#include "cuda_utils.h"
+#include "logging/logging.h"
+#include "pixel_shuffle/pixel_shuffle.hpp"
+#include "preprocess/preprocess.hpp"
+
+static Logger gLogger;
+
+using namespace nvinfer1;
+
+// TensorRT weight files have a simple space delimited format:
+// [type] [size] <data x size in hex>
+std::map<std::string, Weights> loadWeights(const std::string file) {
+    std::cout << "Loading weights: " << file << std::endl;
+    std::map<std::string, Weights> weightMap;
+
+    // Open weights file
+    std::ifstream input(file);
+    assert(input.is_open() && "Unable to load weight file. please check if the .wts file path is right!!!!!!");
+
+    // Read number of weight blobs
+    int32_t count;
+    input >> count;
+    assert(count > 0 && "Invalid weight map file.");
+
+    while (count--) {
+        Weights wt{DataType::kFLOAT, nullptr, 0};
+        uint32_t size;
+
+        // Read name and type of blob
+        std::string name;
+        input >> name >> std::dec >> size;
+        wt.type = DataType::kFLOAT;
+
+        // Load blob
+        uint32_t* val = reinterpret_cast<uint32_t*>(malloc(sizeof(val) * size));
+        for (uint32_t x = 0, y = size; x < y; ++x) {
+            input >> std::hex >> val[x];
+        }
+        wt.values = val;
+
+        wt.count = size;
+        weightMap[name] = wt;
+    }
+
+    return weightMap;
+}
+
+auto* ConvPRelu(INetworkDefinition* network, std::map<std::string, Weights>& weightMap, ITensor& input, int conv_nb,
+                int index) {
+
+    IConvolutionLayer* conv = network->addConvolutionNd(input, conv_nb, DimsHW{3, 3},
+                                                        weightMap["body." + std::to_string(index) + ".weight"],
+                                                        weightMap["body." + std::to_string(index) + ".bias"]);
+    assert(conv);
+    conv->setName(("body." + std::to_string(index) + ".weight").c_str());
+    conv->setStrideNd(DimsHW{1, 1});
+    conv->setPaddingNd(DimsHW{1, 1});
+    auto conv_res = conv->getOutput(0);
+
+    // add prelu layer
+    // slope 64 number
+
+    //auto slope = network->addConstant( {64}, weightMap["body." + std::to_string(index + 1) + ".weight"] );
+    auto slope = network->addConstant(Dims4{1, 64, 1, 1}, weightMap["body." + std::to_string(index + 1) + ".weight"]);
+    assert(slope);
+    slope->setName(("body." + std::to_string(index + 1) + ".weight").c_str());
+
+    auto prelu = network->addParametricReLU(*conv_res, *slope->getOutput(0));
+    assert(prelu);
+
+    return prelu;
+}
+
+void build_engine(DataType dt, std::string& wts_path) {
+
+    std::map<std::string, Weights> weightMap = loadWeights(wts_path);
+
+    nvinfer1::IBuilder* builder = nvinfer1::createInferBuilder(gLogger);
+    nvinfer1::IBuilderConfig* config = builder->createBuilderConfig();
+
+    nvinfer1::INetworkDefinition* network = builder->createNetworkV2(1U);
+
+    auto data = network->addInput(INPUT_BLOB_NAME, nvinfer1::DataType::kFLOAT,
+                                  nvinfer1::Dims4{BATCH_SIZE, INPUT_C, INPUT_H, INPUT_W});
+
+    // first
+    auto layer = ConvPRelu(network, weightMap, *data, 64, 0);
+
+    for (int i = 0; i < 32; ++i) {
+        layer = ConvPRelu(network, weightMap, *layer->getOutput(0), 64, 2 * i + 2);
+    }
+
+    auto conv_last = network->addConvolutionNd(*layer->getOutput(0), 48, DimsHW{3, 3}, weightMap["body.66.weight"],
+                                               weightMap["body.66.bias"]);
+    assert(conv_last);
+    conv_last->setName("body.66.weight");
+    conv_last->setStrideNd(DimsHW{1, 1});
+    conv_last->setPaddingNd(DimsHW{1, 1});
+    auto conv_last_res = conv_last->getOutput(0);
+
+    // add pixel shuffle layer by plugin
+    IPluginCreator* creator = getPluginRegistry()->getPluginCreator("PixelShufflePlugin", "1");
+    const PluginFieldCollection* pluginFC = creator->getFieldNames();
+    std::vector<PluginField> pluginData;
+    int upscaleFactor = 4;
+    pluginData.emplace_back(PluginField{"upscaleFactor", &upscaleFactor, PluginFieldType::kINT32, 1});
+    PluginFieldCollection pluginFCWithData = {static_cast<int>(pluginData.size()), pluginData.data()};
+    auto pluginObj = creator->createPlugin("PixelShuffle", &pluginFCWithData);
+
+    auto pixelShuffleLayer = network->addPluginV2(&conv_last_res, 1, *pluginObj);
+
+    // the input "data" interpolate 4x and add to pixelShuffleLayer->getOutput(0)
+
+    auto interpolateLayer = network->addResize(*data);
+    interpolateLayer->setResizeMode(ResizeMode::kNEAREST);
+    // Define scale factors
+    float scales[] = {1.0f, 1.0f, 1.0 * OUT_SCALE, 1.0 * OUT_SCALE};  // scale_factor=4 for height and width
+    interpolateLayer->setScales(scales, OUT_SCALE);
+
+    // Add the two tensor as output
+    auto addLayer = network->addElementWise(*interpolateLayer->getOutput(0), *pixelShuffleLayer->getOutput(0),
+                                            ElementWiseOperation::kSUM);
+
+    // output
+    addLayer->getOutput(0)->setName(OUTPUT_BLOB_NAME);
+    network->markOutput(*addLayer->getOutput(0));
+
+    // fp16
+    if (USE_FP16) {
+        config->setFlag(BuilderFlag::kFP16);
+    }
+
+    std::cout << "Building engine, please wait for a while..." << std::endl;
+    IHostMemory* serialized_model = builder->buildSerializedNetwork(*network, *config);
+    std::cout << "Build engine successfully!" << std::endl;
+
+    // Don't need the network any more
+    delete network;
+
+    // Release host memory
+    for (auto& mem : weightMap) {
+        free((void*)(mem.second.values));
+    }
+
+    std::ofstream ofs("../weights/real-esrgan.engine", std::ios::binary);
+
+    assert(serialized_model != nullptr);
+    ofs.write(reinterpret_cast<const char*>(serialized_model->data()), serialized_model->size());
+
+    delete config;
+    delete serialized_model;
+    delete builder;
+}
+
+static inline int read_files_in_dir(const char* p_dir_name, std::vector<std::string>& file_names) {
+    DIR* p_dir = opendir(p_dir_name);
+    if (p_dir == nullptr) {
+        return -1;
+    }
+
+    struct dirent* p_file = nullptr;
+    while ((p_file = readdir(p_dir)) != nullptr) {
+        if (strcmp(p_file->d_name, ".") != 0 && strcmp(p_file->d_name, "..") != 0) {
+            //std::string cur_file_name(p_dir_name);
+            //cur_file_name += "/";
+            //cur_file_name += p_file->d_name;
+            std::string cur_file_name(p_file->d_name);
+            file_names.push_back(cur_file_name);
+        }
+    }
+
+    closedir(p_dir);
+    return 0;
+}
+
+void doInference(IExecutionContext& context, cudaStream_t& stream, void** buffers, float* output) {
+    context.setBindingDimensions(0, Dims4(BATCH_SIZE, INPUT_C, INPUT_H, INPUT_W));
+    context.enqueueV2(buffers, stream, nullptr);
+    CUDA_CHECK(cudaMemcpyAsync(output, buffers[1], BATCH_SIZE * OUTPUT_SIZE * sizeof(float), cudaMemcpyDeviceToHost,
+                               stream));
+    CUDA_CHECK(cudaStreamSynchronize(stream));
+}
+
+int main(int argc, char** argv) {
+    std::string img_dir;
+
+    if (argc < 2) {
+        std::cerr << "Usage: " << argv[0] << " <image_dir>" << std::endl;
+        return -1;
+    } else {
+        img_dir = argv[1];
+    }
+
+    std::string wts_path = "../weights/real-esrgan.wts";
+    build_engine(DataType::kFLOAT, wts_path);
+
+    std::string engine_name = "../weights/real-esrgan.engine";
+    // deserialize the .engine and run inference
+    std::ifstream file(engine_name, std::ios::binary);
+    if (!file.good()) {
+        std::cerr << "read " << engine_name << " error!" << std::endl;
+        return -1;
+    }
+    char* trtModelStream = nullptr;
+    size_t size = 0;
+    file.seekg(0, file.end);
+    size = file.tellg();
+    file.seekg(0, file.beg);
+    trtModelStream = new char[size];
+    assert(trtModelStream);
+    file.read(trtModelStream, size);
+    file.close();
+
+    IRuntime* runtime = createInferRuntime(gLogger);
+    assert(runtime != nullptr);
+    ICudaEngine* engine = runtime->deserializeCudaEngine(trtModelStream, size);
+    assert(engine != nullptr);
+    IExecutionContext* context = engine->createExecutionContext();
+    assert(context != nullptr);
+    delete[] trtModelStream;
+    assert(engine->getNbBindings() == 2);
+    void* buffers[2];
+    // In order to bind the buffers, we need to know the names of the input and output tensors.
+    // Note that indices are guaranteed to be less than IEngine::getNbBindings()
+    const int inputIndex = engine->getBindingIndex(INPUT_BLOB_NAME);
+    const int outputIndex = engine->getBindingIndex(OUTPUT_BLOB_NAME);
+    assert(inputIndex == 0);
+    assert(outputIndex == 1);
+
+    // Create GPU buffers on device
+    CUDA_CHECK(cudaMalloc(&buffers[inputIndex], BATCH_SIZE * INPUT_C * INPUT_H * INPUT_W * sizeof(float)));
+    CUDA_CHECK(cudaMalloc(&buffers[outputIndex], BATCH_SIZE * OUTPUT_SIZE * sizeof(float)));
+
+    std::vector<float> data;
+    std::vector<float> output;
+    //std::vector<float> res;
+
+    //data.resize(BATCH_SIZE * 3 * INPUT_H * INPUT_W);
+    data.resize(BATCH_SIZE * INPUT_C * INPUT_H * INPUT_W);
+    output.resize(BATCH_SIZE * OUTPUT_SIZE);
+
+    // Create stream
+    cudaStream_t stream;
+    CUDA_CHECK(cudaStreamCreate(&stream));
+
+    std::vector<std::string> file_names;
+    if (read_files_in_dir(img_dir.c_str(), file_names) < 0) {
+        std::cerr << "read_files_in_dir failed." << std::endl;
+        return -1;
+    }
+
+    for (int index = 0; index < file_names.size(); ++index) {
+
+        auto img = cv::imread(img_dir + "/" + file_names[index]);
+        auto begin = std::chrono::high_resolution_clock::now();
+
+        // BATCH_SIZE = 1
+        for (int b = 0; b < BATCH_SIZE; b++) {
+            int i = 0;
+            for (int row = 0; row < INPUT_H; ++row) {
+                uchar* uc_pixel = img.data + row * img.step;
+                for (int col = 0; col < INPUT_W; ++col) {
+                    //    static float data[BATCH_SIZE * 3 * INPUT_H * INPUT_W];
+                    // BGR2RGB and normalization
+                    data[b * 3 * INPUT_H * INPUT_W + i] = (float)uc_pixel[2] / 255.0;
+                    data[b * 3 * INPUT_H * INPUT_W + i + INPUT_H * INPUT_W] = (float)uc_pixel[1] / 255.0;
+                    data[b * 3 * INPUT_H * INPUT_W + i + 2 * INPUT_H * INPUT_W] = (float)uc_pixel[0] / 255.0;
+                    uc_pixel += 3;
+                    ++i;
+                }
+            }
+        }
+        CUDA_CHECK(cudaMemcpyAsync(buffers[inputIndex], data.data(),
+                                   BATCH_SIZE * INPUT_C * INPUT_H * INPUT_W * sizeof(float), cudaMemcpyHostToDevice,
+                                   stream));
+        doInference(*context, stream, (void**)buffers, output.data());
+        auto end = std::chrono::high_resolution_clock::now();
+        std::cout << "Inference time: " << std::chrono::duration_cast<std::chrono::milliseconds>(end - begin).count()
+                  << " ms" << std::endl;
+
+        int OUTPUT_C = 3;
+        int OUTPUT_H = INPUT_H * OUT_SCALE;
+        int OUTPUT_W = INPUT_W * OUT_SCALE;
+
+        for (int b = 0; b < BATCH_SIZE; b++) {
+            cv::Mat img_res(OUTPUT_H, OUTPUT_W, CV_8UC3);
+            int i = 0;
+            for (int row = 0; row < OUTPUT_H; ++row) {
+                uchar* uc_pixel = img_res.data + row * img_res.step;
+                for (int col = 0; col < OUTPUT_W; ++col) {
+                    // RGB2BGR and de_normalization
+                    auto r2 = std::round(output[b * OUTPUT_C * OUTPUT_H * OUTPUT_W + i] * 255.0);
+                    if (r2 < 0)
+                        r2 = 0;
+                    if (r2 > 255)
+                        r2 = 255;
+                    auto g2 = std::round(output[b * OUTPUT_C * OUTPUT_H * OUTPUT_W + i + 1 * OUTPUT_H * OUTPUT_W] *
+                                         255.0);
+                    if (g2 < 0)
+                        g2 = 0;
+                    if (g2 > 255)
+                        g2 = 255;
+                    auto b2 = std::round(output[b * OUTPUT_C * OUTPUT_H * OUTPUT_W + i + 2 * OUTPUT_H * OUTPUT_W] *
+                                         255.0);
+                    if (b2 < 0)
+                        b2 = 0;
+                    if (b2 > 255)
+                        b2 = 255;
+
+                    uc_pixel[0] = static_cast<uchar>(b2);  // B
+                    uc_pixel[1] = static_cast<uchar>(g2);  // G
+                    uc_pixel[2] = static_cast<uchar>(r2);  // R
+
+                    // uc_pixel[0] = static_cast<uchar>(std::round(output[b * OUTPUT_C * OUTPUT_H * OUTPUT_W + i + 2 * OUTPUT_H * OUTPUT_W] * 255.0)); // B
+                    // uc_pixel[1] = static_cast<uchar>(std::round(output[b * OUTPUT_C * OUTPUT_H * OUTPUT_W + i + 1 * OUTPUT_H * OUTPUT_W] * 255.0)); // G
+                    // uc_pixel[2] = static_cast<uchar>(std::round(output[b * OUTPUT_C * OUTPUT_H * OUTPUT_W + i] * 255.0)); // R
+                    uc_pixel += 3;
+                    ++i;
+                }
+            }
+            cv::imwrite("_" + file_names[index] + ".jpg", img_res);
+        }
+    }
+
+    // Release stream and buffers
+    cudaStreamDestroy(stream);
+    CUDA_CHECK(cudaFree(buffers[0]));
+    CUDA_CHECK(cudaFree(buffers[1]));
+    // Destroy the engine
+    delete context;
+    delete engine;
+    delete runtime;
+}

real-esrgan/general-x4v3/src/include/config/config.hpp

@@ -0,0 +1,22 @@
+#ifndef REAL_ESRGAN_TRT_CONFIG_HPP
+#define REAL_ESRGAN_TRT_CONFIG_HPP
+
+#include <string>
+
+//std::string INPUT_BLOB_NAME = "input";
+//std::string OUTPUT_BLOB_NAME = "output";
+
+const char* INPUT_BLOB_NAME = "input_0";
+const char* OUTPUT_BLOB_NAME = "output_0";
+
+const bool USE_FP16 = false;
+
+static const int BATCH_SIZE = 1;
+static const int INPUT_C = 3;
+static const int INPUT_H = 450;
+static const int INPUT_W = 300;
+static const int OUT_SCALE = 4;
+//static const int OUTPUT_SIZE = INPUT_C * INPUT_H * OUT_SCALE * INPUT_W * OUT_SCALE;
+static const int OUTPUT_SIZE = BATCH_SIZE * 48 * 450 * 300;
+//INPUT_C * INPUT_H * OUT_SCALE * INPUT_W * OUT_SCALE;
+#endif  //REAL_ESRGAN_TRT_CONFIG_HPP

real-esrgan/general-x4v3/src/include/cuda_utils.h

@@ -0,0 +1,21 @@
+#ifndef TRTX_CUDA_UTILS_H_
+#define TRTX_CUDA_UTILS_H_
+
+#include <cuda_runtime_api.h>
+#include <stdint.h>
+#include <cstdio>
+#include <iostream>
+#include <vector>
+
+#ifndef CUDA_CHECK
+#define CUDA_CHECK(callstr)                                                                    \
+    {                                                                                          \
+        cudaError_t error_code = callstr;                                                      \
+        if (error_code != cudaSuccess) {                                                       \
+            std::cerr << "CUDA error " << error_code << " at " << __FILE__ << ":" << __LINE__; \
+            assert(0);                                                                         \
+        }                                                                                      \
+    }
+#endif  // CUDA_CHECK
+
+#endif  // TRTX_CUDA_UTILS_H_

real-esrgan/general-x4v3/src/include/logging/logging.h

@@ -0,0 +1,455 @@
+/*
+ * Copyright (c) 2019, NVIDIA CORPORATION. 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.
+ */
+
+#ifndef TENSORRT_LOGGING_H
+#define TENSORRT_LOGGING_H
+
+#include <cassert>
+#include <ctime>
+#include <iomanip>
+#include <iostream>
+#include <ostream>
+#include <sstream>
+#include <string>
+#include "NvInferRuntimeCommon.h"
+
+using Severity = nvinfer1::ILogger::Severity;
+
+class LogStreamConsumerBuffer : public std::stringbuf {
+   public:
+    LogStreamConsumerBuffer(std::ostream& stream, const std::string& prefix, bool shouldLog)
+        : mOutput(stream), mPrefix(prefix), mShouldLog(shouldLog) {}
+
+    LogStreamConsumerBuffer(LogStreamConsumerBuffer&& other) : mOutput(other.mOutput) {}
+
+    ~LogStreamConsumerBuffer() {
+        // std::streambuf::pbase() gives a pointer to the beginning of the buffered part of the output sequence
+        // std::streambuf::pptr() gives a pointer to the current position of the output sequence
+        // if the pointer to the beginning is not equal to the pointer to the current position,
+        // call putOutput() to log the output to the stream
+        if (pbase() != pptr()) {
+            putOutput();
+        }
+    }
+
+    // synchronizes the stream buffer and returns 0 on success
+    // synchronizing the stream buffer consists of inserting the buffer contents into the stream,
+    // resetting the buffer and flushing the stream
+    virtual int sync() {
+        putOutput();
+        return 0;
+    }
+
+    void putOutput() {
+        if (mShouldLog) {
+            // prepend timestamp
+            std::time_t timestamp = std::time(nullptr);
+            tm* tm_local = std::localtime(&timestamp);
+            std::cout << "[";
+            std::cout << std::setw(2) << std::setfill('0') << 1 + tm_local->tm_mon << "/";
+            std::cout << std::setw(2) << std::setfill('0') << tm_local->tm_mday << "/";
+            std::cout << std::setw(4) << std::setfill('0') << 1900 + tm_local->tm_year << "-";
+            std::cout << std::setw(2) << std::setfill('0') << tm_local->tm_hour << ":";
+            std::cout << std::setw(2) << std::setfill('0') << tm_local->tm_min << ":";
+            std::cout << std::setw(2) << std::setfill('0') << tm_local->tm_sec << "] ";
+            // std::stringbuf::str() gets the string contents of the buffer
+            // insert the buffer contents pre-appended by the appropriate prefix into the stream
+            mOutput << mPrefix << str();
+            // set the buffer to empty
+            str("");
+            // flush the stream
+            mOutput.flush();
+        }
+    }
+
+    void setShouldLog(bool shouldLog) { mShouldLog = shouldLog; }
+
+   private:
+    std::ostream& mOutput;
+    std::string mPrefix;
+    bool mShouldLog;
+};
+
+//!
+//! \class LogStreamConsumerBase
+//! \brief Convenience object used to initialize LogStreamConsumerBuffer before std::ostream in LogStreamConsumer
+//!
+class LogStreamConsumerBase {
+   public:
+    LogStreamConsumerBase(std::ostream& stream, const std::string& prefix, bool shouldLog)
+        : mBuffer(stream, prefix, shouldLog) {}
+
+   protected:
+    LogStreamConsumerBuffer mBuffer;
+};
+
+//!
+//! \class LogStreamConsumer
+//! \brief Convenience object used to facilitate use of C++ stream syntax when logging messages.
+//!  Order of base classes is LogStreamConsumerBase and then std::ostream.
+//!  This is because the LogStreamConsumerBase class is used to initialize the LogStreamConsumerBuffer member field
+//!  in LogStreamConsumer and then the address of the buffer is passed to std::ostream.
+//!  This is necessary to prevent the address of an uninitialized buffer from being passed to std::ostream.
+//!  Please do not change the order of the parent classes.
+//!
+class LogStreamConsumer : protected LogStreamConsumerBase, public std::ostream {
+   public:
+    //! \brief Creates a LogStreamConsumer which logs messages with level severity.
+    //!  Reportable severity determines if the messages are severe enough to be logged.
+    LogStreamConsumer(Severity reportableSeverity, Severity severity)
+        : LogStreamConsumerBase(severityOstream(severity), severityPrefix(severity), severity <= reportableSeverity),
+          std::ostream(&mBuffer)  // links the stream buffer with the stream
+          ,
+          mShouldLog(severity <= reportableSeverity),
+          mSeverity(severity) {}
+
+    LogStreamConsumer(LogStreamConsumer&& other)
+        : LogStreamConsumerBase(severityOstream(other.mSeverity), severityPrefix(other.mSeverity), other.mShouldLog),
+          std::ostream(&mBuffer)  // links the stream buffer with the stream
+          ,
+          mShouldLog(other.mShouldLog),
+          mSeverity(other.mSeverity) {}
+
+    void setReportableSeverity(Severity reportableSeverity) {
+        mShouldLog = mSeverity <= reportableSeverity;
+        mBuffer.setShouldLog(mShouldLog);
+    }
+
+   private:
+    static std::ostream& severityOstream(Severity severity) {
+        return severity >= Severity::kINFO ? std::cout : std::cerr;
+    }
+
+    static std::string severityPrefix(Severity severity) {
+        switch (severity) {
+            case Severity::kINTERNAL_ERROR:
+                return "[F] ";
+            case Severity::kERROR:
+                return "[E] ";
+            case Severity::kWARNING:
+                return "[W] ";
+            case Severity::kINFO:
+                return "[I] ";
+            case Severity::kVERBOSE:
+                return "[V] ";
+            default:
+                assert(0);
+                return "";
+        }
+    }
+
+    bool mShouldLog;
+    Severity mSeverity;
+};
+
+//! \class Logger
+//!
+//! \brief Class which manages logging of TensorRT tools and samples
+//!
+//! \details This class provides a common interface for TensorRT tools and samples to log information to the console,
+//! and supports logging two types of messages:
+//!
+//! - Debugging messages with an associated severity (info, warning, error, or internal error/fatal)
+//! - Test pass/fail messages
+//!
+//! The advantage of having all samples use this class for logging as opposed to emitting directly to stdout/stderr is
+//! that the logic for controlling the verbosity and formatting of sample output is centralized in one location.
+//!
+//! In the future, this class could be extended to support dumping test results to a file in some standard format
+//! (for example, JUnit XML), and providing additional metadata (e.g. timing the duration of a test run).
+//!
+//! TODO: For backwards compatibility with existing samples, this class inherits directly from the nvinfer1::ILogger
+//! interface, which is problematic since there isn't a clean separation between messages coming from the TensorRT
+//! library and messages coming from the sample.
+//!
+//! In the future (once all samples are updated to use Logger::getTRTLogger() to access the ILogger) we can refactor the
+//! class to eliminate the inheritance and instead make the nvinfer1::ILogger implementation a member of the Logger
+//! object.
+
+class Logger : public nvinfer1::ILogger {
+   public:
+    Logger(Severity severity = Severity::kWARNING) : mReportableSeverity(severity) {}
+
+    //!
+    //! \enum TestResult
+    //! \brief Represents the state of a given test
+    //!
+    enum class TestResult {
+        kRUNNING,  //!< The test is running
+        kPASSED,   //!< The test passed
+        kFAILED,   //!< The test failed
+        kWAIVED    //!< The test was waived
+    };
+
+    //!
+    //! \brief Forward-compatible method for retrieving the nvinfer::ILogger associated with this Logger
+    //! \return The nvinfer1::ILogger associated with this Logger
+    //!
+    //! TODO Once all samples are updated to use this method to register the logger with TensorRT,
+    //! we can eliminate the inheritance of Logger from ILogger
+    //!
+    nvinfer1::ILogger& getTRTLogger() { return *this; }
+
+    //!
+    //! \brief Implementation of the nvinfer1::ILogger::log() virtual method
+    //!
+    //! Note samples should not be calling this function directly; it will eventually go away once we eliminate the
+    //! inheritance from nvinfer1::ILogger
+    //!
+    void log(Severity severity, const char* msg) noexcept override {
+        LogStreamConsumer(mReportableSeverity, severity) << "[TRT] " << std::string(msg) << std::endl;
+    }
+
+    //!
+    //! \brief Method for controlling the verbosity of logging output
+    //!
+    //! \param severity The logger will only emit messages that have severity of this level or higher.
+    //!
+    void setReportableSeverity(Severity severity) { mReportableSeverity = severity; }
+
+    //!
+    //! \brief Opaque handle that holds logging information for a particular test
+    //!
+    //! This object is an opaque handle to information used by the Logger to print test results.
+    //! The sample must call Logger::defineTest() in order to obtain a TestAtom that can be used
+    //! with Logger::reportTest{Start,End}().
+    //!
+    class TestAtom {
+       public:
+        TestAtom(TestAtom&&) = default;
+
+       private:
+        friend class Logger;
+
+        TestAtom(bool started, const std::string& name, const std::string& cmdline)
+            : mStarted(started), mName(name), mCmdline(cmdline) {}
+
+        bool mStarted;
+        std::string mName;
+        std::string mCmdline;
+    };
+
+    //!
+    //! \brief Define a test for logging
+    //!
+    //! \param[in] name The name of the test.  This should be a string starting with
+    //!                  "TensorRT" and containing dot-separated strings containing
+    //!                  the characters [A-Za-z0-9_].
+    //!                  For example, "TensorRT.sample_googlenet"
+    //! \param[in] cmdline The command line used to reproduce the test
+    //
+    //! \return a TestAtom that can be used in Logger::reportTest{Start,End}().
+    //!
+    static TestAtom defineTest(const std::string& name, const std::string& cmdline) {
+        return TestAtom(false, name, cmdline);
+    }
+
+    //!
+    //! \brief A convenience overloaded version of defineTest() that accepts an array of command-line arguments
+    //!        as input
+    //!
+    //! \param[in] name The name of the test
+    //! \param[in] argc The number of command-line arguments
+    //! \param[in] argv The array of command-line arguments (given as C strings)
+    //!
+    //! \return a TestAtom that can be used in Logger::reportTest{Start,End}().
+    static TestAtom defineTest(const std::string& name, int argc, char const* const* argv) {
+        auto cmdline = genCmdlineString(argc, argv);
+        return defineTest(name, cmdline);
+    }
+
+    //!
+    //! \brief Report that a test has started.
+    //!
+    //! \pre reportTestStart() has not been called yet for the given testAtom
+    //!
+    //! \param[in] testAtom The handle to the test that has started
+    //!
+    static void reportTestStart(TestAtom& testAtom) {
+        reportTestResult(testAtom, TestResult::kRUNNING);
+        assert(!testAtom.mStarted);
+        testAtom.mStarted = true;
+    }
+
+    //!
+    //! \brief Report that a test has ended.
+    //!
+    //! \pre reportTestStart() has been called for the given testAtom
+    //!
+    //! \param[in] testAtom The handle to the test that has ended
+    //! \param[in] result The result of the test. Should be one of TestResult::kPASSED,
+    //!                   TestResult::kFAILED, TestResult::kWAIVED
+    //!
+    static void reportTestEnd(const TestAtom& testAtom, TestResult result) {
+        assert(result != TestResult::kRUNNING);
+        assert(testAtom.mStarted);
+        reportTestResult(testAtom, result);
+    }
+
+    static int reportPass(const TestAtom& testAtom) {
+        reportTestEnd(testAtom, TestResult::kPASSED);
+        return EXIT_SUCCESS;
+    }
+
+    static int reportFail(const TestAtom& testAtom) {
+        reportTestEnd(testAtom, TestResult::kFAILED);
+        return EXIT_FAILURE;
+    }
+
+    static int reportWaive(const TestAtom& testAtom) {
+        reportTestEnd(testAtom, TestResult::kWAIVED);
+        return EXIT_SUCCESS;
+    }
+
+    static int reportTest(const TestAtom& testAtom, bool pass) {
+        return pass ? reportPass(testAtom) : reportFail(testAtom);
+    }
+
+    Severity getReportableSeverity() const { return mReportableSeverity; }
+
+   private:
+    //!
+    //! \brief returns an appropriate string for prefixing a log message with the given severity
+    //!
+    static const char* severityPrefix(Severity severity) {
+        switch (severity) {
+            case Severity::kINTERNAL_ERROR:
+                return "[F] ";
+            case Severity::kERROR:
+                return "[E] ";
+            case Severity::kWARNING:
+                return "[W] ";
+            case Severity::kINFO:
+                return "[I] ";
+            case Severity::kVERBOSE:
+                return "[V] ";
+            default:
+                assert(0);
+                return "";
+        }
+    }
+
+    //!
+    //! \brief returns an appropriate string for prefixing a test result message with the given result
+    //!
+    static const char* testResultString(TestResult result) {
+        switch (result) {
+            case TestResult::kRUNNING:
+                return "RUNNING";
+            case TestResult::kPASSED:
+                return "PASSED";
+            case TestResult::kFAILED:
+                return "FAILED";
+            case TestResult::kWAIVED:
+                return "WAIVED";
+            default:
+                assert(0);
+                return "";
+        }
+    }
+
+    //!
+    //! \brief returns an appropriate output stream (cout or cerr) to use with the given severity
+    //!
+    static std::ostream& severityOstream(Severity severity) {
+        return severity >= Severity::kINFO ? std::cout : std::cerr;
+    }
+
+    //!
+    //! \brief method that implements logging test results
+    //!
+    static void reportTestResult(const TestAtom& testAtom, TestResult result) {
+        severityOstream(Severity::kINFO) << "&&&& " << testResultString(result) << " " << testAtom.mName << " # "
+                                         << testAtom.mCmdline << std::endl;
+    }
+
+    //!
+    //! \brief generate a command line string from the given (argc, argv) values
+    //!
+    static std::string genCmdlineString(int argc, char const* const* argv) {
+        std::stringstream ss;
+        for (int i = 0; i < argc; i++) {
+            if (i > 0)
+                ss << " ";
+            ss << argv[i];
+        }
+        return ss.str();
+    }
+
+    Severity mReportableSeverity;
+};
+
+namespace {
+
+//!
+//! \brief produces a LogStreamConsumer object that can be used to log messages of severity kVERBOSE
+//!
+//! Example usage:
+//!
+//!     LOG_VERBOSE(logger) << "hello world" << std::endl;
+//!
+inline LogStreamConsumer LOG_VERBOSE(const Logger& logger) {
+    return LogStreamConsumer(logger.getReportableSeverity(), Severity::kVERBOSE);
+}
+
+//!
+//! \brief produces a LogStreamConsumer object that can be used to log messages of severity kINFO
+//!
+//! Example usage:
+//!
+//!     LOG_INFO(logger) << "hello world" << std::endl;
+//!
+inline LogStreamConsumer LOG_INFO(const Logger& logger) {
+    return LogStreamConsumer(logger.getReportableSeverity(), Severity::kINFO);
+}
+
+//!
+//! \brief produces a LogStreamConsumer object that can be used to log messages of severity kWARNING
+//!
+//! Example usage:
+//!
+//!     LOG_WARN(logger) << "hello world" << std::endl;
+//!
+inline LogStreamConsumer LOG_WARN(const Logger& logger) {
+    return LogStreamConsumer(logger.getReportableSeverity(), Severity::kWARNING);
+}
+
+//!
+//! \brief produces a LogStreamConsumer object that can be used to log messages of severity kERROR
+//!
+//! Example usage:
+//!
+//!     LOG_ERROR(logger) << "hello world" << std::endl;
+//!
+inline LogStreamConsumer LOG_ERROR(const Logger& logger) {
+    return LogStreamConsumer(logger.getReportableSeverity(), Severity::kERROR);
+}
+
+//!
+//! \brief produces a LogStreamConsumer object that can be used to log messages of severity kINTERNAL_ERROR
+//         ("fatal" severity)
+//!
+//! Example usage:
+//!
+//!     LOG_FATAL(logger) << "hello world" << std::endl;
+//!
+inline LogStreamConsumer LOG_FATAL(const Logger& logger) {
+    return LogStreamConsumer(logger.getReportableSeverity(), Severity::kINTERNAL_ERROR);
+}
+
+}  // anonymous namespace
+
+#endif  // TENSORRT_LOGGING_H

real-esrgan/general-x4v3/src/include/pixel_shuffle/pixel_shuffle.hpp

@@ -0,0 +1,156 @@
+#ifndef REAL_ESRGAN_TRT_PIXEL_SHUFFLE_HPP
+#define REAL_ESRGAN_TRT_PIXEL_SHUFFLE_HPP
+
+#include <string>
+#include <vector>
+#include "NvInfer.h"
+
+class PixelShufflePlugin : public nvinfer1::IPluginV2DynamicExt {
+   public:
+    PixelShufflePlugin(int upscaleFactor) : mUpscaleFactor(upscaleFactor) {}
+
+    PixelShufflePlugin(const void* data, size_t length) { memcpy(&mUpscaleFactor, data, sizeof(mUpscaleFactor)); }
+
+    const char* getPluginType() const noexcept override { return "PixelShufflePlugin"; }
+
+    const char* getPluginVersion() const noexcept override { return "1"; }
+
+    int getNbOutputs() const noexcept override { return 1; }
+
+    // nvinfer1::DimsExprs getOutputDimensions(int outputIndex, const nvinfer1::DimsExprs* inputs, int nbInputs, nvinfer1::IExprBuilder& exprBuilder) noexcept override
+    // {
+    //     assert(outputIndex == 0);
+    //     auto* in = &inputs[0];
+    //     nvinfer1::DimsExprs outputDims = *in;
+    //     int channels = in->d[0];
+    //     int height = in->d[1];
+    //     int width = in->d[2];
+    //     int upscaleFactor = mUpscaleFactor;
+    //     outputDims.d[0] = exprBuilder.constant(channels / (upscaleFactor * upscaleFactor));
+    //     outputDims.d[1] = exprBuilder.operation(nvinfer1::DimensionOperation::kPROD, {height, exprBuilder.constant(upscaleFactor)});
+    //     outputDims.d[2] = exprBuilder.operation(nvinfer1::DimensionOperation::kPROD, {width, exprBuilder.constant(upscaleFactor)});
+    //     return outputDims;
+    // }
+    nvinfer1::DimsExprs getOutputDimensions(int32_t outputIndex, nvinfer1::DimsExprs const* inputs, int32_t nbInputs,
+                                            nvinfer1::IExprBuilder& exprBuilder) noexcept override {
+        // assert(nbInputs == 1);
+        auto inDims = inputs[0];
+        // assert(inDims.nbDims == 4);
+        int c = inDims.d[1]->getConstantValue() / (mUpscaleFactor * mUpscaleFactor);
+        int h = inDims.d[2]->getConstantValue() * mUpscaleFactor;
+        int w = inDims.d[3]->getConstantValue() * mUpscaleFactor;
+        nvinfer1::DimsExprs outDims;
+        outDims.nbDims = 4;
+        outDims.d[0] = inDims.d[0];
+        outDims.d[1] = exprBuilder.constant(c);
+        outDims.d[2] = exprBuilder.constant(h);
+        outDims.d[3] = exprBuilder.constant(w);
+        return outDims;
+    }
+
+    bool supportsFormatCombination(int pos, const nvinfer1::PluginTensorDesc* inOut, int nbInputs,
+                                   int nbOutputs) noexcept override {
+        return inOut[pos].format == nvinfer1::TensorFormat::kLINEAR && inOut[pos].type == nvinfer1::DataType::kFLOAT;
+    }
+
+    nvinfer1::DataType getOutputDataType(int index, const nvinfer1::DataType* inputTypes,
+                                         int nbInputs) const noexcept override {
+
+        return inputTypes[0];
+    }
+
+    // bool canBroadcastInputAcrossBatch(int inputIndex) const noexcept override
+    // {
+    //     return false;
+    // }
+
+    void configurePlugin(const nvinfer1::DynamicPluginTensorDesc* inputs, int nbInputs,
+                         const nvinfer1::DynamicPluginTensorDesc* outputs, int nbOutputs) noexcept override {}
+
+    // void configurePlugin(const nvinfer1::DynamicPluginTensorDesc* in, int nbInputs, const nvinfer1::DynamicPluginTensorDesc* out, int nbOutputs) noexcept override
+    // {
+    //     // Optionally configure plugin if necessary
+    // }
+
+    size_t getWorkspaceSize(const nvinfer1::PluginTensorDesc* inputs, int nbInputs,
+                            const nvinfer1::PluginTensorDesc* outputs, int nbOutputs) const noexcept override {
+        return 0;
+    }
+
+    size_t getSerializationSize() const noexcept override { return sizeof(mUpscaleFactor); }
+
+    void serialize(void* buffer) const noexcept override { memcpy(buffer, &mUpscaleFactor, sizeof(mUpscaleFactor)); }
+
+    void destroy() noexcept override {
+        // delete this;
+    }
+
+    nvinfer1::IPluginV2DynamicExt* clone() const noexcept override { return new PixelShufflePlugin(mUpscaleFactor); }
+
+    void setPluginNamespace(const char* pluginNamespace) noexcept override { mNamespace = pluginNamespace; }
+
+    const char* getPluginNamespace() const noexcept override { return mNamespace.c_str(); }
+
+    int initialize() noexcept override { return 0; }
+
+    int32_t enqueue(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc,
+                    void const* const* inputs, void* const* outputs, void* workspace,
+                    cudaStream_t stream) noexcept override;
+
+    void terminate() noexcept override {}
+
+   private:
+    int mUpscaleFactor;
+    std::string mNamespace;
+};
+
+class PixelShufflePluginCreator : public nvinfer1::IPluginCreator {
+   public:
+    PixelShufflePluginCreator() {
+        mPluginAttributes.clear();
+        mPluginAttributes.emplace_back(
+                nvinfer1::PluginField("upscaleFactor", nullptr, nvinfer1::PluginFieldType::kINT32, 1));
+        mFC.nbFields = mPluginAttributes.size();
+        mFC.fields = mPluginAttributes.data();
+    }
+
+    ~PixelShufflePluginCreator() override = default;
+
+    const char* getPluginName() const noexcept override { return "PixelShufflePlugin"; }
+
+    const char* getPluginVersion() const noexcept override { return "1"; }
+
+    const nvinfer1::PluginFieldCollection* getFieldNames() noexcept override { return &mFC; }
+
+    nvinfer1::IPluginV2* createPlugin(const char* name, const nvinfer1::PluginFieldCollection* fc) noexcept override {
+        int upscaleFactor = 0;
+        for (int i = 0; i < fc->nbFields; ++i) {
+            if (strcmp(fc->fields[i].name, "upscaleFactor") == 0) {
+                upscaleFactor = *static_cast<const int*>(fc->fields[i].data);
+            }
+        }
+        return new PixelShufflePlugin(upscaleFactor);
+    }
+
+    nvinfer1::IPluginV2* deserializePlugin(const char* name, const void* serialData,
+                                           size_t serialLength) noexcept override {
+        return new PixelShufflePlugin(serialData, serialLength);
+    }
+
+    void setPluginNamespace(const char* pluginNamespace) noexcept override { mNamespace = pluginNamespace; }
+
+    const char* getPluginNamespace() const noexcept override { return mNamespace.c_str(); }
+
+   private:
+    static nvinfer1::PluginFieldCollection mFC;
+    static std::vector<nvinfer1::PluginField> mPluginAttributes;
+    std::string mNamespace;
+};
+
+nvinfer1::PluginFieldCollection PixelShufflePluginCreator::mFC{};
+std::vector<nvinfer1::PluginField> PixelShufflePluginCreator::mPluginAttributes{
+        nvinfer1::PluginField{"upscaleFactor", nullptr, nvinfer1::PluginFieldType::kINT32, 1}};
+
+REGISTER_TENSORRT_PLUGIN(PixelShufflePluginCreator);
+
+#endif  //REAL_ESRGAN_TRT_PIXEL_SHUFFLE_HPP

real-esrgan/general-x4v3/src/include/preprocess/preprocess.hpp

@@ -0,0 +1,11 @@
+#ifndef REAL_ESRGAN_TRT_PREPROCESS_HPP
+#define REAL_ESRGAN_TRT_PREPROCESS_HPP
+
+struct PreprocessStruct {
+    int N;
+    int C;
+    int H;
+    int W;
+};
+
+#endif  //REAL_ESRGAN_TRT_PREPROCESS_HPP

real-esrgan/general-x4v3/src/pixel_shuffle/pixel_shuffle.cpp

@@ -0,0 +1,128 @@
+// PixelShufflePlugin.cpp
+//
+// #include "pixel_shuffle/pixel_shuffle.hpp"
+// #include <cstring>
+// #include <cassert>
+//
+// PixelShufflePlugin::PixelShufflePlugin(int upscaleFactor)
+//         : mUpscaleFactor(upscaleFactor) {
+//     // Initialize other members
+// }
+//
+// PixelShufflePlugin::PixelShufflePlugin(const void* data, size_t length) {
+//     // Deserialize data to initialize members
+//     const char* d = static_cast<const char*>(data);
+//     mUpscaleFactor = *reinterpret_cast<const int*>(d);
+//     d += sizeof(int);
+//     mInputVolume = *reinterpret_cast<const size_t*>(d);
+//     d += sizeof(size_t);
+//     mOutputVolume = *reinterpret_cast<const size_t*>(d);
+// }
+//
+// int PixelShufflePlugin::getNbOutputs() const {
+//     return 1;
+// }
+//
+// nvinfer1::Dims PixelShufflePlugin::getOutputDimensions(int index, const nvinfer1::Dims* inputs, int nbInputDims) {
+//     assert(index == 0);
+//     assert(nbInputDims == 1);
+//     int c = inputs[0].d[0];
+//     int h = inputs[0].d[1];
+//     int w = inputs[0].d[2];
+//     int upscaleFactor = mUpscaleFactor;
+//
+//     assert(c % (upscaleFactor * upscaleFactor) == 0);
+//     int newC = c / (upscaleFactor * upscaleFactor);
+//     int newH = h * upscaleFactor;
+//     int newW = w * upscaleFactor;
+//
+//     return nvinfer1::Dims3(newC, newH, newW);
+// }
+//
+// int PixelShufflePlugin::initialize() {
+//     return 0;
+// }
+//
+// void PixelShufflePlugin::terminate() {
+//     // Clean up
+// }
+//
+// size_t PixelShufflePlugin::getWorkspaceSize(int maxBatchSize) const {
+//     return 0;
+// }
+//
+// int PixelShufflePlugin::enqueue(int batchSize, const void* const* inputs, void** outputs, void* workspace, cudaStream_t stream) {
+//     // Launch CUDA kernel for PixelShuffle
+//     // Assume inputs[0] and outputs[0] are pointers to device memory
+//     const float* input = static_cast<const float*>(inputs[0]);
+//     float* output = static_cast<float*>(outputs[0]);
+//
+//     int c = mInputVolume / (mUpscaleFactor * mUpscaleFactor);
+//     int h = mOutputVolume / (c * mUpscaleFactor);
+//     int w = h; // Assuming square input for simplicity
+//     int upscaleFactor = mUpscaleFactor;
+//
+//     // Launch CUDA kernel (to be implemented)
+//     // pixelShuffleKernel(input, output, c, h, w, upscaleFactor, stream);
+//
+//     return 0;
+// }
+//
+// size_t PixelShufflePlugin::getSerializationSize() const {
+//     return sizeof(int) + sizeof(size_t) * 2;
+// }
+//
+// void PixelShufflePlugin::serialize(void* buffer) const {
+//     char* d = static_cast<char*>(buffer);
+//     *reinterpret_cast<int*>(d) = mUpscaleFactor;
+//     d += sizeof(int);
+//     *reinterpret_cast<size_t*>(d) = mInputVolume;
+//     d += sizeof(size_t);
+//     *reinterpret_cast<size_t*>(d) = mOutputVolume;
+// }
+//
+// void PixelShufflePlugin::destroy() {
+//     delete this;
+// }
+//
+// const char* PixelShufflePlugin::getPluginType() const {
+//     return "PixelShufflePlugin";
+// }
+//
+// const char* PixelShufflePlugin::getPluginVersion() const {
+//     return "1";
+// }
+//
+// void PixelShufflePlugin::setPluginNamespace(const char* pluginNamespace) {
+//     mPluginNamespace = pluginNamespace;
+// }
+//
+// const char* PixelShufflePlugin::getPluginNamespace() const {
+//     return mPluginNamespace;
+// }
+//
+// nvinfer1::IPluginV2IOExt* PixelShufflePlugin::clone() const {
+//     return new PixelShufflePlugin(mUpscaleFactor);
+// }
+//
+// bool PixelShufflePlugin::supportsFormatCombination(int pos, const nvinfer1::PluginTensorDesc* inOut, int nbInputs, int nbOutputs) const {
+//     return inOut[pos].format == nvinfer1::TensorFormat::kLINEAR && inOut[pos].type == nvinfer1::DataType::kFLOAT;
+// }
+//
+// void PixelShufflePlugin::configurePlugin(const nvinfer1::DynamicPluginTensorDesc* in, int nbInputs, const nvinfer1::DynamicPluginTensorDesc* out, int nbOutputs) {
+//     // Configure the plugin based on the input and output descriptions
+//     mInputVolume = in[0].desc.volume();
+//     mOutputVolume = out[0].desc.volume();
+// }
+//
+// nvinfer1::DataType PixelShufflePlugin::getOutputDataType(int index, const nvinfer1::DataType* inputTypes, int nbInputs) const {
+//     return inputTypes[0];
+// }
+//
+// bool PixelShufflePlugin::isOutputBroadcastAcrossBatch(int outputIndex, const bool* inputIsBroadcasted, int nbInputs) const {
+//     return false;
+// }
+//
+// bool PixelShufflePlugin::canBroadcastInputAcrossBatch(int inputIndex) const {
+//     return false;
+// }

real-esrgan/general-x4v3/src/pixel_shuffle/pixel_shuffle.cu

@@ -0,0 +1,50 @@
+#include <cuda_runtime.h>
+#include <string>
+#include "pixel_shuffle/pixel_shuffle.hpp"
+
+// CUDA kernel for PixelShuffle
+__global__ void PixelShuffleKernel(const float* input, float* output, int batchSize, int channels, int height,
+                                   int width, int upscaleFactor) {
+    int outHeight = height * upscaleFactor;
+    int outWidth = width * upscaleFactor;
+    int outChannels = channels / (upscaleFactor * upscaleFactor);
+
+    int idx = blockIdx.x * blockDim.x + threadIdx.x;
+    if (idx >= batchSize * outChannels * outHeight * outWidth)
+        return;
+
+    int out_w = idx % outWidth;
+    int out_h = (idx / outWidth) % outHeight;
+    int out_c = (idx / outWidth / outHeight) % outChannels;
+    int b = idx / (outWidth * outHeight * outChannels);
+
+    int in_c =
+            out_c * upscaleFactor * upscaleFactor + (out_h % upscaleFactor) * upscaleFactor + (out_w % upscaleFactor);
+    int in_h = out_h / upscaleFactor;
+    int in_w = out_w / upscaleFactor;
+
+    output[idx] = input[((b * channels + in_c) * height + in_h) * width + in_w];
+}
+
+int32_t PixelShufflePlugin::enqueue(nvinfer1::PluginTensorDesc const* inputDesc,
+                                    nvinfer1::PluginTensorDesc const* outputDesc, void const* const* inputs,
+                                    void* const* outputs, void* workspace, cudaStream_t stream) noexcept {
+    const float* input = static_cast<const float*>(inputs[0]);
+    float* output = static_cast<float*>(outputs[0]);
+
+    int batchSize = inputDesc[0].dims.d[0];
+    int channels = inputDesc[0].dims.d[1];
+    int height = inputDesc[0].dims.d[2];
+    int width = inputDesc[0].dims.d[3];
+    int upscaleFactor = mUpscaleFactor;
+
+    int outChannels = channels / (upscaleFactor * upscaleFactor);
+    int outHeight = height * upscaleFactor;
+    int outWidth = width * upscaleFactor;
+
+    int numElements = batchSize * outChannels * outHeight * outWidth;
+
+    PixelShuffleKernel<<<(numElements + 255) / 256, 256>>>(input, output, batchSize, channels, height, width,
+                                                           upscaleFactor);
+    return cudaGetLastError() != cudaSuccess;
+}