SeFiMLXComfyNode

ComfyUI-SeFi-MLX

A single node that downloads a SeFi-Image-*-turbo checkpoint, converts its DiT to Apple MLX, runs the semantic-first dual-time turbo denoise in MLX, and outputs an image — on Apple Silicon.

SeFi Turbo (MLX, all-in-one)
  inputs : model, prompt, width, height, steps(4/8/10), seed, quantize(q8/q4/bf16),
           [clip], [vae], [negative_prompt], [guidance_scale], [delta_t],
           [timestep_shift_alpha], [hf_token], [group_size], [dual_time]
  outputs: IMAGE, texture LATENT

One node does the whole pipeline (each stage is cached, so re-runs are fast):

1. Download the gated repo via huggingface_hub (needs a token + accepted license).
2. Build the reference torch SeFiTransformer (vendored under sefi_mlx/vendored/) from the diffusers transformer/ shards → a comfy-named, qkv-fused state dict.
3. Convert that to an MLX file under models/mlx/sefi/<model>-<dtype>/ (bf16, or affine-group-quantized q8/q4).
4. Sample: load the MLX DiT and run the dual-time turbo loop (steps 4/8/10, guidance 1.0 by default).
5. Decode: encode the prompt (wired CLIP or auto-built) and VAE-decode the texture latent (wired VAE or auto-built) → IMAGE.

Requirements

• Apple Silicon + pip install -r requirements.txt.
• ComfyUI itself — used for comfy.ops/comfy.utils during the one-time DiT conversion. No other custom nodes are required. This node is self-contained: the Qwen3-VL MLX encoder, the SeFi torch DiT, and the diffusers→comfy component converter are all vendored under sefi_mlx/vendored/ (MIT, same author).
• The model is gated (CC-BY-NC, non-commercial). Accept the license at https://huggingface.co/SeFi-Image/SeFi-Image-5B-turbo while logged in, then pass a token via the hf_token input or the HF_TOKEN env var.

Fully torch-free by default (MLX encoder + MLX VAE)

By default the node runs everything in MLX:

• DiT — faithful MLX port of the SeFi/Flux2 dual-stream backbone.
• Text encoder (mlx_encoder=True) — the Qwen3-VL-4B language tower in MLX (vendored from ComfyUI-Krea2-MLX), with SeFi's bare chat template and the [9,18,27]→7680 layer tap.
• VAE (mlx_vae=True) — the 32-ch AutoencoderKL decoder in MLX; it 2×2-unshuffles the packed texture (128→32) and decodes to RGB.

All three are converted once and cached under models/mlx/sefi/. The encoder still uses an HF tokenizer (CPU, no torch) loaded from the gated snapshot.

Overrides

If you'd rather use comfy's torch components, wire them and they take precedence:

• CLIPLoader → clip input (overrides the MLX encoder).
• VAELoader → vae input (overrides the MLX VAE).

Set mlx_encoder=False / mlx_vae=False to fall back to auto-built comfy components without wiring. Don't point the VAE at a generic 16-ch Flux VAE — the SeFi texture VAE is 32-ch; the node selects the right file by slug.

Turbo defaults (match the diffusers reference)

steps=4, guidance_scale=1.0 (no CFG), delta_t=0.1, timestep_shift_alpha=1.0 (identity), scheduler_shift=1.0 (linear flow-match sigmas). These mirror the official turbo recipe and the native node's FAITHFUL euler path.

⚠️ Verify on-device before trusting outputs

The MLX DiT (sefi_mlx/flux2_mlx.py) is a faithful port of the exact Flux2 backbone the native node uses (global modulation, fully bias-free, RMSNorm q/k, no-affine LayerNorms, 4-axis RoPE θ=2000, SwiGLU MLPs, dual-time embed). The torch↔MLX key mapping is verified (no extra/missing params). But the numeric forward could not be executed in the build environment (no Apple Silicon / no gated weights), so run the parity harness on your Mac first:

cd ComfyUI
python custom_nodes/ComfyUI-SeFi-MLX/dev/parity.py

It builds a tiny model, pushes it through the real converter+loader, runs the torch and MLX forwards on identical weights/input, and prints the max/mean abs difference (expect ≈0 in fp32; large/structural diffs = a bug to fix). A pass also proves the end-to-end key mapping. (The dev/ parity/replay harnesses are developer tools; some compare against the reference SeFi repos and are not needed to run the node.)

Known tuning points (flip if outputs look off)

• Dual-time order — vec = concat([semantic, texture]) (semantic first), matching the native DUAL_TIME_ORDER_SEM_FIRST=True.
• Time input scale — t is fed as timestep/1000 ∈ [0,1]; the sinusoid's internal time_factor=1000 restores the range (native TIME_INPUT_SCALE=1.0).
• Noise RNG — the initial latent is seeded with torch.randn (matching the SeFi reference), not mx.random.normal. This matters: MLX's normal RNG gives systematically worse compositions for this model, so torch noise is used when available (falls back to MLX only if torch is missing). Images still won't bit-match the native node at the same seed, but quality/adherence do.
• Resolution-aware (dynamic) flow shift is not wired here (turbo uses linear sigmas); add it if you need the base/RL schedules.

Files

| file | role | |------|------| | nodes.py | the single all-in-one node | | sefi_mlx/flux2_mlx.py | MLX port of the SeFi/Flux2 dual-time DiT | | sefi_mlx/convert.py | torch (comfy-named) → MLX safetensors + quantize | | sefi_mlx/loader.py | build + load the MLX DiT | | sefi_mlx/sample.py | dual-time turbo denoise loop (MLX) | | sefi_mlx/schedule.py | semantic-first schedule (numpy) | | sefi_mlx/encoder_mlx.py | SeFi Qwen3-VL-4B conditioner (reuses Krea2's port) | | sefi_mlx/vae_mlx.py | 32-ch AutoencoderKL decoder in MLX | | sefi_mlx/download.py | gated HF download + shard loading | | sefi_mlx/vendored/ | vendored MIT deps (Qwen3-VL MLX, SeFi torch DiT, converter) | | dev/parity*.py, dev/dump_*.py | developer verification/replay harnesses |

Verify the MLX parts on-device

The DiT, VAE, and encoder key mappings are all statically verified (no extra/missing params). Run the numeric parity checks on your Mac:

python custom_nodes/ComfyUI-SeFi-MLX/dev/parity.py          # DiT
python custom_nodes/ComfyUI-SeFi-MLX/dev/parity_vae.py      # VAE
python custom_nodes/ComfyUI-SeFi-MLX/dev/parity_encoder.py  # encoder

Tuning knobs if a check fails: VAE — the bn latent normalization is OFF by default (matches comfy's decode); encoder — rope_theta (5e6), max_length (1024), the [9,18,27] layer indices, then the chat template.

License

MIT — see LICENSE. Vendored code under sefi_mlx/vendored/ is from the author's own MIT-licensed ComfyUI-Krea2-MLX and ComfyUI-SeFi-Image projects. Model weights are gated and licensed separately by their authors (non-commercial) and are not included here.