DiTSampler

The whole pack leads here. The loaders get you a model, the label nodes get you a class, and DiTSampler is the bit that runs the denoising loop and produces an image. If you've used a KSampler in ComfyUI the shape is familiar - seed, steps, cfg, denoise - but this one is a different animal under the hood. It's not ComfyUI's sampler at all. It's the original diffusion loop from the Meta DiT repo, itself descended from OpenAI's guided-diffusion, ported in wholesale. It doesn't accept ComfyUI's scheduler or sampler dropdowns. Your steps value respaces the fixed 1,000-step linear schedule, and that's your entire control surface.

How it works

The model predicts noise plus learned variance (Meta's learn_sigma=True), the schedule is a plain linear beta curve from the DiT codebase, and classifier-free guidance runs inside the model's forward_with_cfg rather than in the sampler itself. That last bit changes what cfg means, and it's worth knowing: DiT applies guidance only to the noise-prediction channels and leaves the variance channels alone - a deliberate choice the Meta code comments on, made for exact reproducibility. So don't map your KSampler habits onto this 1:1. The default is 8, which is what the DiT repo ships with, and it runs hot; if things come out oversaturated, work down toward 2–4. A cfg of 0 disables guidance entirely, which mostly ignores your label.

The empty class is handled for you. The sampler appends the "empty" token - index num_classes, the extra row the Simple loader's math accounted for - as the unconditional half of the batch, then chunks the result in half at the end to discard it. You never supply it, which is exactly why DiTLabelSelect has no "None" option.

One confession lives in the source, and it's worth a warning. denoise interpolates your input latent toward fresh noise with a plain torch.lerp to fake img2img support, and the author left a comment on that line that just says "this is wrong." It's a crude lerp rather than a proper partial-noising pass, so treat img2img here as experimental. For class-to-image generation you'll keep it at 1.0 anyway.

The inputs that matter

• model - the DIT output from either loader.
• class_labels - the DITLAB output from DiTLabelSelect (or DiTLabelCombine for multiple classes).
• latent_image - and this is the trap. The pack wires the model's latent format to SD 1.5's, so this latent must come from VAEEncode using an SD 1.5 VAE, at exactly the resolution you told the loader: image_size 256 or 512, where the latent grid is image_size / 8. Wrong resolution, and you get a tensor shape error from the patch embedding at best, silent garbage at worst.
• seed, steps (default 20), cfg (default 8), denoise (default 1).

The output is a LATENT, decoded with the same SD 1.5 VAE via VAEDecode. Full graph: loader → sampler, label → sampler, VAEEncode → sampler → VAEDecode. Four nodes and an image.

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 stack, no extra pip packages. Model files go in ComfyUI/models/dit/ (.pt weights from the facebookresearch/DiT repo).

Where people get burned

Watch VRAM with DiT-XL/2 at 512 - the model itself is roughly SD 1.5-sized, but the token count at patch size 2 adds up fast, and the sampler holds the guidance batch (cond + null) in memory for the whole run. Check your image_size against your latent before blaming the sampler, and remember cfg here isn't your KSampler cfg. And the standing caveat applies hardest to this node: it's WIP, the README says to use ComfyUI_ExtraModels instead, and the author's own code has a "this is wrong" comment in it. Treat it as a curiosity, a history lesson, and a genuinely fun way to see what class-conditioned diffusion looked like before text took over.