İvme-Denoise-v1 (İvmetron-Guanine1): 2.97M-param Bayesian shrinkage denoiser

README.md

@@ -0,0 +1,179 @@
+---
+license: apache-2.0
+tags:
+- ivme
+- image-denoising
+- denoising
+- image-restoration
+- gaussian-scale-mixture
+- bayesian
+- from-scratch
+- tiny
+- small
+- efficient
+- custom_code
+pipeline_tag: image-to-image
+library_name: pytorch
+datasets:
+- DIV2K
+- BSD500
+---
+
+# İvme-Denoise-v1 (Codenamed Margay 1)
+
+**İvme** (Turkish: *acceleration*) is a series of stupidly small models built to punch above their weight. This is the first vision release: a **2.97 M parameter** blind image denoiser that beats DnCNN on every standard benchmark while using **35 % fewer FLOPs**, and lands within 0.22–0.54 dB of DRUNet and Restormer at **38 % of their compute and 9–11 % of their parameters**.
+
+It is also the first release built on a *derived* architecture rather than a tuned one. Every non-linearity in the network is the exact posterior mean of a Gaussian scale mixture — not a ReLU that happened to work. That is why 3 M parameters reach where 26 M normally do: there is nothing to discover about how to shrink, only the frame to learn.
+
+The codename is a margay, a small nocturnal wild cat with the largest eyes-to-body ratio of any cat — it denoises the dark for a living.
+
+---
+
+## Model Details
+
+| Parameter | Value |
+|---|---|
+| Architecture | İvmetron-Guanine1 (multi-scale Bayesian shrinkage U-Net) |
+| Parameters | 2,971,342 |
+| Compute | 56.8 GFLOP @ 256², 158 GFLOP @ 512² |
+| Scales | 3 (full, 1/2, 1/4) |
+| Blocks per scale | 3 / 3 / 3 Bayesian shrinkage blocks |
+| Width | 96 channels |
+| Dilations | 1 / 2 / 4 (≈150 px receptive field) |
+| Non-linearity | GSM posterior gain — monotone in one scalar energy |
+| Noise conditioning | FiLM + a spatial σ map input (scalar or per-pixel) |
+| Colour space | sRGB, gamma encoded, `[0,1]` |
+| Training | 300 k steps, batch 32, 128² crops, ~7.5 GPU-hours on one GPU |
+| Latency | 17.9 ms / 512² frame (RTX PRO 6000, bf16) |
+| File size | 12 MB (safetensors, fp32) |
+
+### İvmetron-Guanine1
+
+The whole network is built from one block. For an observation `y = x + n` with `n ~ N(0, σ²I)`, the optimal estimator is the posterior mean, and Tweedie's identity turns it into a score:
+
+```
+E[x|y] = y + σ² ∇_y log p_σ(y)
+```
+
+Model a patch of coefficients as a Gaussian scale mixture, `c = √z · w` with `w ~ N(0, C)` and a hidden positive scalar `z`. Then
+
+```
+E[c|y] = E_z[ C (C + σ²/z · I)⁻¹ y ]
+```
+
+Because `C` commutes with `I`, diagonalising `C = UΛUᵀ` gives
+
+```
+E[c|y] = U diag(g₁(e), …, g_d(e)) Uᵀ y ,    g_i(e) = E[ λ_i / (λ_i + σ²/z) | e ]
+```
+
+where the posterior over `z` depends on `y` **only through one scalar** — the local energy `e`. That is the entire structure of the Bayes-optimal denoiser: *a linear map whose gains are monotone functions of a single energy statistic*. Wiener filtering, wavelet shrinkage, BLS-GSM, EPLL and learned shrinkage fields are all corners of that one family; İvmetron sits at its centre and learns only the frame.
+
+The block is 12 lines:
+
+```
+u   = DW3x3( Lin(x) )                 # analysis in a learned frame
+e   = mean_c( u² ) * 3x3              # the shared scalar energy  (ẑ)
+g   = 2 · sigmoid( w · log(e) + b )   # GSM posterior gain, monotone in log e
+v   = u * g                           # multiplicative shrinkage — no bias
+out = x + Proj( SimpleGate(v) )       # synthesis + residual
+```
+
+Note what is absent: no additive bias after the gain, no normalisation. A Bayes gain can only attenuate. Imposing that constraint is worth several dB of sample efficiency.
+
+The upsampling path uses `Conv1×1 → bilinear ×2 → Conv1×1` rather than `PixelShuffle`, because `PixelShuffle` maps a low-resolution tensor onto a 2×2 periodic pattern *by construction* — invisible while signal passes through, and exposed as a hard grid the moment the gated gains zero everything out.
+
+---
+
+## Benchmarks
+
+PSNR in dB, measured with identical noise seeds and the standard protocol (8-pixel border crop for colour, single-channel for Set12).
+
+| dataset | σ | DnCNN | DRUNet | Restormer | **İvme-Denoise-v1** |
+|---|---|---|---|---|---|
+| CBSD68 | 15 | 33.90 | 34.30 | 34.40 | **34.08** |
+| CBSD68 | 25 | 31.24 | 31.69 | 31.79 | **31.44** |
+| CBSD68 | 50 | 27.95 | 28.51 | 28.60 | **28.20** |
+| Kodak24 | 15 | 34.60 | 35.31 | 35.47 | **34.99** |
+| Kodak24 | 25 | 32.14 | 32.89 | 33.04 | **32.56** |
+| Kodak24 | 50 | 28.95 | 29.86 | 30.01 | **29.48** |
+| Set12 (grey) | 15 | 32.67 | 33.25 | 33.42 | **32.86** |
+| Set12 (grey) | 25 | 30.35 | 30.94 | 31.08 | **30.52** |
+| Set12 (grey) | 50 | 27.18 | 27.90 | 28.00 | **27.36** |
+| BSD100 | 25 | — | — | — | **31.25** |
+
+**No cell in bold loses to DnCNN.** Against the large models it is 0.22–0.54 dB behind — at a fraction of the cost:
+
+| | params | GFLOP @256² | train | CBSD68 σ25 |
+|---|---|---|---|---|
+| DnCNN | 0.56 M | ~87 | ~1 GPU-day | 31.24 |
+| DRUNet | 32.6 M | ~150 | — | 31.69 |
+| Restormer | 26.1 M | 141 | 22.5 GPU-h | 31.79 |
+| **İvme-Denoise-v1** | **2.97 M** | **56.8** | **7.5 GPU-h** | **31.44** |
+
+Accuracy per GFLOP: **0.553 dB/GFLOP** versus Restormer's 0.225 — a **2.5× efficiency advantage**.
+
+### Robustness to a wrong noise level
+
+One feature no other denoiser reports: the conditioning is a *hint*, not a command. With the true noise fixed at σ = 5 and the network told something else:
+
+| σ given | 5 | 15 | 25 | 35 | 50 |
+|---|---|---|---|---|---|
+| PSNR (Kodak24) | 40.63 | 40.05 | 38.86 | **37.66** | **35.75** |
+
+A 10× over-estimate costs 5 dB, not 14. This is trained in: 25 % of crops see a deliberately wrong σ, and 40 % see a smooth *spatial* σ field that generates the noise and is fed back as the conditioning map.
+
+---
+
+## Usage
+
+```bash
+pip install torch safetensors numpy pillow
+```
+
+```python
+import numpy as np
+from PIL import Image
+from modeling_ivmetron import load_model, denoise, denoise_tiled, estimate_sigma_map
+
+model = load_model(".", device="cuda")          # or "cpu"
+
+img = np.asarray(Image.open("photo.jpg").convert("RGB"), dtype=np.float32) / 255.0
+
+# 1. you know the noise level (0-255 scale)
+out = denoise(model, img, sigma=15, device="cuda")
+
+# 2. you do not - let it estimate a per-pixel noise map
+out = denoise_tiled(model, img, estimate_sigma_map(img), device="cuda")
+
+Image.fromarray((np.clip(out, 0, 1) * 255).round().astype(np.uint8)).save("clean.png")
+```
+
+`denoise_tiled` is the safe default for anything above ~5 MP: the blocks are 192 channels wide, so a 12 MP frame exceeds 2³¹ elements and CUDA indexing overflows. It blends cosine-windowed 896 px tiles with 96 px overlap.
+
+A 12.3 MP frame takes 6.5 s including tiling, on one GPU.
+
+**Choosing σ.** The model is blind over σ ∈ [5, 50]. If you are not sure, pass `estimate_sigma_map(img)` — that is the intended path for real photographs. Do not pass σ ≥ 30 for a frame whose real noise is well under 10; the network will no longer invent structure (that was the v3 failure mode, fixed), but it will smooth detail away.
+
+---
+
+## Limitations
+
+* **Not state of the art on raw PSNR.** 0.22–0.54 dB behind DRUNet/Restormer. Closing that needs 2–3× the FLOPs, which would remove the point of the model.
+* **No SIDD/DND numbers.** Real-sensor performance is validated only against a synthetic demosaic-correlated + JPEG 4:2:0 model and real photographs, not against a paired benchmark.
+* **sRGB only.** Assumes gamma-encoded 8-bit-style input; do not feed linear RAW or HDR.
+* **Trained on 1 200 natural images** (DIV2K train + valid, BSD500). Faces, text and synthetic content are out of distribution.
+* **JPEG q < 50** is outside the trained range.
+* Training on JPEG-correlated noise can drift hue slightly in saturated regions.
+
+## Citation
+
+```bibtex
+@misc{ivme_denoise_v1,
+  title  = {İvme-Denoise-v1 (İvmetron-Guanine1)},
+  author = {Ekşi, Eren},
+  year   = {2026},
+  note   = {IvmeLabs},
+  url    = {https://huggingface.co/IvmeLabs/Ivme-Denoise-v1}
+}
+```

config.json

@@ -0,0 +1,17 @@
+{
+  "architectures": ["BSNet"],
+  "arch_name": "İvmetron-Guanine1",
+  "model_type": "ivmetron-denoise",
+  "C": 96,
+  "blocks": [3, 3, 3],
+  "dilations": [1, 2, 4],
+  "mix_k": 1,
+  "in_channels": 4,
+  "out_channels": 3,
+  "num_parameters": 2971342,
+  "gflops_256": 56.8,
+  "noise_conditioning": "scalar sigma (0-255) or per-pixel sigma map",
+  "noise_range": [5, 50],
+  "color_space": "sRGB, gamma encoded, [0,1]",
+  "license": "apache-2.0"
+}

example.py

@@ -0,0 +1,10 @@
+"""Minimal example: denoise a photograph with İvme-Denoise-v1."""
+import numpy as np
+from PIL import Image
+from modeling_ivmetron import load_model, denoise_tiled, estimate_sigma_map
+
+model = load_model(".", device="cuda")
+img = np.asarray(Image.open("photo.jpg").convert("RGB"), dtype=np.float32) / 255.0
+out = denoise_tiled(model, img, estimate_sigma_map(img), device="cuda")
+Image.fromarray((np.clip(out, 0, 1) * 255).round().astype(np.uint8)).save("clean.png", quality=95)
+print("wrote clean.png")

model.safetensors

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

modeling_ivmetron.py

@@ -0,0 +1,238 @@
+"""İvmetron-Guanine1 — reference implementation.
+
+A blind single-image denoiser for sRGB photographs whose only non-linearity is
+the exact posterior mean of a Gaussian-scale-mixture prior:
+
+    E[c|y] = E_z[ C (C + sigma^2/z I)^{-1} y ]  =  U diag(g_1(e) ... g_d(e)) U^T y
+
+with the gains monotone in a single scalar local-energy statistic `e`.
+
+Self-contained: needs torch, numpy and Pillow only.
+"""
+from __future__ import annotations
+
+import json
+import math
+import os
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from PIL import Image
+
+__all__ = ["BSNet", "load_model", "denoise", "denoise_tiled", "estimate_sigma_map"]
+
+
+class FiLM(nn.Module):
+    def __init__(self, C, hidden=64):
+        super().__init__()
+        self.net = nn.Sequential(nn.Linear(1, hidden), nn.SiLU(), nn.Linear(hidden, 2 * C))
+        nn.init.zeros_(self.net[-1].weight)
+        nn.init.zeros_(self.net[-1].bias)
+
+    def forward(self, x, s):
+        g, b = self.net(s).chunk(2, -1)
+        return x * (1 + g[..., None, None]) + b[..., None, None]
+
+
+class BSB(nn.Module):
+    """Bayesian Shrinkage Block.
+
+        u   = DW3x3( Lin(x) )                 analysis in a learned frame
+        e   = mean_c(u^2) * 3x3               the shared scalar energy  (z-hat)
+        g   = 2 * sigmoid( w * log(e) + b )   GSM posterior gain, monotone in log e
+        v   = u * g                           multiplicative shrinkage -- no bias
+        out = x + Proj( SimpleGate(v) )
+    """
+
+    def __init__(self, C, exp=2, energy_k=3, dilation=1, mix_k=1):
+        super().__init__()
+        E = C * exp
+        self.mix_k = mix_k
+        self.lin = nn.Conv2d(C, E, 1)
+        self.dw = nn.Conv2d(E, E, 3, padding=dilation, groups=E, dilation=dilation)
+        self.energy = nn.Conv2d(E, 1, energy_k, padding=energy_k // 2, bias=True)
+        self.gate = nn.Conv2d(1, E * mix_k, 1)
+        self.mix = nn.Conv2d(1, mix_k, 1) if mix_k > 1 else None
+        self.proj = nn.Conv2d(E // 2, C, 1)
+        nn.init.zeros_(self.proj.weight)
+        nn.init.zeros_(self.proj.bias)
+        nn.init.constant_(self.gate.bias, 1.0)
+        if self.mix is not None:
+            nn.init.zeros_(self.mix.weight)
+            nn.init.zeros_(self.mix.bias)
+
+    def forward(self, x):
+        u = self.dw(self.lin(x))
+        e = F.avg_pool2d(u * u, 3, 1, 1)
+        e = e.mean(1, keepdim=True)
+        le = torch.log(e.clamp_min(1e-6))
+        if self.mix is None:
+            g = 2.0 * torch.sigmoid(self.gate(le))
+        else:
+            K = self.mix_k
+            w = torch.softmax(self.mix(le).clamp(-8.0, 8.0), 1)
+            gall = 2.0 * torch.sigmoid(self.gate(le))
+            b, ke, h, wd = gall.shape
+            gall = gall.view(b, K, ke // K, h, wd)
+            g = (gall * w[:, :, None]).sum(1)
+        v = u * g
+        a, b = v.chunk(2, 1)
+        return x + self.proj(a * b)
+
+
+class BSNet(nn.Module):
+    def __init__(self, C=96, blocks=(3, 3, 3), in_ch=4, dilations=(1, 2, 4),
+                 legacy=False, mix_k=1):
+        super().__init__()
+        d1, d2, d3 = dilations
+        self.stem = nn.Conv2d(in_ch, C, 3, padding=1)
+        self.film = FiLM(C)
+        self.e1 = nn.Sequential(*[BSB(C, dilation=d1, mix_k=mix_k) for _ in range(blocks[0])])
+        self.d1 = nn.Conv2d(C, 2 * C, 3, stride=2, padding=1)
+        self.e2 = nn.Sequential(*[BSB(2 * C, dilation=d2, mix_k=mix_k) for _ in range(blocks[1])])
+        self.d2 = nn.Conv2d(2 * C, 4 * C, 3, stride=2, padding=1)
+        self.e3 = nn.Sequential(*[BSB(4 * C, dilation=d3, mix_k=mix_k) for _ in range(blocks[2])])
+        if legacy:
+            self.u2 = nn.Sequential(nn.Conv2d(4 * C, 8 * C, 1), nn.PixelShuffle(2),
+                                    nn.Conv2d(2 * C, 2 * C, 1))
+            self.u1 = nn.Sequential(nn.Conv2d(2 * C, 4 * C, 1), nn.PixelShuffle(2),
+                                    nn.Conv2d(C, C, 1))
+        else:
+            self.u2 = nn.Sequential(nn.Conv2d(4 * C, 2 * C, 1),
+                                    nn.Upsample(scale_factor=2, mode="bilinear", align_corners=False),
+                                    nn.Conv2d(2 * C, 2 * C, 1))
+            self.u1 = nn.Sequential(nn.Conv2d(2 * C, C, 1),
+                                    nn.Upsample(scale_factor=2, mode="bilinear", align_corners=False),
+                                    nn.Conv2d(C, C, 1))
+        self.m2 = BSB(2 * C, dilation=d2, mix_k=mix_k)
+        self.m1 = BSB(C, dilation=d1, mix_k=mix_k)
+        self.head = nn.Conv2d(C, 3, 3, padding=1)
+        nn.init.zeros_(self.head.weight)
+        nn.init.zeros_(self.head.bias)
+
+    def forward(self, y, sigma):
+        """sigma: (B,) scalar per image, or (B,1,H,W) / (H,W) spatial noise map."""
+        if sigma.dim() == 1:
+            smap = sigma.view(-1, 1, 1, 1).expand(-1, 1, y.shape[-2], y.shape[-1])
+        elif sigma.dim() == 3:
+            smap = sigma[:, None]
+        elif sigma.dim() == 2:
+            smap = sigma[None, None]
+        else:
+            smap = sigma
+        if sigma.dim() > 1 and smap.shape[-2:] != y.shape[-2:]:
+            smap = F.interpolate(smap, size=y.shape[-2:], mode="bilinear", align_corners=False)
+        s = (smap.mean(dim=(1, 2, 3)).view(-1, 1) / 50.0).clamp(0, 4)
+        f = self.stem(torch.cat([y, (smap / 50.0).clamp(0, 4)], 1))
+        f = self.film(f, s)
+        f1 = self.e1(f)
+        f2 = self.e2(self.d1(f1))
+        f3 = self.e3(self.d2(f2))
+        g2 = self.m2(self.u2(f3) + f2)
+        g1 = self.m1(self.u1(g2) + f1)
+        return y + self.head(g1)
+
+
+# ---------------------------------------------------------------- loading
+def load_model(path_or_repo=".", device="cpu"):
+    """Load the model from a directory (or HF repo id) holding model.safetensors."""
+    from safetensors.torch import load_file
+
+    if os.path.isdir(path_or_repo):
+        cfg_path = os.path.join(path_or_repo, "config.json")
+        st_path = os.path.join(path_or_repo, "model.safetensors")
+    else:
+        from huggingface_hub import hf_hub_download   # only needed for repo ids
+        cfg_path = hf_hub_download(path_or_repo, "config.json")
+        st_path = hf_hub_download(path_or_repo, "model.safetensors")
+
+    cfg = json.load(open(cfg_path)) if os.path.exists(cfg_path) else {}
+    sd, meta = load_file(st_path, device="cpu"), None
+    from safetensors import safe_open
+    with safe_open(st_path, framework="pt") as f:
+        meta = f.metadata() or {}
+
+    C = int(meta.get("C", cfg.get("C", 96)))
+    blocks = tuple(int(x) for x in meta.get("blocks", "3,3,3").split(","))
+    dilations = tuple(int(x) for x in meta.get("dilations", "1,2,4").split(","))
+    mix_k = int(meta.get("mix_k", 1))
+
+    model = BSNet(C=C, blocks=blocks, dilations=dilations, mix_k=mix_k)
+    model.load_state_dict(sd)
+    return model.to(device).eval()
+
+
+# ---------------------------------------------------------------- inference
+@torch.no_grad()
+def denoise(model, img, sigma, device="cuda", pad=8):
+    """img: HxWx3 float32 in [0,1].  sigma: scalar, or HxW noise map in 0-255 units."""
+    x = torch.from_numpy(np.ascontiguousarray(img.transpose(2, 0, 1)))[None].float().to(device)
+    H, W = x.shape[-2:]
+    ph, pw = (pad - H % pad) % pad, (pad - W % pad) % pad
+    xp = F.pad(x, (0, pw, 0, ph), mode="reflect")
+    if np.ndim(sigma) == 2:
+        st = torch.from_numpy(np.asarray(sigma, np.float32))[None, None].to(device)
+        s = F.pad(st, (0, pw, 0, ph), mode="reflect")
+    else:
+        s = torch.tensor([float(sigma)], device=device)
+    if device == "cuda":
+        with torch.autocast("cuda", dtype=torch.bfloat16):
+            out = model(xp, s)
+    else:
+        out = model(xp, s)
+    return out[..., :H, :W].clamp(0, 1)[0].float().cpu().numpy().transpose(1, 2, 0)
+
+
+def estimate_sigma_map(img, block=32, pct=0.10, smooth=49):
+    """Robust per-block noise estimate -> HxW map in 0-255 units."""
+    x = torch.from_numpy(np.ascontiguousarray(img.transpose(2, 0, 1)))[None].float()
+    H, W = x.shape[-2:]
+    g = x.mean(1, keepdim=True)
+    k = torch.tensor([[0, 1, 0], [1, -4, 1], [0, 1, 0]], dtype=torch.float32).view(1, 1, 3, 3)
+    hp = F.conv2d(F.pad(g, (1, 1, 1, 1), mode="reflect"), k).abs() / 4.4721
+    B = block
+    gh, gw = H // B, W // B
+    blk = hp[..., :gh * B, :gw * B].reshape(1, 1, gh, B, gw, B).permute(0, 1, 2, 4, 3, 5)
+    blk = blk.reshape(1, 1, gh, gw, B * B)
+    sig = blk.sort(-1).values[..., int(pct * B * B)] / 0.1257
+    sig = F.interpolate(sig[0, 0][None, None], size=(H, W), mode="bilinear", align_corners=False)
+    sig = F.avg_pool2d(sig, smooth, 1, smooth // 2)
+    return (sig * 255.0).clamp(1.5, 60.0)[0, 0].numpy()
+
+
+def _cos_window(n):
+    return torch.hann_window(n, periodic=False).clamp_min(1e-3)
+
+
+def denoise_tiled(model, img, sigma, tile=896, overlap=96, device="cuda"):
+    """Tiled inference with cosine blending, for frames above ~5 MP."""
+    H, W = img.shape[:2]
+    if H * W <= 5_000_000:
+        return denoise(model, img, sigma, device=device)
+    step = tile - overlap
+    xs = sorted(set(min(x, max(0, W - tile)) for x in list(range(0, max(1, W - tile), step)) + [max(0, W - tile)]))
+    ys = sorted(set(min(y, max(0, H - tile)) for y in list(range(0, max(0, H - tile), step)) + [max(0, H - tile)]))
+    th, tw = min(tile, H), min(tile, W)
+    win = (_cos_window(th)[:, None] * _cos_window(tw)[None, :]).clamp_min(1e-4)
+    acc = torch.zeros(3, H, W)
+    wsum = torch.zeros(1, H, W)
+    smap = None if np.ndim(sigma) != 2 else np.asarray(sigma, np.float32)
+    for y in ys:
+        for x in xs:
+            s = sigma if smap is None else smap[y:y + th, x:x + tw]
+            out = denoise(model, img[y:y + th, x:x + tw], s, device=device)
+            t = torch.from_numpy(np.ascontiguousarray(out.transpose(2, 0, 1)))
+            acc[:, y:y + th, x:x + tw] += t * win
+            wsum[:, y:y + th, x:x + tw] += win
+    return (acc / wsum.clamp_min(1e-6)).clamp(0, 1).permute(1, 2, 0).numpy()
+
+
+def load_image(path):
+    return np.asarray(Image.open(path).convert("RGB"), dtype=np.float32) / 255.0
+
+
+def save_image(path, arr, quality=95):
+    Image.fromarray((np.clip(arr, 0, 1) * 255).round().astype(np.uint8)).save(
+        path, quality=quality, subsampling=0)

requirements.txt

@@ -0,0 +1,4 @@
+torch>=2.0
+safetensors>=0.4
+numpy
+pillow