Idol / python /export_graphs.py

Faster default form in the root; the first release's form moves to 262k/

74728e5 verified about 3 hours ago

15 kB

"""Export the shipped Idol graph set (idol_graphs.py) with litert-torch, after checking each part's GPU form in
PyTorch float32 against the reference run's own tensors (ref/<input>/ref.npz from ref_pipeline.py).
  python export_graphs.py --parts den,oct,gs,dino,vae --out models/v1 --ref ref/armchair [--no-export]
"""
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)
sys.path.insert(0, os.path.expanduser(os.environ.get('IDOL_SDK', 'idol-sdk')))  # github.com/speridlabs/idol-sdk @4702a7f
import f0_export as F0  # noqa: E402
import idol_graphs as IG  # noqa: E402
import idol_glue as GL  # noqa: E402


def cos_rel(got, ref):
    a, b = np.asarray(got, np.float64).ravel(), np.asarray(ref, np.float64).ravel()
    return {'cos': float(a @ b / (np.linalg.norm(a) * np.linalg.norm(b))),
            'rel_l2': float(np.linalg.norm(a - b) / np.linalg.norm(b)), 'max_abs_diff': float(np.abs(a - b).max()),
            'max_abs_ref': float(np.abs(b).max())}


def show(name, r):
    print(f'[check] {name}: cos {r["cos"]:.9f} rel_l2 {r["rel_l2"]:.3e} max|d| {r["max_abs_diff"]:.3e} '
          f'max|ref| {r["max_abs_ref"]:.3e}', flush=True)
    return r


def t32(a):
    return torch.from_numpy(np.ascontiguousarray(a, dtype=np.float32))


def export(module, args, path):
    t0 = time.time()
    edge = __import__('litert_torch').convert(module.eval(), args)
    edge.export(path)
    info = F0.scan_tflite(path)
    info['export_s'] = round(time.time() - t0, 1)
    print(f'[export] {os.path.basename(path)}: {info["bytes"]:,} B, {info["ops"]} ops, rank<= {info["max_rank"]}, '
          f'BMM ranks {info["bmm_ranks"]}, {info["export_s"]} s', flush=True)
    return info


def export_multi(named, path):
    import litert_torch
    t0 = time.time()
    conv = None
    for name, module, args in named:
        conv = litert_torch.signature(name, module.eval(), args) if conv is None else conv.signature(name, module, args)
    conv.convert().export(path)
    info = F0.scan_tflite(path)
    info['export_s'] = round(time.time() - t0, 1)
    print(f'[export] {os.path.basename(path)}: {info["bytes"]:,} B, {info["ops"]} ops ({len(named)} signatures), '
          f'{info["export_s"]} s', flush=True)
    return info


# ---------------------------------------------------------------- denoiser

def part_den(out, ref, do_export, chunk=512):
    den, _ = F0.load_denoiser()
    den = den.float()
    lora = load_file(F0.snap('speridlabs/idol-turbo', 'idol_turbo_lora.safetensors'))
    ctx, dinp = IG.DenCtx(den, lora, chunk).eval(), IG.DenIn(den, lora, chunk).eval()
    joints = [IG.DenJoint(den, k, lora, chunk).eval() for k in range(len(den.blocks))]
    pos_emb = den.pos_embedder(den.pos_pe).float()
    pos_emb.numpy().astype(np.float32).tofile(os.path.join(out, 'den_pos_emb.f32'))
    res = {}
    f1, f2 = t32(ref['feature1']), t32(ref['feature2'])
    tg = ref['t_grid']
    for i in (0, len(tg) - 2):
        t = float(tg[i])
        gate = IG.lora_gate('slow' if t >= GL.PHASE_THRESHOLD else 'fast')
        h_cond = ctx(f1, f2, gate)
        h_x, h_cam, t_mod, t_emb = dinp(t32(ref[f'x_latent_{i}']), t32(ref[f'x_camera_{i}']),
                                        IG.sinusoid(1000 * t), pos_emb, gate)
        h = joints[0](h_x, h_cond, h_cam, t_mod, gate)
        for k in range(1, len(joints) - 1):
            h = joints[k](h, t_mod, gate)
        v, vc = joints[-1](h, t_mod, t_emb, gate)
        res[f'step{i}_latent'] = show(f'denoiser step {i} (t={t:.4f}) latent velocity', cos_rel(v, ref[f'v_latent_{i}']))
        res[f'step{i}_camera'] = show(f'denoiser step {i} camera velocity', cos_rel(vc, ref[f'v_camera_{i}']))
    if do_export:
        g = IG.lora_gate('slow')
        L = IG.DEN_TOKENS + IG.COND_TOKENS + 1
        res['export'] = {
            'den_ctx': export(ctx, (f1, f2, g), os.path.join(out, 'den_ctx.tflite')),
            'den_in': export(dinp, (torch.randn(1, 8192, 16), torch.randn(1, 1, 5), IG.sinusoid(700.0), pos_emb, g),
                             os.path.join(out, 'den_in.tflite')),
        }
        t_mod = torch.randn(1, 6144)
        for k, jb in enumerate(joints):
            if k == 0:
                args = (torch.randn(1, 8192, 1024), torch.randn(1, 4101, 1024), torch.randn(1, 1, 1024), t_mod, g)
            elif k == len(joints) - 1:
                args = (torch.randn(1, L, 1024), t_mod, torch.randn(1, 1024), g)
            else:
                args = (torch.randn(1, L, 1024), t_mod, g)
            res['export'][f'den_joint_{k:02d}'] = export(jb, args, os.path.join(out, f'den_joint_{k:02d}.tflite'))
    return res


# ---------------------------------------------------------------- octree

def load_decoder(name):
    from idol.models.decoder import GaussianFixedlenDecoder, OctreeProbabilityFixedlenDecoder
    cfg = json.load(open(F0.snap('speridlabs/idol', 'config.json')))
    cls = OctreeProbabilityFixedlenDecoder if name == 'decoder' else GaussianFixedlenDecoder
    m = cls(**cfg[name]['args'])
    m.load_state_dict(load_file(F0.snap('speridlabs/idol', f'{name}.safetensors')), strict=True)
    m = m.float().eval()
    m.dtype = torch.float32
    return m


def bucket(n):
    for b in IG.OCT_BUCKETS:
        if n <= b:
            return b
    raise ValueError(n)


def oct_inputs(cells, level):
    """cells [1, n, 3] -> padded graph inputs at the bucket size."""
    n = cells.shape[1]
    b = bucket(n)
    c = torch.zeros(1, b, 3)
    c[:, :n] = cells
    return (c, IG.point_pe(c[0])[None], IG.sinusoid(float(1 << level), max_period=1024.0, two_pi=True)), b


def part_oct(out, ref, do_export, chunk=512):
    dec = load_decoder('decoder')
    g = IG.Octree(dec, chunk).eval()
    latent = t32(ref['latent'])
    res = {}
    for l in range(1, GL.MAX_LEVEL + 1):
        cells = t32(ref[f'oct_cells_{l}'])
        (c, pe, lf), b = oct_inputs(cells, l)
        logits = g(c, pe, lf, latent)[:, :cells.shape[1]]
        res[f'level{l}'] = show(f'octree level {l} ({cells.shape[1]} rows, bucket {b}) probs',
                                cos_rel(torch.softmax(logits, -1), ref[f'oct_probs_{l}']))
    if do_export:
        named = [(f'L{b}', g, (torch.rand(1, b, 3), torch.randn(1, b, 1024), IG.sinusoid(64.0, max_period=1024.0,
                  two_pi=True), latent)) for b in IG.OCT_BUCKETS]
        res['export'] = {'oct': export_multi(named, os.path.join(out, 'oct.tflite'))}
    return res


# ---------------------------------------------------------------- Gaussian decoder

GS_SPLIT = [(2 * i, 2 * i + 2) for i in range(8)]


def part_gs(out, ref, do_export, chunk=512):
    dgs = load_decoder('decoder_gs')
    parts = [IG.Gauss(dgs, lo, hi, chunk).eval() for lo, hi in GS_SPLIT]
    anchors, latent = t32(ref['anchors']), t32(ref['latent'])
    pe = IG.point_pe(anchors[0])[None]
    h = parts[0](anchors, pe, latent)
    for p in parts[1:]:
        h = p(h, latent)
    res = {'features': show('gaussian decoder features', cos_rel(h, ref['gs_features']))}
    gg = GL.to_gaussians(anchors, h, dgs.rep_config, dgs.points_offset_perturbation.float(),
                         dgs.base_offset_scale.float())
    for k in ('xyz', 'features_dc', 'scaling', 'rotation', 'opacity'):
        res[k] = show(f'gaussians {k}', cos_rel(gg[k], ref[f'g_{k}']))
    np.save(os.path.join(out, 'gs_perturbation.npy'), dgs.points_offset_perturbation.float().numpy())
    json.dump({'base_offset_scale': float(dgs.base_offset_scale), 'rep': dgs.rep_config},
              open(os.path.join(out, 'gs_rep.json'), 'w'), indent=1)
    if do_export:
        res['export'] = {}
        for i, p in enumerate(parts):
            args = (anchors, pe, latent) if i == 0 else (torch.randn(1, anchors.shape[1], 1024), latent)
            res['export'][f'gs_{i}'] = export(p, args, os.path.join(out, f'gs_{i}.tflite'))
    return res


# ---------------------------------------------------------------- DINOv3

DINO_SPLIT = [(8 * i, 8 * i + 8) for i in range(4)]


def load_dino():
    from transformers import DINOv3ViTModel
    cfg = F0.dino_config()
    m = DINOv3ViTModel(cfg)
    sd = load_file(F0.snap('Comfy-Org/Idol', 'clip_vision/dino_v3_vit_h_bf16.safetensors'))
    sd = {('model.' + k if k.startswith('layer.') else k): v for k, v in sd.items()}
    m.load_state_dict(sd, strict=False)
    return m.float().eval(), cfg


def part_dino(out, ref, do_export, chunk=1024):
    model, cfg = load_dino()
    _, _, cos, sin = F0.dino_rope(cfg, 64)
    cos.numpy().astype(np.float32).tofile(os.path.join(out, 'dino_rope_cos.f32'))
    sin.numpy().astype(np.float32).tofile(os.path.join(out, 'dino_rope_sin.f32'))
    parts = [IG.Dino(model, lo, hi, chunk).eval() for lo, hi in DINO_SPLIT]
    pixels = t32(ref['pixels'])
    h = pixels
    for p in parts:
        h = p(h, cos, sin)
    res = {'feature1': show('dinov3 feature1', cos_rel(h, ref['feature1']))}
    if do_export:
        res['export'] = {}
        for i, p in enumerate(parts):
            args = (pixels if i == 0 else torch.randn(1, 4101, 1280), cos, sin)
            res['export'][f'dino_{i}'] = export(p, args, os.path.join(out, f'dino_{i}.tflite'))
    return res


# ---------------------------------------------------------------- BiRefNet (the author's background remover)

class BirefnetGraph(torch.nn.Module):
    """BiRefNet.forward: normalised image [1, 3, 1024, 1024] -> sigmoid map [1, 1, 1024, 1024]."""

    def __init__(self, bir):
        super().__init__()
        self.m = bir

    def forward(self, x):
        return self.m.forward(x)


def part_bir(out, ref, do_export):
    """The author's BiRefNet with its 20 deformable convolutions in gather form (deform_conv.py) and the memory-lean
    rewrites of bir_lean.py, one graph for the CPU (the gather is GATHER_ND, which the LiteRT GPU does not run).
    Checked on the reference's own input image:
    the gather form vs the author's forward, then the uint8 alpha after the glue's host steps vs the reference's."""
    from PIL import Image
    import idol.models.birefnet as BR
    import bir_lean
    root = os.path.dirname(F0.snap('speridlabs/idol', 'config.json'))
    bir = BR.BiRefNet()
    bir.load_state_dict(load_file(os.path.join(root, 'birefnet.safetensors')), strict=True)
    bir = bir.eval().float()
    meta = json.load(open(os.path.join(os.path.dirname(os.path.abspath(ref['__path__'])), 'meta.json')))
    raw = Image.open(meta['image'])
    w, h = raw.size
    scale = GL.IMAGE_SIZE / min(w, h)
    rgb = raw.resize((max(1, int(round(w * scale))), max(1, int(round(h * scale)))), Image.LANCZOS).convert('RGB')
    captured = {}

    def run_author(x):
        captured['x'] = torch.from_numpy(x)
        captured['map'] = bir.forward(captured['x'])
        return captured['map'].numpy()

    a_author = GL.birefnet_alpha(rgb, run_author)
    # the memory-lean form (bir_lean.py): row-tiled gathers, no ASPP concatenation, the last 1x1 conv folded
    # through the resize; the first form (one-tile gathers, the author's concatenations) peaked at 9.3 GB per call
    restore = bir_lean.apply(BR)
    try:
        graph = BirefnetGraph(bir).eval()
        res = {'map': show('birefnet map, gather form vs author', cos_rel(graph(captured['x']), captured['map']))}
        a_glue = GL.birefnet_alpha(rgb, lambda x: graph(torch.from_numpy(x)).numpy())
        d = np.abs(a_glue.astype(int) - ref['alpha_birefnet'].astype(int))
        res['alpha_uint8_vs_reference'] = {'pixels_differing': int((d > 0).sum()), 'max_abs_diff': int(d.max()),
                                           'author_forward_equal_reference': bool((a_author == ref['alpha_birefnet']).all())}
        print(f'[check] birefnet uint8 alpha vs reference: {res["alpha_uint8_vs_reference"]}', flush=True)
        if do_export:
            res['export'] = {'bir': export(graph, (captured['x'],), os.path.join(out, 'bir.tflite'))}
    finally:
        restore()
    return res


# ---------------------------------------------------------------- VAE encoder

def latent_tokens(moments, noise, vae):
    """ImageEncoder._flux2_tokens after the encoder, with the shared posterior noise -> [1, 4096, 128]."""
    m, logvar = torch.chunk(moments, 2, dim=1)
    lat = m + torch.exp(0.5 * torch.clamp(logvar, -30.0, 20.0)) * noise
    b, c, h, w = lat.shape
    lat = lat.view(b, c, h // 2, 2, w // 2, 2).permute(0, 1, 3, 5, 2, 4).reshape(b, c * 4, h // 2, w // 2)
    bm = vae.bn.running_mean.view(1, -1, 1, 1)
    bs = torch.sqrt(vae.bn.running_var.view(1, -1, 1, 1) + vae.config.batch_norm_eps)
    return ((lat - bm) / bs).flatten(2).transpose(1, 2)


def part_vae(out, ref, noise, do_export, chunk=1024):
    vae = F0.load_vae()
    json.dump({'bn_mean': vae.bn.running_mean.float().tolist(), 'bn_var': vae.bn.running_var.float().tolist(),
               'bn_eps': vae.config.batch_norm_eps}, open(os.path.join(out, 'vae_bn.json'), 'w'))
    parts = [IG.Vae(F0.load_vae(), s, chunk).eval() for s in IG.Vae.PARTS]
    x = t32(ref['pixels'])
    h, ins = x, []
    for p in parts:
        ins.append(h)
        h = p(h)
    res = {'moments': show('vae moments', cos_rel(h, ref['vae_moments']))}
    tok = latent_tokens(h, t32(noise['vae_noise']), vae)
    res['tokens'] = show('vae latent tokens', cos_rel(tok, ref['feature2'][:, 5:]))
    if do_export:
        res['export'] = {}
        for s, p, a in zip(IG.Vae.PARTS, parts, ins):
            res['export'][f'vae_{s}'] = export(p, (a,), os.path.join(out, f'vae_{s}.tflite'))
    return res


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--parts', default='den,oct,gs,dino,vae')
    ap.add_argument('--out', required=True)
    ap.add_argument('--ref', required=True)
    ap.add_argument('--no-export', action='store_true')
    a = ap.parse_args()
    os.makedirs(a.out, exist_ok=True)
    torch.set_grad_enabled(False)
    ref = dict(np.load(os.path.join(a.ref, 'ref.npz')))
    ref['__path__'] = os.path.join(a.ref, 'ref.npz')
    noise = dict(np.load(os.path.join(a.ref, 'noise.npz')))
    jpath = os.path.join(a.out, 'export.json')
    results = json.load(open(jpath)) if os.path.exists(jpath) else {}
    for p in a.parts.split(','):
        t0 = time.time()
        if p == 'den':
            r = part_den(a.out, ref, not a.no_export)
        elif p == 'oct':
            r = part_oct(a.out, ref, not a.no_export)
        elif p == 'gs':
            r = part_gs(a.out, ref, not a.no_export)
        elif p == 'dino':
            r = part_dino(a.out, ref, not a.no_export)
        elif p == 'vae':
            r = part_vae(a.out, ref, noise, not a.no_export)
        elif p == 'bir':
            r = part_bir(a.out, ref, not a.no_export)
        else:
            raise SystemExit(p)
        r['wall_s'] = round(time.time() - t0, 1)
        results[p] = r
        json.dump(results, open(jpath, 'w'), indent=1)
        print(f'== {p} done in {r["wall_s"]} s', flush=True)


if __name__ == '__main__':
    main()

	"""Export the shipped Idol graph set (idol_graphs.py) with litert-torch, after checking each part's GPU form in
	PyTorch float32 against the reference run's own tensors (ref/<input>/ref.npz from ref_pipeline.py).

	python export_graphs.py --parts den,oct,gs,dino,vae --out models/v1 --ref ref/armchair [--no-export]
	"""
	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)
	sys.path.insert(0, os.path.expanduser(os.environ.get('IDOL_SDK', 'idol-sdk'))) # github.com/speridlabs/idol-sdk @4702a7f
	import f0_export as F0 # noqa: E402
	import idol_graphs as IG # noqa: E402
	import idol_glue as GL # noqa: E402


	def cos_rel(got, ref):
	a, b = np.asarray(got, np.float64).ravel(), np.asarray(ref, np.float64).ravel()
	return {'cos': float(a @ b / (np.linalg.norm(a) * np.linalg.norm(b))),
	'rel_l2': float(np.linalg.norm(a - b) / np.linalg.norm(b)), 'max_abs_diff': float(np.abs(a - b).max()),
	'max_abs_ref': float(np.abs(b).max())}


	def show(name, r):
	print(f'[check] {name}: cos {r["cos"]:.9f} rel_l2 {r["rel_l2"]:.3e} max\|d\| {r["max_abs_diff"]:.3e} '
	f'max\|ref\| {r["max_abs_ref"]:.3e}', flush=True)
	return r


	def t32(a):
	return torch.from_numpy(np.ascontiguousarray(a, dtype=np.float32))


	def export(module, args, path):
	t0 = time.time()
	edge = __import__('litert_torch').convert(module.eval(), args)
	edge.export(path)
	info = F0.scan_tflite(path)
	info['export_s'] = round(time.time() - t0, 1)
	print(f'[export] {os.path.basename(path)}: {info["bytes"]:,} B, {info["ops"]} ops, rank<= {info["max_rank"]}, '
	f'BMM ranks {info["bmm_ranks"]}, {info["export_s"]} s', flush=True)
	return info


	def export_multi(named, path):
	import litert_torch
	t0 = time.time()
	conv = None
	for name, module, args in named:
	conv = litert_torch.signature(name, module.eval(), args) if conv is None else conv.signature(name, module, args)
	conv.convert().export(path)
	info = F0.scan_tflite(path)
	info['export_s'] = round(time.time() - t0, 1)
	print(f'[export] {os.path.basename(path)}: {info["bytes"]:,} B, {info["ops"]} ops ({len(named)} signatures), '
	f'{info["export_s"]} s', flush=True)
	return info


	# ---------------------------------------------------------------- denoiser

	def part_den(out, ref, do_export, chunk=512):
	den, _ = F0.load_denoiser()
	den = den.float()
	lora = load_file(F0.snap('speridlabs/idol-turbo', 'idol_turbo_lora.safetensors'))
	ctx, dinp = IG.DenCtx(den, lora, chunk).eval(), IG.DenIn(den, lora, chunk).eval()
	joints = [IG.DenJoint(den, k, lora, chunk).eval() for k in range(len(den.blocks))]
	pos_emb = den.pos_embedder(den.pos_pe).float()
	pos_emb.numpy().astype(np.float32).tofile(os.path.join(out, 'den_pos_emb.f32'))
	res = {}
	f1, f2 = t32(ref['feature1']), t32(ref['feature2'])
	tg = ref['t_grid']
	for i in (0, len(tg) - 2):
	t = float(tg[i])
	gate = IG.lora_gate('slow' if t >= GL.PHASE_THRESHOLD else 'fast')
	h_cond = ctx(f1, f2, gate)
	h_x, h_cam, t_mod, t_emb = dinp(t32(ref[f'x_latent_{i}']), t32(ref[f'x_camera_{i}']),
	IG.sinusoid(1000 * t), pos_emb, gate)
	h = joints[0](h_x, h_cond, h_cam, t_mod, gate)
	for k in range(1, len(joints) - 1):
	h = joints[k](h, t_mod, gate)
	v, vc = joints[-1](h, t_mod, t_emb, gate)
	res[f'step{i}_latent'] = show(f'denoiser step {i} (t={t:.4f}) latent velocity', cos_rel(v, ref[f'v_latent_{i}']))
	res[f'step{i}_camera'] = show(f'denoiser step {i} camera velocity', cos_rel(vc, ref[f'v_camera_{i}']))
	if do_export:
	g = IG.lora_gate('slow')
	L = IG.DEN_TOKENS + IG.COND_TOKENS + 1
	res['export'] = {
	'den_ctx': export(ctx, (f1, f2, g), os.path.join(out, 'den_ctx.tflite')),
	'den_in': export(dinp, (torch.randn(1, 8192, 16), torch.randn(1, 1, 5), IG.sinusoid(700.0), pos_emb, g),
	os.path.join(out, 'den_in.tflite')),
	}
	t_mod = torch.randn(1, 6144)
	for k, jb in enumerate(joints):
	if k == 0:
	args = (torch.randn(1, 8192, 1024), torch.randn(1, 4101, 1024), torch.randn(1, 1, 1024), t_mod, g)
	elif k == len(joints) - 1:
	args = (torch.randn(1, L, 1024), t_mod, torch.randn(1, 1024), g)
	else:
	args = (torch.randn(1, L, 1024), t_mod, g)
	res['export'][f'den_joint_{k:02d}'] = export(jb, args, os.path.join(out, f'den_joint_{k:02d}.tflite'))
	return res


	# ---------------------------------------------------------------- octree

	def load_decoder(name):
	from idol.models.decoder import GaussianFixedlenDecoder, OctreeProbabilityFixedlenDecoder
	cfg = json.load(open(F0.snap('speridlabs/idol', 'config.json')))
	cls = OctreeProbabilityFixedlenDecoder if name == 'decoder' else GaussianFixedlenDecoder
	m = cls(**cfg[name]['args'])
	m.load_state_dict(load_file(F0.snap('speridlabs/idol', f'{name}.safetensors')), strict=True)
	m = m.float().eval()
	m.dtype = torch.float32
	return m


	def bucket(n):
	for b in IG.OCT_BUCKETS:
	if n <= b:
	return b
	raise ValueError(n)


	def oct_inputs(cells, level):
	"""cells [1, n, 3] -> padded graph inputs at the bucket size."""
	n = cells.shape[1]
	b = bucket(n)
	c = torch.zeros(1, b, 3)
	c[:, :n] = cells
	return (c, IG.point_pe(c[0])[None], IG.sinusoid(float(1 << level), max_period=1024.0, two_pi=True)), b


	def part_oct(out, ref, do_export, chunk=512):
	dec = load_decoder('decoder')
	g = IG.Octree(dec, chunk).eval()
	latent = t32(ref['latent'])
	res = {}
	for l in range(1, GL.MAX_LEVEL + 1):
	cells = t32(ref[f'oct_cells_{l}'])
	(c, pe, lf), b = oct_inputs(cells, l)
	logits = g(c, pe, lf, latent)[:, :cells.shape[1]]
	res[f'level{l}'] = show(f'octree level {l} ({cells.shape[1]} rows, bucket {b}) probs',
	cos_rel(torch.softmax(logits, -1), ref[f'oct_probs_{l}']))
	if do_export:
	named = [(f'L{b}', g, (torch.rand(1, b, 3), torch.randn(1, b, 1024), IG.sinusoid(64.0, max_period=1024.0,
	two_pi=True), latent)) for b in IG.OCT_BUCKETS]
	res['export'] = {'oct': export_multi(named, os.path.join(out, 'oct.tflite'))}
	return res


	# ---------------------------------------------------------------- Gaussian decoder

	GS_SPLIT = [(2 * i, 2 * i + 2) for i in range(8)]


	def part_gs(out, ref, do_export, chunk=512):
	dgs = load_decoder('decoder_gs')
	parts = [IG.Gauss(dgs, lo, hi, chunk).eval() for lo, hi in GS_SPLIT]
	anchors, latent = t32(ref['anchors']), t32(ref['latent'])
	pe = IG.point_pe(anchors[0])[None]
	h = parts[0](anchors, pe, latent)
	for p in parts[1:]:
	h = p(h, latent)
	res = {'features': show('gaussian decoder features', cos_rel(h, ref['gs_features']))}
	gg = GL.to_gaussians(anchors, h, dgs.rep_config, dgs.points_offset_perturbation.float(),
	dgs.base_offset_scale.float())
	for k in ('xyz', 'features_dc', 'scaling', 'rotation', 'opacity'):
	res[k] = show(f'gaussians {k}', cos_rel(gg[k], ref[f'g_{k}']))
	np.save(os.path.join(out, 'gs_perturbation.npy'), dgs.points_offset_perturbation.float().numpy())
	json.dump({'base_offset_scale': float(dgs.base_offset_scale), 'rep': dgs.rep_config},
	open(os.path.join(out, 'gs_rep.json'), 'w'), indent=1)
	if do_export:
	res['export'] = {}
	for i, p in enumerate(parts):
	args = (anchors, pe, latent) if i == 0 else (torch.randn(1, anchors.shape[1], 1024), latent)
	res['export'][f'gs_{i}'] = export(p, args, os.path.join(out, f'gs_{i}.tflite'))
	return res


	# ---------------------------------------------------------------- DINOv3

	DINO_SPLIT = [(8 * i, 8 * i + 8) for i in range(4)]


	def load_dino():
	from transformers import DINOv3ViTModel
	cfg = F0.dino_config()
	m = DINOv3ViTModel(cfg)
	sd = load_file(F0.snap('Comfy-Org/Idol', 'clip_vision/dino_v3_vit_h_bf16.safetensors'))
	sd = {('model.' + k if k.startswith('layer.') else k): v for k, v in sd.items()}
	m.load_state_dict(sd, strict=False)
	return m.float().eval(), cfg


	def part_dino(out, ref, do_export, chunk=1024):
	model, cfg = load_dino()
	_, _, cos, sin = F0.dino_rope(cfg, 64)
	cos.numpy().astype(np.float32).tofile(os.path.join(out, 'dino_rope_cos.f32'))
	sin.numpy().astype(np.float32).tofile(os.path.join(out, 'dino_rope_sin.f32'))
	parts = [IG.Dino(model, lo, hi, chunk).eval() for lo, hi in DINO_SPLIT]
	pixels = t32(ref['pixels'])
	h = pixels
	for p in parts:
	h = p(h, cos, sin)
	res = {'feature1': show('dinov3 feature1', cos_rel(h, ref['feature1']))}
	if do_export:
	res['export'] = {}
	for i, p in enumerate(parts):
	args = (pixels if i == 0 else torch.randn(1, 4101, 1280), cos, sin)
	res['export'][f'dino_{i}'] = export(p, args, os.path.join(out, f'dino_{i}.tflite'))
	return res


	# ---------------------------------------------------------------- BiRefNet (the author's background remover)

	class BirefnetGraph(torch.nn.Module):
	"""BiRefNet.forward: normalised image [1, 3, 1024, 1024] -> sigmoid map [1, 1, 1024, 1024]."""

	def __init__(self, bir):
	super().__init__()
	self.m = bir

	def forward(self, x):
	return self.m.forward(x)


	def part_bir(out, ref, do_export):
	"""The author's BiRefNet with its 20 deformable convolutions in gather form (deform_conv.py) and the memory-lean
	rewrites of bir_lean.py, one graph for the CPU (the gather is GATHER_ND, which the LiteRT GPU does not run).
	Checked on the reference's own input image:
	the gather form vs the author's forward, then the uint8 alpha after the glue's host steps vs the reference's."""
	from PIL import Image
	import idol.models.birefnet as BR
	import bir_lean
	root = os.path.dirname(F0.snap('speridlabs/idol', 'config.json'))
	bir = BR.BiRefNet()
	bir.load_state_dict(load_file(os.path.join(root, 'birefnet.safetensors')), strict=True)
	bir = bir.eval().float()
	meta = json.load(open(os.path.join(os.path.dirname(os.path.abspath(ref['__path__'])), 'meta.json')))
	raw = Image.open(meta['image'])
	w, h = raw.size
	scale = GL.IMAGE_SIZE / min(w, h)
	rgb = raw.resize((max(1, int(round(w * scale))), max(1, int(round(h * scale)))), Image.LANCZOS).convert('RGB')
	captured = {}

	def run_author(x):
	captured['x'] = torch.from_numpy(x)
	captured['map'] = bir.forward(captured['x'])
	return captured['map'].numpy()

	a_author = GL.birefnet_alpha(rgb, run_author)
	# the memory-lean form (bir_lean.py): row-tiled gathers, no ASPP concatenation, the last 1x1 conv folded
	# through the resize; the first form (one-tile gathers, the author's concatenations) peaked at 9.3 GB per call
	restore = bir_lean.apply(BR)
	try:
	graph = BirefnetGraph(bir).eval()
	res = {'map': show('birefnet map, gather form vs author', cos_rel(graph(captured['x']), captured['map']))}
	a_glue = GL.birefnet_alpha(rgb, lambda x: graph(torch.from_numpy(x)).numpy())
	d = np.abs(a_glue.astype(int) - ref['alpha_birefnet'].astype(int))
	res['alpha_uint8_vs_reference'] = {'pixels_differing': int((d > 0).sum()), 'max_abs_diff': int(d.max()),
	'author_forward_equal_reference': bool((a_author == ref['alpha_birefnet']).all())}
	print(f'[check] birefnet uint8 alpha vs reference: {res["alpha_uint8_vs_reference"]}', flush=True)
	if do_export:
	res['export'] = {'bir': export(graph, (captured['x'],), os.path.join(out, 'bir.tflite'))}
	finally:
	restore()
	return res


	# ---------------------------------------------------------------- VAE encoder

	def latent_tokens(moments, noise, vae):
	"""ImageEncoder._flux2_tokens after the encoder, with the shared posterior noise -> [1, 4096, 128]."""
	m, logvar = torch.chunk(moments, 2, dim=1)
	lat = m + torch.exp(0.5 * torch.clamp(logvar, -30.0, 20.0)) * noise
	b, c, h, w = lat.shape
	lat = lat.view(b, c, h // 2, 2, w // 2, 2).permute(0, 1, 3, 5, 2, 4).reshape(b, c * 4, h // 2, w // 2)
	bm = vae.bn.running_mean.view(1, -1, 1, 1)
	bs = torch.sqrt(vae.bn.running_var.view(1, -1, 1, 1) + vae.config.batch_norm_eps)
	return ((lat - bm) / bs).flatten(2).transpose(1, 2)


	def part_vae(out, ref, noise, do_export, chunk=1024):
	vae = F0.load_vae()
	json.dump({'bn_mean': vae.bn.running_mean.float().tolist(), 'bn_var': vae.bn.running_var.float().tolist(),
	'bn_eps': vae.config.batch_norm_eps}, open(os.path.join(out, 'vae_bn.json'), 'w'))
	parts = [IG.Vae(F0.load_vae(), s, chunk).eval() for s in IG.Vae.PARTS]
	x = t32(ref['pixels'])
	h, ins = x, []
	for p in parts:
	ins.append(h)
	h = p(h)
	res = {'moments': show('vae moments', cos_rel(h, ref['vae_moments']))}
	tok = latent_tokens(h, t32(noise['vae_noise']), vae)
	res['tokens'] = show('vae latent tokens', cos_rel(tok, ref['feature2'][:, 5:]))
	if do_export:
	res['export'] = {}
	for s, p, a in zip(IG.Vae.PARTS, parts, ins):
	res['export'][f'vae_{s}'] = export(p, (a,), os.path.join(out, f'vae_{s}.tflite'))
	return res


	def main():
	ap = argparse.ArgumentParser()
	ap.add_argument('--parts', default='den,oct,gs,dino,vae')
	ap.add_argument('--out', required=True)
	ap.add_argument('--ref', required=True)
	ap.add_argument('--no-export', action='store_true')
	a = ap.parse_args()
	os.makedirs(a.out, exist_ok=True)
	torch.set_grad_enabled(False)
	ref = dict(np.load(os.path.join(a.ref, 'ref.npz')))
	ref['__path__'] = os.path.join(a.ref, 'ref.npz')
	noise = dict(np.load(os.path.join(a.ref, 'noise.npz')))
	jpath = os.path.join(a.out, 'export.json')
	results = json.load(open(jpath)) if os.path.exists(jpath) else {}
	for p in a.parts.split(','):
	t0 = time.time()
	if p == 'den':
	r = part_den(a.out, ref, not a.no_export)
	elif p == 'oct':
	r = part_oct(a.out, ref, not a.no_export)
	elif p == 'gs':
	r = part_gs(a.out, ref, not a.no_export)
	elif p == 'dino':
	r = part_dino(a.out, ref, not a.no_export)
	elif p == 'vae':
	r = part_vae(a.out, ref, noise, not a.no_export)
	elif p == 'bir':
	r = part_bir(a.out, ref, not a.no_export)
	else:
	raise SystemExit(p)
	r['wall_s'] = round(time.time() - t0, 1)
	results[p] = r
	json.dump(results, open(jpath, 'w'), indent=1)
	print(f'== {p} done in {r["wall_s"]} s', flush=True)


	if __name__ == '__main__':
	main()