FIG-002 Figure

"""FIG-002: DiT block with adaLN-Zero (paper Fig. 2 right).

Self-contained generator for knowledge card FIG-002 (SYSTEM.md D19):
solve -> render -> overwrite FIG-002.svg next to this script.
Run: kblayout/.venv/bin/python figures/FIG-002.py
"""
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "kblayout"))

from kblayout.core.figure import Figure
from kblayout.ext import layout
from kblayout.ext.line import edge
from kblayout.ext.rect import rect
from kblayout.ext.text import text_block, measure as text_block_measure

FIGURE = Figure(id="dit-block-adaln")


def node(nid, label, cx, cy, *,
         font_size: float = 14.0, gap: float = 10.0):
    tw, th = text_block_measure(label, font_size, pad=0.0)
    w, h = tw + 2 * gap, th + 2 * gap
    r = rect(FIGURE, nid, w, h, cx=cx, cy=cy)
    text_block(FIGURE, r, label, pad=0.0)
    return r


# ---- main chain (bottom -> top) ----
n_in = node("input_tokens", "Input Tokens", 0, 0)
n_ln1 = node("ln1", "Layer Norm", 0, -100)
n_ss1 = node("ss1", "Scale, Shift", 0, -185)   # gamma1/beta1: LN1 out
n_mha = node("mha", "Multi-Head\nSelf-Attention", 0, -305)
n_sc1 = node("sc1", "Scale", 0, -370)           # alpha1: before resid
n_add1 = node("add1", "⊕", 0, -440, gap=20.0)
n_ln2 = node("ln2", "Layer Norm", 0, -540)
n_ss2 = node("ss2", "Scale, Shift", 0, -605)   # gamma2/beta2: LN2 out
n_ff = node("ff", "Pointwise\nFeedforward", 0, -725)
n_sc2 = node("sc2", "Scale", 0, -795)           # alpha2: before resid
n_add2 = node("add2", "⊕", 0, -865, gap=20.0)

# ---- conditioning side ----
n_cond = node("cond", "Conditioning", 130, -10)
n_mlp = node("mlp", "MLP", 130, -100)          # aligned with Layer Norm 1

# ---- chain edges (vertical) ----
chain = [n_in, n_ln1, n_ss1, n_mha, n_sc1, n_add1,
         n_ln2, n_ss2, n_ff, n_sc2, n_add2]
for src, dst in zip(chain, chain[1:]):
    edge(FIGURE, src.top, dst.bottom, bends=0, U0=(0, -1), U1=(0, -1))

# ---- residual #1: Input Tokens -> (+) #1 via the LEFT corridor ----
edge(FIGURE, n_in.left, n_add1.left, bends=2,
     U0=(-1, 0), U1=(1, 0),
     waypoints=[(-90.0, -235.0)])
# ---- residual #2: (+) #1 -> (+) #2 via the LEFT corridor ----
edge(FIGURE, n_add1.left, n_add2.left, bends=2,
     U0=(-1, 0), U1=(1, 0),
     waypoints=[(-90.0, -680.0)])

# ---- conditioning column ----
edge(FIGURE, n_cond.top, n_mlp.bottom, bends=0, U0=(0, -1), U1=(0, -1))

# ---- parameter lines: MLP top -> gates, L-shaped per line
# (vertical out of the MLP top, then horizontal into the gate right
# edge - the paper's fan is drawn as slanted lines; we keep folds to
# stay axis-aligned) ----
for gate, gy in [(n_ss1, -185.0), (n_sc1, -370.0),
                 (n_ss2, -605.0), (n_sc2, -795.0)]:
    edge(FIGURE, n_mlp.top, gate.right, bends=1,
         U0=(0, -1), U1=(-1, 0), w_dir=60.0)

# ---- layout relations ----
flow = [(n_in, n_ln1), (n_ln1, n_ss1), (n_ss1, n_mha), (n_mha, n_sc1),
        (n_sc1, n_add1), (n_add1, n_ln2), (n_ln2, n_ss2), (n_ss2, n_ff),
        (n_ff, n_sc2), (n_sc2, n_add2), (n_cond, n_mlp)]
for src, dst in flow:
    FIGURE.constraints.append(layout.flow_progress(src=src, dst=dst, gap=24.0))
FIGURE.constraints.append(
    layout.align_x(refs=[r.center for r in chain], name="layout.main_col"))
FIGURE.constraints.append(
    layout.align_x(refs=[n_cond.center, n_mlp.center],
                   name="layout.cond_col"))
for a, b in [(n_in, n_cond), (n_ln1, n_mlp), (n_mha, n_mlp),
             (n_add1, n_cond), (n_ff, n_cond), (n_ff, n_mlp)]:
    FIGURE.constraints.append(layout.pair_no_overlap(a=a, b=b, gap=16.0))
FIGURE.constraints.append(
    layout.origin_anchor(refs=[n_in.center, n_add1.center, n_add2.center]))


if __name__ == "__main__":
    from kblayout.kgrad.solver import solve_kgrad as solve
    from kblayout.render.svg import render
    from kblayout.report import backfill_card, print_report

    # latest engine: whole-solve loop in one Mojo call (A5: requires
    # KGRAD_MOJO_LIB pointing at kblayout/kgrad-mojo; no silent fallback)
    res = solve(FIGURE, steps=3000, lr=0.1, backend="mojo")
    svg = render(res, FIGURE)
    out = Path(__file__).with_suffix(".svg")
    out.write_text(svg, encoding="utf-8")
    print(f"{out.name}: loss {float(res.loss_final):.2f}")
    print_report(res)
    print(backfill_card(Path(__file__).with_suffix(".md"), res, svg))