Nodes/ComfyUI_FFTNet/Load FFTNet Model
ComfyUI Node

Load FFTNet Model

The LLM loader that trades self-attention for an FFT

By bemoregt·Created 6 months ago·Updated 6 months ago· 1
Load FFTNet Model
    • fftnet_model
    checkpoint_path/path/to/fftnet_model.pt
    device
    vocab_size50257
    d_model512
    n_layers6
    d_ff2048
    max_seq_len512
    window_size4
    tokenizer_path

    Why this node exists in your image tool

    ComfyUI is for pictures. Load FFTNet Model is for text. It's one half of bemoregt/ComfyUI_FFTNet, a small two-node pack that runs a language model inside the image workflow - the kind of pack you only meet when a downloaded workflow is missing a node, or when you found the paper it's built on and got curious.

    The paper is "The FFT Strikes Back: An Efficient Alternative to Self-Attention" (arXiv:2502.18394, February 2025). The pitch: transformer attention is O(n²), which hurts as context grows. FFTNet swaps it for O(n log n) global token mixing done in the frequency domain. This node is the on-ramp for that model - it takes a trained checkpoint and turns it into a FFTNET_MODEL you can hand to the pack's generate node.

    Honest framing up front: this is research-adjacent, not a mainstream tool. There are no pretrained weights bundled and essentially no community footprint behind it. If you loaded this expecting an easy local chatbot, you came to the wrong node. If you have (or want to train) an FFTNet checkpoint, this is the only place I know of to run it inside ComfyUI.

    How it works

    Each FFTNet block splits into two branches. A global branch runs the sequence through rfft, multiplies it by a learned complex filter, applies a frequency-domain modReLU (gates the magnitude, keeps the phase), then irfft back to tokens. A local branch runs a depthwise 1D convolution to catch short-range structure. A sigmoid gate fuses the two, and the block wraps in pre-norm + GELU feed-forward - the usual transformer skeleton, just with attention replaced.

    The loader does two different jobs depending on what you point it at:

    • Standalone FFTNet - a .pt/.pth with a model_state_dict (plus optional config).
    • Upgraded LLaMA - an existing HuggingFace model (LLaMA, Mistral, Qwen2, Falcon…) whose self-attention was swapped for FFTNet by the pack's upgrade/ tooling. Those checkpoints carry an fftnet_upgraded marker and take a different code path.

    Config resolution is the part that will bite you. Priority is: checkpoint's embedded config dict > config.json sitting next to the file > your node inputs > defaults. Notice where your node inputs land. If the checkpoint already knows its hyperparameters, the numbers you type here are silently ignored - you can't override a baked-in config from the UI. If it doesn't know them, the six optional integer inputs are your only way to match the architecture. Get them wrong and load_state_dict(strict=False) quietly skips keys and prints warnings to the console; the model loads, the output is garbage.

    The inputs that actually matter

    Only two are required:

    • checkpoint_path - the .pt/.pth file. There's no model picker; it's a raw path string.
    • device - auto, cpu, cuda, mps. auto picks CUDA if present, then Apple's MPS, then CPU, so on a Mac you can just leave it.

    The optional vocab_size, d_model, n_layers, d_ff, max_seq_len, window_size are only for checkpoints without a config, and tokenizer_path lets you point at a HuggingFace tokenizer directory (otherwise the pack falls back to tiktoken's GPT-2 encoding, then a character-level tokenizer you don't want). Output is a single fftnet_model wire - that goes into FFTNet Generate, the other node in the pack.

    Installing it

    Easiest is ComfyUI Manager - search for ComfyUI_FFTNet and install, or:

    cd ComfyUI/custom_nodes
    git clone https://github.com/bemoregt/ComfyUI_FFTNet
    # then restart ComfyUI
    

    Its requirements.txt pulls torch>=2.0, plus tiktoken (the sensible default tokenizer), and transformers/datasets which are only needed for the upgraded-LLaMA path and HuggingFace tokenizers. If you skip transformers, standalone mode still works. Remember this installs into your shared Python environment - ComfyUI's custom nodes have no isolation, so a transformers pin here could in theory collide with another pack that wants a different version. That's ecosystem-wide, not specific to this node.

    Two gotchas worth knowing: the README's example checkpoint_path is a placeholder - there's no bundled model, and none of the paper's official weights are wired in, so expect to bring your own or train one (the pack ships upgrade/train.py for that). And if you loaded a workflow with this node and nothing else, that workflow came from someone who already had an FFTNet checkpoint on disk. You're replicating their setup, not downloading a model.

    CategoryFFTNet

    Inputs (9)

    NameTypeDefaultDescription
    checkpoint_pathSTRING/path/to/fftnet_model.pt
    deviceCOMBO4 options: auto, cpu, cuda, mps
    vocab_sizeoptINT50257256–256000
    d_modeloptINT51264–4096
    n_layersoptINT61–64
    d_ffoptINT204864–16384
    max_seq_lenoptINT51264–8192
    window_sizeoptINT41–64
    tokenizer_pathoptSTRING

    Outputs (1)

    NameTypeDescription
    fftnet_modelFFTNET_MODEL