ComfyUI Node

TF Load Pipeline

Where the JAX model, the compile, and the 2 GB downloads all land

By KorayUlusan·Created 3 days ago·Updated about 21 hours ago· 1
TF Load Pipeline
    • pipeline
    • info
    checkpointauto (download TF_L_edit)
    config
    warmuptrue

    Every TF workflow starts here. TF Load Pipeline loads the Trajectory Forcing flow model and its RAE decoder, and it's the one node in the pack where you genuinely wait - think minutes, not seconds, on a cold start. This is also where most of the pack's install trouble surfaces, so it's worth understanding what's actually happening behind the progress bar before you blame your machine.

    The checkpoint dropdown lists whatever sits in ComfyUI/models/trajectory_forcing/. The default is auto (download TF_L_edit), which fetches the released editing checkpoint (~2 GB) into that folder on first use. config is advanced and usually stays on edit_env_config.yml, which is the one that matches that checkpoint - the others are training and eval configs you don't want. Leave warmup on. Here's the thing warmup is buying you: the first sample after load triggers an XLA compile that takes one to two minutes, and a ComfyUI graph that sits idle that long reads as hung. With warmup on, the node samples and decodes one throwaway image at load, paying the compile and the decoder build where its own progress bar can show them. Turning it off doesn't save those minutes - it just moves them to your first TF Generate, which can't show progress for an opaque compile. People report a "hung" graph over exactly this.

    The pipeline output is a handle, not a thing you see. It's cached for the life of the ComfyUI process, so re-queueing never re-reads the checkpoint - that's why the node body reports "ready in Ns" on re-runs rather than re-loading. The info string tells you what you actually got: which config, which checkpoint, decoder path, number of levels (4), number of classes (1000), steps per level, and load time. Read it once; it's the fastest way to confirm you're on the intended checkpoint rather than whatever a config mixup loaded.

    Because the flow checkpoint and the RAE decoder download on first use, a brand-new install's first load is really three waits stacked: restoring the checkpoint, downloading the decoder (~1.6 GB) if it isn't already in TrajectoryForcing/checkpoints/rae/, then compiling. The progress bar labels each stage, and the node makes a point of sending a bar update before anything has happened, because the stage that most needs feedback is the one where the bar used to be silent.

    VRAM numbers worth having before you queue it, from the pack's own measurement script: ~2.5 GiB with the model loaded, ~4.6 GiB after the sampler compiles, ~6.6 GiB once the decoder is built. JAX and torch share the card and ComfyUI's memory management can't see JAX's allocations, so if something else is in the graph and you OOM, the fix is TF_XLA_MEM_FRACTION=0.3 (and a matching --reserve-vram for ComfyUI), not buying a bigger card. The stated floor is 8 GB VRAM, 12 GB comfortable, Linux or WSL2 only - jax[cuda12] has no native-Windows wheels.

    Install

    This is the install-heavy one, because everything else depends on it working. The honest summary: Manager → search Trajectory Forcing → Install sometimes just works (when install.py finds your torch compatible), and sometimes declines with a reason. Either way nothing you had is changed - requirements.txt is empty on purpose so the Manager can't rewrite your torch. When it declines, or when the nodes register but TF Load Pipeline errors with a missing jax/flax/orbax message, the README's separate-environment route is the answer:

    cd ComfyUI/custom_nodes/ComfyUI-TrajectoryForcing
    python -m tf_nodes.doctor   # asks the pack directly; each bad row names its fix
    

    The pack needs Python 3.11, CUDA 12 with torch 2.8+, and ~25 GB free (11 GB environment plus ~13 GB weights). If a row in doctor says the environment is fine and the problem is upstream, don't fight it - tf_nodes.doctor was written exactly so you don't have to guess.

    CategoryTrajectoryForcing/generate

    Inputs (3)

    NameTypeDefaultDescription
    checkpointCOMBOauto (download TF_L_edit)Flow checkpoint from models/trajectory_forcing/. 'auto (download TF_L_edit)' fetches the released editing checkpoint into that folder.
    configCOMBOTrajectoryForcing config. edit_env_config.yml is the one that matches the TF_L_edit checkpoint; the others are training/eval configs.
    warmupBOOLEANtrueSample and decode one throwaway image at load, paying the XLA compile and the decoder build here -- where this node's progress bar shows them -- instead of on your first real prompt. Turning it off does not save the 1-2 minutes, it moves them to the first TF Generate or TF Decode, which cannot show progress for a single opaque compile. Leave it on unless you are loading the model without intending to sample.

    Outputs (2)

    NameTypeDescription
    pipelineTF_PIPELINE
    infoSTRING