2056 lines
76 KiB
Python
2056 lines
76 KiB
Python
"""
|
|
ComfyUI-PiD
|
|
|
|
Experimental ComfyUI nodes for NVIDIA PiD (Pixel Diffusion Decoder).
|
|
|
|
PiD is not a native ComfyUI VAE object. NVIDIA's released latent-conditioned
|
|
checkpoints primarily need:
|
|
latent + sigma + caption -> PiD -> final 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 types import MethodType
|
|
from typing import Callable, 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 comfy.utils as comfy_utils
|
|
except Exception: # pragma: no cover - only available inside ComfyUI
|
|
comfy_utils = 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"
|
|
DEFAULT_PID_MODEL_DIR = COMFYUI_DIR / "models" / "nvidia_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"
|
|
PID_PROGRESS_SENTINEL = "__COMFYUI_PID_PROGRESS__"
|
|
AE_REL_PATH = "checkpoints/ae.safetensors"
|
|
AE_FILENAME = "ae.safetensors"
|
|
PID_REQUIRED_SOURCE_FILES = (
|
|
"pid/_src/configs/pid/experiment_2kto4k/sdxl.py",
|
|
"pid/_src/configs/pid/experiment_2kto4k/qwenimage.py",
|
|
"pid/_src/tokenizers/sdxl_vae.py",
|
|
"pid/_src/tokenizers/qwenimage_vae.py",
|
|
)
|
|
|
|
|
|
@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
|
|
latent_downscale: int
|
|
ckpt_types: Tuple[str, ...]
|
|
aux_files: Tuple[str, ...] = ()
|
|
needs_flux_ae: bool = False
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class PiDHFSnapshot:
|
|
repo_id: str
|
|
required_files: Tuple[str, ...]
|
|
|
|
|
|
# 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_2606/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,
|
|
),
|
|
("sdxl", "2kto4k"): PiDCheckpoint(
|
|
experiment="PiD_res2kto4k_sr4x_official_sdxl_distill_4step",
|
|
relpath="checkpoints/PiD_res2kto4k_sr4x_official_sdxl_distill_4step/model_ema_bf16.pth",
|
|
scale=4,
|
|
),
|
|
("qwenimage", "2kto4k"): PiDCheckpoint(
|
|
experiment="PiD_res2kto4k_sr4x_official_qwenimage_distill_4step",
|
|
relpath="checkpoints/PiD_res2kto4k_sr4x_official_qwenimage_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, 8, ("2k", "2kto4k"), needs_flux_ae=True),
|
|
"zimage-turbo": PiDBackbone(
|
|
"Z-Image-Turbo", "zimage", 16, 4, 8, ("2k", "2kto4k"), needs_flux_ae=True
|
|
),
|
|
"flux": PiDBackbone("Flux", "flux", 16, 4, 8, ("2k", "2kto4k"), needs_flux_ae=True),
|
|
"flux2": PiDBackbone(
|
|
"Flux2",
|
|
"flux2",
|
|
128,
|
|
4,
|
|
16,
|
|
("2k", "2kto4k"),
|
|
aux_files=("checkpoints/flux2_ae.safetensors",),
|
|
),
|
|
"flux2-klein-4b": PiDBackbone(
|
|
"Flux2-Klein-4B",
|
|
"flux2",
|
|
128,
|
|
4,
|
|
16,
|
|
("2k", "2kto4k"),
|
|
aux_files=("checkpoints/flux2_ae.safetensors",),
|
|
),
|
|
"flux2-klein-9b": PiDBackbone(
|
|
"Flux2-Klein-9B",
|
|
"flux2",
|
|
128,
|
|
4,
|
|
16,
|
|
("2k", "2kto4k"),
|
|
aux_files=("checkpoints/flux2_ae.safetensors",),
|
|
),
|
|
"sd3": PiDBackbone(
|
|
"SD3",
|
|
"sd3",
|
|
16,
|
|
4,
|
|
8,
|
|
("2k", "2kto4k"),
|
|
aux_files=("checkpoints/sd3_vae/vae/diffusion_pytorch_model.safetensors",),
|
|
),
|
|
"sdxl": PiDBackbone(
|
|
"SDXL",
|
|
"sdxl",
|
|
4,
|
|
4,
|
|
8,
|
|
("2kto4k",),
|
|
aux_files=("checkpoints/sdxl_vae.safetensors",),
|
|
),
|
|
"qwenimage": PiDBackbone(
|
|
"Qwen-Image",
|
|
"qwenimage",
|
|
16,
|
|
4,
|
|
8,
|
|
("2kto4k",),
|
|
aux_files=("checkpoints/QwenImage_VAE_2d.pth",),
|
|
),
|
|
"qwenimage-2512": PiDBackbone(
|
|
"Qwen-Image-2512",
|
|
"qwenimage",
|
|
16,
|
|
4,
|
|
8,
|
|
("2kto4k",),
|
|
aux_files=("checkpoints/QwenImage_VAE_2d.pth",),
|
|
),
|
|
"dinov2": PiDBackbone(
|
|
"DINOv2 RAE",
|
|
"rae",
|
|
768,
|
|
4,
|
|
16,
|
|
("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,
|
|
16,
|
|
("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())
|
|
SEQUENTIAL_OFFLOAD_CHOICES = [
|
|
"auto_low_vram",
|
|
"disabled",
|
|
"sequential_blocks",
|
|
"sequential_blocks_aggressive",
|
|
]
|
|
PID_WEIGHT_PRECISION_CHOICES = ["fp32_compatible", "bf16_weights_experimental"]
|
|
_PIXEL_PATCH_SIZE = 16
|
|
|
|
GEMMA_SNAPSHOT = PiDHFSnapshot(
|
|
repo_id="Efficient-Large-Model/gemma-2-2b-it",
|
|
required_files=(
|
|
"config.json",
|
|
"generation_config.json",
|
|
"model-00001-of-00002.safetensors",
|
|
"model-00002-of-00002.safetensors",
|
|
"model.safetensors.index.json",
|
|
"special_tokens_map.json",
|
|
"tokenizer.json",
|
|
"tokenizer.model",
|
|
"tokenizer_config.json",
|
|
),
|
|
)
|
|
DINOv2_SNAPSHOT = PiDHFSnapshot(
|
|
repo_id="facebook/dinov2-with-registers-base",
|
|
required_files=("config.json", "model.safetensors", "preprocessor_config.json"),
|
|
)
|
|
SIGLIP_SNAPSHOT = PiDHFSnapshot(
|
|
repo_id="google/siglip2-so400m-patch14-224",
|
|
required_files=(
|
|
"config.json",
|
|
"model.safetensors",
|
|
"preprocessor_config.json",
|
|
"special_tokens_map.json",
|
|
"tokenizer.json",
|
|
"tokenizer.model",
|
|
"tokenizer_config.json",
|
|
),
|
|
)
|
|
|
|
_MODEL_CACHE: Dict[Tuple[object, ...], object] = {}
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class PiDMemoryPlan:
|
|
requested_mode: str
|
|
pid_weight_precision: str
|
|
infer_image_size: Tuple[int, int]
|
|
total_pixel_patches: int
|
|
pixel_chunk_patches: int
|
|
chunk_pixel_blocks: bool
|
|
cpu_cache_pixel_positions: bool
|
|
skip_vae_encoder: bool
|
|
offload_patch_blocks: bool
|
|
offload_pixel_blocks: bool
|
|
|
|
|
|
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
|
|
latent to CPU 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 _trim_cuda_cache() -> None:
|
|
"""Release unused CUDA allocator cache without synchronizing live work."""
|
|
if torch.cuda.is_available():
|
|
try:
|
|
torch.cuda.empty_cache()
|
|
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 resolve_pid_memory_plan(
|
|
sequential_offload: str,
|
|
pid_weight_precision: str,
|
|
pixel_chunk_patches: int,
|
|
infer_image_size: Tuple[int, int],
|
|
total_vram_gb: Optional[float] = None,
|
|
supplies_lq_latent: bool = True,
|
|
) -> PiDMemoryPlan:
|
|
mode = str(sequential_offload or "auto_low_vram").strip().lower()
|
|
mode = {"none": "disabled", "off": "disabled"}.get(mode, mode)
|
|
if mode not in SEQUENTIAL_OFFLOAD_CHOICES:
|
|
raise PiDNodeError(
|
|
f"Unknown sequential_offload={mode!r}; expected one of {SEQUENTIAL_OFFLOAD_CHOICES}"
|
|
)
|
|
|
|
precision = str(pid_weight_precision or "fp32_compatible").strip().lower()
|
|
if precision not in PID_WEIGHT_PRECISION_CHOICES:
|
|
raise PiDNodeError(
|
|
f"Unknown pid_weight_precision={precision!r}; expected one of {PID_WEIGHT_PRECISION_CHOICES}"
|
|
)
|
|
|
|
try:
|
|
requested_chunk_patches = int(pixel_chunk_patches)
|
|
except Exception as exc:
|
|
raise PiDNodeError(f"pixel_chunk_patches must be an integer, got {pixel_chunk_patches!r}") from exc
|
|
if requested_chunk_patches < 0:
|
|
raise PiDNodeError("pixel_chunk_patches must be zero (automatic) or a positive integer.")
|
|
|
|
h, w = int(infer_image_size[0]), int(infer_image_size[1])
|
|
if h <= 0 or w <= 0:
|
|
raise PiDNodeError(f"infer_image_size must be positive, got {(h, w)!r}")
|
|
total_pixel_patches = max(1, (h // _PIXEL_PATCH_SIZE) * (w // _PIXEL_PATCH_SIZE))
|
|
|
|
low_vram = mode in ("auto_low_vram", "sequential_blocks_aggressive")
|
|
balanced = mode == "sequential_blocks"
|
|
capacity_gb = _vram_total_gb() if total_vram_gb is None else float(total_vram_gb)
|
|
if requested_chunk_patches > 0:
|
|
chunk_patches = min(total_pixel_patches, requested_chunk_patches)
|
|
elif low_vram:
|
|
automatic_chunk = 4096 if capacity_gb <= 16.5 else 8192
|
|
chunk_patches = min(total_pixel_patches, automatic_chunk)
|
|
elif balanced:
|
|
chunk_patches = min(total_pixel_patches, 8192)
|
|
else:
|
|
chunk_patches = 0
|
|
|
|
chunk_pixel_blocks = 0 < chunk_patches < total_pixel_patches
|
|
return PiDMemoryPlan(
|
|
requested_mode=mode,
|
|
pid_weight_precision=precision,
|
|
infer_image_size=(h, w),
|
|
total_pixel_patches=total_pixel_patches,
|
|
pixel_chunk_patches=chunk_patches,
|
|
chunk_pixel_blocks=chunk_pixel_blocks,
|
|
cpu_cache_pixel_positions=low_vram and chunk_pixel_blocks,
|
|
skip_vae_encoder=bool(supplies_lq_latent) and mode != "disabled",
|
|
offload_patch_blocks=mode != "disabled",
|
|
offload_pixel_blocks=low_vram,
|
|
)
|
|
|
|
|
|
def _log_pid_memory_plan(plan: PiDMemoryPlan) -> None:
|
|
capacity_gb = _vram_total_gb()
|
|
h, w = plan.infer_image_size
|
|
print(
|
|
"[ComfyUI-PiD] memory plan: "
|
|
f"mode={plan.requested_mode}, image={w}x{h}, gpu_capacity={capacity_gb:.1f} GiB, "
|
|
f"pixel_patches={plan.total_pixel_patches}, chunk_patches={plan.pixel_chunk_patches or 'disabled'}, "
|
|
f"precision={plan.pid_weight_precision}, cpu_pos_cache={plan.cpu_cache_pixel_positions}, "
|
|
f"skip_vae={plan.skip_vae_encoder}, patch_offload={plan.offload_patch_blocks}, "
|
|
f"pixel_offload={plan.offload_pixel_blocks}",
|
|
flush=True,
|
|
)
|
|
|
|
|
|
def _reset_cuda_peak_memory_stats() -> None:
|
|
if torch.cuda.is_available():
|
|
try:
|
|
torch.cuda.reset_peak_memory_stats()
|
|
except Exception:
|
|
pass
|
|
|
|
|
|
def _log_cuda_peak_memory(label: str) -> None:
|
|
if not torch.cuda.is_available():
|
|
return
|
|
try:
|
|
allocated_mib = float(torch.cuda.max_memory_allocated()) / (1024 ** 2)
|
|
reserved_mib = float(torch.cuda.max_memory_reserved()) / (1024 ** 2)
|
|
print(
|
|
f"[ComfyUI-PiD] {label}: peak allocated={allocated_mib:.1f} MiB, "
|
|
f"peak reserved={reserved_mib:.1f} MiB",
|
|
flush=True,
|
|
)
|
|
except Exception:
|
|
pass
|
|
|
|
|
|
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. Large outputs should normally use "
|
|
"auto_low_vram so live pixel-block tensors are chunked before they can spill into shared memory.\n\n"
|
|
"Try these first:\n"
|
|
" 1. Set sequential_offload=auto_low_vram and pixel_chunk_patches=0.\n"
|
|
" 2. If memory is still tight, set pixel_chunk_patches=2048.\n"
|
|
" 3. Keep unload_comfy_before_pid=true and aggressive_cleanup=true.\n"
|
|
" 4. Lower scale or use a smaller base latent.\n"
|
|
" 5. Restart ComfyUI after a CUDA allocator crash.\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 _resolve_pid_model_dir() -> Path:
|
|
"""Return the ComfyUI-managed root for NVIDIA PiD weights and assets."""
|
|
if folder_paths is not None:
|
|
models_dir = getattr(folder_paths, "models_dir", None)
|
|
if models_dir:
|
|
return (Path(models_dir) / "nvidia_pid").expanduser().resolve()
|
|
return DEFAULT_PID_MODEL_DIR.resolve()
|
|
|
|
|
|
def _migrate_legacy_checkpoints(model_dir: Path) -> List[str]:
|
|
"""Move bundled legacy checkpoints into ComfyUI's shared model directory."""
|
|
legacy_dir = DEFAULT_PID_DIR / "checkpoints"
|
|
target_dir = model_dir / "checkpoints"
|
|
messages: List[str] = []
|
|
if not legacy_dir.is_dir():
|
|
return messages
|
|
|
|
for src in sorted(path for path in legacy_dir.rglob("*") if path.is_file()):
|
|
relpath = src.relative_to(legacy_dir)
|
|
target = target_dir / relpath
|
|
if target.exists():
|
|
messages.append(f"Legacy PiD asset kept because destination exists: {src}")
|
|
continue
|
|
target.parent.mkdir(parents=True, exist_ok=True)
|
|
shutil.move(str(src), str(target))
|
|
messages.append(f"Moved legacy PiD asset: {src} -> {target}")
|
|
|
|
for directory in sorted((path for path in legacy_dir.rglob("*") if path.is_dir()), reverse=True):
|
|
try:
|
|
directory.rmdir()
|
|
except OSError:
|
|
pass
|
|
try:
|
|
legacy_dir.rmdir()
|
|
except OSError:
|
|
pass
|
|
return messages
|
|
|
|
|
|
def _pid_is_present(pid_dir: Path) -> bool:
|
|
return (pid_dir / "pid" / "_src" / "utils" / "model_loader.py").is_file()
|
|
|
|
|
|
def _missing_pid_source_files(pid_dir: Path) -> List[str]:
|
|
return [relpath for relpath in PID_REQUIRED_SOURCE_FILES if not (pid_dir / relpath).is_file()]
|
|
|
|
|
|
def _required_pid_source_files_for_backbone(backbone: str) -> Tuple[str, ...]:
|
|
if backbone == "sdxl":
|
|
return (
|
|
"pid/_src/configs/pid/experiment_2kto4k/sdxl.py",
|
|
"pid/_src/tokenizers/sdxl_vae.py",
|
|
)
|
|
if backbone in ("qwenimage", "qwenimage-2512"):
|
|
return (
|
|
"pid/_src/configs/pid/experiment_2kto4k/qwenimage.py",
|
|
"pid/_src/tokenizers/qwenimage_vae.py",
|
|
)
|
|
return ()
|
|
|
|
|
|
def _missing_required_pid_source_files(pid_dir: Path, required_files: Tuple[str, ...]) -> List[str]:
|
|
return [relpath for relpath in required_files if not (pid_dir / relpath).is_file()]
|
|
|
|
|
|
def _try_update_managed_pid_source(pid_dir: Path) -> List[str]:
|
|
messages: List[str] = []
|
|
if pid_dir.resolve() != DEFAULT_PID_DIR.resolve():
|
|
return messages
|
|
git = shutil.which("git")
|
|
if not git or not (pid_dir / ".git").is_dir():
|
|
return messages
|
|
messages.append(f"Updating bundled PiD source in {pid_dir} ...")
|
|
_run([git, "-C", str(pid_dir), "pull", "--ff-only"])
|
|
return messages
|
|
|
|
|
|
def _ensure_pid_source(
|
|
pid_dir: Path,
|
|
allow_download: bool = True,
|
|
required_files: Tuple[str, ...] = (),
|
|
) -> List[str]:
|
|
messages: List[str] = []
|
|
if _pid_is_present(pid_dir):
|
|
missing = _missing_required_pid_source_files(pid_dir, required_files)
|
|
if missing and allow_download:
|
|
messages.extend(_try_update_managed_pid_source(pid_dir))
|
|
missing = _missing_required_pid_source_files(pid_dir, required_files)
|
|
if missing:
|
|
missing_text = "\n".join(f" - {pid_dir / relpath}" for relpath in missing)
|
|
raise PiDNodeError(
|
|
"The local NVIDIA PiD source is older than the selected backbone.\n"
|
|
"Update the PiD source checkout, or delete the bundled vendor/PiD folder "
|
|
"and run again with auto_download=true.\n\n"
|
|
f"Missing source files:\n{missing_text}"
|
|
)
|
|
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}")
|
|
missing = _missing_required_pid_source_files(pid_dir, required_files)
|
|
if missing:
|
|
missing_text = "\n".join(f" - {pid_dir / relpath}" for relpath in missing)
|
|
raise PiDNodeError(f"PiD source download finished but selected-backbone files are missing:\n{missing_text}")
|
|
|
|
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>=0.36,<1.0\""
|
|
) from exc
|
|
|
|
|
|
def _hf_snapshot_dir(model_dir: Path, snapshot: PiDHFSnapshot) -> Path:
|
|
return model_dir / "huggingface" / Path(snapshot.repo_id)
|
|
|
|
|
|
def _missing_hf_snapshot_files(model_dir: Path, snapshot: PiDHFSnapshot) -> List[str]:
|
|
snapshot_dir = _hf_snapshot_dir(model_dir, snapshot)
|
|
return [relpath for relpath in snapshot.required_files if not (snapshot_dir / relpath).is_file()]
|
|
|
|
|
|
def _required_hf_snapshots(backbone: str) -> Tuple[PiDHFSnapshot, ...]:
|
|
if backbone not in PID_BACKBONES:
|
|
raise PiDNodeError(f"Unknown backbone={backbone!r}; expected one of {BACKBONE_CHOICES}")
|
|
snapshots = [GEMMA_SNAPSHOT]
|
|
if backbone == "dinov2":
|
|
snapshots.append(DINOv2_SNAPSHOT)
|
|
elif backbone == "siglip":
|
|
snapshots.append(SIGLIP_SNAPSHOT)
|
|
return tuple(snapshots)
|
|
|
|
|
|
def _download_hf_snapshot(model_dir: Path, snapshot: PiDHFSnapshot) -> str:
|
|
snapshot_dir = _hf_snapshot_dir(model_dir, snapshot)
|
|
missing = _missing_hf_snapshot_files(model_dir, snapshot)
|
|
if not missing:
|
|
return f"Hugging Face snapshot already exists: {snapshot_dir}"
|
|
|
|
_ensure_hf_download_available()
|
|
from huggingface_hub import snapshot_download
|
|
|
|
token = os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN")
|
|
snapshot_dir.mkdir(parents=True, exist_ok=True)
|
|
snapshot_download(
|
|
repo_id=snapshot.repo_id,
|
|
local_dir=str(snapshot_dir),
|
|
allow_patterns=list(snapshot.required_files),
|
|
token=token,
|
|
)
|
|
missing = _missing_hf_snapshot_files(model_dir, snapshot)
|
|
if missing:
|
|
missing_text = "\n".join(f" - {snapshot_dir / relpath}" for relpath in missing)
|
|
raise PiDNodeError(
|
|
f"Hugging Face snapshot download finished but required files are missing for {snapshot.repo_id}:\n"
|
|
f"{missing_text}"
|
|
)
|
|
return f"Downloaded Hugging Face snapshot: {snapshot.repo_id} -> {snapshot_dir}"
|
|
|
|
|
|
def _ensure_hf_snapshot(model_dir: Path, snapshot: PiDHFSnapshot, allow_download: bool = True) -> Path:
|
|
snapshot_dir = _hf_snapshot_dir(model_dir, snapshot)
|
|
missing = _missing_hf_snapshot_files(model_dir, snapshot)
|
|
if not missing:
|
|
return snapshot_dir
|
|
if allow_download:
|
|
try:
|
|
_download_hf_snapshot(model_dir, snapshot)
|
|
return snapshot_dir
|
|
except Exception as exc:
|
|
if isinstance(exc, PiDNodeError):
|
|
raise
|
|
raise PiDNodeError(
|
|
f"Could not download required Hugging Face snapshot {snapshot.repo_id} to:\n"
|
|
f" {snapshot_dir}\n\n"
|
|
f"Original download error: {exc}"
|
|
) from exc
|
|
|
|
missing_text = "\n".join(f" - {snapshot_dir / relpath}" for relpath in missing)
|
|
raise PiDNodeError(
|
|
f"PiD requires the local Hugging Face snapshot {snapshot.repo_id} at:\n"
|
|
f" {snapshot_dir}\n\n"
|
|
"auto_download=false enables strict local-only mode. Download or copy these missing files:\n"
|
|
f"{missing_text}"
|
|
)
|
|
|
|
|
|
def _ensure_hf_snapshots(model_dir: Path, backbone: str, allow_download: bool = True) -> None:
|
|
for snapshot in _required_hf_snapshots(backbone):
|
|
_ensure_hf_snapshot(model_dir, snapshot, allow_download=allow_download)
|
|
|
|
|
|
|
|
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. 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(model_dir: Path) -> Optional[str]:
|
|
target = model_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(model_dir: Path) -> str:
|
|
target = model_dir / AE_REL_PATH
|
|
if target.is_file():
|
|
return f"Flux AE VAE already exists: {target}"
|
|
|
|
_ensure_hf_download_available()
|
|
from huggingface_hub import hf_hub_download
|
|
|
|
token = os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN")
|
|
model_dir.mkdir(parents=True, exist_ok=True)
|
|
hf_hub_download(
|
|
repo_id=HF_REPO_ID,
|
|
filename=AE_REL_PATH,
|
|
local_dir=str(model_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(model_dir: Path, allow_download: bool = True) -> Path:
|
|
"""Ensure PiD's required Flux AE exists in the shared model directory."""
|
|
target = model_dir / AE_REL_PATH
|
|
if target.is_file():
|
|
return target
|
|
|
|
copied = _copy_local_ae_to_pid(model_dir)
|
|
if copied and target.is_file():
|
|
return target
|
|
|
|
if allow_download:
|
|
try:
|
|
_download_ae_safetensors(model_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(model_dir: Path, relpath: str) -> str:
|
|
target = model_dir / relpath
|
|
if target.is_file():
|
|
return f"Asset already exists: {target}"
|
|
|
|
_ensure_hf_download_available()
|
|
from huggingface_hub import hf_hub_download
|
|
|
|
token = os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN")
|
|
model_dir.mkdir(parents=True, exist_ok=True)
|
|
hf_hub_download(
|
|
repo_id=HF_REPO_ID,
|
|
filename=relpath,
|
|
local_dir=str(model_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 _normalize_scale_for_checkpoint(backbone: str, ckpt: PiDCheckpoint, scale: int) -> int:
|
|
"""Return a safe scale value while preserving manual low-VRAM overrides."""
|
|
if int(scale) <= 0:
|
|
info = PID_BACKBONES.get(backbone)
|
|
return int(ckpt.scale or (info.default_scale if info is not None else 4))
|
|
scale = int(scale)
|
|
info = PID_BACKBONES.get(backbone)
|
|
expected = int(ckpt.scale or (info.default_scale if info is not None else scale))
|
|
if scale != expected:
|
|
label = info.label if info is not None else backbone
|
|
print(
|
|
f"[ComfyUI-PiD] warning: {label} {ckpt.experiment} was trained for scale={expected}; "
|
|
f"using manual scale={scale}. This can be useful for VRAM tests but is out-of-distribution.",
|
|
flush=True,
|
|
)
|
|
return scale
|
|
|
|
|
|
def _warn_if_non_distilled_step_count(pid_steps: int) -> None:
|
|
try:
|
|
steps = int(pid_steps)
|
|
except Exception:
|
|
return
|
|
if steps != 4:
|
|
print(
|
|
"[ComfyUI-PiD] warning: NVIDIA's released PiD checkpoints are 4-step distilled; "
|
|
f"pid_steps={steps} is experimental/out-of-distribution.",
|
|
flush=True,
|
|
)
|
|
|
|
|
|
def _warn_if_qwen_bf16_precision(backbone: str, pid_weight_precision: str) -> None:
|
|
if (
|
|
str(backbone).strip() in ("qwenimage", "qwenimage-2512")
|
|
and str(pid_weight_precision or "").strip().lower() == "bf16_weights_experimental"
|
|
):
|
|
print(
|
|
"[ComfyUI-PiD] warning: Qwen-Image PiD is running with "
|
|
"pid_weight_precision=bf16_weights_experimental. This option changes output; "
|
|
"use fp32_compatible for the recommended/reference behavior.",
|
|
flush=True,
|
|
)
|
|
|
|
|
|
def _latent_pid_source_backbone(latent: dict) -> Optional[str]:
|
|
if not isinstance(latent, dict):
|
|
return None
|
|
source = latent.get("pid_source_backbone")
|
|
if source is None:
|
|
return None
|
|
source = str(source).strip().lower()
|
|
return source or None
|
|
|
|
|
|
def _validate_latent_source_backbone(latent: dict, backbone: str) -> None:
|
|
source = _latent_pid_source_backbone(latent)
|
|
target = str(backbone).strip()
|
|
if source == "qwenimage" and target not in ("qwenimage", "qwenimage-2512"):
|
|
raise PiDNodeError(
|
|
"This latent was captured from a Qwen-Image model, but PiD is set to "
|
|
f"backbone={target!r}. Select the Qwen-Image PiD backbone "
|
|
"('qwenimage' or 'qwenimage-2512') for this latent."
|
|
)
|
|
|
|
|
|
def _sdxl_vp_from_ve_latent(samples: torch.Tensor, sigma: float) -> Tuple[torch.Tensor, float]:
|
|
"""Convert SDXL Comfy/k-diffusion x_t from VE frame to PiD's VP frame.
|
|
|
|
NVIDIA's PiD SDXL capture path rescales both x_t and sigma by
|
|
sqrt(sigma**2 + 1). For sigma=0 this is an exact no-op.
|
|
"""
|
|
sigma_f = float(sigma)
|
|
if abs(sigma_f) < 1e-12:
|
|
return samples, sigma_f
|
|
denom = float((sigma_f * sigma_f + 1.0) ** 0.5)
|
|
return samples / denom, sigma_f / denom
|
|
|
|
|
|
def _prepare_latent_for_pid_backbone(samples: torch.Tensor, sigma: float, backbone: str) -> Tuple[torch.Tensor, float]:
|
|
if str(backbone).strip() == "sdxl":
|
|
return _sdxl_vp_from_ve_latent(samples, sigma)
|
|
return samples, float(sigma)
|
|
|
|
|
|
def _ensure_checkpoint(model_dir: Path, backbone: str, ckpt_type: str, allow_download: bool = True) -> Path:
|
|
ckpt = _checkpoint_for(backbone, ckpt_type)
|
|
ckpt_path = model_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(model_dir, ckpt.relpath)
|
|
return ckpt_path
|
|
|
|
|
|
def _ensure_backbone_assets(model_dir: Path, backbone: str, allow_download: bool = True) -> None:
|
|
info = PID_BACKBONES[backbone]
|
|
if info.needs_flux_ae:
|
|
_ensure_ae_safetensors(model_dir, allow_download=allow_download)
|
|
|
|
missing = [relpath for relpath in info.aux_files if not (model_dir / relpath).is_file()]
|
|
if missing and not allow_download:
|
|
missing_text = "\n".join(f" - {model_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(model_dir, relpath)
|
|
_ensure_hf_snapshots(model_dir, backbone, allow_download=allow_download)
|
|
|
|
|
|
def _redirect_pid_text_encoder(model_dir: Path) -> None:
|
|
"""Point NVIDIA PiD's built-in Gemma mapping at the local snapshot."""
|
|
try:
|
|
from pid._src.models import pixeldit_model
|
|
except Exception as exc:
|
|
raise PiDNodeError(f"Could not configure NVIDIA PiD's local text encoder path: {exc}") from exc
|
|
pixeldit_model._TEXT_ENCODER_DICT["gemma-2-2b-it"] = _hydra_path(_hf_snapshot_dir(model_dir, GEMMA_SNAPSHOT))
|
|
|
|
|
|
def _hydra_path(path: Path) -> str:
|
|
"""Format an absolute local path for NVIDIA PiD's Hydra overrides."""
|
|
return path.expanduser().resolve().as_posix()
|
|
|
|
|
|
def _pid_asset_experiment_opts(model_dir: Path, backbone: str) -> List[str]:
|
|
"""Redirect NVIDIA PiD's config-relative tokenizer assets to ComfyUI models."""
|
|
def override(field: str, relpath: str) -> str:
|
|
return f"++model.config.tokenizer.{field}={_hydra_path(model_dir / relpath)}"
|
|
|
|
if backbone in ("zimage", "zimage-turbo", "flux"):
|
|
return [override("vae_pth", AE_REL_PATH)]
|
|
if backbone in ("flux2", "flux2-klein-4b", "flux2-klein-9b"):
|
|
return [override("vae_pth", "checkpoints/flux2_ae.safetensors")]
|
|
if backbone == "sd3":
|
|
return [override("vae_pth", "checkpoints/sd3_vae/vae/diffusion_pytorch_model.safetensors")]
|
|
if backbone == "sdxl":
|
|
return [override("vae_pth", "checkpoints/sdxl_vae.safetensors")]
|
|
if backbone in ("qwenimage", "qwenimage-2512"):
|
|
return [override("vae_pth", "checkpoints/QwenImage_VAE_2d.pth")]
|
|
if backbone == "dinov2":
|
|
return [
|
|
override("pretrained_path", f"huggingface/{DINOv2_SNAPSHOT.repo_id}"),
|
|
override("pretrained_decoder_path", "checkpoints/rae/decoders/dinov2/wReg_base/ViTXL_n08_i512/model.pt"),
|
|
override("normalization_stat_path", "checkpoints/rae/stats/dinov2/wReg_base/imagenet1k_512/stat.pt"),
|
|
]
|
|
if backbone == "siglip":
|
|
return [
|
|
override("pretrained_path", f"huggingface/{SIGLIP_SNAPSHOT.repo_id}"),
|
|
override("pretrained_decoder_path", "checkpoints/scale_rae/decoder/siglip2_sop14_i224_web73M_ganw3_decXL.pt"),
|
|
override("decoder_config_path", "checkpoints/scale_rae/decoder/XL_decoder_config.json"),
|
|
]
|
|
raise PiDNodeError(f"Unknown backbone={backbone!r}; expected one of {BACKBONE_CHOICES}")
|
|
|
|
|
|
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,
|
|
model_dir: Path,
|
|
backbone: str,
|
|
ckpt_type: str,
|
|
checkpoint_path: Path,
|
|
dtype_choice: str,
|
|
load_ema_to_reg: bool = False,
|
|
skip_vae_encoder: bool = False,
|
|
pid_weight_precision: str = "fp32_compatible",
|
|
):
|
|
precision = str(pid_weight_precision or "fp32_compatible").strip().lower()
|
|
if precision not in PID_WEIGHT_PRECISION_CHOICES:
|
|
raise PiDNodeError(
|
|
f"Unknown pid_weight_precision={precision!r}; expected one of {PID_WEIGHT_PRECISION_CHOICES}"
|
|
)
|
|
experiment = _checkpoint_for(backbone, ckpt_type).experiment
|
|
config_file = "pid/_src/configs/pid/config.py"
|
|
cache_key = (
|
|
str(pid_dir),
|
|
str(model_dir),
|
|
backbone,
|
|
ckpt_type,
|
|
str(checkpoint_path),
|
|
experiment,
|
|
bool(load_ema_to_reg),
|
|
dtype_choice,
|
|
bool(skip_vae_encoder),
|
|
precision,
|
|
)
|
|
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)
|
|
_redirect_pid_text_encoder(model_dir)
|
|
experiment_opts = _pid_asset_experiment_opts(model_dir, backbone)
|
|
if skip_vae_encoder:
|
|
experiment_opts = [*experiment_opts, "model.config.tokenizer=null"]
|
|
|
|
# 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=experiment_opts,
|
|
strict=False,
|
|
load_ema_to_reg=load_ema_to_reg,
|
|
)
|
|
model.eval()
|
|
if precision == "bf16_weights_experimental":
|
|
net = getattr(model, "net", None)
|
|
if net is None:
|
|
raise PiDNodeError("Loaded PiD model has no network to cast to BF16 weights.")
|
|
net.to(dtype=torch.bfloat16)
|
|
_trim_cuda_cache()
|
|
_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 _pid_lq_condition_type(model: object) -> str:
|
|
config = getattr(model, "config", None)
|
|
return str(getattr(config, "lq_condition_type", "latent") or "latent").strip().lower()
|
|
|
|
|
|
def _pid_uses_lq_image(model: object) -> bool:
|
|
return _pid_lq_condition_type(model) in ("image", "image_latent")
|
|
|
|
|
|
def _pid_uses_lq_latent(model: object) -> bool:
|
|
return _pid_lq_condition_type(model) in ("latent", "image_latent")
|
|
|
|
|
|
def _move_module_like_to_cpu(obj: object) -> None:
|
|
if obj is None:
|
|
return
|
|
nested_model = getattr(obj, "model", None)
|
|
nested_nested_model = getattr(nested_model, "model", None)
|
|
for target in (obj, nested_model, nested_nested_model):
|
|
if target is None:
|
|
continue
|
|
try:
|
|
to = getattr(target, "to", None)
|
|
if callable(to):
|
|
to("cpu")
|
|
except Exception:
|
|
pass
|
|
|
|
|
|
def _offload_unused_pid_conditioners(model: object, include_vae: bool = True) -> None:
|
|
_move_module_like_to_cpu(getattr(model, "text_encoder", None))
|
|
null_caption = getattr(model, "_null_caption_embs", None)
|
|
if isinstance(null_caption, torch.Tensor):
|
|
try:
|
|
setattr(model, "_null_caption_embs", null_caption.cpu())
|
|
except Exception:
|
|
pass
|
|
if include_vae:
|
|
_move_module_like_to_cpu(getattr(model, "vae_encoder", None))
|
|
_trim_cuda_cache()
|
|
|
|
|
|
class _LazyLQFeatureList:
|
|
"""Compute PiD LQ projection heads just-in-time instead of keeping all heads."""
|
|
|
|
def __init__(self, tokens: torch.Tensor, output_heads, pit_head):
|
|
self.tokens = tokens
|
|
self.output_heads = output_heads
|
|
self.pit_head = pit_head
|
|
|
|
def __len__(self) -> int:
|
|
return len(self.output_heads) + (1 if self.pit_head is not None else 0)
|
|
|
|
def __getitem__(self, idx: int) -> torch.Tensor:
|
|
length = len(self)
|
|
if idx < 0:
|
|
idx += length
|
|
if idx < 0 or idx >= length:
|
|
raise IndexError(idx)
|
|
if idx < len(self.output_heads):
|
|
return self.output_heads[idx](self.tokens)
|
|
if self.pit_head is None:
|
|
raise IndexError(idx)
|
|
return self.pit_head(self.tokens)
|
|
|
|
|
|
def _enable_lazy_lq_projection(model: object) -> None:
|
|
"""Patch PiD's LQ projector so 4K runs do not cache every block feature."""
|
|
net = getattr(model, "net", None)
|
|
lq_proj = getattr(net, "lq_proj", None)
|
|
if lq_proj is None or getattr(lq_proj, "_comfy_pid_lazy_lq", False):
|
|
return
|
|
if getattr(net, "_cp_group", None) is not None:
|
|
return
|
|
needed = (
|
|
"image_conv",
|
|
"latent_proj",
|
|
"merge",
|
|
"output_heads",
|
|
"pit_head",
|
|
"_align_image_to_patch_grid",
|
|
"_align_latent_to_patch_grid",
|
|
)
|
|
if any(not hasattr(lq_proj, name) for name in needed):
|
|
return
|
|
|
|
def lazy_forward(
|
|
lq_video_or_image=None,
|
|
lq_latent=None,
|
|
target_pH: int = 0,
|
|
target_pW: int = 0,
|
|
):
|
|
if target_pH <= 0 or target_pW <= 0:
|
|
raise AssertionError("Must provide target_pH and target_pW")
|
|
|
|
features = []
|
|
if lq_proj.image_conv is not None and lq_video_or_image is not None:
|
|
features.append(lq_proj._align_image_to_patch_grid(lq_video_or_image, target_pH, target_pW))
|
|
if lq_proj.latent_proj is not None and lq_latent is not None:
|
|
features.append(lq_proj._align_latent_to_patch_grid(lq_latent, target_pH, target_pW))
|
|
|
|
if len(features) == 2 and lq_proj.merge is not None:
|
|
merged = lq_proj.merge(torch.cat(features, dim=1))
|
|
elif len(features) == 1:
|
|
merged = features[0]
|
|
else:
|
|
ref = lq_video_or_image if lq_video_or_image is not None else lq_latent
|
|
if ref is None:
|
|
raise PiDNodeError("PiD low-VRAM LQ projection received neither LQ image nor latent.")
|
|
b, device, dtype = ref.shape[0], ref.device, ref.dtype
|
|
hidden_dim = int(getattr(lq_proj, "hidden_dim", 0) or getattr(lq_proj, "out_dim", 0))
|
|
tokens = torch.zeros(
|
|
b,
|
|
int(target_pH) * int(target_pW),
|
|
hidden_dim,
|
|
device=device,
|
|
dtype=dtype,
|
|
)
|
|
return _LazyLQFeatureList(tokens, lq_proj.output_heads, lq_proj.pit_head)
|
|
|
|
tokens = merged.flatten(2).transpose(1, 2)
|
|
return _LazyLQFeatureList(tokens, lq_proj.output_heads, lq_proj.pit_head)
|
|
|
|
lq_proj.forward = lazy_forward
|
|
lq_proj._comfy_pid_lazy_lq = True
|
|
|
|
|
|
def _make_pid_progress_bar(pid_steps: int):
|
|
if comfy_utils is None or not getattr(comfy_utils, "PROGRESS_BAR_ENABLED", True):
|
|
return None
|
|
try:
|
|
steps = max(1, int(pid_steps))
|
|
return comfy_utils.ProgressBar(steps)
|
|
except Exception:
|
|
return None
|
|
|
|
|
|
def _update_pid_progress_bar(pbar, current: int, total: int) -> None:
|
|
if pbar is None:
|
|
return
|
|
try:
|
|
total = max(1, int(total))
|
|
current = min(max(0, int(current)), total)
|
|
pbar.update_absolute(current, total)
|
|
except Exception:
|
|
pass
|
|
|
|
|
|
class _PiDProgressNet:
|
|
def __init__(self, net: object, total_steps: int, progress_callback: Callable[[int, int], None], state: dict):
|
|
self._net = net
|
|
self._total_steps = max(1, int(total_steps))
|
|
self._progress_callback = progress_callback
|
|
self._state = state
|
|
|
|
def __getattr__(self, name: str):
|
|
return getattr(self._net, name)
|
|
|
|
def __call__(self, *args, **kwargs):
|
|
result = self._net(*args, **kwargs)
|
|
current = min(int(self._state.get("current", 0)) + 1, self._total_steps)
|
|
self._state["current"] = current
|
|
try:
|
|
self._progress_callback(current, self._total_steps)
|
|
except Exception:
|
|
pass
|
|
return result
|
|
|
|
|
|
@contextlib.contextmanager
|
|
def _pid_step_progress(model: object, total_steps: int, progress_callback: Optional[Callable[[int, int], None]]):
|
|
if progress_callback is None:
|
|
yield
|
|
return
|
|
|
|
try:
|
|
total_steps = max(1, int(total_steps))
|
|
except Exception:
|
|
yield
|
|
return
|
|
|
|
state = {"current": 0}
|
|
original_maybe_compile_net = getattr(model, "_maybe_compile_net", None)
|
|
if callable(original_maybe_compile_net):
|
|
|
|
def maybe_compile_net_with_progress(*args, **kwargs):
|
|
net = original_maybe_compile_net(*args, **kwargs)
|
|
return _PiDProgressNet(net, total_steps, progress_callback, state)
|
|
|
|
try:
|
|
setattr(model, "_maybe_compile_net", maybe_compile_net_with_progress)
|
|
yield
|
|
finally:
|
|
try:
|
|
setattr(model, "_maybe_compile_net", original_maybe_compile_net)
|
|
except Exception:
|
|
pass
|
|
return
|
|
|
|
original_net = getattr(model, "net", None)
|
|
if original_net is None:
|
|
yield
|
|
return
|
|
|
|
try:
|
|
setattr(model, "net", _PiDProgressNet(original_net, total_steps, progress_callback, state))
|
|
yield
|
|
finally:
|
|
try:
|
|
setattr(model, "net", original_net)
|
|
except Exception:
|
|
pass
|
|
|
|
|
|
def _generate_samples_low_vram(model: object, data_batch: dict, progress_callback: Optional[Callable[[int, int], None]] = None, **kwargs):
|
|
"""Run PiD with text/VAE offload and lazy LQ conditioning."""
|
|
_enable_lazy_lq_projection(model)
|
|
|
|
batch = dict(data_batch)
|
|
condition_type = _pid_lq_condition_type(model)
|
|
if condition_type == "latent":
|
|
batch.pop("LQ_video_or_image", None)
|
|
elif condition_type == "image":
|
|
batch.pop("LQ_latent", None)
|
|
|
|
original_encode = getattr(model, "_encode_text_raw", None)
|
|
if not callable(original_encode):
|
|
total_steps = kwargs.get("num_steps", getattr(getattr(model, "config", None), "student_sample_steps", 1))
|
|
with _pid_step_progress(model, total_steps, progress_callback):
|
|
return model.generate_samples_from_batch(batch, **kwargs)
|
|
|
|
def encode_then_offload(captions):
|
|
text_encoder = getattr(model, "text_encoder", None)
|
|
try:
|
|
to = getattr(text_encoder, "to", None)
|
|
if callable(to):
|
|
to("cuda")
|
|
except Exception:
|
|
pass
|
|
result = original_encode(captions)
|
|
_offload_unused_pid_conditioners(model, include_vae="LQ_latent" in batch)
|
|
return result
|
|
|
|
try:
|
|
setattr(model, "_encode_text_raw", encode_then_offload)
|
|
total_steps = kwargs.get("num_steps", getattr(getattr(model, "config", None), "student_sample_steps", 1))
|
|
with _pid_step_progress(model, total_steps, progress_callback):
|
|
return model.generate_samples_from_batch(batch, **kwargs)
|
|
finally:
|
|
try:
|
|
setattr(model, "_encode_text_raw", original_encode)
|
|
except Exception:
|
|
pass
|
|
|
|
|
|
def _make_pid_data_batch(
|
|
model: object,
|
|
caption: str,
|
|
sigma: float,
|
|
latent_cpu: torch.Tensor,
|
|
device: str,
|
|
) -> dict:
|
|
# Free heavyweight PiD conditioners before moving LQ inputs onto CUDA. The
|
|
# text encoder is moved back only for caption encoding inside the sampler.
|
|
if _pid_uses_lq_image(model) or not _pid_uses_lq_latent(model):
|
|
raise PiDNodeError(
|
|
"This ComfyUI-PiD build only supports NVIDIA latent-conditioned PiD checkpoints. "
|
|
"The loaded checkpoint is configured for image LQ conditioning, which this node no longer exposes."
|
|
)
|
|
_offload_unused_pid_conditioners(model, include_vae=True)
|
|
batch = int(latent_cpu.shape[0])
|
|
return {
|
|
getattr(model.config, "input_caption_key", "caption"): [caption or ""] * batch,
|
|
"degrade_sigma": torch.full((batch,), float(sigma), device=device, dtype=torch.float32),
|
|
"LQ_latent": latent_cpu.to(device=device, dtype=torch.bfloat16),
|
|
}
|
|
|
|
|
|
def _infer_lq_size_from_latent(samples: torch.Tensor, backbone: str) -> Tuple[int, int]:
|
|
info = PID_BACKBONES.get(backbone, PID_BACKBONES["zimage"])
|
|
return int(samples.shape[-2]) * int(info.latent_downscale), int(samples.shape[-1]) * int(info.latent_downscale)
|
|
|
|
|
|
def _apply_adaln(x: torch.Tensor, shift: torch.Tensor, scale: torch.Tensor) -> torch.Tensor:
|
|
return x * (1 + scale) + shift
|
|
|
|
|
|
def _enable_cpu_cached_pixel_positions(model: object, chunk_patches: int) -> None:
|
|
"""Patch the pixel embedder to keep full-image position data in system RAM."""
|
|
net = getattr(model, "net", None)
|
|
embedder = getattr(net, "pixel_embedder", None)
|
|
if embedder is None:
|
|
raise PiDNodeError("Loaded PiD model has no pixel embedder for low-VRAM position caching.")
|
|
if getattr(embedder, "_comfy_pid_cpu_position_chunks", 0) == int(chunk_patches):
|
|
return
|
|
if not hasattr(embedder, "_fetch_pixel_pos_image") or not hasattr(embedder, "proj"):
|
|
raise PiDNodeError("Loaded PiD pixel embedder is incompatible with low-VRAM position caching.")
|
|
|
|
original_forward = getattr(embedder, "_comfy_pid_original_forward", None) or embedder.forward
|
|
setattr(embedder, "_comfy_pid_original_forward", original_forward)
|
|
chunk_patches = max(1, int(chunk_patches))
|
|
|
|
def chunked_forward(self, inputs, img_height=None, img_width=None, patch_size=None):
|
|
if inputs.dim() != 4:
|
|
return original_forward(inputs, img_height=img_height, img_width=img_width, patch_size=patch_size)
|
|
if img_height is None or img_width is None or patch_size is None:
|
|
raise AssertionError("Need H, W, patch_size for image mode")
|
|
|
|
b, _, h, w = inputs.shape
|
|
if h != img_height or w != img_width:
|
|
raise AssertionError("Input spatial size mismatch")
|
|
if (h % patch_size) != 0 or (w % patch_size) != 0:
|
|
raise AssertionError("H and W must be divisible by patch_size")
|
|
|
|
hs, ws = h // patch_size, w // patch_size
|
|
patch_count = hs * ws
|
|
p2 = patch_size * patch_size
|
|
pos_key = ("image_patches_cpu", h, w, patch_size, str(inputs.dtype))
|
|
pos_patches = self._pos_cache.get(pos_key)
|
|
if pos_patches is None:
|
|
pos_full = self._fetch_pixel_pos_image(h, w, "cpu", inputs.dtype)
|
|
pos_patches = pos_full.view(hs, patch_size, ws, patch_size, self.hidden_size_output)
|
|
pos_patches = pos_patches.permute(0, 2, 1, 3, 4).contiguous()
|
|
pos_patches = pos_patches.view(patch_count, p2, self.hidden_size_output)
|
|
self._pos_cache[pos_key] = pos_patches
|
|
self._pos_cache.pop(("image", h, w), None)
|
|
|
|
projected = inputs.permute(0, 2, 3, 1).contiguous()
|
|
projected = self.proj(projected)
|
|
projected = projected.view(b, hs, patch_size, ws, patch_size, self.hidden_size_output)
|
|
projected = projected.permute(0, 1, 3, 2, 4, 5).contiguous()
|
|
projected = projected.view(b * patch_count, p2, self.hidden_size_output)
|
|
output = None
|
|
for batch_index in range(b):
|
|
batch_offset = batch_index * patch_count
|
|
for patch_start in range(0, patch_count, chunk_patches):
|
|
patch_end = min(patch_count, patch_start + chunk_patches)
|
|
flat_start = batch_offset + patch_start
|
|
flat_end = batch_offset + patch_end
|
|
pos = pos_patches[patch_start:patch_end].to(device=projected.device, dtype=inputs.dtype)
|
|
chunk = projected[flat_start:flat_end] + pos
|
|
if output is None:
|
|
output = chunk.new_empty(projected.shape)
|
|
output[flat_start:flat_end] = chunk
|
|
_trim_cuda_cache()
|
|
return output
|
|
|
|
embedder.forward = MethodType(chunked_forward, embedder)
|
|
embedder._comfy_pid_cpu_position_chunks = chunk_patches
|
|
|
|
|
|
def _enable_chunked_pixel_blocks(model: object, chunk_patches: int) -> None:
|
|
"""Chunk pixel-block AdaLN and MLP while retaining one global attention pass."""
|
|
net = getattr(model, "net", None)
|
|
blocks = getattr(net, "pixel_blocks", None)
|
|
if blocks is None:
|
|
raise PiDNodeError("Loaded PiD model has no pixel blocks for low-VRAM chunking.")
|
|
chunk_patches = max(1, int(chunk_patches))
|
|
|
|
for block in blocks:
|
|
if getattr(block, "_comfy_pid_chunk_patches", 0) == chunk_patches:
|
|
continue
|
|
needed = (
|
|
"pixel_dim",
|
|
"patch_size",
|
|
"attn_dim",
|
|
"adaLN_modulation",
|
|
"norm1",
|
|
"compress_to_attn",
|
|
"attn",
|
|
"expand_from_attn",
|
|
"norm2",
|
|
"mlp",
|
|
"_fetch_pos",
|
|
)
|
|
if any(not hasattr(block, name) for name in needed):
|
|
raise PiDNodeError("Loaded PiD pixel block is incompatible with exact low-VRAM chunking.")
|
|
|
|
original_forward = getattr(block, "_comfy_pid_original_forward", None) or block.forward
|
|
setattr(block, "_comfy_pid_original_forward", original_forward)
|
|
|
|
def chunked_forward(
|
|
self,
|
|
x,
|
|
s_cond,
|
|
image_height,
|
|
image_width,
|
|
patch_size,
|
|
mask=None,
|
|
_chunk_patches=chunk_patches,
|
|
_original_forward=original_forward,
|
|
):
|
|
if getattr(self, "_cp_group", None) is not None:
|
|
return _original_forward(x, s_cond, image_height, image_width, patch_size, mask)
|
|
|
|
bl, p2, channels = x.shape
|
|
if channels != self.pixel_dim:
|
|
raise ValueError(f"PiTBlock expected pixel_dim={self.pixel_dim}, got {channels}")
|
|
if patch_size != self.patch_size:
|
|
raise AssertionError("PiTBlock expects fixed patch_size")
|
|
if p2 != patch_size * patch_size:
|
|
raise AssertionError("Token count per patch must equal patch_size^2")
|
|
if (image_height % patch_size) != 0 or (image_width % patch_size) != 0:
|
|
raise AssertionError("H and W must be divisible by patch_size")
|
|
|
|
hs, ws = image_height // patch_size, image_width // patch_size
|
|
local_patch_count = hs * ws
|
|
if bl % local_patch_count != 0:
|
|
raise AssertionError("Total sequences must be a multiple of local patch count")
|
|
batch = bl // local_patch_count
|
|
|
|
normalized = torch.empty_like(x)
|
|
for start in range(0, bl, _chunk_patches):
|
|
end = min(bl, start + _chunk_patches)
|
|
params = self.adaLN_modulation(s_cond[start:end]).view(end - start, p2, 6 * self.pixel_dim)
|
|
shift_msa, scale_msa, _, _, _, _ = torch.chunk(params, 6, dim=-1)
|
|
normalized[start:end] = _apply_adaln(self.norm1(x[start:end]), shift_msa, scale_msa)
|
|
|
|
compressed = self.compress_to_attn(normalized.view(bl, p2 * self.pixel_dim))
|
|
del normalized
|
|
pos = self._fetch_pos(hs, ws, x.device)
|
|
attention = self.attn(compressed.view(batch, local_patch_count, self.attn_dim), pos, mask)
|
|
del compressed
|
|
expanded = self.expand_from_attn(attention.reshape(bl, self.attn_dim)).view(bl, p2, self.pixel_dim)
|
|
del attention
|
|
output = torch.empty_like(x)
|
|
for start in range(0, bl, _chunk_patches):
|
|
end = min(bl, start + _chunk_patches)
|
|
params = self.adaLN_modulation(s_cond[start:end]).view(end - start, p2, 6 * self.pixel_dim)
|
|
shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = torch.chunk(params, 6, dim=-1)
|
|
residual = x[start:end] + gate_msa * expanded[start:end]
|
|
mlp = self.mlp(_apply_adaln(self.norm2(residual), shift_mlp, scale_mlp))
|
|
output[start:end] = residual + gate_mlp * mlp
|
|
return output
|
|
|
|
block.forward = MethodType(chunked_forward, block)
|
|
block._comfy_pid_chunk_patches = chunk_patches
|
|
|
|
|
|
|
|
class _SequentialBlockOffloader:
|
|
"""Move selected PiD block stacks to CUDA only for their forward calls."""
|
|
|
|
def __init__(
|
|
self,
|
|
model,
|
|
mode: str,
|
|
device: str = "cuda",
|
|
offload_patch_blocks: bool = True,
|
|
offload_pixel_blocks: bool = False,
|
|
):
|
|
self.model = model
|
|
self.mode = str(mode or "disabled").strip().lower()
|
|
self.device = device
|
|
self.handles = []
|
|
self.blocks = []
|
|
self.container_names = []
|
|
self.pixel_block_ids = set()
|
|
if self.mode == "disabled":
|
|
return
|
|
if self.mode not in SEQUENTIAL_OFFLOAD_CHOICES:
|
|
raise PiDNodeError(
|
|
f"Unknown sequential_offload={self.mode!r}; expected one of {SEQUENTIAL_OFFLOAD_CHOICES}"
|
|
)
|
|
net = getattr(model, "net", None)
|
|
if offload_patch_blocks:
|
|
self._append_stack(getattr(net, "patch_blocks", None), "net.patch_blocks")
|
|
if offload_pixel_blocks:
|
|
pixel_blocks = self._append_stack(getattr(net, "pixel_blocks", None), "net.pixel_blocks")
|
|
self.pixel_block_ids.update(id(block) for block in pixel_blocks)
|
|
if not self.blocks:
|
|
fallback, container_name = self._find_largest_block_stack(model)
|
|
self._append_stack(fallback, container_name)
|
|
if not self.blocks:
|
|
raise PiDNodeError(
|
|
"Sequential block offload could not find a transformer/DiT block stack in the loaded PiD model. "
|
|
"Set sequential_offload=disabled for this model."
|
|
)
|
|
self._install()
|
|
|
|
def _append_stack(self, stack, name: str):
|
|
if stack is None:
|
|
return []
|
|
try:
|
|
blocks = list(stack)
|
|
except Exception:
|
|
return []
|
|
existing = {id(block) for block in self.blocks}
|
|
appended = [block for block in blocks if id(block) not in existing]
|
|
if appended:
|
|
self.blocks.extend(appended)
|
|
self.container_names.append(name)
|
|
return appended
|
|
|
|
@staticmethod
|
|
def _module_param_count(module) -> int:
|
|
try:
|
|
return sum(int(p.numel()) for p in module.parameters(recurse=True))
|
|
except Exception:
|
|
return 0
|
|
|
|
@staticmethod
|
|
def _children(module):
|
|
try:
|
|
return list(module.children())
|
|
except Exception:
|
|
return []
|
|
|
|
def _find_largest_block_stack(self, model):
|
|
candidates = []
|
|
for name, module in model.named_modules():
|
|
children = self._children(module)
|
|
if len(children) < 4:
|
|
continue
|
|
cls_name = module.__class__.__name__.lower()
|
|
name_l = str(name).lower()
|
|
child_classes = " ".join(c.__class__.__name__.lower() for c in children[:8])
|
|
looks_like_blocks = (
|
|
"blocks" in name_l
|
|
or "block" in cls_name
|
|
or "block" in child_classes
|
|
or "transformer" in name_l
|
|
or "dit" in name_l
|
|
)
|
|
if not looks_like_blocks:
|
|
continue
|
|
child_param_counts = [self._module_param_count(c) for c in children]
|
|
total_params = sum(child_param_counts)
|
|
# Avoid tiny helper containers; a real PiD block stack is large.
|
|
if total_params < 10_000_000:
|
|
continue
|
|
candidates.append((total_params, name or module.__class__.__name__, children))
|
|
|
|
if not candidates:
|
|
return [], ""
|
|
candidates.sort(key=lambda item: item[0], reverse=True)
|
|
_total, name, blocks = candidates[0]
|
|
return blocks, name
|
|
|
|
def _install(self) -> None:
|
|
# Move selected blocks to CPU before the sample loop starts. The hooks
|
|
# bring each block back only for its own forward pass.
|
|
for block in self.blocks:
|
|
try:
|
|
block.to("cpu")
|
|
except Exception:
|
|
pass
|
|
_trim_cuda_cache()
|
|
|
|
for block in self.blocks:
|
|
self.handles.append(block.register_forward_pre_hook(self._pre_forward))
|
|
self.handles.append(block.register_forward_hook(self._post_forward))
|
|
|
|
def _pre_forward(self, module, inputs):
|
|
try:
|
|
module.to(self.device)
|
|
except Exception as exc:
|
|
raise PiDNodeError(f"Sequential block offload failed moving a PiD block to CUDA: {exc}") from exc
|
|
return None
|
|
|
|
def _post_forward(self, module, inputs, output):
|
|
try:
|
|
module.to("cpu")
|
|
except Exception:
|
|
pass
|
|
if id(module) in self.pixel_block_ids:
|
|
_trim_cuda_cache()
|
|
return output
|
|
|
|
def cleanup(self) -> None:
|
|
for handle in self.handles:
|
|
try:
|
|
handle.remove()
|
|
except Exception:
|
|
pass
|
|
self.handles.clear()
|
|
_trim_cuda_cache()
|
|
|
|
|
|
def configure_pid_runtime(model: object, plan: PiDMemoryPlan, device: str = "cuda"):
|
|
if plan.cpu_cache_pixel_positions:
|
|
_enable_cpu_cached_pixel_positions(model, plan.pixel_chunk_patches)
|
|
if plan.chunk_pixel_blocks:
|
|
_enable_chunked_pixel_blocks(model, plan.pixel_chunk_patches)
|
|
if not plan.offload_patch_blocks and not plan.offload_pixel_blocks:
|
|
return None
|
|
return _SequentialBlockOffloader(
|
|
model,
|
|
plan.requested_mode,
|
|
device=device,
|
|
offload_patch_blocks=plan.offload_patch_blocks,
|
|
offload_pixel_blocks=plan.offload_pixel_blocks,
|
|
)
|
|
|
|
|
|
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 == 5:
|
|
if samples.shape[2] == 1:
|
|
samples = samples[:, :, 0, :, :]
|
|
elif samples.shape[1] == 1:
|
|
samples = samples[:, 0, :, :, :]
|
|
else:
|
|
raise PiDNodeError(
|
|
"Expected latent samples as [B,C,H,W], [B,C,1,H,W], or [B,1,C,H,W], "
|
|
f"got shape {list(samples.shape)}"
|
|
)
|
|
if samples.ndim != 4:
|
|
raise PiDNodeError(f"Expected latent samples as [B,C,H,W], got shape {list(samples.shape)}")
|
|
return samples
|
|
|
|
|
|
def _latent_pid_sigma(latent: dict, fallback: float) -> float:
|
|
if isinstance(latent, dict) and "pid_sigma" in latent and abs(float(fallback)) < 1e-12:
|
|
try:
|
|
return float(latent["pid_sigma"])
|
|
except Exception:
|
|
pass
|
|
return float(fallback)
|
|
|
|
|
|
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)}")
|
|
|
|
|
|
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": {
|
|
"pid_source_dir": ("STRING", {"default": "", "multiline": False}),
|
|
"sequential_offload": (SEQUENTIAL_OFFLOAD_CHOICES, {"default": "auto_low_vram"}),
|
|
"pid_weight_precision": (PID_WEIGHT_PRECISION_CHOICES, {"default": "fp32_compatible"}),
|
|
"pixel_chunk_patches": ("INT", {"default": 0, "min": 0, "max": 65536, "step": 1024}),
|
|
},
|
|
}
|
|
|
|
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,
|
|
pid_source_dir: str = "",
|
|
sequential_offload: str = "auto_low_vram",
|
|
pid_weight_precision: str = "fp32_compatible",
|
|
pixel_chunk_patches: int = 0,
|
|
):
|
|
backbone = str(backbone).strip()
|
|
pid_ckpt_type = str(pid_ckpt_type).strip()
|
|
sequential_offload = str(sequential_offload or "auto_low_vram").strip().lower()
|
|
|
|
if sequential_offload not in SEQUENTIAL_OFFLOAD_CHOICES:
|
|
raise PiDNodeError(
|
|
f"Unknown sequential_offload={sequential_offload!r}; expected one of {SEQUENTIAL_OFFLOAD_CHOICES}"
|
|
)
|
|
|
|
if backbone not in PID_BACKBONES:
|
|
raise PiDNodeError(f"Unknown backbone={backbone!r}; expected one of {BACKBONE_CHOICES}")
|
|
_validate_latent_source_backbone(latent, backbone)
|
|
backbone_info = PID_BACKBONES[backbone]
|
|
ckpt = _checkpoint_for(backbone, pid_ckpt_type)
|
|
scale = _normalize_scale_for_checkpoint(backbone, ckpt, int(scale))
|
|
|
|
pid_dir = _resolve_pid_dir(pid_source_dir)
|
|
model_dir = _resolve_pid_model_dir()
|
|
_ensure_pid_source(
|
|
pid_dir,
|
|
allow_download=bool(auto_download),
|
|
required_files=_required_pid_source_files_for_backbone(backbone),
|
|
)
|
|
_migrate_legacy_checkpoints(model_dir)
|
|
checkpoint_path = _ensure_checkpoint(model_dir, backbone, pid_ckpt_type, allow_download=bool(auto_download))
|
|
_ensure_backbone_assets(model_dir, backbone, allow_download=bool(auto_download))
|
|
|
|
samples = _latent_samples(latent)
|
|
sigma = _latent_pid_sigma(latent, sigma)
|
|
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."
|
|
)
|
|
|
|
# Keep only a CPU latent copy before unloading Z-Image / CLIP / VAE.
|
|
# PiD's released latent-conditioned checkpoints do not need any image input.
|
|
samples_cpu = samples.detach().to("cpu").contiguous()
|
|
samples_cpu, sigma = _prepare_latent_for_pid_backbone(samples_cpu, sigma, backbone)
|
|
h, w = _infer_lq_size_from_latent(samples, backbone)
|
|
del samples
|
|
latent = None
|
|
infer_image_size = (int(h) * int(scale), int(w) * int(scale))
|
|
plan = resolve_pid_memory_plan(
|
|
sequential_offload=sequential_offload,
|
|
pid_weight_precision=pid_weight_precision,
|
|
pixel_chunk_patches=pixel_chunk_patches,
|
|
infer_image_size=infer_image_size,
|
|
)
|
|
_warn_if_qwen_bf16_precision(backbone, plan.pid_weight_precision)
|
|
_log_pid_memory_plan(plan)
|
|
|
|
# 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,
|
|
model_dir=model_dir,
|
|
backbone=backbone,
|
|
ckpt_type=pid_ckpt_type,
|
|
checkpoint_path=checkpoint_path,
|
|
dtype_choice="bf16",
|
|
load_ema_to_reg=False,
|
|
skip_vae_encoder=plan.skip_vae_encoder,
|
|
pid_weight_precision=plan.pid_weight_precision,
|
|
)
|
|
|
|
device = "cuda"
|
|
|
|
caption = caption or ""
|
|
data_batch = _make_pid_data_batch(model, caption, float(sigma), samples_cpu, device)
|
|
del samples_cpu
|
|
|
|
if model_management is not None:
|
|
try:
|
|
model_management.throw_exception_if_processing_interrupted()
|
|
except Exception:
|
|
pass
|
|
|
|
_free_cuda_memory(aggressive=bool(aggressive_cleanup))
|
|
offloader = configure_pid_runtime(model, plan, device=device)
|
|
_reset_cuda_peak_memory_stats()
|
|
try:
|
|
_warn_if_non_distilled_step_count(pid_steps)
|
|
pbar = _make_pid_progress_bar(pid_steps)
|
|
|
|
def update_progress(current: int, total: int) -> None:
|
|
_update_pid_progress_bar(pbar, current, total)
|
|
|
|
with torch.inference_mode():
|
|
out = _generate_samples_low_vram(
|
|
model,
|
|
data_batch,
|
|
cfg_scale=float(cfg_scale),
|
|
num_steps=int(pid_steps),
|
|
seed=int(seed),
|
|
shift=None,
|
|
image_size=infer_image_size,
|
|
progress_callback=update_progress if pbar is not None else None,
|
|
)
|
|
_log_cuda_peak_memory("direct decode")
|
|
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.
|
|
if offloader is not None:
|
|
offloader.cleanup()
|
|
offloader = None
|
|
del data_batch
|
|
_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
|
|
|
|
if offloader is not None:
|
|
offloader.cleanup()
|
|
offloader = None
|
|
|
|
_unload_pid_model(model, aggressive=bool(aggressive_cleanup))
|
|
del model
|
|
|
|
return (image,)
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"PiDDecode": PiDDecode,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"PiDDecode": "PiD Decode",
|
|
}
|