deploy sam3-zerogpu

NeuFlow/backbone_v7.py

@@ -0,0 +1,85 @@
+import torch
+import torch.nn.functional as F
+
+
+class ConvBlock(torch.nn.Module):
+    def __init__(self, in_planes, out_planes, kernel_size, stride, padding):
+        super(ConvBlock, self).__init__()
+
+        self.conv1 = torch.nn.Conv2d(in_planes, out_planes, kernel_size=kernel_size, stride=stride, padding=padding, padding_mode='zeros', bias=False)
+
+        self.conv2 = torch.nn.Conv2d(out_planes, out_planes, kernel_size=3, stride=1, padding=1, bias=False)
+
+        self.relu = torch.nn.LeakyReLU(negative_slope=0.1, inplace=False)
+
+        self.norm1 = torch.nn.BatchNorm2d(out_planes)
+
+        self.norm2 = torch.nn.BatchNorm2d(out_planes)
+
+        # self.dropout = torch.nn.Dropout(p=0.1)
+
+    def forward(self, x):
+
+        # x = self.dropout(x)
+
+        x = self.relu(self.norm1(self.conv1(x)))
+        x = self.relu(self.norm2(self.conv2(x)))
+        # x = self.relu(self.conv1(x))
+        # x = self.relu(self.conv2(x))
+
+        return x
+
+    def forward_fuse(self, x):
+
+        x = self.relu(self.conv1(x))
+        x = self.relu(self.conv2(x))
+
+        return x
+
+
+class CNNEncoder(torch.nn.Module):
+    def __init__(self, feature_dim_s16, context_dim_s16, feature_dim_s8, context_dim_s8):
+        super(CNNEncoder, self).__init__()
+
+        self.block_8_1 = ConvBlock(3, feature_dim_s8 * 2, kernel_size=8, stride=4, padding=2)
+
+        self.block_8_2 = ConvBlock(3, feature_dim_s8, kernel_size=6, stride=2, padding=2)
+
+        self.block_cat_8 = ConvBlock(feature_dim_s8 * 3, feature_dim_s8 + context_dim_s8, kernel_size=3, stride=1, padding=1)
+
+        self.block_16_1 = ConvBlock(3, feature_dim_s16, kernel_size=6, stride=2, padding=2)
+
+        self.block_8_16 = ConvBlock(feature_dim_s8 + context_dim_s8, feature_dim_s16, kernel_size=6, stride=2, padding=2)
+
+        self.block_cat_16 = ConvBlock(feature_dim_s16 * 2, feature_dim_s16 + context_dim_s16 - 2, kernel_size=3, stride=1, padding=1)
+
+    def init_pos(self, batch_size, height, width, device, amp):
+        ys, xs = torch.meshgrid(torch.arange(height, dtype=torch.half if amp else torch.float, device=device), torch.arange(width, dtype=torch.half if amp else torch.float, device=device), indexing='ij')
+        ys = (ys-height/2)
+        xs = (xs-width/2)
+        pos = torch.stack([ys, xs])
+        return pos[None].repeat(batch_size,1,1,1)
+
+    def init_bhwd(self, batch_size, height, width, device, amp):
+        self.pos_s16 = self.init_pos(batch_size, height, width, device, amp)
+
+    def forward(self, img):
+
+        img = F.avg_pool2d(img, kernel_size=2, stride=2)
+        x_8 = self.block_8_1(img)
+
+        img = F.avg_pool2d(img, kernel_size=2, stride=2)
+        x_8_2 = self.block_8_2(img)
+
+        x_8 = self.block_cat_8(torch.cat([x_8, x_8_2], dim=1))
+
+        img = F.avg_pool2d(img, kernel_size=2, stride=2)
+        x_16 = self.block_16_1(img)
+
+        x_16_2 = self.block_8_16(x_8)
+
+        x_16 = self.block_cat_16(torch.cat([x_16, x_16_2], dim=1))
+
+        x_16 = torch.cat([x_16, self.pos_s16], dim=1)
+
+        return x_16, x_8

NeuFlow/config.py

@@ -0,0 +1,7 @@
+feature_dim_s16 = 128
+context_dim_s16 = 64
+iter_context_dim_s16 = 64
+feature_dim_s8 = 128
+context_dim_s8 = 64
+iter_context_dim_s8 = 64
+feature_dim_s1 = 128

NeuFlow/corr.py

@@ -0,0 +1,122 @@
+import torch
+import torch.nn.functional as F
+import math
+
+from NeuFlow import utils
+
+
+def bilinear_sample(img, coords):
+    """ Wrapper for grid_sample, uses pixel coordinates """
+    H, W = img.shape[-2:]
+    xgrid, ygrid = coords.split([1,1], dim=-1)
+    xgrid = 2*xgrid/(W-1) - 1
+    ygrid = 2*ygrid/(H-1) - 1
+
+    grid = torch.cat([xgrid, ygrid], dim=-1)
+
+    with torch.backends.cudnn.flags(enabled=False):
+        img = F.grid_sample(img, grid, align_corners=True)
+
+    return img
+
+
+class CorrBlock:
+    def __init__(self, radius, levels):
+
+        self.radius = radius
+        self.levels = levels
+
+    def init_bhwd(self, batch_size, height, width, device, amp):
+
+        xy_range = torch.linspace(-self.radius, self.radius, 2*self.radius+1, dtype=torch.half if amp else torch.float, device=device)
+
+        delta = torch.stack(torch.meshgrid(xy_range, xy_range, indexing='ij'), axis=-1)
+        delta = delta.view(1, 2*self.radius+1, 2*self.radius+1, 2)
+
+        self.grid = utils.coords_grid(batch_size, height, width, device, amp)
+        self.delta = delta.repeat(batch_size * height * width, 1, 1, 1)
+
+    def __call__(self, corr_pyramid, flow):
+
+        b, _, h, w = flow.shape
+
+        coords = (self.grid + flow).permute(0, 2, 3, 1)
+        coords = coords.reshape(b*h*w, 1, 1, 2)
+
+        out_list = []
+
+        for level, corr in enumerate(corr_pyramid):
+            curr_coords = coords / 2**level + self.delta
+            corr = bilinear_sample(corr, curr_coords)
+            corr = corr.view(b, h, w, -1)
+            out_list.append(corr)
+
+        out = torch.cat(out_list, dim=-1)
+
+        return out.permute(0, 3, 1, 2).contiguous()
+
+    def init_corr_pyr(self, feature0, feature1):
+        b, c, h, w = feature0.shape
+        feature0 = feature0.view(b, c, h*w)
+        feature1 = feature1.view(b, c, h*w)
+        
+        corr = torch.matmul(feature0.transpose(1,2), feature1)
+        corr = corr.view(b*h*w, 1, h, w) / math.sqrt(c)
+
+        corr_pyramid = [corr]
+        for i in range(self.levels-1):
+            corr = F.avg_pool2d(corr, kernel_size=2, stride=2)
+            corr_pyramid.append(corr)
+
+        return corr_pyramid
+
+
+# class RAFTCorrBlock:
+#     def __init__(self, fmap1, fmap2, num_levels=4, radius=4):
+#         self.num_levels = num_levels
+#         self.radius = radius
+#         self.corr_pyramid = []
+
+#         # all pairs correlation
+#         corr = RAFTCorrBlock.corr(fmap1, fmap2)
+
+#         batch, h1, w1, dim, h2, w2 = corr.shape
+#         corr = corr.reshape(batch*h1*w1, dim, h2, w2)
+        
+#         self.corr_pyramid.append(corr)
+#         for i in range(self.num_levels-1):
+#             corr = F.avg_pool2d(corr, 2, stride=2)
+#             self.corr_pyramid.append(corr)
+
+#     def __call__(self, coords):
+#         r = self.radius
+#         coords = coords.permute(0, 2, 3, 1)
+#         batch, h1, w1, _ = coords.shape
+
+#         out_pyramid = []
+#         for i in range(self.num_levels):
+#             corr = self.corr_pyramid[i]
+#             dx = torch.linspace(-r, r, 2*r+1, device=coords.device)
+#             dy = torch.linspace(-r, r, 2*r+1, device=coords.device)
+#             delta = torch.stack(torch.meshgrid(dy, dx), axis=-1)
+
+#             centroid_lvl = coords.reshape(batch*h1*w1, 1, 1, 2) / 2**i
+#             delta_lvl = delta.view(1, 2*r+1, 2*r+1, 2)
+#             coords_lvl = centroid_lvl + delta_lvl
+
+#             corr = bilinear_sample(corr, coords_lvl)
+#             corr = corr.view(batch, h1, w1, -1)
+#             out_pyramid.append(corr)
+
+#         out = torch.cat(out_pyramid, dim=-1)
+#         return out.permute(0, 3, 1, 2).contiguous().float()
+
+#     @staticmethod
+#     def corr(fmap1, fmap2):
+#         batch, dim, ht, wd = fmap1.shape
+#         fmap1 = fmap1.view(batch, dim, ht*wd)
+#         fmap2 = fmap2.view(batch, dim, ht*wd) 
+        
+#         corr = torch.matmul(fmap1.transpose(1,2), fmap2)
+#         corr = corr.view(batch, ht, wd, 1, ht, wd)
+#         return corr  / torch.sqrt(torch.tensor(dim).float())

NeuFlow/matching.py

@@ -0,0 +1,25 @@
+import torch.nn.functional as F
+
+from NeuFlow import utils
+
+
+class Matching:
+
+    def init_bhwd(self, batch_size, height, width, device, amp):
+        self.grid = utils.coords_grid(batch_size, height, width, device, amp)  # [B, 2, H, W]
+        self.flatten_grid = self.grid.view(batch_size, 2, -1).permute(0, 2, 1)  # [B, H*W, 2]
+
+    def global_correlation_softmax(self, feature0, feature1):
+
+        b, c, h, w = feature0.shape
+
+        feature0 = feature0.flatten(-2).permute(0, 2, 1)
+        feature1 = feature1.flatten(-2).permute(0, 2, 1)
+
+        correspondence = F.scaled_dot_product_attention(feature0, feature1, self.flatten_grid)
+
+        correspondence = correspondence.view(b, h, w, 2).permute(0, 3, 1, 2)  # [B, 2, H, W]
+
+        flow = correspondence - self.grid
+
+        return flow

NeuFlow/neuflow.py

@@ -0,0 +1,150 @@
+import torch
+import torch.nn.functional as F
+
+from NeuFlow import backbone_v7
+from NeuFlow import transformer
+from NeuFlow import matching
+from NeuFlow import corr
+from NeuFlow import refine
+from NeuFlow import upsample
+from NeuFlow import config
+
+from huggingface_hub import PyTorchModelHubMixin
+
+
+class NeuFlow(torch.nn.Module,
+              PyTorchModelHubMixin,
+              repo_url="https://github.com/neufieldrobotics/NeuFlow_v2", license="apache-2.0", pipeline_tag="image-to-image"):
+    def __init__(self):
+        super(NeuFlow, self).__init__()
+
+        self.backbone = backbone_v7.CNNEncoder(config.feature_dim_s16, config.context_dim_s16, config.feature_dim_s8, config.context_dim_s8)
+        
+        self.cross_attn_s16 = transformer.FeatureAttention(config.feature_dim_s16+config.context_dim_s16, num_layers=2, ffn=True, ffn_dim_expansion=1, post_norm=True)
+        
+        self.matching_s16 = matching.Matching()
+
+        # self.flow_attn_s16 = transformer.FlowAttention(config.feature_dim_s16)
+
+        self.corr_block_s16 = corr.CorrBlock(radius=4, levels=1)
+        self.corr_block_s8 = corr.CorrBlock(radius=4, levels=1)
+        
+        self.merge_s8 = torch.nn.Sequential(torch.nn.Conv2d(config.feature_dim_s16 + config.feature_dim_s8, config.feature_dim_s8, kernel_size=3, stride=1, padding=1, bias=False),
+                                              torch.nn.GELU(),
+                                              torch.nn.Conv2d(config.feature_dim_s8, config.feature_dim_s8, kernel_size=3, stride=1, padding=1, bias=False),
+                                              torch.nn.BatchNorm2d(config.feature_dim_s8))
+
+        self.context_merge_s8 = torch.nn.Sequential(torch.nn.Conv2d(config.context_dim_s16 + config.context_dim_s8, config.context_dim_s8, kernel_size=3, stride=1, padding=1, bias=False),
+                                           torch.nn.GELU(),
+                                           torch.nn.Conv2d(config.context_dim_s8, config.context_dim_s8, kernel_size=3, stride=1, padding=1, bias=False),
+                                           torch.nn.BatchNorm2d(config.context_dim_s8))
+
+        self.refine_s16 = refine.Refine(config.context_dim_s16, config.iter_context_dim_s16, num_layers=5, levels=1, radius=4, inter_dim=128)
+        self.refine_s8 = refine.Refine(config.context_dim_s8, config.iter_context_dim_s8, num_layers=5, levels=1, radius=4, inter_dim=96)
+
+        self.conv_s8 = backbone_v7.ConvBlock(3, config.feature_dim_s1, kernel_size=8, stride=8, padding=0)
+        self.upsample_s8 = upsample.UpSample(config.feature_dim_s1, upsample_factor=8)
+
+        for p in self.parameters():
+            if p.dim() > 1:
+                torch.nn.init.xavier_uniform_(p)
+
+    def init_bhwd(self, batch_size, height, width, device, amp=True):
+
+        self.backbone.init_bhwd(batch_size*2, height//16, width//16, device, amp)
+
+        self.matching_s16.init_bhwd(batch_size, height//16, width//16, device, amp)
+
+        self.corr_block_s16.init_bhwd(batch_size, height//16, width//16, device, amp)
+        self.corr_block_s8.init_bhwd(batch_size, height//8, width//8, device, amp)
+
+        self.refine_s16.init_bhwd(batch_size, height//16, width//16, device, amp)
+        self.refine_s8.init_bhwd(batch_size, height//8, width//8, device, amp)
+
+        self.init_iter_context_s16 = torch.zeros(batch_size, config.iter_context_dim_s16, height//16, width//16, device=device, dtype=torch.half if amp else torch.float)
+        self.init_iter_context_s8 = torch.zeros(batch_size, config.iter_context_dim_s8, height//8, width//8, device=device, dtype=torch.half if amp else torch.float)
+
+    def split_features(self, features, context_dim, feature_dim):
+
+        context, features = torch.split(features, [context_dim, feature_dim], dim=1)
+
+        context, _ = context.chunk(chunks=2, dim=0)
+        feature0, feature1 = features.chunk(chunks=2, dim=0)
+
+        return features, torch.relu(context)
+
+    def forward(self, img0, img1, iters_s16=1, iters_s8=8):
+
+        flow_list = []
+
+        img0 /= 255.
+        img1 /= 255.
+
+        features_s16, features_s8 = self.backbone(torch.cat([img0, img1], dim=0))
+
+        features_s16 = self.cross_attn_s16(features_s16)
+
+        features_s16, context_s16 = self.split_features(features_s16, config.context_dim_s16, config.feature_dim_s16)
+        features_s8, context_s8 = self.split_features(features_s8, config.context_dim_s8, config.feature_dim_s8)
+
+        feature0_s16, feature1_s16 = features_s16.chunk(chunks=2, dim=0)
+
+        flow0 = self.matching_s16.global_correlation_softmax(feature0_s16, feature1_s16)
+
+        # flow0 = self.flow_attn_s16(feature0_s16, flow0)
+
+        corr_pyr_s16 = self.corr_block_s16.init_corr_pyr(feature0_s16, feature1_s16)
+
+        iter_context_s16 = self.init_iter_context_s16
+
+        for i in range(iters_s16):
+
+            if self.training and i > 0:
+                flow0 = flow0.detach()
+                # iter_context_s16 = iter_context_s16.detach()
+
+            corrs = self.corr_block_s16(corr_pyr_s16, flow0)
+
+            iter_context_s16, delta_flow = self.refine_s16(corrs, context_s16, iter_context_s16, flow0)
+
+            flow0 = flow0 + delta_flow
+
+            if self.training:
+                up_flow0 = F.interpolate(flow0, scale_factor=16, mode='bilinear') * 16
+                flow_list.append(up_flow0)
+
+        flow0 = F.interpolate(flow0, scale_factor=2, mode='nearest') * 2
+
+        features_s16 = F.interpolate(features_s16, scale_factor=2, mode='nearest')
+
+        features_s8 = self.merge_s8(torch.cat([features_s8, features_s16], dim=1))
+
+        feature0_s8, feature1_s8 = features_s8.chunk(chunks=2, dim=0)
+
+        corr_pyr_s8 = self.corr_block_s8.init_corr_pyr(feature0_s8, feature1_s8)
+
+        context_s16 = F.interpolate(context_s16, scale_factor=2, mode='nearest')
+
+        context_s8 = self.context_merge_s8(torch.cat([context_s8, context_s16], dim=1))
+
+        iter_context_s8 = self.init_iter_context_s8
+
+        for i in range(iters_s8):
+
+            if self.training and i > 0:
+                flow0 = flow0.detach()
+                # iter_context_s8 = iter_context_s8.detach()
+
+            corrs = self.corr_block_s8(corr_pyr_s8, flow0)
+
+            iter_context_s8, delta_flow = self.refine_s8(corrs, context_s8, iter_context_s8, flow0)
+
+            flow0 = flow0 + delta_flow
+
+            if self.training or i == iters_s8 - 1:
+
+                feature0_s1 = self.conv_s8(img0)
+                up_flow0 = self.upsample_s8(feature0_s1, flow0) * 8
+                flow_list.append(up_flow0)
+
+        return flow_list

NeuFlow/refine.py

@@ -0,0 +1,50 @@
+import torch
+from NeuFlow import utils
+
+
+class ConvBlock(torch.nn.Module):
+    def __init__(self, in_planes, out_planes, kernel_size, stride, padding):
+        super(ConvBlock, self).__init__()
+
+        self.conv = torch.nn.Conv2d(in_planes, out_planes, kernel_size=kernel_size, stride=stride, padding=padding, padding_mode='zeros', bias=False)
+        self.relu = torch.nn.LeakyReLU(negative_slope=0.1, inplace=False)
+
+    def forward(self, x):
+        return self.relu(self.conv(x))
+
+class Refine(torch.nn.Module):
+    def __init__(self, context_dim, iter_context_dim, num_layers, levels, radius, inter_dim):
+        super(Refine, self).__init__()
+
+        self.radius = radius
+
+        self.conv1 = ConvBlock((radius*2+1)**2*levels+context_dim+iter_context_dim+2+1, context_dim+iter_context_dim, kernel_size=3, stride=1, padding=1)
+
+        self.conv2 = ConvBlock(context_dim+iter_context_dim, inter_dim, kernel_size=3, stride=1, padding=1)
+
+        self.conv_layers = torch.nn.ModuleList([ConvBlock(inter_dim, inter_dim, kernel_size=3, stride=1, padding=1)
+                                                for i in range(num_layers)])
+
+        self.conv3 = torch.nn.Conv2d(inter_dim, iter_context_dim+2, kernel_size=3, stride=1, padding=1, padding_mode='zeros', bias=True)
+
+        # self.hidden_act = torch.nn.Tanh()
+        self.hidden_act = torch.nn.Hardtanh(min_val=-4.0, max_val=4.0)
+        # self.hidden_norm = torch.nn.BatchNorm2d(feature_dim)
+
+    def init_bhwd(self, batch_size, height, width, device, amp):
+        self.radius_emb = torch.tensor(self.radius, dtype=torch.half if amp else torch.float, device=device).view(1,-1,1,1).expand([batch_size,1,height,width])
+
+    def forward(self, corrs, context, iter_context, flow0):
+
+        x = torch.cat([corrs, context, iter_context, flow0, self.radius_emb], dim=1)
+
+        x = self.conv1(x)
+
+        x = self.conv2(x)
+
+        for layer in self.conv_layers:
+            x = layer(x)
+
+        x = self.conv3(x)
+
+        return self.hidden_act(x[:,2:]), x[:,:2]

NeuFlow/transformer.py

@@ -0,0 +1,118 @@
+import torch
+import torch.nn.functional as F
+
+class TransformerLayer(torch.nn.Module):
+    def __init__(self,
+                 feature_dim,
+                 ffn=True,
+                 ffn_dim_expansion=1
+                 ):
+        super(TransformerLayer, self).__init__()
+
+        # multi-head attention
+        self.q_proj = torch.nn.Linear(feature_dim, feature_dim)
+        self.k_proj = torch.nn.Linear(feature_dim, feature_dim)
+        self.v_proj = torch.nn.Linear(feature_dim, feature_dim)
+
+        self.merge = torch.nn.Linear(feature_dim, feature_dim)
+
+        # self.multi_head_attn = torch.nn.MultiheadAttention(feature_dim, 2, batch_first=True, device='cuda')
+
+        self.norm1 = torch.nn.LayerNorm(feature_dim)
+
+        self.ffn = ffn
+
+        if self.ffn:
+            in_channels = feature_dim * 2
+            self.mlp = torch.nn.Sequential(
+                torch.nn.Linear(in_channels, in_channels * ffn_dim_expansion, bias=False),
+                torch.nn.GELU(),
+                torch.nn.Linear(in_channels * ffn_dim_expansion, feature_dim, bias=False),
+            )
+
+            self.norm2 = torch.nn.LayerNorm(feature_dim)
+
+    def forward(self, source, target):
+        # source, target: [B, L, C]
+        query, key, value = source, target, target
+
+        # single-head attention
+        query = self.q_proj(query)  # [B, L, C]
+        key = self.k_proj(key)  # [B, L, C]
+        value = self.v_proj(value)  # [B, L, C]
+
+        message = F.scaled_dot_product_attention(query, key, value, dropout_p=0.0)
+
+        message = self.merge(message)
+
+        # message, _ = self.multi_head_attn(query, key, value, need_weights=False)
+        message = self.norm1(message)
+
+        if self.ffn:
+            message = self.mlp(torch.cat([source, message], dim=-1))
+            message = self.norm2(message)
+
+        return source + message
+
+class FeatureAttention(torch.nn.Module):
+    def __init__(self, feature_dim, num_layers, ffn=True, ffn_dim_expansion=1, post_norm=False):
+        super(FeatureAttention, self).__init__()
+
+        self.layers = torch.nn.ModuleList([
+            TransformerLayer(feature_dim, ffn=ffn, ffn_dim_expansion=ffn_dim_expansion
+                             )
+            for i in range(num_layers)])
+
+        self.post_norm = post_norm
+
+        if self.post_norm:
+            self.norm = torch.nn.BatchNorm2d(feature_dim)
+
+    def forward(self, concat_features0):
+
+        b, c, h, w = concat_features0.shape
+
+        concat_features0 = concat_features0.flatten(-2).permute(0, 2, 1)  # [B, H*W, C]
+        concat_features1 = torch.cat(concat_features0.chunk(chunks=2, dim=0)[::-1], dim=0)
+
+        for layer in self.layers:
+            concat_features0 = layer(concat_features0, concat_features1)
+            concat_features1 = torch.cat(concat_features0.chunk(chunks=2, dim=0)[::-1], dim=0)
+
+        # reshape back
+        concat_features0 = concat_features0.view(b, h, w, c).permute(0, 3, 1, 2).contiguous()  # [B, C, H, W]
+        
+        if self.post_norm:
+            concat_features0 = self.norm(concat_features0)
+
+        return concat_features0
+
+
+class FlowAttention(torch.nn.Module):
+    """
+    flow propagation with self-attention on feature
+    query: feature0, key: feature0, value: flow
+    """
+
+    def __init__(self, feature_dim):
+        super(FlowAttention, self).__init__()
+
+        self.q_proj = torch.nn.Linear(feature_dim, feature_dim)
+        self.k_proj = torch.nn.Linear(feature_dim, feature_dim)
+
+    def forward(self, feature, flow):
+        # q, k: feature [B, C, H, W], v: flow [B, 2, H, W]
+        b, c, h, w = feature.size()
+
+        feature = feature.flatten(-2).permute(0, 2, 1)  # [B, H*W, C]
+
+        flow = flow.flatten(-2).permute(0, 2, 1)
+
+        query = self.q_proj(feature)  # [B, H*W, C]
+        key = self.k_proj(feature)  # [B, H*W, C]
+
+        flow = F.scaled_dot_product_attention(query, key, flow)
+
+        flow = flow.view(b, h, w, 2).permute(0, 3, 1, 2)
+
+        return flow

NeuFlow/upsample.py

@@ -0,0 +1,36 @@
+import torch
+import torch.nn.functional as F
+# from spatial_correlation_sampler import SpatialCorrelationSampler
+
+class UpSample(torch.nn.Module):
+    def __init__(self, feature_dim, upsample_factor):
+        super(UpSample, self).__init__()
+
+        self.upsample_factor = upsample_factor
+
+        self.conv1 = torch.nn.Conv2d(2 + feature_dim, 256, 3, 1, 1)
+        self.conv2 = torch.nn.Conv2d(256, 512, 3, 1, 1)
+        self.conv3 = torch.nn.Conv2d(512, upsample_factor ** 2 * 9, 1, 1, 0)
+        self.relu = torch.nn.ReLU(inplace=True)
+
+    def forward(self, feature, flow):
+
+        concat = torch.cat((flow, feature), dim=1)
+
+        mask = self.conv3(self.relu(self.conv2(self.relu(self.conv1(concat)))))
+
+        b, _, h, w = flow.shape
+
+        mask = mask.view(b, 1, 9, self.upsample_factor, self.upsample_factor, h, w)  # [B, 1, 9, K, K, H, W]
+        mask = torch.softmax(mask, dim=2)
+
+        # up_flow = F.unfold(self.upsample_factor * flow, [3, 3], padding=1)
+        up_flow = F.unfold(flow, [3, 3], padding=1)
+        up_flow = up_flow.view(b, 2, 9, 1, 1, h, w)  # [B, 2, 9, 1, 1, H, W]
+
+        up_flow = torch.sum(mask * up_flow, dim=2)  # [B, 2, K, K, H, W]
+        up_flow = up_flow.permute(0, 1, 4, 2, 5, 3)  # [B, 2, K, H, K, W]
+        up_flow = up_flow.reshape(b, 2, self.upsample_factor * h,
+                                  self.upsample_factor * w)  # [B, 2, K*H, K*W]
+
+        return up_flow

NeuFlow/utils.py

@@ -0,0 +1,37 @@
+import torch
+import torch.nn.functional as F
+
+# def normalize(x):
+#     x_min = x.min()
+#     return (x - x_min) / (x.max() - x_min)
+
+def coords_grid(b, h, w, device, amp):
+    ys, xs = torch.meshgrid(torch.arange(h, dtype=torch.half if amp else torch.float, device=device), torch.arange(w, dtype=torch.half if amp else torch.float, device=device), indexing='ij')  # [H, W]
+
+    grid = torch.stack([xs, ys], dim=0)  # [2, H, W] or [3, H, W]
+
+    grid = grid[None].repeat(b, 1, 1, 1)  # [B, 2, H, W] or [B, 3, H, W]
+
+    return grid
+
+def bilinear_sample(img, sample_coords):
+
+    b, _, h, w = sample_coords.shape
+
+    # Normalize to [-1, 1]
+    x_grid = 2 * sample_coords[:, 0] / (w - 1) - 1
+    y_grid = 2 * sample_coords[:, 1] / (h - 1) - 1
+
+    grid = torch.stack([x_grid, y_grid], dim=-1)  # [B, H, W, 2]
+
+    img = F.grid_sample(img, grid, mode='bilinear', padding_mode='zeros', align_corners=True)
+
+    return img
+
+def flow_warp(feature, flow):
+
+    b, c, h, w = feature.size()
+
+    grid = coords_grid(b, h, w).to(flow.device) + flow  # [B, 2, H, W]
+
+    return bilinear_sample(feature, grid)

app.py

@@ -119,10 +119,12 @@ def _get_video_predictor() -> tuple[str, Any]:
     print("[GPU worker] Falling back to HF Transformers Sam3VideoModel...")
     from transformers import Sam3VideoModel, Sam3VideoProcessor
 
-    _video_hf_processor = Sam3VideoProcessor.from_pretrained("facebook/sam3", token=HF_TOKEN)
-    _video_hf_model = Sam3VideoModel.from_pretrained("facebook/sam3", token=HF_TOKEN, torch_dtype=torch.bfloat16).to(
-        "cuda"
+    _video_hf_processor = Sam3VideoProcessor.from_pretrained(
+        "facebook/sam3", token=HF_TOKEN
     )
+    _video_hf_model = Sam3VideoModel.from_pretrained(
+        "facebook/sam3", token=HF_TOKEN, torch_dtype=torch.bfloat16
+    ).to("cuda")
     print("[GPU worker] HF Sam3VideoModel loaded")
     return ("hf_sam3", (_video_hf_model, _video_hf_processor))
 
@@ -145,7 +147,9 @@ def _get_image_model_and_processor() -> tuple[Any, Any]:
         _image_model = model
         _image_processor = Sam3Processor(_image_model)
         dtype_sample = next(model.parameters()).dtype
-        print(f"[GPU worker] SAM 3 image model ready (first-param dtype={dtype_sample})")
+        print(
+            f"[GPU worker] SAM 3 image model ready (first-param dtype={dtype_sample})"
+        )
     return _image_model, _image_processor
 
 
@@ -327,16 +331,28 @@ def _run_tracking_gpu(video_path: str, query: str) -> dict[str, Any]:
             video_storage_device="cpu",
             dtype=torch.bfloat16,
         )
-        inference_session = processor.add_text_prompt(inference_session=inference_session, text=query)
+        inference_session = processor.add_text_prompt(
+            inference_session=inference_session, text=query
+        )
         outputs_per_frame: dict[int, dict[str, Any]] = {}
-        for model_outputs in model.propagate_in_video_iterator(inference_session=inference_session):
+        for model_outputs in model.propagate_in_video_iterator(
+            inference_session=inference_session
+        ):
             processed = processor.postprocess_outputs(inference_session, model_outputs)
             outputs_per_frame[model_outputs.frame_idx] = processed
 
-        visible_indices = [fi for fi, payload in outputs_per_frame.items() if (payload.get("object_ids") or [])]
+        visible_indices = [
+            fi
+            for fi, payload in outputs_per_frame.items()
+            if (payload.get("object_ids") or [])
+        ]
         visible_indices.sort()
 
-        fps_est = len(frames) / _video_duration_s(video_path) if _video_duration_s(video_path) > 0 else SAMPLE_FPS
+        fps_est = (
+            len(frames) / _video_duration_s(video_path)
+            if _video_duration_s(video_path) > 0
+            else SAMPLE_FPS
+        )
         tracked_segments: list[dict[str, float]] = []
         if visible_indices:
             current_start = visible_indices[0]
@@ -429,7 +445,9 @@ def _download_image_bytes(url: str, *, timeout_s: float = 30.0) -> bytes:
         for chunk in response.iter_bytes(chunk_size=256 * 1024):
             buf.write(chunk)
             if buf.tell() > MAX_IMAGE_BYTES:
-                raise ValueError(f"Image exceeds {MAX_IMAGE_BYTES // (1024 * 1024)} MB limit")
+                raise ValueError(
+                    f"Image exceeds {MAX_IMAGE_BYTES // (1024 * 1024)} MB limit"
+                )
     return buf.getvalue()
 
 
@@ -507,7 +525,9 @@ def _masks_to_rle(output: Any) -> list[dict[str, Any] | None]:
                 masks = candidate
                 break
         if masks is None:
-            print(f"[GPU worker] mask extraction: no mask-like key in output (got {sorted(output.keys())})")
+            print(
+                f"[GPU worker] mask extraction: no mask-like key in output (got {sorted(output.keys())})"
+            )
     if masks is None:
         return []
 
@@ -565,7 +585,9 @@ def _binary_mask_to_rle(mask: np.ndarray) -> dict[str, Any]:
         prev = 1 - prev
     return {
         "size": [int(h), int(w)],
-        "counts": base64.b64encode(",".join(str(r) for r in runs).encode("ascii")).decode("ascii"),
+        "counts": base64.b64encode(
+            ",".join(str(r) for r in runs).encode("ascii")
+        ).decode("ascii"),
         "format": "uncompressed_rle_b64",
     }
 
@@ -747,7 +769,9 @@ def _bbox_to_rle(
 
     return {
         "size": [int(img_h), int(img_w)],
-        "counts": base64.b64encode(",".join(str(r) for r in runs).encode("ascii")).decode("ascii"),
+        "counts": base64.b64encode(
+            ",".join(str(r) for r in runs).encode("ascii")
+        ).decode("ascii"),
         "format": "bbox_rle_b64",
     }
 
@@ -922,7 +946,9 @@ SAPIENS_MAX_PERSONS_PER_FRAME = 12
 SAPIENS_INPUT_HEIGHT = 1024
 SAPIENS_INPUT_WIDTH = 768
 
-SAPIENS_KEYPOINT_NAMES_PATH = os.path.join(os.path.dirname(__file__), "goliath_keypoints.txt")
+SAPIENS_KEYPOINT_NAMES_PATH = os.path.join(
+    os.path.dirname(__file__), "goliath_keypoints.txt"
+)
 
 # Worker-local model cache — separate from the SAM caches above so the
 # two model families don't compete for the same global slot.
@@ -937,7 +963,9 @@ def _sapiens_load_keypoint_names() -> list[str]:
         return _sapiens_worker_keypoint_names
     if os.path.exists(SAPIENS_KEYPOINT_NAMES_PATH):
         with open(SAPIENS_KEYPOINT_NAMES_PATH) as f:
-            names = [line.strip() for line in f if line.strip() and not line.startswith("#")]
+            names = [
+                line.strip() for line in f if line.strip() and not line.startswith("#")
+            ]
             if names:
                 _sapiens_worker_keypoint_names = names
                 return _sapiens_worker_keypoint_names
@@ -1045,12 +1073,19 @@ def _sapiens_preprocess(image: np.ndarray) -> tuple[Any, tuple[int, int]]:
     from PIL import Image
 
     orig_h, orig_w = image.shape[:2]
-    pil = Image.fromarray(image).resize((SAPIENS_INPUT_WIDTH, SAPIENS_INPUT_HEIGHT), Image.BILINEAR)
+    pil = Image.fromarray(image).resize(
+        (SAPIENS_INPUT_WIDTH, SAPIENS_INPUT_HEIGHT), Image.BILINEAR
+    )
     arr = np.array(pil).astype(np.float32) / 255.0
     mean = np.array([0.485, 0.456, 0.406], dtype=np.float32)
     std = np.array([0.229, 0.224, 0.225], dtype=np.float32)
     arr = (arr - mean) / std
-    tensor = torch.from_numpy(arr).permute(2, 0, 1).unsqueeze(0).to("cuda", dtype=torch.float32)
+    tensor = (
+        torch.from_numpy(arr)
+        .permute(2, 0, 1)
+        .unsqueeze(0)
+        .to("cuda", dtype=torch.float32)
+    )
     return tensor, (orig_h, orig_w)
 
 
@@ -1180,7 +1215,9 @@ def _sapiens_extract_video_frames(
                 subprocess.run(cmd, check=True, capture_output=True, timeout=30)
                 frame_bgr = cv2.imread(out_path)
                 if frame_bgr is not None:
-                    extracted_ff.append((ts, cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB)))
+                    extracted_ff.append(
+                        (ts, cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB))
+                    )
             except (subprocess.CalledProcessError, subprocess.TimeoutExpired) as exc:
                 print(f"[sapiens] ffmpeg seek failed for ts={ts}: {exc}")
                 continue
@@ -1227,7 +1264,10 @@ def _sapiens_track_persons_across_frames(
         # ByteTrack drops rejected detections from its output, so output
         # length ≤ input length, and order isn't preserved. Per-input-row
         # bbox-equality matching is the only safe way to assign IDs.
-        if getattr(tracked, "tracker_id", None) is not None and getattr(tracked, "xyxy", None) is not None:
+        if (
+            getattr(tracked, "tracker_id", None) is not None
+            and getattr(tracked, "xyxy", None) is not None
+        ):
             for out_idx in range(len(tracked.tracker_id)):
                 tid = tracked.tracker_id[out_idx]
                 if tid is None:
@@ -1236,7 +1276,10 @@ def _sapiens_track_persons_across_frames(
                 for p in frame_persons:
                     if "id" in p:
                         continue  # already matched to a tracker_id
-                    if abs(p["bbox"]["x1"] - float(out_x1)) < 1.0 and abs(p["bbox"]["y1"] - float(out_y1)) < 1.0:
+                    if (
+                        abs(p["bbox"]["x1"] - float(out_x1)) < 1.0
+                        and abs(p["bbox"]["y1"] - float(out_y1)) < 1.0
+                    ):
                         p["id"] = f"p_{int(tid)}"
                         break
 
@@ -1269,7 +1312,9 @@ def api_pose_image(
         image_array = _sapiens_decode_image_np(raw)
         h, w = image_array.shape[:2]
 
-        person_bboxes = _sapiens_detect_persons(image_array, detector_conf=detector_conf)
+        person_bboxes = _sapiens_detect_persons(
+            image_array, detector_conf=detector_conf
+        )
         if not person_bboxes:
             return {
                 "ok": True,
@@ -1283,7 +1328,9 @@ def api_pose_image(
             }
 
         pose_model = _sapiens_load_pose_model()
-        persons = _sapiens_infer_pose_for_persons(pose_model, image_array, person_bboxes, confidence_threshold)
+        persons = _sapiens_infer_pose_for_persons(
+            pose_model, image_array, person_bboxes, confidence_threshold
+        )
         for j, p in enumerate(persons):
             p["id"] = f"p_{j}"
 
@@ -1363,10 +1410,14 @@ def api_pose_video_frames(
 
         for ts, image_array in frames:
             h, w = image_array.shape[:2]
-            person_bboxes = _sapiens_detect_persons(image_array, detector_conf=detector_conf)
+            person_bboxes = _sapiens_detect_persons(
+                image_array, detector_conf=detector_conf
+            )
             if len(person_bboxes) >= SAPIENS_MAX_PERSONS_PER_FRAME:
                 any_capped = True
-            persons = _sapiens_infer_pose_for_persons(pose_model, image_array, person_bboxes, confidence_threshold)
+            persons = _sapiens_infer_pose_for_persons(
+                pose_model, image_array, person_bboxes, confidence_threshold
+            )
             per_frame_persons.append(persons)
             per_frame_meta.append({"timestamp_s": ts, "width": w, "height": h})
 
@@ -1433,45 +1484,62 @@ NEUFLOW_MAX_TIMESTAMPS = 32
 NEUFLOW_DEFAULT_PAIR_GAP_S = 0.1
 NEUFLOW_DEFAULT_DOWNSAMPLE = 384
 
+# NeuFlow needs init_bhwd(B, H, W, device) called once per resolution.
+# Cache the (H, W) pair the model was last initialised at so we re-init
+# only when the request resolution actually changes.
 _neuflow_worker_model: Any = None
+_neuflow_last_init_hw: tuple[int, int] | None = None
 
 
-def _neuflow_load_model() -> Any:
-    global _neuflow_worker_model
-    if _neuflow_worker_model is not None:
-        return _neuflow_worker_model
+def _neuflow_load_model(image_h: int, image_w: int) -> Any:
+    """Load NeuFlow v2 from the vendored ./NeuFlow/ package + HF weights.
 
-    # NeuFlow v2 publishes a NeuFlow class via the neufieldrobotics
-    # repo that registers PyTorchModelHubMixin. We avoid pinning the
-    # GitHub install in requirements.txt (would force a heavy build);
-    # instead resolve the class at runtime so the Space build stays
-    # cheap and any neuflow API drift surfaces here, not at import.
-    try:
-        from neuflow_v2 import NeuFlow  # type: ignore[import-not-found]
-    except ImportError:
-        from neuflow.neuflow import NeuFlow  # type: ignore[import-not-found,no-redef]
+    Calls init_bhwd whenever the requested (H, W) differs from the last
+    init (NeuFlow allocates internal correlation/upsample buffers per
+    resolution, so init must run before inference at any new size).
+    """
+    global _neuflow_worker_model, _neuflow_last_init_hw
+    if _neuflow_worker_model is None:
+        # Vendored class definition lives in ./NeuFlow/neuflow.py
+        from NeuFlow.neuflow import NeuFlow  # type: ignore[import-not-found]
+
+        print(f"[neuflow] Loading {NEUFLOW_REPO} into GPU memory...")
+        model = NeuFlow.from_pretrained(NEUFLOW_REPO).to("cuda")
+        model.eval()
+        model.half()  # NeuFlow's reference inference path runs fp16
+        _neuflow_worker_model = model
+        print("[neuflow] NeuFlow v2 ready on CUDA")
+
+    if _neuflow_last_init_hw != (image_h, image_w):
+        _neuflow_worker_model.init_bhwd(1, image_h, image_w, "cuda")
+        _neuflow_last_init_hw = (image_h, image_w)
 
-    print(f"[neuflow] Loading {NEUFLOW_REPO} into GPU memory...")
-    model = NeuFlow.from_pretrained(NEUFLOW_REPO).to("cuda")
-    model.eval()
-    _neuflow_worker_model = model
-    print("[neuflow] NeuFlow v2 ready on CUDA")
     return _neuflow_worker_model
 
 
 def _neuflow_preprocess(image: np.ndarray, max_dim: int) -> tuple[Any, tuple[int, int]]:
-    """Downsample to max_dim on the long side; convert to (1,3,H,W) [0,1] tensor."""
+    """Downsample to max_dim on the long side; convert to (1,3,H,W) half tensor.
+
+    NeuFlow v2's reference inference takes raw uint8 BGR (no ImageNet
+    normalisation) → half tensor. NeuFlow also requires H and W to be
+    multiples of 16 for its hierarchical pyramid; we floor-round.
+    """
     h, w = image.shape[:2]
     scale = max_dim / max(h, w)
     if scale < 1.0:
         new_h, new_w = int(h * scale), int(w * scale)
-        image_small = cv2.resize(image, (new_w, new_h), interpolation=cv2.INTER_AREA)
     else:
-        image_small = image
-
-    arr = image_small.astype(np.float32) / 255.0
-    tensor = torch.from_numpy(arr).permute(2, 0, 1).unsqueeze(0).to("cuda", dtype=torch.float32)
-    return tensor, (image_small.shape[0], image_small.shape[1])
+        new_h, new_w = h, w
+    new_h = max(16, (new_h // 16) * 16)
+    new_w = max(16, (new_w // 16) * 16)
+    image_small = cv2.resize(image, (new_w, new_h), interpolation=cv2.INTER_AREA)
+
+    # cv2 default is BGR — match the reference infer_hf.py exactly (no
+    # RGB conversion, no /255 scaling, no mean/std).
+    tensor = (
+        torch.from_numpy(image_small).permute(2, 0, 1).unsqueeze(0).half().to("cuda")
+    )
+    return tensor, (new_h, new_w)
 
 
 def _neuflow_aggregate_flow(
@@ -1483,7 +1551,11 @@ def _neuflow_aggregate_flow(
     pair_timestamp_s: float,
     pair_gap_s: float,
 ) -> dict[str, Any]:
-    """Reduce a (1, 2, H', W') flow tensor to scalar motion stats.
+    """Reduce a (2, H', W') flow tensor to scalar motion stats.
+
+    Caller must already strip the batch dim — NeuFlow's reference
+    inference returns multi-scale flow as a list of (B, 2, H, W) tensors;
+    the wrapper takes [-1][0] to get the highest-res un-batched flow.
 
     Magnitudes are reported in ORIGINAL image pixel space (multiplied by
     src_h / downsampled_h) so they're comparable across calls regardless
@@ -1511,7 +1583,9 @@ def _neuflow_aggregate_flow(
     # Dominant direction: angle of the mean flow vector
     mean_fx = float(np.mean(fx_orig))
     mean_fy = float(np.mean(fy_orig))
-    dominant_dir_deg = float((np.degrees(np.arctan2(-mean_fy, mean_fx)) + 360.0) % 360.0)
+    dominant_dir_deg = float(
+        (np.degrees(np.arctan2(-mean_fy, mean_fx)) + 360.0) % 360.0
+    )
 
     # Direction consistency: |mean(unit_vectors)|
     eps = 1e-6
@@ -1599,14 +1673,19 @@ def api_optical_flow(
                 "elapsed_s": round(time.monotonic() - started, 3),
             }
 
-        model = _neuflow_load_model()
         results: list[dict[str, Any]] = []
         for ts, img1, img2 in pairs:
             src_h, src_w = img1.shape[:2]
             t1, ds_dims = _neuflow_preprocess(img1, downsample_to)
             t2, _ = _neuflow_preprocess(img2, downsample_to)
+            # init_bhwd is per-resolution; load_model takes (H, W) so it
+            # re-inits when we switch sizes (typically constant per call,
+            # but caller can mix resolutions across pairs in theory).
+            model = _neuflow_load_model(ds_dims[0], ds_dims[1])
             with torch.no_grad():
-                flow = model(t1, t2)
+                # NeuFlow returns multi-scale flow; [-1] is highest-res,
+                # [0] strips batch dim → shape (2, H, W)
+                flow = model(t1, t2)[-1][0]
             results.append(
                 _neuflow_aggregate_flow(
                     flow,

requirements.txt

@@ -45,9 +45,9 @@ ultralytics>=8.3.0
 supervision>=0.25.0
 
 # ── BUNDLED: NeuFlow v2 optical flow (Phase 2 spatial-tools) ─────────────
-# NeuFlow v2 doesn't ship as a pip package. The model class is in the
-# upstream repo and uses PyTorchModelHubMixin to load the weights from
-# Study-is-happy/neuflow-v2. May need post-deploy iteration if the
-# install path/import path differs at runtime — see _neuflow_load_model
-# in app.py for the dual-path fallback.
-git+https://github.com/neufieldrobotics/NeuFlow_v2.git
+# NeuFlow v2 doesn't ship as a pip package or have a setup.py — its
+# model class lives in the ./NeuFlow/ subdir which we vendor into this
+# Space (copied from neufieldrobotics/NeuFlow_v2 commit 204b5e3 +
+# loaded by app.py via `from NeuFlow.neuflow import NeuFlow`). Weights
+# are pulled at first inference from Study-is-happy/neuflow-v2 via
+# PyTorchModelHubMixin.from_pretrained. No additional pip dep needed.