main.cpp

#include <iostream>
#include <fstream>
#include <vector>
#include <memory>
#include <random>
#include <cmath>
#include <chrono>
#include <opencv2/opencv.hpp>
#include <cuda_runtime_api.h>
#include <NvInfer.h>

using namespace nvinfer1;

// Logger class for TensorRT
class TRTLogger : public ILogger {
    void log(Severity severity, const char* msg) noexcept override {
        if (severity <= Severity::kINFO) {
            std::cout << "[TensorRT] " << msg << std::endl;
        }
    }
} gLogger;

// RAII wrappers for CUDA memory
template <typename T>
struct CudaDeleter {
    void operator()(T* ptr) const {
        if (ptr) cudaFree(ptr);
    }
};

template <typename T>
using CudaUniquePtr = std::unique_ptr<T, CudaDeleter<T>>;

template <typename T>
CudaUniquePtr<T> make_cuda_unique(size_t size) {
    T* ptr = nullptr;
    if (cudaMalloc((void**)&ptr, size * sizeof(T)) != cudaSuccess) {
        throw std::runtime_error("Failed to allocate CUDA memory!");
    }
    return CudaUniquePtr<T>(ptr, CudaDeleter<T>());
}

// Struct to hold a TensorRT engine and execution context
struct TRTEngine {
    std::shared_ptr<IRuntime> runtime;
    std::shared_ptr<ICudaEngine> engine;
    std::shared_ptr<IExecutionContext> context;

    bool load(const std::string& path) {
        std::ifstream file(path, std::ios::binary | std::ios::ate);
        if (!file.good()) {
            std::cerr << "Error: Unable to open engine file: " << path << std::endl;
            return false;
        }
        size_t size = file.tellg();
        file.seekg(0, std::ios::beg);
        std::vector<char> modelData(size);
        file.read(modelData.data(), size);

        runtime = std::shared_ptr<IRuntime>(createInferRuntime(gLogger));
        if (!runtime) return false;

        engine = std::shared_ptr<ICudaEngine>(runtime->deserializeCudaEngine(modelData.data(), size));
        if (!engine) return false;

        context = std::shared_ptr<IExecutionContext>(engine->createExecutionContext());
        if (!context) return false;

        std::cout << "Engine has " << engine->getNbIOTensors() << " IO tensors:" << std::endl;
        for (int i = 0; i < engine->getNbIOTensors(); ++i) {
            const char* name = engine->getIOTensorName(i);
            TensorIOMode mode = engine->getTensorIOMode(name);
            Dims dims = engine->getTensorShape(name);
            DataType type = engine->getTensorDataType(name);
            std::cout << "  - Tensor " << i << ": Name=" << name
                      << ", Mode=" << (mode == TensorIOMode::kINPUT ? "INPUT" : "OUTPUT")
                      << ", Shape=";
            for (int d = 0; d < dims.nbDims; ++d) {
                std::cout << dims.d[d] << (d == dims.nbDims - 1 ? "" : "x");
            }
            std::cout << ", Type=" << static_cast<int>(type) << std::endl;
        }

        return true;
    }
};

// Fill a host vector with normal random values matching PyTorch's seed behavior
void fill_normal(std::vector<float>& vec, float mean, float stddev, unsigned int seed) {
    std::mt19937 gen(seed);
    std::normal_distribution<float> dist(mean, stddev);
    for (auto& val : vec) {
        val = dist(gen);
    }
}

int main(int argc, char** argv) {
    if (argc < 4) {
        std::cout << "Usage: " << argv[0] << " <input_image> <vae_engine_path> <pid_engine_path> [output_image] [seed]" << std::endl;
        return -1;
    }

    std::string inputPath = argv[1];
    std::string vaeEnginePath = argv[2];
    std::string pidEnginePath = argv[3];
    std::string outputPath = (argc >= 5) ? argv[4] : "upscaled_output.png";
    unsigned int seed = (argc >= 6) ? std::stoul(argv[5]) : 42;

    std::cout << "---------------------------------------------------" << std::endl;
    std::cout << "Starting C++ TensorRT PiD Upscaler..." << std::endl;
    std::cout << "Input path: " << inputPath << std::endl;
    std::cout << "VAE Engine path: " << vaeEnginePath << std::endl;
    std::cout << "PiD Engine path: " << pidEnginePath << std::endl;
    std::cout << "Seed: " << seed << std::endl;
    std::cout << "---------------------------------------------------" << std::endl;

    // 1. Load Engines
    TRTEngine vaeModel;
    if (!vaeModel.load(vaeEnginePath)) {
        std::cerr << "Failed to load VAE Engine!" << std::endl;
        return -1;
    }
    std::cout << "[OK] Loaded SDXL VAE Encoder engine." << std::endl;

    TRTEngine pidModel;
    if (!pidModel.load(pidEnginePath)) {
        std::cerr << "Failed to load PiD Engine!" << std::endl;
        return -1;
    }
    std::cout << "[OK] Loaded PiD Pixel Diffusion engine." << std::endl;

    // 2. Load and preprocess input image using OpenCV
    cv::Mat inputImg = cv::imread(inputPath, cv::IMREAD_COLOR);
    if (inputImg.empty()) {
        std::cerr << "Failed to load image: " << inputPath << std::endl;
        return -1;
    }

    // Force resize to 1024x1024 if needed
    if (inputImg.rows != 1024 || inputImg.cols != 1024) {
        std::cout << "Resizing input image from " << inputImg.cols << "x" << inputImg.rows << " to 1024x1024." << std::endl;
        cv::resize(inputImg, inputImg, cv::Size(1024, 1024), 0, 0, cv::INTER_CUBIC);
    }

    // Convert BGR (OpenCV default) to RGB, then normalize to [-1, 1], and HWC -> CHW
    cv::Mat inputImgRGB;
    cv::cvtColor(inputImg, inputImgRGB, cv::COLOR_BGR2RGB);

    std::vector<float> inputBuffer(1 * 3 * 1024 * 1024);
    for (int c = 0; c < 3; ++c) {
        for (int h = 0; h < 1024; ++h) {
            for (int w = 0; w < 1024; ++w) {
                float val = static_cast<float>(inputImgRGB.at<cv::Vec3b>(h, w)[c]);
                // Scale to [-1, 1] matching PyTorch VAE input
                inputBuffer[c * 1024 * 1024 + h * 1024 + w] = (val / 127.5f) - 1.0f;
            }
        }
    }

    // Print preprocessed input stats to make sure we don't have NaN/Inf in inputBuffer
    double in_sum = 0, in_sq_sum = 0;
    for (float val : inputBuffer) {
        if (std::isnan(val) || std::isinf(val)) {
            std::cerr << "[WARNING] Found NaN/Inf in inputBuffer!" << std::endl;
        }
        in_sum += val;
        in_sq_sum += val * val;
    }
    double in_mean = in_sum / inputBuffer.size();
    double in_std = std::sqrt(in_sq_sum / inputBuffer.size() - in_mean * in_mean);
    std::cout << "[DEBUG] Preprocessed VAE input stats: Mean = " << in_mean << ", Std = " << in_std << std::endl;
    std::cout << "[DEBUG] Corner pixels: [0,0] RGB = (" 
              << (int)inputImgRGB.at<cv::Vec3b>(0, 0)[0] << ", " 
              << (int)inputImgRGB.at<cv::Vec3b>(0, 0)[1] << ", " 
              << (int)inputImgRGB.at<cv::Vec3b>(0, 0)[2] << ") -> normalized: ("
              << inputBuffer[0] << ", " << inputBuffer[1024*1024] << ", " << inputBuffer[2*1024*1024] << ")" << std::endl;

    // 3. Allocate CUDA memory & run VAE Encoder
    size_t vaeInputSize = 1 * 3 * 1024 * 1024;
    size_t vaeOutputSize = 1 * 4 * 128 * 128;
    auto d_vaeInput = make_cuda_unique<float>(vaeInputSize);
    auto d_vaeOutput = make_cuda_unique<float>(vaeOutputSize);

    // Copy VAE inputs
    cudaMemcpy(d_vaeInput.get(), inputBuffer.data(), vaeInputSize * sizeof(float), cudaMemcpyHostToDevice);

    // Run VAE Inference
    void* vaeBindings[] = { d_vaeInput.get(), d_vaeOutput.get() };
    cudaStream_t stream;
    cudaStreamCreate(&stream);

    std::cout << "Running VAE Encoder..." << std::endl;
    // Use enqueueV3 bindings approach
    vaeModel.context->setInputTensorAddress("image", d_vaeInput.get());
    vaeModel.context->setOutputTensorAddress("latent", d_vaeOutput.get());
    
    auto t_start = std::chrono::high_resolution_clock::now();
    vaeModel.context->enqueueV3(stream);
    cudaStreamSynchronize(stream);
    auto t_end = std::chrono::high_resolution_clock::now();
    std::cout << "[OK] VAE Encoder finished in " 
              << std::chrono::duration_cast<std::chrono::milliseconds>(t_end - t_start).count() 
              << " ms." << std::endl;

    // Debug print VAE output stats
    std::vector<float> h_vaeOutput(vaeOutputSize);
    cudaMemcpy(h_vaeOutput.data(), d_vaeOutput.get(), vaeOutputSize * sizeof(float), cudaMemcpyDeviceToHost);
    double vae_sum = 0, vae_sq_sum = 0;
    for (float val : h_vaeOutput) {
        vae_sum += val;
        vae_sq_sum += val * val;
    }
    double vae_mean = vae_sum / vaeOutputSize;
    double vae_std = std::sqrt(vae_sq_sum / vaeOutputSize - vae_mean * vae_mean);
    std::cout << "[DEBUG] VAE Latents Stats: Mean = " << vae_mean << ", Std = " << vae_std << std::endl;

    // 4. Initialize PiD Loop
    // PiD output is 4096x4096px
    size_t pidImageSize = 1 * 3 * 4096 * 4096;
    std::vector<float> h_x(pidImageSize);
    std::cout << "Initializing starting noise x_T with seed: " << seed << std::endl;
    fill_normal(h_x, 0.0f, 1.0f, seed); // Initial Gaussian Noise

    double x_sum = 0, x_sq_sum = 0;
    for (float val : h_x) {
        x_sum += val;
        x_sq_sum += val * val;
    }
    double x_mean = x_sum / pidImageSize;
    double x_std = std::sqrt(x_sq_sum / pidImageSize - x_mean * x_mean);
    std::cout << "[DEBUG] Starting Noise x_T Stats: Mean = " << x_mean << ", Std = " << x_std << std::endl;

    // Setup PiD Tensors
    auto d_x = make_cuda_unique<float>(pidImageSize);
    auto d_v_pred = make_cuda_unique<float>(pidImageSize);
    auto d_t = make_cuda_unique<float>(1);
    auto d_degrade_sigma = make_cuda_unique<float>(1);

    // Copy initial noise to device
    cudaMemcpy(d_x.get(), h_x.data(), pidImageSize * sizeof(float), cudaMemcpyHostToDevice);

    // Set constant degrade_sigma = 0.0f
    float degrade_sigma_val = 0.0f;
    cudaMemcpy(d_degrade_sigma.get(), &degrade_sigma_val, sizeof(float), cudaMemcpyHostToDevice);

    // Setup input/output addresses for PiD
    pidModel.context->setInputTensorAddress("x", d_x.get());
    pidModel.context->setInputTensorAddress("t", d_t.get());
    pidModel.context->setInputTensorAddress("lq_latent", d_vaeOutput.get());
    pidModel.context->setInputTensorAddress("degrade_sigma", d_degrade_sigma.get());
    pidModel.context->setOutputTensorAddress("v_pred", d_v_pred.get());

    // Timeline steps: student_t_list = [0.999, 0.866, 0.634, 0.342, 0.0]
    const float t_list[] = { 0.999f, 0.866f, 0.634f, 0.342f, 0.0f };
    const float timescale = 1000.0f;

    std::vector<float> h_v_pred(pidImageSize);
    std::vector<float> h_eps(pidImageSize);

    std::cout << "Starting 4-step Pixel Diffusion decoding..." << std::endl;
    auto t_loop_start = std::chrono::high_resolution_clock::now();

    for (int step = 0; step < 4; ++step) {
        float t_cur = t_list[step];
        float t_next = t_list[step + 1];
        float t_cur_scaled = t_cur * timescale;

        std::cout << "  - Step " << (step + 1) << "/4: t_cur=" << t_cur << ", t_next=" << t_next << " (scaled_t=" << t_cur_scaled << ")" << std::endl;

        // Copy scaled t to device
        cudaMemcpy(d_t.get(), &t_cur_scaled, sizeof(float), cudaMemcpyHostToDevice);

        // Run PiD inference step
        pidModel.context->enqueueV3(stream);
        cudaStreamSynchronize(stream);

        // Copy current v_pred and image state back to host to calculate SDE update
        cudaMemcpy(h_v_pred.data(), d_v_pred.get(), pidImageSize * sizeof(float), cudaMemcpyDeviceToHost);
        cudaMemcpy(h_x.data(), d_x.get(), pidImageSize * sizeof(float), cudaMemcpyDeviceToHost);

        double v_sum = 0, v_sq_sum = 0;
        double cx_sum = 0, cx_sq_sum = 0;
        for (size_t idx = 0; idx < pidImageSize; ++idx) {
            float v_val = h_v_pred[idx];
            float x_val = h_x[idx];
            v_sum += v_val;
            v_sq_sum += v_val * v_val;
            cx_sum += x_val;
            cx_sq_sum += x_val * x_val;
        }
        double v_mean = v_sum / pidImageSize;
        double v_std = std::sqrt(v_sq_sum / pidImageSize - v_mean * v_mean);
        double cx_mean = cx_sum / pidImageSize;
        double cx_std = std::sqrt(cx_sq_sum / pidImageSize - cx_mean * cx_mean);

        std::cout << "[DEBUG]   Step " << (step + 1) << " before update: x mean = " << cx_mean 
                  << ", std = " << cx_std << " | v_pred mean = " << v_mean << ", std = " << v_std << std::endl;

        // Update math:
        // x0_pred = x - t_cur * v_pred
        // if t_next > 0:
        //     x = (1.0 - t_next) * x0_pred + t_next * eps
        // else:
        //     x = x0_pred
        if (t_next > 0.0f) {
            // Generate random normal eps for SDE path
            fill_normal(h_eps, 0.0f, 1.0f, seed + step + 1);
            for (size_t idx = 0; idx < pidImageSize; ++idx) {
                float x0_pred = h_x[idx] - t_cur * h_v_pred[idx];
                h_x[idx] = (1.0f - t_next) * x0_pred + t_next * h_eps[idx];
            }
        } else {
            for (size_t idx = 0; idx < pidImageSize; ++idx) {
                h_x[idx] = h_x[idx] - t_cur * h_v_pred[idx]; // final step x is just x0_pred
            }
        }

        // Copy updated x back to device for the next step
        cudaMemcpy(d_x.get(), h_x.data(), pidImageSize * sizeof(float), cudaMemcpyHostToDevice);
    }

    auto t_loop_end = std::chrono::high_resolution_clock::now();
    std::cout << "[OK] PiD loop finished in " 
              << std::chrono::duration_cast<std::chrono::milliseconds>(t_loop_end - t_loop_start).count() 
              << " ms." << std::endl;

    // 5. Postprocess and save using OpenCV
    cv::Mat outImg(4096, 4096, CV_8UC3);
    for (int c = 0; c < 3; ++c) {
        for (int h = 0; h < 4096; ++h) {
            for (int w = 0; w < 4096; ++w) {
                float val = h_x[c * 4096 * 4096 + h * 4096 + w];
                // Clamp to [-1, 1]
                val = std::max(-1.0f, std::min(1.0f, val));
                // Convert [-1, 1] to [0, 255]
                uint8_t pixelVal = static_cast<uint8_t>((val + 1.0f) * 127.5f);
                // OpenCV BGR write: OpenCV colors are BGR so c=0 -> R, c=1 -> G, c=2 -> B
                // R is channels index 2, G is channels index 1, B is channels index 0
                outImg.at<cv::Vec3b>(h, w)[2 - c] = pixelVal;
            }
        }
    }

    if (cv::imwrite(outputPath, outImg)) {
        std::cout << "[OK] Saved upscaled image to: " << outputPath << std::endl;
    } else {
        std::cerr << "Failed to save output image!" << std::endl;
    }

    // Clean up CUDA stream
    cudaStreamDestroy(stream);

    std::cout << "Done!" << std::endl;
    return 0;
}