text: # BSD 3-Clause License
  #
  # Copyright (c) 2020, princeton-vl
  # All rights reserved.
  #
  # Redistribution and use in source and binary forms, with or without
  # modification, are permitted provided that the following conditions are met:
  #
  # * Redistributions of source code must retain the above copyright notice, this
  #   list of conditions and the following disclaimer.
  #
  # * Redistributions in binary form must reproduce the above copyright notice,
  #   this list of conditions and the following disclaimer in the documentation
  #   and/or other materials provided with the distribution.
  #
  # * Neither the name of the copyright holder nor the names of its
  #   contributors may be used to endorse or promote products derived from
  #   this software without specific prior written permission.
  #
  # THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
  # AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
  # IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
  # DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
  # FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
  # DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
  # SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
  # CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
  # OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
  # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
  
  
  from typing import Tuple
  
  import torch
  import torch.nn as nn
  import torch.nn.functional as F
  from torch import Tensor
  
  
  class FlowHead(nn.Module):
      def __init__(self, input_dim: int = 128, hidden_dim: int = 256) -> None:
          super(FlowHead, self).__init__()
          self.conv1 = nn.Conv2d(input_dim, hidden_dim, 3, padding=1)
          self.conv2 = nn.Conv2d(hidden_dim, 2, 3, padding=1)
          self.relu = nn.ReLU(inplace=True)
  
      def forward(self, x: Tensor) -> Tensor:
          return self.conv2(self.relu(self.conv1(x)))
  
  
  class ConvGRU(nn.Module):
      def __init__(self, hidden_dim: int = 128, input_dim: int = (192 + 128)) -> None:
          super(ConvGRU, self).__init__()
          self.convz = nn.Conv2d(hidden_dim + input_dim, hidden_dim, 3, padding=1)
          self.convr = nn.Conv2d(hidden_dim + input_dim, hidden_dim, 3, padding=1)
          self.convq = nn.Conv2d(hidden_dim + input_dim, hidden_dim, 3, padding=1)
  
      def forward(self, h: Tensor, x: Tensor) -> Tensor:
          hx = torch.cat([h, x], dim=1)
  
          z = torch.sigmoid(self.convz(hx))
          r = torch.sigmoid(self.convr(hx))
          q = torch.tanh(self.convq(torch.cat([r * h, x], dim=1)))
  
          h = (1 - z) * h + z * q
          return h
  
  
  class SepConvGRU(nn.Module):
      def __init__(self, hidden_dim: int = 128, input_dim: int = (192 + 128)) -> None:
          super(SepConvGRU, self).__init__()
          self.convz1 = nn.Conv2d(
              hidden_dim + input_dim, hidden_dim, (1, 5), padding=(0, 2)
          )
          self.convr1 = nn.Conv2d(
              hidden_dim + input_dim, hidden_dim, (1, 5), padding=(0, 2)
          )
          self.convq1 = nn.Conv2d(
              hidden_dim + input_dim, hidden_dim, (1, 5), padding=(0, 2)
          )
  
          self.convz2 = nn.Conv2d(
              hidden_dim + input_dim, hidden_dim, (5, 1), padding=(2, 0)
          )
          self.convr2 = nn.Conv2d(
              hidden_dim + input_dim, hidden_dim, (5, 1), padding=(2, 0)
          )
          self.convq2 = nn.Conv2d(
              hidden_dim + input_dim, hidden_dim, (5, 1), padding=(2, 0)
          )
  
      def forward(self, h: Tensor, x: Tensor) -> Tensor:
          # horizontal
          hx = torch.cat([h, x], dim=1)
          z = torch.sigmoid(self.convz1(hx))
          r = torch.sigmoid(self.convr1(hx))
          q = torch.tanh(self.convq1(torch.cat([r * h, x], dim=1)))
          h = (1 - z) * h + z * q
  
          # vertical
          hx = torch.cat([h, x], dim=1)
          z = torch.sigmoid(self.convz2(hx))
          r = torch.sigmoid(self.convr2(hx))
          q = torch.tanh(self.convq2(torch.cat([r * h, x], dim=1)))
          h = (1 - z) * h + z * q
  
          return h
  
  
  class BasicMotionEncoder(nn.Module):
      def __init__(self, corr_levels: int, corr_radius: int) -> None:
          super(BasicMotionEncoder, self).__init__()
          corr_planes = corr_levels * (2 * corr_radius + 1) ** 2
          self.convc1 = nn.Conv2d(corr_planes, 256, 1, padding=0)
          self.convc2 = nn.Conv2d(256, 192, 3, padding=1)
          self.convf1 = nn.Conv2d(2, 128, 7, padding=3)
          self.convf2 = nn.Conv2d(128, 64, 3, padding=1)
          self.conv = nn.Conv2d(64 + 192, 128 - 2, 3, padding=1)
  
      def forward(self, flow: Tensor, corr: Tensor) -> Tensor:
          cor = F.relu(self.convc1(corr))
          cor = F.relu(self.convc2(cor))
          flo = F.relu(self.convf1(flow))
          flo = F.relu(self.convf2(flo))
  
          cor_flo = torch.cat([cor, flo], dim=1)
          out = F.relu(self.conv(cor_flo))
          return torch.cat([out, flow], dim=1)
  
  
  class BasicUpdateBlock(nn.Module):
      def __init__(
          self, corr_levels: int, corr_radius: int, hidden_dim: int = 128
      ) -> None:
          super(BasicUpdateBlock, self).__init__()
          self.encoder = BasicMotionEncoder(corr_levels, corr_radius)
          self.gru = SepConvGRU(hidden_dim=hidden_dim, input_dim=(128 + hidden_dim))
          self.flow_head = FlowHead(hidden_dim, hidden_dim=256)
  
          self.mask = nn.Sequential(
              nn.Conv2d(128, 256, 3, padding=1),
              nn.ReLU(inplace=True),
              nn.Conv2d(256, 64 * 9, 1, padding=0),
          )
  
      def forward(
          self,
          net: Tensor,
          inp: Tensor,
          corr: Tensor,
          flow: Tensor,
      ) -> Tuple[Tensor, Tensor, Tensor]:
          motion_features = self.encoder(flow, corr)
          inp = torch.cat([inp, motion_features], dim=1)
  
          net = self.gru(net, inp)
          delta_flow = self.flow_head(net)
  
          # scale mask to balence gradients
          mask = 0.25 * self.mask(net)
          return net, mask, delta_flow
  