HEART release package: arch, configs, ONNX, README, license

HEART.md

@@ -0,0 +1,583 @@
+# HEART — Efficient Attention with Rank-factorized bias Transformer
+
+**TL;DR:** window-attention SR transformer. HAT-iLN quality, FlashAttention
+speed, none of the complexity. **One architecture, one config. Trains in bf16.
+Exports to dynamic ONNX with a single command** — no fused/unfused checkpoint
+pairs, no custom CUDA kernels, no relative-position tables.
+
+## What is HEART?
+
+HEART is what HAT-iLN becomes when you strip out the parts nobody actually
+deploys:
+
+- **RIB** replaces the relative-position-bias *table* — learned low-rank
+  position features concatenated onto Q/K, so attention runs on standard
+  FlashAttention/SDPA kernels. No table, no index gather, no mask.
+- **OCAB removed** — HAT's most complex, most memory-hungry block. Our
+  ablations showed it isn't needed at window 32.
+- **Keeps the proven bits:** i-LN (input-adaptive normalization, bf16-stable),
+  window attention, a convolutional branch (CAB).
+
+The result: a ~16.7M single-image SR architecture that **ties HAT on PSNR,
+trains in stable bf16 AMP, and exports to ONNX that runs in TensorRT / ONNX
+Runtime / DirectML.**
+
+## Why HEART exists
+
+Most SR papers chase benchmark numbers by adding modules that are annoying to
+train and painful to deploy — dynamic routing, deformable attention, blocks
+that only work with a specific resolution, or a "fused" checkpoint for
+inference and an "unfused" one for fine-tuning. The result looks good on a
+leaderboard and gathers dust in practice.
+
+HEART has a different goal: **real-world simplicity.** Train it anywhere,
+convert it to dynamic-shape ONNX, run it everywhere. No fusion step. No second
+model file. No custom kernels.
+
+## When to use HEART
+
+- You want HAT-class quality on real images (denoise, deblur, upscale) without
+  babysitting the training.
+- You need weights you can convert to ONNX and ship — no fusion step, no "use
+  variant B for inference."
+- You are tired of benchmark-chasing networks that add modules no one can
+  actually deploy.
+- You need desktop quality at reasonable inference cost. For mobile/edge, use
+  NERVE (coming in the BODY suite).
+
+## Quickstart
+
+1. Copy `traiNNer/archs/heart_arch.py` into your traiNNer-redux clone
+   (`traiNNer/archs/`).
+2. Copy the release config from `options/train/RELEASE/4x_HEART_release.yml`
+   and point `dataroot_gt` / `dataroot_lq` at your data.
+3. Train:
+   ```bash
+   python train.py -opt options/train/RELEASE/4x_HEART_release.yml --auto_resume
+   ```
+   (`--auto_resume` is **mandatory** — without it an interrupted run restarts
+   from iter 0 and archives the experiment dir.)
+4. Convert to dynamic ONNX (each model, one at a time):
+   ```bash
+   ulimit -v 14000000
+   ./venv/bin/python scripts/heart/export_minimal.py \
+       experiments/4x_HEART_pretrain600k/models/net_g_ema_615000.safetensors \
+       onnx/4x_HEART_1x3xHxW_fp32_op17.onnx 4
+   ```
+5. Run inference anywhere ONNX Runtime or TensorRT runs.
+
+Full architecture docs, ablation results, and training recipes are below.
+
+> **Training crop size: use `lq_size: 96` (a multiple of the 32 px window).**
+> HEART partitions the image into 32×32 windows. Training at a crop size that
+> is *not* a multiple of 32 (e.g. 80) makes the model reflect-pad its own
+> bottom/right edge every batch, so it learns a false "edge = mirrored
+> content" response — visible as a bright bottom band and a tile-grid pattern,
+> and strongly amplified by a GAN loss. All HEART ablations and release runs
+> used `lq_size: 96` (= 3×32); a larger multiple (128) also works. This
+> matters most for GAN/OTF finetunes, which are most sensitive to boundary
+> statistics.
+
+## BODY suite status
+
+| Model | Role | Params | Status |
+|---|---|---|---|
+| HEART | desktop quality tier | 16.7M | trained + ONNX |
+| NERVE | mobile/edge tier | ~270K | trained |
+
+## Design
+
+```
+conv_first → [6× RHAG-RIB] → conv_after_body (+ global residual)
+           → pixelshuffle upsampler
+
+RHAG-RIB (residual group):
+  [6× HAB-RIB blocks] → 3x3 conv → group residual
+
+HAB-RIB block:
+  i-LN → RIB window attention (window 32, FlashAttention/SDPA)
+       + CAB conv-attention branch (scale 0.01)
+  i-LN → MLP (mlp_ratio 2)
+```
+
+- **i-LN** (Image Restoration tailored LayerNorm, arXiv:2504.06629) — stable
+  training in bf16/AMP, no divergence.
+- **RIB window attention** — FlashAttention-compatible, window 32, shifted
+  windows via non-wrapping pad-and-partition (no masks, no position tables).
+- **CAB conv branch** — HAT's proven local inductive bias, kept.
+- **OCAB removed** — the most complex and memory-hungry part of HAT is gone.
+
+## Goals
+
+| Goal | How |
+|---|---|
+| Great results | HAT's proven hybrid block (window + conv) + large-window global context |
+| Fast train & inference | FlashAttention via RIB (SST measured 2.1x train / 2.9x inference vs RPB) |
+| Stable training | i-LN + EMA std clamping, bf16/AMP safe |
+| Simple code | ~640 lines, reuses existing blocks, no masks/tables/OCAB |
+| Easy deployment | Pure PyTorch, ONNX-exportable, no custom CUDA kernels |
+
+Note on speed: the FlashAttention win only exists in bf16/fp16 (measured
+~1.6x faster than HAT-iLN-M at bf16 inference). In fp32, HEART is slightly
+slower than HAT-iLN — fp32 has no flash kernel. Use reduced precision.
+
+## Status
+
+Implemented, verified, and releasing. Architecture is frozen; all weights and
+ONNX files below are measured, not extrapolated.
+
+Final architecture (post-ablation, 2026-08-31): RIB window-32 attention + CAB,
+`attention_freq=2`, rank 8, linear MLP — the 20k sweep winner. RoPE,
+no-positional, GDFN, SwiGLU, LayerScale, window 48/64, and rank 18/32 were all
+measured and rejected; the losing code paths were removed from `heart_arch.py`
+(646 lines, 16/16 tests, lint + typecheck clean).
+
+Release artifacts (all exist, paths verified):
+
+| Artifact | Checkpoint | ONNX |
+|---|---|---|
+| 4x bicubic | `experiments/4x_HEART_pretrain600k/models/net_g_ema_615000.safetensors` | `onnx/4x_HEART_1x3xHxW_fp32_op17.onnx` |
+| 2x | `experiments/2x_HEART/models/net_g_ema_200000.safetensors` | `onnx/2x_HEART_1x3xHxW_fp32_op17.onnx` |
+| 4x OTF fidelity (v1.1) | `experiments/backup_otf_v2/net_g_ema_150000.safetensors` | `onnx/4x_HEART_OTF_1x3xHxW_fp32_op17.onnx` |
+| 4x OTF perceptual (GAN) | training | — |
+
+All checkpoints backed up to a private HF repo
+(`Phips/heart-sisr-pretrain`). ONNX files pass `onnx.checker` and match
+PyTorch numerically (ORT diff < 1e-2).
+
+In flight: **4x OTF perceptual (GAN) finetune** (~3 days remaining). When it
+finishes, its weights + ONNX join the release set.
+
+Known limitation: a faint **window-32 grid artifact** (128 px period in 4x
+output) on smooth gradients and fine lines in OTF outputs. Diagnosed via FFT;
+root cause is window boundaries (not JPEG blocks or PixelShuffle). Overlap
+attention contradicts the frozen arch, so it is documented here, not fixed in
+1.0.
+
+Release-gate checks passed: `--auto_resume` (CLI flag) resume verified,
+tiled inference seam-free, ONNX legacy-tracer opset 17 + onnxruntime,
+fp16/bf16 inference, bf16 AMP training stability, torch.compile validated.
+  the bare model, so checkpoints and val are unaffected. Kernel profile of
+  one step: ~23% flash attention, ~30% small elementwise ops (what compile
+  fuses), ~12% convs — no architectural hotspot left for custom kernels.
+
+Dev notes for anyone touching the attention code:
+
+- Shifted windows are non-wrapping: the feature map is padded top/left by half
+  a window (reflect) and windows are partitioned from the padded origin. Border
+  windows see reflected padding, never content from the opposite image border.
+  Do NOT reintroduce `torch.roll` cyclic shifting — without an attention mask
+  it couples opposite borders, and adding a mask would force SDPA off the flash
+  kernel (the whole point of RIB).
+- Logits are pre-scaled per SST eq. 5 (`q/√D`, `pos/√R`) and SDPA is called
+  with `scale=1.0`. Re-adding default SDPA scaling silently divides logits
+  by √48 ≈ 7 and shifts attention toward uniform — shape-only tests will
+  not catch it. Without this fix the model trains around it and loses
+  quality for no reason.
+- RIB features are computed in fp32 then cast to the input dtype. Skipping
+  the cast breaks pure fp16 inference (buffer/param dtype mismatch).
+
+Next step is a short ablation against the baselines before the full training
+schedule:
+
+- `heart` (this network, 19.2M params, RIB window 32)
+- `atd` (20.3M, the efficiency champion)
+- `hat_m` (20.8M, classic HAT)
+
+The ablation trains each for ~20-50k iterations on DIV2K and compares
+Urban100 PSNR and wall-clock speed. HEART proceeds to full DF2K training only
+if it wins.
+
+## Ablation results (2026-08-15, post shifted-window fix)
+
+30k iterations each, identical settings (post-fix HEART code). RTX 3060 12 GB,
+DIV2K → Urban100:
+
+| arch  | best Urban100 PSNR | @iter | final SSIM | train it/s | peak VRAM |
+|-------|--------------------|-------|------------|------------|-----------|
+| HEART | 25.0235            | 30000 | 0.7497     | 0.86       | 1.44 GB   |
+| HAT_M | 25.0524            | 30000 | 0.7517     | 0.95       | 2.84 GB   |
+| ATD   | OOM at iter 1 — did not run on the same config |
+
+Reading the table honestly:
+
+- **Quality:** HAT_M leads by +0.03 dB / +0.002 SSIM at the 30k cutoff —
+  within noise for a short schedule, and both were still improving on their
+  final validation. Effectively a tie; only a full-length run separates them.
+- **VRAM:** HEART uses ~half the peak memory of HAT_M (1.44 vs 2.84 GB).
+- **Train speed:** HAT_M trains ~10% faster (0.95 vs 0.86 it/s). Window 32
+  attention is 4x the per-pixel attention work of window 16; the flash kernel
+  makes it feasible, not free. HEART's speed advantage is at *inference*
+  (~1.6x vs HAT-iLN-M bf16), not training throughput.
+- **ATD OOM caveat:** HEART/HAT_M ran with `use_checkpoint: true`; ATD's
+  factory has no such option. The OOM is a real deployment-relevant result
+  (ATD does not train on a 12 GB card at lq_size 64 / batch 4 as configured),
+  but the memory comparison is not fully symmetric because of that.
+
+### Inference benchmark (same trained EMA checkpoints, RTX 3060)
+
+Median of 5 runs after warmup, `cudnn.benchmark=False`,
+`/tmp/kilo/inference_bench.py`:
+
+| arch  | dtype | input    | time    | peak VRAM |
+|-------|-------|----------|---------|-----------|
+| HEART | bf16  | 320x180  | 0.884s  | 0.45 GB   |
+| HAT_M | bf16  | 320x180  | 1.092s  | 1.20 GB   |
+| HEART | bf16  | 480x270  | 2.024s  | 0.90 GB   |
+| HAT_M | bf16  | 480x270  | 2.351s  | 2.47 GB   |
+| HEART | fp32  | 320x180  | 2.350s  | 0.89 GB   |
+| HAT_M | fp32  | 320x180  | 1.816s  | 2.37 GB   |
+| HEART | fp32  | 480x270  | 5.347s  | 1.78 GB   |
+| HAT_M | fp32  | 480x270  | 3.973s  | 4.94 GB   |
+
+In bf16 (the deployment-relevant precision) HEART wins both speed (1.16-1.24x)
+and VRAM (~2.7x less). In fp32 HAT_M is faster (no flash kernel for fp32), but
+HEART still uses 2.7-2.8x less VRAM.
+
+### Training speed benchmarks (RTX 3060, bf16 AMP, channels-last)
+
+torch.compile (`use_compile: true`) validated in the training harness:
+~1.5x training speed, quality-neutral (300-iter smoke: 21.88 dB vs 21.80 dB
+uncompiled).
+
+| config | it/s | 800k iters | peak VRAM |
+|---|---|---|---|
+| lq64 bs4, no compile (ablation) | 0.86 | ~10.7 days | 1.44 GB |
+| lq64 bs4 + compile | 1.26 | ~7.4 days | 1.44 GB |
+| lq96 bs2 + compile | 1.19 | ~7.8 days | 0.98 GB |
+| lq96 bs4 + compile | ~0.63 | ~14.7 days | 1.70 GB |
+| lq96 bs4 / lq128, no checkpoint | OOM | — | — |
+
+compile_mode comparison (@ lq96 bs2, 3060): `reduce-overhead` 1.23 it/s >
+`default` 1.18 > `dynamic=False` 1.19 ≈ `default+TF32` 1.19 > no compile 0.77.
+`max-autotune` is not worth it on small GPUs (inductor warns: not enough SMs).
+At the ablation config (lq64 bs4), `max-autotune-no-cudagraphs` was measured
+equal to `reduce-overhead` (both 1.28 it/s) with 2x the compile time and
+slightly more VRAM — it gains nothing on a 28-SM card because the tuned GEMM
+search falls back. `fast_matmul` (TF32) is a no-op under bf16 AMP — measured.
+
+Kernel profile of one training step (uncompiled): flash attention ~23%,
+small elementwise ops ~30% (what compile fuses — the source of the 1.5x),
+convs ~12%, MLP gemms ~4%. No hotspot remains that a custom Triton kernel
+could profitably target.
+
+### Scoreboard vs HAT_M (30k ablation)
+
+| Metric | Winner |
+|---|---|
+| Quality | tie (HAT_M +0.03 dB, within noise) |
+| Inference speed bf16 | HEART (1.16-1.24x) |
+| Inference VRAM | HEART (~2.7x less) |
+| Training VRAM | HEART (1.44 vs 2.84 GB) |
+| Training speed (no compile) | HAT_M (+10%) |
+| ATD on 12 GB | does not run (OOM) |
+
+### Rank ablation (done, 2026-08-17)
+
+RIB's per-head positional-bias matrix has rank <= `rank`. Three 20k-iteration
+runs compared rank 8 / 18 / 32 on the identical ablation config **with
+torch.compile reduce-overhead** (`HEART_rank{8,18,32}_ablation.yml`, lr decay
+at 16k). All three used compile so the comparison is not confounded (compile
+shifts short-run PSNR ~0.08 dB, larger than the 0.03 dB decision threshold).
+
+| rank | 20k PSNR |
+|---|---|
+| **8** | **24.7820** |
+| 32 | 24.7586 |
+| 18 | 24.7346 |
+
+**Winner: rank 8** (smallest within 0.03 dB of the best). Rank 8 beating
+rank 18 means positional capacity is *not* the bottleneck at rank 8 — a
+lower-rank positional mechanism is enough, and it is leaner (concat dim 40
+vs 48) and slightly faster. The canonical default in `heart_arch.py` is now
+`rank=8`. The rank-18 uncompiled 30k reference (25.0235) was historical only.
+The release configs are patched to rank 8.
+
+### Architecture ablation: window size & attention frequency (2026-08-19)
+
+Ran because an 8-day release was at stake: if window 48/64 or sparser
+attention changed the quality-per-second picture, it had to be found before
+committing. Setup: lq96 bs2 (the release crop), 20k iters, compile
+reduce-overhead, rank 8, DIV2K, seed 1024 — identical settings except the
+variable (`HEART_{w32,w48,w64,attn2}_ablation.yml`). The release run was
+halted at 75k iters (checkpointed, resumable) while this ran.
+
+| config | 20k PSNR | it/s (compiled) | params |
+|---|---|---|---|
+| **attn2** (`attention_freq=2`, RIB) | **24.6282** | **2.02** | 16.68M |
+| attn3 (`attention_freq=3`) | 24.5606 | ~2.3 | 15.87M |
+| w32 (baseline, RIB rank 8) | 24.5063 | 1.29 | 19.10M |
+| none (no positional, freq=1) | 24.4965 | ~1.4 | 18.94M |
+| w48 | 24.4347 | 0.80 | 19.10M |
+| rope (RoPE, freq=1) | 24.3615 | ~1.4 | 18.94M |
+| gdfn (Gated-Dconv FFN, ratio 2.0, freq=2) | 24.2546 | ~1.9 | 19.28M |
+| swiglu (SwiGLU FFN, ratio 4/3 param-matched, freq=2) | 23.7403 (best @12k; collapsed to 23.21 @20k) | ~1.9 | ~16.7M |
+| qknorm (per-head QK LayerNorm, freq=2) | 24.6334 (tie with attn2, +0.005) | ~2.0 | ~16.7M |
+| swiglu + qknorm | 23.9576 | ~1.9 | ~16.7M |
+| layerscale (per-block residual scaling, init 1e-6) | 24.0288 | ~2.0 | ~16.7M |
+| w64 | skipped (w48 already lost; 4.2x slower) | 0.31 | 19.10M |
+
+Reading:
+
+- **attention_freq=2 wins on both axes**: +0.12 dB over the baseline AND
+  1.56x faster training, with 2.4M fewer params. HEART has more attention
+  than it needs at 20M scale.
+- **GDFN rejected** (-0.37 dB vs attn2): doubling the per-block 3x3 conv
+  density (CAB already provides spatial convs) hurt. The FFN stays a linear
+  MLP.
+- **Positional barely matters**: no-positional ties RIB at freq=1 (-0.01 dB,
+  noise), while RoPE's fixed rotation actively hurt (-0.15 dB). RIB stays as
+  cheap insurance; it is not a quality driver.
+- **Window 48/64 closed**: window size is not a quality bottleneck at this
+  crop size (w48 -0.07 dB, 38% slower). SST's large-window gains needed
+  scaled crops+datasets this regime doesn't have.
+- **Phase 2 confirmation** (in progress): w32 vs attn2 head-to-head at 100k
+  iters (`HEART_{w32,attn2}_confirm.yml`, lr decay at 50k). The 20k ordering
+  must survive to 100k before the release config changes.
+
+### Release training (actual pipeline, 2026-08-25..28)
+
+The release trained in five chained phases, all unattended with
+`--auto_resume`. This is the pipeline that actually ran (not the earlier
+3-phase design, which was superseded).
+
+| Phase | Config | Iters | Result |
+|---|---|---|---|
+| 1. 4x pretrain (bicubic) | `options/train/RELEASE/4x_HEART_release.yml` | 615k (early-stopped) | 26.2139 dB |
+| 2. 2x finetune | `options/train/RELEASE/2x_HEART_release.yml` | 200k | 33.9959 dB |
+| 3. 4x OTF v1.0 (real-world) | `options/train/RELEASE/4x_HEART_OTF.yml` | 300k | 22.8473 dB |
+| 4. 4x OTF v1.1 (anti-alias) | `options/train/RELEASE/4x_HEART_OTF_v11_finetune.yml` | 150k | 22.9529 dB |
+| 5. 4x OTF perceptual (GAN) | `options/train/RELEASE/4x_HEART_OTF_gan.yml` | 150k | in progress |
+
+Key operational lessons learned the hard way:
+
+- **`--auto_resume` is mandatory** (CLI flag, not YAML). Without it an
+  interrupted run restarts from iter 0 and archives the experiment dir.
+- **Dataset:** LUCID CC0 v2 HC-512 (100,866 tiles, 512x512, CC0). LR generated
+  with chainner CubicCatrom bicubic (the chaiNNer ecosystem standard).
+- **OTF degradation:** Real-ESRGAN-style two-pass blur + noise + JPEG, applied
+  on-the-fly by `realesrgandataset`.
+- **ONNX export froze the desktop twice** (see Part 1 of this plan) — the fix
+  is a per-process memory cap, not `nice`/`ionice`.
+
+Pipeline scripts (durable, in `scripts/heart/`): `download_lucid.py`,
+`prepare_lucid.py`, `run_release_sequence.sh`, `watch_download_then_prepare.sh`,
+`export_one.py`, `export_minimal.py`, `infer_otf_comparison.py`,
+`grid_diagnosis.py`. Logs in `/tmp/kilo/`.
+
+## Variants
+
+HEART 1.0 ships as a single architecture (`heart`, 19.2M params, embed 180,
+6×6 blocks, window 32, CAB compress/squeeze 3/30). A principled S/M/L scaling
+family (width, depth, heads, rank, CAB) is a HEART 2.0 item — deliberately
+not added before the release model proves the architecture.
+
+## Positioning
+
+HEART deliberately does **not** compete on benchmark PSNR. Research
+architectures (HAT, DAT, RGT, Restormer, DRCT, Mamba-based SSMs) exist to
+move leaderboard numbers, and they will out-score HEART on some val sets —
+that is expected and accepted.
+
+HEART competes on everything that happens *after* the paper: training
+without babysitting, exporting without a fuse step, running in chaiNNer
+tomorrow, and being understood by a new maintainer in 2028.
+
+The rules that define the position:
+
+- **One configuration.** No S/M/L variants, no fused/unfused checkpoint
+  pairs, no "lite" forks. One architecture, one way to run it.
+- **Zero custom kernels.** Everything a user needs ships in vanilla
+  PyTorch; ONNX opset-17 export is verified, not promised.
+- **Anything that loses the ablation protocol is removed from the code.**
+  14 variants were measured; the winners stayed, the losers were deleted —
+  not kept behind flags.
+- **The architecture is boring on purpose.** The block is i-LN -> [RIB
+  window attention | CAB] -> MLP. No dynamic routing, no deformable
+  anything, no state-space branches.
+
+If a technique requires a fuse step, a custom kernel, a shape assumption,
+or a "use variant B only for inference" note, it does not belong in HEART.
+
+## Deployment notes
+
+- **Arbitrary input sizes:** `check_image_size` pads the input to a window
+  multiple (reflect, replicate fallback for tiny images) and crops the
+  output back. Verified at 32x32, 64x64, 128x128, 323x711, 777x65,
+  720x1280. Users never handle padding themselves.
+- **ONNX (verified, 2026-08-21, trained 20k checkpoint):** legacy tracer
+  (`dynamo=False`), opset 17, dynamic H/W + batch. `onnx.checker` passes;
+  onnxruntime output matches PyTorch (max abs diff 0.00026 at 128x128).
+  Export pipeline for releases: export -> onnxsim -> ORT optimizer ->
+  validate at 64x64 / 256x256 / 720x1280. Note: ORT *CPU* inference of
+  large images needs several GB of RAM (a 16.7M transformer is not a CPU
+  workload) — deploy on GPU providers (CUDA/TensorRT/DirectML).
+- **NCNN: not supported, by design.** The fused SDPA node does not convert
+  to NCNN, and window attention is not an NCNN-style workload anyway.
+  HEART targets ONNX Runtime / TensorRT / DirectML; a mobile-conv sibling
+  (NERVE) is the future NCNN member of the suite.
+- **Large images — single pass vs tiling:** single pass fits up to ~720x1280
+  LR on a 12 GB card (5.55 GB measured, bf16). 1080x1920 LR single-pass
+  OOMs (~10.6 GB); tiled inference (tile 256, overlap 32, weighted blend)
+  runs it at 0.86 GB and is seam-free (measured). Use chaiNNer seamless
+  tiling / the repo's `tile_size` option for anything bigger than 720x1280.
+- **Real-world degradations:** the backbone is degradation-agnostic. For
+  denoising/JPEG cleanup train at `scale=1` (verified working) with a
+  Real-ESRGAN-style degradation pipeline, same as HAT-R relates to HAT.
+- **spandrel:** registered in-repo via `SPANDREL_REGISTRY` +
+  `store_hyperparameters`, so traiNNer-redux/chainer forks that carry this
+  repo load it directly. **Upstream-PR decision (made):** at PR time,
+  `heart_arch.py`'s imports from `arch_util` (iLN) and `hat_iln_arch`
+  (CAB, AffineTransform, DropPath, Mlp, PatchEmbed, PatchUnEmbed, Upsample)
+  get vendored into the self-contained spandrel copy (~200 lines) rather
+  than landing HAT-iLN in spandrel first. No code change in this repo until
+  then.
+
+## Experiment log (everything that shaped HEART)
+
+Chronological record of every experiment, success and failure. The failures
+matter as much as the wins — each one closed a direction.
+
+| When | Experiment | Result | Decision / lesson |
+|---|---|---|---|
+| design | HAT-iLN base: RPB table -> RIB, OCAB removed, window 32 | — | Attention runs on FlashAttention/SDPA; simpler than HAT |
+| 08-14 | Shifted windows via `torch.roll` | BUG: cyclic wrap coupled opposite borders (no mask) | Fixed: non-wrapping pad-and-partition |
+| 08-14 | `F.pad` for V concat | minor | cleanup |
+| 08-15 | 30k arch comparison HEART vs HAT_M vs ATD | 25.02 vs 25.05 vs OOM | tie with HAT_M at equal settings -> proceeded |
+| 08-15 | Inference bench vs HAT_M | bf16: 1.16-1.24x faster, ~2.7x less VRAM; fp32 slower | deployment story is bf16 |
+| 08-15 | torch.compile (smoke + harness) | +46-54% training, quality-neutral | adopted (`use_compile: true`) |
+| 08-15 | compile_mode sweep | reduce-overhead 1.23 > default 1.18 > dynamic=False/TF32 1.19; max-autotune-no-cudagraphs == reduce-overhead with 2x compile time | reduce-overhead adopted; max-autotune rejected (28 SMs) |
+| 08-15 | `fast_matmul` (TF32) | no-op under bf16 AMP | rejected |
+| 08-15 | Crop/batch probes | lq96 bs4 OOM without checkpoint; lq96 bs2 fits, 9 windows | release crop = lq96 bs2 |
+| 08-15 | `heart_light` (8.7M variant) | created, then removed | S/M/L family deferred to v2; 1.0 = one architecture |
+| 08-17 | Rank ablation 8/18/32 @20k | 8 wins (24.782) > 32 (24.759) > 18 (24.735) | rank 8 default — positional capacity is NOT the bottleneck |
+| 08-18 | Window ablation 48/64 @20k | w48 loses 0.07 dB, 38% slower; w64 skipped | spatial context is NOT the bottleneck at lq96 |
+| 08-18 | Attention frequency 1/2/3 @20k | freq=2 wins: 24.628 (+0.12 dB) at 1.56x speed; freq=3 drops | attention density sweet spot = every 2nd block |
+| 08-19 | RoPE instead of RIB @20k | loses 0.15 dB (24.36 vs 24.51) | learned low-rank positional beats fixed rotation |
+| 08-19 | No positional at all @20k | ties RIB at freq=1 (24.497 vs 24.506) | RIB stays as cheap insurance; not a quality driver |
+| 08-19 | GDFN (gated depthwise FFN) @20k | loses 0.37 dB (24.25 vs 24.63) | more per-block conv density hurts (CAB already convs) |
+| 08-20 | SwiGLU FFN @20k | loses decisively: best 23.74 @12k, declined to 23.21 @20k | gated FFN destabilizes the i-LN-rescaled block — MLP stays |
+| 08-20 | QK-Norm @20k | 24.6334 — tie with attn2 (+0.005) | tie-breaker: plain (simplest) — redundant with i-LN + RIB |
+| 08-20 | SwiGLU + QK-Norm @20k | 23.9576 | gate instability again |
+| 08-20 | LayerScale @20k | 24.0288 (−0.60) | 1e-6 init over-damps the branches; existing i-LN+conv_scale damping is right |
+| 08-20 | Final cleanup | all test switches (ffn_type/qk_norm/layer_scale) stripped; canonical default attention_freq=2 | the file is the single-path measured-best architecture |
+| 08-16 | Resume test | `--auto_resume` is a CLI flag, not YAML | all training commands pass it |
+| 08-16 | Tiled inference | seam-free (seam diff == interior diff) | large-image tiling safe |
+| 08-14..20 | ONNX export | legacy tracer opset 17 + checker + onnxruntime pass | deployment path verified |
+
+Process lessons (drivers/pipeline): /tmp gets wiped on reboot — logs restored
+from durable experiment copies; `ls dir/*.png | wc -l` breaks at ~100k files
+(ARG_MAX) — use `find`; drivers are idempotent/self-healing and training runs
+deprioritized (nice/ionice) so the machine stays usable.
+
+## Considered-and-decided register
+
+Every idea that was considered, with the decision and reason. If someone asks
+"why didn't you use X" — the answer is here. Status: ✅ integrated (and
+measured), ❌ rejected by measurement, 🚫 discarded by analysis (never run),
+⏳ pending, 📌 deferred (not forgotten).
+
+### Integrated
+
+| Idea | Why |
+|---|---|
+| ✅ RIB instead of RPB table | Concatenated position features make attention FlashAttention/SDPA-compatible; measured (ties HAT_M, rank 8 suffices) |
+| ✅ i-LN normalization | Input-adaptive holistic norm; bf16 stability is the paper's core claim, verified in training |
+| ✅ CAB conv branch | Local inductive bias; the measured carrier of quality at this scale |
+| ✅ Non-wrapping shifted windows | Fixes the cyclic-wrap border bug without the Swin mask (which would kill the flash kernel) |
+| ✅ `attention_freq=2` | Measured winner: +0.12 dB at 1.56x speed vs attention-every-block |
+| ✅ rank 8 RIB | Measured winner of 8/18/32 sweep |
+| ✅ Window 32 | w48/64 measured losers at lq96 |
+| ✅ torch.compile reduce-overhead | +46% training speed, quality-neutral; best of the mode sweep |
+| ✅ bf16 AMP + gradient checkpointing | Verified stable; 1.4-1.7 GB training VRAM |
+| ✅ Chainer CubicCatrom LR (release data) | The chaiNNer ecosystem standard — models behave as users generate LR |
+
+### Rejected by measurement (ran, lost)
+
+| Idea | Result | Why rejected |
+|---|---|---|
+| ❌ RoPE positional | −0.15 dB @20k | fixed rotation worse than learned low-rank positional |
+| ❌ No positional | ties RIB at freq=1, but attn2+RIB wins overall | RIB stays as cheap insurance |
+| ❌ GDFN gated-FFN | −0.37 dB @20k | more per-block conv density hurts (CAB already provides spatial convs) |
+| ❌ SwiGLU FFN (param-matched 4/3) | best 23.74 @12k, declined to 23.21 @20k (−0.9 to −1.4 dB vs attn2) | the multiplicative gate destabilizes the i-LN-rescaled block (loss degraded after 12k) |
+| ❌ SwiGLU + QK-Norm | 23.9576 @20k | same gate instability |
+| ❌ LayerScale (init 1e-6) | 24.0288 @20k (−0.60 dB) | the near-zero init suppresses the branches through warmup; i-LN rescale can't compensate — the block's existing damping (i-LN + conv_scale 0.01) is the right amount |
+| ⚖️ QK-Norm | 24.6334 — tie with attn2 (24.6282, +0.005 dB) | tie-breaker: plain wins (no component for zero measured gain) |
+| ❌ Window 48 / 64 | w48 −0.07 dB + 38% slower; w64 skipped | spatial context is not the bottleneck at lq96 |
+| ❌ `attention_freq=3` | −0.07 dB vs freq=2 | density sweet spot is every 2nd block |
+| ❌ rank 18 / 32 | rank 8 wins | positional capacity is not the bottleneck |
+| ❌ max-autotune compile modes | == reduce-overhead, 2x compile time | tuned GEMM search falls back on 28-SM GPUs |
+| ❌ `fast_matmul` (TF32) | no-op | everything already runs bf16 on tensor cores |
+| ❌ `compile(dynamic=False)` | no gain | shapes were effectively static anyway |
+| ❌ `heart_light` 8.7M variant | created, then removed | hand-picked config, not a principled family; S/M/L is a v2 item |
+
+### Discarded by analysis (never run — the reason is the evidence)
+
+| Idea | Why discarded |
+|---|---|
+| 🚫 OCAB (HAT's overlapping cross-attention) | RPB tables + gather construction + memory are exactly what RIB/HEART removed; the 30k ablation already ties HAT_M *with* OCAB — the missing piece isn't needed |
+| 🚫 Swin-style attention mask (shift fix alternative) | an additive mask forces SDPA off the flash kernel — the whole point of RIB |
+| 🚫 RepCAB (reparameterized CAB) | `fuse()`-before-export creates two weight formats and breaks the drop-in property, for ~2-3% inference; quality hypothesis contradicted by the GDFN result |
+| 🚫 MDTA transposed channel attention | global channel-covariance context — falsified direction (see w48) |
+| 🚫 Stripe/multi-axis attention (HMA) | full-row/column context — same falsified hypothesis, plus partitioning complexity |
+| 🚫 Dense residuals (DRCT) | groups and blocks already have residuals; dense skips grow activation memory |
+| 🚫 Mamba/SSM backbones | custom CUDA kernels (`selective_scan`) violate the zero-custom-kernel constraint |
+| 🚫 IET adaptive token selection | dynamic gather/indexing breaks dynamic ONNX tracing |
+| 🚫 UCAN / SAT / FPLIA | sub-1M-parameter tier techniques; irrelevant at ~16.7M |
+| 🚫 Muon optimizer | its gains are LLM-pretraining-scale; the RIB net is ~4k params |
+| 🚫 FlashBias Triton kernel | HEART already *is* the SDPA variant of FlashBias; no bias-add attention exists to accelerate |
+| 🚫 LAformer swap-in | different architecture, not an optimization; the deployability premise is already satisfied by HEART |
+| 🚫 Triton-fused RIB concat | ~2-4% of a step at the cost of custom kernels; torch.compile already fuses the elementwise ops |
+| 🚫 4D-native block refactor | single-digit % after compile's 46% landed; complexity not justified |
+| 🚫 Hybrid window schedule (16/32 mix) | smaller windows contradict the measured landscape (w48 lost; context isn't the lever) |
+| 🚫 lq128 crops | OOM without checkpointing and pathological compile on this GPU |
+| 🚫 Progressive patch training (LQ64 -> LQ96) | measured apples-to-apples (bs2, compiled): lq64 3.86 it/s vs lq96 1.52 it/s (2.5x per iteration), but pixel throughput only 31.6 vs 28.0 kpx/s (+13%) — and lq64 crops see 4 windows vs lq96's 9, weakening cross-window mixing. ~8% total wall-time for real pipeline complexity and quality risk |
+
+### Pending
+
+None — the ablation program is complete. 14 variants measured; the core
+(RIB rank-8 + attention_freq=2 + window 32 + CAB + plain MLP + i-LN) beat or
+tied all of them. The architecture is at a measured local optimum.
+
+### Deferred (not forgotten — see the 2.0 backlog)
+
+📌 INT8 quantization, TensorRT export CI, RIB positional caching in eval, RIB
+dtype sweep, ENAF-style early exits for 4K+ images, S/M/L family, larger
+training crops (bigger GPU), video adaptation, rank 2/4 micro-sweep,
+attention-freq bs4 batch check, MLP ratio sweep.
+
+## HEART 2.0 backlog
+
+Frozen for 1.0 (do not touch the architecture while the release model trains).
+Ideas that need a fully trained 1.0 baseline before they can be judged
+properly — the rule is: only adopt what wins on *quality per second* against
+the 1.0 model, not what moves a benchmark number.
+
+| Idea | Why it's on the list | Effort |
+|---|---|---|
+| Gated-Dconv FFN (GDFN, Restormer) | MEASURED (20k, ratio 2.0, freq=2): lost 0.37 dB (24.25 vs 24.63) — doubling per-block 3x3 conv density (CAB already provides spatial convs) hurt. Code removed. | — |
+| SwiGLU FFN + QK-Norm | External review (2026-08-20): modernize the one remaining vanilla component (the MLP) + per-head QK LayerNorm for logit stability. TESTING: three 20k runs on the attn2 base (swiglu-only, qknorm-only, both). SwiGLU param-matched at ratio 4/3 (their 8/3 would have doubled params). | Low-Med |
+| RepCAB (reparameterized CAB) | Train-time multi-branch conv fused at export. REJECTED for 1.0: the fuse()-before-export step conflicts with the trivial-export property that is part of the maintainability goal. | — |
+| Transposed channel attention (MDTA) | Global channel-covariance attention — chases global context, which the w48 result already falsified as the lever. Backlog only if GDFN-style spatial enrichment proves the direction. | Medium |
+| Multi-axis stripe attention (HMA) | Full-row/column context — same global-context rationale, contradicted by w48 losing at 20k. Skip unless a crop-size scaling study revives the context hypothesis. | Medium |
+| Dense residual connections (DRCT) | Zero-param dense skips; NTIRE-2024 proven, but our blocks/groups already have residuals, and dense skips grow activation memory. Low expected value; cheap to test one day. | Low |
+| ENAF early-exit routing | Inference-time wrapper for huge images, not an architecture change; post-release option only. | High |
+| IET adaptive token selection | Dynamic gather/indexing risks breaking dynamic ONNX tracing — rejected on deployability grounds. | — |
+| Mamba/SSM backbones | Custom CUDA kernels (selective_scan) violate the zero-custom-kernel constraint — rejected. | — |
+| Larger windows (48 / 64) | MEASURED (20k, lq96): window 48 lost to 32 (-0.07 dB) and cost 38% more time. Closed for this crop/GPU regime unless crops scale too. | — |
+| Attention frequency (H C H C) | MEASURED (20k, lq96): `attention_freq=2` beat the baseline (+0.12 dB) at 1.56x speed. Live candidate for the release config; longer-run confirmation pending. | Low-Med |
+| RoPE instead of RIB | MEASURED (20k, lq96): RoPE lost 0.15 dB to RIB (24.36 vs 24.51) — learned low-rank positional wins; fixed rotations can't substitute. Closed. | — |
+| Rank beyond release range (4/12/24/…) | The release gate only tests 8/18/32. Fine-grained rank scaling could shave the positional capacity further. | Low |
+| RIB positional caching in eval | q_pos/k_pos recomputed every forward in every block; tiny but free to cache for inference. Micro-opt, measure before doing. | Low |
+| RIB dtype sweep (bf16/fp16 vs fp32) | fp32 cast is for fp16-inference stability; whether bf16 RIB changes quality/speed is unmeasured. | Low |
+| CAB variants / S/M/L family | An earlier hand-picked `heart_light` (8.7M) was removed for 1.0; a principled S/M/L scaling rule (width, depth, heads, rank, CAB) is the v2 way to extend the family. | Low-Med |
+| TensorRT + export CI | ONNX verified; TensorRT untested. A CI export test (PyTorch→ONNX→ORT→TensorRT, compare max-abs-error/PSNR/runtime) would harden the deployment claim. | Medium |
+| Quantization (INT8) | Untested; a 3060-class GPU model could get much faster. | Medium |
+| Larger training crops | lq 128+ needs checkpointing + slower it/s; only test with a bigger GPU. | Medium |
+| Video adaptation | temporal stability (like TSPAN) — separate scope. | High |
+| Cyclic vs reflect shift | Deliberately closed for 1.0 (cyclic + mask kills the flash kernel). Revisit only if a flash-compatible exact-cyclic scheme emerges. | — |
+
+Priority order after 1.0: confirm attention_freq=2 on a longer run (if the
+release adopts it, this is moot) -> export hardening (TensorRT/CI) ->
+RoPE swap-in benchmark -> rank/capacity scaling. Everything is measured
+against the frozen 1.0 model, not against papers.

LICENSE

@@ -0,0 +1,202 @@
+
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

README.md

@@ -0,0 +1,86 @@
+---
+license: apache-2.0
+language:
+- en
+tags:
+- super-resolution
+- image-restoration
+- pytorch
+- onnx
+- spandrel
+- chaiNNer
+pipeline_tag: image-to-image
+base_model:
+- Phips/heart-sisr-pretrain
+---
+
+# HEART — Efficient Attention with Rank-factorized bias Transformer
+
+Single-image super-resolution (SISR) architecture for traiNNer-redux and
+spandrel. Simple to train, fast to run, and ships as a dynamic ONNX model that
+works in ONNX Runtime, TensorRT, and chaiNNer.
+
+## Highlights
+
+- **16.7M params, 4x upscale**, window attention with Rank-factorized
+  Implicit Bias (RIB) — FlashAttention/SDPA friendly, no attention masks.
+- **Pure PyTorch, single file** — drop `heart_arch.py` into any traiNNer-redux
+  clone and train.
+- **Stable in reduced precision** — i-LN + bf16 AMP validated, no divergence.
+- **Fast** — beats HAT-M in speed and VRAM (see HEART.md in the source repo).
+- **Dynamic ONNX** — one export, any input size, runs anywhere.
+
+## Quick start (train)
+
+1. Copy `heart_arch.py` into `traiNNer/archs/`.
+2. Use one of the example configs in `configs/` (start with
+   `4x_HEART_release.yml`).
+3. Train:
+   ```bash
+   python train.py -opt options/train/RELEASE/4x_HEART_release.yml --auto_resume
+   ```
+
+> **Important**: keep the training crop (`lq_size`) a multiple of the 32 px
+> attention window — `96` (= 3×32) is recommended and used for all released
+> models. A non-multiple size (e.g. 80) teaches the model a false "edges are
+> mirrored" response that shows up as a bright bottom band / tile grid,
+> especially with a GAN loss.
+
+## Quick start (infer)
+
+Load the `.safetensors` in chaiNNer via spandrel, or run the `.onnx` with any
+ONNX Runtime / TensorRT engine. Inputs are RGB, dynamic H×W.
+
+## Models
+
+| File | Role |
+|---|---|
+| `models/heart_4x_pretrain.safetensors` | 4x pretrain (DF2K-style) |
+| `models/heart_2x_finetune.safetensors` | 2x finetune |
+| `models/heart_4x_otf_v2.safetensors` | 4x OTF fidelity (real-world, clean) |
+| `models/heart_4x_otf_gan.safetensors` | 4x OTF GAN (real-world, sharp) |
+| `onnx/heart_4x_otf_gan_1x3xHxW_fp32_op17.onnx` | dynamic ONNX, onnxslim-optimized |
+
+## Contents
+
+- `heart_arch.py` — the architecture (spandrel-registered).
+- `configs/` — example training configs (release, OTF, OTF GAN).
+- `models/` — pretrained checkpoints.
+- `onnx/` — dynamic fp32 ONNX (opset 17), onnxslim-slimmed.
+- `HEART.md` — full architecture docs, ablations, recipes.
+- `export_heart.py` — export a checkpoint to dynamic ONNX.
+
+## 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. Full design rationale
+and ablation data are in `HEART.md`.

configs/4x_HEART_OTF_gan.yml

@@ -0,0 +1,165 @@
+# yaml-language-server: $schema=https://raw.githubusercontent.com/the-database/traiNNer-redux/refs/heads/master/schemas/redux-config.schema.json
+#########################################################################################
+# HEART 4x OTF GAN finetune v2 (lq 96) — removes learned edge boundary response
+#
+# Root-cause fix for the bottom/right edge brightening artifact seen in the
+# lq-80 GAN run: an lq crop of 80px is NOT a multiple of the 32px attention
+# window, so every batch the model reflect-pads its own bottom/right edge and
+# learns "edges mirror content", which the GAN loss amplifies into a visible
+# band + tile grid. Training at lq_size 96 (= 3 x 32) keeps every window on
+# real content, so the boundary response is never learned.
+#
+# Starts from the finished lq-80 GAN weights (150k). Degradation + losses
+# unchanged from 4x_HEART_OTF_gan.yml. ~50-100k iters is enough.
+#########################################################################################
+name: 4x_HEART_OTF_gan_lq96
+scale: 4
+use_amp: true
+amp_bf16: true
+use_channels_last: true
+fast_matmul: false
+use_compile: false  # inductor errors on discriminator step — keep default
+compile_mode: default
+num_gpu: auto
+manual_seed: 1024
+
+high_order_degradation: true
+high_order_degradations_debug: false
+high_order_degradations_debug_limit: 100
+
+lq_usm: true
+lq_usm_radius_range: [1, 25]
+
+# Degradation matches the v2 fidelity pretrain (mild noise)
+blur_prob: 0.6
+resize_prob: [0.2, 0.7, 0.1]
+resize_mode_list: ['bilinear', 'bicubic', 'nearest-exact', 'lanczos']
+resize_mode_prob: [0.25, 0.25, 0.25, 0.25]
+resize_range: [0.4, 1.5]
+gaussian_noise_prob: 0.4
+noise_range: [1, 15]
+poisson_scale_range: [0.05, 2]
+gray_noise_prob: 0.1
+jpeg_prob: 0.4
+jpeg_range: [40, 95]
+
+blur_prob2: 0.3
+resize_prob2: [0.3, 0.4, 0.3]
+resize_mode_list2: ['bilinear', 'bicubic', 'nearest-exact', 'lanczos']
+resize_mode_prob2: [0.25, 0.25, 0.25, 0.25]
+resize_range2: [0.6, 1.2]
+gaussian_noise_prob2: 0.3
+noise_range2: [1, 10]
+poisson_scale_range2: [0.05, 1.5]
+gray_noise_prob2: 0.05
+jpeg_prob2: 0.3
+jpeg_range2: [40, 95]
+
+resize_mode_list3: ['bilinear', 'bicubic', 'nearest-exact', 'lanczos']
+resize_mode_prob3: [0.25, 0.25, 0.25, 0.25]
+
+queue_size: 120
+
+datasets:
+  train:
+    name: LUCID OTF Train
+    type: realesrgandataset
+    dataroot_gt: [datasets/train/lucid_hc512/hr]
+
+    blur_kernel_size: 12
+    kernel_list: ['iso', 'aniso', 'generalized_iso', 'generalized_aniso', 'plateau_iso', 'plateau_aniso']
+    kernel_prob: [0.45, 0.25, 0.12, 0.03, 0.12, 0.03]
+    kernel_range: [5, 17]
+    sinc_prob: 0.0
+    blur_sigma: [0.2, 2]
+    betag_range: [0.5, 4]
+    betap_range: [1, 2]
+
+    blur_kernel_size2: 12
+    kernel_list2: ['iso', 'aniso', 'generalized_iso', 'generalized_aniso', 'plateau_iso', 'plateau_aniso']
+    kernel_prob2: [0.45, 0.25, 0.12, 0.03, 0.12, 0.03]
+    kernel_range2: [5, 17]
+    sinc_prob2: 0.0
+    blur_sigma2: [0.2, 1]
+    betag_range2: [0.5, 4]
+    betap_range2: [1, 2]
+
+    final_sinc_prob: 0.0
+    final_kernel_range: [5, 17]
+
+    lq_size: 96          # = 3 x 32px window — window-aligned, no learned edge response
+    use_hflip: true
+    use_rot: true
+    num_worker_per_gpu: 8
+    batch_size_per_gpu: 2
+    accum_iter: 1
+
+  val:
+    name: Urban100 OTF Val
+    type: pairedimagedataset
+    dataroot_gt: [datasets/val/urban100/hr]
+    dataroot_lq: [datasets/val/urban100_otf_x4]
+
+network_g:
+  type: heart
+  use_checkpoint: false
+
+network_d:
+  type: dunet
+
+path:
+  pretrain_network_g: experiments/4x_HEART_OTF_gan/models/net_g_ema_latest.safetensors
+  param_key_g: ~
+  strict_load_g: true
+  resume_state: ~
+
+train:
+  ema_decay: 0.999
+  ema_power: 0.75
+  grad_clip: false
+  grad_clip_max_norm: 100
+  optim_g:
+    type: AdamW
+    lr: !!float 1e-4
+    weight_decay: 0
+    betas: [0.9, 0.99]
+  optim_d:
+    type: AdamW
+    lr: !!float 1e-4
+    weight_decay: 0
+    betas: [0.9, 0.99]
+  scheduler:
+    type: MultiStepLR
+    milestones: [75000, 120000]
+    gamma: 0.5
+  total_iter: 150000
+  warmup_iter: -1
+  losses:
+    - type: mssimloss
+      loss_weight: 0.5
+    - type: perceptualloss
+      criterion: charbonnier
+      loss_weight: 0.01
+    - type: hsluvloss
+      criterion: charbonnier
+      loss_weight: 1.0
+    - type: ganloss
+      gan_type: vanilla
+      loss_weight: 0.1
+
+val:
+  val_enabled: true
+  val_freq: 5000
+  save_img: true
+  tile_size: 0
+  tile_overlap: 8
+  metrics_enabled: true
+  metrics:
+    topiq:
+      type: calculate_topiq
+
+logger:
+  print_freq: 100
+  save_checkpoint_freq: 5000
+  save_checkpoint_format: safetensors
+  use_tb_logger: true

configs/4x_HEART_OTF_v2.yml

@@ -0,0 +1,155 @@
+# yaml-language-server: $schema=https://raw.githubusercontent.com/the-database/traiNNer-redux/refs/heads/master/schemas/redux-config.schema.json
+#########################################################################################
+# HEART 4x OTF v2 (lq 96) — window-aligned fidelity pretrain for GAN finetunes
+#
+# Replaces the v1.1 lq-80 run as the canonical OTF fidelity pretrain. Root-cause
+# fix: lq_size 96 (= 3 x 32px window) keeps every attention window on real
+# content during training. The lq-80 crop forced the model to reflect-pad its
+# own bottom/right edge every batch, teaching a "mirror at edge" response that
+# shows up as a bright band / tile grid (and that the GAN loss amplifies).
+#
+# Starts from the finished v1.1 weights (backup_otf_v2, 150k). Same degradation
+# + losses as v1.1, only the crop size changes.
+#########################################################################################
+name: 4x_HEART_OTF_v2
+scale: 4
+use_amp: true
+amp_bf16: true
+use_channels_last: true
+fast_matmul: false
+use_compile: true
+compile_mode: reduce-overhead
+num_gpu: auto
+manual_seed: 1024
+
+high_order_degradation: true
+high_order_degradations_debug: false
+high_order_degradations_debug_limit: 100
+
+lq_usm: true
+lq_usm_radius_range: [1, 25]
+
+# Degradation identical to v1.1
+blur_prob: 0.6
+resize_prob: [0.2, 0.7, 0.1]
+resize_mode_list: ['bilinear', 'bicubic', 'nearest-exact', 'lanczos']
+resize_mode_prob: [0.25, 0.25, 0.25, 0.25]
+resize_range: [0.4, 1.5]
+gaussian_noise_prob: 0.4
+noise_range: [1, 15]
+poisson_scale_range: [0.05, 2]
+gray_noise_prob: 0.1
+jpeg_prob: 0.4
+jpeg_range: [40, 95]
+
+blur_prob2: 0.3
+resize_prob2: [0.3, 0.4, 0.3]
+resize_mode_list2: ['bilinear', 'bicubic', 'nearest-exact', 'lanczos']
+resize_mode_prob2: [0.25, 0.25, 0.25, 0.25]
+resize_range2: [0.6, 1.2]
+gaussian_noise_prob2: 0.3
+noise_range2: [1, 10]
+poisson_scale_range2: [0.05, 1.5]
+gray_noise_prob2: 0.05
+jpeg_prob2: 0.3
+jpeg_range2: [40, 95]
+
+resize_mode_list3: ['bilinear', 'bicubic', 'nearest-exact', 'lanczos']
+resize_mode_prob3: [0.25, 0.25, 0.25, 0.25]
+
+queue_size: 120
+
+datasets:
+  train:
+    name: LUCID OTF Train
+    type: realesrgandataset
+    dataroot_gt: [datasets/train/lucid_hc512/hr]
+
+    blur_kernel_size: 12
+    kernel_list: ['iso', 'aniso', 'generalized_iso', 'generalized_aniso', 'plateau_iso', 'plateau_aniso']
+    kernel_prob: [0.45, 0.25, 0.12, 0.03, 0.12, 0.03]
+    kernel_range: [5, 17]
+    blur_sigma: [0.2, 2]
+    betag_range: [0.5, 4]
+    betap_range: [1, 2]
+
+    blur_kernel_size2: 12
+    kernel_list2: ['iso', 'aniso', 'generalized_iso', 'generalized_aniso', 'plateau_iso', 'plateau_aniso']
+    kernel_prob2: [0.45, 0.25, 0.12, 0.03, 0.12, 0.03]
+    kernel_range2: [5, 17]
+    blur_sigma2: [0.2, 1]
+    betag_range2: [0.5, 4]
+    betap_range2: [1, 2]
+
+    final_kernel_range: [5, 17]
+
+    lq_size: 96          # = 3 x 32px window — window-aligned, no learned edge response
+    use_hflip: true
+    use_rot: true
+    num_worker_per_gpu: 8
+    batch_size_per_gpu: 2
+    accum_iter: 1
+
+  val:
+    name: Urban100 OTF Val
+    type: pairedimagedataset
+    dataroot_gt: [datasets/val/urban100/hr]
+    dataroot_lq: [datasets/val/urban100_otf_x4]
+
+network_g:
+  type: heart
+  use_checkpoint: false
+
+path:
+  pretrain_network_g: experiments/backup_otf_v2/net_g_ema_150000.safetensors
+  param_key_g: ~
+  strict_load_g: true
+  resume_state: ~
+
+train:
+  ema_decay: 0.999
+  ema_power: 0.75
+  grad_clip: false
+  grad_clip_max_norm: 100
+  optim_g:
+    type: AdamW
+    lr: !!float 5e-5
+    weight_decay: 0
+    betas: [0.9, 0.99]
+  scheduler:
+    type: MultiStepLR
+    milestones: [50000, 100000]
+    gamma: 0.5
+  total_iter: 150000
+  warmup_iter: -1
+  losses:
+    - type: charbonnierloss
+      loss_weight: 1.0
+    - type: perceptualloss
+      criterion: charbonnier
+      loss_weight: 0.01
+    - type: mssimloss
+      loss_weight: 0.5
+
+val:
+  val_enabled: true
+  val_freq: 5000
+  save_img: true
+  tile_size: 0
+  tile_overlap: 8
+  metrics_enabled: true
+  metrics:
+    psnr:
+      type: calculate_psnr
+      crop_border: 4
+      test_y_channel: true
+    ssim:
+      type: calculate_ssim
+      crop_border: 4
+      test_y_channel: true
+
+logger:
+  print_freq: 100
+  save_checkpoint_freq: 5000
+  save_checkpoint_format: safetensors
+  use_tb_logger: true

configs/4x_HEART_release.yml

@@ -0,0 +1,91 @@
+# yaml-language-server: $schema=https://raw.githubusercontent.com/the-database/traiNNer-redux/refs/heads/master/schemas/redux-config.schema.json
+#########################################################################################
+# HEART 4x release training
+# LUCID CC0 v2 HC-512 (100,866 tiles, 512x512, CC0) - license-clean pretrain.
+# lq 96 / batch 2: 9 windows per crop so the shifted-window mixing actually works.
+# torch.compile reduce-overhead: +46% training speed, quality-neutral (validated).
+# rank is set by the release driver after the rank ablation picks a winner.
+#########################################################################################
+name: 4x_HEART
+scale: 4
+use_amp: true
+amp_bf16: true
+use_channels_last: true
+fast_matmul: false
+use_compile: true
+compile_mode: reduce-overhead
+num_gpu: auto
+manual_seed: 1024
+
+datasets:
+  train:
+    name: LUCID Train
+    type: pairedimagedataset
+    dataroot_gt: [datasets/train/lucid_hc512/hr]
+    dataroot_lq: [datasets/train/lucid_hc512/lr_x4]
+    lq_size: 96
+    use_hflip: true
+    use_rot: true
+    num_worker_per_gpu: 8
+    batch_size_per_gpu: 2
+    accum_iter: 1
+  val:
+    name: Urban100 Val
+    type: pairedimagedataset
+    dataroot_gt: [datasets/val/urban100/hr]
+    dataroot_lq: [datasets/val/urban100/lr_x4]
+    filename_tmpl: "{}x4"
+
+network_g:
+  type: heart
+  use_checkpoint: false
+  rank: 8
+  attention_freq: 2
+
+path:
+  param_key_g: ~
+  strict_load_g: true
+  resume_state: ~
+
+train:
+  ema_decay: 0.999
+  ema_power: 0.75
+  grad_clip: false
+  grad_clip_max_norm: 100
+  optim_g:
+    type: AdamW
+    lr: !!float 2e-4
+    weight_decay: 0
+    betas: [0.9, 0.99]
+  scheduler:
+    type: MultiStepLR
+    milestones: [300000, 500000, 650000, 700000, 750000]
+    gamma: 0.5
+  total_iter: 800000
+  warmup_iter: -1
+  losses:
+    - type: charbonnierloss
+      loss_weight: 1.0
+
+val:
+  val_enabled: true
+  val_freq: 5000
+  save_img: true
+  tile_size: 0
+  tile_overlap: 8
+  metrics_enabled: true
+  metrics:
+    psnr:
+      type: calculate_psnr
+      crop_border: 4
+      test_y_channel: true
+    ssim:
+      type: calculate_ssim
+      crop_border: 4
+      test_y_channel: true
+
+logger:
+  print_freq: 100
+  save_checkpoint_freq: 5000
+  save_checkpoint_format: safetensors
+  use_tb_logger: true

heart_arch.py

@@ -0,0 +1,645 @@
+# ruff: noqa
+# type: ignore
+# HEART - Hybrid Efficient Attention with Rank-factorized bias Transformer
+# based on HAT-iLN (arXiv:2504.06629) and SST's RIB (arXiv:2603.06738)
+from __future__ import annotations
+
+from collections.abc import Sequence
+from typing import Literal
+
+import torch
+from einops import rearrange
+from spandrel.util import store_hyperparameters
+from spandrel.util.timm import trunc_normal_
+from torch import nn
+from torch.nn import functional as F  # noqa: N812
+from torch.utils import checkpoint
+
+from traiNNer.archs.arch_util import iLN
+from traiNNer.archs.hat_iln_arch import (
+    CAB,
+    AffineTransform,
+    DropPath,
+    Mlp,
+    PatchEmbed,
+    PatchUnEmbed,
+    Upsample,
+)
+from traiNNer.utils.registry import ARCH_REGISTRY, SPANDREL_REGISTRY
+
+
+class RIBWindowAttention(nn.Module):
+    """Window attention with Rank-factorized Implicit Neural Bias (RIB).
+
+    RIB (SST, arXiv 2603.06738) replaces the relative position bias table:
+    position features from a small implicit net are concatenated onto Q/K,
+    turning bias addition into a dot product, so attention runs on
+    FlashAttention/SDPA kernels with no bias table and no masks.
+    Channels-first (B, C, H, W); shifted windows use non-wrapping
+    pad-and-partition (no cyclic wrap).
+    """
+
+    def __init__(
+        self,
+        dim: int,
+        window_size: int,
+        num_heads: int,
+        rank: int = 8,
+        rib_hidden_dim: int = 32,
+        rib_n_freqs: int = 10,
+        shift: bool = False,
+    ) -> None:
+        super().__init__()
+        self.window_size = (window_size, window_size)
+        self.num_heads = num_heads
+        assert dim % num_heads == 0, "dim must be divisible by num_heads"
+
+        self.to_qkv = nn.Conv2d(dim, dim * 3, 1, 1, 0)
+        self.to_out = nn.Conv2d(dim, dim, 1, 1, 0)
+
+        self.rank = rank
+
+        # normalized window coordinates
+        wh, ww = self.window_size
+        yy, xx = torch.meshgrid(torch.arange(wh), torch.arange(ww), indexing="ij")
+        coords = torch.stack([xx, yy], dim=-1).reshape(-1, 2).float()  # (N, 2)
+        coords[:, 0] = (2.0 * (coords[:, 0] + 0.5) / ww) - 1.0
+        coords[:, 1] = (2.0 * (coords[:, 1] + 0.5) / wh) - 1.0
+
+        self.n_freqs = rib_n_freqs
+        if self.n_freqs > 0:
+            base_coords = coords.clone()
+            for i in range(self.n_freqs):
+                freq = 2**i
+                coords = torch.cat(
+                    [
+                        coords,
+                        torch.sin(base_coords * freq),
+                        torch.cos(base_coords * freq),
+                    ],
+                    dim=-1,
+                )
+        self.register_buffer("rib_coords", coords, persistent=False)
+
+        n_input = 2 + 4 * self.n_freqs
+        hidden_d = rib_hidden_dim
+        self.to_hidden = nn.Parameter(torch.empty(n_input, hidden_d))
+        self.hidden_b = nn.Parameter(torch.zeros(1, hidden_d))
+        self.to_q = nn.Parameter(torch.empty(num_heads, hidden_d, self.rank))
+        self.to_k = nn.Parameter(torch.empty(num_heads, hidden_d, self.rank))
+
+        nn.init.normal_(self.to_hidden, mean=0.0, std=0.05)
+        nn.init.normal_(self.to_q, mean=0.0, std=0.05)
+        nn.init.normal_(self.to_k, mean=0.0, std=0.05)
+
+        self.shift = shift
+
+    def pad_to_win(self, x: torch.Tensor, h: int, w: int) -> torch.Tensor:
+        # inputs are pre-aligned to window multiples by HEART.check_image_size
+        # (which owns the replicate fallback), so reflect padding is safe here
+        pad_h = (self.window_size[0] - h % self.window_size[0]) % self.window_size[0]
+        pad_w = (self.window_size[1] - w % self.window_size[1]) % self.window_size[1]
+        return F.pad(x, (0, pad_w, 0, pad_h), mode="reflect")
+
+    def _flash_cat_attn(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        pos_q: torch.Tensor,
+        pos_k: torch.Tensor,
+    ) -> torch.Tensor:
+        # q, k, v: (Bwin, N, heads, D); pos_q, pos_k: (Bwin, heads, N, R)
+        d = q.shape[-1]
+        pos_q = pos_q.transpose(1, 2)  # (Bwin, N, heads, R)
+        pos_k = pos_k.transpose(1, 2)
+        r = pos_q.shape[-1]
+
+        q = q * (d**-0.5)
+        pos_q = pos_q * (r**-0.5)
+        q_cat = torch.cat([q, pos_q], dim=-1)
+        k_cat = torch.cat([k, pos_k], dim=-1)
+
+        d_total = q_cat.shape[-1]
+        pad = (8 - (d_total % 8)) % 8
+        if pad:
+            q_cat = F.pad(q_cat, (0, pad))
+            k_cat = F.pad(k_cat, (0, pad))
+        v_cat = F.pad(v, (0, r + pad))
+
+        # q and pos_q are pre-scaled (SST eq. 5); scale=1.0 keeps SDPA from
+        # rescaling them again
+        out = F.scaled_dot_product_attention(
+            q_cat.transpose(1, 2),
+            k_cat.transpose(1, 2),
+            v_cat.transpose(1, 2),
+            scale=1.0,
+        ).transpose(1, 2)[:, :, :, :d]
+        return out
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # x: (B, C, H, W)
+        _, _, h, w = x.shape
+        x = self.pad_to_win(x, h, w)
+
+        # shifted windows: pad top/left by half a window, partition from the
+        # padded origin. x is window-aligned here, so padding (sh, sh) per
+        # axis keeps alignment (one extra window row/col).
+        sh = sw = 0
+        if self.shift:
+            sh, sw = self.window_size[0] // 2, self.window_size[1] // 2
+            x = F.pad(x, (sw, sw, sh, sh), mode="reflect")
+
+        h_div, w_div = (
+            x.shape[2] // self.window_size[0],
+            x.shape[3] // self.window_size[1],
+        )
+
+        qkv = self.to_qkv(x)
+        qkv = rearrange(
+            qkv,
+            "b (qkv heads c) (h wh) (w ww) -> qkv (b h w) (wh ww) heads c",
+            heads=self.num_heads,
+            wh=self.window_size[0],
+            ww=self.window_size[1],
+            qkv=3,
+        )
+        q, k, v = qkv[0], qkv[1], qkv[2]  # (Bwin, N, heads, D)
+
+        # RIB position features (fp32 for fp16-inference stability)
+        intermediate = F.relu(
+            self.rib_coords.to(dtype=torch.float32)
+            @ self.to_hidden.to(dtype=torch.float32)
+            + self.hidden_b.to(dtype=torch.float32)
+        )  # (N, hidden_d)
+        q_pos = torch.einsum(
+            "nd,hdr->hnr", intermediate, self.to_q.to(dtype=torch.float32)
+        )  # (heads, N, R)
+        k_pos = torch.einsum(
+            "nd,hdr->hnr", intermediate, self.to_k.to(dtype=torch.float32)
+        )
+        bwin = q.shape[0]
+        q_pos = (
+            q_pos.to(q.dtype).unsqueeze(0).expand(bwin, -1, -1, -1)
+        )  # (Bwin, heads, N, R)
+        k_pos = k_pos.to(k.dtype).unsqueeze(0).expand(bwin, -1, -1, -1)
+
+        out = self._flash_cat_attn(q, k, v, q_pos, k_pos)  # (Bwin, N, heads, D)
+
+        out = rearrange(
+            out,
+            "(b h w) (wh ww) heads c -> b (heads c) (h wh) (w ww)",
+            h=h_div,
+            w=w_div,
+            wh=self.window_size[0],
+            ww=self.window_size[1],
+        )
+        out = out[:, :, sh : sh + h, sw : sw + w]
+        return self.to_out(out)
+
+
+class HAB_RIB(nn.Module):
+    """Hybrid Attention Block: RIB window attention + CAB conv branch + i-LN.
+
+    The block unpacks ``(x, std)`` from norm1/norm2, so norm_layer must be
+    i-LN-compatible; the parameter exists for HAT symmetry, not as a knob.
+    """
+
+    def __init__(
+        self,
+        dim: int,
+        num_heads: int,
+        window_size: int = 32,
+        compress_ratio: int = 3,
+        squeeze_factor: int = 30,
+        conv_scale: float = 0.01,
+        mlp_ratio: float = 2.0,
+        drop_path: float = 0.0,
+        norm_layer: type[nn.Module] = iLN,
+        rank: int = 8,
+        rib_hidden_dim: int = 32,
+        rib_n_freqs: int = 10,
+        shift: bool = False,
+        use_attention: bool = True,
+    ) -> None:
+        super().__init__()
+        self.conv_scale = conv_scale
+        self.use_attention = use_attention
+
+        self.norm1 = norm_layer(dim)
+        if use_attention:
+            self.attn = RIBWindowAttention(
+                dim=dim,
+                window_size=window_size,
+                num_heads=num_heads,
+                rank=rank,
+                rib_hidden_dim=rib_hidden_dim,
+                rib_n_freqs=rib_n_freqs,
+                shift=shift,
+            )
+
+        self.conv_block = CAB(
+            num_feat=dim, compress_ratio=compress_ratio, squeeze_factor=squeeze_factor
+        )
+
+        self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
+        self.norm2 = norm_layer(dim)
+        mlp_hidden_dim = int(dim * mlp_ratio)
+        self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim)
+
+    def forward(self, x: torch.Tensor, x_size: tuple[int, int]) -> torch.Tensor:
+        h, w = x_size
+        b, _, c = x.shape
+
+        shortcut = x
+        x, std1 = self.norm1(x)
+        x_2d = x.transpose(1, 2).contiguous().view(b, c, h, w)
+
+        # conv branch
+        conv_x = self.conv_block(x_2d).view(b, c, -1).transpose(1, 2)
+        residual = conv_x * self.conv_scale
+
+        # RIB window attention (channels-first), optional per attention_freq
+        if self.use_attention:
+            attn_out = self.attn(x_2d)
+            attn_x = attn_out.view(b, c, -1).transpose(1, 2)
+            residual = residual + self.drop_path(attn_x)
+
+        # i-LN: rescale by std from norm1
+        x = shortcut + std1.view(b, 1, 1) * residual
+
+        # FFN with i-LN rescaling
+        x_normed, std2 = self.norm2(x)
+        ffn_out = self.mlp(x_normed)
+        x = x + std2.view(b, 1, 1) * self.drop_path(ffn_out)
+
+        return x
+
+
+class AttenBlocksRIB(nn.Module):
+    """A stack of HAB_RIB blocks with optional gradient checkpointing."""
+
+    def __init__(
+        self,
+        dim: int,
+        depth: int,
+        num_heads: int,
+        window_size: int,
+        compress_ratio: int,
+        squeeze_factor: int,
+        conv_scale: float,
+        mlp_ratio: float,
+        drop_path: float | list[float],
+        norm_layer: type[nn.Module] = iLN,
+        use_checkpoint: bool = False,
+        rank: int = 8,
+        rib_hidden_dim: int = 32,
+        rib_n_freqs: int = 10,
+        attention_freq: int = 2,
+    ) -> None:
+        super().__init__()
+        self.use_checkpoint = use_checkpoint
+        blocks: list[HAB_RIB] = []
+        attn_idx = 0
+        for i in range(depth):
+            use_attention = i % attention_freq == 0
+            blocks.append(
+                HAB_RIB(
+                    dim=dim,
+                    num_heads=num_heads,
+                    window_size=window_size,
+                    compress_ratio=compress_ratio,
+                    squeeze_factor=squeeze_factor,
+                    conv_scale=conv_scale,
+                    mlp_ratio=mlp_ratio,
+                    drop_path=drop_path[i]
+                    if isinstance(drop_path, list)
+                    else drop_path,
+                    norm_layer=norm_layer,
+                    rank=rank,
+                    rib_hidden_dim=rib_hidden_dim,
+                    rib_n_freqs=rib_n_freqs,
+                    shift=(attn_idx % 2 == 1) if use_attention else False,
+                    use_attention=use_attention,
+                )
+            )
+            if use_attention:
+                attn_idx += 1
+        self.blocks = nn.ModuleList(blocks)
+
+    def forward(self, x: torch.Tensor, x_size: tuple[int, int]) -> torch.Tensor:
+        use_chk = self.use_checkpoint and self.training
+        for blk in self.blocks:
+            if use_chk:
+                x = checkpoint.checkpoint(blk, x, x_size, use_reentrant=False)
+            else:
+                x = blk(x, x_size)
+        return x
+
+
+class RHAG_RIB(nn.Module):
+    """Residual Hybrid Attention Group: AttenBlocksRIB -> 3x3 conv -> residual."""
+
+    def __init__(
+        self,
+        dim: int,
+        depth: int,
+        num_heads: int,
+        window_size: int,
+        compress_ratio: int,
+        squeeze_factor: int,
+        conv_scale: float,
+        mlp_ratio: float,
+        drop_path: float | list[float],
+        norm_layer: type[nn.Module] = iLN,
+        use_checkpoint: bool = False,
+        img_size: int = 64,
+        patch_size: int = 1,
+        resi_connection: str = "1conv",
+        rank: int = 8,
+        rib_hidden_dim: int = 32,
+        rib_n_freqs: int = 10,
+        attention_freq: int = 2,
+    ) -> None:
+        super().__init__()
+        self.dim = dim
+
+        self.residual_group = AttenBlocksRIB(
+            dim=dim,
+            depth=depth,
+            num_heads=num_heads,
+            window_size=window_size,
+            compress_ratio=compress_ratio,
+            squeeze_factor=squeeze_factor,
+            conv_scale=conv_scale,
+            mlp_ratio=mlp_ratio,
+            drop_path=drop_path,
+            norm_layer=norm_layer,
+            use_checkpoint=use_checkpoint,
+            rank=rank,
+            rib_hidden_dim=rib_hidden_dim,
+            rib_n_freqs=rib_n_freqs,
+            attention_freq=attention_freq,
+        )
+
+        if resi_connection == "1conv":
+            self.conv = nn.Conv2d(dim, dim, 3, 1, 1)
+        elif resi_connection == "identity":
+            self.conv = nn.Identity()
+        else:
+            raise ValueError(f"unsupported resi_connection: {resi_connection}")
+
+        self.patch_embed = PatchEmbed(
+            img_size=img_size,
+            patch_size=patch_size,
+            in_chans=0,
+            embed_dim=dim,
+            norm_layer=None,
+        )
+        self.patch_unembed = PatchUnEmbed(
+            img_size=img_size,
+            patch_size=patch_size,
+            in_chans=0,
+            embed_dim=dim,
+            norm_layer=None,
+        )
+
+    def forward(self, x: torch.Tensor, x_size: tuple[int, int]) -> torch.Tensor:
+        return (
+            self.patch_embed(
+                self.conv(self.patch_unembed(self.residual_group(x, x_size), x_size))
+            )
+            + x
+        )
+
+
+@store_hyperparameters()
+class HEART(nn.Module):
+    """HEART - Hybrid Efficient Attention with Rank-factorized bias Transformer.
+
+    Simplified HAT-iLN: RIB window attention (FlashAttention-friendly) replaces
+    the relative position bias table and OCAB is removed. i-LN keeps training
+    stable in reduced precision. Pure PyTorch, ONNX-exportable.
+    """
+
+    hyperparameters = {}
+
+    def __init__(
+        self,
+        *,
+        img_size: int = 64,
+        patch_size: int = 1,
+        in_chans: int = 3,
+        embed_dim: int = 180,
+        depths: Sequence[int] = (6, 6, 6, 6, 6, 6),
+        num_heads: Sequence[int] = (6, 6, 6, 6, 6, 6),
+        window_size: int = 32,
+        compress_ratio: int = 3,
+        squeeze_factor: int = 30,
+        conv_scale: float = 0.01,
+        mlp_ratio: float = 2.0,
+        drop_rate: float = 0.0,
+        drop_path_rate: float = 0.1,
+        norm_layer: type[nn.Module] = iLN,
+        ape: bool = False,
+        patch_norm: bool = True,
+        use_checkpoint: bool = False,
+        upscale: int = 1,
+        upsampler: Literal["pixelshuffle"] = "pixelshuffle",
+        resi_connection: str = "1conv",
+        num_feat: int = 64,
+        rank: int = 8,
+        rib_hidden_dim: int = 32,
+        rib_n_freqs: int = 10,
+        attention_freq: int = 2,
+    ) -> None:
+        super().__init__()
+
+        self.window_size = window_size
+
+        num_in_ch = in_chans
+        num_out_ch = in_chans
+        self.upscale = upscale
+        self.upsampler = upsampler
+
+        # ------------------------- 1, shallow feature extraction ------------------------- #
+        self.conv_first = nn.Conv2d(num_in_ch, embed_dim, 3, 1, 1)
+
+        # ------------------------- 2, deep feature extraction ------------------------- #
+        self.num_layers = len(depths)
+        self.embed_dim = embed_dim
+        self.ape = ape
+        self.patch_norm = patch_norm
+        self.num_features = embed_dim
+
+        self.patch_embed = PatchEmbed(
+            img_size=img_size,
+            patch_size=patch_size,
+            in_chans=embed_dim,
+            embed_dim=embed_dim,
+            norm_layer=AffineTransform if self.patch_norm else None,
+        )
+        num_patches = self.patch_embed.num_patches
+        patches_resolution = self.patch_embed.patches_resolution
+        self.patches_resolution = patches_resolution
+
+        self.patch_unembed = PatchUnEmbed(
+            img_size=img_size,
+            patch_size=patch_size,
+            in_chans=embed_dim,
+            embed_dim=embed_dim,
+            norm_layer=AffineTransform if self.patch_norm else None,
+        )
+
+        if self.ape:
+            self.absolute_pos_embed = nn.Parameter(
+                torch.zeros(1, num_patches, embed_dim)
+            )
+            trunc_normal_(self.absolute_pos_embed, std=0.02)
+
+        self.pos_drop = nn.Dropout(p=drop_rate)
+
+        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))]
+
+        self.layers = nn.ModuleList()
+        for i_layer in range(self.num_layers):
+            layer = RHAG_RIB(
+                dim=embed_dim,
+                depth=depths[i_layer],
+                num_heads=num_heads[i_layer],
+                window_size=window_size,
+                compress_ratio=compress_ratio,
+                squeeze_factor=squeeze_factor,
+                conv_scale=conv_scale,
+                mlp_ratio=mlp_ratio,
+                drop_path=dpr[sum(depths[:i_layer]) : sum(depths[: i_layer + 1])],
+                norm_layer=norm_layer,
+                use_checkpoint=use_checkpoint,
+                img_size=img_size,
+                patch_size=patch_size,
+                resi_connection=resi_connection,
+                rank=rank,
+                rib_hidden_dim=rib_hidden_dim,
+                rib_n_freqs=rib_n_freqs,
+                attention_freq=attention_freq,
+            )
+            self.layers.append(layer)
+        self.norm = AffineTransform(self.num_features)
+
+        if resi_connection == "1conv":
+            self.conv_after_body = nn.Conv2d(embed_dim, embed_dim, 3, 1, 1)
+        elif resi_connection == "identity":
+            self.conv_after_body = nn.Identity()
+        else:
+            raise ValueError(f"unsupported resi_connection: {resi_connection}")
+
+        # ------------------------- 3, high quality image reconstruction ------------------------- #
+        if self.upsampler == "pixelshuffle":
+            self.conv_before_upsample = nn.Sequential(
+                nn.Conv2d(embed_dim, num_feat, 3, 1, 1), nn.LeakyReLU(inplace=True)
+            )
+            self.upsample = Upsample(upscale, num_feat)
+            self.conv_last = nn.Conv2d(num_feat, num_out_ch, 3, 1, 1)
+
+        self.apply(self._init_weights)
+
+    def _init_weights(self, m: nn.Module) -> None:
+        if isinstance(m, nn.Linear):
+            trunc_normal_(m.weight, std=0.02)
+            if m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        elif isinstance(m, (nn.LayerNorm, iLN)):
+            nn.init.constant_(m.bias, 0)
+            nn.init.constant_(m.weight, 1.0)
+
+    def check_image_size(self, x: torch.Tensor) -> torch.Tensor:
+        _, _, h, w = x.shape
+        pad_h = (self.window_size - h % self.window_size) % self.window_size
+        pad_w = (self.window_size - w % self.window_size) % self.window_size
+        can_reflect = (pad_h == 0 or pad_h < h) and (pad_w == 0 or pad_w < w)
+        mode = "reflect" if can_reflect else "replicate"
+        return F.pad(x, (0, pad_w, 0, pad_h), mode=mode)
+
+    def forward_features(self, x: torch.Tensor) -> torch.Tensor:
+        x_size = (x.shape[2], x.shape[3])
+
+        x = self.patch_embed(x)
+        if self.ape:
+            x = x + self.absolute_pos_embed
+        x = self.pos_drop(x)
+
+        for layer in self.layers:
+            x = layer(x, x_size)
+
+        x = self.norm(x)
+        x = self.patch_unembed(x, x_size)
+
+        return x
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        h, w = x.shape[2:]
+        x = self.check_image_size(x)
+
+        if self.upsampler == "pixelshuffle":
+            x = self.conv_first(x)
+            x = self.conv_after_body(self.forward_features(x)) + x
+            x = self.conv_before_upsample(x)
+            x = self.conv_last(self.upsample(x))
+
+        return x[:, :, : h * self.upscale, : w * self.upscale]
+
+
+@SPANDREL_REGISTRY.register()
+@ARCH_REGISTRY.register()
+def heart(
+    scale: int = 4,
+    img_size: int = 64,
+    patch_size: int = 1,
+    in_chans: int = 3,
+    embed_dim: int = 180,
+    depths: Sequence[int] = (6, 6, 6, 6, 6, 6),
+    num_heads: Sequence[int] = (6, 6, 6, 6, 6, 6),
+    window_size: int = 32,
+    compress_ratio: int = 3,
+    squeeze_factor: int = 30,
+    conv_scale: float = 0.01,
+    mlp_ratio: float = 2.0,
+    drop_rate: float = 0.0,
+    drop_path_rate: float = 0.1,
+    ape: bool = False,
+    patch_norm: bool = True,
+    use_checkpoint: bool = False,
+    upsampler: Literal["pixelshuffle"] = "pixelshuffle",
+    resi_connection: str = "1conv",
+    num_feat: int = 64,
+    rank: int = 8,
+    rib_hidden_dim: int = 32,
+    rib_n_freqs: int = 10,
+    attention_freq: int = 2,
+) -> HEART:
+    return HEART(
+        upscale=scale,
+        img_size=img_size,
+        patch_size=patch_size,
+        in_chans=in_chans,
+        embed_dim=embed_dim,
+        depths=depths,
+        num_heads=num_heads,
+        window_size=window_size,
+        compress_ratio=compress_ratio,
+        squeeze_factor=squeeze_factor,
+        conv_scale=conv_scale,
+        mlp_ratio=mlp_ratio,
+        drop_rate=drop_rate,
+        drop_path_rate=drop_path_rate,
+        ape=ape,
+        patch_norm=patch_norm,
+        use_checkpoint=use_checkpoint,
+        upsampler=upsampler,
+        resi_connection=resi_connection,
+        num_feat=num_feat,
+        rank=rank,
+        rib_hidden_dim=rib_hidden_dim,
+        rib_n_freqs=rib_n_freqs,
+        attention_freq=attention_freq,
+    )

models/.gitignore

@@ -0,0 +1 @@
+*.safetensors

onnx/heart_2x_fp32_op17.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:08ee34f60303ac086c9b50012fdc3d261086d907756640623035cec1b6030102
+size 71690367

onnx/heart_4x_otf_fp32_op17.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c2fd8a518b96d987dabaf7bbe284194fa26f14d2f972e34c8a4370c1d00df8ab
+size 72281693

onnx/heart_4x_otf_gan_fp32_op17.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:cf494433c0873c6893d324baff4d843afb2b96efcf54e4efca07a033a8308427
+size 72281667

onnx/heart_4x_pretrain_fp32_op17.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c3b58cce1b18ba65c602f8937df40da75dfe1dbd208d661f48eb85f8380c4369
+size 72281693

scripts/export_heart.py

@@ -0,0 +1,27 @@
+#!/usr/bin/env python3
+"""Minimal memory ONNX export. No validation, no onnxsim, tiny dummy shape.
+If this fits in RAM, great. Otherwise the export must run on another machine."""
+import gc, os, sys
+import torch
+from safetensors.torch import load_file
+from traiNNer.archs.heart_arch import heart
+
+ckpt, out, scale = sys.argv[1], sys.argv[2], int(sys.argv[3])
+torch.set_num_threads(2)
+torch.set_grad_enabled(False)
+
+sd = load_file(ckpt, device="cpu")
+m = heart(scale=scale, ape=False, use_checkpoint=False).eval()
+m.load_state_dict(sd, strict=True)
+del sd; gc.collect()
+
+# tiny dummy to minimize activation memory during trace
+x = torch.randn(1, 3, 64, 64)
+torch.onnx.export(
+    m, x, out,
+    input_names=["input"], output_names=["output"],
+    dynamic_axes={"input": {0:"batch",2:"height",3:"width"}, "output": {0:"batch",2:"height_out",3:"width_out"}},
+    opset_version=17, do_constant_folding=True, dynamo=False,
+)
+del m, x; gc.collect()
+print(f"OK -> {out}")