Engine: native preset L (rrlite unfolded input, expansion and reconstruction stages)

engine/l_model.py

@@ -0,0 +1,50 @@
+#!/usr/bin/env python3
+"""Preset L (rrlite, unfolded): plain weights of the input stage (rrlite_enc0_4x4_*) and the expansion stage
+(rrlite_dec0_4x4), the two stages L does not share with preset M. See L_NOTES.md.
+
+Both are a Swin block of preset M's shape at render-pixel resolution (C 32, 2 heads, 4x4-pixel windows); their weight
+allocations are laid out like a network layer's, so m_model.plain_layer reads them:
+  enc0: the dither table (4096 bytes), the embedding (W1 [32][32] FP8, b1, W2 [32][32] FP8, b2) at 4096, then the
+        block and its patch merge (128 -> 64 channels) at 6272
+  dec0: the patch expand (64 -> 4 x 32 channels) and its bias, then the block (m_model.plain_layer with CIN 64)
+"""
+import os, sys
+import numpy as np
+HERE = os.path.dirname(os.path.abspath(__file__))
+sys.path.insert(0, HERE)
+import m_model as mm
+import model_enc3 as me
+import swin_model as sm
+
+E_EMB, E_BLOCK = 4096, 6272
+mm.LAYERS.update({'enc0l': (32, 2, 64, 0), 'dec0l': (32, 2, 0, 64)})
+
+
+class Blob:
+    """the part of a weight allocation that m_model.plain_layer reads (model_enc3.Params' wbuf and f16vec)"""
+    def __init__(self, wbuf):
+        self.wbuf = wbuf
+
+    def f16vec(self, off, n):
+        return self.wbuf[off:off + 2 * n].view(np.float16).copy()
+
+
+def enc0_layer(wbuf):
+    """plain tensors of the input stage: the block, its merge and the embedding (emb_*), natural channel order"""
+    t = mm.plain_layer(Blob(wbuf[E_BLOCK:]), 'enc0l')
+    fp8 = lambda T: me.E4M3[wbuf[T]]
+    b = Blob(wbuf)
+    g = mm.gp(32)
+    # the embedding has the layout of an MLP chunk of a C 32 block: fc1's columns in K-slot order, fc2's rows in pair order
+    t['emb_w1'] = np.asarray(fp8(sm.woff_table(E_EMB, 512, 16 * 32, 32, 32)), np.float16)
+    t['emb_b1'] = b.f16vec(E_EMB + 32 * 32, 32)
+    raw = fp8(sm.woff_table(E_EMB + 32 * 32 + 64, 0, 512, 32, 32))
+    w2 = np.zeros((32, 32)); w2[g] = raw[:, g]
+    t['emb_w2'] = np.asarray(w2, np.float16)
+    t['emb_b2'] = b.f16vec(E_EMB + 2 * 32 * 32 + 64, 32)[g]
+    return t
+
+
+def dec0_layer(wbuf):
+    """plain tensors of the expansion stage: the patch expand and the block"""
+    return mm.plain_layer(Blob(wbuf), 'dec0l')