114 lines
5.8 KiB
Python
114 lines
5.8 KiB
Python
"""Encoder parity: MLX Qwen3-VL conditioner vs comfy CLIP conditioning.
|
|
|
|
Run on your Mac (needs torch+comfy+mlx+transformers). Compares per-token cosine
|
|
similarity of the 7680-dim conditioning for a prompt. ~1.0 means the MLX encoder
|
|
matches comfy. If it's low, sweep (in order): rope_theta, max_length, the layer
|
|
indices [9,18,27], then the chat template.
|
|
|
|
cd ComfyUI
|
|
python custom_nodes/ComfyUI-SeFi-MLX/sefi_mlx/parity_encoder.py \
|
|
models/text_encoders/sefi-image-5b-turbo_qwen3vl.safetensors \
|
|
"a blue ceramic mug on a white desk"
|
|
|
|
Provide a tokenizer dir (snapshot) as a 3rd arg if the HF id isn't cached.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import os
|
|
import sys
|
|
|
|
|
|
def main():
|
|
here = os.path.abspath(__file__)
|
|
custom_nodes = os.path.dirname(os.path.dirname(os.path.dirname(here)))
|
|
comfy_root = os.path.dirname(custom_nodes)
|
|
sys.path.insert(0, comfy_root)
|
|
sys.path.insert(0, os.path.dirname(os.path.dirname(here)))
|
|
|
|
import mlx.core as mx
|
|
import numpy as np
|
|
|
|
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"
|
|
tok_dir = sys.argv[3] if len(sys.argv) > 3 else ""
|
|
# 4th arg: max_length to pad BOTH sides to (e.g. 512). Default 0 = no padding
|
|
# (real tokens only). Use a padded length to verify the PADDING-token parity,
|
|
# which is what governs adherence (the DiT attends to padding unmasked).
|
|
pad_len = int(sys.argv[4]) if len(sys.argv) > 4 else 0
|
|
|
|
from sefi_mlx.encoder_mlx import load_sefi_encoder_mlx
|
|
ml = pad_len if pad_len > 0 else 1024
|
|
cond = load_sefi_encoder_mlx(ckpt, tokenizer_dir=tok_dir, max_length=ml)
|
|
mlx_c = np.array(cond.encode(prompt, pad_to_max=pad_len > 0).astype(mx.float32))[0]
|
|
print(f"MLX conditioning: {mlx_c.shape} (pad_to_max={pad_len > 0}, len={ml})")
|
|
|
|
# comfy reference (optional; needs a loaded CLIP)
|
|
try:
|
|
import comfy.sd
|
|
import folder_paths
|
|
ctype = getattr(comfy.sd.CLIPType, "FLUX2", None) or comfy.sd.CLIPType.FLUX
|
|
clip = comfy.sd.load_clip(ckpt_paths=[ckpt], clip_type=ctype,
|
|
embedding_directory=folder_paths.get_folder_paths("embeddings"))
|
|
if pad_len > 0: # force comfy to pad to the same length (native trick)
|
|
for _n, sub in vars(clip.tokenizer).items():
|
|
if hasattr(sub, "min_length"):
|
|
sub.min_length = pad_len
|
|
tokens = clip.tokenize(prompt)
|
|
# --- token-id diagnostic: compare comfy's padded ids to the MLX path's ---
|
|
try:
|
|
from sefi_mlx.encoder_mlx import SEFI_TEMPLATE
|
|
key = next(iter(tokens))
|
|
comfy_ids = [t[0] for t in tokens[key][0]]
|
|
text = SEFI_TEMPLATE.format(prompt) # the actual MLX template
|
|
enc = cond.tokenizer(text, truncation=True, padding="max_length",
|
|
max_length=pad_len or 512, return_tensors="np")
|
|
mlx_ids = enc["input_ids"][0].tolist()
|
|
print(f"comfy ids: n={len(comfy_ids)} first16={comfy_ids[:16]} last5={comfy_ids[-5:]}")
|
|
print(f"mlx ids: n={len(mlx_ids)} first16={mlx_ids[:16]} last5={mlx_ids[-5:]}")
|
|
nn = min(len(comfy_ids), len(mlx_ids))
|
|
diff_pos = [i for i in range(nn) if comfy_ids[i] != mlx_ids[i]]
|
|
print(f"matching ids: {nn - len(diff_pos)}/{nn}")
|
|
for i in diff_pos[:12]:
|
|
print(f" pos {i}: comfy={comfy_ids[i]} mlx={mlx_ids[i]}")
|
|
globals()["_diff_pos"] = set(diff_pos)
|
|
except Exception as de:
|
|
print(f"(id diagnostic skipped: {de})")
|
|
c = clip.encode_from_tokens_scheduled(tokens)
|
|
ref = c[0][0].float().cpu().numpy()[0] # (L, 7680)
|
|
n = min(ref.shape[0], mlx_c.shape[0])
|
|
a, b = ref[:n], mlx_c[:n]
|
|
cos = (a * b).sum(-1) / (np.linalg.norm(a, axis=-1) * np.linalg.norm(b, axis=-1) + 1e-8)
|
|
print(f"compared {n} tokens (incl. padding if padded)")
|
|
print(f"per-token cosine: mean={cos.mean():.4f} min={cos.min():.4f}")
|
|
if n > 20: # separate real vs padding-region parity
|
|
print(f" real(first 16): mean={cos[:16].mean():.4f} "
|
|
f"padding(rest): mean={cos[16:].mean():.4f} min={cos[16:].min():.4f}")
|
|
# split padding parity by whether the token id matched, to prove
|
|
# whether the divergence is the 4 id diffs or a systemic forward diff
|
|
dp = globals().get("_diff_pos", set())
|
|
pad_match = [cos[i] for i in range(16, n) if i not in dp]
|
|
pad_mismatch = [cos[i] for i in range(16, n) if i in dp]
|
|
if pad_match:
|
|
pm = np.array(pad_match)
|
|
print(f" padding @ MATCHING ids: mean={pm.mean():.4f} min={pm.min():.4f} "
|
|
f"(n={len(pm)}) <-- if low, it's a FORWARD diff, not tokenization")
|
|
if pad_mismatch:
|
|
print(f" padding @ differing ids: mean={np.array(pad_mismatch).mean():.4f} "
|
|
f"(n={len(pad_mismatch)})")
|
|
# per-position cosine across the padding region: if it falls with
|
|
# position -> rope/position issue; if flat/random -> context/forward.
|
|
samples = [p for p in (20, 24, 40, 80, 160, 256, 384, 511) if p < n]
|
|
print(" per-position cosine:",
|
|
" ".join(f"@{p}={cos[p]:.3f}" for p in samples))
|
|
print("ENCODER PARITY OK" if cos.mean() > 0.99 else
|
|
"ENCODER MISMATCH — check position ids / mask / layers")
|
|
except Exception as e:
|
|
print(f"(comfy reference unavailable: {e}); MLX encoder ran — shape looks "
|
|
f"{'correct' if mlx_c.shape[-1] == 7680 else 'WRONG'} (expect 7680).")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|