| """ |
| SimFeatUp upsamplers for dense feature restoration. |
| From SegEarth-OV/OV-2 simfeatup_dev. Used by CLIP-based variants (OV, OV-2). |
| """ |
| import math |
| from typing import Optional |
|
|
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
|
|
| try: |
| from featup.adaptive_conv_cuda.adaptive_conv import AdaptiveConv |
| except Exception: |
| AdaptiveConv = None |
|
|
|
|
| def adaptive_conv_py_simple(input, filters): |
| """Pure PyTorch fallback when featup CUDA is unavailable.""" |
| b, c, h1, w1 = input.shape |
| b, h2, w2, f1, f2 = filters.shape |
| assert f1 == f2 |
| t_filters = filters.reshape(b, h2, w2, f1 * f2) |
| patches = torch.nn.Unfold(f1)(input).view((b, c, f1 * f2, h2, w2)) |
| return torch.einsum("bhwf,bcfhw->bchw", t_filters, patches) |
|
|
|
|
| def _meshgrid(device, diameter): |
| dist_range = torch.linspace(-1, 1, diameter, device=device) |
| x, y = torch.meshgrid(dist_range, dist_range, indexing="ij") |
| return torch.cat([x.unsqueeze(0), y.unsqueeze(0)], dim=0) |
|
|
|
|
| class Bilinear(torch.nn.Module): |
| def forward(self, source, guidance): |
| _, _, h, w = guidance.shape |
| return F.interpolate(source, (h, w), mode="bilinear") |
|
|
|
|
| class LayeredResizeConv(torch.nn.Module): |
| def __init__(self, dim, kernel_size=1, *args, **kwargs): |
| super().__init__(*args, **kwargs) |
| self.conv1 = nn.Conv2d(dim + 3, dim, kernel_size, padding="same") |
| self.conv2 = nn.Conv2d(dim + 3, dim, kernel_size, padding="same") |
| self.conv3 = nn.Conv2d(dim + 3, dim, kernel_size, padding="same") |
| self.conv4 = nn.Conv2d(dim + 3, dim, kernel_size, padding="same") |
|
|
| def apply_conv(self, source, guidance, conv, activation): |
| big_source = F.interpolate(source, scale_factor=2, mode="bilinear") |
| _, _, h, w = big_source.shape |
| small_guidance = F.interpolate(guidance, (h, w), mode="bilinear") |
| output = activation(conv(torch.cat([big_source, small_guidance], dim=1))) |
| return big_source + output |
|
|
| def forward(self, source, guidance): |
| source_2 = self.apply_conv(source, guidance, self.conv1, F.relu) |
| source_4 = self.apply_conv(source_2, guidance, self.conv2, F.relu) |
| source_8 = self.apply_conv(source_4, guidance, self.conv3, F.relu) |
| source_16 = self.apply_conv(source_8, guidance, self.conv4, lambda x: x) |
| return source_16 |
|
|
|
|
| class SimpleImplicitFeaturizer(torch.nn.Module): |
| def __init__(self, n_freqs=20): |
| super().__init__() |
| self.n_freqs = n_freqs |
| self.dim_multiplier = 2 |
|
|
| def forward(self, x): |
| b, c, h, w = x.shape |
| dtype = x.dtype |
| grid_h = torch.linspace(-1, 1, h, device=x.device, dtype=dtype) |
| grid_w = torch.linspace(-1, 1, w, device=x.device, dtype=dtype) |
| feats = torch.stack(torch.meshgrid(grid_h, grid_w, indexing="ij")).unsqueeze(0) |
| feats = feats.broadcast_to((b, feats.shape[1], h, w)) |
| freqs = torch.exp(torch.linspace(-2, 10, self.n_freqs, device=x.device)).to(dtype).reshape( |
| 1, self.n_freqs, 1, 1, 1 |
| ) |
| feats = (feats.unsqueeze(1) * freqs).reshape(b, self.n_freqs * self.dim_multiplier, h, w) |
| return torch.cat([torch.sin(feats), torch.cos(feats), x], dim=1) |
|
|
|
|
| class IFA(torch.nn.Module): |
| def __init__(self, feat_dim, num_scales=20): |
| super().__init__() |
| self.feat_dim = feat_dim |
| self.sin_feats = SimpleImplicitFeaturizer() |
| self.mlp = nn.Sequential( |
| nn.Conv2d(feat_dim + (num_scales * 4) + 2, feat_dim, 1), |
| nn.BatchNorm2d(feat_dim), |
| nn.LeakyReLU(), |
| nn.Conv2d(feat_dim, feat_dim, 1), |
| ) |
|
|
| def _upsample_2x(self, source): |
| b, c, h, w = source.shape |
| dtype = source.dtype |
| up_source = F.interpolate(source, (h * 2, w * 2), mode="nearest") |
| lr_cord = torch.linspace(0, h, steps=h, device=source.device, dtype=dtype) |
| hr_cord = torch.linspace(0, h, steps=2 * h, device=source.device, dtype=dtype) |
| lr_coords = torch.stack(torch.meshgrid(lr_cord, lr_cord, indexing="ij")).unsqueeze(0) |
| hr_coords = torch.stack(torch.meshgrid(hr_cord, hr_cord, indexing="ij")).unsqueeze(0) |
| up_lr_coords = F.interpolate(lr_coords, (h * 2, w * 2), mode="nearest") |
| coord_diff = up_lr_coords - hr_coords |
| coord_diff_feats = self.sin_feats(coord_diff).to(dtype) |
| bcast_coord_feats = coord_diff_feats.broadcast_to((b, coord_diff_feats.shape[1], h * 2, w * 2)) |
| return self.mlp(torch.cat([up_source, bcast_coord_feats], dim=1)) |
|
|
| def forward(self, source, guidance): |
| _, _, gh, gw = guidance.shape |
| x = source |
| while x.shape[2] < gh or x.shape[3] < gw: |
| x = self._upsample_2x(x) |
| if x.shape[2] != gh or x.shape[3] != gw: |
| x = F.interpolate(x, (gh, gw), mode="bilinear") |
| return x |
|
|
|
|
| class JBULearnedRange(torch.nn.Module): |
| def __init__(self, guidance_dim, feat_dim, key_dim, scale=2, radius=3): |
| super().__init__() |
| self.scale = scale |
| self.radius = radius |
| self.diameter = self.radius * 2 + 1 |
| self.guidance_dim = guidance_dim |
| self.key_dim = key_dim |
| self.feat_dim = feat_dim |
| self.range_temp = nn.Parameter(torch.tensor(0.0)) |
| self.range_proj = nn.Sequential( |
| nn.Conv2d(guidance_dim, key_dim, 1, 1), |
| nn.GELU(), |
| nn.Dropout2d(0.1), |
| nn.Conv2d(key_dim, key_dim, 1, 1), |
| ) |
| self.fixup_proj = nn.Sequential( |
| nn.Conv2d(guidance_dim + self.diameter ** 2, self.diameter ** 2, 1, 1), |
| nn.GELU(), |
| nn.Dropout2d(0.1), |
| nn.Conv2d(self.diameter ** 2, self.diameter ** 2, 1, 1), |
| ) |
| self.sigma_spatial = nn.Parameter(torch.tensor(1.0)) |
|
|
| def get_range_kernel(self, x): |
| GB, GC, GH, GW = x.shape |
| proj_x = self.range_proj(x) |
| proj_x_padded = F.pad(proj_x, pad=[self.radius] * 4, mode="reflect") |
| queries = ( |
| torch.nn.Unfold(self.diameter)(proj_x_padded) |
| .reshape((GB, self.key_dim, self.diameter * self.diameter, GH, GW)) |
| .permute(0, 1, 3, 4, 2) |
| ) |
| pos_temp = self.range_temp.exp().clamp_min(1e-4).clamp_max(1e4) |
| return F.softmax(pos_temp * torch.einsum("bchwp,bchw->bphw", queries, proj_x), dim=1) |
|
|
| def get_spatial_kernel(self, device): |
| patch = _meshgrid(device, self.diameter) |
| return torch.exp(-patch.square().sum(0) / (2 * self.sigma_spatial ** 2)).reshape( |
| 1, self.diameter * self.diameter, 1, 1 |
| ) |
|
|
| def forward(self, source, guidance): |
| GB, GC, GH, GW = guidance.shape |
| SB, SC, SH, SQ = source.shape |
| assert SB == GB |
| dtype = source.dtype |
| guidance = guidance.to(dtype) |
| spatial_kernel = self.get_spatial_kernel(source.device).to(dtype) |
| range_kernel = self.get_range_kernel(guidance).to(dtype) |
| combined_kernel = (range_kernel * spatial_kernel).to(dtype) |
| combined_kernel /= combined_kernel.sum(1, keepdim=True).clamp(1e-7) |
| combined_kernel += 0.1 * self.fixup_proj(torch.cat([combined_kernel, guidance], dim=1)) |
| combined_kernel = combined_kernel.permute(0, 2, 3, 1).reshape( |
| GB, GH, GW, self.diameter, self.diameter |
| ) |
| hr_source = F.interpolate(source, size=(GH, GW), mode="bicubic", align_corners=False) |
| hr_source_padded = F.pad(hr_source, pad=[self.radius] * 4, mode="reflect") |
| combined_kernel = combined_kernel.to(hr_source_padded.dtype) |
| if AdaptiveConv is not None: |
| result = AdaptiveConv.apply(hr_source_padded, combined_kernel) |
| else: |
| result = adaptive_conv_py_simple(hr_source_padded, combined_kernel) |
| return result |
|
|
|
|
| class JBUStack(torch.nn.Module): |
| def __init__(self, feat_dim, *args, **kwargs): |
| super().__init__(*args, **kwargs) |
| self.up1 = JBULearnedRange(3, feat_dim, 32, radius=3) |
| self.up2 = JBULearnedRange(3, feat_dim, 32, radius=3) |
| self.up3 = JBULearnedRange(3, feat_dim, 32, radius=3) |
| self.up4 = JBULearnedRange(3, feat_dim, 32, radius=3) |
| self.fixup_proj = nn.Sequential( |
| nn.Dropout2d(0.2), |
| nn.Conv2d(feat_dim, feat_dim, kernel_size=1), |
| ) |
|
|
| def upsample(self, source, guidance, up): |
| _, _, h, w = source.shape |
| small_guidance = F.adaptive_avg_pool2d(guidance, (h * 2, w * 2)) |
| return up(source, small_guidance) |
|
|
| def forward(self, source, guidance): |
| source_2 = self.upsample(source, guidance, self.up1) |
| source_4 = self.upsample(source_2, guidance, self.up2) |
| source_8 = self.upsample(source_4, guidance, self.up3) |
| source_16 = self.upsample(source_8, guidance, self.up4) |
| return self.fixup_proj(source_16) * 0.1 + source_16 |
|
|
|
|
| class JBUOne(torch.nn.Module): |
| def __init__(self, feat_dim, *args, **kwargs): |
| super().__init__(*args, **kwargs) |
| self.up = JBULearnedRange(3, feat_dim, 32, radius=5) |
| self.fixup_proj = nn.Sequential( |
| nn.Dropout2d(0.2), |
| nn.Conv2d(feat_dim, feat_dim, kernel_size=1), |
| ) |
|
|
| def upsample(self, source, guidance, up): |
| _, _, h, w = source.shape |
| small_guidance = F.adaptive_avg_pool2d(guidance, (h * 2, w * 2)) |
| return up(source, small_guidance) |
|
|
| def forward(self, source, guidance): |
| source_2 = self.upsample(source, guidance, self.up) |
| source_4 = self.upsample(source_2, guidance, self.up) |
| source_8 = self.upsample(source_4, guidance, self.up) |
| source_16 = self.upsample(source_8, guidance, self.up) |
| return self.fixup_proj(source_16) * 0.1 + source_16 |
|
|
|
|
| FEATUP_CHECKPOINTS = { |
| "jbu_one": "simfeatup/xclip_jbu_one_million_aid.ckpt", |
| "jbu_stack": "simfeatup/clip_jbu_stack_cocostuff.ckpt", |
| "jbu_stack_maskclip": "simfeatup/maskclip_jbu_stack_cocostuff.ckpt", |
| } |
|
|
|
|
| def get_upsampler(name: str, feat_dim: int): |
| if name == "bilinear": |
| return Bilinear() |
| elif name == "jbu_one": |
| return JBUOne(feat_dim) |
| elif name == "jbu_stack": |
| return JBUStack(feat_dim) |
| elif name == "resize_conv": |
| return LayeredResizeConv(feat_dim, 1) |
| elif name == "ifa": |
| return IFA(feat_dim) |
| else: |
| raise ValueError(f"Unknown upsampler: {name}. Use: bilinear, jbu_one, jbu_stack, resize_conv, ifa") |
|
|
