RRDBNet fine-tuned for low-light 4x denoising + super-resolution

Fine-tuned from the official Real-ESRGAN x4plus generator on a low-light denoising + 4x super-resolution dataset (paired noisy-LR -> clean-HR images). Trained end-to-end directly on the actual noisy inputs (not synthetically clean ones), so it learns denoising and super-resolution jointly - no separate denoising stage is required.

Best validation PSNR: 39.55 dB


Repository files

File Description
model.py Model definition (RRDBNet) plus get_model() / load_pretrained() helpers
config.json Architecture and training hyperparameters in machine-readable form
best_model.pth Fine-tuned weights (state_dict, best validation checkpoint)

Architecture

RRDBNet (the ESRGAN / Real-ESRGAN generator): a stack of Residual-in-Residual Dense Blocks (RRDB) for feature extraction, followed by two nearest-neighbor-upsample + conv steps (not pixel-shuffle - Real-ESRGAN's deliberate choice to avoid pixel-shuffle's characteristic checkerboard artifacts) to reach 4x.

Forward pass, stage by stage:

  1. Shallow feature extraction - conv_first: one 3x3 conv, 3 -> num_feat (64) channels.
  2. Deep feature trunk - body: num_block (23) stacked RRDB blocks, all operating at 64 channels. Each RRDB block is itself 3 stacked Residual Dense Blocks (RDB):
    • Each RDB has 5 conv layers (3x3). Layer i takes the concatenation of the block's input and the outputs of all previous layers within the same RDB ("dense" connections), and produces num_grow_ch (32) new channels - except the 5th conv, which maps back down to num_feat (64) channels to close the block.
    • LeakyReLU(0.2) follows convs 1-4; conv 5 has no activation.
    • The RDB's output is scaled by 0.2 and added back to the block's input (local residual).
    • The RRDB's output is likewise the third RDB's output scaled by 0.2, added back to the RRDB's own input (residual-in-residual).
  3. Trunk skip connection - conv_body: one 3x3 conv on the trunk output, added element-wise to the shallow features from step 1 (a long skip around the entire 23-block trunk).
  4. Upsampling (x4 total, in two x2 stages) - for each stage: nearest-neighbor interpolate (scale_factor=2) -> 3x3 conv (conv_up1, then conv_up2) -> LeakyReLU(0.2).
  5. HR reconstruction - conv_hr: 3x3 conv + LeakyReLU(0.2) at the final (4x) resolution, then conv_last: one 3x3 conv, num_feat -> 3 channels, producing the final RGB output.
input (3ch)
  -> conv_first (3x3)                         [64ch]
  -> [RRDB x 23]  (each: 3x RDB, dense 5-conv, 0.2-scaled residual)
  -> conv_body (3x3) -> + (skip from conv_first output)
  -> upsample x2 (nearest) -> conv_up1 -> LReLU
  -> upsample x2 (nearest) -> conv_up2 -> LReLU
  -> conv_hr (3x3) -> LReLU
  -> conv_last (3x3)                          [3ch]
output (3ch, 4x spatial size)

Configuration

Full architecture and training hyperparameters (also shipped as config.json in this repo):

{
  "architecture": "RRDBNet",
  "base_model": "Real-ESRGAN x4plus (ESRGAN generator)",
  "base_model_source": "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth",
  "task": "low-light image denoising + 4x super-resolution",
  "num_in_ch": 3,
  "num_out_ch": 3,
  "scale": 4,
  "num_feat": 64,
  "num_block": 23,
  "num_grow_ch": 32,
  "upsampling": "nearest-neighbor-interpolate + conv (2x2 steps)",
  "num_parameters": 16700000,
  "loss": "L1",
  "optimizer": "Adam",
  "learning_rate": 2e-4,
  "weight_decay": 0,
  "lr_schedule": "cosine annealing",
  "epochs": 50,
  "hr_patch_size": 256,
  "batch_size": 8,
  "best_val_psnr_db": 39.55
}
Setting Value
num_feat 64
num_block 23
num_grow_ch 32
scale 4
Parameters 16.70M

Training

  • Initialization: the official RealESRGAN_x4plus.pth checkpoint, loaded with a perfect match (0 missing, 0 unexpected keys) before fine-tuning - this confirms the architecture above is identical to the official Real-ESRGAN x4plus generator.
  • Data: 1,105 paired low-light noisy-LR / clean-HR training images, 267 validation pairs.
  • Loss: L1 (mean absolute error) between the predicted and ground-truth HR image.
  • Optimizer: Adam, initial lr=2e-4, weight_decay=0, cosine-annealing schedule over 50 epochs.
  • Patch size: 256x256 HR crops (64x64 LR) per training step, batch size 8.
  • Random flips/transpose augmentation (8-fold dihedral symmetry) applied per crop.

Training curve (validation PSNR)

Epoch val PSNR
1 39.07 dB
8 39.42 dB
22 39.46 dB
30 39.52 dB
37 39.54 dB
48 39.55 dB (best)
50 39.55 dB

Validation PSNR largely plateaued from around epoch 30 onward.

Quickstart

import torch
import numpy as np
from PIL import Image
from model import load_pretrained  # model.py in this repo

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = load_pretrained('best_model.pth', device=device)  # filename as uploaded to this repo

image = Image.open('input_lowlight.png').convert('RGB')
input_tensor = torch.from_numpy(np.array(image).astype(np.float32) / 255.0).permute(2, 0, 1).unsqueeze(0).to(device)

with torch.no_grad():
    output_tensor = model(input_tensor).clamp(0.0, 1.0)

output_np = (output_tensor[0].cpu().permute(1, 2, 0).numpy() * 255.0).round().astype(np.uint8)
Image.fromarray(output_np).save('output_enhanced_4x.png')

For large images that don't fit in GPU memory in one pass, tile the input and stitch the outputs back together (see the tiled-inference pattern in the training notebook this model came from).

Limitations

  • Fine-tuned and validated on a single, specific low-light dataset - it is not a general-purpose denoiser/upscaler, and performance on out-of-distribution images (different noise characteristics, different lighting conditions) is untested.
  • Optimized purely for PSNR via an L1 loss (no perceptual/adversarial loss), so outputs prioritize pixel-accuracy over subjective sharpness - this is a deliberate choice for a PSNR-scored task, not a general recommendation for visual quality use cases.
  • The validation split comes from the same source/collection process as the training data, so this PSNR number does not guarantee similar performance on images collected differently.
Downloads last month

-

Downloads are not tracked for this model. How to track
Inference Providers NEW
This model isn't deployed by any Inference Provider. 🙋 Ask for provider support