UIT Sampler

Here's the thing about HiDream-O1 that trips up every first-time user: it has no VAE and no separate text encoder. It's a pixel-space model, an 8B Pixel-level Unified Transformer (UiT) that encodes raw pixels, text, and task conditions in one shared token space. In a normal ComfyUI workflow that reads as a broken graph - no VAE loader, no CLIP loader, and the "missing component" error isn't your fault. This node exists to kill that whole class of confusion: one UIT Sampler node takes the checkpoint and the prompts and runs the entire generation pipeline internally.

Don't mistake the one-node convenience for a beginner toy. HiDream-O1 itself is the "mediocre generator, very good editor" model the community landed on - great at multi-reference editing and style transfer (no latent round-trip means edits don't accumulate compression loss), slow enough to cost minutes per image, and locked to ~2048×2048. UIT Sampler is the practical on-ramp to that architecture, and it also happens to be a nice way to see pixel-space sampling happen, thanks to the step_images output.

How it works

Under the hood this is a mini-pipeline that normally takes four nodes. It encodes your prompts through the CLIP object, applies the model's noise scale, attaches up to two reference images to the conditioning as reference_latents, builds a latent (or encodes your input image), runs the sampler with a callback that decodes every intermediate step, and hands you the final image plus all the steps.

The two things worth knowing mechanically:

• noise_scale is a real HiDream knob, not a hack. The node clones the model, rebuilds the model_sampling object with your shift/multiplier preserved, and calls set_noise_scale(). This is exactly what ComfyUI's ModelNoiseScale node does. Base checkpoint wants 8.0, Dev wants 7.5 - that's the author's own guidance and the default follows it.
• Input images get auto-rescaled to 4 MP (~2048×2048) via Lanczos, rounded to a multiple of 32. That matches the model's training resolution, which is the whole reason you're told to stop fighting for smaller sizes - O1 just doesn't do 1024 well.

The inputs that matter

The two that bite people are the ones that look optional:

• clip and vae - the schema marks them optional, but the node genuinely needs them. Wire the CLIP and VAE outputs straight out of your Load Checkpoint node into these. Leave them dangling and it will crash the moment it tries to encode your prompt or decode the result. "Dummy connection" in the README means the checkpoint carries them, not that you can skip them.
• positive_prompt / negative_prompt - plain text strings, or wire in a text node.
• input_image - optional; connect it for img2img or editing. When it's connected, width and height are ignored.
• reference_image1 / reference_image2 - up to two reference images for multi-reference editing (the model itself supports four; this node caps at two).
• noise_scale, steps, cfg, sampler (44 of ComfyUI's samplers), scheduler, seed, denoise - the usual suspects. Defaults are euler / normal / 12 steps / cfg 3.

Outputs: image (the final render) and step_images - every intermediate denoising step decoded and stacked into one batch tensor, which is a genuinely nice touch for debugging and for seeing the model "think."

Installing it

The easy way is ComfyUI Manager - search "UIT Sampler HiDream-O1" and install. Or manually:

cd ComfyUI/custom_nodes
git clone https://github.com/easygoing0114/ComfyUI-uit-hidream-o1.git

Restart ComfyUI and it shows up under sampling/uit. There's no requirements.txt and the pyproject lists no extra deps - this node is pure ComfyUI, no pip pain. The real download is the model: grab a HiDream-O1 checkpoint from the pack README's list (Comfy-Org's on Hugging Face is the recommended one), drop it where Load Checkpoint finds it, and you're set.

Troubleshooting

• Red "missing" VAE/CLIP: that's the model, not you - connect the checkpoint's own CLIP and VAE outputs to the node.
• Something's slow. Yes. This is the pixel-space tax: compute scales with pixel count, so expect minutes at 2048×2048 even on a strong card. Not a bug, not fixable by dropping steps to 4.
• Dev checkpoint and weird output: check noise_scale - 7.5 for Dev, not the 8.0 default. And if you're getting a subtle grid artifact pattern, that's a known Dev-model quirk, not the node.
• step_images eats VRAM: it decodes every step at full res into one batch. If you don't need the preview, leave it unplugged.