super-resolution

PyTorch Lightning training code for image super-resolution with a Swin2SR backbone.

What this repository includes

• A training pipeline built on pytorch-lightning
• A Swin2SR model wrapper initialized from caidas/swin2SR-realworld-sr-x4-64-bsrgan-psnr
• Dataset support for:
  • local PNG image folders
  • LAION high-resolution data through Hugging Face Datasets
• Super-resolution specific losses and image quality metrics (PSNR, SSIM, MS-SSIM)
• Weights & Biases logging and checkpointing

Project layout

src/super_resolution/
├── scripts/
│   ├── train.py         # main training entrypoint
│   └── download_hf.py   # helper script to download LAION HD images locally
├── modeling/
│   └── swin2sr/model.py # Swin2SR model wrapper
├── datasets/            # dataset implementations
├── datamodule.py        # Lightning DataModule
├── lightning_module.py  # training/validation logic
└── losses.py            # custom SR losses

Requirements

• Python 3.12+
• CUDA GPU recommended for training

Installation

Option 1: pip

pip install -e .

Option 2: uv

uv sync

Data options

1) Train from local PNG files

Point --data_root to a directory containing PNG files (recursively discovered).

2) Stream from LAION high-resolution dataset

If --data_root is omitted, training uses laion/laion-high-resolution via Hugging Face Datasets.

Training

From the repository root:

python -m super_resolution.scripts.train <run_name>

Examples:

# Train from LAION stream
python -m super_resolution.scripts.train sr-laion-exp1

# Train from local files
python -m super_resolution.scripts.train sr-local-exp1 --data_root /path/to/images

# Multi-GPU run
python -m super_resolution.scripts.train sr-ddp-exp --num_devices 2

Useful flags:

• --batch_size (default: 8)
• --num_workers (default: 8)
• --num_devices (default: 1)
• --log_path (default: logs)
• --ckpt_path to resume from a Lightning checkpoint
• --weights_path to load model weights
• --debug to run W&B in offline mode
• --overfit to overfit a small subset for debugging

Optional: download LAION HD subset locally

python -m super_resolution.scripts.download_hf /path/to/data_root

Images are saved under /path/to/data_root/laion_hd.

Linting and tests

ruff check src
pytest

Notes

• Training logs are written under logs/<run_name> by default.
• Validation currently attempts to push the model to Hugging Face Hub (swin2sr-laion-hd) at validation end.
  • Authenticate first with huggingface-cli login if you want this upload behavior.
  • If you do not want automatic uploads, remove or gate the push_to_hub(...) call in src/super_resolution/lightning_module.py.

License

This project is licensed under the terms in LICENSE.