feat: Support NVIDIA PixelDiT and PiD (CORE-201)

[kijai]

Adds support for Nvidia PixelDiT T2I image model, as well as the new PiD models.

Models (nsclv1 license):

https://huggingface.co/Comfy-Org/PixelDiT

---

PixelDiT text to image

pixeldit_test_01.json

---

PiD

Encode - decode example:

pid_512-2048_flux1_upscale_example_02.json

Z-image to 4096 example:

z_image_turbo_to_pid_03.json

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Run ID: 88ec5108-8a46-4d54-a2cc-0210aa722af0

📥 Commits

Reviewing files that changed from the base of the PR and between adcf14d and 81ba159.

📒 Files selected for processing (1)

• comfy/ldm/pixeldit/pid.py

🚧 Files skipped from review as they are similar to previous changes (1)

• comfy/ldm/pixeldit/pid.py

---

📝 Walkthrough

Walkthrough

This pull request adds PixelDiT, a pixel-space multimodal diffusion model, and PiD, a pixel-diffusion decoder variant, to ComfyUI. The implementation includes new pixel- and patch-level transformer modules, a Gemma2-2B-based text encoder/tokenizer, PiD low-quality latent injection and gating, updates to model detection/registration, model wrappers, and a ComfyUI node for PiD conditioning.

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 5.71% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title 'feat: Support NVIDIA PixelDiT and PiD (CORE-201)' accurately describes the main change—adding support for two NVIDIA image models—and is clear and specific.
Description check	✅ Passed	The description is directly related to the changeset, referencing the NVIDIA PixelDiT and PiD models being added, providing context about licensing, and including workflow examples.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[iperide]

Hi, quick question about NVIDIA licensing.
For the new PixelDiT model on HuggingFace, NVIDIA explicitly lists NSCLv1 as the license, which allows redistribution as long as it remains non‑commercial:

redistribution allowed under the same license
derivative works allowed
non‑commercial use only

I didn’t find the restrictive clause (“may not be distributed, deployed…”) on the HF page or in the LICENSE file.

Since NVIDIA sometimes applies multiple overlapping licenses to their research models, could you clarify whether we should always assume the more restrictive NVIDIA Research license applies, even when the HuggingFace page only shows NSCLv1?

Just trying to understand the policy you follow for ComfyUI integration.

Thanks.

[jprsyt5]

Really want to try PiD, but I can't get access to the TE model yet.
Does it also work with Gemma 2 2B from lumina repackaged? https://huggingface.co/Comfy-Org/Lumina_Image_2.0_Repackaged/tree/main/split_files/text_encoders

[kijai]

Hi, quick question about NVIDIA licensing. For the new PixelDiT model on HuggingFace, NVIDIA explicitly lists NSCLv1 as the license, which allows redistribution as long as it remains non‑commercial:

redistribution allowed under the same license
derivative works allowed
non‑commercial use only

I didn’t find the restrictive clause (“may not be distributed, deployed…”) on the HF page or in the LICENSE file.

Since NVIDIA sometimes applies multiple overlapping licenses to their research models, could you clarify whether we should always assume the more restrictive NVIDIA Research license applies, even when the HuggingFace page only shows NSCLv1?

Just trying to understand the policy you follow for ComfyUI integration.

Thanks.

It's two different repos, the new PiD model uses the PixelDiT as base, this PR implements both.

PixelDiT license: https://huggingface.co/nvidia/PixelDiT-1300M-1024px/blob/main/LICENSE

3.1 Redistribution. You may reproduce or distribute the Work only if (a) you do so under this license, (b) you include a complete copy of this license with your distribution, and (c) you retain without modification any copyright, patent, trademark, or attribution notices that are present in the Work.

So okay to repackage and share when license is included.

---

PiD license however is this: https://www.nvidia.com/en-us/agreements/enterprise-software/nvidia-internal-scientific-research-and-development-model-license/

3.1 The Model and any Derivative Model may not be distributed, deployed, sublicensed, publicly displayed, publicly performed, or sublicensed by You. You may not use the Model or a Derivative Model in a production environment or for the purpose of generating works for sale or distribution.

Edit: They have now changed the license to match PixelDiT: https://huggingface.co/nvidia/PiD/commit/b87dba45e5a2b2a18bac9515fca883f52b957558

[kijai]

Really want to try PiD, but I can't get access to the TE model yet. Does it also work with Gemma 2 2B from lumina repackaged? https://huggingface.co/Comfy-Org/Lumina_Image_2.0_Repackaged/tree/main/split_files/text_encoders

I've made the repo public now: https://huggingface.co/Comfy-Org/PixelDiT

[jprsyt5]

Really want to try PiD, but I can't get access to the TE model yet. Does it also work with Gemma 2 2B from lumina repackaged? https://huggingface.co/Comfy-Org/Lumina_Image_2.0_Repackaged/tree/main/split_files/text_encoders

I've made the repo public now: https://huggingface.co/Comfy-Org/PixelDiT

Thank you!

Just tested it, but I'm still a bit confused.

Is PiD sensitive to the dimensions/image size?

I tried with a vertical image, the source is 1088×1440. Then I downscaled it by 0.5 using the Upscale Image By node.

After that, I used another Image Resize node and set it to 2048×2048 with the method set to keep proportion, and used the height & width output from this node to adjust the latent width/height, but I got weird results.

If I increase it to something like 2880×2880, the image looks fine.

I also noticed it always adds a green color at the bottom part of image, even though I'm not seeing that in your example results?

[kijai]

I added better explanation for the different models, it's important to choose the correct one closest to your resolution, their naming is a bit confusing.

[TheNeObr]

I wrote my own implementation, but for some reason yours just doesn't work...
FLUX
FLUX2

Any idea what might be causing this problem?

[kijai]

Any idea what might be causing this problem?

Input image is too large maybe? That selected model works with 512x512 inputs.

Also make sure to use the ELM variant of the Gemma2... at least with the PixelDit T2I the standard Gemma2 gives broken outputs.

[ssugar008-maker]

I have used the same workflow - [z_image_turbo_to_pid_02.json] but kept returning with error on the incorrect dimensions, e.g. Even after the PID conditioning, it does not work.
weight [512, 128, 3, 3] → PiD expects 128 channels
input [1, 16, 256, 256] → Z-Image latent has 16 channels

Not sure what ight be the exact issue, thanks

[kijai]

I have used the same workflow - [z_image_turbo_to_pid_02.json] but kept returning with error on the incorrect dimensions, e.g. Even after the PID conditioning, it does not work. weight [512, 128, 3, 3] → PiD expects 128 channels input [1, 16, 256, 256] → Z-Image latent has 16 channels

Not sure what ight be the exact issue, thanks

Z-image uses flux1 VAE, do you have that selected in the conditioning node?
And the PiD model itself being the flux1?

[ssugar008-maker]

I have used the same workflow - [z_image_turbo_to_pid_02.json] but kept returning with error on the incorrect dimensions, e.g. Even after the PID conditioning, it does not work. weight [512, 128, 3, 3] → PiD expects 128 channels input [1, 16, 256, 256] → Z-Image latent has 16 channels
Not sure what ight be the exact issue, thanks

Z-image uses flux1 VAE, do you have that selected in the conditioning node? And the PiD model itself being the flux1?

Thank you for the prompt reply. Yes, Flux1 VAE is selected in the conditioning node (latent_format = flux1) and the PID model itself being flux1 (e.g. [PiD_res2kto4k_sr4x_official_flux_distill_4step]) but same issue; perhaps I'll reinstall ComfyUI for a re-test thanks!

[kijai]

This PR adds preliminary context window support, so we can do sliding tiling on 1D axis to avoid the aspect ratio issues of the model:

#14136 (comment)

[jtreminio]

Heads' up for anyone following along: the PiD Conditioning.latent_format dropdown options were changed to simply flux and sd3, removing flux1 and flux2

[kijai]

Heads' up for anyone following along: the PiD Conditioning.latent_format dropdown options were changed to simply flux and sd3, removing flux1 and flux2

Yes, sorry for not mentioning that here, for the actually merged version it's just "flux" and "sd3" currently since it's easy to autodetect flux2, this makes the node automatic for anything but sd3.

[jtreminio]

Thanks @kijai - I believe the attached JSON workflows from first post still use the old values.

[yeyingxian]

Hi @kijai . Why PiD for flux2, will make the picture turn whitish? I use ERNIE-Image to generate the origin picture

[kijai]

Hi @kijai . Why PiD for flux2, will make the picture turn whitish? I use ERNIE-Image to generate the origin picture

It's just how the model is, I don't know why, could be the distillation (undistilled models are not released yet) or just their training in general.. all their examples are bright scenes too. The ComfyUI implementation matches what I get from their code:

[Heliumrich]

What would have the most "real" details, generating at 1536x1536 natively, or using PiD to get from 1024x1024 to 4096x4096 ?
Not talking of overall sharpness, but having real fine details that are not sharpened random blobs

[kijai]

What would have the most "real" details, generating at 1536x1536 natively, or using PiD to get from 1024x1024 to 4096x4096 ? Not talking of overall sharpness, but having real fine details that are not sharpened random blobs

I suppose it depends on the model, but probably the best use is to generate at given models native resolution, ending it early and passing the latent with the noise to PiD, then using the degrade_sigma to finish and decode the image. I don't know the optimal workflow for this yet, but seems to work great with Z-image at least.

[yifanlu0227]

Hey guys, we have a new checkpoint for FLUX.2 2kto4k that solves the color drifting problem. We will release the fixed checkpoint soon!!

Thanks for merging PiD to ComfyUI!!

[seghier]

Is this available in Comfyui studio?

[seghier]

How you get the final result and the model input is disconnected?

[iminternet]

Hey guys!

Is there a working method for Flux 2 yet?

[sajmon83]

I have used the same workflow - [z_image_turbo_to_pid_02.json] but kept returning with error on the incorrect dimensions, e.g. Even after the PID conditioning, it does not work. weight [512, 128, 3, 3] → PiD expects 128 channels input [1, 16, 256, 256] → Z-Image latent has 16 channels

Not sure what ight be the exact issue, thanks

In VAE Decode after SamplerCustom you use pixel_space, not FLAX1, FLAX2 VAE

[MoRanYue]

I got very blurry result when using Flux.2 VAE, the image was just like not fully denoised.
Ok, when it received condtioned latent larger than it can deal with, then blurry will appear.

[Astral793]

I get the output image, but it is pretty much the same as the input. I tried upscaling a 1024x1024 image using pid_flux1_1024_to_4096_4step_bf16.safetensors, but got this error:

[INFO] got prompt
[ERROR] Failed to validate prompt for output 48:
[ERROR] * (prompt):
[ERROR] - Required input is missing: images
[ERROR] * PreviewImage 48:
[ERROR] - Required input is missing: images
[ERROR] Output will be ignored
[INFO] Model PixelDiTTE_ prepared for dynamic VRAM loading. 4986MB Staged. 0 patches attached. Force pre-loaded 105 weights: 472 KB.
[MultiGPU Runtime] Using runtime device cuda:0 (comfy.sample.sample_custom:ModelPatcherDynamic)
[INFO] Model PiD prepared for dynamic VRAM loading. 2597MB Staged. 0 patches attached. Force pre-loaded 150 weights: 207 KB.
0%| | 0/4 [00:00<?, ?it/s, Model Initializing ... ]aimdo: src/model-vbar.c:74:WARNING:VBAR 000002E34BCB89A0: Page 4 pin_count=2
25%|██████████████████████████████████▌ | 1/4 [00:02<00:08, 2.95s/it, Model Initialization complete! ]aimdo: src/model-vbar.c:74:WARNING:VBAR 000002E34BCB89A0: Page 4 pin_count=2
50%|██████████████████████████████████████████████████████████████████████████████████████▌ | 2/4 [00:05<00:05, 2.87s/it]aimdo: src/model-vbar.c:74:WARNING:VBAR 000002E34BCB89A0: Page 4 pin_count=2
75%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████▊ | 3/4 [00:08<00:02, 2.92s/it]aimdo: src/model-vbar.c:74:WARNING:VBAR 000002E34BCB89A0: Page 4 pin_count=2
100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 4/4 [00:11<00:00, 2.88s/it]
[INFO] Requested to load PixelspaceConversionVAE
[INFO] loaded completely; 0.00 MB loaded, full load: True
aimdo: src/model-vbar.c:74:WARNING:VBAR 000002E34BCB89A0: Page 4 pin_count=2
[INFO] Prompt executed in 13.91 seconds

[titusfx]

Im using the latest version desktop, 0.22.3

Node 'ID #8' has no class_type. The workflow may be corrupted or a custom node is missing.
Node 'ID #8' has no class_type. The workflow may be corrupted or a custom node is missing.

[AIisCool]

Perhaps I'm doing something wrong, but the output is not as good as expected?

Note that I resized the input to 1024x1024 for processing, and just resized it back to the original input size for comparison.

Using the workflow here: #14103 (comment) with flux1 upscaling from 1024 to 4096

[jprsyt5]

Im using the latest version desktop, 0.22.3

Node 'ID #8' has no class_type. The workflow may be corrupted or a custom node is missing.

Use git and the master branch if you want the latest fixes/features immediately.

And since you're using Desktop, wait a few extra days for it to catch up.

[yifanlu0227]

Hi @kijai,
Just a quick heads-up: we’ve dropped new checkpoints for SDXL, Qwen-Image, plus an updated FLUX2 with the color bugs fixed!

We name the new flux2 checkpoint with 2606 suffix. Would you mind updating / replacing your huggingface model with the new one?

[kayselmecnun]

I have used the same workflow - [z_image_turbo_to_pid_02.json] but kept returning with error on the incorrect dimensions, e.g. Even after the PID conditioning, it does not work. weight [512, 128, 3, 3] → PiD expects 128 channels input [1, 16, 256, 256] → Z-Image latent has 16 channels

Not sure what ight be the exact issue, thanks

You are very likely downloaded the flux.2 diffusion model instead of flux.1. so make sure it starts with "pid_flux1...", not "pid_flux_2..."

[dnelbggr00]

does anybody has the workflow for FLUX2 ?

[luisclement]

Is there no way to use this with a Load Image first node and not generate t2i first, just to upscale? Great work as usual @kijai

[jorismak]

There is, this just works on an empty chromaradiance latent, and the conditioning node uses a latent as input.
the output latent is 'pixel space'.

So, you can build it up 'just to upscale' quite like a normal img2img worfklow.

Note that there are different models for flux1 input latent (and zimage/zimageturbo input latent), flux2 input latent, sd3.5 input latent.. (and sdxl and qwen coming or already there?).

And there are models made for 512 -> 2048 upscaling, and 1024 -> 4096 upscaling. If you go stray from those, your results may vary (but can also work ok).

for instance, using the flux1 model, 512 to 2048 version:

1. load 512x512 image, load flux1 vae.
2. vae encode the image to a flux1 latent.
3. load diffusion model, load clip.
4. add clip-text-encode to enter a prompt (just a positive prompt).
5. after the prompt, add pid conditioning with the flux1 latent you created earlier, and set it to flux mode.
6. from the prompt again, add 'conditioning zero out'.
7. Add a ksampler
8. Connect pid-conditioning to the ksampler in the positive slot.
9. Connect the 'conditioning zero out' to the negative slot of the ksampler, and set cfg to 1 (cfg = 1 means no CFG scaling so no negative prompt).
10. Set sampler to 'lcm', set scheduler to 'simple', set 4 steps.
11. Connect a 2048x2048 'EmptyChromaRadianceLatentImage'.
12. The output goes to VAE Decode into preview-image / save-image / whatever.
13. Add a 'load vae' and set it to 'pixel_space' and not a real Vae model! connect it to the 'vae decode' of the output.

Seems to work better on an ultra clean output of a diffusion model than a real image though.
It's basically meant as an 'upscaling vae', it takes the part of decoding vae to pixel space, just it does it with a diffusion model instead of 'a simple vae'.
This means that imperfections, noise, compression artifacts, etc... is absolutely NOT what this model was trained on.

for flux2, just use a flux2 latent as input. PidConditioning set to flux-mode will autodetect which to use.

Most people getting size mismatches from 16 -> 128 or something similar, are either loading the wrong Pid-model for their input latent, or are trying to vae-decode the Pid output with a real vae, instead of picking 'pixel_space'.