diff --git a/README.md b/README.md index 407005a..af0d60a 100644 --- a/README.md +++ b/README.md @@ -1,7 +1,7 @@ # ComfyUI-Anima-PiD NVIDIA **PiD** (Pixel Diffusion Decoder) as a drop-in replacement for **VAE Decode** -on Anima / Qwen-Image latents. It takes a `LATENT` and emits a **4× super-resolved +on Anima / Qwen-Image latents and on MOSS-Image / FLUX.2 latents. It takes a `LATENT` and emits a **4× super-resolved `IMAGE`** in a single 4-step pass — decode *and* upscale fused into one node. The gemma text encoder is **never loaded** — no ~5 GB download and no prompt @@ -32,7 +32,21 @@ PiD replaces VAE Decode. Drop `AnimaPiDDecode` where `VAEDecode` was: There is **no second KSampler** and your Anima model does **not** connect to PiD — PiD runs its own internal 4-step pixel diffusion. Output size = `latent_grid × 8 × 4` -(e.g. a 64×64 latent → 2048×2048; a 128×128 latent → 4096×4096). +for Qwen latents (a 128×128 latent → 4096×4096) and `latent_grid × 16 × 4` for FLUX.2 / +MOSS latents (a 64×64 latent → 4096×4096) — 4× the native VAE decode either way. + +## MOSS-Image / FLUX.2 + +PiD v1.5 ships a FLUX.2 checkpoint +(`PiD_v1pt5_res2kto4k_sr4x_official_flux2_distill_4step`). Pick +**`(auto-download) flux2 / MOSS-Image`** in the loader and wire the MOSS KSampler's +`LATENT` into the decode node. It takes the packed, BatchNorm-normalized +`(B, 128, H/16, W/16)` latent — the MOSS DiT latent and ComfyUI's Flux2 `LATENT` — +as is, and unpatchifies it 2×2 to `(B, 32, H/8, W/8)` inside its LQ projection +(upstream `PID_SR4X_V1PT5_FOR_FLUX2`). It shares the qwenimage checkpoint's prompt +prefix and 4-step schedule, so the same bundled null caption is used. The loader +reads the family from the checkpoint (input width of the first LQ conv: 16 → Qwen, +32 → FLUX.2), and the decode node refuses a latent of the other family. ## Checkpoint: PiD v1.5 @@ -70,11 +84,15 @@ Anima generation → 4096px out, which is the intended operating point. Feeding ## Install 1. Copy/clone this folder into `ComfyUI/custom_nodes/`. -2. **No manual download needed** — in **Anima PiD Loader**, pick the - `…distill_4step (auto-download)` entry and the official checkpoint is pulled - from the public `nvidia/PiD` repo into `ComfyUI/models/pid/` on first run - (one-time). To use your own checkpoint instead, drop a `.pth`/`.safetensors` - into `ComfyUI/models/pid/` and select it from the dropdown: +2. **No manual download needed** — in **Anima PiD Loader**, pick + `(auto-download)` (Qwen-Image / Anima) or `(auto-download) flux2 / MOSS-Image` + and the official checkpoint is pulled from the public `nvidia/PiD` repo into + `ComfyUI/models/pid/` on first run (one-time, ~2.8 GB each). To keep the + checkpoints on another disk, add a `pid:` entry to `extra_model_paths.yaml`; + the first registered `pid` folder takes the downloads, and a checkpoint already + present in any of them is not fetched again. To use your own checkpoint + instead, drop a `.pth`/`.safetensors` into a `pid` folder and select it from + the dropdown: ```bash hf download nvidia/PiD --local-dir /tmp/pid \ --include "checkpoints/PiD_v1pt5_res2kto4k_sr4x_official_qwenimage_distill_4step/*" @@ -87,15 +105,17 @@ Anima generation → 4096px out, which is the intended operating point. Feeding ## Nodes -- **Anima PiD Loader** — `ckpt_name` (from `models/pid/`), `dtype` → `ANIMA_PID`. +- **Anima PiD Loader** — `ckpt_name` (from the `pid` folders), `dtype` → `ANIMA_PID`. - **Anima PiD Decode (4x SR)** — `ANIMA_PID` + `LATENT` → `IMAGE`. - `steps` (default 4) — distilled student steps. - `sigma` (default 0.0) — assumed latent degradation; 0 = clean, higher = more synthesized detail. - `tile_latent` (default 64) — `0` decodes the whole image at once (4K may OOM - on ≤16 GB); `>0` tiles the latent (each tile → `tile×32` px) with feather - blending. **64 → 2048 px tiles** (~7 GB peak in bf16). - - `tile_overlap` (default 16) — latent overlap between tiles (px = `overlap×32`). + on ≤16 GB); `>0` tiles with feather blending, each tile `tile×32` output px. + **64 → 2048 px tiles** (~7 GB peak in bf16). The unit is one Qwen latent cell + (32 output px), so a FLUX.2 / MOSS latent tiles at `tile/2` of its 16-px cells + and the pixel size stays the same. + - `tile_overlap` (default 16) — overlap between tiles, same unit (px = `overlap×32`). - `compile` (default off) — `torch.compile` the PiD net. ~1.8× faster warm (e.g. 3.8s → 2.1s for a 2048px tile), after a one-time ~37s compilation **per output resolution**. With tiling on, every tile is the same size so it @@ -106,12 +126,15 @@ Anima generation → 4096px out, which is the intended operating point. Feeding **pre-v1.5** `qwenimage` checkpoint, which decoded **flat and desaturated** vs the native Qwen VAE. **PiD v1.5 fixes colour accuracy upstream**, so the transform is off by default — stacking it on v1.5 would over-correct. Turn it - on only if you are running a hand-placed v1 checkpoint. The transform is a + on only if you are running a hand-placed v1 checkpoint. Ignored for the FLUX.2 + checkpoint. The transform is a static linear `out = (rgb @ M.T + b)` fit against native-VAE decodes (contrast ≈ ×1.11, saturation ≈ ×1.12, no hue tint). See **Provenance**. ## Latent convention +FLUX.2 / MOSS latents go in unchanged (they are already BatchNorm-normalized). + ComfyUI stores raw Qwen/Wan VAE latents; PiD wants the per-channel **normalized** latent. The node applies `(latent − mean)/std` (the standard Qwen `latents_mean/std`, `scale_factor=1.0`) internally — the same convention diff --git a/node.zip b/node.zip index e7b679c..3fa649b 100644 Binary files a/node.zip and b/node.zip differ diff --git a/nodes.py b/nodes.py index a7613a1..1a06e35 100644 --- a/nodes.py +++ b/nodes.py @@ -1,11 +1,13 @@ -"""ComfyUI nodes: NVIDIA PiD pixel-diffusion decoder for Anima / Qwen-Image latents. +"""ComfyUI nodes: NVIDIA PiD pixel-diffusion decoder for Anima / Qwen-Image and +MOSS-Image / FLUX.2 latents. Two nodes: - * AnimaPiDLoader — load a PiD qwenimage checkpoint -> ANIMA_PID socket + * AnimaPiDLoader — load a PiD qwenimage or flux2 checkpoint -> ANIMA_PID socket * AnimaPiDDecode — LATENT (+ PiD model) -> IMAGE, 4x super-resolved -PiD REPLACES VAE Decode: it consumes the (normalized) Qwen latent and emits RGB -pixels directly, upscaling 4x in the same pass (latent_grid*8 -> *4). The gemma +PiD REPLACES VAE Decode: it consumes the (normalized) latent and emits RGB pixels +directly, upscaling 4x in the same pass (latent_grid*8 -> *4 for Qwen, +latent_grid*16 -> *4 for FLUX.2). The gemma text encoder is not loaded at decode time — the net conditions on a fixed, pre-baked null caption gemma(chi_prompt+"") bundled with the node (~1.4 MB), so no ~5GB download and no prompt input. Drop AnimaPiDDecode where VAEDecode was: @@ -32,15 +34,16 @@ from comfy.utils import ProgressBar from .pid_core import ( COLOR_CALIB_FILENAME, - NET_KWARGS_V1PT5, + NET_KWARGS_V1, NULL_CAPTION_FILENAME, SR_SCALE, - VAE_DOWN, apply_color_calib, build_pid_net, categorize_load_keys, comfy_latent_to_lq, count_tiles, + latent_family, + latent_to_pixel_factor, load_color_calib, load_null_caption_embs, load_pid_weights, @@ -52,10 +55,19 @@ from .pid_core import ( from .pid_net.attention_backend import set_attention_backend # Register a ComfyUI models/pid folder for PiD checkpoints (.pth / .safetensors). +# A `pid:` entry in extra_model_paths.yaml is registered before this runs, so it +# comes first in folder_paths and takes the auto-downloads (see _download_dir). _PID_DIR = os.path.join(folder_paths.models_dir, "pid") os.makedirs(_PID_DIR, exist_ok=True) folder_paths.add_model_folder_path("pid", _PID_DIR) + +def _download_dir() -> str: + """First registered `pid` folder: an extra_model_paths.yaml entry if there is + one, else models/pid.""" + paths = folder_paths.get_folder_paths("pid") + return paths[0] if paths else _PID_DIR + # Faithful null caption gemma(chi_prompt + "") — bundled with the node (~1.4 MB, # derived data), so gemma is never needed. Regen recipe in README "Provenance". _NULL_CAPTION_PATH = os.path.join(os.path.dirname(__file__), NULL_CAPTION_FILENAME) @@ -71,43 +83,57 @@ _COLOR_CALIB_PATH = os.path.join(os.path.dirname(__file__), COLOR_CALIB_FILENAME _DTYPES = {"bf16": torch.bfloat16, "fp16": torch.float16, "fp32": torch.float32} -# Official PiD v1.5 4-step qwenimage checkpoint on the (ungated, public) nvidia/PiD -# repo. Auto-fetched into models/pid/ on first use, flattened + renamed to a flat +# Official PiD v1.5 4-step checkpoints on the (ungated, public) nvidia/PiD repo. +# Auto-fetched into models/pid/ on first use, flattened + renamed to a flat # filename so the dropdown contract stays uniform. # # NB the pre-v1.5 checkpoint this node used to fetch now lives under -# checkpoints_deprecated/ upstream — the old path 404s, hence the pin below. +# checkpoints_deprecated/ upstream — the old path 404s, hence the pins below. _HF_PID_REPO = "nvidia/PiD" -_HF_PID_FILE = "checkpoints/PiD_v1pt5_res2kto4k_sr4x_official_qwenimage_distill_4step/model_ema_bf16.pth" -_OFFICIAL_CKPT = "PiD_v1pt5_res2kto4k_sr4x_official_qwenimage_distill_4step.pth" -# Stable dropdown sentinel for the official auto-download. It is ALWAYS present in -# the list (whether or not the file has been fetched), so a saved workflow that -# selected it stays valid across restarts. Kept short and filename-free — load() -# matches on this exact string, and resolves it to the real official checkpoint. -_AUTODL_ENTRY = "(auto-download)" +_OFFICIAL = { + "qwenimage": "PiD_v1pt5_res2kto4k_sr4x_official_qwenimage_distill_4step", + "flux2": "PiD_v1pt5_res2kto4k_sr4x_official_flux2_distill_4step", +} +# Stable dropdown sentinels for the official auto-downloads. They are ALWAYS present +# in the list (whether or not the file has been fetched), so a saved workflow that +# selected one stays valid across restarts. The bare "(auto-download)" is the +# qwenimage checkpoint (what workflows saved before flux2 support selected). +_AUTODL_ENTRIES = { + "(auto-download)": "qwenimage", + "(auto-download) flux2 / MOSS-Image": "flux2", +} # Only these are real checkpoints — everything else under models/pid/ (e.g. the # HF cache's .gitignore / *.metadata) is filtered out of the dropdown. _CKPT_EXTS = (".pth", ".safetensors", ".pt", ".ckpt", ".bin") -def _download_official_ckpt() -> str: - """Fetch the official PiD qwenimage checkpoint into models/pid/ (one-time). +def _official_ckpt_name(family: str) -> str: + return _OFFICIAL[family] + ".pth" + + +def _download_official_ckpt(family: str) -> str: + """Fetch an official PiD v1.5 checkpoint into models/pid/ (one-time). Downloads into the shared HF hub cache (NOT ``local_dir=models/pid``, which would litter ``models/pid/.cache/huggingface/`` with ``.gitignore`` / ``*.metadata`` files that then show up in the loader dropdown), then copies the blob out to a flat, descriptive filename in ``models/pid/`` so it lists like any - hand-placed checkpoint. Returns the local path.""" - dest = os.path.join(_PID_DIR, _OFFICIAL_CKPT) - if os.path.exists(dest): - return dest + hand-placed checkpoint. Any registered `pid` folder already holding it counts. + Returns the local path.""" + name = _official_ckpt_name(family) + found = folder_paths.get_full_path("pid", name) + if found is not None: + return found + dest = os.path.join(_download_dir(), name) + os.makedirs(os.path.dirname(dest), exist_ok=True) from huggingface_hub import hf_hub_download + hf_file = f"checkpoints/{_OFFICIAL[family]}/model_ema_bf16.pth" print( - f"[AnimaPiD] fetching {_HF_PID_REPO}/{_HF_PID_FILE} -> {dest} (one-time, ~public download).\n" + f"[AnimaPiD] fetching {_HF_PID_REPO}/{hf_file} -> {dest} (one-time, ~2.8 GB public download).\n" f"[AnimaPiD] NOTE: PiD weights are NVIDIA NSCLv1 — non-commercial (research/evaluation) use only." ) - downloaded = hf_hub_download(repo_id=_HF_PID_REPO, filename=_HF_PID_FILE) + downloaded = hf_hub_download(repo_id=_HF_PID_REPO, filename=hf_file) shutil.copyfile(downloaded, dest) return dest @@ -119,11 +145,13 @@ class AnimaPiDModel: VRAM residency exactly like a UNET/VAE: it rests on the offload (CPU) device and is pulled onto the GPU — evicting other models as needed — via load_models_gpu at decode time, then evicted on-demand when something else needs the VRAM. `.net` is - the underlying module (what the pid_core decode fns consume).""" + the underlying module (what the pid_core decode fns consume). `family` is the + latent family the checkpoint decodes ("qwenimage" or "flux2").""" - def __init__(self, patcher, dtype): + def __init__(self, patcher, dtype, family="qwenimage"): self.patcher = patcher self.dtype = dtype + self.family = family @property def net(self): @@ -139,11 +167,10 @@ class AnimaPiDLoader: files = [f for f in folder_paths.get_filename_list("pid") if f.lower().endswith(_CKPT_EXTS) and os.path.basename(f) not in (NULL_CAPTION_FILENAME, COLOR_CALIB_FILENAME)] - # The auto-download sentinel is ALWAYS the first entry, present whether or - # not the official checkpoint has been fetched. This keeps a saved workflow - # that selected it valid across restarts — the old behaviour dropped the - # sentinel once the real file existed, invalidating saved graphs. - files.insert(0, _AUTODL_ENTRY) + # The auto-download sentinels are ALWAYS the first entries, present whether + # or not the official checkpoint has been fetched. This keeps a saved + # workflow that selected one valid across restarts. + files = list(_AUTODL_ENTRIES) + files return { "required": { "ckpt_name": (files,), @@ -157,19 +184,22 @@ class AnimaPiDLoader: CATEGORY = "Anima/PiD" def load(self, ckpt_name, dtype): - # The auto-download sentinel (and the bare official filename) resolve to - # the official checkpoint, fetching it on first use. Tolerate any - # "(auto-download)"-suffixed value so older saved workflows keep working. - if ckpt_name == _AUTODL_ENTRY or ckpt_name == _OFFICIAL_CKPT or "(auto-download)" in ckpt_name: - path = _download_official_ckpt() - ckpt_name = _OFFICIAL_CKPT + # An auto-download sentinel resolves to its official checkpoint, fetching it + # on first use. Any other "(auto-download)"-suffixed value (older saved + # workflows embedded the filename) maps to qwenimage. + family = _AUTODL_ENTRIES.get(ckpt_name) + if family is None and "(auto-download)" in ckpt_name: + family = "flux2" if "flux2" in ckpt_name else "qwenimage" + if family is not None: + path = _download_official_ckpt(family) + ckpt_name = _official_ckpt_name(family) else: path = folder_paths.get_full_path("pid", ckpt_name) if path is None: raise FileNotFoundError( - f"PiD checkpoint {ckpt_name!r} not found under {_PID_DIR}. " - f"Download nvidia/PiD {_HF_PID_FILE} and place it there, " - f"or select the auto-download entry." + f"PiD checkpoint {ckpt_name!r} not found under {folder_paths.get_folder_paths('pid')}. " + f"Download it from nvidia/PiD and place it there, " + f"or select an auto-download entry." ) dt = _DTYPES[dtype] # Build the net on the OFFLOAD (CPU) device, not the GPU — the same load/ @@ -186,7 +216,8 @@ class AnimaPiDLoader: # a hand-placed v1 file. Read the state dict first, then build to match it. sd = read_pid_state_dict(path) net_kwargs = net_kwargs_for_state_dict(sd) - is_v1pt5 = net_kwargs is NET_KWARGS_V1PT5 + arch = "v1" if net_kwargs is NET_KWARGS_V1 else "v1.5" + family = latent_family(net_kwargs) net = build_pid_net(offload_device, dt, net_kwargs) missing, unexpected = load_pid_weights(net, sd) expected_missing, suspect_missing, unexpected = categorize_load_keys(missing, unexpected) @@ -210,9 +241,9 @@ class AnimaPiDLoader: patcher = comfy.model_patcher.ModelPatcher(net, load_device=load_device, offload_device=offload_device) print( f"[AnimaPiD] loaded {ckpt_name} as {dtype} " - f"(arch={'v1.5' if is_v1pt5 else 'v1'}, offload={offload_device}, load={load_device}; ComfyUI-managed)" + f"(arch={arch}, latent={family}, offload={offload_device}, load={load_device}; ComfyUI-managed)" ) - return (AnimaPiDModel(patcher, dt),) + return (AnimaPiDModel(patcher, dt, family),) class AnimaPiDDecode: @@ -229,10 +260,11 @@ class AnimaPiDDecode: "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}), "tile_latent": ("INT", {"default": 64, "min": 0, "max": 256, "step": 8, "tooltip": "0 = decode whole image at once (4K output may OOM on <=16GB). " - ">0 = tile the latent (each tile -> tile*32 px) with feather " - "blending. 64 -> 2048px tiles."}), + ">0 = tile with feather blending; each tile is tile*32 output " + "px for either latent family (64 -> 2048px tiles; a FLUX.2 / " + "MOSS latent tiles at half this many cells)."}), "tile_overlap": ("INT", {"default": 16, "min": 0, "max": 64, "step": 4, - "tooltip": "Latent-space overlap between tiles (pixels = overlap*32). " + "tooltip": "Overlap between tiles, in the same units (pixels = overlap*32). " "Larger = fewer seams, slower."}), "compile": ("BOOLEAN", {"default": False, "tooltip": "Per-block torch.compile of the PiD net: each transformer block " @@ -320,17 +352,24 @@ class AnimaPiDDecode: cap = self._null_caption(load_device, dt) - lq = comfy_latent_to_lq(latent["samples"], load_device, dt) # (B,16,h,w) normalized + family = getattr(pid_model, "family", "qwenimage") + lq = comfy_latent_to_lq(latent["samples"], load_device, dt, + latent_channels=net.lq_proj.latent_channels) lh, lw = lq.shape[-2], lq.shape[-1] - out_h, out_w = lh * VAE_DOWN * SR_SCALE, lw * VAE_DOWN * SR_SCALE - print(f"[AnimaPiD] decode latent {lh}x{lw} -> {out_h}x{out_w} ({SR_SCALE}x), " - f"steps={steps} sigma={sigma} tile={tile_latent or 'off'} compile={compile}") + up = latent_to_pixel_factor(net) + out_h, out_w = lh * up, lw * up + # tile_latent / tile_overlap count 32-px units (one Qwen latent cell), so a + # setting keeps its pixel size across latent families. + tile = tile_latent * 32 // up + overlap = tile_overlap * 32 // up + print(f"[AnimaPiD] decode {family} latent {lh}x{lw} -> {out_h}x{out_w} ({SR_SCALE}x), " + f"steps={steps} sigma={sigma} tile={f'{tile} cells' if tile else 'off'} compile={compile}") - use_tiling = bool(tile_latent) and (lh > tile_latent or lw > tile_latent) + use_tiling = bool(tile) and (lh > tile or lw > tile) # Drive ComfyUI's node progress bar: one tick per SDE step, summed over # tiles. (PiD runs its own sampler loop, so without this the node shows # no progress.) - n_tiles = count_tiles(lq, tile_latent, tile_overlap) if use_tiling else 1 + n_tiles = count_tiles(lq, tile, overlap) if use_tiling else 1 pbar = ProgressBar(steps * n_tiles) step_cb = lambda: pbar.update(1) # noqa: E731 if use_tiling: @@ -340,7 +379,7 @@ class AnimaPiDDecode: # fit through this same path (see metadata global_lq/tile_latent). px = pid_decode_latent_tiled( net, lq, steps=steps, sigma=sigma, seed=seed, - tile=tile_latent, overlap=tile_overlap, dtype=dt, compile=compile, + tile=tile, overlap=overlap, dtype=dt, compile=compile, caption_embs=cap, step_cb=step_cb, ) else: @@ -349,7 +388,9 @@ class AnimaPiDDecode: # (B,3,H,W) in [-1,1] -> ComfyUI IMAGE (B,H,W,3) in [0,1] img = ((px.float() + 1.0) / 2.0).clamp(0, 1).permute(0, 2, 3, 1).contiguous() - if use_calib: + if use_calib and family != "qwenimage": + print("[AnimaPiD] use_calib ignored: the bundled calib was fit on the Qwen-Image checkpoint.") + elif use_calib: M, b = self._color_calib(load_device) img = apply_color_calib(img, M, b) # PiD->native-VAE color match (channel-last) return (img.cpu(),) diff --git a/pid_core.py b/pid_core.py index 9790bb6..bcda383 100644 --- a/pid_core.py +++ b/pid_core.py @@ -16,10 +16,16 @@ nv-tlabs PiD `qwenimage` 2kto4k checkpoint config (captured by introspection): (NaN-safe — the y_embedder is `Linear(2304→D, bias=True)` then RMSNorm — but off-distribution; the node always passes the faithful null). -The PiD net consumes a *normalized* Qwen latent (LQ_latent) and emits RGB pixels -directly — there is NO VAE decode at the end, and the Qwen VAE is not needed at -all (we feed the latent straight in). Output spatial size = latent_grid * 8 * -sr_scale(=4). +The PiD net consumes a *normalized* latent (LQ_latent) and emits RGB pixels +directly — there is NO VAE decode at the end, and the VAE is not needed at all (we +feed the latent straight in). Output spatial size = latent_grid * vae_down * 4. + +Two latent families, picked from the checkpoint: + * qwenimage — 16-ch Qwen-Image VAE latent at /8 (Anima, Qwen-Image), per-channel + mean/std normalized here. + * flux2 — FLUX.2 VAE latent packed 2x2 to 128-ch at /16 and BatchNorm-normalized + (MOSS-Image's DiT latent, ComfyUI's Flux2 LATENT). Fed as is; the net + unpatchifies it to 32-ch at /8 inside its LQ projection. """ from __future__ import annotations @@ -69,22 +75,52 @@ NET_KWARGS_V1PT5 = dict( pit_lq_inject=True, ) +# PiD v1.5 for FLUX.2 latents (upstream PID_SR4X_V1PT5_FOR_FLUX2): takes the packed, +# BN-normalized (B,128,H/16,W/16) latent and unpatchifies it to (B,32,H/8,W/8) +# before the LQ trunk. Same text prompt / schedule as qwenimage. +NET_KWARGS_V1PT5_FLUX2 = dict( + NET_KWARGS_V1PT5, + lq_latent_channels=128, + latent_spatial_down_factor=16, + lq_latent_unpatchify_factor=2, +) + # Default for a freshly-fetched official checkpoint. NET_KWARGS = NET_KWARGS_V1PT5 +_LATENT_PROJ_IN_KEY = "lq_proj.latent_proj.0.weight" + def net_kwargs_for_state_dict(sd: dict) -> dict: """Pick the PidNet kwargs matching a checkpoint's actual architecture. v1.5 is identified by its dedicated PiT injection head (`lq_proj.pit_head`), which - v1 does not have. This keeps a hand-placed v1 checkpoint (the deprecated Hub copy) - loadable alongside the current official one, instead of failing with a shape error. - """ + v1 does not have. Within v1.5, the first LQ conv's input width names the latent + family: 16 for qwenimage / flux, 32 for flux2 (128-ch packed latent unpatchified + 2x2). The distilled student may omit lq_proj keys; without the conv weight it + falls back to the qwenimage layout.""" keys = {k[len("net."):] if k.startswith("net.") else k for k in sd} - return NET_KWARGS_V1PT5 if "lq_proj.pit_head.weight" in keys else NET_KWARGS_V1 + if "lq_proj.pit_head.weight" not in keys: + return NET_KWARGS_V1 + w = sd.get(_LATENT_PROJ_IN_KEY, sd.get("net." + _LATENT_PROJ_IN_KEY)) + in_ch = int(w.shape[1]) if w is not None else 16 + if in_ch == 16: + return NET_KWARGS_V1PT5 + if in_ch == 32: + return NET_KWARGS_V1PT5_FLUX2 + raise ValueError(f"unsupported PiD v1.5 latent projection with {in_ch} input channels") + + +def latent_family(net_kwargs: dict) -> str: + return "flux2" if net_kwargs.get("lq_latent_channels") == 128 else "qwenimage" + + +def latent_to_pixel_factor(net: PidNet) -> int: + """Output pixels per source-latent cell: vae_down * sr_scale (32 for qwenimage, + 64 for flux2).""" + return net.lq_proj.latent_spatial_down_factor * net.sr_scale SR_SCALE = 4 -VAE_DOWN = 8 # latent grid -> vae-native pixels MODEL_MAX_LENGTH = 300 CAPTION_CHANNELS = 2304 @@ -211,12 +247,26 @@ def apply_color_calib(img01: torch.Tensor, M: torch.Tensor, b: torch.Tensor) -> return out.clamp(0, 1).to(img01.dtype) -def comfy_latent_to_lq(samples: torch.Tensor, device, dtype=torch.bfloat16) -> torch.Tensor: - """ComfyUI LATENT['samples'] (raw Qwen VAE latent, 4D or 5D) -> PiD LQ_latent - (per-channel normalized (mu-mean)/std), 4D (B,16,h,w).""" +def comfy_latent_to_lq(samples: torch.Tensor, device, dtype=torch.bfloat16, + latent_channels: int = 16) -> torch.Tensor: + """ComfyUI LATENT['samples'] -> PiD LQ_latent, 4D. + + 16-ch (raw Qwen VAE latent, 4D or 5D): per-channel normalized (mu-mean)/std. + 128-ch (FLUX.2 packed latent, already BN-normalized — MOSS-Image's DiT latent and + ComfyUI's Flux2 LATENT): passed through unchanged. + `latent_channels` is what the loaded net expects; a mismatch raises.""" x = samples if x.ndim == 5: # (B,C,T,h,w) -> drop singleton frame x = x[:, :, 0] + if x.shape[1] != latent_channels: + family = {16: "Qwen-Image / Anima", 128: "FLUX.2 / MOSS-Image"} + raise ValueError( + f"latent has {x.shape[1]} channels ({family.get(x.shape[1], 'unknown')}) but the " + f"loaded PiD checkpoint expects {latent_channels} " + f"({family.get(latent_channels, 'unknown')}). Load the matching PiD checkpoint." + ) + if latent_channels == 128: + return x.to(device=device, dtype=dtype) mean = torch.tensor(QWEN_LATENTS_MEAN, device=x.device, dtype=torch.float32).view(1, 16, 1, 1) std = torch.tensor(QWEN_LATENTS_STD, device=x.device, dtype=torch.float32).view(1, 16, 1, 1) lq = (x.float() - mean) / std @@ -344,8 +394,9 @@ def pid_decode_latent(net: PidNet, lq_latent: torch.Tensor, *, steps: int = 4, sigma: float = 0.0, seed: int = 0, dtype=torch.bfloat16, compile: bool = False, caption_embs: torch.Tensor = None, _runner=None, step_cb=None, lq_features=None) -> torch.Tensor: - """Run the 4-step SDE student. lq_latent: (B,16,h,w) normalized. - Returns pixels (B,3,H,W) in [-1,1] with H=h*8*4, W=w*8*4. + """Run the 4-step SDE student. lq_latent: normalized (B,16,h,w) or (B,128,h,w). + Returns pixels (B,3,H,W) in [-1,1] with H=h*up, W=w*up, up=latent_to_pixel_factor + (32 qwenimage, 64 flux2). `caption_embs` is the fixed null caption y the net conditions on. Pass None for the zeros null (off-distribution but NaN-safe); pass a (1 or B, 300, 2304) tensor @@ -360,7 +411,8 @@ def pid_decode_latent(net: PidNet, lq_latent: torch.Tensor, *, steps: int = 4, features (see `pid_decode_latent_tiled`).""" device = lq_latent.device B, _, lh, lw = lq_latent.shape - H, W = lh * VAE_DOWN * SR_SCALE, lw * VAE_DOWN * SR_SCALE + up = latent_to_pixel_factor(net) + H, W = lh * up, lw * up run = _runner if _runner is not None else get_runner(net, dtype, compile) if caption_embs is None: @@ -436,7 +488,7 @@ def pid_decode_latent_tiled(net: PidNet, lq_latent: torch.Tensor, *, steps: int step_cb=None, global_lq: bool = True) -> torch.Tensor: """Tiled SR decode for latents larger than `tile` (memory bound). Decodes overlapping latent tiles and feather-blends them in pixel space. Output - (B,3, h*32, w*32) in [-1,1]. `step_cb` ticks once per SDE step across all + (B,3, h*up, w*up) in [-1,1]; `tile` / `overlap` are in this latent's cells. `step_cb` ticks once per SDE step across all tiles (total ticks = steps * count_tiles()). `global_lq` (default True) makes the LQ conditioning tile-invariant: the LQ @@ -450,7 +502,7 @@ def pid_decode_latent_tiled(net: PidNet, lq_latent: torch.Tensor, *, steps: int per-tile behavior (matches the pre-fix shipped color calibration).""" device = lq_latent.device B, _, Hh, Ww = lq_latent.shape - up = VAE_DOWN * SR_SCALE # 32 + up = latent_to_pixel_factor(net) # 32 qwenimage, 64 flux2 Hout, Wout = Hh * up, Ww * up stride = max(1, tile - overlap) ys, th = _tile_positions(Hh, tile, stride) diff --git a/pid_net/lq_projection_2d.py b/pid_net/lq_projection_2d.py index 66851b0..7cbc03c 100644 --- a/pid_net/lq_projection_2d.py +++ b/pid_net/lq_projection_2d.py @@ -195,6 +195,7 @@ class LQProjection2D(nn.Module): patch_size: int = 16, sr_scale: int = 4, latent_spatial_down_factor: int = 8, + latent_unpatchify_factor: int = 1, num_res_blocks: int = 4, num_outputs: int = 1, gate_type: str = _SUPPORTED_GATE_TYPE, @@ -214,6 +215,12 @@ class LQProjection2D(nn.Module): self.patch_size = patch_size self.sr_scale = sr_scale self.latent_spatial_down_factor = latent_spatial_down_factor + if latent_spatial_down_factor % latent_unpatchify_factor != 0: + raise ValueError( + f"latent_spatial_down_factor ({latent_spatial_down_factor}) must be divisible by " + f"latent_unpatchify_factor ({latent_unpatchify_factor})" + ) + self.latent_unpatchify_factor = latent_unpatchify_factor self.num_outputs = num_outputs self.interval = interval self.zero_init = zero_init @@ -247,7 +254,12 @@ class LQProjection2D(nn.Module): # --- Latent branch --- # Spatial alignment (fold / upsample) → Conv proj → ResBlocks if latent_channels > 0: - z_to_patch_ratio = (sr_scale * latent_spatial_down_factor) / patch_size + uf = latent_unpatchify_factor + if latent_channels % (uf * uf) != 0: + raise ValueError(f"latent_channels ({latent_channels}) not divisible by unpatchify factor^2 ({uf * uf})") + # The trunk sees the unpatchified latent: C/uf² channels at lsdf/uf. + latent_channels = latent_channels // (uf * uf) + z_to_patch_ratio = (sr_scale * (latent_spatial_down_factor // uf)) / patch_size self.z_to_patch_ratio = z_to_patch_ratio if z_to_patch_ratio > 1: @@ -399,6 +411,13 @@ class LQProjection2D(nn.Module): Returns [B, hidden_dim, pH, pW]. """ + f = self.latent_unpatchify_factor + if f > 1: + # Flux2 normalized latent: [B, C*f*f, H, W] -> [B, C, H*f, W*f], channel + # order (c, pi, pj) as in diffusers' _patchify_latents. No BN inverse. + B, C, H, W = lq_latent.shape + lq_latent = lq_latent.reshape(B, C // (f * f), f, f, H, W) + lq_latent = lq_latent.permute(0, 1, 4, 2, 5, 3).reshape(B, C // (f * f), H * f, W * f) B, z_dim = lq_latent.shape[:2] if self.z_to_patch_ratio > 1: diff --git a/pid_net/pid_net.py b/pid_net/pid_net.py index 2e93a1c..4da31a9 100644 --- a/pid_net/pid_net.py +++ b/pid_net/pid_net.py @@ -88,6 +88,9 @@ class PidNet(PixDiT_T2I): train_lq_proj_only: bool = False, sr_scale: int = 4, latent_spatial_down_factor: int = 8, + # PiD v1.5 Flux2: 2 unpatchifies the packed (B,128,H/16,W/16) latent to + # (B,32,H/8,W/8) before the LQ trunk. 1 = off. + lq_latent_unpatchify_factor: int = 1, # --- PiT LQ injection args --- # Inject LQ features into PiT pixel blocks via a dedicated output head # from the same LQ projection CNN backbone. Added to s_cond before PiT loop. @@ -141,6 +144,7 @@ class PidNet(PixDiT_T2I): patch_size=patch_size, sr_scale=sr_scale, latent_spatial_down_factor=latent_spatial_down_factor, + latent_unpatchify_factor=lq_latent_unpatchify_factor, num_res_blocks=lq_num_res_blocks, num_outputs=num_lq_outputs, gate_type=lq_gate_type, diff --git a/pyproject.toml b/pyproject.toml index 7798919..1ed0442 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,7 +1,7 @@ [project] name = "comfyui-anima-pid" -description = "NVIDIA PiD (Pixel Diffusion Decoder) as a drop-in VAE-Decode replacement for Anima / Qwen-Image latents: LATENT -> 4x super-resolved IMAGE in one 4-step pass, with the gemma text encoder skipped (no multi-GB download). Self-contained vendored PiD net (no hydra/imaginaire)." -version = "0.2.0" +description = "NVIDIA PiD (Pixel Diffusion Decoder) as a drop-in VAE-Decode replacement for Anima / Qwen-Image and MOSS-Image / FLUX.2 latents: LATENT -> 4x super-resolved IMAGE in one 4-step pass, with the gemma text encoder skipped (no multi-GB download). Self-contained vendored PiD net (no hydra/imaginaire)." +version = "0.3.0" license = { file = "LICENSE" } # torch / numpy / pillow ship with any ComfyUI install; the vendored PiD net # needs nothing beyond torch. Listed for a bare `pip install .` outside ComfyUI.