Support the PiD v1.5 flux2 checkpoint (MOSS-Image / FLUX.2 latents), 0.3.0

The flux2 student takes the packed, BatchNorm-normalized (B,128,H/16,W/16)
latent as is and unpatchifies it 2x2 inside the LQ projection (upstream
PID_SR4X_V1PT5_FOR_FLUX2). Same prompt prefix and 4-step schedule as
qwenimage, so the bundled null caption carries over.

- pid_net: latent_unpatchify_factor
- pid_core: flux2 kwargs, family detected from the first LQ conv's input
  width (16 / 32), output scale read from the net (32 / 64 px per cell),
  128-ch latents passed through unnormalized, family mismatch raises
- loader: "(auto-download) flux2 / MOSS-Image"; the bare "(auto-download)"
  stays qwenimage. Downloads go to the first registered pid folder
  (extra_model_paths.yaml first) and skip a checkpoint already present in
  any of them
- decode: tile_latent / tile_overlap keep their 32-px unit across families;
  use_calib is ignored for flux2

Checked on two MOSS cached latents: clean load (0 missing / 0 unexpected),
area-downsampled output 31.6 / 34.7 dB PSNR against the FLUX.2 VAE decode,
8.2 GB peak with 2048 px tiles.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

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

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(),)

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)

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:

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,

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.