107 lines
4.2 KiB
Python
107 lines
4.2 KiB
Python
"""Localize the divergence between my MLX encoder and the OFFICIAL HF SeFi encoder.
|
|
|
|
Dumps, from the real HF Qwen3-VL-4B encoder (sefi/modeling/qwen3vl_text_encoder.py):
|
|
1. the token IDs it feeds (apply_chat_template) vs my SEFI_TEMPLATE,
|
|
2. the position_ids shape it uses (1-axis vs 3-axis M-RoPE),
|
|
3. its RoPE cos table vs my standard rope (theta 5e6),
|
|
so I know whether to fix the template, the rope (interleaved M-RoPE), or both.
|
|
|
|
cd ComfyUI
|
|
python custom_nodes/ComfyUI-SeFi-MLX/sefi_mlx/dump_official.py \
|
|
models/text_encoders/sefi-image-5b-turbo_qwen3vl.safetensors \
|
|
"a blue ceramic mug on a white desk"
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import glob
|
|
import os
|
|
import sys
|
|
|
|
|
|
def main():
|
|
here = os.path.abspath(__file__)
|
|
cn = os.path.dirname(os.path.dirname(os.path.dirname(here)))
|
|
comfy_root = os.path.dirname(cn)
|
|
sys.path.insert(0, comfy_root)
|
|
sys.path.insert(0, os.path.dirname(os.path.dirname(here)))
|
|
sys.path.insert(0, os.path.join(cn, "ComfyUI-SeFi-Image", "SeFi-Image"))
|
|
|
|
import numpy as np
|
|
import torch
|
|
import mlx.core as mx
|
|
|
|
ckpt = sys.argv[1] if len(sys.argv) > 1 else os.path.join(
|
|
comfy_root, "models", "text_encoders", "sefi-image-5b-turbo_qwen3vl.safetensors")
|
|
prompt = sys.argv[2] if len(sys.argv) > 2 else "a blue ceramic mug on a white desk"
|
|
max_length = 1024
|
|
|
|
pat = os.path.join(comfy_root, "models", "sefi", "models--SeFi-Image*", "snapshots", "*")
|
|
snaps = [d for d in glob.glob(pat) if os.path.isdir(os.path.join(d, "Qwen3-VL-4B-Instruct"))]
|
|
snap = snaps[0]
|
|
print("snapshot:", snap)
|
|
|
|
from sefi.modeling.qwen3vl_text_encoder import Qwen3VLTextEncoder
|
|
enc = Qwen3VLTextEncoder("qwen3vl_4b", weights_root=snap, max_length=max_length,
|
|
hidden_layers=(9, 18, 27))
|
|
|
|
# ---- 1) official token ids vs SEFI_TEMPLATE ----
|
|
chat = enc._build_chat_text(prompt)
|
|
off = enc.tokenizer([chat], return_tensors="pt", padding="max_length",
|
|
truncation=True, max_length=max_length)["input_ids"][0].tolist()
|
|
print("\nofficial chat_text repr:\n", repr(chat))
|
|
from sefi_mlx.encoder_mlx import SEFI_TEMPLATE
|
|
mine = enc.tokenizer([SEFI_TEMPLATE.format(prompt)], return_tensors="pt",
|
|
padding="max_length", truncation=True, max_length=max_length)["input_ids"][0].tolist()
|
|
print("official ids first24:", off[:24])
|
|
print("mine ids first24:", mine[:24])
|
|
n = min(len(off), len(mine))
|
|
same = sum(1 for i in range(n) if off[i] == mine[i])
|
|
print(f"token ids match: {same}/{n}")
|
|
|
|
# ---- 2/3) capture HF rope (cos table) + position_ids ----
|
|
caps = {}
|
|
rot = None
|
|
for name, mod in enc.model.named_modules():
|
|
if name.endswith("rotary_emb"):
|
|
rot = mod
|
|
print("rotary module:", name)
|
|
break
|
|
if rot is not None:
|
|
def pre(_m, args, kwargs):
|
|
pid = kwargs.get("position_ids")
|
|
if pid is None and len(args) > 1:
|
|
pid = args[1]
|
|
if pid is not None:
|
|
caps["pos_shape"] = tuple(pid.shape)
|
|
caps["pos_sample"] = pid.detach().float().cpu().numpy().reshape(pid.shape[0], -1)[:, :6].tolist()
|
|
|
|
def post(_m, i, o):
|
|
if "cos" not in caps:
|
|
cos = o[0] if isinstance(o, (tuple, list)) else o
|
|
caps["cos"] = cos.detach().float().cpu().numpy()
|
|
caps["cos_shape"] = tuple(cos.shape)
|
|
rot.register_forward_pre_hook(pre, with_kwargs=True)
|
|
rot.register_forward_hook(post)
|
|
|
|
with torch.no_grad():
|
|
enc.encode([prompt])
|
|
|
|
print("\nposition_ids shape:", caps.get("pos_shape"), "sample:", caps.get("pos_sample"))
|
|
print("HF rope cos shape:", caps.get("cos_shape"))
|
|
|
|
if "cos" in caps:
|
|
hf = caps["cos"]
|
|
while hf.ndim > 2:
|
|
hf = hf[0]
|
|
from sefi_mlx.encoder_mlx import sefi_rope_tables
|
|
my_cos = np.array(sefi_rope_tables(max_length, 128, 5_000_000.0)[0].astype(mx.float32))
|
|
for p in (16, 256):
|
|
d = np.abs(hf[p] - my_cos[p])
|
|
print(f" rope-cos @pos{p}: HF[:8]={np.round(hf[p,:8],4)}")
|
|
print(f" std[:8]={np.round(my_cos[p,:8],4)} max|d|={d.max():.4f}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|