Sail2Dream/phase-rlfn-codec-v1

Phase-RLFN 3× Video Super-Resolution (MLVC codec-aware, RKNN-ready)

A hardware-friendly, BatchNorm-free 3× video super-resolution model trained on frozen MLVC-S real P-frame quantization reconstructions, targeting Rockchip RK3576 / RK3588 / RV1126B NPUs. The SR core is fully quantizable — the training graph is the deployment graph (no BN, no separate fuse step).

This repository hosts the best QAT checkpoint (best_ema.pth / best.pth), the int8 quantized weights (quantized_state_dict.pth), the FP32 reference model (float/best.pth), and the exact training config (config.yaml).

中文说明：本项目是一个 3× 视频超分模型，基于冻结的 MLVC-S 真实 P 帧量化重建训练， 面向 RKNN（瑞芯微）NPU 部署。训练图即部署图、无 BatchNorm，可整网量化。本仓库 收录了最佳 QAT 权重、int8 量化权重、FP32 参考权重与训练配置。

Model description

The model is a Phase-RLFN residual network that predicts a signed residual added on top of a fixed RGA bicubic base. It works per-frame (independent single frame SR) and optionally fuses the MLVC decoder ref_feature through a zero-initialized, bias-free codec adapter to add codec-aware temporal context.

Pipeline:

OpenVidHD mp4
  -> TorchCodec NVDEC + 360x640 canvas
  -> frozen MLVC-S (DPB + round quantization)
  -> current YUV444 frame + optional decoder ref_feature
  -> RGA bicubic base + BN-free Phase-RLFN residual
  -> 1080x1920

• Input / target: MLVC BT.709 full-range YUV444 [0,255].
• The image path always works on its own; the codec adapter is zero-initialized and trained with 25% codec-feature dropout, so the model retains full base SR when no codec context is available.
• Loss is L1; the same vmaf/psnr validation protocol selects the best checkpoint in every training phase.

Architecture

• PixelUnshuffle(phase_factor=2) → PhaseRLFNCore
  • stem (3x3 conv)
  • 4 × ResidualLocalBlock (BN-free two-conv residual, LeakyReLU 0.1)
  • feature_fuse (1x1 conv, dense residuals) + shallow skip
  • optional codec_expand → PixelShuffle(4) → codec_fuse (zero-init, bias-free)
  • residual_head (1x1 conv → signed residual)
• PixelShuffle(core_scale=6) → clip to [0,255]
• RGA bicubic base and clipping are outside the exported NPU core.

Hyper-parameters: num_channels=32, num_blocks=4, phase_factor=2, scale 3, codec_feature_channels=96, codec_project_channels=16, codec_upsample_factor=4.

NPU (RKNN) interface

• RKNN core input: 12x180x320 phase tensor (3ch × phase_factor² × LR space).
• Optional second input: 96x46x80 MLVC decoder feature.
• RKNN core output: 108x180x320 signed residual.
• RGA bicubic, PixelUnshuffle/PixelShuffle, residual addition and clip are placed outside the NPU graph.

Main results

Best QAT checkpoint (best_ema.pth, step ~52000):

Metric	Value
VMAF	70.6694
PSNR	34.9340 dB

Protocol: MLVC-reconstructed YUV444 360×640 → 1080×1920, vmaf/psnr (per config.yaml val_metric: vmaf). Same protocol used throughout training for model selection.

Usage

The checkpoints are raw torch state dicts for the model defined in src/rknn_super_resolution/models/phase_rlfn_sr.py. Load the config and weights with:

import torch
from rknn_super_resolution.config import load_config
from rknn_super_resolution.models import build_model

cfg = load_config("config.yaml")
model = build_model(cfg["model"])           # PhaseRLFNSR
state = torch.load("best_ema.pth", map_location="cpu")
model.load_state_dict(state, strict=True)
model.eval()

The input is the phase-concatenated current frame (or (phases, codec_feature) tuple). Predicts a signed residual; add the RGA bicubic base and clip to [0,255] outside the model to obtain the final SR frame.

Files

File	Description
best_ema.pth	Best QAT checkpoint (EMA), 318 keys, fake-quant.
best.pth	Best QAT checkpoint (raw), 318 keys, fake-quant.
quantized_state_dict.pth	int8 (RKNN-style) quantized weights, 80 keys.
float/best.pth	FP32 reference weights (pre-quantization).
config.yaml	Exact training / model / deploy configuration.
preview.png	SR preview grid.

Training

• 2-GPU DDP, step-based float → observer → QAT timeline (train/loop.py + train/unified.py).
• float_lr=1e-3, qat_lr=1e-5, EMA decay 0.999, QAT backend qnnpack.
• Full training entrypoint: torchrun -m rknn_super_resolution.train.unified. See scripts/run_train.sh.

License

MIT. See LICENSE.