MangaJPEGLQ as a 1× Core AI asset from the original weights

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+ESRGAN-MangaJPEGLQ.aimodel/main.mlirb filter=lfs diff=lfs merge=lfs -text

ESRGAN-MangaJPEGLQ.aimodel/main.hash

@@ -0,0 +1 @@
+�RI\�JF8M}@~��Ѥ�vyqU��d(�n��q

ESRGAN-MangaJPEGLQ.aimodel/main.mlirb

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

ESRGAN-MangaJPEGLQ.aimodel/metadata.json

@@ -0,0 +1,8 @@
+{
+  "description" : "1x-MangaJPEGLQ: x1 super-resolution, 256x256 -> 256x256 tile, float16, [0,1] RGB in\/out.",
+  "assetVersion" : "2.0",
+  "producer" : "coreai-core 1.0.0b3",
+  "license" : "CC-BY-NC-SA-4.0",
+  "author" : "bunzero (MangaJPEGLQ)",
+  "creationDate" : "20261012T085404Z"
+}

README.md

@@ -31,3 +31,14 @@ VincentGOURBIN/RealESRGAN-CoreML) labels them BSD-3; that is wrong for these two
 None to the models: each archive matched the SHA-256 Hugging Face listed upstream and is republished
 unzipped. Attribution: 4x-UltraSharp © Kim2091; MangaJPEGLQ © bunzero; Core ML conversion by
 VincentGOURBIN. This repository is shared under the same licence (CC BY-NC-SA 4.0).
+
+## MangaJPEGLQ on Core AI (`ESRGAN-MangaJPEGLQ.aimodel`, swift-upscale 2.1+)
+
+`upscale --variant manga` crashed in swift-upscale 2.0 / upscale 0.3.0: **MangaJPEGLQ is a 1× model**
+(it cleans JPEG blocks out of greyscale manga without enlarging — OpenModelDB lists `1x-MangaJPEGLQ`,
+bunzero, CC BY-NC-SA 4.0), and the Core ML build was run as if it were 4×. This Core AI asset is
+exported from bunzero's original weights (`1x-MangaJPEGLQ.pth`, OpenModelDB's file host) with
+`export_esrgan.py`: a 256-px tile, float16, RGB in → Rec. 601 luma → the 1-channel model → copied to RGB.
+Parity against the original PyTorch weights on six stock-photo tiles: **min 56.6 dB, mean 61.0 dB**
+(`parity.py`). `upscale` now runs it as a clean-up pass before the default 4× model.
+The Core ML `.mlpackage` builds above stay for older versions.

export_esrgan.py

@@ -0,0 +1,56 @@
+# /// script
+# requires-python = ">=3.11"
+# dependencies = ["coreai-core==1.0.0b3", "coreai-torch==0.4.3", "spandrel==0.4.1", "numpy", "pillow"]
+# [tool.uv]
+# index-url = "https://pypi.org/simple"
+# prerelease = "allow"
+# index-strategy = "unsafe-best-match"
+# ///
+"""Export an ESRGAN-family upscaler (RRDBNet / SRVGGNetCompact, via spandrel) as a fixed-tile
+Core AI asset with the same contract swift-upscale's CoreAIUpscaler runs for SwinIR-L:
+input "x" [0,1] RGB, output "output" [0,1] RGB, clamped.
+    uv run export_esrgan.py <ckpt.pth> <out.aimodel> <tile> <float16|float32> <author> <licence>
+"""
+import sys, time
+from pathlib import Path
+import torch
+from spandrel import ModelLoader
+from coreai.runtime import AIModelAssetMetadata
+from coreai_torch import TorchConverter, get_decomp_table
+
+ckpt, out, tile, dtype_name, author, licence = sys.argv[1], Path(sys.argv[2]), int(sys.argv[3]), sys.argv[4], sys.argv[5], sys.argv[6]
+dtype = {"float16": torch.float16, "float32": torch.float32}[dtype_name]
+desc = ModelLoader().load_from_file(ckpt)
+print("arch", type(desc.model).__name__, "scale", desc.scale, flush=True)
+
+class Wrapped(torch.nn.Module):
+    """RGB in, RGB out. A greyscale model (one input channel — MangaJPEGLQ is built for
+    black-and-white manga) sees Rec. 601 luma, and its single output channel is copied to
+    all three, so the Swift side feeds and reads RGB either way."""
+    def __init__(self, m, in_ch, out_ch):
+        super().__init__(); self.m = m; self.in_ch = in_ch; self.out_ch = out_ch
+        self.register_buffer("luma", torch.tensor([0.299, 0.587, 0.114]).view(1, 3, 1, 1))
+    def forward(self, x):
+        if self.in_ch == 1:
+            x = (x * self.luma.to(x.dtype)).sum(dim=1, keepdim=True)
+        y = self.m(x)
+        if self.out_ch == 1:
+            y = y.repeat(1, 3, 1, 1)
+        return y.clamp(0, 1)
+
+print("channels", desc.input_channels, "->", desc.output_channels, flush=True)
+model = Wrapped(desc.model.eval(), desc.input_channels, desc.output_channels).to(dtype)
+example = {"x": torch.rand(1, 3, tile, tile).to(dtype)}
+t0 = time.time()
+with torch.no_grad(), torch.autocast(device_type="cpu", dtype=dtype, enabled=dtype != torch.float32):
+    exported = torch.export.export(model, args=(), kwargs=example)
+exported = exported.run_decompositions(get_decomp_table())
+program = TorchConverter(mode=TorchConverter.Mode.RELEASE).add_exported_program(
+    exported_program=exported, input_names=["x"], output_names=["output"]).to_coreai()
+md = AIModelAssetMetadata()
+md.author = author
+md.license = licence
+md.model_description = (f"{Path(ckpt).stem}: x{desc.scale} super-resolution, {tile}x{tile} -> "
+                        f"{tile*desc.scale}x{tile*desc.scale} tile, {dtype_name}, [0,1] RGB in/out.")
+program.save_asset(out, metadata=md)
+print(f"saved {out} in {time.time()-t0:.1f} s", flush=True)

parity.py

@@ -0,0 +1,51 @@
+# /// script
+# requires-python = ">=3.11"
+# dependencies = ["coreai-core==1.0.0b3", "coreai-torch==0.4.3", "spandrel==0.4.1", "numpy", "pillow"]
+# [tool.uv]
+# index-url = "https://pypi.org/simple"
+# prerelease = "allow"
+# index-strategy = "unsafe-best-match"
+# ///
+"""Core AI export vs the original PyTorch weights (float32), tile by tile, on stock photos."""
+import asyncio, sys, time
+import numpy as np, torch
+from PIL import Image
+from spandrel import ModelLoader
+import coreai.runtime as rt
+
+ckpt, path, tile = sys.argv[1], sys.argv[2], int(sys.argv[3])
+images = sys.argv[4:]
+desc = ModelLoader().load_from_file(ckpt)
+net = desc.model.eval().float()
+luma = torch.tensor([0.299, 0.587, 0.114]).view(1, 3, 1, 1)
+def reference(x):
+    if desc.input_channels == 1: x = (x * luma).sum(dim=1, keepdim=True)
+    y = net(x)
+    if desc.output_channels == 1: y = y.repeat(1, 3, 1, 1)
+    return y.clamp(0, 1)
+def psnr(a, b):
+    mse = float(np.mean((a.astype(np.float64) - b.astype(np.float64)) ** 2))
+    return 99.0 if mse == 0 else 10 * np.log10(1.0 / mse)
+async def main():
+    rng = np.random.default_rng(0); tiles = []
+    for p in images:
+        im = np.asarray(Image.open(p).convert("RGB")).astype(np.float32) / 255
+        h, w, _ = im.shape
+        for _ in range(3):
+            y = int(rng.integers(0, max(1, h - tile))); x = int(rng.integers(0, max(1, w - tile)))
+            c = im[y:y+tile, x:x+tile]
+            if c.shape[:2] == (tile, tile): tiles.append(c)
+    m = await rt.AIModel.load(path)
+    fn = m.load_function("main")
+    scores, times = [], []
+    for c in tiles:
+        x = torch.from_numpy(c).permute(2, 0, 1)[None].contiguous()
+        with torch.no_grad(): ref = reference(x)[0].permute(1, 2, 0).numpy()
+        t = time.time()
+        out = await fn({"x": rt.NDArray(data=x.half().contiguous())})
+        times.append(time.time() - t)
+        got = out["output"].numpy().astype(np.float32)[0].transpose(1, 2, 0)
+        scores.append(psnr(ref, got))
+    print(f"{path.split('/')[-1]:34} scale x{desc.scale}  PSNR vs PyTorch fp32 over {len(scores)} tiles: "
+          f"min {min(scores):.1f} dB, mean {np.mean(scores):.1f} dB; {np.median(times)*1000:.0f} ms/tile")
+asyncio.run(main())