dr_guide.cpp

#define STB_IMAGE_WRITE_IMPLEMENTATION
#include <torch/script.h>
#include <torch/torch.h>
#include "./stb_image_write.h"
using namespace torch::indexing;


torch::Tensor tensor_to_grid(torch::Tensor input, int nrow=8, int padding=2) {
    TORCH_CHECK(input.dim() == 4, "Input must be 4D tensor (B, C, H, W)");

    input = input.cpu().to(torch::kFloat32);
    int64_t B = input.size(0);
    int64_t C = input.size(1);
    int64_t H = input.size(2);
    int64_t W = input.size(3);

    nrow = nrow == 0 ? static_cast<int>(std::sqrt(B)) : nrow;
    int64_t ncol = (B + nrow - 1) / nrow;

    torch::Tensor grid = torch::full(
        {C, H * ncol + padding * (ncol - 1), 
         W * nrow + padding * (nrow - 1)}, 
        1.0f);

    for (int64_t b = 0; b < B; ++b) {
        int64_t row = b / nrow;
        int64_t col = b % nrow;
        
        int64_t y_start = row * (H + padding);
        int64_t x_start = col * (W + padding);
        
        grid.slice(1, y_start, y_start + H)
            .slice(2, x_start, x_start + W)
            .copy_(input[b]);
    }

    grid = grid.mul(255).clamp(0, 255).to(torch::kUInt8);
    return grid.permute({1, 2, 0});
}


int main() {
    torch::Device device = torch::kCPU;
    if (torch::cuda::is_available()) {
        device = torch::kCUDA;
    }

    {
        torch::NoGradGuard no_grad;

        ////////////////////////////////
        // // // basic settings // // //
        ////////////////////////////////

        constexpr int base_dx = 105, base_dy = 105;
        constexpr int guide_dz = 9, guide_dx = 128, guide_dy = 128, cond_dx = 128, cond_dy = 128;
        int num_show_each_sample = 10;
        float threshold = 0.95;
        bool visualize = true;
        std::string save_path = "./result";


        //////////////////////////////
        // // // build models // // //
        //////////////////////////////

        std::cout << "start building models..." << std::endl;
        torch::jit::Module fm;
        fm = torch::jit::load("../model/traced_fm_s128_h512_0.pt");
        fm.to(device);
        torch::jit::Module ae;
        ae = torch::jit::load("../model/traced_ae_s128_h512_0.pt");
        ae.to(device);
        ae.to(torch::kHalf); // half-precision to save runtime
        std::cout << "model building finished!" << std::endl;


        /////////////////////////////////////////////////
        // // // specify your sample inputs here // // //
        /////////////////////////////////////////////////

        // example net info
        std::vector<std::string> net_name = {"net_1", "net_2"};
        int num_net = static_cast<int>(net_name.size());

        // example pin input, format: [net_id, layer_num, x_cor, y_cor]
        auto pin_map = torch::zeros({num_net, guide_dz, base_dx, base_dy}, device);
        // net 1
        pin_map.index_put_({0, 0, 40, 40}, 1.0f);   // physical-pin 1
        pin_map.index_put_({0, 0, 40, 41}, 1.0f);
        pin_map.index_put_({0, 0, 40, 42}, 1.0f);
        pin_map.index_put_({0, 0, 80, 40}, 1.0f);   // physical-pin 2
        pin_map.index_put_({0, 0, 80, 41}, 1.0f);
        pin_map.index_put_({0, 0, 81, 40}, 1.0f);
        pin_map.index_put_({0, 0, 81, 41}, 1.0f);
        pin_map.index_put_({0, 0, 80, 80}, 1.0f);   // physical-pin 3
        pin_map.index_put_({0, 0, 80, 81}, 1.0f);
        pin_map.index_put_({0, 0, 80, 82}, 1.0f);
        pin_map.index_put_({0, 0, 80, 83}, 1.0f);
        // net 2
        pin_map.index_put_({1, 0, 30, 70}, 1.0f);   // physical-pin 1
        pin_map.index_put_({1, 0, 30, 71}, 1.0f);
        pin_map.index_put_({1, 0, 30, 72}, 1.0f);
        pin_map.index_put_({1, 6, 60, 0}, 1.0f);    // boundary-pin 1
        pin_map.index_put_({1, 6, 60, 104}, 1.0f);  // boundary-pin 2

        pin_map = torch::nn::functional::interpolate(
        pin_map,
        torch::nn::functional::InterpolateFuncOptions()
            .size(std::vector<int64_t>{cond_dx, cond_dy})
            .mode(torch::kArea)
        );
        pin_map = torch::where(pin_map > 0.0f, 1.0f, 0.0f);
        pin_map = (pin_map - 0.5f) * 2.0f;
        pin_map = pin_map.unsqueeze(0);

        // example blockage input, format: [layer_num, x_cor, y_cor]
        auto block_map = torch::zeros({guide_dz, base_dx, base_dy}, device);
        for (int x = 0; x < 105; ++x) {                 // blockage 1
            for (int y = 50; y < 60; ++y) {
                block_map.index_put_({1, x, y}, 1.0f);
            }
        }
        block_map = block_map.unsqueeze(0);
        block_map = torch::nn::functional::interpolate(
            block_map,
            torch::nn::functional::InterpolateFuncOptions()
                .size(std::vector<int64_t>{cond_dx, cond_dy})
                .mode(torch::kBilinear).align_corners(false)
        );
        block_map = (block_map - 0.5f) * 2.0f;
        
        auto mask = torch::ones({num_net, num_net}, torch::device(device).dtype(torch::kBool));
        mask = mask.unsqueeze(0);
        
        
        ////////////////////////////////////
        // // // perform generation // // //
        ////////////////////////////////////

        // flow-matching model sampling
        std::cout << "start generation..." << std::endl;
        auto zs = torch::randn({1, num_net, 512}, torch::device(device));
        auto t = torch::zeros({1,}, torch::device(device));
        auto final_sample = zs + fm.forward({zs, t, pin_map, block_map, mask}).toTensor();
        // auto-encoder decoding
        final_sample = final_sample.to(torch::kHalf).view({num_net, 32, 4, 4}) / 0.18215;
        auto samples_decode = ae.forward({final_sample}).toTensor().to(torch::kFloat);
        // binarization
        samples_decode = torch::where(samples_decode > threshold, 1.0f, -1.0f);
        samples_decode = (samples_decode + 1) / 2;
        // // further refinement (optional)
        // auto samples_decode_sum = torch::sum(samples_decode, 0, true);
        // auto samples_decode_mask = torch::where(samples_decode_sum <= 1, 1.0f, 0.0f);
        // samples_decode = samples_decode * samples_decode_mask;
        std::cout << "generation finished!" << std::endl;


        ///////////////////////////////
        // // // visualization // // //
        ///////////////////////////////
        
        if (visualize) {
            std::cout << "start visualization..." << std::endl;
            auto samples_vis = samples_decode.view({-1, guide_dz, guide_dx, guide_dy}).index({Slice(0, num_show_each_sample)}).cpu();
            samples_vis = samples_vis.view({-1, 1, guide_dx, guide_dy});

            auto phypin_vis = pin_map.view({-1, guide_dz, cond_dx, cond_dy}).index({Slice(0, num_show_each_sample)}).cpu();
            phypin_vis = (phypin_vis + 1) / 2;
            phypin_vis = phypin_vis.view({-1, 1, cond_dx, cond_dy});
            phypin_vis = torch::nn::functional::interpolate(
                phypin_vis,
                torch::nn::functional::InterpolateFuncOptions()
                    .size(std::vector<int64_t>{guide_dx, guide_dy})
                    .mode(torch::kBilinear).align_corners(false)
            );

            block_map = block_map.repeat({1, num_net, 1, 1, 1});
            auto block_vis = block_map.view({-1, guide_dz, cond_dx, cond_dy}).index({Slice(0, num_show_each_sample)}).cpu();
            block_vis = (block_vis + 1) / 2;
            block_vis = block_vis.view({-1, 1, cond_dx, cond_dy});
            block_vis = torch::nn::functional::interpolate(
                block_vis,
                torch::nn::functional::InterpolateFuncOptions()
                    .size(std::vector<int64_t>{guide_dx, guide_dy})
                    .mode(torch::kBilinear).align_corners(false)
            );

            auto vis = torch::cat({samples_vis, phypin_vis, block_vis}, 1);
            auto grid = tensor_to_grid(vis, guide_dz);

            int H = grid.size(0);
            int W = grid.size(1);
            int C = grid.size(2);
            grid = grid.contiguous();
            std::string filepath = save_path + std::string("/") + "route.png";
            stbi_write_png(filepath.c_str(), W, H, C, 
                        grid.data_ptr<uint8_t>(), W * C);
            std::cout << "visualization finished!" << std::endl;
        }
    }
    return 0;
}