Skip to content

README.md

Latest commit

 

History

History
232 lines (173 loc) · 10.1 KB

README.md

File metadata and controls

232 lines (173 loc) · 10.1 KB
Raw

python-upscaler

A small, standalone AI image upscaler — one Python file, one CLI, any super-resolution model.

python upscale.py input.png output.png --weights weights/4x-UltraSharp.pth
  • Any architecture. Loads ESRGAN, Real-ESRGAN, DAT, HAT, SwinIR, SRFormer, OmniSR and friends through spandrel — it detects the architecture and the model's native scale straight from the weights file. No per-model code.
  • Alpha-safe. RGBA in, RGBA out. The model upscales RGB; the alpha channel is resized with Lanczos and reattached, so transparency survives (important for UI art, sprites, logos).
  • Self-ensemble. --ensemble averages the 8 dihedral orientations for a cleaner, less speckly result at 8× the GPU cost.
  • No heavy deps. Just torch + numpy + pillow + spandrel. If spandrel isn't installed, a self-contained RRDBNet in rrdbnet.py still runs ESRGAN-family weights.
  • CUDA when available, CPU otherwise — same command either way.

Why a whole repo for this? So the heavy Python/CUDA toolchain lives in exactly one place and the projects consuming the upscaled assets stay lean. The recipes/ folder is where that pays off: each recipe is a committed script pinning the model and flags for one repeatable job.

Install

python3 -m venv .venv
. .venv/bin/activate
pip install -r requirements.txt

Recent torch ships CUDA wheels on PyPI, so the plain install usually gets you GPU support. If you need a specific CUDA version:

pip install -r requirements.txt --extra-index-url https://download.pytorch.org/whl/cu124

CPU-only works fine for occasional images — expect seconds instead of milliseconds for ESRGAN models, and considerably longer for the transformer ones.

Tested on Python 3.10 with torch 2.x. numpy is pinned <2 because some torch/spandrel builds still expect the 1.x ABI.

Quick start

. .venv/bin/activate

# 4x with the model's native scale
python upscale.py input.png output.png --weights weights/4x-UltraSharp.pth

# upscale with the model, then resize down to exactly 2x the original
python upscale.py input.png output.png --weights weights/4x-UltraSharp.pth --scale 2

# cleanest result, 8x the GPU work
python upscale.py input.png output.png --weights weights/4xNomos8kDAT.safetensors --ensemble

CLI

Argument Default What it does
input — Source image. Anything Pillow reads (PNG, JPG, WebP, …).
output — Destination. The extension picks the format; use .png to keep alpha.
--weights weights/4x-UltraSharp.pth Path to a .pth or .safetensors model.
--scale model's native Final scale relative to the input. The model always runs at its own scale first; if the result doesn't match, it's Lanczos-resized to the target.
--ensemble off Geometric self-ensemble: run 8 flipped/transposed variants and average them. De-speckles flat areas and firms up edges, slightly softer, 8× slower.

The run prints the device, the detected architecture and scale, and the final size:

device: cuda (NVIDIA GeForce RTX 4090)
model: DAT x4 (4xNomos8kDAT.safetensors)
saved: output.png  (1024, 768)  ((256, 192) -> 4.0x)

Loading new weights

Nothing is bundled — models are hundreds of megabytes and carry their own licenses, so you fetch the ones you want. weights/ is gitignored.

mkdir -p weights

# Real-ESRGAN x4plus — the general-purpose baseline
curl -L -o weights/RealESRGAN_x4plus.pth \
  https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth

# 4x-UltraSharp — sharper ESRGAN variant, keeps fine detail
curl -L -o weights/4x-UltraSharp.pth \
  https://huggingface.co/lokCX/4x-Ultrasharp/resolve/main/4x-UltraSharp.pth

Then point --weights at any file you like — the path is arbitrary, weights/ is only a convention:

python upscale.py in.png out.png --weights /some/other/dir/4xNomos8kDAT.safetensors

Where to find more models. OpenModelDB is the community catalogue — filter by scale and by what the model was trained for (photo, anime, pixel art, text, faces, compressed sources). Most entries link to a HuggingFace or GitHub download. Grab the .pth or .safetensors, drop it in weights/, done.

What actually happens when you load one. load_model() hands the file to spandrel, which sniffs the state dict, picks the right architecture, and reports the scale — that scale then drives the whole pipeline, so non-4x models (2x, 3x, 8x) just work with no flags. spandrel_extra_arches is installed alongside and registers the architectures spandrel keeps out of the core package. If spandrel is missing entirely, the code falls back to the bundled RRDBNet, including a remap of legacy old-ESRGAN state dicts (model.1.sub.N.…) via _convert_old_esrgan().

Troubleshooting.

Symptom Fix
UnsupportedModelError spandrel doesn't know that architecture. Check spandrel_extra_arches is installed and both packages are current.
size mismatch / Missing key(s) You're on the RRDBNet fallback with non-ESRGAN weights. Install spandrel.
CUDA out of memory Drop --ensemble, or upscale a smaller crop, or force CPU with CUDA_VISIBLE_DEVICES="". Transformer models (DAT/HAT) want far more VRAM than ESRGAN ones.
Output has a black/white halo where it was transparent Save as .png. JPEG can't hold an alpha channel.
numpy.dtype size changed numpy 2 crept in — reinstall with the numpy<2 pin.

Licensing. The MIT license in this repo covers the code only. Model weights are separately licensed by their authors, and several architectures shipped in spandrel_extra_arches are non-commercial — check the model's page before shipping its output commercially.

Model notes

Model Character Speed
RealESRGAN_x4plus Cleans JPEG artifacts and noise well; can over-smooth fine detail like small icons or text. sub-second
4x-UltraSharp Sharper ESRGAN variant, retains detail; crisp, occasionally a bit crunchy. sub-second
4xNomos8kDAT (DAT) Modern transformer. Cleaner, more coherent detail — good on ornate/illustrated art. ~6–12 s
4xNomos8kSCHAT-L (HAT) Transformer, heaviest; similar tier to DAT, different flavour. ~6–12 s

Times are for a small (few-hundred-pixel) image on a mid-range consumer GPU; they scale with pixel count and multiply by 8 under --ensemble.

Rules of thumb: noisy or JPEG-sourced input → Real-ESRGAN. Clean art where you want detail held → UltraSharp. Ornate illustration, UI art, anything where ESRGAN output looks busy → DAT/HAT with --ensemble. When in doubt, run recipes/compare-models.sh on one representative image and look.

Recipes

A recipe is a small bash script in recipes/ that pins the model, the flags and the paths for one repeatable job. Committing it means the exact command that produced your assets is reviewable, diffable and rerunnable months later — instead of living in someone's shell history.

Two ship as examples and as templates.

recipes/batch-folder.sh — upscale every image in a folder.

./recipes/batch-folder.sh art/ art-hd/

# pick a model, turn on the self-ensemble
MODEL=weights/4xNomos8kDAT.safetensors ENSEMBLE=1 ./recipes/batch-folder.sh art/ art-hd/

# force exactly 2x output and redo files that already exist
SCALE=2 FORCE=1 ./recipes/batch-folder.sh art/ art-hd/

Walks *.png *.jpg *.jpeg *.webp, writes <name>.png into the output folder, and skips anything already there unless FORCE=1 — so an interrupted batch resumes instead of starting over.

recipes/compare-models.sh — one image, every model, side by side.

./recipes/compare-models.sh sample.png out/
MODELS="weights/4x-UltraSharp.pth weights/4xNomos8kDAT.safetensors" \
  ./recipes/compare-models.sh sample.png out/

With no MODELS, it uses every weights file present. Writes out/sample--4x-UltraSharp.png, out/sample--4xNomos8kDAT.png, … so you can flip between them at 1:1 and pick a winner.

Write your own

Copy this skeleton into recipes/ and chmod +x it:

#!/usr/bin/env bash
# One line saying what this recipe regenerates and why.
#   ./recipes/my-recipe.sh [OUTPUT_DIR]
set -euo pipefail

HERE="$(cd "$(dirname "$0")" && pwd)"
ROOT="$(dirname "$HERE")"

OUT="${1:-./out}"
MODEL="${MODEL:-$ROOT/weights/4x-UltraSharp.pth}"
ENS=(); [ "${ENSEMBLE:-0}" = "1" ] && ENS=(--ensemble)

PY="${PYTHON:-}"
[ -z "$PY" ] && [ -x "$ROOT/.venv/bin/python" ] && PY="$ROOT/.venv/bin/python"
[ -z "$PY" ] && PY="python3"

for pair in "icon.png:icon-hd.png" "banner.jpg:banner-hd.png"; do
  "$PY" "$ROOT/upscale.py" "$OUT/${pair%%:*}" "$OUT/${pair##*:}" --weights "$MODEL" "${ENS[@]}"
done
echo "wrote HD assets to $OUT (model: $(basename "$MODEL"))"

Four conventions worth keeping:

  1. set -euo pipefail — a failed upscale halfway through a batch should stop the script, not quietly leave you with a half-updated asset folder.
  2. Resolve paths from the script, not the caller's cwd (HERE / ROOT above), so the recipe works from anywhere.
  3. Override with env vars, defaults via ${VAR:-default} — the recipe encodes the decision (this model, this ensemble setting), while still letting you A/B it without editing the file.
  4. Take paths as arguments. Never hardcode an absolute path from your machine; a recipe with /home/you/... in it is broken for everyone else, including future you.

Layout

upscale.py             the CLI — model loading, self-ensemble, alpha handling
rrdbnet.py             self-contained Real-ESRGAN generator (the no-spandrel fallback)
recipes/               committed scripts for repeatable jobs
requirements.txt
weights/               your models        (gitignored)
out/                   scratch output     (gitignored)
.venv/                 virtualenv         (gitignored)

License

MIT — for the code in this repository. Model weights you download are licensed by their respective authors; check each model's terms before redistributing it or its output.