DiTCheckpointLoaderSimple

If you're going to run any of Meta's ImageNet DiT models in ComfyUI, this is the loader you'll reach for. DiTCheckpointLoaderSimple does exactly what DiTCheckpointLoader does minus one knob - and the knob it removes is the one that trips people up. Instead of making you type num_classes and hope you got it right, it reads the answer straight off the checkpoint.

The pack context matters here, so one paragraph: this is city96's WIP lab bench for class-conditional DiT, the README points you to ComfyUI_ExtraModels as the maintained successor, and "simple" in the name is relative - the whole thing still only speaks ImageNet class labels, never text.

How it works

The loader peeks at the state dict's y_embedder.embedding_table.weight tensor and counts its rows. That count is the number of classes plus one - the extra slot is the "empty" token DiT reserves for the unconditional pass during classifier-free guidance - so it subtracts one and it's done. That's the entire trick, and it's why stock checkpoints just work: the loader can't be wrong, because it's counting from the exact file you loaded rather than trusting you to remember. It also unwraps a "model" key if the file wraps its weights, so both raw and wrapped state dicts load cleanly.

The inputs that matter

• ckpt_name - the file in ComfyUI/models/dit/.
• model - the twelve-entry architecture list, DiT-XL/2 down to DiT-S/8. Letter is model size, number is patch size.
• image_size - 256 or 512, and it must match what the checkpoint was trained at.

The names on HuggingFace encode all of this for you: facebook/DiT-XL-2-256 is XL, patch 2, 256px. Match the model and image_size fields to the file name and you're set.

What it feeds

The single output, model (DIT), goes into DiTSampler. Your full graph is tiny: this loader plus DiTLabelSelect for a class, a VAEEncode with the SD 1.5 VAE (the model's latent format is SD 1.5's, at a resolution matching image_size), and the sampler in the middle.

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

No requirements.txt - this pack runs on ComfyUI's own dependencies and nothing else, which is refreshing in a custom-node ecosystem where half your installs are dependency tetris. The model files are the actual download: .pt weights from the facebookresearch/DiT repo, dropped into ComfyUI/models/dit/.

Where people get burned

Point it at a non-DiT checkpoint and it dies looking for y_embedder - that key is the fingerprint of a class-conditional DiT, and nothing else has one. Pick the wrong model size or image_size and you get a load_state_dict shape error. And it will happily start loading a 256 checkpoint even if you set image_size to 512, right up until the position embeddings disagree and the whole thing falls over. Match the three fields to the file and the loader basically can't hurt you.