DiTCheckpointLoader

Between the two checkpoint loaders in city96's DiT pack, this is the one you'll reach for almost never. DiTCheckpointLoader is the same model browser as DiTCheckpointLoaderSimple plus one extra integer - num_classes - and for every stock Meta checkpoint you're likely to grab, the Simple version reads that number correctly on its own. So why does the full one exist? As the escape hatch. If you ever run a DiT that was trained on a different label count than the standard 1,000 (people do fine-tune these), a wrong count makes the weights fail to load, and this is the loader that lets you speak up.

First, the pack-level honesty, because it applies to all five nodes here. This is a WIP test platform. The README opens with "Old repo, use ExtraModels instead for DiT support" and warns that none of the code is stable, so expect breaking changes. It's city96's - the same person behind ComfyUI-GGUF, the pack that turned quantized Flux into the standard low-VRAM path - but this repo is an early lab bench for Meta's class-conditional ImageNet DiT models, not a production text-to-image tool. You're here to poke research checkpoints, and that's a fine reason to be here. Just don't build anything permanent on it.

How it works

Pick a file from ComfyUI/models/dit/ (the pack registers that folder the first time it loads), pick an architecture name from a list of twelve, and the node constructs the matching DiT, loads the state dict, and hands you a ComfyUI model patcher. It transparently unwraps a "model" key if the file wraps its weights that way, so both raw state dicts and wrapped ones work.

The architecture list is the part to understand. DiT-XL/2 through DiT-S/8: the letter is the model size, the number after the slash is the patch size. Smaller patch means the image gets chopped into more tokens, which means more compute and VRAM. DiT-XL/2 at 512 is the flagship and a half; DiT-S at 256 is tiny and quick.

The inputs that matter

• ckpt_name - the file sitting in ComfyUI/models/dit/.
• model - must match the checkpoint; mismatched size and load_state_dict fails with a shape error.
• image_size - 256 or 512. The latent grid is image_size / 8, and the position embeddings were baked for one resolution, so this has to line up with what the checkpoint was trained at.
• num_classes - default 1,000. This is the knob that distinguishes this loader from Simple. Leave it alone for stock models; change it for custom-trained ones.

The single output, model (DIT), plugs straight into DiTSampler. From there the rest of the graph is DiTLabelSelect for a class, a VAEEncode with the SD 1.5 VAE, and the sampler - which expects latents at exactly the image_size you chose here.

Installing it

# ComfyUI Manager → Install Custom Nodes → search "ComfyUI_DiT"
# or manually:
cd ComfyUI/custom_nodes
git clone https://github.com/city96/ComfyUI_DiT
# then restart ComfyUI

There's no requirements.txt - the pack runs entirely on ComfyUI's own stack, so no pip dance and no dependency conflicts to untangle. The model files are the real download: grab the .pt weights from the facebookresearch/DiT repo on HuggingFace (facebook/DiT-XL-2-256 and friends) and drop them into ComfyUI/models/dit/ (create it if it's not there).

Where people get burned

Both classic failures are shape mismatches. Wrong model choice or wrong image_size, and the state dict won't line up. Wrong num_classes, and the embedding table doesn't fit. None of these give you a friendly message - you get a long torch traceback about sizes. And remember the downstream chain: this model wants SD 1.5 VAE latents, so if your VAEEncode is on the wrong VAE or the wrong resolution, the failure will look like it's the sampler's fault when it's actually upstream of it.