satlink: add RS-Tokenizer backend (256x18b, ~114x, ~41dB) alongside One-D-Piece; pixel path only

satlink/API.md

@@ -8,6 +8,16 @@
 | `satlink.Ground` | 码字 → 像素 | 337.8M 参数 | 地面 |
 | `satlink.GroundVLM` | 码字 → 答案 | 解码器 + 视觉塔 + LLM，一个整体 | 地面 |
 
+支持两个 tokenizer 后端，由权重文件自动识别，不用传标志：
+
+| 后端 | 权重 | 每片 | 单片下行 | 压缩比 | PSNR | 通路 |
+|---|---|---|---|---|---|---|
+| `opd` One-D-Piece | `sat_encoder.pt` / `ground_decoder.pt` | 48 tok × 12 bit | 72 B | ~910× | ~21 dB | 像素 + 隐空间 |
+| `rstok` RS-Tokenizer | `rstok_sat.pt` / `rstok_ground.pt` | 256 tok × 18 bit | 576 B | ~114× | ~41 dB | **仅像素** |
+
+多 8 倍带宽换约 20 dB。选哪个取决于链路和画质哪个更紧张。`rstok` 只接了像素通路——
+隐空间通路的 bridge 是针对 One-D-Piece 那个特定的中间层特征训出来的，换 tokenizer 能跑但是错的。
+
 两个值类型在它们之间传递：`satlink.VQFrame`（码字 + 摆放几何）和 `satlink.LinkProfile`
 （帧长宽 / 切片大小 / 每片 token 数——星地**约定好的常量，不上下行**）。
 
@@ -48,6 +58,26 @@ sat.save(frame, "codes.json")                    # 落盘（原子写入）
 附近**的码字翻转，实测约 0.2%–0.5%。平均重建影响 < 0.4/255，但局部单点最大可达 ~118/255。
 结论是：不同 batch 的两次运行**不保证逐位相同**，验收判据应当看重建像素，而不是码字。
 
+**关于精度（重要，两个后端的约束相反）。**
+
+| | 编码器 | 解码器 |
+|---|---|---|
+| `opd` | fp16 / bf16 / fp32 均可 | fp16 / bf16 / fp32 均可（默认 bf16） |
+| `rstok` | **不能 bf16**（用 fp16） | **不能 fp16**（用 bf16，默认） |
+
+`rstok` 编码器在 bf16 下，8 张参考图里有 1 张从 41.4 dB 塌到 22.5 dB——19 万码字在 24 维单位球面上
+近邻的余弦间隔极小，bf16 的 8 位尾数分不开，argmin 会整块选错。fp16 全程与 fp32 差 <0.5 dB。
+
+`rstok` 解码器反过来：1024 宽 24 层，激活会冲出 fp16 的动态范围，**整帧变成 NaN**，
+然后被 cast 成一张全黑图。
+
+**所以解码器统一走 bf16**——实测两个后端的 bf16 与 fp32 差 0.00 dB，显存减半，也绕开了上面这个坑。
+`Ground` 和 `run_ground` 的默认值都是 `bfloat16`。保留两道护栏：`RSTokDecoder` 构造时拒绝 fp16，
+`Ground.decode` 检查输出是否有限——非有限直接报错，而不是输出一张看起来「成功」的黑图。
+
+（官方参考输出是 bf16 跑的，所以继承了编码器那个塌陷。`test_rstok.py` 把精确匹配钉在 bf16 上，
+那是**溯源检查**，不是推荐用 bf16。）
+
 **关于图像格式。** PNG 和 BMP 都支持，且**结果逐位一致**（已验证）。两者读法不同：
 BMP 未压缩、行有 4 字节对齐填充，可以 `np.memmap` 按需读某一条；PNG 是单条 zlib 流、
 行间有 filter 依赖，无法定位到第 N 行，只能整幅解码一次再切。接口一样，代价不一样。
@@ -176,7 +206,9 @@ python -m satlink.accept.run_ground --codes results/case1/codes \
 | PNG 与 BMP 重建 PSNR | 均为 20.80 dB |
 | PNG 与 BMP 的 LLM 回答 | 完全相同 |
 | 流式编码 vs 整幅 `encode()` | batch ≥8 时逐位一致 |
-| 32768×30720 单帧 | 2.9 s，5365 crop/s，下行 1.11 MB，910.2×，峰值显存 20.16 GB |
+| 32768×30720 单帧 (opd) | 2.9 s，5365 crop/s，下行 1.11 MB，910.2×，峰值显存 20.16 GB |
+| 同一张 1024×768 上两个后端 | opd 864 B / 910.2× / 20.80 dB；rstok 6912 B / 113.8× / 27.56 dB |
+| rstok vendor 正确性 | bf16 下码字与官方参考 **100.00%** 一致；用官方码字解码与官方重建差 0.36/255 |
 
 `batch=1` 是唯一的例外：约 0.2%–0.5% 的码字会翻转（见上文「关于 batch」），
 平均重建影响 <0.4/255。判据看的是重建像素，不是码字。

satlink/accept/run_ground.py

@@ -59,9 +59,11 @@ def main():
     ap = argparse.ArgumentParser()
     ap.add_argument("--codes", required=True, help="directory of codes .json from the satellite")
     ap.add_argument("--out", required=True)
-    ap.add_argument("--decoder", default="satlink/weights/ground_decoder.pt")
+    ap.add_argument("--decoder", default="satlink/weights/ground_decoder.pt",
+                    help="ground_decoder.pt (One-D-Piece) or rstok_ground.pt (RS-Tokenizer)")
     ap.add_argument("--source", default=None, help="original frames, for PSNR (optional)")
-    ap.add_argument("--dtype", default="float16", choices=["float16", "bfloat16", "float32"])
+    ap.add_argument("--dtype", default="bfloat16", choices=["float16", "bfloat16", "float32"],
+                    help="decoder precision; bfloat16 matches float32 to 0.01 dB on both backends")
     ap.add_argument("--device", default="auto")
     ap.add_argument("--psnr-min", type=float, default=15.0,
                     help="floor for the reconstruction check; see the note in the report")
@@ -91,6 +93,11 @@ def main():
     # ------------------------------------------------------------ reconstruction
     s = Step("R1", "地面重建与视觉检查")
     s.check(len(files) > 0, f"读取到 {len(files)} 个码流文件")
+    # bfloat16 throughout: identical to float32 within 0.01 dB on both backends, and it avoids
+    # the float16 overflow that turns an RS-Tokenizer decode into NaN.
+    from ..api import _sniff
+    backend = _sniff(a.decoder)
+    print(f"  backend  {backend}")
     g = Ground(a.decoder, device=device, dtype=a.dtype)
     rows = []
     from PIL import Image

satlink/accept/run_sat.py

@@ -32,7 +32,6 @@ import torch
 
 from ..device import name as dev_name, resolve
 from ..protocol import LinkProfile
-from ..sat.encoder import SatEncoder
 from ..api import encode_one
 from ..imageio import StripReader
 from .power import PowerMonitor
@@ -96,8 +95,11 @@ def main():
     ap.add_argument("--case", type=int, required=True, choices=[1, 2])
     ap.add_argument("--data", required=True, help="input directory on the SSD")
     ap.add_argument("--out", required=True, help="output directory for codes and the report")
-    ap.add_argument("--weights", default="satlink/weights/sat_encoder.pt")
-    ap.add_argument("--n-tok", type=int, default=48)
+    ap.add_argument("--weights", default="satlink/weights/sat_encoder.pt",
+                    help="sat_encoder.pt (One-D-Piece) or rstok_sat.pt (RS-Tokenizer); the "
+                         "backend and its token geometry are read from the file")
+    ap.add_argument("--n-tok", type=int, default=0,
+                    help="override tokens per crop; 0 = whatever the backend requires")
     ap.add_argument("--dtype", default="float16", choices=["float16", "bfloat16", "float32"])
     ap.add_argument("--batch", type=int, default=64, help="crops per encoder call")
     ap.add_argument("--device", default="auto", help="auto | musa | cuda | cpu")
@@ -153,9 +155,23 @@ def main():
         _finish(a, steps, t_start, device, 1)
         return 1
 
-    sat = SatEncoder(a.weights, LinkProfile(*readable[0][1].im.size, n_tok=a.n_tok),
-                     device=device, dtype=getattr(torch, a.dtype))
+    # Which tokenizer this is comes from the checkpoint, not from a flag the operator has to keep
+    # in sync with --weights. Each backend pins its own tokens-per-crop and codebook width.
+    from ..api import TOKENIZERS, _sniff, profile_for
+    backend = _sniff(a.weights)
+    n_tok = a.n_tok or TOKENIZERS[backend]["n_tok"]
+    if backend == "rstok":
+        from ..sat.rstok import RSTokEncoder as E
+    else:
+        from ..sat.encoder import SatEncoder as E
+    sat = E(a.weights, profile_for(backend, *readable[0][1].im.size, n_tok=n_tok),
+            device=device, dtype=getattr(torch, a.dtype))
+    print(f"  backend  {backend}: {n_tok} tok x {sat.profile.bits} b = "
+          f"{sat.profile.crop_bytes} B/crop")
     print(f"  model    {sum(p.numel() for p in sat.parameters())/1e6:.2f}M params on {device}")
+    if backend == "rstok" and a.dtype == "bfloat16":
+        raise SystemExit("RS-Tokenizer 在 bf16 下会有整块 crop 塌陷(实测 8 张里 1 张 41.4->22.5 dB)"
+                         "，请用 --dtype float16")
 
     # ---------------------------------------------------------------- F02 / loop
     s2 = Step("F02", "离散Token压缩")
@@ -166,7 +182,7 @@ def main():
         while True:
             rounds += 1
             for path, reader in readable:
-                prof = LinkProfile(reader.width, reader.height, n_tok=a.n_tok)
+                prof = profile_for(backend, reader.width, reader.height, n_tok=n_tok)
                 sat.profile = prof
                 frame, secs, peak = encode_one(sat, reader, a.batch)
                 raw_bits = reader.width * reader.height * 8          # grayscale, the conservative
@@ -176,6 +192,7 @@ def main():
                 rec = {"input": os.path.abspath(path), "width": reader.width,
                        "height": reader.height, "grid": list(prof.grid),
                        "n_crops": prof.n_tiles, "n_tok": prof.n_tok, "tile": prof.tile,
+                       "bits": prof.bits, "backend": backend,
                        "raw_bits": raw_bits, "token_bits": tok_bits, "ratio": ratio,
                        "seconds": secs, "crops_per_s": prof.n_tiles / max(secs, 1e-9),
                        "peak_bytes": peak, "round": rounds}
@@ -185,6 +202,7 @@ def main():
                     _atomic_write_json(out, {
                         "image": os.path.basename(path), "width": reader.width,
                         "height": reader.height, "n_tok": prof.n_tok, "dtype": a.dtype,
+                        "bits": prof.bits, "tile": prof.tile, "backend": backend,
                         "tile_origins": frame.tile_origins(), "codes": frame.tolist()})
                     rec["codes"] = os.path.abspath(out)
                 records.append(rec)
@@ -248,12 +266,12 @@ def main():
             j = json.load(open(r["codes"]))
             n = len(j["codes"])
             assert n == r["n_crops"] and all(len(c) == r["n_tok"] for c in j["codes"])
-            assert all(0 <= v < 4096 for c in j["codes"] for v in c)
+            assert all(0 <= v < (1 << r["bits"]) for c in j["codes"] for v in c)
             ok_files += 1
         except Exception as e:
             parse_bad.append(f"{os.path.basename(r['codes'])}: {type(e).__name__}")
     s3.check(ok_files == len(first),
-             f"{ok_files}/{len(first)} 个结果文件可解析、字段完整、码字在 0..4095 内"
+             f"{ok_files}/{len(first)} 个结果文件可解析、字段完整、码字在码本范围内"
              + ("" if not parse_bad else f" 失败: {parse_bad[:3]}"))
     s3.check(not any(f.endswith(".tmp") for f in os.listdir(codes_dir)),
              "无遗留 .tmp -- 全部为原子写入(先写临时文件再 rename)")

satlink/api.py

@@ -43,8 +43,49 @@ from .frame import VQFrame
 from .imageio import StripReader
 from .protocol import LinkProfile, pack_frame, unpack_frame
 
-__all__ = ["Satellite", "Ground", "GroundVLM", "VQFrame", "LinkProfile",
-           "save_codes", "load_codes", "frame_from_wire", "read_image", "encode_one"]
+__all__ = ["Satellite", "Ground", "GroundVLM", "VQFrame", "LinkProfile", "TOKENIZERS",
+           "save_codes", "load_codes", "frame_from_wire", "read_image", "encode_one",
+           "profile_for"]
+
+
+# ---------------------------------------------------------------------------- tokenizer backends
+# Two tokenizers are supported and they are NOT interchangeable at a given profile -- each pins its
+# own tokens-per-crop and codebook width, so the profile is derived from the backend rather than
+# passed in and hoped for.
+#
+#   opd    One-D-Piece      48 tokens x 12 bits =  72 B/crop, ~910x, ~21 dB
+#   rstok  RS-Tokenizer     256 tokens x 18 bits = 576 B/crop, ~114x, ~41 dB
+#
+# 8x the downlink buys about 20 dB. Which one is right depends on whether the link or the imagery
+# is the scarce resource, so both ship.
+TOKENIZERS = {
+    "opd":   {"n_tok": 48,  "bits": 12, "tile": 256},
+    "rstok": {"n_tok": 256, "bits": 18, "tile": 256},
+}
+
+
+def profile_for(tokenizer, width, height, **kw):
+    """-> a LinkProfile with the token geometry that `tokenizer` actually requires."""
+    if tokenizer not in TOKENIZERS:
+        raise ValueError(f"unknown tokenizer {tokenizer!r}; expected one of {sorted(TOKENIZERS)}")
+    spec = dict(TOKENIZERS[tokenizer])
+    spec.update(kw)
+    return LinkProfile(width, height, **spec)
+
+
+def _sniff(weights):
+    """-> "opd" | "rstok", read off the checkpoint's own config.
+
+    Sniffing rather than making the caller declare it: the two checkpoints are structurally
+    distinct, and a mismatched declaration would fail deep inside a state-dict load with an
+    unhelpful message."""
+    import torch
+    cfg = torch.load(weights, map_location="cpu", weights_only=False).get("config", {})
+    if "token_grid" in cfg or "quantizer" in cfg:
+        return "rstok"
+    if "model" in cfg:
+        return "opd"
+    raise ValueError(f"cannot tell which tokenizer {weights} holds")
 
 
 # ---------------------------------------------------------------------------- codes on disk
@@ -77,14 +118,20 @@ def load_codes(path):
     return VQFrame(torch.tensor(j["codes"], dtype=torch.long), prof)
 
 
-def frame_from_wire(wire, width, height, n_tok=48, tile=256, bits=12):
+def frame_from_wire(wire, width=None, height=None, n_tok=48, tile=256, bits=12, profile=None):
     """-> VQFrame, from packed downlink bytes.
 
-    width/height are supplied by the caller because the stream does NOT carry them: frame size,
-    tile size and token count are deployment constants agreed by both ends, so spending downlink
-    bytes on a header that always says the same thing is pure overhead."""
+    Nothing in the stream says how big the frame is or how it was tokenized: frame size, tile size,
+    token count and codebook width are deployment constants both ends already agree on, so
+    spending downlink bytes on a header that always says the same thing is pure overhead. The
+    caller therefore supplies the geometry.
+
+    Pass `profile=` when you already have one (from the encoder, or from load_codes) -- it carries
+    all four constants together. The loose width/height form defaults to One-D-Piece's geometry,
+    which is wrong for the RS-Tokenizer, and getting only n_tok right while leaving bits at 12 is
+    exactly the mistake that produces a wrong-length frame."""
     import torch
-    prof = LinkProfile(width, height, tile=tile, n_tok=n_tok, bits=bits)
+    prof = profile or LinkProfile(width, height, tile=tile, n_tok=n_tok, bits=bits)
     return VQFrame(torch.tensor(unpack_frame(prof, wire), dtype=torch.long), prof)
 
 
@@ -171,13 +218,19 @@ class Satellite:
     two runs at different batch sizes are not required to be bit-identical.
     """
 
-    def __init__(self, weights, width, height, n_tok=48, tile=256, bits=12,
-                 device="auto", dtype="float16"):
+    def __init__(self, weights, width, height, tokenizer="auto", n_tok=None, tile=None,
+                 bits=None, device="auto", dtype="float16"):
         import torch
-        from .sat.encoder import SatEncoder
-        self.profile = LinkProfile(width, height, tile=tile, n_tok=n_tok, bits=bits)
+        self.tokenizer = _sniff(weights) if tokenizer == "auto" else tokenizer
+        over = {k: v for k, v in (("n_tok", n_tok), ("tile", tile), ("bits", bits))
+                if v is not None}
+        self.profile = profile_for(self.tokenizer, width, height, **over)
         dt = getattr(torch, dtype) if isinstance(dtype, str) else dtype
-        self.encoder = SatEncoder(weights, self.profile, device=device, dtype=dt)
+        if self.tokenizer == "rstok":
+            from .sat.rstok import RSTokEncoder as E
+        else:
+            from .sat.encoder import SatEncoder as E
+        self.encoder = E(weights, self.profile, device=device, dtype=dt)
         self.device, self.dtype = self.encoder.device, dt
 
     @property
@@ -221,15 +274,24 @@ class Satellite:
 
 # ---------------------------------------------------------------------------- ground
 class Ground:
-    """Ground side, pixel reconstruction only. Codes in, image out."""
+    """Ground side, pixel reconstruction only. Codes in, image out.
 
-    def __init__(self, weights, device="auto", dtype="float32"):
+    bfloat16 by default: measured identical to float32 to 0.01 dB on both backends, at half the
+    memory. float16 is NOT the default because the RS-Tokenizer decoder overflows in it -- the
+    whole frame comes back NaN, which casts to a black image rather than to an error.
+    """
+
+    def __init__(self, weights, tokenizer="auto", device="auto", dtype="bfloat16"):
         import torch
         from .device import resolve
-        from .ground.decoder import GroundDecoder
+        self.tokenizer = _sniff(weights) if tokenizer == "auto" else tokenizer
         dt = getattr(torch, dtype) if isinstance(dtype, str) else dtype
         self.device = resolve(device)
-        self.decoder = GroundDecoder(weights, dtype=dt).to(self.device)
+        if self.tokenizer == "rstok":
+            from .ground.rstok_ground import RSTokDecoder as D
+        else:
+            from .ground.decoder import GroundDecoder as D
+        self.decoder = D(weights, dtype=dt).to(self.device)
 
     def decode(self, frame):
         """-> float tensor [H, W] in [0, 1], overlaps averaged."""
@@ -243,7 +305,16 @@ class Ground:
         # discarding two thirds of the output.
         while x.ndim > 3:
             x = x[0]
-        return x.mean(0) if x.ndim == 3 else x
+        x = x.mean(0) if x.ndim == 3 else x
+        # A non-finite decode must be an error, not a picture. numpy casts NaN to 0 without
+        # complaint, so an overflowed decoder silently produces a plausible-looking all-black PNG
+        # that passes every check except a PSNR comparison the caller may not be running.
+        if not torch.isfinite(x).all():
+            raise RuntimeError(
+                f"decoder produced non-finite output ({int((~torch.isfinite(x)).sum())} of "
+                f"{x.numel()} values). The RS-Tokenizer decoder overflows in float16 -- "
+                f"use dtype='float32' or 'bfloat16'.")
+        return x
 
     def reconstruct(self, frame):
         """-> PIL grayscale image."""
@@ -268,6 +339,10 @@ class GroundVLM:
         from .device import resolve
         if bool(decoder) == bool(latent_ground):
             raise ValueError("give exactly one of decoder= or latent_ground=")
+        # Only the pixel path supports the RS-Tokenizer: the latent path projects a PARTIAL decode
+        # into the vision tower, through a bridge trained against One-D-Piece's particular
+        # intermediate features. Pointing it at another tokenizer would run, and would be wrong.
+        self.tokenizer = _sniff(decoder) if decoder else "opd"
         self.device = resolve(device)
         self.model_dir = model_dir
         self.latent = bool(latent_ground)

satlink/bench/test_isolation.py

@@ -120,13 +120,23 @@ def main():
         else:
             print(f"PASS satlink/{m}.py: deps within {sorted(SHARED_DEPS.get(m, set())) or 'none'}")
 
-    # the satellite bundle must not contain decoder weights or decoder classes
-    enc = open(os.path.join(ROOT, "sat", "opd_encoder.py"), encoding="utf-8").read()
-    if "class TiTokDecoder" in enc or "PixelDecoder" in enc:
-        fails += 1
-        print("FAIL satlink/sat/opd_encoder.py contains decoder classes")
-    else:
-        print("PASS satlink/sat/opd_encoder.py: no decoder classes")
+    # The satellite bundle must not contain decoder classes. Checked per vendored file rather
+    # than by scanning the package, because each tokenizer's decoder has its own class name and a
+    # generic "*Decoder" grep would both miss names and fire on unrelated ones.
+    for fn, banned in (("opd_encoder.py", ("class TiTokDecoder", "PixelDecoder")),
+                       ("rstok_encoder.py", ("class LlamaTransformerDecoder",))):
+        path = os.path.join(ROOT, "sat", fn)
+        if not os.path.isfile(path):
+            fails += 1
+            print(f"FAIL satlink/sat/{fn} is missing")
+            continue
+        src = open(path, encoding="utf-8").read()
+        hit = [b for b in banned if b in src]
+        if hit:
+            fails += 1
+            print(f"FAIL satlink/sat/{fn} contains decoder classes: {hit}")
+        else:
+            print(f"PASS satlink/sat/{fn}: no decoder classes")
 
     for sub in ("sat", "ground"):
         deps = set()

satlink/bench/test_rstok.py

@@ -0,0 +1,168 @@
+"""RS-Tokenizer backend: does the vendored split reproduce the published reference?
+
+The source repo (NJU-Star/RS-Tokenizer-10M-100x-0701) ships its own encode/decode outputs for
+eight crops. That is the only thing worth testing here -- the model code was vendored verbatim, so
+what can break is the SPLIT: wrong weight partition, a dropped tensor, the normalisation applied
+on one side and not the other.
+
+    python -m satlink.bench.test_rstok --ref /path/to/RS-Tokenizer-10M-100x-0701
+
+Without --ref the reference-dependent checks are skipped and reported as skipped, not as passes.
+
+Note on precision: the published token grids were produced in bf16. In bf16 this implementation
+reproduces them EXACTLY. In fp32 about 14% of indices differ -- a 196560-entry codebook in 24
+dimensions has very dense nearest neighbours, so the argmin flips readily -- while the resulting
+image moves by ~0.2/255. That is why the code check is pinned to bf16 and the pixel check is not.
+"""
+import argparse
+import glob
+import os
+import sys
+
+import torch
+
+from ..frame import VQFrame
+from ..protocol import LinkProfile
+
+N_TOK, BITS = 256, 18
+
+
+def _psnr(a, b):
+    m = ((a - b) ** 2).mean()
+    return float("inf") if m <= 0 else float(10 * torch.log10(255.0 ** 2 / m))
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--ref", default=None, help="clone of the source repo, for reference outputs")
+    ap.add_argument("--sat", default="satlink/weights/rstok_sat.pt")
+    ap.add_argument("--ground", default="satlink/weights/rstok_ground.pt")
+    ap.add_argument("--device", default="auto")
+    a = ap.parse_args()
+    from ..device import resolve
+    dev = resolve(a.device)
+    fails = []
+
+    # ---------------------------------------------------------------- profile
+    p = LinkProfile(1024, 768, n_tok=N_TOK, bits=BITS)
+    ok = p.crop_bytes == 576 and p.n_tiles == 12
+    print(f"  {'PASS' if ok else 'FAIL'} profile: {p.n_tok}x{p.bits}b -> {p.crop_bytes} B/crop, "
+          f"{p.n_tiles} crops, ratio {p.width * p.height / p.frame_bytes:.1f}x")
+    fails += [] if ok else ["profile"]
+
+    # 256 x 18 = 4608 bits is a whole number of bytes, so the pack must not add padding.
+    from ..protocol import pack_indices, unpack_indices
+    idx = [(i * 7919) % 196560 for i in range(N_TOK)]
+    w = pack_indices(idx, BITS)
+    ok = len(w) == 576 and unpack_indices(w, N_TOK, BITS) == idx
+    print(f"  {'PASS' if ok else 'FAIL'} 18-bit packing round-trips with no padding ({len(w)} B)")
+    fails += [] if ok else ["packing"]
+
+    # An index at the top of the codebook must survive; one past the 18-bit ceiling must not.
+    try:
+        pack_indices([196559], BITS)
+        pack_indices([1 << BITS], BITS)
+        print("  FAIL out-of-range index accepted")
+        fails.append("range")
+    except ValueError:
+        print("  PASS codebook bounds enforced (196559 fits, 2^18 rejected)")
+
+    # ---------------------------------------------------------------- weights
+    from ..sat.rstok import RSTokEncoder
+    sat = RSTokEncoder(a.sat, LinkProfile(256, 256, n_tok=N_TOK, bits=BITS),
+                       device=dev, dtype=torch.bfloat16)
+    n_sat = sum(x.numel() for x in sat.parameters())
+    ok = 13e6 < n_sat < 15e6
+    print(f"  {'PASS' if ok else 'FAIL'} satellite half is {n_sat/1e6:.1f}M params (encoder+VQ only)")
+    fails += [] if ok else ["sat-size"]
+
+    # The satellite file must not carry decoder weights -- that is the whole point of the split.
+    sd = torch.load(a.sat, map_location="cpu", weights_only=False)["state_dict"]
+    stray = [k for k in sd if k.startswith(("decoder.", "post_quant_embed."))]
+    print(f"  {'PASS' if not stray else 'FAIL'} satellite checkpoint holds no decoder tensors"
+          + (f" (found {stray[:3]})" if stray else ""))
+    fails += [] if not stray else ["sat-leak"]
+
+    # n_tok is not a rate knob for this tokenizer; truncating must be refused, not silently done.
+    try:
+        RSTokEncoder(a.sat, LinkProfile(256, 256, n_tok=128, bits=BITS), device=dev)
+        print("  FAIL truncated n_tok accepted (these tokens are a grid, not a prefix)")
+        fails.append("truncate")
+    except ValueError:
+        print("  PASS truncated n_tok refused")
+
+    # ---------------------------------------------------------------- reference
+    if not a.ref:
+        print("  SKIP reference checks: pass --ref <clone of the source repo>")
+    else:
+        import numpy as np
+        from PIL import Image
+        from ..ground.rstok_ground import RSTokDecoder
+        base = os.path.join(a.ref, "outputs_center_crop_no_enhance_check")
+        ins = sorted(glob.glob(os.path.join(base, "encoded", "input", "*_input.png")))
+        if not ins:
+            print(f"  FAIL no reference inputs under {base}")
+            fails.append("ref-missing")
+        else:
+            dec = RSTokDecoder(a.ground).to(dev)
+            prof = LinkProfile(256, 256, n_tok=N_TOK, bits=BITS)
+            agree, dpix = [], []
+            for f in ins:
+                stem = os.path.basename(f)[:-len("_input.png")]
+                im = Image.open(f).convert("RGB")
+                ref = np.load(os.path.join(base, "encoded", "tokens",
+                                           f"{stem}_tokens.npy")).astype(np.int64).reshape(-1)
+                mine = sat.encode(im).tokens[0].numpy()
+                agree.append((mine == ref).mean() * 100)
+                rr = torch.tensor(np.asarray(Image.open(
+                    os.path.join(base, "decoded", "recon", f"{stem}_recon.png")
+                ).convert("L")).astype("float32"))
+                got = dec(VQFrame(torch.tensor(ref).view(1, -1), prof).to(dev))[0].cpu() * 255
+                dpix.append(float((rr - got).abs().mean()))
+            m = float(np.mean(agree))
+            ok = m == 100.0
+            print(f"  {'PASS' if ok else 'FAIL'} bf16 codes match the published grids exactly "
+                  f"({m:.2f}% over {len(ins)} crops)")
+            fails += [] if ok else ["ref-codes"]
+            # Decoding the REFERENCE indices must reproduce the reference image; this isolates the
+            # decoder split from any encoder drift.
+            worst = max(dpix)
+            ok = worst < 3.0
+            print(f"  {'PASS' if ok else 'FAIL'} reference indices decode to the reference image "
+                  f"(worst {worst:.2f}/255, mean {sum(dpix)/len(dpix):.2f})")
+            fails += [] if ok else ["ref-pixels"]
+
+    # ---------------------------------------------------------------- precision
+    # bf16 is not safe for this tokenizer and the test says so in numbers rather than in a comment.
+    if a.ref:
+        import numpy as np
+        from PIL import Image
+        from ..ground.rstok_ground import RSTokDecoder
+        base = os.path.join(a.ref, "outputs_center_crop_no_enhance_check")
+        ins = sorted(glob.glob(os.path.join(base, "encoded", "input", "*_input.png")))
+        if ins:
+            dec = locals().get("dec") or RSTokDecoder(a.ground).to(dev)
+            prof = LinkProfile(256, 256, n_tok=N_TOK, bits=BITS)
+            worst = {}
+            for dt, nm in ((torch.bfloat16, "bf16"), (torch.float16, "fp16")):
+                e = RSTokEncoder(a.sat, prof, device=dev, dtype=dt)
+                lo = 1e9
+                for f in ins:
+                    im = Image.open(f).convert("RGB")
+                    src = torch.tensor(np.asarray(im.convert("L")).astype("float32"))
+                    r = dec(e.encode(im).to(dev))[0].cpu() * 255
+                    lo = min(lo, _psnr(src, r))
+                worst[nm] = lo
+            ok = worst["fp16"] > 35.0
+            print(f"  {'PASS' if ok else 'FAIL'} fp16 holds up (worst crop "
+                  f"{worst['fp16']:.1f} dB); bf16 worst is {worst['bf16']:.1f} dB "
+                  f"-- bf16 is unsafe for a 196560-entry codebook, use fp16")
+            fails += [] if ok else ["fp16-quality"]
+
+    print()
+    print("OK" if not fails else f"{len(fails)} FAILURE(S): {fails}")
+    return 0 if not fails else 1
+
+
+if __name__ == "__main__":
+    sys.exit(main())

satlink/examples/quickstart.py

@@ -42,8 +42,10 @@ def main():
     print(f"      落盘 {codes}")
 
     # ---------------------------------------------------------------- 链路
-    # 下行只有码字。帧尺寸不在流里，是两端约定的常量，所以这里要显式给出。
-    frame = satlink.frame_from_wire(wire, w, h, n_tok=frame.profile.n_tok)
+    # 下行只有码字。帧尺寸和 token 几何都不在流里，是两端约定的常量，所以这里要显式给出。
+    # 这里直接把 profile 传过去，它把四个常量一起带上——只对齐 n_tok 而漏掉 bits，
+    # 正是会解出错误长度的那种失误。
+    frame = satlink.frame_from_wire(wire, profile=frame.profile)
 
     # ---------------------------------------------------------------- 地面
     recon = os.path.join(a.out, "recon.png")

satlink/ground/modeling_vqvlm.py

@@ -31,6 +31,23 @@ import torch.nn.functional as F
 
 from ..protocol import LinkProfile
 from .decoder import GroundDecoder
+
+
+def _build_decoder(weights, profile):
+    """-> the ground tokenizer matching this checkpoint.
+
+    Two backends are supported and they are told apart by the checkpoint's own config rather than
+    by a flag the caller has to keep in sync: One-D-Piece (48 tokens x 12 bits per crop) and the
+    RS-Tokenizer (256 x 18). Both expose the same forward(VQFrame) -> [3, H, W], so everything
+    downstream -- patchify, the vision tower, the LLM -- is unchanged. Only the pixel path is
+    wired up for the RS-Tokenizer; the latent path needs a bridge trained against its particular
+    partial decode, and no such bridge exists."""
+    import torch as _torch
+    cfg = _torch.load(weights, map_location="cpu", weights_only=False).get("config", {})
+    if "token_grid" in cfg or "quantizer" in cfg:
+        from .rstok_ground import RSTokDecoder
+        return RSTokDecoder(weights, profile)
+    return GroundDecoder(weights, profile)
 from ..frame import VQFrame
 from .glm4v_text import Glm4vTextModel
 from .glm4v_vision import Glm4vVisionModel
@@ -66,7 +83,7 @@ class VQGlm4vForConditionalGeneration(nn.Module):
         # The tokenizer decoder is a submodule, not a neighbour. It is created here so that
         # .to(), .eval(), state_dict() and autocast all cover it -- the ways a bolt-on module
         # silently diverges from the model it feeds.
-        self.opd = GroundDecoder(decoder_weights, profile) if decoder_weights else None
+        self.opd = _build_decoder(decoder_weights, profile) if decoder_weights else None
         self.register_buffer("_img_mean", torch.tensor(IMAGE_MEAN).view(3, 1, 1), persistent=False)
         self.register_buffer("_img_std", torch.tensor(IMAGE_STD).view(3, 1, 1), persistent=False)
 

satlink/ground/rstok_decoder.py

@@ -0,0 +1,221 @@
+"""RS-Tokenizer (NJU-Star/RS-Tokenizer-10M-100x-0701) -- DECODER + codebook.
+
+The ground half of satlink/sat/rstok_encoder.py. Vendored from the same infer.py: the shared
+transformer blocks, LlamaTransformerDecoder, and the VectorQuantizer (needed only for its
+codebook lookup -- the ground never runs an argmin).
+
+336M parameters against the encoder's 13M, which is the whole reason the split is worth making.
+
+VENDORING EDITS: same as the satellite file -- unused third-party imports dropped, encoder class
+removed. The maths is untouched.
+"""
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+def build_2d_rope_cache(grid_h, grid_w, head_dim, theta=10000.0, device=None, dtype=torch.float32):
+    if head_dim % 4 != 0:
+        raise ValueError(f"head_dim must be divisible by 4, got {head_dim}")
+    dim_half = head_dim // 2
+    dim_quarter = head_dim // 4
+    inv_freq = 1.0 / (
+        theta ** (
+            torch.arange(0, dim_quarter, device=device, dtype=torch.float32) * 2 / dim_half
+        )
+    )
+    h_pos = torch.arange(grid_h, device=device, dtype=torch.float32)
+    w_pos = torch.arange(grid_w, device=device, dtype=torch.float32)
+    angles_h = torch.outer(h_pos, inv_freq)
+    angles_w = torch.outer(w_pos, inv_freq)
+    angles_h = angles_h[:, None, :].expand(grid_h, grid_w, dim_quarter).reshape(-1, dim_quarter)
+    angles_w = angles_w[None, :, :].expand(grid_h, grid_w, dim_quarter).reshape(-1, dim_quarter)
+    emb_h = torch.cat([angles_h, angles_h], dim=-1)
+    emb_w = torch.cat([angles_w, angles_w], dim=-1)
+    cos = torch.cat([emb_h.cos(), emb_w.cos()], dim=-1)
+    sin = torch.cat([emb_h.sin(), emb_w.sin()], dim=-1)
+    return cos.to(dtype=dtype), sin.to(dtype=dtype)
+
+
+def rotate_half(x):
+    x1, x2 = x.chunk(2, dim=-1)
+    return torch.cat([-x2, x1], dim=-1)
+
+
+def rotate_half_2d(x):
+    dim = x.shape[-1]
+    return torch.cat([
+        rotate_half(x[..., : dim // 2]),
+        rotate_half(x[..., dim // 2:]),
+    ], dim=-1)
+
+
+def apply_2d_rope(q, k, cos, sin):
+    cos = cos[None, None, :, :].to(q.dtype)
+    sin = sin[None, None, :, :].to(q.dtype)
+    return q * cos + rotate_half_2d(q) * sin, k * cos + rotate_half_2d(k) * sin
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim, eps=1e-6):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim))
+
+    def forward(self, x):
+        x_float = x.float()
+        normed = x_float * torch.rsqrt(x_float.pow(2).mean(-1, keepdim=True) + self.eps)
+        return normed.to(dtype=x.dtype) * self.weight.to(dtype=x.dtype)
+
+
+class SwiGLUMLP(nn.Module):
+    def __init__(self, dim, mlp_ratio=4.0):
+        super().__init__()
+        raw_hidden = int(dim * mlp_ratio * 2 / 3)
+        hidden = (raw_hidden + 63) // 64 * 64
+        self.gate_proj = nn.Linear(dim, hidden, bias=False)
+        self.up_proj = nn.Linear(dim, hidden, bias=False)
+        self.down_proj = nn.Linear(hidden, dim, bias=False)
+
+    def forward(self, x):
+        return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))
+
+
+class LlamaSelfAttention(nn.Module):
+    def __init__(self, dim, num_heads):
+        super().__init__()
+        if dim % num_heads != 0:
+            raise ValueError(f"dim {dim} must be divisible by heads {num_heads}")
+        self.num_heads = int(num_heads)
+        self.head_dim = int(dim) // int(num_heads)
+        self.dim = int(dim)
+        self.q_proj = nn.Linear(dim, dim, bias=False)
+        self.k_proj = nn.Linear(dim, dim, bias=False)
+        self.v_proj = nn.Linear(dim, dim, bias=False)
+        self.out_proj = nn.Linear(dim, dim, bias=False)
+
+    def forward(self, x, cos, sin):
+        batch, seqlen, _ = x.shape
+        q = self.q_proj(x).reshape(batch, seqlen, self.num_heads, self.head_dim).transpose(1, 2)
+        k = self.k_proj(x).reshape(batch, seqlen, self.num_heads, self.head_dim).transpose(1, 2)
+        v = self.v_proj(x).reshape(batch, seqlen, self.num_heads, self.head_dim).transpose(1, 2)
+        q, k = apply_2d_rope(q, k, cos, sin)
+        out = F.scaled_dot_product_attention(q, k, v, dropout_p=0.0, is_causal=False)
+        out = out.transpose(1, 2).reshape(batch, seqlen, self.dim)
+        return self.out_proj(out)
+
+
+class LlamaViTBlock(nn.Module):
+    def __init__(self, dim, num_heads, mlp_ratio=4.0):
+        super().__init__()
+        self.norm1 = RMSNorm(dim)
+        self.attn = LlamaSelfAttention(dim, num_heads)
+        self.norm2 = RMSNorm(dim)
+        self.mlp = SwiGLUMLP(dim, mlp_ratio)
+
+    def forward(self, x, cos, sin):
+        x = x + self.attn(self.norm1(x), cos, sin)
+        x = x + self.mlp(self.norm2(x))
+        return x
+
+
+class LlamaTransformerDecoder(nn.Module):
+    def __init__(
+        self,
+        image_size=256,
+        patch_size=16,
+        width=1024,
+        layers=24,
+        heads=16,
+        mlp_ratio=4,
+        drop_rate=0.0,
+        use_ffn_output=True,
+        dim_ffn_output=None,
+        **_,
+    ):
+        super().__init__()
+        del image_size, drop_rate
+        self.patch_size = int(patch_size)
+        self.width = int(width)
+        self.head_dim = self.width // int(heads)
+        self.rope_theta = 10000.0
+        self.blocks = nn.ModuleList([
+            LlamaViTBlock(self.width, heads, mlp_ratio) for _ in range(int(layers))
+        ])
+        self.norm = RMSNorm(self.width)
+        self.use_ffn_output = bool(use_ffn_output)
+        if dim_ffn_output is None:
+            dim_ffn_output = self.width * 4
+        if self.use_ffn_output:
+            self.ffn = nn.Sequential(nn.Linear(self.width, int(dim_ffn_output)), nn.Tanh())
+            conv_in_features = int(dim_ffn_output)
+        else:
+            self.ffn = nn.Identity()
+            conv_in_features = self.width
+        self.conv_out = nn.Linear(conv_in_features, 3 * self.patch_size * self.patch_size, bias=True)
+
+    def forward(self, x, grid_sizes):
+        if len(set(tuple(item) for item in grid_sizes)) != 1:
+            raise ValueError("Standalone decoder expects a uniform batch.")
+        gh, gw = tuple(int(v) for v in grid_sizes[0])
+        cos, sin = build_2d_rope_cache(
+            gh,
+            gw,
+            self.head_dim,
+            self.rope_theta,
+            device=x.device,
+            dtype=x.dtype,
+        )
+        for block in self.blocks:
+            x = block(x, cos, sin)
+        x = self.norm(x)
+        x = self.ffn(x)
+        x = self.conv_out(x)
+        batch = x.shape[0]
+        x = x.reshape(batch, gh, gw, 3, self.patch_size, self.patch_size)
+        x = x.permute(0, 3, 1, 4, 2, 5).reshape(batch, 3, gh * self.patch_size, gw * self.patch_size)
+        return x
+
+
+class VectorQuantizer(nn.Module):
+    def __init__(self, n_embed=196560, embed_dim=24, l2_norm=True, beta=0.25):
+        super().__init__()
+        self.n_embed = int(n_embed)
+        self.embed_dim = int(embed_dim)
+        self.l2_norm = bool(l2_norm)
+        self.beta = float(beta)
+        self.embedding = nn.Embedding(self.n_embed, self.embed_dim)
+        self.bits_per_index = int(math.ceil(math.log2(self.n_embed)))
+
+    def encode_indices(self, z):
+        if self.l2_norm:
+            z = F.normalize(z, dim=-1)
+            z_flat = z.reshape(-1, self.embed_dim)
+            embedding_weight = F.normalize(self.embedding.weight, dim=-1)
+            distance = -z_flat @ embedding_weight.t()
+        else:
+            z_flat = z.reshape(-1, self.embed_dim)
+            distance = (
+                torch.sum(z_flat ** 2, dim=1, keepdim=True)
+                + torch.sum(self.embedding.weight ** 2, dim=1)
+                - 2 * z_flat @ self.embedding.weight.t()
+            )
+        indices = torch.argmin(distance.detach(), dim=1)
+        return indices.view(z.shape[0], -1)
+
+    def decode_indices(self, indices):
+        z_q = self.embedding(indices.long())
+        if self.l2_norm:
+            z_q = F.normalize(z_q, dim=-1)
+        return z_q
+
+    def forward(self, z):
+        indices = self.encode_indices(z)
+        z_q = self.embedding(indices).view_as(z)
+        if self.l2_norm:
+            z_q = F.normalize(z_q, dim=-1)
+        loss = self.beta * torch.mean(((z_q.detach() - z) ** 2).sum(dim=-1))
+        z_q = z + (z_q - z).detach()
+        return z_q, loss, {"indices": indices}

satlink/ground/rstok_ground.py

@@ -0,0 +1,121 @@
+"""Ground side, RS-Tokenizer backend. Same interface as satlink.ground.decoder.GroundDecoder.
+
+    dec = RSTokDecoder("satlink/weights/rstok_ground.pt")
+    image = dec(frame)            # [3, H, W] in [0,1]
+
+340.9M parameters against the satellite's 13.9M -- the asymmetry the split exists to exploit.
+
+Like the One-D-Piece decoder this is an nn.Module, not a preprocessing script, so it can be owned
+by the VLM as a submodule rather than bolted on beside it.
+
+Deliberately no reversible-preprocess inverse here. The published pipeline offers a logit
+unsharp/affine enhancement whose inverse needs a per-crop `center` scalar, and the zero-header
+protocol has nowhere to put per-crop side information. Running the inverse WITHOUT that scalar
+would be worse than not running it at all -- it would apply the wrong affine and shift the whole
+crop -- so the encoder does not apply the forward transform either, and the two halves stay
+consistent. See the note in satlink/tools/export_rstok_weights.py for what that costs.
+"""
+import torch
+import torch.nn as nn
+
+from ..config import Cfg
+from ..frame import VQFrame
+from ..protocol import LinkProfile
+
+
+class RSTokDecoder(nn.Module):
+    """RS-Tokenizer decoder + codebook."""
+
+    def __init__(self, weights, profile: LinkProfile = None, dtype=torch.float32, n_layers=None):
+        super().__init__()
+        from .rstok_decoder import LlamaTransformerDecoder, VectorQuantizer
+        ck = torch.load(weights, map_location="cpu", weights_only=False)
+        cfg = Cfg(ck["config"])
+        self.profile = profile
+        self.patch_size, self.image_size = cfg.patch_size, cfg.image_size
+        self.grid = tuple(cfg.token_grid)
+        q, d = cfg.quantizer, cfg.decoder
+        self.quantize = VectorQuantizer(n_embed=q.n_embed, embed_dim=q.embed_dim,
+                                        l2_norm=q.l2_norm)
+        self.post_quant_embed = nn.Linear(q.embed_dim, d.width)
+        self.decoder = LlamaTransformerDecoder(patch_size=cfg.patch_size, width=d.width,
+                                               layers=d.layers, heads=d.heads,
+                                               mlp_ratio=d.mlp_ratio)
+        got = self.load_state_dict(ck["state_dict"], strict=False)
+        if got.missing_keys or got.unexpected_keys:
+            raise RuntimeError(f"checkpoint mismatch: missing={got.missing_keys[:5]} "
+                               f"unexpected={got.unexpected_keys[:5]}")
+        if n_layers is not None:
+            # Truncating costs quality and is only meaningful for experiments; the full stack is
+            # the deployment configuration.
+            self.decoder.blocks = self.decoder.blocks[:n_layers]
+        if dtype == torch.float16:
+            # Measured, not defensive: at 1024 width over 24 layers this decoder's activations
+            # leave float16's range and the whole frame comes back NaN (which then casts to a
+            # black image). bfloat16 has the exponent for it and matches float32 to 0.01 dB.
+            raise ValueError("RSTokDecoder overflows in float16; use float32 or bfloat16")
+        self.zero_mean = bool(cfg.get("zero_mean", True))
+        self.eval()
+        for p in self.parameters():
+            p.requires_grad_(False)
+        self.to(dtype)
+
+    def codebook_device(self):
+        return self.quantize.embedding.weight.device
+
+    # ------------------------------------------------------------------ tiles
+    @torch.no_grad()
+    def decode_tiles(self, indices, chunk=8):
+        """[n_tiles, 256] int -> [n_tiles, 3, 256, 256] float in [0,1].
+
+        Chunked because the decoder is 341M parameters at 1024 width over 256 tokens; a whole
+        survey frame's worth of crops at once is a needless memory spike when the crops are
+        independent anyway."""
+        dev = self.codebook_device()
+        dt = self.post_quant_embed.weight.dtype
+        gh, gw = self.grid
+        outs = []
+        for i in range(0, indices.shape[0], chunk):
+            idx = indices[i:i + chunk].to(dev).long()
+            z = self.quantize.decode_indices(idx).to(dt)
+            h = self.post_quant_embed(z)
+            x = self.decoder(h, grid_sizes=[(gh, gw)] * idx.shape[0]).clamp(-1.0, 1.0)
+            outs.append(self._denormalize(x.float()))
+        return torch.cat(outs)
+
+    def _denormalize(self, x):
+        """model output -> [0,1]."""
+        return ((x + 1.0) / 2.0 if self.zero_mean else x).clamp(0, 1)
+
+    # ------------------------------------------------------------------ stitch
+    @torch.no_grad()
+    def stitch(self, tiles, profile):
+        """[n_tiles, 3, tile, tile] -> [3, H, W], averaging the overlap.
+
+        Averaging rather than last-writer-wins: crops are coded independently, so the seam where
+        two reconstructions meet is the most visible artifact of the scheme.
+
+        Channels are collapsed by mean and replicated, as in the One-D-Piece decoder. The source
+        imagery is strictly grayscale (the published reference inputs have zero channel spread)
+        while the decoder's three outputs disagree by 3-11/255 -- they are three noisy estimates of
+        one luminance, so averaging them is a denoiser, not a loss of colour."""
+        p = profile
+        t = p.tile
+        xs, ys = p.starts(0), p.starts(1)
+        cw, ch = max(p.width, max(xs) + t), max(p.height, max(ys) + t)
+        acc = torch.zeros(ch, cw, dtype=torch.float32, device=tiles.device)
+        cnt = torch.zeros(ch, cw, dtype=torch.float32, device=tiles.device)
+        k = 0
+        for y in ys:
+            for x in xs:
+                acc[y:y + t, x:x + t] += tiles[k].float().mean(0)
+                cnt[y:y + t, x:x + t] += 1.0
+                k += 1
+        out = (acc / cnt.clamp(min=1)).clamp(0, 1)[:p.height, :p.width]
+        return out[None].expand(3, -1, -1).contiguous()
+
+    # ------------------------------------------------------------------ api
+    @torch.no_grad()
+    def forward(self, frame: VQFrame):
+        """VQFrame -> [3, H, W]. Geometry comes from the frame, so one instance serves any size."""
+        return self.stitch(self.decode_tiles(frame.tokens), frame.profile)

satlink/sat/rstok.py

@@ -0,0 +1,138 @@
+"""Satellite side, RS-Tokenizer backend. Same interface as satlink.sat.encoder.SatEncoder.
+
+    prof = LinkProfile(1024, 768, n_tok=256, bits=18)      # agreed, not transmitted
+    sat  = RSTokEncoder("satlink/weights/rstok_sat.pt", prof)
+    frame = sat.encode(image)
+
+Two things differ from the One-D-Piece backend and both are visible in the profile rather than
+hidden in the model:
+
+  n_tok=256, bits=18   a 16x16 grid into a 196560-entry codebook -> 576 B/crop, ~114x
+  RGB input            One-D-Piece was trained on grayscale and this one was not, so the frame is
+                       fed as 3 channels. A grayscale source is replicated across them, which is
+                       what the published reference does too (it opens every input as RGB).
+
+DO NOT RUN THIS BACKEND IN bfloat16. The codebook has 196560 entries in 24 dimensions, all L2
+normalised, so neighbouring codewords are separated by very small cosine gaps; bf16's 8-bit
+mantissa is not enough to resolve them and the argmin lands on the wrong codeword often enough to
+destroy a crop. Measured over the eight published reference crops: seven are unaffected, and one
+drops from 41.4 dB (fp16) to 22.5 dB (bf16) -- a 19 dB collapse on 1 image in 8, not a gentle
+average degradation. fp16 tracks fp32 within 0.5 dB everywhere, so fp16 is the default and is what
+the acceptance runs use. (The upstream reference outputs were generated in bf16 and therefore
+carry this collapse; that is why satlink/bench/test_rstok.py pins its exact-match check to bf16 --
+matching the reference is a provenance check, not an endorsement of the dtype.)
+
+n_tok is NOT a rate knob here. One-D-Piece is trained with tail-token drop, so any prefix of its
+48 tokens is a valid encoding; this tokenizer's 256 tokens are a 2-D grid with no importance
+ordering, and dropping any of them removes a patch of the image rather than some detail. So n_tok
+is fixed at 256 and the constructor refuses anything else, instead of silently returning a frame
+the decoder cannot use.
+"""
+import torch
+
+from ..config import Cfg
+from ..device import autocast, resolve
+from ..frame import VQFrame
+from ..protocol import LinkProfile, pack_frame
+from .rstok_encoder import LlamaTransformerEncoder, VectorQuantizer
+
+N_TOK = 256          # 16x16 grid
+BITS = 18            # ceil(log2(196560)); entropy floor is 17.585, so 0.42 bits/index is spent
+
+
+class RSTokEncoder(torch.nn.Module):
+    """RS-Tokenizer encoder + VQ. 13.9M parameters; the 340.9M decoder stays on the ground."""
+
+    def __init__(self, weights, profile: LinkProfile, device="auto", dtype=torch.float16):
+        super().__init__()
+        ck = torch.load(weights, map_location="cpu", weights_only=False)
+        cfg = Cfg(ck["config"])
+        if profile.n_tok != N_TOK or profile.bits != BITS:
+            raise ValueError(f"RS-Tokenizer needs n_tok={N_TOK}, bits={BITS}; profile says "
+                             f"n_tok={profile.n_tok}, bits={profile.bits}. Unlike One-D-Piece "
+                             f"these tokens are a spatial grid, not an importance-ordered "
+                             f"sequence, so they cannot be truncated.")
+        if profile.tile != cfg.image_size:
+            raise ValueError(f"tile {profile.tile} but the tokenizer is trained at "
+                             f"{cfg.image_size}")
+        self.profile = profile
+        e = cfg.encoder
+        self.encoder = LlamaTransformerEncoder(image_size=cfg.image_size,
+                                               patch_size=cfg.patch_size, width=e.width,
+                                               layers=e.layers, heads=e.heads,
+                                               mlp_ratio=e.mlp_ratio)
+        q = cfg.quantizer
+        self.quant_embed = torch.nn.Linear(e.width, q.embed_dim)
+        self.quantize = VectorQuantizer(n_embed=q.n_embed, embed_dim=q.embed_dim,
+                                        l2_norm=q.l2_norm)
+        got = self.load_state_dict(ck["state_dict"], strict=False)
+        if got.missing_keys or got.unexpected_keys:
+            raise RuntimeError(f"checkpoint mismatch: missing={got.missing_keys[:5]} "
+                               f"unexpected={got.unexpected_keys[:5]}")
+        self.zero_mean = bool(cfg.get("zero_mean", True))
+        device = resolve(device)
+        self.eval().to(device=device)
+        self.device, self.dtype = device, dtype
+        for p in self.parameters():
+            p.requires_grad_(False)
+
+    # ---------------------------------------------------------------- tiling
+    def tiles(self, img):
+        """-> [n_tiles, 3, tile, tile] in profile raster order, values in [0,1]."""
+        from .encoder import _to_tensor
+        p = self.profile
+        x = img if torch.is_tensor(img) else _to_tensor(img)
+        x = x.to(self.device, non_blocking=True)
+        if x.ndim == 2:
+            x = x[None]
+        if x.shape[0] == 1:
+            x = x.expand(3, -1, -1)
+        elif x.shape[0] != 3:
+            raise ValueError(f"expected 1 or 3 channels, got {x.shape[0]}")
+        _, H, W = x.shape
+        if (W, H) != (p.width, p.height):
+            raise ValueError(f"frame is {W}x{H}, profile says {p.width}x{p.height}")
+        t = p.tile
+        if H < t or W < t:
+            x = torch.nn.functional.pad(x[None], (0, max(0, t - W), 0, max(0, t - H)),
+                                        mode="replicate")[0]
+        xs, ys = p.starts(0), p.starts(1)
+        return torch.stack([x[:, y:y + t, x0:x0 + t] for y in ys for x0 in xs]).contiguous()
+
+    def _normalize(self, x01):
+        """[0,1] -> what the model was trained on. zero_mean maps to [-1,1]."""
+        return x01 * 2.0 - 1.0 if self.zero_mean else x01
+
+    # ---------------------------------------------------------------- encode
+    @torch.no_grad()
+    def indices_batch(self, crops, chunk=64):
+        """-> LongTensor [B, 256] for crops already cut: [B, 3, 256, 256] in [0,1]."""
+        t = self.profile.tile
+        if crops.ndim != 4 or crops.shape[1] != 3 or crops.shape[-2:] != (t, t):
+            raise ValueError(f"expected [B, 3, {t}, {t}] crops, got {tuple(crops.shape)}")
+        outs = []
+        for i in range(0, crops.shape[0], chunk):
+            b = self._normalize(crops[i:i + chunk].to(self.device, non_blocking=True))
+            with autocast(self.device, self.dtype):
+                h, _ = self.encoder(b)
+                h = self.quant_embed(h)
+                idx = self.quantize.encode_indices(h)
+            outs.append(idx.reshape(b.shape[0], -1))
+        return torch.cat(outs)
+
+    @torch.no_grad()
+    def indices(self, img, chunk=64):
+        return self.indices_batch(self.tiles(img), chunk=chunk)
+
+    @torch.no_grad()
+    def encode(self, img, chunk=64):
+        return VQFrame(self.indices(img, chunk=chunk).cpu(), self.profile)
+
+    @torch.no_grad()
+    def encode_wire(self, img, chunk=64):
+        return self.pack(self.encode(img, chunk=chunk))
+
+    def pack(self, frame: VQFrame):
+        if frame.profile.crop_bytes != self.profile.crop_bytes:
+            raise ValueError("frame profile does not match this encoder's profile")
+        return pack_frame(frame.profile, frame.tolist())

satlink/sat/rstok_encoder.py

@@ -0,0 +1,208 @@
+"""RS-Tokenizer (NJU-Star/RS-Tokenizer-10M-100x-0701) -- ENCODER + VQ only.
+
+Vendored from that repo's infer.py, truncated so this file physically cannot contain a decoder:
+the shared transformer blocks, LlamaTransformerEncoder and VectorQuantizer are here;
+LlamaTransformerDecoder is not, and lives in satlink/ground/rstok_decoder.py. Same discipline as
+opd_encoder.py -- an import-path check would still pass while dragging 336M parameters into orbit,
+so the decoder is simply absent from the file.
+
+Shape of the thing: a 256x256 crop becomes a 16x16 grid of indices into a 196560-entry codebook,
+so 256 tokens x 18 bits = 576 bytes per crop, about 114x smaller than the 65536 raw grayscale
+bytes. That is 8x more downlink than One-D-Piece's 72 bytes, bought with roughly 20 dB of
+reconstruction PSNR (~41 dB against ~21 dB on the same crops).
+
+196560 is the kissing number of the Leech lattice and embed_dim is 24 -- this is a 24-dimensional
+lattice quantiser, which is why the codebook is not a power of two. 18 bits per index therefore
+wastes 0.42 bits against the 17.585-bit entropy floor; packing at 18 is still exact at the crop
+boundary (256 x 18 = 4608 bits = 576 bytes) so no padding is spent.
+
+VENDORING EDITS, all mechanical:
+  - numpy / omegaconf / PIL / safetensors / torchvision imports dropped; unused here
+  - the decoder class and the combined StandaloneL24SQTokenizer removed
+Nothing in the maths is touched; satlink/bench/test_rstok.py checks the codes against the
+reference token grids published in the source repo.
+"""
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+def build_2d_rope_cache(grid_h, grid_w, head_dim, theta=10000.0, device=None, dtype=torch.float32):
+    if head_dim % 4 != 0:
+        raise ValueError(f"head_dim must be divisible by 4, got {head_dim}")
+    dim_half = head_dim // 2
+    dim_quarter = head_dim // 4
+    inv_freq = 1.0 / (
+        theta ** (
+            torch.arange(0, dim_quarter, device=device, dtype=torch.float32) * 2 / dim_half
+        )
+    )
+    h_pos = torch.arange(grid_h, device=device, dtype=torch.float32)
+    w_pos = torch.arange(grid_w, device=device, dtype=torch.float32)
+    angles_h = torch.outer(h_pos, inv_freq)
+    angles_w = torch.outer(w_pos, inv_freq)
+    angles_h = angles_h[:, None, :].expand(grid_h, grid_w, dim_quarter).reshape(-1, dim_quarter)
+    angles_w = angles_w[None, :, :].expand(grid_h, grid_w, dim_quarter).reshape(-1, dim_quarter)
+    emb_h = torch.cat([angles_h, angles_h], dim=-1)
+    emb_w = torch.cat([angles_w, angles_w], dim=-1)
+    cos = torch.cat([emb_h.cos(), emb_w.cos()], dim=-1)
+    sin = torch.cat([emb_h.sin(), emb_w.sin()], dim=-1)
+    return cos.to(dtype=dtype), sin.to(dtype=dtype)
+
+
+def rotate_half(x):
+    x1, x2 = x.chunk(2, dim=-1)
+    return torch.cat([-x2, x1], dim=-1)
+
+
+def rotate_half_2d(x):
+    dim = x.shape[-1]
+    return torch.cat([
+        rotate_half(x[..., : dim // 2]),
+        rotate_half(x[..., dim // 2:]),
+    ], dim=-1)
+
+
+def apply_2d_rope(q, k, cos, sin):
+    cos = cos[None, None, :, :].to(q.dtype)
+    sin = sin[None, None, :, :].to(q.dtype)
+    return q * cos + rotate_half_2d(q) * sin, k * cos + rotate_half_2d(k) * sin
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim, eps=1e-6):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim))
+
+    def forward(self, x):
+        x_float = x.float()
+        normed = x_float * torch.rsqrt(x_float.pow(2).mean(-1, keepdim=True) + self.eps)
+        return normed.to(dtype=x.dtype) * self.weight.to(dtype=x.dtype)
+
+
+class SwiGLUMLP(nn.Module):
+    def __init__(self, dim, mlp_ratio=4.0):
+        super().__init__()
+        raw_hidden = int(dim * mlp_ratio * 2 / 3)
+        hidden = (raw_hidden + 63) // 64 * 64
+        self.gate_proj = nn.Linear(dim, hidden, bias=False)
+        self.up_proj = nn.Linear(dim, hidden, bias=False)
+        self.down_proj = nn.Linear(hidden, dim, bias=False)
+
+    def forward(self, x):
+        return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))
+
+
+class LlamaSelfAttention(nn.Module):
+    def __init__(self, dim, num_heads):
+        super().__init__()
+        if dim % num_heads != 0:
+            raise ValueError(f"dim {dim} must be divisible by heads {num_heads}")
+        self.num_heads = int(num_heads)
+        self.head_dim = int(dim) // int(num_heads)
+        self.dim = int(dim)
+        self.q_proj = nn.Linear(dim, dim, bias=False)
+        self.k_proj = nn.Linear(dim, dim, bias=False)
+        self.v_proj = nn.Linear(dim, dim, bias=False)
+        self.out_proj = nn.Linear(dim, dim, bias=False)
+
+    def forward(self, x, cos, sin):
+        batch, seqlen, _ = x.shape
+        q = self.q_proj(x).reshape(batch, seqlen, self.num_heads, self.head_dim).transpose(1, 2)
+        k = self.k_proj(x).reshape(batch, seqlen, self.num_heads, self.head_dim).transpose(1, 2)
+        v = self.v_proj(x).reshape(batch, seqlen, self.num_heads, self.head_dim).transpose(1, 2)
+        q, k = apply_2d_rope(q, k, cos, sin)
+        out = F.scaled_dot_product_attention(q, k, v, dropout_p=0.0, is_causal=False)
+        out = out.transpose(1, 2).reshape(batch, seqlen, self.dim)
+        return self.out_proj(out)
+
+
+class LlamaViTBlock(nn.Module):
+    def __init__(self, dim, num_heads, mlp_ratio=4.0):
+        super().__init__()
+        self.norm1 = RMSNorm(dim)
+        self.attn = LlamaSelfAttention(dim, num_heads)
+        self.norm2 = RMSNorm(dim)
+        self.mlp = SwiGLUMLP(dim, mlp_ratio)
+
+    def forward(self, x, cos, sin):
+        x = x + self.attn(self.norm1(x), cos, sin)
+        x = x + self.mlp(self.norm2(x))
+        return x
+
+
+class LlamaTransformerEncoder(nn.Module):
+    def __init__(self, image_size=256, patch_size=16, width=384, layers=5, heads=6, mlp_ratio=4, drop_rate=0.0, **_):
+        super().__init__()
+        del drop_rate
+        self.image_size = int(image_size)
+        self.patch_size = int(patch_size)
+        self.width = int(width)
+        self.head_dim = self.width // int(heads)
+        self.rope_theta = 10000.0
+        self.patch_embed = nn.Conv2d(3, self.width, kernel_size=self.patch_size, stride=self.patch_size, bias=True)
+        self.blocks = nn.ModuleList([
+            LlamaViTBlock(self.width, heads, mlp_ratio) for _ in range(int(layers))
+        ])
+        self.norm = RMSNorm(self.width)
+
+    def forward(self, x):
+        feat = self.patch_embed(x)
+        batch, channels, gh, gw = feat.shape
+        tokens = feat.flatten(2).transpose(1, 2)
+        cos, sin = build_2d_rope_cache(
+            gh,
+            gw,
+            self.head_dim,
+            self.rope_theta,
+            device=tokens.device,
+            dtype=tokens.dtype,
+        )
+        for block in self.blocks:
+            tokens = block(tokens, cos, sin)
+        return self.norm(tokens), [(gh, gw)] * batch
+
+
+class VectorQuantizer(nn.Module):
+    def __init__(self, n_embed=196560, embed_dim=24, l2_norm=True, beta=0.25):
+        super().__init__()
+        self.n_embed = int(n_embed)
+        self.embed_dim = int(embed_dim)
+        self.l2_norm = bool(l2_norm)
+        self.beta = float(beta)
+        self.embedding = nn.Embedding(self.n_embed, self.embed_dim)
+        self.bits_per_index = int(math.ceil(math.log2(self.n_embed)))
+
+    def encode_indices(self, z):
+        if self.l2_norm:
+            z = F.normalize(z, dim=-1)
+            z_flat = z.reshape(-1, self.embed_dim)
+            embedding_weight = F.normalize(self.embedding.weight, dim=-1)
+            distance = -z_flat @ embedding_weight.t()
+        else:
+            z_flat = z.reshape(-1, self.embed_dim)
+            distance = (
+                torch.sum(z_flat ** 2, dim=1, keepdim=True)
+                + torch.sum(self.embedding.weight ** 2, dim=1)
+                - 2 * z_flat @ self.embedding.weight.t()
+            )
+        indices = torch.argmin(distance.detach(), dim=1)
+        return indices.view(z.shape[0], -1)
+
+    def decode_indices(self, indices):
+        z_q = self.embedding(indices.long())
+        if self.l2_norm:
+            z_q = F.normalize(z_q, dim=-1)
+        return z_q
+
+    def forward(self, z):
+        indices = self.encode_indices(z)
+        z_q = self.embedding(indices).view_as(z)
+        if self.l2_norm:
+            z_q = F.normalize(z_q, dim=-1)
+        loss = self.beta * torch.mean(((z_q.detach() - z) ** 2).sum(dim=-1))
+        z_q = z + (z_q - z).detach()
+        return z_q, loss, {"indices": indices}

satlink/tools/export_rstok_weights.py

@@ -0,0 +1,70 @@
+"""Convert the published RS-Tokenizer safetensors into the two .pt files satlink loads.
+
+    python -m satlink.tools.export_rstok_weights --src /tmp/rstok --out satlink/weights
+
+Why convert at all: satlink/sat is held to a torch-only dependency budget (checked by
+satlink/bench/test_isolation.py), and reading .safetensors would add safetensors to the flight
+bundle for one load call. The ground has no such constraint but is converted too, so both halves
+load the same way.
+
+The split mirrors the source repo's own: encoder + quant_embed + codebook go to the satellite,
+post_quant_embed + decoder + codebook go to the ground. The codebook appears on BOTH sides on
+purpose -- the satellite needs it to run the argmin, the ground needs it to look indices back up,
+and it is 196560 x 24 fp32 = 18.9 MB, small next to either half.
+"""
+import argparse
+import os
+import sys
+
+import torch
+
+SAT_PREFIXES = ("encoder.", "quant_embed.", "quantize.")
+GROUND_PREFIXES = ("post_quant_embed.", "decoder.", "quantize.")
+
+# decoder.distill_mlp.* is a 25.2M-parameter auxiliary head used during training only. The
+# published inference script silently drops it (it lands in load_partition's skipped_unexpected
+# bucket), so it is dropped HERE instead -- at export, once, visibly -- rather than by relaxing
+# the runtime's state-dict check, which would then also swallow a genuine mismatch.
+DROP_PREFIXES = ("decoder.distill_mlp.",)
+
+# From the source repo's config.yaml. Recorded in the checkpoint so the runtime does not need the
+# yaml (and so omegaconf stays out of the deployment bundle).
+CONFIG = {"image_size": 256, "patch_size": 16, "token_grid": [16, 16],
+          "encoder": {"width": 384, "layers": 5, "heads": 6, "mlp_ratio": 4},
+          "decoder": {"width": 1024, "layers": 24, "heads": 16, "mlp_ratio": 4},
+          "quantizer": {"n_embed": 196560, "embed_dim": 24, "l2_norm": True},
+          "zero_mean": True,
+          # The user's deployment choice: no reversible preprocess. The enhanced variant needs a
+          # per-crop `center` scalar as side information, which the zero-header protocol has no
+          # place for; without it the average cost is 2.5 dB, though one of the eight published
+          # reference crops drops 42.5 -> 22.7 dB, so this is a tail risk, not a uniform 2.5 dB.
+          "reversible_preprocess": "none"}
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--src", required=True, help="a clone of NJU-Star/RS-Tokenizer-10M-100x-0701")
+    ap.add_argument("--out", default="satlink/weights")
+    a = ap.parse_args()
+    from safetensors.torch import load_file
+
+    enc = load_file(os.path.join(a.src, "encoder_vq.safetensors"), device="cpu")
+    dec = load_file(os.path.join(a.src, "decoder.safetensors"), device="cpu")
+    full = dict(enc)
+    full.update(dec)
+    os.makedirs(a.out, exist_ok=True)
+
+    for name, prefixes in (("rstok_sat.pt", SAT_PREFIXES), ("rstok_ground.pt", GROUND_PREFIXES)):
+        sd = {k: v for k, v in full.items()
+              if k.startswith(prefixes) and not k.startswith(DROP_PREFIXES)}
+        if not sd:
+            raise RuntimeError(f"no tensors matched {prefixes}")
+        p = os.path.join(a.out, name)
+        torch.save({"config": CONFIG, "state_dict": sd}, p)
+        n = sum(v.numel() for v in sd.values())
+        print(f"  {p}  {len(sd)} tensors, {n/1e6:.1f}M params, {os.path.getsize(p)/1e6:.1f} MB")
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())