--- license: apache-2.0 language: - en library_name: pytorch pipeline_tag: image-to-image tags: - super-resolution - image-restoration - image-to-image - sisr - real-world-restoration - pytorch - onnx - tensorrt - spandrel - chainner - transformer - attention datasets: - Phips/lucid-cc0-v2-hc-512 --- # 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](#relation-to-nerve) for which one to pick. All HEART models were trained on a single consumer GPU (**RTX 3060, 12 GB**) — this network trains and runs on ordinary consumer hardware, no server-class GPU required. **Links:** [NERVE](https://huggingface.co/Phips/NERVE) · [HEART](https://huggingface.co/Phips/HEART) · [NERVE demo](https://huggingface.co/spaces/Phips/nerve) · [HEART demo](https://huggingface.co/spaces/Phips/heart) · [convert a checkpoint](https://huggingface.co/spaces/Phips/nerve-convert) · [local app](https://github.com/Phhofm/nerve-app) · [train NERVE](https://github.com/Phhofm/nerve-train) · [train HEART](https://github.com/Phhofm/heart-train) > Short version: copy `heart_arch.py` into traiNNer-redux, train with > `lq_size: 96`, export to ONNX, run it (ONNX Runtime verified). Details below. ![HEART 4x on anime — input | bicubic | release | OTF GAN](examples/ani_compare.png) *Input (as-is) | bicubic x4 | HEART 4x release | HEART 4x OTF GAN.* ![HEART 4x on a noisy image](examples/noisy_compare.png) ![HEART 4x on a Stable-Diffusion image](examples/stablediffusion_compare.png) --- ## 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](https://github.com/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`](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](https://github.com/umzi2) and [Kim2091](https://github.com/Kim2091), fdat by [sharekhan](https://github.com/stinkybread), and [the-database](https://github.com/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](https://github.com/the-database/traiNNer-redux) architecture — no special steps. ```bash # 1. Add the architecture (auto-registers: traiNNer scans *_arch.py) cp heart_arch.py /traiNNer/archs/ # 2. Copy a training config and point it at your data cp configs/4x_HEART_release.yml /options/ # 3. Train cd 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 /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`](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`](https://huggingface.co/datasets/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`](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. ## Update — 2026-09-30: speed benchmark (HEART vs HAT_M) Measured on an **RTX 3060**, input `1x3x256x256` (a typical SR tile), fp16, warmup 5, mean of 20 runs. ONNX is static-shape; TensorRT uses the same builder options as NERVE's published config (`builderOptimizationLevel=5` + CUDA graphs). | model | params | PyTorch fp16 | ONNX fp16 | **TensorRT fp16** | peak VRAM | |---|---|---|---|---|---| | **HEART 4x** | 16.7M | 816 ms | 745 ms | **266 ms** (208 FPS) | **679 MB** | | HEART 4x OTF GAN | 16.7M | 809 ms | 753 ms | **270 ms** | 679 MB | | HAT_M 4x | 20.8M | 1485 ms | 849 ms | 598 ms (107 FPS) | 2428 MB | ![HEART vs HAT_M — PyTorch / ONNX / TensorRT speed and VRAM](docs/heart_speed.png) **Verdict.** At a 256px tile HEART is **~1.8x faster than HAT_M** in PyTorch and uses **~3.6x less VRAM** (679 MB vs 2.4 GB) — a bigger margin than our earlier 30k-iteration note suggested. Under **TensorRT** both get dramatically faster, and HEART keeps a **~2.2x** lead (266 ms vs 598 ms). **TensorRT is a big win for HEART** (3.1x over PyTorch), much larger than for NERVE (1.05x) — a bigger model gives the builder more to optimize. Use it: ```python import onnxruntime as ort sess = ort.InferenceSession( "4x_HEART_fp16.onnx", # static shape, e.g. 1x3x256x256 providers=[ ("TensorrtExecutionProvider", { "trt_fp16_enable": True, "trt_builder_optimization_level": 5, "trt_cuda_graph_enable": True, "trt_engine_cache_enable": True, "trt_engine_cache_path": "trt_cache", "trt_timing_cache_enable": True, }), "CUDAExecutionProvider", ], ) ``` Caveats: this is a **single size (256px)** on one GPU, and ONNX/TensorRT need a **static shape** — export one per size you use. HEART is **ONNX-only** (no ncnn; its attention ops don't convert cleanly).