README.md

---
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-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`](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).