HEART

Hybrid Efficient Attention with Rank-factorized bias Transformer

A super-resolution architecture that is simple to train and fast to run. Released models are verified with traiNNer-redux (training) and ONNX Runtime (inference); TensorRT and other runtimes are expected to work (same standard op set) but are not tested here.

Part of the BODY suite by Philip Hofmann — networks built to be real-world usable first, not benchmark-chasing. See Relation to NERVE for which one to pick.

Short version: copy heart_arch.py into traiNNer-redux, train with lq_size: 96, export to ONNX, run it (ONNX Runtime verified). Details below.

Input (as-is) | bicubic x4 | HEART 4x release | HEART 4x OTF GAN.

---

Why HEART?

There are strong and established SR networks already. This is the honest case for when HEART is worth your time — and when it is not.

This project comes out of a trainer's frustration rather than a benchmark goal. After training hundreds of SISR models (see Phhofm/models), the recurring annoyance was networks that felt like paper networks: they chase leaderboard metrics with techniques that leave you juggling an unfused and a fused checkpoint, or whose dynamic-ONNX conversion is fragile. This network is the opposite of that — built so a trainer can pick it up, train it, and convert it to dynamic ONNX without headaches.

Strengths of HEART

• Real-world usability over metric chasing. OTF models are trained on the full Real-ESRGAN degradation chain and restore real photos, not just clean bicubic downsamples.
• Competitive quality at lower cost. In our tests it matches HAT_M on PSNR (25.02 vs 25.05 dB, within noise) while using ~half the training VRAM (1.44 vs 2.84 GB), and at bf16 inference it is ~1.2x faster with ~2.7x less VRAM.
• One file, exports cleanly. Single-file pure PyTorch; dynamic ONNX on the first try; no fusion step or fused/unfused checkpoint pairs.
• Stable, simple training. i-LN keeps bf16 AMP stable without fp32 gymnastics.
• FlashAttention / SDPA friendly. No attention masks, no position-bias table lookups — RIB replaces them with a dot product.

When to pick something else. For maximum benchmark PSNR or a small mobile model, use NERVE or one of the community networks below. HEART is the desktop / server quality tier.

Full architecture details, design rationale and ablation results: docs/EXPERIMENTS.md.

Credit where it is due. None of this exists in a vacuum. SISR has many strong architectures and generous community members, for example HAT, ATD and DAT2, the community network work of umzi and Kim2091, fdat by sharekhan, and the-database (traiNNer-redux, its benchmarks and its models). This list is far from complete — apologies to anyone not named; the SISR community has a great many people contributing a great deal.

---

Quickstart

As a trainer (want to train your own model)

HEART is a normal traiNNer-redux architecture — no special steps.

# 1. Add the architecture (auto-registers: traiNNer scans *_arch.py)
cp heart_arch.py <traiNNer-redux>/traiNNer/archs/

# 2. Copy a training config and point it at your data
cp configs/4x_HEART_release.yml <traiNNer-redux>/options/

# 3. Train
cd <traiNNer-redux>
python train.py -opt options/4x_HEART_release.yml --auto_resume

# 4. Export to optimized dynamic ONNX (official converter)
cp configs/4x_HEART_onnx.yml <traiNNer-redux>/options/
python convert_to_onnx.py -opt options/4x_HEART_onnx.yml

Warm-start from the released pretrains by setting path.pretrain_network_g in your training config (e.g. models/heart_4x_pretrain.safetensors for a 4x run).

Rules for good results (learned the hard way — all HEART models follow these):

• lq_size must be a multiple of 32 (the window size). Use 96. Other sizes (e.g. 80) teach the model a fake "edges are mirrored" response → visible bright band + tile grid, worse with GAN loss.
• Mild noise: noise_range: [1,15], noise_range2: [1,10].
• Batch 2 is fine on a 12 GB card at lq96.
• --auto_resume is recommended when running train.py.

As a user (want to upscale images)

• ONNX Runtime (verified): run the onnx/*.onnx files (dynamic H/W). TensorRT should also work (same ops) but is untested here.

• chaiNNer: load the onnx/*.onnx files directly. Loading the .safetensors via spandrel needs HEART registered in your chaiNNer/spandrel build (not upstream yet).

• Input: RGB, any size (dynamic H/W).

• Scale: 4x (heart_4x_*), 2x (heart_2x).

---

Models

File	Role	Best for
models/heart_4x_pretrain.safetensors	4x official pretrain	starting any 4x training
models/heart_2x.safetensors	2x official pretrain	starting 2x training
models/heart_4x_otf_v2.safetensors	4x OTF fidelity	the pretrain for OTF GAN finetunes; clean real-world
models/heart_4x_otf_gan.safetensors	4x OTF GAN	sharpest real-world output
onnx/heart_4x_pretrain_fp32_op17.onnx	4x pretrain, dynamic ONNX	inference
onnx/heart_2x_fp32_op17.onnx	2x, dynamic ONNX	inference
onnx/heart_4x_otf_v2_fp32_op17.onnx	OTF fidelity, dynamic ONNX	inference
onnx/heart_4x_otf_gan_fp32_op17.onnx	OTF GAN, dynamic ONNX	inference

All ONNX files are dynamic-shape fp32, opset 17, onnxslim-optimized.

Pretrain chain (how the models relate):

4x pretrain ──→ OTF v2 (fidelity) ──→ OTF GAN
     └──→ 2x finetune

---

Visual examples

examples/ has full-resolution comparisons on the Real-ESRGAN test set (input fed as-is), each showing: input vs HEART 4x release vs HEART 4x OTF fidelity vs HEART 4x OTF GAN.

• examples/00003_compare.png (skyline)
• examples/ADE_val_00000114_compare.png (scene)
• examples/OST_009_compare.png
• examples/wolf_gray_compare.png

examples/inputs/ also contains the original Real-ESRGAN test images, so you can inspect full-size inputs and run the models yourself.

A second set (realphoto, ani, manga, text, noisy, stablediffusion) feeds real images in as-is (not downscaled):

• examples/*_compare.png — whole image: input | bicubic x4 | HEART 4x release | HEART 4x OTF GAN.
• examples/*_1to1.png — a 1:1 native-pixel center crop (detail comparison). The source inputs are in examples/inputs/.

---

What we actually tested

Every design choice came from a measured experiment, not a guess. The full log (including dead ends, benchmarks and the "considered and decided" register) is in docs/EXPERIMENTS.md. Highlights:

• HEART vs HAT_M (30k, identical settings): effectively a tie on quality (25.02 vs 25.05 dB, within noise) at ~half the peak VRAM (1.44 vs 2.84 GB), and ~1.16-1.24x faster + ~2.7x less VRAM at inference in bf16 (the deployment-relevant precision).
• ATD: OOM'd at iter 1 on a 12 GB card on the same config — a real deployment result.
• torch.compile: ~1.5x training speed, quality-neutral (validated).
• i-LN (norm choice) for stable bf16 AMP with no fp32 gymnastics.
• RIB replaces HAT's relative-position-bias table + OCAB with a FlashAttention-friendly dot product.
• Shifted-window bug found and fixed during ablations (it was quietly costing quality).
• Runtime validation: every released .safetensors loads (strict) and runs; every .onnx loads in ONNX Runtime.
• Not tested here (community welcome): TensorRT / DirectML speed and fixed-shape static exports, CoreML, and on-device benchmarks.
• OTF GAN lq96 window alignment: diagnosed the bright band + tile grid and fixed it via the lq_size multiple-of-32 rule.

---

Training data & license

All released HEART models were trained only on Phips/lucid-cc0-v2-hc-512, a CC0 dataset — the released weights are clean to use commercially.

HEART ties HAT on quality at a much lower cost, but we did not chase the benchmark leaderboard. If higher PSNR/LPIPS is your goal, training on larger academic datasets (DIV2K, Flickr2K, LSDIR, ...) is expected to help.

---

Relation to NERVE

HEART and NERVE are the two members of the BODY suite (by Philip Hofmann): same philosophy — real-world usability over benchmark scores — different tiers.

	HEART	NERVE
size	~16.7M params	~1.8M params
design	attention (FlashAttention-friendly)	pure convolution
use when	you want the best quality and have the compute	you want speed, small size, and dead-simple deployment (mobile/edge/web)
cost	heavier to run	a fraction of the cost

Rule of thumb: HEART for maximum quality, NERVE for maximum practicality. Both export to dynamic ONNX; NERVE additionally converts to ncnn cleanly because it has no attention ops.

Side-by-side comparisons on the same inputs (input | bicubic | both release models | both OTF GAN models), each with a 1:1 crop: examples/OST_009_nerve_vs_heart_* and examples/00003_nerve_vs_heart_* (Real-ESRGAN test set), plus examples/realphoto_nerve_vs_heart_* and examples/noisy_nerve_vs_heart_*.

---

Repository layout

heart_arch.py            the architecture (single file)
configs/                 example training configs (release, OTF, OTF GAN, ONNX export)
models/                  pretrained checkpoints (see table above)
onnx/                    dynamic fp32 ONNX (opset 17, onnxslim-slimmed)
examples/                visual comparisons on the Real-ESRGAN test set
docs/EXPERIMENTS.md      everything that shaped the architecture
scripts/export_heart.py  checkpoint → dynamic ONNX exporter

---

Help wanted (community welcome)

The goal here is real-world usability — easy training, hassle-free dynamic ONNX, simple and robust to maintain — not leaderboard chasing. There's plenty worth measuring. Pick anything below and open an issue/PR/discussion with your numbers.

Benchmarks

• Full-length training on an academic dataset (DIV2K / Flickr2K / LSDIR) and a head-to-head against HAT / HAT-iLN and other transformers.
• Perceptual metrics (LPIPS / TopIQ / MUSIQ) on real-world degraded images.

Speed & deployment (not benchmarked here)

• TensorRT / DirectML latency + peak VRAM on the dynamic ONNX.
• Static-shape ONNX (e.g. 1x3x256x256) vs dynamic: speed + memory.
• fp16 vs fp32 ONNX.
• CoreML conversion + on-device test.

Models & code

• More finetunes (illustration, manga, denoise).
• Architecture experiments — see docs/EXPERIMENTS.md for what we already tested (and rejected).

Small, measured contributions are welcome.

---

License & credits

Apache-2.0 (this repository's code and pretrained weights).

HEART is a re-implementation/simplification inspired by:

• HAT (arXiv:2205.04437) and HAT-iLN (arXiv:2504.06629) — architecture family and i-LN.
• SST / RIB (arXiv:2603.06738) — Rank-factorized Implicit Bias, replacing HAT's relative-position-bias table + OCAB with a FlashAttention-friendly dot-product position bias.

The architecture file is original code; verify attribution/compatibility against those papers before commercial redistribution.