torchvision.transforms.functional_tensor

from torchvision.transforms import functional as F

class AdaptiveAvgPool2d_ONNX(nn.Module):
    def __init__(self, output_size):
        super().__init__()
        if isinstance(output_size, int):
            self.output_size = (output_size, output_size)
        else:
            self.output_size = tuple(output_size)

    def forward(self, x):
        # Step 1: Global average pool to 1x1
        x = F.adaptive_avg_pool2d(x, (1, 1))
        # Step 2: Upsample to target size
        return F.upsample(x, size=self.output_size, mode='nearest')

class PPM2(nn.Module):
    def __init__(self, in_dim, reduction_dim, bins):
        super(PPM2, self).__init__()
        self.features = []
        for bin in bins: #[1,2,3,6]
            self.features.append(nn.Sequential(
                AdaptiveAvgPool2d_ONNX(bin),
                nn.Conv2d(in_dim, reduction_dim, kernel_size=1, bias=False),
                nn.PReLU()
            ))
        self.features = nn.ModuleList(self.features)
        self.fuse = nn.Sequential(
                nn.Conv2d(in_dim+reduction_dim*4, in_dim, kernel_size=3, padding=1, bias=False),
                nn.PReLU())

    def forward(self, x):
        x_size = x.size()
        out = [x]
        for f in self.features:
            out.append(F.interpolate(f(x), x_size[2:], mode='bilinear', align_corners=True))
        out_feat = self.fuse(torch.cat(out, 1))
        return out_feat

class KernelConv2D_V2(nn.Module):
    def __init__(self, ksize=5, act=True):
        super(KernelConv2D_V2, self).__init__()
        self.ksize = ksize
        self.act = act

    def forward(self, feat_in, kernel):
        channels = feat_in.size(1)
        N, kernels, H, W = kernel.size()
        pad = (self.ksize - 1) // 2

        feat_in = F.pad(feat_in, (pad, pad, pad, pad), mode="replicate")
        
        #modified unfold layer to onnx compatible functions
        feat_in = onnx_unfold_alternative_height(feat_in, self.ksize, 1)
        feat_in = onnx_unfold_alternative_width(feat_in, self.ksize, 1)

        feat_in = feat_in.permute(0, 2, 3, 1, 4, 5).contiguous()
        feat_in = feat_in.reshape(N, H, W, channels, -1)

        kernel = kernel.permute(0, 2, 3, 1).reshape(N, H, W, channels, -1)
        feat_out = torch.sum(feat_in * kernel, -1)
        feat_out = feat_out.permute(0, 3, 1, 2).contiguous()
        
        if self.act: 
            feat_out = F.leaky_relu(feat_out, negative_slope=0.2, inplace=True)
        return feat_out

net = ARCH_REGISTRY.get('LEDNet')(channels=[32, 64, 128, 128], connection=False, ppm_version = 2).to(device)

try:
        engine, engine_path = onnx2trt(
            model_path=onnx_path,
            int8=False,
            fp16=use_fp16,  # feel free to experiment. Model did not seem to be trained with FP16 but the INT8s may compensate.
            bf16=False, # feel free to try this also. 
            optimization_level=3,
            shape_profiles=[
                {
                    input_names[0]: {
                        "min": [1, 3, 640, 1120],
                        "opt": [1, 3, 640, 1120],
                        "max": [1, 3, 640, 1120]
                    }
                }
            ]
        )