Load Checkpoint with FLATTEN model

A checkpoint loader that isn't loading what you think

Load Checkpoint with FLATTEN model looks like the standard ComfyUI checkpoint loader - same ckpt_name dropdown, same MODEL, CLIP, VAE outputs - but it isn't doing the same job. It loads your SD1.5 checkpoint with FLATTEN's 3D UNet in place of the normal 2D one. That one swap is what lets a video of frames flow through a model that was only ever trained on stills.

FLATTEN (optical FLow-guided ATTENtion, CVPR 2024) edits video with temporal consistency by making the UNet aware of motion across frames. The 3D UNet is the structural half of that: where the standard UNet reshapes a batch of frames into flat 2D latents, this one keeps a frame dimension and lets the flow-guided attention (see Apply Flatten Attention) look along it. You'd reach for this node whenever you're running the pack's full pipeline - load the checkpoint here, feed it into the Unsampler and KSampler, and your video frames get the 3D treatment the rest of the graph assumes.

How it works

The node temporarily swaps ComfyUI's model_base.BaseModel for a patched subclass whose UNet is FLATTEN's UNetModel, then loads the checkpoint through the usual guess-config path - so your ordinary SD1.5 safetensors work, with no conversion. It wraps the model's forward pass to rearrange the latent from (b f) c h w into b c f h w, restoring the frame dimension that the 3D blocks need, and it handles FLATTEN's context-windowed batching for you (grabbing the right trajectory window based on AnimateDiff's sub_idxs).

The one input

ckpt_name - a dropdown of everything in your models/checkpoints folder. Pick any SD1.5 checkpoint.

Outputs are MODEL, CLIP, and VAE, wired exactly like the core loader. The catch, and the README says it plainly: the loaded model only works with the Flatten KSampler. If you want to sample with a standard KSampler, load a checkpoint with the ordinary loader instead. This node's model and the standard sampler don't speak the same language.

Gotchas

It's SD1.5 only, like the whole pack. There's an sdxl branch, and the author is upfront that SDXL experiments produced "somewhat consistent videos" but were never up to par, with no further development planned. So treat any "FLATTEN on SDXL" tutorial you find as archaeology. Also note the author: logtd is the same person behind ComfyUI-InstanceDiffusion, ComfyUI-OtherVAEs and ComfyUI-Fluxtapoz - a known quantity in the ComfyUI port scene, which counts for something when you're running code that monkeypatches your sampler internals.

Installing

Search ComfyUI-FLATTEN in ComfyUI Manager, or:

cd ComfyUI/custom_nodes
git clone https://github.com/logtd/ComfyUI-FLATTEN

Then restart ComfyUI. The pack's requirements.txt is empty - nothing beyond ComfyUI's bundled torch/torchvision, since the optical flow comes from torchvision's built-in RAFT. You bring your own SD 1.5 checkpoint.