nodes/Merserk-ComfyUI-PiD
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Add files via upload

Browse Source
This commit is contained in:
Merserk
2026-05-25 14:26:50 +03:00
committed by GitHub
parent 75c65cfc77
commit b4edd87012
9 changed files with 1965 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.gitignore
+4
View File
@@ -0,0 +1,4 @@
__pycache__/
*.pyc
vendor/PiD/
vendor/PiD-main.zip
CHANGELOG.md
+17
View File
@@ -0,0 +1,17 @@
# Changelog
## VRAM offload lean build
- Removed the separate automatic settings module from the package.
- Node registration now exposes only **PiD Decode** and **PiD Text Prompt**.
- Updated README to document manual **PiD Decode** settings only.
- Kept the VRAM offload behavior from the previous build: CPU staging before PiD, Comfy model unload before PiD, and PiD unload after decode.
## VRAM offload build
- Before loading PiD, `PiD Decode` now detaches the input latent and baseline image to CPU.
- `PiD Decode` clears Python references to upstream latent, VAE, and baseline tensors before ComfyUI model unload.
- `_free_cuda_memory()` now also runs Python garbage collection before Comfy/CUDA cache cleanup.
- After a successful PiD decode, the PiD model is removed from this node's private cache and moved off CUDA so the next Z-Image generation does not keep PiD in VRAM.
- The same PiD unload cleanup is attempted after PiD inference errors.
- Added an example workflow that sends `VAEDecode IMAGE` into `PiD Decode baseline_image` and leaves `PiD Decode vae` disconnected.
README.md
+178
View File
@@ -0,0 +1,178 @@
# ComfyUI-PiD
Experimental ComfyUI custom node for using NVIDIA **PiD** as a pixel diffusion decoder.
PiD is not a normal `.safetensors` VAE. NVIDIA's official path gives PiD a latent, the native decoder/VAE baseline image, a sigma value, and the prompt:
```text
LATENT + baseline IMAGE + prompt + sigma -> PiD -> IMAGE
```
The node can create the baseline through a connected ComfyUI `VAE`, or you can connect a pre-decoded `baseline_image` from another workflow component.
## Node
- **PiD Decode**: decodes a PiD-supported latent and outputs `IMAGE`.
- **PiD Text Prompt**: one prompt box with outputs for both CLIP text encoding and PiD caption conditioning.
There are no separate setup/download/unload nodes. PiD source, checkpoints, and required asset files are prepared lazily when **PiD Decode** runs with `auto_download=true`.
## Supported Backbones
The backbone list follows NVIDIA's official PiD checkpoint registry:
| Node value | Official backbone | Latent channels | Checkpoints |
| --- | --- | ---: | --- |
| `zimage` | Z-Image, reuses Flux weights | 16 | `2k`, `2kto4k` |
| `flux` | Flux | 16 | `2k`, `2kto4k` |
| `flux2` | Flux2 | 128 | `2k`, `2kto4k` |
| `sd3` | Stable Diffusion 3 | 16 | `2k`, `2kto4k` |
| `dinov2` | DINOv2 RAE | 768 | `2k` |
| `siglip` | SigLIP Scale-RAE | 1152 | `2k` |
`scale=0` means "use NVIDIA's default scale for that checkpoint": 4x for Flux, Flux2, SD3, Z-Image, and DINOv2; 8x for SigLIP Scale-RAE.
## Install From GitHub
Clone this custom node into ComfyUI:
```bash
cd ComfyUI/custom_nodes
git clone https://github.com/Merserk/ComfyUI-PiD.git
cd ComfyUI-PiD
python -m pip install -r requirements.txt
```
Restart ComfyUI.
Do not download PiD model weights during install. The node downloads only what it needs the first time you run it.
## Automatic Downloads
With `auto_download=true`, **PiD Decode** will:
1. Clone NVIDIA's PiD source into `vendor/PiD` if it is missing.
2. Download the selected PiD checkpoint from `nvidia/PiD`.
3. Download the extra tokenizer/VAE assets needed by the selected backbone.
For Flux/Z-Image, the node first tries to copy your existing:
```text
ComfyUI/models/vae/ae.safetensors
```
to:
```text
ComfyUI/custom_nodes/ComfyUI-PiD/vendor/PiD/checkpoints/ae.safetensors
```
If no local file is found, it downloads `checkpoints/ae.safetensors` from `nvidia/PiD`.
If Hugging Face requires a token, set one before starting ComfyUI:
```bash
export HF_TOKEN=hf_your_token_here
```
Windows PowerShell:
```powershell
$env:HF_TOKEN = "hf_your_token_here"
```
## Basic Workflow
For Z-Image/Flux-style workflows:
```text
PiD Text Prompt text -> CLIP Text Encode
PiD Text Prompt caption -> PiD Decode caption
KSampler LATENT -> VAEDecode -> PiD Decode baseline_image
KSampler LATENT -> PiD Decode latent
PiD Decode image -> Save Image
```
For the lowest VRAM peak, prefer connecting a pre-decoded `baseline_image` and leave the optional `vae` input on **PiD Decode** disconnected. The direct `VAE -> PiD Decode vae` path still works, but pre-decoding the baseline image makes the PiD-only stage easier to isolate.
Recommended first test settings on **PiD Decode**:
```text
backbone = zimage
pid_ckpt_type = 2k
pid_steps = 4
scale = 1 or 2
cfg_scale = 1.0
sigma = 0.0
auto_download = true
unload_comfy_before_pid = true
aggressive_cleanup = true
```
For backbones where the matching VAE is not available as a ComfyUI `VAE`, connect a pre-decoded `baseline_image` instead.
### VRAM offload behavior
This version is optimized so **PiD Decode** keeps only CPU copies of the input latent and baseline image before the PiD stage. It then asks ComfyUI to unload the previously loaded Z-Image/CLIP/VAE models and clears CUDA cache before loading PiD. After the PiD image is converted back to a ComfyUI `IMAGE`, the PiD model is also removed from this node's private cache and moved off CUDA.
For the lowest VRAM peak, prefer this workflow shape:
```text
KSampler LATENT -> VAEDecode -> PiD Decode baseline_image
KSampler LATENT -> PiD Decode latent
PiD Text Prompt caption -> PiD Decode caption
Do not connect VAE -> PiD Decode vae when baseline_image is already connected.
```
The old direct `VAE -> PiD Decode vae` path still works, but pre-decoding the baseline image makes the PiD-only stage easier to isolate.
## Notes
1. This is a practical ComfyUI wrapper around NVIDIA's public PiD code, not an official NVIDIA or ComfyUI node.
2. This node outputs `IMAGE`, not a ComfyUI `VAE`, because PiD is a conditional pixel diffusion decoder.
3. NVIDIA's best generated-image demos use captured intermediate latents, for example Z-Image around step 46 of 50. A final ComfyUI latent with `sigma=0.0` can work, but it is not identical to the official capture recipe.
4. PiD currently expects CUDA and significant VRAM, especially with `2kto4k` and high output scales.
5. NVIDIA's PiD weights have their own license/terms. Check the Hugging Face model card before using them.
## Troubleshooting
### Widget values shift after switching browser tabs
This build uses ComfyUI's canonical combo syntax for `backbone` and `pid_ckpt_type`. If a workflow was already saved while the node was broken, delete and re-add the **PiD Decode** node once, or manually set the values again after installing this build.
### Missing dependencies
Install the custom node requirements and restart ComfyUI:
```bash
python -m pip install -r requirements.txt
```
### Missing checkpoint or asset
Set `auto_download=true` on **PiD Decode**. The node downloads the selected checkpoint and the required VAE/tokenizer assets on first run.
### Output looks wrong
Check that:
- `backbone` matches the latent you are feeding into PiD.
- The latent channel count matches the selected backbone.
- Your connected `VAE` can decode that latent, or `baseline_image` is the correct native baseline image.
- Start with `sigma=0.0`, `cfg_scale=1.0`, `pid_steps=4`, and `scale=1` or `2`.
### VRAM or cudaMallocAsync errors
This build unloads ComfyUI models before PiD and unloads PiD again after the decode. If VRAM is still high, check the workflow first: connect a pre-decoded `baseline_image` into **PiD Decode** and disconnect the optional `vae` input on **PiD Decode**.
For 16GB GPUs, start with one of these:
- 1024x1024 base latent + `pid_ckpt_type=2k` + `scale=1` or `scale=2`
- 512x512 base latent + `pid_ckpt_type=2k` + `scale=4`
Avoid 1024x1024 base latent + `pid_ckpt_type=2kto4k` + `scale=4` on 16GB cards. That asks PiD to generate roughly 4096x4096 inside ComfyUI and commonly triggers CUDA allocator/VRAM failures.
If a CUDA allocator internal assert occurs, restart ComfyUI before trying again. The CUDA process can remain unstable after that error.
__init__.py
+3
View File
@@ -0,0 +1,3 @@
from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS
__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]
examples/image_z_image_pid_baseline_offload.json
+792
View File
@@ -0,0 +1,792 @@
{
"id": "9ae6082b-c7f4-433c-9971-7a8f65a3ea65",
"revision": 0,
"last_node_id": 127,
"last_link_id": 199,
"nodes": [
{
"id": 9,
"type": "SaveImage",
"pos": [
1542.0543179682477,
-828.1065784841086
],
"size": [
730,
570
],
"flags": {},
"order": 12,
"mode": 0,
"inputs": [
{
"name": "images",
"type": "IMAGE",
"link": 198
}
],
"outputs": [],
"properties": {
"cnr_id": "comfy-core",
"ver": "0.3.64",
"Node name for S&R": "SaveImage",
"enableTabs": false,
"tabWidth": 65,
"tabXOffset": 10,
"hasSecondTab": false,
"secondTabText": "Send Back",
"secondTabOffset": 80,
"secondTabWidth": 65
},
"widgets_values": [
"z-image-base"
]
},
{
"id": 98,
"type": "CLIPTextEncode",
"pos": [
89.82489218494345,
-635.688141888832
],
"size": [
557.03125,
189.28125
],
"flags": {},
"order": 7,
"mode": 0,
"inputs": [
{
"name": "clip",
"type": "CLIP",
"link": 126
},
{
"name": "text",
"type": "STRING",
"widget": {
"name": "text"
},
"link": 144
}
],
"outputs": [
{
"name": "CONDITIONING",
"type": "CONDITIONING",
"links": [
128
]
}
],
"title": "CLIP Text Encode (Positive Prompt)",
"properties": {
"cnr_id": "comfy-core",
"ver": "0.3.73",
"Node name for S&R": "CLIPTextEncode",
"enableTabs": false,
"tabWidth": 65,
"tabXOffset": 10,
"hasSecondTab": false,
"secondTabText": "Send Back",
"secondTabOffset": 80,
"secondTabWidth": 65
},
"widgets_values": [
""
],
"color": "#232",
"bgcolor": "#353"
},
{
"id": 99,
"type": "EmptySD3LatentImage",
"pos": [
-854.4052795360801,
-79.8964947543219
],
"size": [
260,
168
],
"flags": {},
"order": 0,
"mode": 0,
"inputs": [],
"outputs": [
{
"name": "LATENT",
"type": "LATENT",
"slot_index": 0,
"links": [
130
]
}
],
"properties": {
"cnr_id": "comfy-core",
"ver": "0.3.64",
"Node name for S&R": "EmptySD3LatentImage",
"enableTabs": false,
"tabWidth": 65,
"tabXOffset": 10,
"hasSecondTab": false,
"secondTabText": "Send Back",
"secondTabOffset": 80,
"secondTabWidth": 65
},
"widgets_values": [
1024,
1024,
1
]
},
{
"id": 100,
"type": "VAELoader",
"pos": [
-860.3984289317032,
-250.04560959620238
],
"size": [
270,
106.65625
],
"flags": {},
"order": 1,
"mode": 0,
"inputs": [],
"outputs": [
{
"name": "VAE",
"type": "VAE",
"links": [
186
]
}
],
"properties": {
"cnr_id": "comfy-core",
"ver": "0.3.73",
"Node name for S&R": "VAELoader",
"models": [
{
"name": "ae.safetensors",
"url": "https://huggingface.co/Comfy-Org/z_image_turbo/resolve/main/split_files/vae/ae.safetensors",
"directory": "vae"
}
],
"enableTabs": false,
"tabWidth": 65,
"tabXOffset": 10,
"hasSecondTab": false,
"secondTabText": "Send Back",
"secondTabOffset": 80,
"secondTabWidth": 65
},
"widgets_values": [
"ae.safetensors"
]
},
{
"id": 101,
"type": "CLIPLoader",
"pos": [
-852.3844599188506,
-448.2477751296032
],
"size": [
270,
144.65625
],
"flags": {},
"order": 2,
"mode": 0,
"inputs": [],
"outputs": [
{
"name": "CLIP",
"type": "CLIP",
"links": [
126,
131
]
}
],
"properties": {
"cnr_id": "comfy-core",
"ver": "0.3.73",
"Node name for S&R": "CLIPLoader",
"models": [
{
"name": "qwen_3_4b.safetensors",
"url": "https://huggingface.co/Comfy-Org/z_image_turbo/resolve/main/split_files/text_encoders/qwen_3_4b.safetensors",
"directory": "text_encoders"
}
],
"enableTabs": false,
"tabWidth": 65,
"tabXOffset": 10,
"hasSecondTab": false,
"secondTabText": "Send Back",
"secondTabOffset": 80,
"secondTabWidth": 65
},
"widgets_values": [
"qwen_3_4b.safetensors",
"lumina2",
"default"
]
},
{
"id": 103,
"type": "ModelSamplingAuraFlow",
"pos": [
666.3468492771774,
-627.2266415976245
],
"size": [
310,
104
],
"flags": {},
"order": 6,
"mode": 0,
"inputs": [
{
"name": "model",
"type": "MODEL",
"link": 135
}
],
"outputs": [
{
"name": "MODEL",
"type": "MODEL",
"slot_index": 0,
"links": [
127
]
}
],
"properties": {
"cnr_id": "comfy-core",
"ver": "0.3.64",
"Node name for S&R": "ModelSamplingAuraFlow",
"enableTabs": false,
"tabWidth": 65,
"tabXOffset": 10,
"hasSecondTab": false,
"secondTabText": "Send Back",
"secondTabOffset": 80,
"secondTabWidth": 65
},
"widgets_values": [
3
]
},
{
"id": 104,
"type": "UNETLoader",
"pos": [
-858.2503268450102,
-644.3815334763085
],
"size": [
341.84375,
147.46875
],
"flags": {},
"order": 3,
"mode": 0,
"showAdvanced": true,
"inputs": [],
"outputs": [
{
"name": "MODEL",
"type": "MODEL",
"links": [
135
]
}
],
"properties": {
"cnr_id": "comfy-core",
"ver": "0.3.73",
"Node name for S&R": "UNETLoader",
"models": [
{
"name": "z_image_bf16.safetensors",
"url": "https://huggingface.co/Comfy-Org/z_image/resolve/main/split_files/diffusion_models/z_image_bf16.safetensors",
"directory": "diffusion_models"
}
],
"enableTabs": false,
"tabWidth": 65,
"tabXOffset": 10,
"hasSecondTab": false,
"secondTabText": "Send Back",
"secondTabOffset": 80,
"secondTabWidth": 65
},
"widgets_values": [
"z_image_bf16.safetensors",
"default"
]
},
{
"id": 105,
"type": "CLIPTextEncode",
"pos": [
88.47352910322218,
-403.69008919228577
],
"size": [
554.78125,
144.65625
],
"flags": {},
"order": 5,
"mode": 0,
"inputs": [
{
"name": "clip",
"type": "CLIP",
"link": 131
},
{
"label": "prompt",
"name": "text",
"type": "STRING",
"widget": {
"name": "text"
},
"link": null
}
],
"outputs": [
{
"name": "CONDITIONING",
"type": "CONDITIONING",
"links": [
129
]
}
],
"title": "CLIP Text Encode (Negative Prompt)",
"properties": {
"cnr_id": "comfy-core",
"ver": "0.3.73",
"Node name for S&R": "CLIPTextEncode",
"enableTabs": false,
"tabWidth": 65,
"tabXOffset": 10,
"hasSecondTab": false,
"secondTabText": "Send Back",
"secondTabOffset": 80,
"secondTabWidth": 65
},
"widgets_values": [
""
],
"color": "#323",
"bgcolor": "#535"
},
{
"id": 106,
"type": "KSampler",
"pos": [
666.6701937790448,
-479.65517806745004
],
"size": [
310,
440
],
"flags": {},
"order": 8,
"mode": 0,
"inputs": [
{
"name": "model",
"type": "MODEL",
"link": 127
},
{
"name": "positive",
"type": "CONDITIONING",
"link": 128
},
{
"name": "negative",
"type": "CONDITIONING",
"link": 129
},
{
"name": "latent_image",
"type": "LATENT",
"link": 130
}
],
"outputs": [
{
"name": "LATENT",
"type": "LATENT",
"slot_index": 0,
"links": [
185,
195
]
}
],
"properties": {
"cnr_id": "comfy-core",
"ver": "0.3.64",
"Node name for S&R": "KSampler",
"enableTabs": false,
"tabWidth": 65,
"tabXOffset": 10,
"hasSecondTab": false,
"secondTabText": "Send Back",
"secondTabOffset": 80,
"secondTabWidth": 65
},
"widgets_values": [
681810010921358,
"randomize",
40,
4,
"euler",
"beta",
1
]
},
{
"id": 115,
"type": "PiDTextPrompt",
"pos": [
-494.421938151345,
-639.7729653465548
],
"size": [
565.625,
379.78125
],
"flags": {},
"order": 4,
"mode": 0,
"inputs": [],
"outputs": [
{
"name": "text",
"type": "STRING",
"links": [
144
]
},
{
"name": "caption",
"type": "STRING",
"links": [
196
]
}
],
"properties": {
"aux_id": "Merserk/ComfyUI-PiD",
"Node name for S&R": "PiDTextPrompt",
"ver": "093723cc9702c2ec7f09d900072c1636a082714c"
},
"widgets_values": [
"A vibrant graphic novel-style cinematic illustration of a confident woman standing in the foreground, framed as the commanding heroine of a sci-fi adventure. She has a sleek jet-black bob haircut, bold red lipstick, and an intense, self-assured expression. She wears a striking crimson bomber jacket over a fitted black tactical bodysuit, with black gloves, a utility belt, and subtle futuristic detailing.\n\nBehind her stretches a vast, idealized ancient Egyptian city blended with sci-fi elements. A winding blue river cuts through the golden desert landscape, reflecting the bright sunlight beneath a clear, vivid blue sky. The environment is filled with monumental pharaonic statues, weathered stone columns, ruined temples, ancient obelisks, and partially collapsed structures. In the far distance, majestic pyramids rise against the horizon, surrounded by towering columns and elegant futuristic spires that hint at an advanced lost civilization.\n\nUse clean, bold linework, dynamic composition, rich warm desert tones, deep shadows, and high-contrast colors. The mood should feel epic, adventurous, mysterious, and cinematic, like a premium comic book cover or concept art for a sci-fi Egyptian fantasy film. Emphasize sharp detail, dramatic scale, atmospheric depth, and a polished graphic novel aesthetic."
]
},
{
"id": 123,
"type": "VAEDecode",
"pos": [
1059.9908164836704,
-132.5232981098975
],
"size": [
225,
71.328125
],
"flags": {},
"order": 9,
"mode": 0,
"inputs": [
{
"name": "samples",
"type": "LATENT",
"link": 185
},
{
"name": "vae",
"type": "VAE",
"link": 186
}
],
"outputs": [
{
"name": "IMAGE",
"type": "IMAGE",
"links": [
184,
199
]
}
],
"properties": {
"cnr_id": "comfy-core",
"ver": "0.22.0",
"Node name for S&R": "VAEDecode"
},
"widgets_values": []
},
{
"id": 124,
"type": "SaveImage",
"pos": [
1523.9206655601654,
-201.9060776691091
],
"size": [
730,
570
],
"flags": {},
"order": 11,
"mode": 0,
"inputs": [
{
"name": "images",
"type": "IMAGE",
"link": 184
}
],
"outputs": [],
"properties": {
"cnr_id": "comfy-core",
"ver": "0.3.64",
"Node name for S&R": "SaveImage",
"enableTabs": false,
"tabWidth": 65,
"tabXOffset": 10,
"hasSecondTab": false,
"secondTabText": "Send Back",
"secondTabOffset": 80,
"secondTabWidth": 65
},
"widgets_values": [
"z-image-base"
]
},
{
"id": 127,
"type": "PiDDecode",
"pos": [
1074.9083062821073,
-830.8458209260093
],
"size": [
309.03125,
492.65625
],
"flags": {},
"order": 10,
"mode": 0,
"inputs": [
{
"name": "latent",
"type": "LATENT",
"link": 195
},
{
"name": "caption",
"type": "STRING",
"link": 196
},
{
"name": "vae",
"shape": 7,
"type": "VAE",
"link": null
},
{
"name": "baseline_image",
"shape": 7,
"type": "IMAGE",
"link": 199
}
],
"outputs": [
{
"name": "image",
"type": "IMAGE",
"links": [
198
]
}
],
"properties": {
"Node name for S&R": "PiDDecode"
},
"widgets_values": [
"zimage",
"2kto4k",
4,
2,
1,
0,
0,
"randomize",
true,
true,
true,
""
]
}
],
"links": [
[
126,
101,
0,
98,
0,
"CLIP"
],
[
127,
103,
0,
106,
0,
"MODEL"
],
[
128,
98,
0,
106,
1,
"CONDITIONING"
],
[
129,
105,
0,
106,
2,
"CONDITIONING"
],
[
130,
99,
0,
106,
3,
"LATENT"
],
[
131,
101,
0,
105,
0,
"CLIP"
],
[
135,
104,
0,
103,
0,
"MODEL"
],
[
144,
115,
0,
98,
1,
"STRING"
],
[
184,
123,
0,
124,
0,
"IMAGE"
],
[
185,
106,
0,
123,
0,
"LATENT"
],
[
186,
100,
0,
123,
1,
"VAE"
],
[
195,
106,
0,
127,
0,
"LATENT"
],
[
196,
115,
1,
127,
1,
"STRING"
],
[
198,
127,
0,
9,
0,
"IMAGE"
],
[
199,
123,
0,
127,
3,
"IMAGE"
]
],
"groups": [],
"config": {},
"extra": {
"ds": {
"scale": 0.7835797005297873,
"offset": [
1310.8456872126858,
1289.9466940639688
]
},
"frontendVersion": "1.44.19",
"workflowRendererVersion": "LG",
"VHS_latentpreview": false,
"VHS_latentpreviewrate": 0,
"VHS_MetadataImage": true,
"VHS_KeepIntermediate": true
},
"version": 0.4
}
nodes.py
+28
View File
@@ -0,0 +1,28 @@
"""ComfyUI registration shim.
Each node implementation lives in its own module. ComfyUI imports this file as
the package entry point, so keep it small and only expose mappings here.
"""
try:
from .pid_decode import NODE_CLASS_MAPPINGS as DECODE_CLASS_MAPPINGS
from .pid_decode import NODE_DISPLAY_NAME_MAPPINGS as DECODE_DISPLAY_NAME_MAPPINGS
from .pid_text_prompt import NODE_CLASS_MAPPINGS as PROMPT_CLASS_MAPPINGS
from .pid_text_prompt import NODE_DISPLAY_NAME_MAPPINGS as PROMPT_DISPLAY_NAME_MAPPINGS
except ImportError: # Allows `python -c "import nodes"` from this folder.
from pid_decode import NODE_CLASS_MAPPINGS as DECODE_CLASS_MAPPINGS
from pid_decode import NODE_DISPLAY_NAME_MAPPINGS as DECODE_DISPLAY_NAME_MAPPINGS
from pid_text_prompt import NODE_CLASS_MAPPINGS as PROMPT_CLASS_MAPPINGS
from pid_text_prompt import NODE_DISPLAY_NAME_MAPPINGS as PROMPT_DISPLAY_NAME_MAPPINGS
NODE_CLASS_MAPPINGS = {
**DECODE_CLASS_MAPPINGS,
**PROMPT_CLASS_MAPPINGS,
}
NODE_DISPLAY_NAME_MAPPINGS = {
**DECODE_DISPLAY_NAME_MAPPINGS,
**PROMPT_DISPLAY_NAME_MAPPINGS,
}
__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]
pid_decode.py
+891
View File
@@ -0,0 +1,891 @@
"""
ComfyUI-PiD
Experimental ComfyUI nodes for NVIDIA PiD (Pixel Diffusion Decoder).
PiD is not a native ComfyUI VAE object. NVIDIA's inference path needs:
latent + native decoder/baseline image + sigma + caption -> PiD -> final image
So the decode node accepts a LATENT plus either a matching ComfyUI VAE or a
pre-decoded baseline IMAGE, then returns IMAGE.
"""
from __future__ import annotations
import contextlib
import gc
import os
import shutil
import subprocess
import sys
import urllib.request
import zipfile
from pathlib import Path
from dataclasses import dataclass
from typing import Dict, List, Optional, Tuple
import torch
try:
import comfy.model_management as model_management
except Exception: # pragma: no cover - only available inside ComfyUI
model_management = None
try:
import folder_paths # ComfyUI model path helper
except Exception: # pragma: no cover - only available inside ComfyUI
folder_paths = None
NODE_DIR = Path(__file__).resolve().parent
# Expected layout: ComfyUI/custom_nodes/ComfyUI-PiD-ZImage/pid_decode.py
COMFYUI_DIR = NODE_DIR.parent.parent
VENDOR_DIR = NODE_DIR / "vendor"
DEFAULT_PID_DIR = VENDOR_DIR / "PiD"
PID_REPO_URL = "https://github.com/nv-tlabs/PiD.git"
PID_ZIP_URL = "https://github.com/nv-tlabs/PiD/archive/refs/heads/main.zip"
HF_REPO_ID = "nvidia/PiD"
AE_REL_PATH = "checkpoints/ae.safetensors"
AE_FILENAME = "ae.safetensors"
@dataclass(frozen=True)
class PiDCheckpoint:
experiment: str
relpath: str
scale: int
@dataclass(frozen=True)
class PiDBackbone:
label: str
registry_key: str
latent_channels: int
default_scale: int
ckpt_types: Tuple[str, ...]
aux_files: Tuple[str, ...] = ()
needs_flux_ae: bool = False
# Mirrors NVIDIA's official pid/_src/inference/checkpoint_registry.py.
PID_CKPTS: Dict[Tuple[str, str], PiDCheckpoint] = {
("flux", "2k"): PiDCheckpoint(
experiment="PiD_res2k_sr4x_official_flux_distill_4step",
relpath="checkpoints/PiD_res2k_sr4x_official_flux_distill_4step/model_ema_bf16.pth",
scale=4,
),
("flux2", "2k"): PiDCheckpoint(
experiment="PiD_res2k_sr4x_official_flux2_distill_4step",
relpath="checkpoints/PiD_res2k_sr4x_official_flux2_distill_4step/model_ema_bf16.pth",
scale=4,
),
("sd3", "2k"): PiDCheckpoint(
experiment="PiD_res2k_sr4x_official_sd3_distill_4step",
relpath="checkpoints/PiD_res2k_sr4x_official_sd3_distill_4step/model_ema_bf16.pth",
scale=4,
),
("zimage", "2k"): PiDCheckpoint(
experiment="PiD_res2k_sr4x_official_flux_distill_4step",
relpath="checkpoints/PiD_res2k_sr4x_official_flux_distill_4step/model_ema_bf16.pth",
scale=4,
),
("rae", "2k"): PiDCheckpoint(
experiment="PiD_res2k_sr4x_official_dinov2_distill_4step",
relpath="checkpoints/PiD_res2k_sr4x_official_dinov2_distill_4step/model_ema_bf16.pth",
scale=4,
),
("scale_rae", "2k"): PiDCheckpoint(
experiment="PiD_res2k_sr8x_official_siglip_distill_4step",
relpath="checkpoints/PiD_res2k_sr8x_official_siglip_distill_4step/model_ema_bf16.pth",
scale=8,
),
("flux", "2kto4k"): PiDCheckpoint(
experiment="PiD_res2kto4k_sr4x_official_flux_distill_4step",
relpath="checkpoints/PiD_res2kto4k_sr4x_official_flux_distill_4step/model_ema_bf16.pth",
scale=4,
),
("flux2", "2kto4k"): PiDCheckpoint(
experiment="PiD_res2kto4k_sr4x_official_flux2_distill_4step",
relpath="checkpoints/PiD_res2kto4k_sr4x_official_flux2_distill_4step/model_ema_bf16.pth",
scale=4,
),
("sd3", "2kto4k"): PiDCheckpoint(
experiment="PiD_res2kto4k_sr4x_official_sd3_distill_4step",
relpath="checkpoints/PiD_res2kto4k_sr4x_official_sd3_distill_4step/model_ema_bf16.pth",
scale=4,
),
("zimage", "2kto4k"): PiDCheckpoint(
experiment="PiD_res2kto4k_sr4x_official_flux_distill_4step",
relpath="checkpoints/PiD_res2kto4k_sr4x_official_flux_distill_4step/model_ema_bf16.pth",
scale=4,
),
}
PID_BACKBONES: Dict[str, PiDBackbone] = {
"zimage": PiDBackbone("Z-Image", "zimage", 16, 4, ("2k", "2kto4k"), needs_flux_ae=True),
"flux": PiDBackbone("Flux", "flux", 16, 4, ("2k", "2kto4k"), needs_flux_ae=True),
"flux2": PiDBackbone(
"Flux2",
"flux2",
128,
4,
("2k", "2kto4k"),
aux_files=("checkpoints/flux2_ae.safetensors",),
),
"sd3": PiDBackbone(
"SD3",
"sd3",
16,
4,
("2k", "2kto4k"),
aux_files=("checkpoints/sd3_vae/vae/diffusion_pytorch_model.safetensors",),
),
"dinov2": PiDBackbone(
"DINOv2 RAE",
"rae",
768,
4,
("2k",),
aux_files=(
"checkpoints/rae/decoders/dinov2/wReg_base/ViTXL_n08_i512/model.pt",
"checkpoints/rae/stats/dinov2/wReg_base/imagenet1k_512/stat.pt",
),
),
"siglip": PiDBackbone(
"SigLIP Scale-RAE",
"scale_rae",
1152,
8,
("2k",),
aux_files=(
"checkpoints/scale_rae/decoder/XL_decoder_config.json",
"checkpoints/scale_rae/decoder/siglip2_sop14_i224_web73M_ganw3_decXL.pt",
),
),
}
BACKBONE_CHOICES = list(PID_BACKBONES.keys())
_MODEL_CACHE: Dict[Tuple[str, str, str, str, str, bool, str], object] = {}
class PiDNodeError(RuntimeError):
pass
@contextlib.contextmanager
def _pushd(path: Path):
old = Path.cwd()
os.chdir(str(path))
try:
yield
finally:
os.chdir(str(old))
def _python_executable() -> str:
return sys.executable or "python"
def _run(cmd: List[str], cwd: Optional[Path] = None) -> str:
proc = subprocess.run(
cmd,
cwd=str(cwd) if cwd else None,
stdout=subprocess.PIPE,
stderr=subprocess.STDOUT,
text=True,
encoding="utf-8",
errors="replace",
)
if proc.returncode != 0:
raise PiDNodeError("Command failed:\n" + " ".join(cmd) + "\n\n" + proc.stdout)
return proc.stdout.strip()
def _free_cuda_memory(aggressive: bool = False) -> None:
"""Best-effort cleanup for running PiD after a large ComfyUI model.
ComfyUI often keeps the Z-Image UNet/TE/VAE staged for dynamic VRAM. PiD is
another 1.36B parameter model, so on 16GB cards it helps to decode the native
baseline image first, then release Comfy's loaded models before PiD inference.
"""
try:
gc.collect()
except Exception:
pass
if model_management is not None:
if aggressive:
for name in ("unload_all_models", "unload_model_clones"):
fn = getattr(model_management, name, None)
if callable(fn):
try:
fn()
except TypeError:
try:
fn(None)
except Exception:
pass
except Exception:
pass
try:
model_management.soft_empty_cache()
except Exception:
pass
if torch.cuda.is_available():
try:
torch.cuda.synchronize()
except Exception:
pass
try:
torch.cuda.empty_cache()
except Exception:
pass
if aggressive:
try:
torch.cuda.ipc_collect()
except Exception:
pass
def _vram_total_gb() -> float:
if not torch.cuda.is_available():
return 0.0
try:
props = torch.cuda.get_device_properties(torch.cuda.current_device())
return float(props.total_memory) / (1024 ** 3)
except Exception:
return 0.0
def _format_pid_runtime_error(exc: BaseException, infer_image_size: Tuple[int, int], ckpt_type: str, scale: int) -> PiDNodeError:
msg = str(exc)
lower = msg.lower()
h, w = int(infer_image_size[0]), int(infer_image_size[1])
advice = (
f"PiD inference failed while generating {w}x{h} with pid_ckpt_type={ckpt_type!r}, scale={scale}.\n\n"
"Most common cause: VRAM/allocator pressure. For a 16GB GPU, 1024 base + scale=4 "
"means roughly 4096x4096 PiD output, which is usually too heavy inside ComfyUI while "
"large generation models were just used.\n\n"
"Try these first:\n"
" 1. Set PiD scale=1 or scale=2 for 1024 base images.\n"
" 2. For the 2k PiD checkpoint, use a 512x512 base latent with scale=4.\n"
" 3. Keep unload_comfy_before_pid=true and aggressive_cleanup=true.\n"
" 4. Restart ComfyUI after a CUDA allocator crash.\n"
" 5. If your launcher supports it, try disabling cudaMallocAsync / CUDA malloc async.\n\n"
f"Original error: {msg}"
)
if "out of memory" in lower or "cudamallocasync" in lower or "internal assert" in lower or "cuda" in lower:
return PiDNodeError(advice)
return PiDNodeError(f"PiD inference failed. Original error: {msg}")
def _resolve_pid_dir(pid_source_dir: str = "") -> Path:
# Priority: explicit UI path -> environment variable -> bundled vendor path.
if pid_source_dir and pid_source_dir.strip():
return Path(pid_source_dir.strip()).expanduser().resolve()
env_path = os.environ.get("PID_REPO_DIR") or os.environ.get("COMFYUI_PID_REPO_DIR")
if env_path:
return Path(env_path).expanduser().resolve()
return DEFAULT_PID_DIR.resolve()
def _pid_is_present(pid_dir: Path) -> bool:
return (pid_dir / "pid" / "_src" / "utils" / "model_loader.py").is_file()
def _ensure_pid_source(pid_dir: Path, allow_download: bool = True) -> List[str]:
messages: List[str] = []
if _pid_is_present(pid_dir):
if str(pid_dir) not in sys.path:
sys.path.insert(0, str(pid_dir))
messages.append(f"PiD source found: {pid_dir}")
return messages
if not allow_download:
raise PiDNodeError(
f"PiD source was not found at {pid_dir}.\n"
"Set auto_download=true, or set PID_REPO_DIR / COMFYUI_PID_REPO_DIR "
"to your local PiD checkout."
)
pid_dir.parent.mkdir(parents=True, exist_ok=True)
git = shutil.which("git")
if git:
messages.append(f"Cloning PiD source into {pid_dir} ...")
_run([git, "clone", "--depth", "1", PID_REPO_URL, str(pid_dir)])
else:
messages.append("git not found; downloading PiD main.zip instead ...")
zip_path = pid_dir.parent / "PiD-main.zip"
urllib.request.urlretrieve(PID_ZIP_URL, zip_path)
with zipfile.ZipFile(zip_path, "r") as zf:
zf.extractall(pid_dir.parent)
extracted = pid_dir.parent / "PiD-main"
if pid_dir.exists():
shutil.rmtree(pid_dir)
extracted.rename(pid_dir)
zip_path.unlink(missing_ok=True)
if not _pid_is_present(pid_dir):
raise PiDNodeError(f"PiD source download finished but expected files were not found under {pid_dir}")
if str(pid_dir) not in sys.path:
sys.path.insert(0, str(pid_dir))
messages.append("PiD source is ready.")
return messages
def _ensure_hf_download_available() -> None:
try:
import huggingface_hub # noqa: F401
except Exception as exc:
raise PiDNodeError(
"Missing dependency: huggingface_hub.\n"
"Install this custom node's requirements, or manually run:\n"
f"{_python_executable()} -m pip install \"huggingface-hub>=1.0\""
) from exc
def _candidate_comfy_ae_paths() -> List[Path]:
"""Return likely local Flux/Z-Image AE VAE paths from ComfyUI.
NVIDIA PiD's own config instantiates a Flux VAE encoder/tokenizer and expects
./checkpoints/ae.safetensors relative to the PiD repository root. Most ComfyUI
Z-Image workflows already have the exact same file under ComfyUI/models/vae/.
"""
candidates: List[Path] = []
# Ask ComfyUI's model registry first. This respects extra_model_paths.yaml.
if folder_paths is not None:
try:
get_full_path = getattr(folder_paths, "get_full_path", None)
if get_full_path is not None:
p = get_full_path("vae", AE_FILENAME)
if p:
candidates.append(Path(p))
except Exception:
pass
try:
get_full_path_or_raise = getattr(folder_paths, "get_full_path_or_raise", None)
if get_full_path_or_raise is not None:
p = get_full_path_or_raise("vae", AE_FILENAME)
if p:
candidates.append(Path(p))
except Exception:
pass
try:
for name in folder_paths.get_filename_list("vae"):
if Path(name).name.lower() == AE_FILENAME:
try:
p = folder_paths.get_full_path("vae", name)
if p:
candidates.append(Path(p))
except Exception:
pass
except Exception:
pass
# Fallback to standard ComfyUI folder locations.
vae_dir = COMFYUI_DIR / "models" / "vae"
candidates.append(vae_dir / AE_FILENAME)
try:
candidates.extend(vae_dir.glob(f"**/{AE_FILENAME}"))
except Exception:
pass
# Deduplicate while preserving order.
seen = set()
out: List[Path] = []
for c in candidates:
try:
cr = c.expanduser().resolve()
except Exception:
continue
if cr in seen:
continue
seen.add(cr)
out.append(cr)
return out
def _copy_local_ae_to_pid(pid_dir: Path) -> Optional[str]:
target = pid_dir / AE_REL_PATH
if target.is_file():
return f"Flux AE VAE already exists: {target}"
for src in _candidate_comfy_ae_paths():
if src.is_file():
target.parent.mkdir(parents=True, exist_ok=True)
shutil.copy2(src, target)
return f"Copied Flux AE VAE from {src} to {target}"
return None
def _download_ae_safetensors(pid_dir: Path) -> str:
_ensure_hf_download_available()
from huggingface_hub import hf_hub_download
target = pid_dir / AE_REL_PATH
if target.is_file():
return f"Flux AE VAE already exists: {target}"
token = os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN")
pid_dir.mkdir(parents=True, exist_ok=True)
hf_hub_download(
repo_id=HF_REPO_ID,
filename=AE_REL_PATH,
local_dir=str(pid_dir),
token=token,
)
if not target.is_file():
raise PiDNodeError(f"Hugging Face download finished but Flux AE VAE is missing: {target}")
return f"Downloaded Flux AE VAE: {target}"
def _ensure_ae_safetensors(pid_dir: Path, allow_download: bool = True) -> Path:
"""Ensure PiD's required ./checkpoints/ae.safetensors exists."""
target = pid_dir / AE_REL_PATH
if target.is_file():
return target
copied = _copy_local_ae_to_pid(pid_dir)
if copied and target.is_file():
return target
if allow_download:
try:
_download_ae_safetensors(pid_dir)
return target
except Exception as exc:
local_candidates = "\n".join(f" - {p}" for p in _candidate_comfy_ae_paths()) or " - none found"
raise PiDNodeError(
"PiD requires Flux's AE VAE at:\n"
f" {target}\n\n"
"I could not find/copy it from ComfyUI and could not download it automatically.\n"
"Manual fix:\n"
f" 1. Make sure {AE_FILENAME} exists in ComfyUI/models/vae/\n"
f" 2. Copy it to {target}\n\n"
"Local paths checked:\n"
f"{local_candidates}\n\n"
f"Original download/copy error: {exc}"
) from exc
local_candidates = "\n".join(f" - {p}" for p in _candidate_comfy_ae_paths()) or " - none found"
raise PiDNodeError(
"PiD requires Flux's AE VAE at:\n"
f" {target}\n\n"
"Set auto_download=true, or manually copy:\n"
f" ComfyUI/models/vae/{AE_FILENAME}\n"
"to:\n"
f" {target}\n\n"
"Local paths checked:\n"
f"{local_candidates}"
)
def _download_hf_file(pid_dir: Path, relpath: str) -> str:
_ensure_hf_download_available()
from huggingface_hub import hf_hub_download
target = pid_dir / relpath
if target.is_file():
return f"Asset already exists: {target}"
token = os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN")
pid_dir.mkdir(parents=True, exist_ok=True)
hf_hub_download(
repo_id=HF_REPO_ID,
filename=relpath,
local_dir=str(pid_dir),
token=token,
)
if not target.is_file():
raise PiDNodeError(f"Hugging Face download finished but asset is missing: {target}")
return f"Downloaded: {target}"
def _checkpoint_for(backbone: str, ckpt_type: str) -> PiDCheckpoint:
if backbone not in PID_BACKBONES:
raise PiDNodeError(f"Unknown backbone={backbone!r}; expected one of {BACKBONE_CHOICES}")
backbone_info = PID_BACKBONES[backbone]
if ckpt_type not in backbone_info.ckpt_types:
supported = ", ".join(backbone_info.ckpt_types)
raise PiDNodeError(
f"{backbone_info.label} does not ship a {ckpt_type!r} checkpoint in NVIDIA's registry. "
f"Supported pid_ckpt_type values: {supported}."
)
try:
return PID_CKPTS[(backbone_info.registry_key, ckpt_type)]
except KeyError as exc:
raise PiDNodeError(f"No PiD checkpoint registered for backbone={backbone!r}, pid_ckpt_type={ckpt_type!r}") from exc
def _ensure_checkpoint(pid_dir: Path, backbone: str, ckpt_type: str, allow_download: bool = True) -> Path:
ckpt = _checkpoint_for(backbone, ckpt_type)
ckpt_path = pid_dir / ckpt.relpath
if ckpt_path.is_file():
return ckpt_path
if not allow_download:
raise PiDNodeError(
f"PiD checkpoint is missing: {ckpt_path}\n"
"Set auto_download=true or download the matching file from nvidia/PiD."
)
_download_hf_file(pid_dir, ckpt.relpath)
return ckpt_path
def _ensure_backbone_assets(pid_dir: Path, backbone: str, allow_download: bool = True) -> None:
info = PID_BACKBONES[backbone]
if info.needs_flux_ae:
_ensure_ae_safetensors(pid_dir, allow_download=allow_download)
missing = [relpath for relpath in info.aux_files if not (pid_dir / relpath).is_file()]
if not missing:
return
if not allow_download:
missing_text = "\n".join(f" - {pid_dir / relpath}" for relpath in missing)
raise PiDNodeError(
f"{info.label} PiD support needs these extra asset files:\n"
f"{missing_text}\n\n"
"Set auto_download=true or download them from nvidia/PiD."
)
for relpath in missing:
_download_hf_file(pid_dir, relpath)
def _import_pid_loader(pid_dir: Path):
if str(pid_dir) not in sys.path:
sys.path.insert(0, str(pid_dir))
try:
from pid._src.utils.model_loader import load_model_from_checkpoint
except Exception as exc:
raise PiDNodeError(
"Could not import NVIDIA PiD. Usually this means PiD dependencies are not installed.\n\n"
"Install the packages in requirements.txt, then restart ComfyUI.\n\n"
f"Original import error: {exc}"
) from exc
return load_model_from_checkpoint
def _load_pid_model(
pid_dir: Path,
backbone: str,
ckpt_type: str,
checkpoint_path: Path,
dtype_choice: str,
load_ema_to_reg: bool = False,
):
experiment = _checkpoint_for(backbone, ckpt_type).experiment
config_file = "pid/_src/configs/pid/config.py"
cache_key = (
str(pid_dir),
backbone,
ckpt_type,
str(checkpoint_path),
experiment,
bool(load_ema_to_reg),
dtype_choice,
)
if cache_key in _MODEL_CACHE:
return _MODEL_CACHE[cache_key]
if not torch.cuda.is_available():
raise PiDNodeError("PiD's official model loader instantiates the decoder on CUDA. CUDA GPU is required.")
load_model_from_checkpoint = _import_pid_loader(pid_dir)
# PiD's config helper expects config_file to be relative to the PiD repo root.
with _pushd(pid_dir):
model, _config = load_model_from_checkpoint(
experiment_name=experiment,
checkpoint_path=str(checkpoint_path),
config_file=config_file,
enable_fsdp=False,
experiment_opts=[],
strict=False,
load_ema_to_reg=load_ema_to_reg,
)
model.eval()
_MODEL_CACHE[cache_key] = model
return model
def _remove_model_from_cache(model: object) -> None:
"""Drop a specific PiD model from the private cache by object identity."""
for key, cached_model in list(_MODEL_CACHE.items()):
if cached_model is model:
_MODEL_CACHE.pop(key, None)
def _unload_pid_model(model: object, aggressive: bool = True) -> None:
"""Move PiD off CUDA and clear the private cache after a decode run."""
_remove_model_from_cache(model)
try:
to = getattr(model, "to", None)
if callable(to):
to("cpu")
except Exception:
pass
_free_cuda_memory(aggressive=aggressive)
def _latent_samples(latent: dict) -> torch.Tensor:
if not isinstance(latent, dict) or "samples" not in latent:
raise PiDNodeError("Expected a ComfyUI LATENT dict containing key 'samples'.")
samples = latent["samples"]
if getattr(samples, "is_nested", False):
samples = samples.unbind()[0]
if samples.ndim != 4:
raise PiDNodeError(f"Expected latent samples as [B,C,H,W], got shape {list(samples.shape)}")
return samples
def _comfy_image_to_bchw_01(image: torch.Tensor) -> torch.Tensor:
# Comfy IMAGE is [B,H,W,C] float 0..1.
if image.ndim != 4 or image.shape[-1] not in (1, 3, 4):
raise PiDNodeError(f"Expected ComfyUI IMAGE [B,H,W,C], got shape {list(image.shape)}")
if image.shape[-1] == 4:
image = image[..., :3]
if image.shape[-1] == 1:
image = image.repeat(1, 1, 1, 3)
return image.permute(0, 3, 1, 2).contiguous().clamp(0, 1)
def _bchw_neg1_to_comfy_image(image: torch.Tensor) -> torch.Tensor:
# PiD output is normally [B,C,H,W] in [-1,1]. Some PiD builds return a
# single-frame video tensor [B,C,1,H,W] or [B,1,C,H,W]. Convert all of
# those to Comfy IMAGE [B,H,W,C] in [0,1].
if image.ndim == 5:
# Observed from PiD/Z-Image: [B, C, T, H, W] with T=1.
if image.shape[2] == 1 and image.shape[1] in (1, 3, 4):
image = image[:, :, 0, :, :]
# Defensive support for [B, T, C, H, W] with T=1.
elif image.shape[1] == 1 and image.shape[2] in (1, 3, 4):
image = image[:, 0, :, :, :]
else:
raise PiDNodeError(
"Expected PiD output [B,C,H,W], [B,C,1,H,W], or [B,1,C,H,W], "
f"got shape {list(image.shape)}"
)
if image.ndim == 3:
image = image.unsqueeze(0)
if image.ndim != 4:
raise PiDNodeError(f"Expected PiD output [B,C,H,W], got shape {list(image.shape)}")
if image.shape[1] == 4:
image = image[:, :3, :, :]
if image.shape[1] == 1:
image = image.repeat(1, 3, 1, 1)
if image.shape[1] != 3:
raise PiDNodeError(f"Expected 1/3/4 output channels, got shape {list(image.shape)}")
image = image.float().clamp(-1, 1)
image = (image + 1.0) * 0.5
return image.permute(0, 2, 3, 1).contiguous().cpu().clamp(0, 1)
def _normalize_pid_samples(samples):
# Official code uses samples[0]; support tensor, list[tensor], or tuple returns.
if isinstance(samples, torch.Tensor):
return samples
if isinstance(samples, (list, tuple)):
if len(samples) == 1 and isinstance(samples[0], torch.Tensor) and samples[0].ndim == 4:
return samples[0]
tensors = [s for s in samples if isinstance(s, torch.Tensor)]
if len(tensors) == 1:
return tensors[0]
if tensors and all(t.ndim == 3 for t in tensors):
return torch.stack(tensors, dim=0)
raise PiDNodeError(f"Unsupported PiD output type: {type(samples)}")
def _decode_baseline_with_comfy_vae(vae, samples: torch.Tensor, backbone: str) -> torch.Tensor:
try:
image = vae.decode(samples)
except Exception as exc:
label = PID_BACKBONES.get(backbone, PID_BACKBONES["zimage"]).label
raise PiDNodeError(
f"ComfyUI VAE decode failed. Make sure the connected VAE can decode {label} "
"latents, or connect a pre-decoded baseline_image instead."
) from exc
return _comfy_image_to_bchw_01(image)
class PiDDecode:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"latent": ("LATENT",),
"caption": ("STRING", {"forceInput": True}),
# Use ComfyUI's canonical combo syntax. The older
# ("COMBO", {"options": ...}) form can cause the frontend to
# rebuild widget_values with the first combo dropped after tab
# switches / node refreshes.
"backbone": (BACKBONE_CHOICES, {"default": "zimage"}),
"pid_ckpt_type": (["2k", "2kto4k"], {"default": "2k"}),
"pid_steps": ("INT", {"default": 4, "min": 1, "max": 64, "step": 1}),
"scale": ("INT", {"default": 0, "min": 0, "max": 8, "step": 1}),
"cfg_scale": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 20.0, "step": 0.1}),
"sigma": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1000.0, "step": 0.001}),
"seed": ("INT", {"default": 0, "min": 0, "max": 2**31 - 1}),
"auto_download": ("BOOLEAN", {"default": True}),
"unload_comfy_before_pid": ("BOOLEAN", {"default": True}),
"aggressive_cleanup": ("BOOLEAN", {"default": True}),
},
"optional": {
"vae": ("VAE",),
"pid_source_dir": ("STRING", {"default": "", "multiline": False}),
"baseline_image": ("IMAGE",),
},
}
RETURN_TYPES = ("IMAGE",)
RETURN_NAMES = ("image",)
FUNCTION = "decode"
CATEGORY = "PiD"
def decode(
self,
latent,
caption: str,
backbone: str,
pid_ckpt_type: str,
pid_steps: int,
scale: int,
cfg_scale: float,
sigma: float,
seed: int,
auto_download: bool,
unload_comfy_before_pid: bool = True,
aggressive_cleanup: bool = True,
vae=None,
pid_source_dir: str = "",
baseline_image=None,
):
backbone = str(backbone).strip()
pid_ckpt_type = str(pid_ckpt_type).strip()
if backbone not in PID_BACKBONES:
raise PiDNodeError(f"Unknown backbone={backbone!r}; expected one of {BACKBONE_CHOICES}")
backbone_info = PID_BACKBONES[backbone]
ckpt = _checkpoint_for(backbone, pid_ckpt_type)
if int(scale) <= 0:
scale = int(ckpt.scale or backbone_info.default_scale)
pid_dir = _resolve_pid_dir(pid_source_dir)
_ensure_pid_source(pid_dir, allow_download=bool(auto_download))
checkpoint_path = _ensure_checkpoint(pid_dir, backbone, pid_ckpt_type, allow_download=bool(auto_download))
_ensure_backbone_assets(pid_dir, backbone, allow_download=bool(auto_download))
samples = _latent_samples(latent)
if samples.shape[1] != backbone_info.latent_channels:
raise PiDNodeError(
f"{backbone_info.label} PiD expects {backbone_info.latent_channels}-channel latents. "
f"Got {samples.shape[1]} channels."
)
if baseline_image is None:
if vae is None:
raise PiDNodeError(
"PiD needs a baseline image. Connect either a matching ComfyUI VAE "
"or a pre-decoded baseline_image."
)
baseline_01 = _decode_baseline_with_comfy_vae(vae, samples, backbone)
else:
baseline_01 = _comfy_image_to_bchw_01(baseline_image)
if baseline_01.shape[0] != samples.shape[0]:
raise PiDNodeError(
f"Batch mismatch: latent batch={samples.shape[0]}, baseline batch={baseline_01.shape[0]}"
)
# Keep only CPU copies before unloading Z-Image / CLIP / VAE.
# This prevents upstream CUDA tensors or a connected ComfyUI VAE object from
# keeping VRAM alive during the PiD-only stage.
samples_cpu = samples.detach().to("cpu").contiguous()
baseline_cpu = baseline_01.detach().to("cpu").contiguous()
del samples
del baseline_01
vae = None
baseline_image = None
latent = None
# Decode or receive the low-res baseline first, then release Comfy's
# Z-Image/TE/VAE models before loading/running PiD.
if unload_comfy_before_pid:
_free_cuda_memory(aggressive=bool(aggressive_cleanup))
if not torch.cuda.is_available():
raise PiDNodeError("CUDA GPU is required for PiD.")
model = _load_pid_model(
pid_dir=pid_dir,
backbone=backbone,
ckpt_type=pid_ckpt_type,
checkpoint_path=checkpoint_path,
dtype_choice="bf16",
load_ema_to_reg=False,
)
b, _c, h, w = baseline_cpu.shape
device = "cuda"
latent_bf16 = samples_cpu.to(device=device, dtype=torch.bfloat16)
baseline_neg1_1 = (baseline_cpu.to(device=device, dtype=torch.bfloat16) * 2.0) - 1.0
del samples_cpu
del baseline_cpu
caption = caption or ""
data_batch = {
model.config.input_caption_key: [caption] * int(b),
"LQ_video_or_image": baseline_neg1_1,
"LQ_latent": latent_bf16,
"degrade_sigma": torch.full((int(b),), float(sigma), device=device, dtype=torch.float32),
}
infer_image_size = (int(h) * int(scale), int(w) * int(scale))
if model_management is not None:
try:
model_management.throw_exception_if_processing_interrupted()
except Exception:
pass
_free_cuda_memory(aggressive=bool(aggressive_cleanup))
try:
with torch.inference_mode():
out = model.generate_samples_from_batch(
data_batch,
cfg_scale=float(cfg_scale),
num_steps=int(pid_steps),
seed=int(seed),
shift=None,
image_size=infer_image_size,
)
except Exception as exc:
# After a CUDA allocator/internal-assert failure, the process may be in a bad state.
# Still try to free memory and return a useful error to the ComfyUI UI.
del data_batch
del latent_bf16
del baseline_neg1_1
_unload_pid_model(model, aggressive=True)
del model
raise _format_pid_runtime_error(exc, infer_image_size, f"{backbone}/{pid_ckpt_type}", int(scale)) from exc
out = _normalize_pid_samples(out)
image = _bchw_neg1_to_comfy_image(out)
del out
del data_batch
del latent_bf16
del baseline_neg1_1
_unload_pid_model(model, aggressive=bool(aggressive_cleanup))
del model
return (image,)
NODE_CLASS_MAPPINGS = {
"PiDDecode": PiDDecode,
}
NODE_DISPLAY_NAME_MAPPINGS = {
"PiDDecode": "PiD Decode",
}
pid_text_prompt.py
+29
View File
@@ -0,0 +1,29 @@
from __future__ import annotations
class PiDTextPrompt:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"prompt": ("STRING", {"multiline": True, "dynamicPrompts": True, "default": ""}),
}
}
RETURN_TYPES = ("STRING", "STRING")
RETURN_NAMES = ("text", "caption")
FUNCTION = "build"
CATEGORY = "PiD"
def build(self, prompt: str = ""):
text = prompt or ""
return (text, text)
NODE_CLASS_MAPPINGS = {
"PiDTextPrompt": PiDTextPrompt,
}
NODE_DISPLAY_NAME_MAPPINGS = {
"PiDTextPrompt": "PiD Text Prompt",
}
requirements.txt
+23
View File
@@ -0,0 +1,23 @@
# NVIDIA PiD utility/runtime dependencies for the diffusers backbones.
# ComfyUI usually already provides torch; do not install torch from here.
hydra-core==1.3.2
omegaconf==2.3.0
attrs
einops
loguru
termcolor
fvcore
iopath
pynvml
wandb
imageio
opencv-python-headless
pandas
safetensors
huggingface-hub>=1.0
sentencepiece
boto3
botocore
accelerate
transformers>=4.57.0
diffusers>=0.37.0