From 56ea9d944a35510ff9ee5c0d986fad0da2027ef5 Mon Sep 17 00:00:00 2001 From: Lemonon <57012106+Lemonononon@users.noreply.github.com> Date: Wed, 21 Aug 2024 11:38:52 +0800 Subject: [PATCH] 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 --- real-esrgan/general-x4v3/CMakeLists.txt | 42 ++ real-esrgan/general-x4v3/README.md | 41 ++ .../general-x4v3/cmake/FindTensorRT.cmake | 87 ++++ real-esrgan/general-x4v3/gen_wts.py | 136 ++++++ real-esrgan/general-x4v3/main.cpp | 342 +++++++++++++ .../src/include/config/config.hpp | 22 + .../general-x4v3/src/include/cuda_utils.h | 21 + .../src/include/logging/logging.h | 455 ++++++++++++++++++ .../include/pixel_shuffle/pixel_shuffle.hpp | 156 ++++++ .../src/include/preprocess/preprocess.hpp | 11 + .../src/pixel_shuffle/pixel_shuffle.cpp | 128 +++++ .../src/pixel_shuffle/pixel_shuffle.cu | 50 ++ real-esrgan/{ => x4plus}/CMakeLists.txt | 0 real-esrgan/{ => x4plus}/README.md | 0 real-esrgan/{ => x4plus}/common.hpp | 0 real-esrgan/{ => x4plus}/cuda_utils.h | 0 real-esrgan/{ => x4plus}/gen_wts.py | 0 real-esrgan/{ => x4plus}/logging.h | 0 real-esrgan/{ => x4plus}/macros.h | 0 real-esrgan/{ => x4plus}/postprocess.cu | 0 real-esrgan/{ => x4plus}/postprocess.hpp | 0 real-esrgan/{ => x4plus}/preprocess.cu | 0 real-esrgan/{ => x4plus}/preprocess.hpp | 0 real-esrgan/{ => x4plus}/real-esrgan.cpp | 0 real-esrgan/{ => x4plus}/utils.h | 0 25 files changed, 1491 insertions(+) create mode 100644 real-esrgan/general-x4v3/CMakeLists.txt create mode 100644 real-esrgan/general-x4v3/README.md create mode 100644 real-esrgan/general-x4v3/cmake/FindTensorRT.cmake create mode 100644 real-esrgan/general-x4v3/gen_wts.py create mode 100644 real-esrgan/general-x4v3/main.cpp create mode 100644 real-esrgan/general-x4v3/src/include/config/config.hpp create mode 100644 real-esrgan/general-x4v3/src/include/cuda_utils.h create mode 100644 real-esrgan/general-x4v3/src/include/logging/logging.h create mode 100644 real-esrgan/general-x4v3/src/include/pixel_shuffle/pixel_shuffle.hpp create mode 100644 real-esrgan/general-x4v3/src/include/preprocess/preprocess.hpp create mode 100644 real-esrgan/general-x4v3/src/pixel_shuffle/pixel_shuffle.cpp create mode 100644 real-esrgan/general-x4v3/src/pixel_shuffle/pixel_shuffle.cu rename real-esrgan/{ => x4plus}/CMakeLists.txt (100%) rename real-esrgan/{ => x4plus}/README.md (100%) rename real-esrgan/{ => x4plus}/common.hpp (100%) rename real-esrgan/{ => x4plus}/cuda_utils.h (100%) rename real-esrgan/{ => x4plus}/gen_wts.py (100%) rename real-esrgan/{ => x4plus}/logging.h (100%) rename real-esrgan/{ => x4plus}/macros.h (100%) rename real-esrgan/{ => x4plus}/postprocess.cu (100%) rename real-esrgan/{ => x4plus}/postprocess.hpp (100%) rename real-esrgan/{ => x4plus}/preprocess.cu (100%) rename real-esrgan/{ => x4plus}/preprocess.hpp (100%) rename real-esrgan/{ => x4plus}/real-esrgan.cpp (100%) rename real-esrgan/{ => x4plus}/utils.h (100%) diff --git a/real-esrgan/general-x4v3/CMakeLists.txt b/real-esrgan/general-x4v3/CMakeLists.txt new file mode 100644 index 0000000..bcb06d5 --- /dev/null +++ b/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) diff --git a/real-esrgan/general-x4v3/README.md b/real-esrgan/general-x4v3/README.md new file mode 100644 index 0000000..4aa674b --- /dev/null +++ b/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 +``` diff --git a/real-esrgan/general-x4v3/cmake/FindTensorRT.cmake b/real-esrgan/general-x4v3/cmake/FindTensorRT.cmake new file mode 100644 index 0000000..abe5d0a --- /dev/null +++ b/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() diff --git a/real-esrgan/general-x4v3/gen_wts.py b/real-esrgan/general-x4v3/gen_wts.py new file mode 100644 index 0000000..2b52368 --- /dev/null +++ b/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() diff --git a/real-esrgan/general-x4v3/main.cpp b/real-esrgan/general-x4v3/main.cpp new file mode 100644 index 0000000..934d2a5 --- /dev/null +++ b/real-esrgan/general-x4v3/main.cpp @@ -0,0 +1,342 @@ +#include +#include +#include +#include +#include +#include +#include + +#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] +std::map loadWeights(const std::string file) { + std::cout << "Loading weights: " << file << std::endl; + std::map 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(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& 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 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 pluginData; + int upscaleFactor = 4; + pluginData.emplace_back(PluginField{"upscaleFactor", &upscaleFactor, PluginFieldType::kINT32, 1}); + PluginFieldCollection pluginFCWithData = {static_cast(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(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& 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] << " " << 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 data; + std::vector output; + //std::vector 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 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(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(b2); // B + uc_pixel[1] = static_cast(g2); // G + uc_pixel[2] = static_cast(r2); // R + + // uc_pixel[0] = static_cast(std::round(output[b * OUTPUT_C * OUTPUT_H * OUTPUT_W + i + 2 * OUTPUT_H * OUTPUT_W] * 255.0)); // B + // uc_pixel[1] = static_cast(std::round(output[b * OUTPUT_C * OUTPUT_H * OUTPUT_W + i + 1 * OUTPUT_H * OUTPUT_W] * 255.0)); // G + // uc_pixel[2] = static_cast(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; +} diff --git a/real-esrgan/general-x4v3/src/include/config/config.hpp b/real-esrgan/general-x4v3/src/include/config/config.hpp new file mode 100644 index 0000000..e7ed3dc --- /dev/null +++ b/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 + +//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 diff --git a/real-esrgan/general-x4v3/src/include/cuda_utils.h b/real-esrgan/general-x4v3/src/include/cuda_utils.h new file mode 100644 index 0000000..a9ca611 --- /dev/null +++ b/real-esrgan/general-x4v3/src/include/cuda_utils.h @@ -0,0 +1,21 @@ +#ifndef TRTX_CUDA_UTILS_H_ +#define TRTX_CUDA_UTILS_H_ + +#include +#include +#include +#include +#include + +#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_ diff --git a/real-esrgan/general-x4v3/src/include/logging/logging.h b/real-esrgan/general-x4v3/src/include/logging/logging.h new file mode 100644 index 0000000..0aa1221 --- /dev/null +++ b/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 +#include +#include +#include +#include +#include +#include +#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(×tamp); + 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 diff --git a/real-esrgan/general-x4v3/src/include/pixel_shuffle/pixel_shuffle.hpp b/real-esrgan/general-x4v3/src/include/pixel_shuffle/pixel_shuffle.hpp new file mode 100644 index 0000000..28a2e38 --- /dev/null +++ b/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 +#include +#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(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 mPluginAttributes; + std::string mNamespace; +}; + +nvinfer1::PluginFieldCollection PixelShufflePluginCreator::mFC{}; +std::vector PixelShufflePluginCreator::mPluginAttributes{ + nvinfer1::PluginField{"upscaleFactor", nullptr, nvinfer1::PluginFieldType::kINT32, 1}}; + +REGISTER_TENSORRT_PLUGIN(PixelShufflePluginCreator); + +#endif //REAL_ESRGAN_TRT_PIXEL_SHUFFLE_HPP diff --git a/real-esrgan/general-x4v3/src/include/preprocess/preprocess.hpp b/real-esrgan/general-x4v3/src/include/preprocess/preprocess.hpp new file mode 100644 index 0000000..7d627c8 --- /dev/null +++ b/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 diff --git a/real-esrgan/general-x4v3/src/pixel_shuffle/pixel_shuffle.cpp b/real-esrgan/general-x4v3/src/pixel_shuffle/pixel_shuffle.cpp new file mode 100644 index 0000000..21250ad --- /dev/null +++ b/real-esrgan/general-x4v3/src/pixel_shuffle/pixel_shuffle.cpp @@ -0,0 +1,128 @@ +// PixelShufflePlugin.cpp +// +// #include "pixel_shuffle/pixel_shuffle.hpp" +// #include +// #include +// +// 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(data); +// mUpscaleFactor = *reinterpret_cast(d); +// d += sizeof(int); +// mInputVolume = *reinterpret_cast(d); +// d += sizeof(size_t); +// mOutputVolume = *reinterpret_cast(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(inputs[0]); +// float* output = static_cast(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(buffer); +// *reinterpret_cast(d) = mUpscaleFactor; +// d += sizeof(int); +// *reinterpret_cast(d) = mInputVolume; +// d += sizeof(size_t); +// *reinterpret_cast(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; +// } diff --git a/real-esrgan/general-x4v3/src/pixel_shuffle/pixel_shuffle.cu b/real-esrgan/general-x4v3/src/pixel_shuffle/pixel_shuffle.cu new file mode 100644 index 0000000..0b80447 --- /dev/null +++ b/real-esrgan/general-x4v3/src/pixel_shuffle/pixel_shuffle.cu @@ -0,0 +1,50 @@ +#include +#include +#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(inputs[0]); + float* output = static_cast(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; +} diff --git a/real-esrgan/CMakeLists.txt b/real-esrgan/x4plus/CMakeLists.txt similarity index 100% rename from real-esrgan/CMakeLists.txt rename to real-esrgan/x4plus/CMakeLists.txt diff --git a/real-esrgan/README.md b/real-esrgan/x4plus/README.md similarity index 100% rename from real-esrgan/README.md rename to real-esrgan/x4plus/README.md diff --git a/real-esrgan/common.hpp b/real-esrgan/x4plus/common.hpp similarity index 100% rename from real-esrgan/common.hpp rename to real-esrgan/x4plus/common.hpp diff --git a/real-esrgan/cuda_utils.h b/real-esrgan/x4plus/cuda_utils.h similarity index 100% rename from real-esrgan/cuda_utils.h rename to real-esrgan/x4plus/cuda_utils.h diff --git a/real-esrgan/gen_wts.py b/real-esrgan/x4plus/gen_wts.py similarity index 100% rename from real-esrgan/gen_wts.py rename to real-esrgan/x4plus/gen_wts.py diff --git a/real-esrgan/logging.h b/real-esrgan/x4plus/logging.h similarity index 100% rename from real-esrgan/logging.h rename to real-esrgan/x4plus/logging.h diff --git a/real-esrgan/macros.h b/real-esrgan/x4plus/macros.h similarity index 100% rename from real-esrgan/macros.h rename to real-esrgan/x4plus/macros.h diff --git a/real-esrgan/postprocess.cu b/real-esrgan/x4plus/postprocess.cu similarity index 100% rename from real-esrgan/postprocess.cu rename to real-esrgan/x4plus/postprocess.cu diff --git a/real-esrgan/postprocess.hpp b/real-esrgan/x4plus/postprocess.hpp similarity index 100% rename from real-esrgan/postprocess.hpp rename to real-esrgan/x4plus/postprocess.hpp diff --git a/real-esrgan/preprocess.cu b/real-esrgan/x4plus/preprocess.cu similarity index 100% rename from real-esrgan/preprocess.cu rename to real-esrgan/x4plus/preprocess.cu diff --git a/real-esrgan/preprocess.hpp b/real-esrgan/x4plus/preprocess.hpp similarity index 100% rename from real-esrgan/preprocess.hpp rename to real-esrgan/x4plus/preprocess.hpp diff --git a/real-esrgan/real-esrgan.cpp b/real-esrgan/x4plus/real-esrgan.cpp similarity index 100% rename from real-esrgan/real-esrgan.cpp rename to real-esrgan/x4plus/real-esrgan.cpp diff --git a/real-esrgan/utils.h b/real-esrgan/x4plus/utils.h similarity index 100% rename from real-esrgan/utils.h rename to real-esrgan/x4plus/utils.h