622 lines
23 KiB
Python
622 lines
23 KiB
Python
"""BasicVSR++ generator (mmengine-free port).
|
||
|
||
This file is a stripped-down port of:
|
||
|
||
- ``lada/models/basicvsrpp/mmagic/basicvsr_plusplus_net.py`` (BasicVSRPlusPlusNet,
|
||
ResidualBlocksWithInputConv, SPyNet, ...)
|
||
- ``lada/models/basicvsrpp/mmagic/flow_warp.py`` (flow_warp helper)
|
||
- ``lada/models/basicvsrpp/deformconv.py`` (the deformable-conv shim)
|
||
|
||
It removes the dependence on ``mmengine.model.BaseModule`` /
|
||
``mmengine.runner.load_checkpoint`` by:
|
||
|
||
1. Replacing ``BaseModule`` with plain ``nn.Module``.
|
||
2. Re-implementing ``constant_init`` / ``kaiming_init`` /
|
||
``default_init_weights`` from mmengine in a few lines.
|
||
3. Loading checkpoints with raw ``torch.load`` and stripping the EMA prefix
|
||
that lada saves.
|
||
|
||
Originally licensed under Apache-2.0 (mmagic, OpenMMLab) and AGPL-3.0 (lada).
|
||
This vendored copy is distributed under AGPL-3.0 as part of
|
||
ComfyUI-MosaicRestore.
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import logging
|
||
import math
|
||
from typing import Optional
|
||
|
||
import torch
|
||
import torch.nn as nn
|
||
import torch.nn.functional as F
|
||
import torchvision
|
||
|
||
logger = logging.getLogger("ComfyUI-MosaicRestore.basicvsrpp")
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# mmengine init shims
|
||
# ---------------------------------------------------------------------------
|
||
|
||
def _kaiming_init(
|
||
module: nn.Module,
|
||
a: float = 0,
|
||
mode: str = "fan_in",
|
||
nonlinearity: str = "leaky_relu",
|
||
bias: float = 0.0,
|
||
) -> None:
|
||
if hasattr(module, "weight") and module.weight is not None:
|
||
nn.init.kaiming_normal_(module.weight, a=a, mode=mode, nonlinearity=nonlinearity)
|
||
if hasattr(module, "bias") and module.bias is not None:
|
||
nn.init.constant_(module.bias, bias)
|
||
|
||
|
||
def _constant_init(module: nn.Module, val: float, bias: float = 0.0) -> None:
|
||
if hasattr(module, "weight") and module.weight is not None:
|
||
nn.init.constant_(module.weight, val)
|
||
if hasattr(module, "bias") and module.bias is not None:
|
||
nn.init.constant_(module.bias, bias)
|
||
|
||
|
||
def _default_init_weights(module: nn.Module, scale: float = 1.0) -> None:
|
||
"""Light-weight stand-in for mmengine's ``default_init_weights``.
|
||
|
||
Walks ``module`` and applies Kaiming init to convs/linears, scaling the
|
||
weights by ``scale`` (used by EDSR-style residual blocks). BatchNorm-like
|
||
layers get a constant 1/0 init to stay well-behaved if any leak in.
|
||
"""
|
||
for m in module.modules():
|
||
if isinstance(m, (nn.Conv2d, nn.Linear)):
|
||
_kaiming_init(m, a=0, mode="fan_in", bias=0)
|
||
m.weight.data *= scale
|
||
elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm, nn.LayerNorm)):
|
||
_constant_init(m, val=1, bias=0)
|
||
|
||
|
||
def _make_layer(block_cls, num_blocks: int, **kwargs) -> nn.Sequential:
|
||
return nn.Sequential(*[block_cls(**kwargs) for _ in range(num_blocks)])
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Optical-flow warp helper (verbatim from lada/mmagic)
|
||
# ---------------------------------------------------------------------------
|
||
|
||
def flow_warp(
|
||
x: torch.Tensor,
|
||
flow: torch.Tensor,
|
||
interpolation: str = "bilinear",
|
||
padding_mode: str = "zeros",
|
||
align_corners: bool = True,
|
||
) -> torch.Tensor:
|
||
"""Warp an image / feature map by an optical flow field.
|
||
|
||
Args:
|
||
x: ``(N, C, H, W)`` source tensor.
|
||
flow: ``(N, H, W, 2)`` flow in pixel units (last dim is dx, dy).
|
||
"""
|
||
if x.size()[-2:] != flow.size()[1:3]:
|
||
raise ValueError(
|
||
f"flow_warp: spatial sizes mismatch (x={tuple(x.size()[-2:])}, "
|
||
f"flow={tuple(flow.size()[1:3])})"
|
||
)
|
||
_, _, h, w = x.size()
|
||
device = flow.device
|
||
grid_y, grid_x = torch.meshgrid(
|
||
torch.arange(0, h, device=device, dtype=x.dtype),
|
||
torch.arange(0, w, device=device, dtype=x.dtype),
|
||
indexing="ij",
|
||
)
|
||
grid = torch.stack((grid_x, grid_y), 2) # (H, W, 2)
|
||
grid.requires_grad_(False)
|
||
|
||
grid_flow = grid + flow
|
||
gx = 2.0 * grid_flow[..., 0] / max(w - 1, 1) - 1.0
|
||
gy = 2.0 * grid_flow[..., 1] / max(h - 1, 1) - 1.0
|
||
grid_flow = torch.stack((gx, gy), dim=3).to(dtype=x.dtype)
|
||
|
||
return F.grid_sample(
|
||
x,
|
||
grid_flow,
|
||
mode=interpolation,
|
||
padding_mode=padding_mode,
|
||
align_corners=align_corners,
|
||
)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Modulated deformable conv 2d
|
||
# ---------------------------------------------------------------------------
|
||
|
||
class _ModulatedDeformConv2d(nn.Module):
|
||
"""The shape-only base used by lada (the actual deform_conv2d call lives
|
||
in :class:`SecondOrderDeformableAlignment`).
|
||
|
||
We only need this so that ``SecondOrderDeformableAlignment`` can subclass
|
||
it and reuse its ``weight`` / ``bias`` parameters.
|
||
"""
|
||
|
||
def __init__(
|
||
self,
|
||
in_channels: int,
|
||
out_channels: int,
|
||
kernel_size: int,
|
||
stride: int = 1,
|
||
padding: int = 0,
|
||
dilation: int = 1,
|
||
groups: int = 1,
|
||
deform_groups: int = 1,
|
||
bias: bool = True,
|
||
) -> None:
|
||
super().__init__()
|
||
self.in_channels = in_channels
|
||
self.out_channels = out_channels
|
||
self.kernel_size = (kernel_size, kernel_size) if isinstance(kernel_size, int) else tuple(kernel_size)
|
||
self.stride = stride
|
||
self.padding = padding
|
||
self.dilation = dilation
|
||
self.groups = groups
|
||
self.deform_groups = deform_groups
|
||
|
||
self.weight = nn.Parameter(
|
||
torch.empty(out_channels, in_channels // groups, *self.kernel_size)
|
||
)
|
||
if bias:
|
||
self.bias = nn.Parameter(torch.empty(out_channels))
|
||
else:
|
||
self.register_parameter("bias", None)
|
||
self._reset_dconv_parameters()
|
||
|
||
def _reset_dconv_parameters(self) -> None:
|
||
n = self.in_channels
|
||
for k in self.kernel_size:
|
||
n *= k
|
||
stdv = 1.0 / math.sqrt(n)
|
||
nn.init.uniform_(self.weight, -stdv, stdv)
|
||
if self.bias is not None:
|
||
nn.init.zeros_(self.bias)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# BasicVSR++ building blocks
|
||
# ---------------------------------------------------------------------------
|
||
|
||
class ResidualBlockNoBN(nn.Module):
|
||
"""ESRGAN-style residual block, no normalization."""
|
||
|
||
def __init__(self, mid_channels: int = 64, res_scale: float = 1.0) -> None:
|
||
super().__init__()
|
||
self.res_scale = res_scale
|
||
self.conv1 = nn.Conv2d(mid_channels, mid_channels, 3, 1, 1, bias=True)
|
||
self.conv2 = nn.Conv2d(mid_channels, mid_channels, 3, 1, 1, bias=True)
|
||
self.relu = nn.ReLU(inplace=True)
|
||
if res_scale == 1.0:
|
||
for m in (self.conv1, self.conv2):
|
||
_default_init_weights(m, 0.1)
|
||
|
||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||
identity = x
|
||
out = self.conv2(self.relu(self.conv1(x)))
|
||
return identity + out * self.res_scale
|
||
|
||
|
||
class ResidualBlocksWithInputConv(nn.Module):
|
||
def __init__(self, in_channels: int, out_channels: int = 64, num_blocks: int = 30) -> None:
|
||
super().__init__()
|
||
layers = [
|
||
nn.Conv2d(in_channels, out_channels, 3, 1, 1, bias=True),
|
||
nn.LeakyReLU(negative_slope=0.1, inplace=True),
|
||
_make_layer(ResidualBlockNoBN, num_blocks, mid_channels=out_channels),
|
||
]
|
||
self.main = nn.Sequential(*layers)
|
||
|
||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||
return self.main(x)
|
||
|
||
|
||
class PixelShufflePack(nn.Module):
|
||
def __init__(
|
||
self,
|
||
in_channels: int,
|
||
out_channels: int,
|
||
scale_factor: int,
|
||
upsample_kernel: int,
|
||
) -> None:
|
||
super().__init__()
|
||
self.in_channels = in_channels
|
||
self.out_channels = out_channels
|
||
self.scale_factor = scale_factor
|
||
self.upsample_kernel = upsample_kernel
|
||
self.upsample_conv = nn.Conv2d(
|
||
in_channels,
|
||
out_channels * scale_factor * scale_factor,
|
||
upsample_kernel,
|
||
padding=(upsample_kernel - 1) // 2,
|
||
)
|
||
_default_init_weights(self, 1)
|
||
|
||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||
x = self.upsample_conv(x)
|
||
x = F.pixel_shuffle(x, self.scale_factor)
|
||
return x
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# SPyNet
|
||
# ---------------------------------------------------------------------------
|
||
|
||
class _SPyNetConv(nn.Module):
|
||
def __init__(self, in_c: int, out_c: int, kernel_size: int, padding: int, activation: bool) -> None:
|
||
super().__init__()
|
||
self.conv = nn.Conv2d(in_c, out_c, kernel_size=kernel_size, stride=1, padding=padding)
|
||
_kaiming_init(self.conv, a=0, nonlinearity="relu")
|
||
self.activate = nn.ReLU(inplace=True) if activation else None
|
||
|
||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||
x = self.conv(x)
|
||
if self.activate is not None:
|
||
x = self.activate(x)
|
||
return x
|
||
|
||
|
||
class _SPyNetBasicModule(nn.Module):
|
||
def __init__(self) -> None:
|
||
super().__init__()
|
||
self.basic_module = nn.Sequential(
|
||
_SPyNetConv(8, 32, 7, 3, True),
|
||
_SPyNetConv(32, 64, 7, 3, True),
|
||
_SPyNetConv(64, 32, 7, 3, True),
|
||
_SPyNetConv(32, 16, 7, 3, True),
|
||
_SPyNetConv(16, 2, 7, 3, False),
|
||
)
|
||
|
||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||
return self.basic_module(x)
|
||
|
||
|
||
class SPyNet(nn.Module):
|
||
"""SPyNet optical-flow estimator (5-level pyramid + initial scale)."""
|
||
|
||
def __init__(self, pretrained: Optional[str] = None) -> None:
|
||
super().__init__()
|
||
self.basic_module = nn.ModuleList([_SPyNetBasicModule() for _ in range(6)])
|
||
self.register_buffer(
|
||
"mean", torch.tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1)
|
||
)
|
||
self.register_buffer(
|
||
"std", torch.tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1)
|
||
)
|
||
# ``pretrained`` is accepted for API parity with lada; the real lada
|
||
# checkpoint contains SPyNet weights inline so we never load them
|
||
# separately. If a path is passed we'll honor it.
|
||
if pretrained is not None:
|
||
sd = torch.load(pretrained, map_location="cpu")
|
||
self.load_state_dict(sd, strict=False)
|
||
|
||
def compute_flow(self, ref: torch.Tensor, supp: torch.Tensor) -> torch.Tensor:
|
||
n, _, h, w = ref.size()
|
||
ref_pyr = [(ref - self.mean) / self.std]
|
||
supp_pyr = [(supp - self.mean) / self.std]
|
||
for _ in range(5):
|
||
ref_pyr.append(F.avg_pool2d(ref_pyr[-1], 2, 2, count_include_pad=False))
|
||
supp_pyr.append(F.avg_pool2d(supp_pyr[-1], 2, 2, count_include_pad=False))
|
||
ref_pyr = ref_pyr[::-1]
|
||
supp_pyr = supp_pyr[::-1]
|
||
|
||
flow = ref_pyr[0].new_zeros(n, 2, h // 32, w // 32)
|
||
for level in range(len(ref_pyr)):
|
||
if level == 0:
|
||
flow_up = flow
|
||
else:
|
||
flow_up = (
|
||
F.interpolate(
|
||
flow, scale_factor=2, mode="bilinear", align_corners=True
|
||
)
|
||
* 2.0
|
||
)
|
||
warped = flow_warp(
|
||
supp_pyr[level],
|
||
flow_up.permute(0, 2, 3, 1),
|
||
padding_mode="border",
|
||
)
|
||
flow = flow_up + self.basic_module[level](
|
||
torch.cat([ref_pyr[level], warped, flow_up], dim=1)
|
||
)
|
||
return flow
|
||
|
||
def forward(self, ref: torch.Tensor, supp: torch.Tensor) -> torch.Tensor:
|
||
h, w = ref.shape[2:4]
|
||
w_up = w if (w % 32) == 0 else 32 * (w // 32 + 1)
|
||
h_up = h if (h % 32) == 0 else 32 * (h // 32 + 1)
|
||
ref = F.interpolate(ref, size=(h_up, w_up), mode="bilinear", align_corners=False)
|
||
supp = F.interpolate(supp, size=(h_up, w_up), mode="bilinear", align_corners=False)
|
||
flow = F.interpolate(
|
||
self.compute_flow(ref, supp),
|
||
size=(h, w),
|
||
mode="bilinear",
|
||
align_corners=False,
|
||
)
|
||
flow[:, 0, :, :] *= float(w) / float(w_up)
|
||
flow[:, 1, :, :] *= float(h) / float(h_up)
|
||
return flow
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Second-order deformable alignment
|
||
# ---------------------------------------------------------------------------
|
||
|
||
class SecondOrderDeformableAlignment(_ModulatedDeformConv2d):
|
||
"""Lada's second-order deformable alignment block.
|
||
|
||
Uses :func:`torchvision.ops.deform_conv2d` for the actual deformable
|
||
convolution, which has been the canonical CUDA implementation since
|
||
torchvision 0.10.
|
||
"""
|
||
|
||
def __init__(self, *args, max_residue_magnitude: int = 10, **kwargs) -> None:
|
||
self.max_residue_magnitude = max_residue_magnitude
|
||
super().__init__(*args, **kwargs)
|
||
self.conv_offset = nn.Sequential(
|
||
nn.Conv2d(3 * self.out_channels + 4, self.out_channels, 3, 1, 1),
|
||
nn.LeakyReLU(negative_slope=0.1, inplace=True),
|
||
nn.Conv2d(self.out_channels, self.out_channels, 3, 1, 1),
|
||
nn.LeakyReLU(negative_slope=0.1, inplace=True),
|
||
nn.Conv2d(self.out_channels, self.out_channels, 3, 1, 1),
|
||
nn.LeakyReLU(negative_slope=0.1, inplace=True),
|
||
nn.Conv2d(self.out_channels, 27 * self.deform_groups, 3, 1, 1),
|
||
)
|
||
_constant_init(self.conv_offset[-1], val=0, bias=0)
|
||
|
||
def forward(
|
||
self,
|
||
x: torch.Tensor,
|
||
extra_feat: torch.Tensor,
|
||
flow_1: torch.Tensor,
|
||
flow_2: torch.Tensor,
|
||
) -> torch.Tensor:
|
||
extra_feat = torch.cat([extra_feat, flow_1, flow_2], dim=1)
|
||
out = self.conv_offset(extra_feat)
|
||
o1, o2, mask = torch.chunk(out, 3, dim=1)
|
||
|
||
offset = self.max_residue_magnitude * torch.tanh(torch.cat([o1, o2], dim=1))
|
||
offset_1, offset_2 = torch.chunk(offset, 2, dim=1)
|
||
offset_1 = offset_1 + flow_1.flip(1).repeat(1, offset_1.size(1) // 2, 1, 1)
|
||
offset_2 = offset_2 + flow_2.flip(1).repeat(1, offset_2.size(1) // 2, 1, 1)
|
||
offset = torch.cat([offset_1, offset_2], dim=1)
|
||
mask = torch.sigmoid(mask)
|
||
|
||
return torchvision.ops.deform_conv2d(
|
||
x,
|
||
offset,
|
||
self.weight,
|
||
self.bias,
|
||
self.stride,
|
||
self.padding,
|
||
self.dilation,
|
||
mask,
|
||
)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Generator
|
||
# ---------------------------------------------------------------------------
|
||
|
||
class BasicVSRPlusPlusGanNet(nn.Module):
|
||
"""BasicVSR++ generator used by lada.
|
||
|
||
Same architecture as ``mmagic`` (with ``spynet.requires_grad_(False)`` on
|
||
instantiation), so the lada training-time GAN wrapper checkpoint loads
|
||
cleanly into this module after stripping the ``generator_ema.``/
|
||
``generator.`` prefix.
|
||
"""
|
||
|
||
def __init__(
|
||
self,
|
||
mid_channels: int = 64,
|
||
num_blocks: int = 15,
|
||
max_residue_magnitude: int = 10,
|
||
spynet_pretrained: Optional[str] = None,
|
||
) -> None:
|
||
super().__init__()
|
||
self.mid_channels = mid_channels
|
||
|
||
self.spynet = SPyNet(pretrained=spynet_pretrained)
|
||
self.spynet.requires_grad_(False)
|
||
|
||
self.feat_extract = nn.Sequential(
|
||
nn.Conv2d(3, mid_channels, 3, 2, 1),
|
||
nn.LeakyReLU(negative_slope=0.1, inplace=True),
|
||
nn.Conv2d(mid_channels, mid_channels, 3, 2, 1),
|
||
nn.LeakyReLU(negative_slope=0.1, inplace=True),
|
||
ResidualBlocksWithInputConv(mid_channels, mid_channels, 5),
|
||
)
|
||
|
||
self.deform_align = nn.ModuleDict()
|
||
self.backbone = nn.ModuleDict()
|
||
modules = ["backward_1", "forward_1", "backward_2", "forward_2"]
|
||
for i, name in enumerate(modules):
|
||
self.deform_align[name] = SecondOrderDeformableAlignment(
|
||
2 * mid_channels,
|
||
mid_channels,
|
||
3,
|
||
padding=1,
|
||
deform_groups=16,
|
||
max_residue_magnitude=max_residue_magnitude,
|
||
)
|
||
self.backbone[name] = ResidualBlocksWithInputConv(
|
||
(2 + i) * mid_channels, mid_channels, num_blocks
|
||
)
|
||
|
||
self.reconstruction = ResidualBlocksWithInputConv(
|
||
5 * mid_channels, mid_channels, 5
|
||
)
|
||
self.upsample1 = PixelShufflePack(mid_channels, mid_channels, 2, upsample_kernel=3)
|
||
self.upsample2 = PixelShufflePack(mid_channels, 64, 2, upsample_kernel=3)
|
||
self.conv_hr = nn.Conv2d(64, 64, 3, 1, 1)
|
||
self.conv_last = nn.Conv2d(64, 3, 3, 1, 1)
|
||
self.lrelu = nn.LeakyReLU(negative_slope=0.1, inplace=True)
|
||
|
||
# ------------------------------------------------------------------
|
||
# Inference
|
||
# ------------------------------------------------------------------
|
||
|
||
def _compute_flow(self, lqs: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
|
||
n, t, c, h, w = lqs.size()
|
||
lqs_1 = lqs[:, :-1].reshape(-1, c, h, w)
|
||
lqs_2 = lqs[:, 1:].reshape(-1, c, h, w)
|
||
flows_backward = self.spynet(lqs_1, lqs_2).view(n, t - 1, 2, h, w)
|
||
flows_forward = self.spynet(lqs_2, lqs_1).view(n, t - 1, 2, h, w)
|
||
return flows_forward, flows_backward
|
||
|
||
def _propagate(
|
||
self,
|
||
feats: dict,
|
||
flows: torch.Tensor,
|
||
module_name: str,
|
||
) -> dict:
|
||
n, t, _, h, w = flows.size()
|
||
frame_idx = list(range(0, t + 1))
|
||
flow_idx = list(range(-1, t))
|
||
mapping_idx = list(range(0, len(feats["spatial"])))
|
||
mapping_idx += mapping_idx[::-1]
|
||
|
||
if "backward" in module_name:
|
||
frame_idx = frame_idx[::-1]
|
||
flow_idx = frame_idx
|
||
|
||
feat_prop = flows.new_zeros(n, self.mid_channels, h, w)
|
||
for i, idx in enumerate(frame_idx):
|
||
feat_current = feats["spatial"][mapping_idx[idx]]
|
||
if i > 0:
|
||
flow_n1 = flows[:, flow_idx[i]]
|
||
cond_n1 = flow_warp(feat_prop, flow_n1.permute(0, 2, 3, 1))
|
||
|
||
feat_n2 = torch.zeros_like(feat_prop)
|
||
flow_n2 = torch.zeros_like(flow_n1)
|
||
cond_n2 = torch.zeros_like(cond_n1)
|
||
|
||
if i > 1:
|
||
feat_n2 = feats[module_name][-2]
|
||
flow_n2 = flows[:, flow_idx[i - 1]]
|
||
flow_n2 = flow_n1 + flow_warp(flow_n2, flow_n1.permute(0, 2, 3, 1))
|
||
cond_n2 = flow_warp(feat_n2, flow_n2.permute(0, 2, 3, 1))
|
||
|
||
cond = torch.cat([cond_n1, feat_current, cond_n2], dim=1)
|
||
feat_prop = torch.cat([feat_prop, feat_n2], dim=1)
|
||
feat_prop = self.deform_align[module_name](feat_prop, cond, flow_n1, flow_n2)
|
||
|
||
feat = (
|
||
[feat_current]
|
||
+ [feats[k][idx] for k in feats if k not in ("spatial", module_name)]
|
||
+ [feat_prop]
|
||
)
|
||
feat = torch.cat(feat, dim=1)
|
||
feat_prop = feat_prop + self.backbone[module_name](feat)
|
||
feats[module_name].append(feat_prop)
|
||
|
||
if "backward" in module_name:
|
||
feats[module_name] = feats[module_name][::-1]
|
||
return feats
|
||
|
||
def _upsample(self, lqs: torch.Tensor, feats: dict) -> torch.Tensor:
|
||
outputs = []
|
||
num = len(feats["spatial"])
|
||
mapping_idx = list(range(0, num))
|
||
mapping_idx += mapping_idx[::-1]
|
||
for i in range(0, lqs.size(1)):
|
||
hr = [feats[k].pop(0) for k in feats if k != "spatial"]
|
||
hr.insert(0, feats["spatial"][mapping_idx[i]])
|
||
hr = torch.cat(hr, dim=1)
|
||
hr = self.reconstruction(hr)
|
||
hr = self.lrelu(self.upsample1(hr))
|
||
hr = self.lrelu(self.upsample2(hr))
|
||
hr = self.lrelu(self.conv_hr(hr))
|
||
hr = self.conv_last(hr)
|
||
hr += lqs[:, i]
|
||
outputs.append(hr)
|
||
return torch.stack(outputs, dim=1)
|
||
|
||
def forward(self, lqs: torch.Tensor) -> torch.Tensor:
|
||
"""Restore an LQ sequence ``(N, T, 3, H, W)`` to ``(N, T, 3, H, W)``.
|
||
|
||
H and W must each be at least 64 pixels (SPyNet pyramid requirement)
|
||
and at least 4-divisible after downsampling — the model crops to a
|
||
multiple of 4 implicitly via the 0.25× interpolate, so any size that
|
||
is a multiple of 4 works without artifacts at the borders.
|
||
"""
|
||
n, t, c, h, w = lqs.size()
|
||
lqs_downsample = F.interpolate(
|
||
lqs.view(-1, c, h, w),
|
||
scale_factor=0.25,
|
||
mode="bicubic",
|
||
).view(n, t, c, h // 4, w // 4)
|
||
|
||
feats = {}
|
||
feats_ = self.feat_extract(lqs.view(-1, c, h, w))
|
||
fh, fw = feats_.shape[2:]
|
||
feats_ = feats_.view(n, t, -1, fh, fw)
|
||
feats["spatial"] = [feats_[:, i] for i in range(0, t)]
|
||
|
||
if lqs_downsample.size(3) < 64 or lqs_downsample.size(4) < 64:
|
||
raise ValueError(
|
||
"BasicVSR++ requires the low-res input (after 0.25x downsample) "
|
||
"to be at least 64x64; got "
|
||
f"{lqs_downsample.size(3)}x{lqs_downsample.size(4)}."
|
||
)
|
||
flows_forward, flows_backward = self._compute_flow(lqs_downsample)
|
||
|
||
for it in (1, 2):
|
||
for direction in ("backward", "forward"):
|
||
module = f"{direction}_{it}"
|
||
feats[module] = []
|
||
flows = flows_backward if direction == "backward" else flows_forward
|
||
feats = self._propagate(feats, flows, module)
|
||
|
||
return self._upsample(lqs, feats)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Checkpoint loading
|
||
# ---------------------------------------------------------------------------
|
||
|
||
def _strip_prefix(sd: dict, prefix: str) -> dict:
|
||
out = {}
|
||
for k, v in sd.items():
|
||
if k.startswith(prefix):
|
||
out[k[len(prefix):]] = v
|
||
return out
|
||
|
||
|
||
def load_basicvsrpp_checkpoint(
|
||
model: BasicVSRPlusPlusGanNet, checkpoint_path: str, device: torch.device
|
||
) -> None:
|
||
"""Load a lada-format BasicVSR++ checkpoint into ``model``.
|
||
|
||
Lada's training script saves an mmengine ``Runner`` checkpoint that wraps
|
||
the generator under ``generator_ema.`` (preferred — EMA weights) or
|
||
``generator.``. Here we just pull whichever is present and strip the
|
||
prefix so a plain ``model.load_state_dict(strict=True)`` works.
|
||
"""
|
||
raw = torch.load(checkpoint_path, map_location="cpu", weights_only=False)
|
||
state = raw.get("state_dict", raw) if isinstance(raw, dict) else raw
|
||
|
||
# Pick EMA generator first, fall back to plain generator.
|
||
sd = _strip_prefix(state, "generator_ema.")
|
||
if not sd:
|
||
sd = _strip_prefix(state, "generator.")
|
||
if not sd:
|
||
# Already a flat state dict (e.g. a raw export).
|
||
sd = state
|
||
|
||
# Some checkpoints record auxiliary EMA bookkeeping keys we should drop
|
||
# before strict-loading (step counters, etc.).
|
||
for junk in ("step_counter", "num_updates"):
|
||
sd.pop(junk, None)
|
||
|
||
missing, unexpected = model.load_state_dict(sd, strict=False)
|
||
if missing:
|
||
logger.warning("BasicVSR++ checkpoint missing keys: %s", missing[:8])
|
||
if unexpected:
|
||
logger.debug("BasicVSR++ checkpoint unexpected keys (ignored): %s", unexpected[:8])
|
||
model.to(device).eval()
|