ayo
7c5d3c9157
Add native K/M/RR inference engines, HIP kernels, replay and validation harnesses, and Windows DLSS setup and runtime support. Include GPU profiling tools, build and packaging scripts, platform integration patches, and implementation documentation.
51 lines
2.1 KiB
Python
51 lines
2.1 KiB
Python
#!/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')
|