text: import os, sys, math, random, itertools
  import numpy as np
  
  import torch
  import torch.nn as nn
  import torch.nn.functional as F
  
  from torchvision import datasets, transforms, models
  from torch.optim.lr_scheduler import MultiStepLR
  from torch.utils.checkpoint import checkpoint
  
  
  # from utils import *
  
  
  class UNet_up_block(nn.Module):
      def __init__(self, prev_channel, input_channel, output_channel, up_sample=True):
          super().__init__()
          self.up_sampling = nn.Upsample(scale_factor=2, mode='bilinear')
          self.conv1 = nn.Conv2d(prev_channel + input_channel, output_channel, 3, padding=1)
          self.bn1 = nn.GroupNorm(8, output_channel)
          self.conv2 = nn.Conv2d(output_channel, output_channel, 3, padding=1)
          self.bn2 = nn.GroupNorm(8, output_channel)
          self.conv3 = nn.Conv2d(output_channel, output_channel, 3, padding=1)
          self.bn3 = nn.GroupNorm(8, output_channel)
          self.relu = torch.nn.ReLU()
          self.up_sample = up_sample
  
      def forward(self, prev_feature_map, x):
          if self.up_sample:
              x = self.up_sampling(x)
          x = torch.cat((x, prev_feature_map), dim=1)
          x = self.relu(self.bn1(self.conv1(x)))
          x = self.relu(self.bn2(self.conv2(x)))
          x = self.relu(self.bn3(self.conv3(x)))
          return x
  
  
  class UNet_down_block(nn.Module):
      def __init__(self, input_channel, output_channel, down_size=True):
          super().__init__()
          self.conv1 = nn.Conv2d(input_channel, output_channel, 3, padding=1)
          self.bn1 = nn.GroupNorm(8, output_channel)
          self.conv2 = nn.Conv2d(output_channel, output_channel, 3, padding=1)
          self.bn2 = nn.GroupNorm(8, output_channel)
          self.conv3 = nn.Conv2d(output_channel, output_channel, 3, padding=1)
          self.bn3 = nn.GroupNorm(8, output_channel)
          self.max_pool = nn.MaxPool2d(2, 2)
          self.relu = nn.ReLU()
          self.down_size = down_size
  
      def forward(self, x):
          x = self.relu(self.bn1(self.conv1(x)))
          x = self.relu(self.bn2(self.conv2(x)))
          x = self.relu(self.bn3(self.conv3(x)))
          if self.down_size:
              x = self.max_pool(x)
          return x
  
  
  class UNet(nn.Module):
      def __init__(self, downsample=6, in_channels=3, out_channels=3):
          super().__init__()
  
          self.in_channels, self.out_channels, self.downsample = in_channels, out_channels, downsample
          self.down1 = UNet_down_block(in_channels, 16, False)
          self.down_blocks = nn.ModuleList(
              [UNet_down_block(2 ** (4 + i), 2 ** (5 + i), True) for i in range(0, downsample)]
          )
  
          bottleneck = 2 ** (4 + downsample)
          self.mid_conv1 = nn.Conv2d(bottleneck, bottleneck, 3, padding=1)
          self.bn1 = nn.GroupNorm(8, bottleneck)
          self.mid_conv2 = nn.Conv2d(bottleneck, bottleneck, 3, padding=1)
          self.bn2 = nn.GroupNorm(8, bottleneck)
          self.mid_conv3 = torch.nn.Conv2d(bottleneck, bottleneck, 3, padding=1)
          self.bn3 = nn.GroupNorm(8, bottleneck)
  
          self.up_blocks = nn.ModuleList(
              [UNet_up_block(2 ** (4 + i), 2 ** (5 + i), 2 ** (4 + i)) for i in range(0, downsample)]
          )
  
          self.last_conv1 = nn.Conv2d(16, 16, 3, padding=1)
          self.last_bn = nn.GroupNorm(8, 16)
          self.last_conv2 = nn.Conv2d(16, out_channels, 1, padding=0)
          self.relu = nn.ReLU()
  
      def forward(self, x):
          x = self.down1(x)
          xvals = [x]
          for i in range(0, self.downsample):
              x = self.down_blocks[i](x)
              xvals.append(x)
  
          x = self.relu(self.bn1(self.mid_conv1(x)))
          x = self.relu(self.bn2(self.mid_conv2(x)))
          x = self.relu(self.bn3(self.mid_conv3(x)))
  
          for i in range(0, self.downsample)[::-1]:
              x = self.up_blocks[i](xvals[i], x)
  
          x = self.relu(self.last_bn(self.last_conv1(x)))
          #         x = self.relu(self.last_conv2(x))
          x = self.last_conv2(x)
          #         x = F.tanh(x)
          x = x.clamp(min=-1.0, max=1.0)  # -1, F.tanh(x)
          return x
  
      def loss(self, pred, target):
          loss = torch.tensor(0.0, device=pred.device)
          return loss, (loss.detach(),)
  
  
  class UNetHeteroscedasticFull(nn.Module):
      def __init__(self, downsample=6, in_channels=3, out_channels=3, eps=1e-5):
          super().__init__()
  
          self.in_channels, self.out_channels, self.downsample = in_channels, out_channels, downsample
          self.down1 = UNet_down_block(in_channels, 16, False)
          self.down_blocks = nn.ModuleList(
              [UNet_down_block(2 ** (4 + i), 2 ** (5 + i), True) for i in range(0, downsample)]
          )
  
          bottleneck = 2 ** (4 + downsample)
          self.mid_conv1 = nn.Conv2d(bottleneck, bottleneck, 3, padding=1)
          self.bn1 = nn.GroupNorm(8, bottleneck)
          self.mid_conv2 = nn.Conv2d(bottleneck, bottleneck, 3, padding=1)
          self.bn2 = nn.GroupNorm(8, bottleneck)
          self.mid_conv3 = torch.nn.Conv2d(bottleneck, bottleneck, 3, padding=1)
          self.bn3 = nn.GroupNorm(8, bottleneck)
  
          self.up_blocks = nn.ModuleList(
              [UNet_up_block(2 ** (4 + i), 2 ** (5 + i), 2 ** (4 + i)) for i in range(0, downsample)]
          )
  
          self.last_conv1 = nn.Conv2d(16, 16, 3, padding=1)
          self.last_bn = nn.GroupNorm(8, 16)
          self.last_conv2 = nn.Conv2d(16, out_channels, 1, padding=0)
  
          self.last_conv1_uncert = nn.Conv2d(16, 16, 3, padding=1)
          self.last_bn_uncert = nn.GroupNorm(8, 16)
          self.last_conv2_uncert = nn.Conv2d(16, ((out_channels + 1) * (out_channels)) / 2, 1, padding=0)
          self.relu = nn.ReLU()
          self.eps = eps
  
      def forward(self, x):
          x = self.down1(x)
          xvals = [x]
          for i in range(0, self.downsample):
              x = self.down_blocks[i](x)
              xvals.append(x)
  
          x = self.relu(self.bn1(self.mid_conv1(x)))
          x = self.relu(self.bn2(self.mid_conv2(x)))
          x = self.relu(self.bn3(self.mid_conv3(x)))
  
          for i in range(0, self.downsample)[::-1]:
              x = self.up_blocks[i](xvals[i], x)
  
          mu = self.relu(self.last_bn(self.last_conv1(x)))
          mu = self.last_conv2(mu)
          #         mu = F.tanh(mu)
          mu = mu.clamp(min=-1.0, max=1.0)
  
          scale = self.relu(self.last_bn_uncert(self.last_conv1_uncert(x)))
          scale = self.last_conv2_uncert(scale)
          scale = F.softplus(scale)
          return mu, scale
  
  
  class UNetHeteroscedasticIndep(nn.Module):
      def __init__(self, downsample=6, in_channels=3, out_channels=3, eps=1e-5):
          super().__init__()
  
          self.in_channels, self.out_channels, self.downsample = in_channels, out_channels, downsample
          self.down1 = UNet_down_block(in_channels, 16, False)
          self.down_blocks = nn.ModuleList(
              [UNet_down_block(2 ** (4 + i), 2 ** (5 + i), True) for i in range(0, downsample)]
          )
  
          bottleneck = 2 ** (4 + downsample)
          self.mid_conv1 = nn.Conv2d(bottleneck, bottleneck, 3, padding=1)
          self.bn1 = nn.GroupNorm(8, bottleneck)
          self.mid_conv2 = nn.Conv2d(bottleneck, bottleneck, 3, padding=1)
          self.bn2 = nn.GroupNorm(8, bottleneck)
          self.mid_conv3 = torch.nn.Conv2d(bottleneck, bottleneck, 3, padding=1)
          self.bn3 = nn.GroupNorm(8, bottleneck)
  
          self.up_blocks = nn.ModuleList(
              [UNet_up_block(2 ** (4 + i), 2 ** (5 + i), 2 ** (4 + i)) for i in range(0, downsample)]
          )
  
          self.last_conv1 = nn.Conv2d(16, 16, 3, padding=1)
          self.last_bn = nn.GroupNorm(8, 16)
          self.last_conv2 = nn.Conv2d(16, out_channels, 1, padding=0)
  
          self.last_conv1_uncert = nn.Conv2d(16, 16, 3, padding=1)
          self.last_bn_uncert = nn.GroupNorm(8, 16)
          self.last_conv2_uncert = nn.Conv2d(16, out_channels, 1, padding=0)
          self.relu = nn.ReLU()
          self.eps = eps
  
      def forward(self, x):
          x = self.down1(x)
          xvals = [x]
          for i in range(0, self.downsample):
              x = self.down_blocks[i](x)
              xvals.append(x)
  
          x = self.relu(self.bn1(self.mid_conv1(x)))
          x = self.relu(self.bn2(self.mid_conv2(x)))
          x = self.relu(self.bn3(self.mid_conv3(x)))
  
          for i in range(0, self.downsample)[::-1]:
              x = self.up_blocks[i](xvals[i], x)
  
          mu = self.relu(self.last_bn(self.last_conv1(x)))
          mu = self.last_conv2(mu)
          #         mu = F.tanh(mu)
          mu = mu.clamp(min=-1.0, max=1.0)
  
          scale = self.relu(self.last_bn_uncert(self.last_conv1_uncert(x)))
          scale = self.last_conv2_uncert(scale)
          scale = F.softplus(scale)
          return mu, scale
  
  
  class UNetHeteroscedasticPooled(nn.Module):
      def __init__(self, downsample=6, in_channels=3, out_channels=3, eps=1e-5, use_clamp=False):
          super().__init__()
  
          self.in_channels, self.out_channels, self.downsample = in_channels, out_channels, downsample
          self.down1 = UNet_down_block(in_channels, 16, False)
          self.down_blocks = nn.ModuleList(
              [UNet_down_block(2 ** (4 + i), 2 ** (5 + i), True) for i in range(0, downsample)]
          )
  
          bottleneck = 2 ** (4 + downsample)
          self.mid_conv1 = nn.Conv2d(bottleneck, bottleneck, 3, padding=1)
          self.bn1 = nn.GroupNorm(8, bottleneck)
          self.mid_conv2 = nn.Conv2d(bottleneck, bottleneck, 3, padding=1)
          self.bn2 = nn.GroupNorm(8, bottleneck)
          self.mid_conv3 = torch.nn.Conv2d(bottleneck, bottleneck, 3, padding=1)
          self.bn3 = nn.GroupNorm(8, bottleneck)
  
          self.up_blocks = nn.ModuleList(
              [UNet_up_block(2 ** (4 + i), 2 ** (5 + i), 2 ** (4 + i)) for i in range(0, downsample)]
          )
  
          self.last_conv1 = nn.Conv2d(16, 16, 3, padding=1)
          self.last_bn = nn.GroupNorm(8, 16)
          self.last_conv2 = nn.Conv2d(16, out_channels, 1, padding=0)
  
          self.last_conv1_uncert = nn.Conv2d(16, 16, 3, padding=1)
          self.last_bn_uncert = nn.GroupNorm(8, 16)
          self.last_conv2_uncert = nn.Conv2d(16, 1, 1, padding=0)
          self.relu = nn.ReLU()
          self.eps = eps
          self.use_clamp = use_clamp
  
      def forward(self, x):
          x = self.down1(x)
          xvals = [x]
          for i in range(0, self.downsample):
              x = self.down_blocks[i](x)
              xvals.append(x)
  
          x = self.relu(self.bn1(self.mid_conv1(x)))
          x = self.relu(self.bn2(self.mid_conv2(x)))
          x = self.relu(self.bn3(self.mid_conv3(x)))
  
          for i in range(0, self.downsample)[::-1]:
              x = self.up_blocks[i](xvals[i], x)
  
          mu = self.relu(self.last_bn(self.last_conv1(x)))
          mu = self.last_conv2(mu)
  
          if self.use_clamp:
              mu = mu.clamp(min=-1.0, max=1.0)
          else:
              mu = F.tanh(mu)
          # mu = mu.clamp(min=-1.0, max=1.0)
  
          scale = self.relu(self.last_bn_uncert(self.last_conv1_uncert(x)))
          scale = self.last_conv2_uncert(scale)
          scale = F.softplus(scale)
          return mu, scale
  
  
  class UNetReshade(nn.Module):
      def __init__(self, downsample=6, in_channels=3, out_channels=3):
          super().__init__()
  
          self.in_channels, self.out_channels, self.downsample = in_channels, out_channels, downsample
          self.down1 = UNet_down_block(in_channels, 16, False)
          self.down_blocks = nn.ModuleList(
              [UNet_down_block(2 ** (4 + i), 2 ** (5 + i), True) for i in range(0, downsample)]
          )
  
          bottleneck = 2 ** (4 + downsample)
          self.mid_conv1 = nn.Conv2d(bottleneck, bottleneck, 3, padding=1)
          self.bn1 = nn.GroupNorm(8, bottleneck)
          self.mid_conv2 = nn.Conv2d(bottleneck, bottleneck, 3, padding=1)
          self.bn2 = nn.GroupNorm(8, bottleneck)
          self.mid_conv3 = torch.nn.Conv2d(bottleneck, bottleneck, 3, padding=1)
          self.bn3 = nn.GroupNorm(8, bottleneck)
  
          self.up_blocks = nn.ModuleList(
              [UNet_up_block(2 ** (4 + i), 2 ** (5 + i), 2 ** (4 + i)) for i in range(0, downsample)]
          )
  
          self.last_conv1 = nn.Conv2d(16, 16, 3, padding=1)
          self.last_bn = nn.GroupNorm(8, 16)
          self.last_conv2 = nn.Conv2d(16, out_channels, 1, padding=0)
          self.relu = nn.ReLU()
  
      def forward(self, x):
          x = self.down1(x)
          xvals = [x]
          for i in range(0, self.downsample):
              x = self.down_blocks[i](x)
              xvals.append(x)
  
          x = self.relu(self.bn1(self.mid_conv1(x)))
          x = self.relu(self.bn2(self.mid_conv2(x)))
          x = self.relu(self.bn3(self.mid_conv3(x)))
  
          for i in range(0, self.downsample)[::-1]:
              x = self.up_blocks[i](xvals[i], x)
  
          x = self.relu(self.last_bn(self.last_conv1(x)))
          x = self.relu(self.last_conv2(x))
          x = x.clamp(max=1, min=0).mean(dim=1, keepdim=True)
          x = x.expand(-1, 3, -1, -1)
          return x
  
      def loss(self, pred, target):
          loss = torch.tensor(0.0, device=pred.device)
          return loss, (loss.detach(),)
  
  
  class ConvBlock(nn.Module):
      def __init__(self, f1, f2, kernel_size=3, padding=1, use_groupnorm=True, groups=8, dilation=1, transpose=False):
          super().__init__()
          self.transpose = transpose
          self.conv = nn.Conv2d(f1, f2, (kernel_size, kernel_size), dilation=dilation, padding=padding * dilation)
          if self.transpose:
              self.convt = nn.ConvTranspose2d(
                  f1, f1, (3, 3), dilation=dilation, stride=2, padding=dilation, output_padding=1
              )
          if use_groupnorm:
              self.bn = nn.GroupNorm(groups, f1)
          else:
              self.bn = nn.BatchNorm2d(f1)
  
      def forward(self, x):
          # x = F.dropout(x, 0.04, self.training)
          x = self.bn(x)
          if self.transpose:
              # x = F.upsample(x, scale_factor=2, mode='bilinear')
              x = F.relu(self.convt(x))
              # x = x[:, :, :-1, :-1]
          x = F.relu(self.conv(x))
          return x
  
  
  def load_from_file(net, checkpoint_path):
      checkpoint = torch.load(checkpoint_path)
      sd = {k.replace("module.", ""): v for k, v in checkpoint['state_dict'].items()}
      net.load_state_dict(sd)
      for p in net.parameters():
          p.requires_grad = False
      return net
  