UpscaleTransformer 3x — pixel-token image upscaler (upscale_model3x_v17.pth)
Input scaled 3x:
Upscaled with model result:
A small autoregressive-style Transformer model for 3× image super-resolution (upscaling).
Instead of predicting continuous values, every pixel channel is treated as a token
(vocabulary = the 256 possible byte values 0–255), and the model classifies what each
high-resolution pixel should be.
What the model does
Given a 5×5 neighborhood of one color channel from a low-resolution image (25 tokens), the model predicts the matching 3×3 block of high-resolution pixels (9 tokens). This is repeated for every pixel and every channel (R, G, B), which reconstructs the full image at 3× the width and height of the input.
LR input (W × H) ──► 5x5 context per pixel/channel ──► Transformer ──► 3x3 HR block ──► HR output (3W × 3H)
Because it is a classifier over 256 values, the raw output can occasionally be
over-confident or noisy. The companion inference script (upscaler_3x_v17_5.py) adds a
robust decoding pipeline that dramatically improves perceived quality:
| Stage | Purpose |
|---|---|
| Soft (expected-value) decoding | Uses the full softmax distribution instead of argmax — kills low-probability spikes |
| Confidence-gated bicubic blending | Pixels with low softmax confidence are blended toward a bicubic baseline |
| Despeckle | Conservative median test removes leftover isolated corrupted pixels |
| Statistics matching | Re-matches output mean/std to the source image so brightness & contrast are unchanged |
| Natural grain synthesis | Optional film/sensor-like grain so the result doesn't look "plastic" |
Architecture
| Component | Detail |
|---|---|
| Input | 25 tokens (5×5 LR context, single channel, values 0–255) |
| Embedding | nn.Embedding(256, 128) + learned positional encoding |
| Encoder | 6 × TransformerEncoderLayer, 8 heads, FF dim 512, pre-LayerNorm, dropout 0.1 |
| Head | nn.Linear(128, 256) — first 9 output positions used (3×3 HR block) |
| Loss | CrossEntropy with label smoothing 0.1 |
| Optimizer | AdamW, lr 3e-4, weight decay 1e-4, ReduceLROnPlateau |
| Size | ~2.5 M parameters (state-dict checkpoint, .pth) |
Training
- Paired data: each HR image is downscaled 3× with PIL bicubic (Catmull-Rom) to build the LR input, so the model learns to invert exactly that downscale.
- 1,000,000 random (patch, channel) samples per epoch, batch size 1024.
- Augmentation: random rotations (90° steps), horizontal/vertical flips, random brightness/contrast jitter on the target.
- Checkpoint is saved every epoch.
Usage
The checkpoint is a plain state_dict for the UpscaleTransformer class defined in
upscaler_3x_v17_5.py.
Upscale an image (recommended settings):
python upscaler_3x_v17_5.py --mode upscale \
--lr "input.png" --out "output.png" \
--model upscale_model3x_v17.pth --conf 0
Train / continue training on your own images (PNGs in a folder, LR pairs are generated automatically):
python upscaler_3x_v17_5.py --mode train --hrdir ./my_images --model upscale_model3x_v17.pth
Upscale a video (requires ffmpeg):
python upscaler_3x_v17_5.py --mode upscale --videoin in.mp4 --videoout out.mp4 --model upscale_model3x_v17.pth
Run as an HTTP server (POST /upscale/image, multipart field lrImage):
python upscaler_3x_v17_5.py --mode server --model upscale_model3x_v17.pth --port 8080
Key inference knobs
| Flag | Default | Effect |
|---|---|---|
--conf |
0.30 |
Confidence below this blends toward bicubic; 0 = trust the model fully (sharpest) |
--no-blend |
off | Disable bicubic confidence blending |
--no-despeckle |
off | Disable impulse-noise cleanup |
--no-stats |
off | Disable brightness/contrast restoration |
--noise |
0.0 |
Add natural grain (sigma as % of 0–255), e.g. 2.0 |
--hard |
off | Original v17 argmax decoding (for comparison — noisier) |
Load the checkpoint in your own code
import torch
from upscaler_3x_v17_5 import UpscaleTransformer
model = UpscaleTransformer()
model.load_state_dict(torch.load("upscale_model3x_v17.pth", map_location="cpu"))
model.eval()
# input: LongTensor (N, 25) of 0..255 pixel values -> logits (N, 9, 256)
Requirements
Python 3, PyTorch, NumPy, OpenCV (cv2), Pillow. Optional: tqdm, ffmpeg (video mode).
Works on CPU; CUDA is used automatically when available.
Limitations
- Fixed 3× scale factor; trained to invert bicubic downscaling specifically.
- Per-channel prediction (no cross-channel context), so saturated colors can show slight channel misalignment.
- Small model — best on photos/illustrations with moderate detail; it will not hallucinate fine detail the way large diffusion-based restorers do.
Model card files
upscale_model3x_v17.pth— model weights (state_dict)- Inference/training script:
upscaler_3x_v17_5.py(model definition must match to load the checkpoint)