docs/WP9_DESIGN.md

WP-9 Design Doc — Projection-Guided Diffusion Sampling (L4)

Author: Claude Fable 5 (frontier design pass), 2026-07-07 Status: DESIGN ONLY — implementation is gated (see §1). This document is WP-9's first deliverable ("Design doc first"). It makes the hard algorithmic decisions up front so the eventual implementation WP is execution, not research. Prereqs to read: MASTERPLAN Part I (I.5, I.5a, I.6), Part III L4. This doc does not restate the physics; it builds on it.

---

1. The gate is not yet open (honest precondition)

WP-9's dependency line reads: "Only start when L0–L3 metrics plateau on the benchmark." They have not plateaued, and this doc does not claim otherwise:

• WP-6 (DIP) just became the strongest offline source at 256px (0.6712 vs 0.7838 LPIPS) — a source that is still improving as we fix the ruler, not a plateau.
• WP-7 (refinement) is net-neutral (4/6 improved, flat mean) and its ≥70% accept item is still open pending the 50-case 256px run.
• The 50-case benchmark at 256px — the actual instrument that would show a plateau — has not been run.
• 50-case 256px verdict (2026-07-07): top sources have NOT plateaued (best-heu=0.7720, DIP=0.6031 pulls clearly ahead, refined-heu=0.8143; spread >>0.05); WP-9 remains gated, cheaper-layer work continues.

Therefore: write the design now (cheap, frontier-appropriate), but do NOT implement until the 50-case 256px benchmark is run and shows L1.2/L1.3/L3 within noise of each other across two successive improvements. Diffusion guidance is the most expensive layer in the plan (GPU, per-step VAE round-trips); spending that budget before the cheaper layers have demonstrably topped out inverts the plan's whole cost discipline. This doc's acceptance criteria (§9) are written so the implementation WP can be picked up unchanged whenever the gate opens.

---

2. Problem, stated as a linear inverse problem

From I.5: densitometry gives us the total exposure field Ĥ = H₁ + H₂ directly (per pixel, on VALID mask pixels). We want two plausible scene images A, B. The measurement operator is:

y = M · x       where   x = (H₁; H₂) stacked,   M = [I  I],   y = Ĥ

M is linear in exposure space — this is the entire reason DDNM-style null-space projection applies. For a single pixel the operator is the 1×2 matrix [1 1]:

• Range space (what the measurement fixes): the sum h₁ + h₂.
• Null space (what the prior must fill): the difference h₁ − h₂, a 1-D line per pixel.

The pseudo-inverse M⁺ = [I; I]/2 distributes a measured sum equally. DDNM's core identity, per pixel:

(h₁, h₂)_consistent  =  M⁺·ĥ  +  (I − M⁺M)·(h₁, h₂)_prior
                     =  (ĥ/2, ĥ/2)  +  ( (h₁−h₂)/2, (h₂−h₁)/2 )_prior

i.e. keep the prior's difference, replace its sum with the measured sum. Equivalently (the form we implement): compute residual r = ĥ − (h₁ + h₂) and split it, h₁ += r/2, h₂ += r/2. These are identical; the residual form is what handles non-negativity and the inequality (shoulder) case cleanly (§6).

The catch that makes this non-trivial (and why it's a T3, not a bolt-on): the diffusion prior does not live in exposure space H. It generates display-referred sRGB images. The map from a display image to its exposure contribution is φ(img) = luminance_from_linear(srgb_to_linear(img)) — pointwise, monotone, nonlinear. So the operator is linear in H but the sampled variable is img, and H = φ(img). DDNM strictly requires a linear M∘(sampling variable). We recover applicability because φ is pointwise monotone and invertible: the projection is done in H-space and mapped back through φ⁻¹. This is exactly L4's "map into exposure space (pointwise monotone maps), project, map back." §5 handles the one place this bites (the VAE).

---

3. Core algorithm: two coupled chains + per-step exposure-space projection

Two independent Stable-Diffusion sampling chains (same UNet weights, different latents, different text prompts). They are coupled only through a per-step projection — no cross-attention, no joint model, no UNet backprop.

prompts_a, prompts_b ← analyze_scan(observed_positive).prompt_a / prompt_b   (reuse WP-5 VLM)
z_a, z_b ← random latents (or img2img-encoded from a heuristic split as a warm start)

for t in reverse diffusion schedule:
    # 1. Standard DDIM/DDPM predicted-x0 from each chain (NO grad through UNet)
    x0_a_lat = predict_x0(unet, z_a, t, prompts_a)
    x0_b_lat = predict_x0(unet, z_b, t, prompts_b)

    # 2. Decode to pixel RGB, map to exposure (luminance) space
    A = vae_decode(x0_a_lat);  B = vae_decode(x0_b_lat)          # sRGB [0,1]
    h1 = φ(A);  h2 = φ(B)                                         # exposure (luminance)

    # 3. Fit nuisance gains g1,g2 (product=1) by LS on VALID (every K steps; §7)
    g1, g2 = fit_gains(h1, h2, Ĥ, mask)         # cheap, no backprop

    # 4. DDNM range-space projection with noisy-consistency scaling (§5)
    r = Ĥ − (g1·h1 + g2·h2)                      # residual on VALID
    step = λ(t) · r                              # λ(t)∈[0,1] ramps 0→1 as t→0 (§5)
    (h1, h2) ← apply_projection(h1, h2, step, mask, curve)   # equal split + non-neg + shoulder ineq (§6)

    # 5. Map back to images, re-encode, and re-noise to level t-1
    A' = render(h1 · scene chroma from A);  B' = render(h2 · chroma from B)   (§4)
    x0_a_lat ← vae_encode(A');  x0_b_lat ← vae_encode(B')
    z_a ← renoise(x0_a_lat, t→t-1);  z_b ← renoise(x0_b_lat, t→t-1)

return A, B  (final decode); register both as one candidate pair

The projection is the physics; the UNet is the prior; they meet once per step in exposure space. This is DDNM (Wang et al. 2022) adapted from "one image, global linear A" to "two images, per-pixel additive A with a pointwise nonlinear sampling map."

---

4. What the constraint applies to: luminance only, chroma from the prior

Decision: the sum constraint is enforced on the luminance/exposure channel only. Chroma (the null space of a luminance-only measurement) is left entirely to the diffusion prior. Rationale:

• The physics constraint Ĥ = H₁ + H₂ is a statement about exposure, which for a single B&W emulsion is luminance, and for the green-channel-driven color path (WP-8) is the green exposure. It says nothing about hue. Forcing chroma would be inventing physics.
• This matches the plan's "B&W first" cross-cutting rule and keeps WP-9 compatible with the color pipeline without a per-channel diffusion model (a future WP if ever).
• Concretely: render(h · chroma) scales each chain's current decoded chroma to the projected luminance (the same luminance-ratio chroma-carry trick already in _replicate_separation / demix's _render_h_to_positive — reuse it, do not reinvent).

So the projected images keep the UNet's colors and textures; only their brightness is pulled onto the physical sum. This is the single most important scoping decision — it is what makes the projection a cheap per-pixel operation instead of an intractable joint constraint over RGB.

---

5. The VAE round-trip is the real cost (and the PSLD caveat)

DDNM's range-space correction must be applied to x̂₀ in the space where the operator is defined = pixel/exposure space. SD samples in a 4-channel latent. So every projected step pays vae_decode → project → vae_encode. Two independent problems:

1. Cost. Decode+encode per step × 2 chains × ~30–50 steps is the dominant runtime. On the user's MPS hardware this is minutes/image (acceptable for a T3 offline candidate, not for interactive use — hence off-by-default, §8).
2. The VAE is not a perfect autoencoder (this is the "PSLD caveat in latent space" the plan names). Re-encoding a projected image injects reconstruction error, and doing it on an early, still-noisy x̂₀ over-commits to a bad estimate.

Mitigations (decided):

• Noisy-consistency scaling λ(t) (from DDNM+): scale the projection step by a schedule that is ~0 at high noise and ramps to 1 as t→0. Early steps let the prior explore; late steps enforce the constraint hard. Concretely λ(t) = (1 − ᾱ_t)^(−½)-normalised to [0,1], or the simpler linear λ = 1 − t/T. Start with linear; the schedule is a tuning knob, not a design fork.
• Sparse early projection: decode/project/encode only every K steps (K≈5) for the first ~70% of the schedule, then every step in the final ~30%. Cuts VAE round-trips ~3× with negligible quality loss (the early corrections are down-weighted by λ anyway).
• The final returned images are a clean decode of the last projected latents (no re-encode after the last projection), so VAE error is not doubled at the output.

---

6. Confidence mask & non-negativity (the honest projection)

The equal-split projection is modified per pixel by the WP-2 confidence mask (TOE=0, VALID=1, SHOULDER=2) — because the measurement Ĥ means different things in each region:

• VALID: full equality projection h₁+h₂ = ĥ, split residual equally.
• Non-negativity: if the equal split would push either hᵢ < 0, clamp that layer to 0 and give the entire residual to the other (a plausible split cannot have negative exposure). This is the "with non-negativity handling" clause in L4.
• SHOULDER (saturation): ĥ is only a lower bound on the true sum (density saturated at D_max). Projection becomes an inequality: only correct upward — if g₁h₁+g₂h₂ < ĥ_lowerbound, push up; if already ≥, leave the prior's values alone. Never pull a shoulder pixel down to a saturated lower bound.
• TOE (fog/noise): ĥ is poorly constrained; apply no projection (mask the residual to 0). Let the prior fill these pixels. This mirrors how the WP-3 physics loss already down-weights TOE via _valid_weight.

All four cases are a single masked, clamped residual update — the implementation is ~15 lines, but getting the shoulder inequality and non-negativity right is the correctness core and must be unit-tested independently (§9).

---

7. Nuisance gains g₁, g₂ (I.5a) inside the loop

The true constraint is g₁·φ(A) + g₂·φ(B) = Ĥ with free per-layer gains (product fixed to 1 to remove redundancy with the curve speed point) absorbing unknown relative exposure / latent-image fading / scanner calibration. Reuse WP-7's exact least-squares gain fit (the learnable-g machinery already in hybrid_loss/latent_optimizer), not a new derivation: fit log₁₀g by LS on VALID pixels against the current (h₁, h₂) estimates, clamp to GAIN_GRID_RANGE. Decision: refit every K steps (same K as the sparse projection), not every step — the gains are slowly varying and the fit is noisy on early x̂₀. Fold g₁, g₂ into the residual (r = Ĥ − (g₁h₁ + g₂h₂)) and into the map-back.

---

8. Architecture, API surface, integration

New module guided_sampling.py (top level, sibling to latent_optimizer.py / demix.py). Mirrors the injected-callable pattern that makes demix/DIP testable offline:

• GuidedSamplingConfig — steps, guidance schedule λ, projection stride K, chroma-carry flag, fallback-mode flag, seed.
• predict_x0: Callable[(latent, t, prompt) -> latent] — injected diffusion denoiser. Real impl wraps an SD-1.5 pipeline from diffusers (MPS-feasible; SDXL optional); the offline stub is a trivial denoiser (e.g. identity-toward-a-fixed-image or a blur) so the projection loop is fully unit-testable with zero weights and no network, exactly like stub_cleanup for demix.
• projection_step(h1, h2, h_total, mask, curve, gains) -> (h1', h2') — pure numpy, the §6 correctness core, independently testable.
• guided_separate(positive_rgb, h_total, confidence_mask, film_curve, predict_x0, analysis, config) -> Optional[SeparationResult] — the §3 loop; returns one candidate pair with method="guided_diffusion", candidate_id="guided_i{steps}", and diagnostics carrying final masked sum-constraint error (reuse the structured-diagnostics pattern from WP-6).
• Reuse, do not reinvent: densitometry.srgb_to_linear/luminance_from_linear/ linear_to_srgb for φ and render; PiecewiseFilmCurve.inverse if a density-space variant is needed; app.demix.analyze_scan for prompts; latent_optimizer's VAE encode/decode helpers (now public) and gain-fit; SeparationResult as the return type.
• Integration: register as a candidate source in generate_candidates behind an include_guided: bool = False flag (off by default — GPU/slow), via the same _append_*-style helper WP-5.1/WP-6 established. One UI checkbox "Guided diffusion (very slow, GPU)". Density path required (skip if density is None), same rule as demix/DIP.

Everything offline-testable except the real-SD run, which is manual/GPU-gated and documented in the Result note — the standing pattern for token/GPU-gated work in this repo.

---

9. Acceptance criteria for the implementation WP (write these into the WP-9 impl spec)

Offline (CI, stub denoiser — the real gate):

1. projection_step unit tests: (a) idempotent on an already-consistent pair; (b) strictly reduces |g₁h₁+g₂h₂ − ĥ| on VALID; (c) non-negativity — never returns a negative layer; (d) shoulder pixels are only pushed up, never down; (e) TOE pixels untouched. Each with teeth (construct an input that violates it, prove the naive version fails).
2. Full loop with the stub denoiser drives two chains to satisfy the masked sum constraint to < 2% (same bar as demix) on the two committed 64×64 fixtures AND the 256px generated cases. This proves the projection math is correct independent of any diffusion prior.
3. Registration: generate_candidates(include_guided=True) (tiny step count, stub denoiser, monkeypatched) yields a method=="guided_diffusion" candidate that rank_candidates scores. Goes through the real generate_candidates surface (per the WP-8 lesson — no direct-call theater).

Manual (GPU-gated, documented, NOT CI): 4. Real SD-1.5 run on the 256px cases: guided candidates ranked against DIP (current best) on per-layer LPIPS. Target: competitive with or beating DIP on ≥50% of cases. Report the real table verbatim; if it loses, that is an honest finding (the fixture-scale lesson applies — judge at 256px, and note real-scan behavior is the true test).

Invariant: the projection must be a genuine projection — verify (2) by asserting the residual is monotone-non-increasing across the last 30% of steps (where λ≈1).

---

10. Fallback path (behind a flag, not primary)

If the VAE re-encode (§5) proves too lossy to hit the <2% offline bar, fall back to gradient guidance (ΠGDM/DPS-style): instead of hard range-space replacement, add ∇_z ‖g₁φ(A)+g₂φ(B) − Ĥ‖²_VALID to each chain's score at each step. This backprops through the VAE decode only (not the UNet — far cheaper than full DPS), which is the "PSLD caveats in latent space" the plan flags. Decision: implement projection as primary, gradient guidance behind config.mode="gradient", and let the offline bar-2 test decide which ships. Do not implement gradient-through-UNet DPS (Chung et al. 2022) — the plan explicitly deprioritizes it as expensive, and the linear operator makes projection the right tool.

---

11. Open questions / risks (for the implementer to surface, not silently resolve)

• VAE reconstruction floor: if vae_encode(vae_decode(x)) ≠ x error exceeds the 2% sum bar on its own, no amount of projection helps — measure this first, before building the loop, and STOP/report if the floor is above bar.
• Warm start: random latents vs img2img-encoding a heuristic split as the starting point. Recommendation: img2img warm start from the best demo/DIP candidate — turns WP-9 into a refiner of the existing pipeline rather than a cold generator, which is both cheaper and more likely to beat the baseline. Decide empirically.
• Schedule λ(t): linear is the starting point; if late-step hard projection causes visible seams, soften. This is tuning, in-scope for the impl WP, not a redesign.
• Model licensing / weights: SD-1.5 weights + diffusers is a real dependency and a first network/GPU dependency for the repo — must stay lazy-imported and fully absent from the offline path, like replicate/anthropic today.

---

12. One-paragraph summary for the WP-9 impl kickoff

Two independent Stable-Diffusion chains, prompted by the WP-5 VLM scene descriptions, coupled only by a per-step DDNM-style null-space projection performed in exposure space on the luminance channel: decode each chain's predicted-x₀ to RGB, map to exposure via the densitometry sRGB→luminance transform, fit the I.5a nuisance gains by least squares on VALID pixels, split the residual Ĥ − (g₁h₁+g₂h₂) equally with non-negativity clamping and a shoulder-inequality/TOE-skip per the confidence mask, carry chroma from the prior, re-encode, and renoise — with a noise-level-scaled guidance strength and sparse early projection to control the VAE round-trip cost. No UNet backprop. Offline-tested to <2% sum-constraint error with a stub denoiser; real SD run is GPU-gated and judged at 256px against the DIP baseline. Gated: do not implement until the 50-case 256px benchmark shows L0–L3 have plateaued.