Nodes/comfyui-snek-nodes/Load AI Toolkit Latent Flux
ComfyUI Node

Load AI Toolkit Latent Flux

Hand a training latent back to ComfyUI

By sneccc·Created 2 years ago·Updated 12 months ago· 1
Load AI Toolkit Latent Flux
    • LATENT
    latent_file

    If you train Flux LoRAs with Ostris's AI Toolkit, you may have noticed it can save the latents from a training run as .safetensors files. There's rarely a good way to get one of those back into ComfyUI. This node is that bridge: it lists .safetensors files sitting in your input folder and loads one as a LATENT, so you can wire a training latent straight into a Flux workflow.

    Why you'd reach for it. Mostly for the odd, specific jobs: recreating or echoing a particular starting latent that AI Toolkit saved, feeding an img2img-style Flux run from a real training latent instead of empty noise, or diffing latents against one another in a debugging workflow. It's not a daily-driver node - but when you need it, the built-in loaders won't do it, because they load models or images, not raw latent tensors. AI Toolkit (the Ostris trainer the KB tracks as the Flux-era default for "how do I train a LoRA on this") stores latents under a tensor key literally named latent, which is exactly what this node knows how to find.

    How it works. Small and honest. On load it scans ComfyUI/input for files ending in .safetensors and offers them as a dropdown. When you pick one it opens the file with safetensors.safe_open(..., framework="pt"), reads the latent tensor, and if it's 3D it adds a batch dimension so you get a 4D tensor. Then it wraps it the way ComfyUI expects - {"samples": latent} - and hands it out as a LATENT. That's the entire mechanism, no model loading, no VAE involved.

    Inputs and outputs. The smallest surface of any node here:

    • latent_file (enum dropdown) - the .safetensors files currently in ComfyUI/input.
    • LATENT (output) - the latent, ready to plug into anything that accepts a latent: a KSampler, a LatentUpscale, a FluxGuidance-style conditioning setup, whatever.

    Installing it. Same pack as the rest of comfyui-snek-nodes:

    cd ComfyUI/custom_nodes
    git clone https://github.com/sneccc/comfyui-snek-nodes
    

    or search "comfyui-snek-nodes" in ComfyUI Manager, then restart. No extra model downloads - this node is pure file plumbing, so it doesn't pull AI Toolkit or any weights. The file itself has to come from a training run.

    Where people get burned.

    • The file has to be in input, not models. The dropdown only scans folder_paths.get_input_directory(). Drop your latent in ComfyUI/input, not next to your checkpoints. If the dropdown doesn't show it, that's the first thing to check.
    • It must contain a tensor named latent. AI Toolkit saves it under that exact key, so files straight out of training work. A random checkpoint file won't - it'll just print Error loading latent to the console and return None, which will silently break the rest of your graph.
    • Dimensions matter downstream. Flux latents are 16 channels at 1/8 resolution. If your downstream node expects a different latent size, the shapes won't line up and the sampler will complain. If you're loading a latent to continue from, make sure your EmptyLatentImage (or whatever precedes it) matches.
    • You have to refresh ComfyUI's input cache. Drop a new file in input while ComfyUI is running and the dropdown may not see it until you reload the frontend.

    Bottom line. It's a one-trick node, but the trick is otherwise awkward to do. If you live in AI Toolkit for Flux training and sometimes want your latents back in the graph, this is the cleanest way to get them there.

    Category🐍 Snek Nodes

    Inputs (1)

    NameTypeDefaultDescription
    latent_fileCOMBO0 options:

    Outputs (1)

    NameTypeDescription
    LATENTLATENT