File size: 3,542 Bytes
74728e5 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 | """The state lever's graph form (round 2): one graph for the denoiser's 22 middle joint blocks, each block's weights as
runtime inputs, so the phone compiles it once per photo instead of once per block (the S26 compiled the joint blocks
96 times per photo, 0.39 s each, because only one or two fit in memory at a time).
Writes <out>/den_joint_mid.tflite (no weights; inputs h, t_mod, gate, then the 30 weight inputs of
idol_graphs.SHARED_WEIGHT_NAMES) and <out>/den_joint_w.f32 (little-endian float32, blocks 1..22 one after another,
each block's inputs in that order, flattened), and checks the shared module against every block's own module in
PyTorch on one random input.
IDOL_SDK=<idol-sdk> python export_shared_joint.py --out models/state
"""
import argparse
import json
import os
import sys
import time
import numpy as np
import torch
from safetensors.torch import load_file
HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, HERE)
import f0_export as F0 # noqa: E402
import idol_graphs as IG # noqa: E402
def main():
ap = argparse.ArgumentParser()
ap.add_argument('--out', required=True)
ap.add_argument('--chunk', type=int, default=512)
a = ap.parse_args()
os.makedirs(a.out, exist_ok=True)
torch.set_grad_enabled(False)
den, _ = F0.load_denoiser()
den = den.float()
lora = load_file(F0.snap('speridlabs/idol-turbo', 'idol_turbo_lora.safetensors'))
blocks = list(range(1, len(den.blocks) - 1))
mids = [IG.DenJoint(den, k, lora, a.chunk).eval() for k in blocks]
shared = IG.DenJointShared(mids[0]).eval()
ref = mids[0].block
shapes = [list(t.shape) for t in IG.shared_block_weights(mids[0])]
L = IG.DEN_TOKENS + IG.COND_TOKENS + 1
torch.manual_seed(0)
h = torch.randn(1, L, 1024)
t_mod = torch.randn(1, 6144) * 0.1
gate = IG.lora_gate('slow')
worst = 0.0
wpath = os.path.join(a.out, 'den_joint_w.f32')
with open(wpath, 'wb') as f:
for k, jb in zip(blocks, mids):
b = jb.block
assert b.modulation and (b.c, b.h, b.d, b.eps1, b.eps2, b.repo.eps) == \
(ref.c, ref.h, ref.d, ref.eps1, ref.eps2, ref.repo.eps), k
w = IG.shared_block_weights(jb)
assert [list(t.shape) for t in w] == shapes, k
got, want = shared(h, t_mod, gate, *w), jb(h, t_mod, gate)
rel = float((got - want).norm() / want.norm())
worst = max(worst, rel)
print(f'block {k:2d}: shared vs own module rel L2 {rel:.2e}', flush=True)
for t in w:
f.write(t.numpy().astype('<f4').tobytes())
assert worst < 1e-5, worst
per_block = sum(int(np.prod(s)) for s in shapes)
assert os.path.getsize(wpath) == per_block * len(blocks) * 4
t0 = time.time()
import litert_torch
path = os.path.join(a.out, 'den_joint_mid.tflite')
litert_torch.convert(shared, (h, t_mod, gate, *IG.shared_block_weights(mids[0]))).export(path)
info = F0.scan_tflite(path)
json.dump({'blocks': blocks, 'names': IG.SHARED_WEIGHT_NAMES, 'shapes': shapes, 'per_block_floats': per_block,
'worst_rel_l2': worst, 'export_s': round(time.time() - t0, 1), 'scan': info},
open(os.path.join(a.out, 'shared_joint.json'), 'w'), indent=1, default=str)
print(f'den_joint_mid.tflite {os.path.getsize(path):,} B; den_joint_w.f32 {os.path.getsize(wpath):,} B '
f'({len(blocks)} blocks x {per_block:,} floats); worst rel L2 {worst:.2e}', flush=True)
if __name__ == '__main__':
main()
|