Commit Graph
1 Commits
Author SHA1 Message Date
Reithan d4ff852491 PR-2: pack-aware per-stream NRS routing + degeneracy tripwire (#36)
## What
Fixes NRS on MiniMax H3 (and any sampler that packs multiple streams).
ComfyUI hands the cfg hook a **flat packed latent** `[B,1,N]` whose
`dim=1` is a singleton, so every NRS reduction runs over a size-1 axis
and the projection geometry collapses — Skew becomes a silent no-op,
Stretch degenerates to elementwise CFG, Squash to a Stretch-reverter
(proven to float precision in Phase 0).

This PR unpacks the flat pack into its real channels-first per-stream
tensors, runs the **unchanged** NRS geometry per stream, then repacks.

## Changes (`NRS/nodes_NRS.py`)
- Guarded `import comfy.utils as _comfy_utils` (module still imports
with no ComfyUI present) + `import math`.
- Module-level pure-Python `_unpack_latents`/`_pack_latents` fallbacks
matching ComfyUI's pack contract (each stream reshaped to `(B,1,-1)` and
concatenated on the last dim; unpack slices `prod(shape[1:])` per
stream).
- Extracted `_apply_guidance(...)` — the per-stream geometry pipeline
(sigma reshape → `_convert_to_v_space` → dot/proj/stretch/skew/squash →
`_finalize_from_v_space`), **byte-for-byte identical** to the prior
inline math.
- Rewired `nrs()`: read `args["model"].latent_shapes`; when present with
`len>1`, unpack cond/uncond/input per stream (comfy-preferred,
pure-Python fallback), map `_apply_guidance`, repack; otherwise the
single-stream path is the **exact prior behavior** (regression no-op for
SDXL etc.).
- Permanent **degeneracy tripwire**: warns once per `patch()` if any
routed stream still has a singleton reduction axis — the silent-failure
class that cost the original Skew investigation days.

## Why this is safe
- Single-stream models are a verified byte-for-byte no-op (two
regression tests compare against a manually-computed `_apply_guidance`).
- Not H3 fingerprinting: any multi-stream packed sampler (LTXV AV
variants) benefits; older ComfyUI without the API falls through cleanly.
- Verified against the real 18MB H3 capture (not committed):
unpack→repack is bit-exact for cond/uncond/x_orig; streams recover as
`[1,24,72,38,22]` (video) / `[1,32,2,405]` (audio); per-stream guidance
produces large live deviations where the flat pack was near-degenerate —
Skew is alive again.

## Tests
New `tests/test_pack_split.py` (9 tests, real-torch via a module-scoped
isolation harness that restores the mock afterward): round-trip
pack/unpack, single- & multi-stream fallback regression, tripwire
fires/doesn't-fire, per-stream reduced shapes, and
non-degenerate-vs-degenerate rejection.

- `uvx ruff check .` → All checks passed!
- `uvx --with torch pytest -q` → **33 passed** (24 existing + 9 new)

Breaking for H3 outputs only. Baseline preserved at tag `pre-flow`.
2026-08-14 01:15:59 -07:00