diff --git a/CHANGELOG.md b/CHANGELOG.md index 6cdb1047d9..4836291770 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -28,10 +28,22 @@ Changes added by this fork to RapidRAW. Upstream application changes remain in t - Embed and verify an sRGB ICC profile in JPEG, PNG, TIFF, and WebP desktop and CLI exports. Profile labeling is independent of capture-metadata retention and GPS removal. - Report unsupported capture-metadata retention requests before writing output, including TIFF inputs or outputs. Remove the ineffective RAW highlight-compression setting without changing RAW development or existing saved edits. +### Nonlocal denoise speed + +- Cut Nonlocal model time about 4× on NVIDIA CUDA. On an RTX 5060 Ti with the shipped settings (TF32 and graph optimization off), a 32.5 MP photograph drops from an estimated 91 s to an estimated 23 s of model time (tile count × measured per-tile time; RAW decoding, noise estimation and DNG writing excluded). Two changes combine: + - The ONNX model is derived from the previous export by `python -m rapidraw_denoise.onnx_rewrite`. It keeps every weight and the FP32 precision, drops the all-zero convolution biases, and runs the 1×1 convolutions (20 dense, 10 grouped) as matrix multiplications, taking a 320×320 tile from 370 ms to 153 ms. It passes the same frozen fixture gates as the previous model on all 18 fixture tiles, on CUDA and on CPU. + - Shrink the tile halo from 64 to 40 pixels, so a 32.5 MP photograph needs 150 tiles instead of 247 (39% fewer, which also shortens CoreML runs; that was not timed end to end). Against the previous halo, 99.9% of pixels in 8-bit test renders of three photographs differ by at most 0.71 of 255 levels (worst pixel 6.8). The largest difference to a converged reference over nine test crops is 0.031 of the noise level. With the noise level halved (a research setting; Intensity does not do this), differences are larger: 99.9% within 2.2 levels. +- Reinstall the Nonlocal model with `install_model kind=nonlocal`. This release accepts only the new model (a single pinned hash; 320 px tiles, 40 px halo, 240 px core), and a model installed by an earlier release is rejected with reinstall instructions. Cached predictions from earlier versions are not reused. Job records and the prediction cache record which model ran. The CoreML package is unchanged and uses the new halo; its per-tile time is unchanged, and full-photograph CoreML parity was last measured with the previous halo. +- Remove the Maximum quality option. Nonlocal always denoises each tile in one pass; four-rotation averaging gained only about 0.04 to 0.06 dB in the patch study at four times the cost. Requests that pass `quality` to `start_denoise` are rejected by the MCP host's schema validation, and the research CLI and Python pipeline no longer take `--ensemble`. + ### MCP mask composition - Add native parent-mask duplication with fresh IDs, optional inversion, and independent local adjustments. Generated AI subject, depth, and other supported AI components can now join an existing parent in additive, subtractive, or intersect mode. The Node MCP host exposes both operations; installed releases require an updated native bridge and host. +### MCP session storage + +- Store each distinct embedded mask image once in a session's manifest instead of repeating it in every history snapshot. Sessions with many AI mask edits shrink from tens or hundreds of megabytes to a few, and previously saved sessions still load. A manifest written by this version cannot be read by an older native bridge. + ### Preview reliability - Clear the previous photo's GPU frame, interactive patch, and mask overlay when switching photos, including cached photos with matching dimensions. Wait for the selected photo's native load before rendering its adjustments. diff --git a/denoise/README.md b/denoise/README.md index 52f5bf102c..90b7eadb22 100644 --- a/denoise/README.md +++ b/denoise/README.md @@ -2,11 +2,11 @@ Process Bayer RAW photographs with a noise-conditioned denoiser, and test improvements using reproducible experiments. RapidRAW's production Nonlocal path runs **in-process through native ONNX Runtime** (CPU or Linux CUDA) or **directly through CoreML.framework on macOS** from a pinned FP32 bundle — no Python, PyTorch, NumPy or SciPy in production. This optional Python package remains as research/export/benchmark tooling: model export and validation, reference inference, and the study CLI. The editor's existing NIND AI denoiser remains the lightweight option. -The GPU path combines the [2026 learned nonlocal model](https://github.com/MIA-UIB/nonlocal-matchfilter) with per-channel blind noise estimation, bounded-memory tiling and optional transform averaging. Native MCP jobs produce a **Bayer DNG** for RapidRAW's normal demosaic and colour pipeline. The research CLI produces a 16-bit sRGB TIFF, review renders, a floating-point packed RAW tensor and a provenance manifest. Both preserve the original photograph. Commercial parity and state-of-the-art performance have **not** been established. +The GPU path combines the [2026 learned nonlocal model](https://github.com/MIA-UIB/nonlocal-matchfilter) with per-channel blind noise estimation and bounded-memory tiling. Native MCP jobs produce a **Bayer DNG** for RapidRAW's normal demosaic and colour pipeline. The research CLI produces a 16-bit sRGB TIFF, review renders, a floating-point packed RAW tensor and a provenance manifest. Both preserve the original photograph. Commercial parity and state-of-the-art performance have **not** been established. ## Native runtime bundle (production) -Production inference reads an ONNX bundle directory containing `model.onnx` + `manifest.json`. The accepted bundle is identified by model SHA-256 `df7f21ddbdfecd0984896a902623f75aa883d25b5f862c49d8c2b3f63a7dd2d9`, source checkpoint lineage `c16747d852b93a95908792cdbac901f89cca98e35b91ea7db6214de42fbd3cad`, tile 320 / halo 64 / FP32 / RGBG. The native backend validates this contract fail-closed before inference. Produce or re-validate bundles with this package's export tooling (`rapidraw_denoise.onnx_export`, `rapidraw_denoise.onnx_validate`); the Python environment below is only needed for that tooling, never for production inference. +Production inference reads an ONNX bundle directory containing `model.onnx` + `manifest.json`. The accepted bundle is identified by model SHA-256 `ba788a97f247058ea03d14bbf2f289adc8f4666039f27d0d530e6f67d23dd73c`, source checkpoint lineage `c16747d852b93a95908792cdbac901f89cca98e35b91ea7db6214de42fbd3cad`, tile 320 / halo 40 / retained core 240 / FP32 / RGBG. The native backend validates this contract fail-closed before inference and rejects any other model hash or tiling. Produce or re-validate bundles with this package's export tooling (`rapidraw_denoise.onnx_export`, `rapidraw_denoise.onnx_rewrite`, `rapidraw_denoise.onnx_validate`); the Python environment below is only needed for that tooling, never for production inference. Normal setup needs no bundle path: call `install_model` with `kind: "nonlocal"` (selects the bundle for `RAPIDRAW_NONLOCAL_PROVIDER`, default `cuda` on Linux, `coreml` on macOS, and `cpu` on other platforms) and the pinned runtime downloads from the [public model repo](https://huggingface.co/sheldonxxxx/RapidRAW-Nonlocal-Denoise) at a pinned commit with every hash verified. `RAPIDRAW_NONLOCAL_BUNDLE` remains an explicit developer override when set: @@ -17,6 +17,30 @@ export RAPIDRAW_NONLOCAL_PROVIDER=cpu Use `cpu`, or `cuda` on a Linux x86_64 host with the pinned NVIDIA runtime pack (recommended — see [ONNX CUDA setup](../docs/mcp/onnx-cuda.md#install-the-pinned-runtime-pack), no `ORT_DYLIB_PATH` needed) or, advanced only, a manual CUDA runtime via `ORT_DYLIB_PATH`. On macOS, `coreml` selects direct native CoreML.framework inference from a pinned `.mlpackage` bundle (see below). There is no provider fallback: `cuda` on an unsupported platform, `coreml` off macOS, or a failing execution provider fails the job explicitly, and CPU never silently substitutes. Graph optimization stays off and CUDA TF32 stays off under the accepted strict policy. Results are cached under the workspace's `nonlocal-onnx-cache`, keyed by backend generation, bundle and checkpoint hashes, provider, ORT build identity and inference configuration; legacy Python-worker cache entries are never reused. +### How the ONNX bundle is derived + +The published ONNX model is the FP32 export of the pinned upstream checkpoint with its graph rewritten by `rapidraw_denoise.onnx_rewrite` for faster execution. The rewrite drops the all-zero convolution biases the exporter materialises and runs every dense and grouped 1×1 convolution with at least 256 input channels as a matrix multiplication. It keeps every weight and FP32 arithmetic, and the runtime keeps graph optimization and TF32 off. The rewritten model's manifest records the source export hash and rewrite counts. On an RTX 5060 Ti with the strict policy (no TF32, no graph optimization, cuDNN heuristic search), a 320×320 tile takes about 153 ms instead of 370 ms. The rewritten graph passes the same frozen fixture gates as the unrewritten export on all 18 fixture tiles, on CUDA and on CPU (elementwise |difference| ≤ 1e-4 + 1e-4·|reference|, mean absolute error ≤ 1e-5, 99th-percentile absolute error ≤ 1e-4). CPU was not timed, and CoreML uses a separate model. + +To reproduce the bundle, first build the unrewritten source bundle with `rapidraw_denoise.onnx_export` (it needs `--checkpoint`, `--fixtures` and `--output`; its `model.onnx` must hash to `df7f21ddbdfecd0984896a902623f75aa883d25b5f862c49d8c2b3f63a7dd2d9`) or take it from the earlier model-repository revision. Then install the two packages the rewrite needs (the model hash assumes `onnx` 1.23.1; other versions may serialise differently and give a different hash, which the runtime rejects) and derive the bundle. The tool never overwrites an existing output directory, and `--fixtures` adds a first-tile CPU check against the source graph: + +```sh +uv pip install --python .venv/bin/python onnx==1.23.1 onnxruntime +.venv/bin/python -m rapidraw_denoise.onnx_rewrite \ + --bundle /path/to/unrewritten-bundle \ + --output /path/to/bundle-packed-gemm-halo40-320-fp32 \ + --fixtures /path/to/fixtures +``` + +Validate the output before use: + +```sh +.venv/bin/python -m rapidraw_denoise.onnx_validate \ + --bundle /path/to/bundle-packed-gemm-halo40-320-fp32 \ + --fixtures /path/to/fixtures --provider cpu --report /path/to/validate-cpu.json +``` + +[`tests/test_onnx_rewrite.py`](tests/test_onnx_rewrite.py) covers the rewrite. + ### Direct CoreML bundle (macOS production) On macOS, production CoreML inference is native: it uses CoreML.framework directly and does not require Python, coremltools, PyTorch or ONNX Runtime. Prefer `install_model kind="nonlocal"` with `RAPIDRAW_NONLOCAL_PROVIDER=coreml`; to use a manual bundle instead, point `RAPIDRAW_NONLOCAL_BUNDLE` at a directory containing `packed.mlpackage` + `manifest.json`: @@ -26,7 +50,7 @@ export RAPIDRAW_NONLOCAL_BUNDLE=/path/to/bundle-packed-coreml-fp32 export RAPIDRAW_NONLOCAL_PROVIDER=coreml ``` -The accepted bundle is identified by package SHA-256 `ecf8f41b72b8e79a7210f61b13b3a21ad2ddbc79c475039be367fc011a98f5d3`, source checkpoint lineage `c16747d852b93a95908792cdbac901f89cca98e35b91ea7db6214de42fbd3cad`, direct TorchScript conversion (`coremltools-direct-torchscript`, FLOAT32, static-packed sampler), tile 320 / halo 64 / RGBG with `raw_with_noise` [1,8,320,320] input and `denoised_raw` [1,4,320,320] output. The native backend validates this contract fail-closed — including the package digest and a runtime tree hash — before compiling the package to `.mlmodelc` and loading one `MLModel` per job (all compute units, no low-precision GPU accumulation). The Python converter (`rapidraw_denoise` tooling / `denoise/scripts/convert_coreml_direct.py`) remains developer-only artifact tooling, never a production dependency. CoreML predictions are cached under the workspace's `nonlocal-coreml-cache` (protocol 3) and cannot collide with ONNX or legacy worker caches; jobs that predate the native runtimes, and cross-family ONNX↔CoreML resumes, are rejected — start a new denoise job instead. +The accepted bundle is identified by package SHA-256 `ecf8f41b72b8e79a7210f61b13b3a21ad2ddbc79c475039be367fc011a98f5d3`, source checkpoint lineage `c16747d852b93a95908792cdbac901f89cca98e35b91ea7db6214de42fbd3cad`, direct TorchScript conversion (`coremltools-direct-torchscript`, FLOAT32, static-packed sampler), tile 320 / halo 40 / RGBG with `raw_with_noise` [1,8,320,320] input and `denoised_raw` [1,4,320,320] output. The native backend validates this contract fail-closed — including the package digest and a runtime tree hash — before compiling the package to `.mlmodelc` and loading one `MLModel` per job (all compute units, no low-precision GPU accumulation). The Python converter (`rapidraw_denoise` tooling / `denoise/scripts/convert_coreml_direct.py`) remains developer-only artifact tooling, never a production dependency. CoreML predictions are cached under the workspace's `nonlocal-coreml-cache` (protocol 3) and cannot collide with ONNX or legacy worker caches; jobs that predate the native runtimes, and cross-family ONNX↔CoreML resumes, are rejected — start a new denoise job instead. ## Research export environment (not production) @@ -68,20 +92,17 @@ Open a Bayer RAW on that host, then call: "session_id": "", "expected_revision": 0, "method": "nonlocal", - "intensity": 100, - "quality": "balanced" + "intensity": 100 } ``` Supply the current revision to `rapidraw_start_denoise`. Poll `rapidraw_get_job` with its `job_id`, then continue editing the completed `result_session_id`. The new session contains the source edits captured at submission; later parent edits do not change the job. The original session and original RAW remain intact. Nonlocal cancellation is cooperative at tile boundaries and cannot abort an active ORT or CoreML call; the editing engine remains alive. `resume_job` restarts an interrupted, cancelled or failed job from the captured input; partial tiles are not checkpoints. Cross-family ONNX↔CoreML resume is rejected; legacy worker jobs cannot resume. -- **Balanced**: one inference pass; the default. -- **Maximum**: average four rotations; roughly four inference passes. The quality benefit in the patch study was small. - **Intensity**: blend the full prediction with the sensor input, from 0 to 100. Zero writes an identity DNG without inference. Nonlocal defaults to 100; existing NIND/BM3D defaults remain 50. Predictions are cached under the workspace's `nonlocal-onnx-cache`, keyed by backend generation, bundle and checkpoint hashes, provider, ORT build identity and inference configuration; legacy Python-worker cache entries are never reused. macOS CoreML predictions use the separate `nonlocal-coreml-cache` (protocol 3), keyed by package/tree/checkpoint hashes, compute policy and OS identity. Changing intensity reuses the full prediction. Every cache read verifies the result manifest, dimensions and checksum. Corrupt entries fail explicitly. Completed DNG sessions are independent of this cache; verified prediction directories may be removed while no job is active to reclaim space. There is no automatic cache eviction. -The native path uses fixed 320×320 packed tiles with a 64-pixel halo. Sensor planes stay in host RAM, so total RAM still scales with photograph size. Output is uncompressed float32 Bayer DNG to avoid integer requantization, approximately 130 MB for a 32.5 MP active sensor. The session and export copies, prediction cache and temporary buffers require additional storage. Jobs reserve 20 GiB plus estimated working space. Inactive sensor borders are removed with CFA/black-level/crop coordinates adjusted; the normal rendered crop and orientation are retained. Signed samples feed noise estimation, and sensor values above white remain unchanged by the prediction blend. +The native path uses fixed 320×320 packed tiles with a 40-pixel halo (240-pixel retained core), one inference pass per tile; a 32.5 MP photograph takes 150 tiles. Sensor planes stay in host RAM, so total RAM still scales with photograph size. Output is uncompressed float32 Bayer DNG to avoid integer requantization, approximately 130 MB for a 32.5 MP active sensor. The session and export copies, prediction cache and temporary buffers require additional storage. Jobs reserve 20 GiB plus estimated working space. Inactive sensor borders are removed with CFA/black-level/crop coordinates adjusted; the normal rendered crop and orientation are retained. Signed samples feed noise estimation, and sensor values above white remain unchanged by the prediction blend. This first integration is an MCP feature on a configured host. The desktop denoise dialog continues to offer its existing methods. A Mac can open the resulting DNG; automatic desktop-to-server submission, transfer and local Metal inference are not included. @@ -102,12 +123,12 @@ Run inside the server's usual headless display environment on Linux. The check e .venv/bin/rapidraw-denoise \ /path/to/photo.CR3 \ --checkpoint /path/to/model-cache/nonlocal-raw-weights.pt \ - --output /path/to/results/photo-v1 --tile 320 --halo 64 + --output /path/to/results/photo-v1 --tile 320 --halo 40 ``` The command works with the default reference sampler. After installing and verifying the optional extension, prefix it with `RAPIDRAW_DENOISE_SAMPLER=cuda` for the faster path. The output directory must be new. The command requires at least 20 GiB free space beyond its estimated outputs. Inference defaults to CUDA and fails if it is unavailable; CPU diagnostic runs require `--device cpu`. Full-resolution arrays stay in host RAM and tiles bound GPU memory. Select smaller tiles if other applications need VRAM. A 320-pixel tile means 640 sensor pixels across, since Bayer channels are packed at half resolution. -`--ensemble 4` or `--ensemble 8` averages transformed predictions and saves their disagreement. Four rotations produced a small, consistent gain in the [expanded patch study](STUDY.md#observed-patch-study-outcome), at roughly four inference passes per photograph. Disagreement measures transform sensitivity, **not** calibrated uncertainty. `--noise-scale` scales the estimated noise standard deviation; it is not a blend with the original. `--noise-profile` accepts a JSON `NoiseProfile` with four `shot` and four `read` values in RGBG order, where `read` is variance in black-subtracted, white-normalized RAW units. Poorly identified estimates are recorded in the manifest; severely clipped inputs can require a measured profile. +`--noise-scale` scales the estimated noise standard deviation; it is not a blend with the original. `--noise-profile` accepts a JSON `NoiseProfile` with four `shot` and four `read` values in RGBG order, where `read` is variance in black-subtracted, white-normalized RAW units. Poorly identified estimates are recorded in the manifest; severely clipped inputs can require a measured profile. Inspect `original.png` and `denoised.png` at the same native magnification. Open `denoised.tiff` as a new photograph in RapidRAW for further editing. For agent editing, use the existing `rapidraw_open_photo` workflow with the returned TIFF path. The new file has denoising and demosaicing baked in; the original RAW and sidecar remain the authoritative source. Existing RAW adjustments are not copied onto the TIFF. diff --git a/denoise/STUDY.md b/denoise/STUDY.md index 35c22afbd3..8154df2d7d 100644 --- a/denoise/STUDY.md +++ b/denoise/STUDY.md @@ -86,7 +86,7 @@ Choose a checkpoint and ensemble using validation only. Require positive mean RA "checkpoint_sha256": "REPLACE_WITH_ACTUAL_SHA256", "ensemble": 4, "tile": 320, - "halo": 64, + "halo": 40, "reason": "Replace with the measured validation decision" } ``` @@ -114,4 +114,4 @@ Held-out scene-balanced RAW PSNR (dB), compared with a single prediction: Every held-out scene mean improved; individual crops did not all improve. The worst crop changes were -0.130 dB for synthetic noise and -0.041 dB for real pairs. Removing the exact duplicate real-pair crop gives a +0.041 dB scene-balanced change, consistent with the primary result. Edge PSNR also improves on average. These are small gains and do not establish a visibly transformative result or commercial parity. -Eight-transform averaging did not materially improve validation over four rotations. Lowering noise conditioning damaged reconstruction accuracy, and larger tile context changed mean validation PSNR by about 0.00001 dB. These results support retaining the tested 320/64 tile/halo settings and treating four rotations as an optional quality-versus-runtime tradeoff. Full-resolution photographic inspection remains a separate acceptance step; ONNX conversion has not been performed. +Eight-transform averaging did not materially improve validation over four rotations. Lowering noise conditioning damaged reconstruction accuracy, and larger tile context changed mean validation PSNR by about 0.00001 dB. These results supported the original 320/64 tile/halo settings. Production has since moved to a 40-pixel halo (see the [README](README.md)) and no longer offers four-rotation averaging, which remains available only to this research study. Full-resolution photographic inspection remains a separate acceptance step; ONNX conversion has not been performed. diff --git a/denoise/rapidraw_denoise/benchmark.py b/denoise/rapidraw_denoise/benchmark.py index e20847c899..3695de123c 100644 --- a/denoise/rapidraw_denoise/benchmark.py +++ b/denoise/rapidraw_denoise/benchmark.py @@ -62,13 +62,13 @@ def main(): noisy=synthesize(clean,[shot]*4,[read]*4,seed=19000+case_index,row_sigma=row,black_bias=bias) oracle=NoiseProfile([shot]*4,[read]*4,[],method="synthetic-oracle-upper-reference") estimated=estimate_noise(noisy) - variants=[("blind",estimated,False,1)] - if a.ablations:variants += [("oracle",oracle,False,1),("pilot",estimated,True,1),("ensemble4",estimated,False,4)] - if a.row_ablation:variants += [("row-corrected",estimated,False,1)] + variants=[("blind",estimated,False)] + if a.ablations:variants += [("oracle",oracle,False),("pilot",estimated,True)] + if a.row_ablation:variants += [("row-corrected",estimated,False)] base={**case,"corruption":corruption,"noise":{"shot":shot,"read":read,"row_sigma":row,"black_bias":bias}, "input":metrics(clean,noisy)} - for name,profile,pilot,ensemble in variants: - output,_,info=denoise(noisy,model,profile,tile=320,halo=64,pilot=pilot,ensemble=ensemble,row_correction=name=="row-corrected") + for name,profile,pilot in variants: + output,info=denoise(noisy,model,profile,tile=320,halo=40,pilot=pilot,row_correction=name=="row-corrected") record={**base,"variant":name,"output":metrics(clean,output),"inference":info} records.append(record) artifact=a.output/f'{case["scene"]}-{case["crop"]}-{corruption}-{name}.npz' diff --git a/denoise/rapidraw_denoise/capture_reference.py b/denoise/rapidraw_denoise/capture_reference.py index ed401aba84..2a69aed0a0 100644 --- a/denoise/rapidraw_denoise/capture_reference.py +++ b/denoise/rapidraw_denoise/capture_reference.py @@ -26,7 +26,7 @@ PINNED_CHECKPOINT_SHA256 = "c16747d852b93a95908792cdbac901f89cca98e35b91ea7db6214de42fbd3cad" NATIVE_TILE = 320 -NATIVE_HALO = 64 +NATIVE_HALO = 40 NATIVE_CORE = NATIVE_TILE - 2 * NATIVE_HALO @@ -124,7 +124,7 @@ def build_arg_parser() -> argparse.ArgumentParser: "The actual digest is recorded and the label is suffixed; output " "must never be presented as pinned-checkpoint acceptance.") p.add_argument("--skip-full", action="store_true", - help="Skip full-photo ensemble assembly (e.g. 32 MP CPU inference " + help="Skip full-photo assembly (e.g. 32 MP CPU inference " "impractical). Tile fixtures are still captured; the report records " "the skip. Full-photo gates remain pending.") return p @@ -156,7 +156,7 @@ def main(argv: list[str] | None = None) -> int: from rapidraw_denoise.worker import validate_request, CHECKPOINT_SHA256 shape = validate_request(request) if tuple(shape[1:]) and (request.get("tile") != NATIVE_TILE or request.get("halo") != NATIVE_HALO): - raise ValueError("Native worker contract requires tile 320 / halo 64") + raise ValueError("Native worker contract requires tile 320 / halo 40") checkpoint_sha = sha256_file(checkpoint) if checkpoint_sha != CHECKPOINT_SHA256 or checkpoint_sha != PINNED_CHECKPOINT_SHA256: if not args.allow_unpinned_checkpoint: @@ -279,13 +279,13 @@ def hook(_m, inputs, _out): hnd.remove() captured_blocks = block_io.get("_captured", {}) - # Full assembled predictions for both native quality settings when - # hardware permits (ensemble 1 and 4). Uses the shared denoise path. + # Full assembled prediction when hardware permits. Uses the shared + # denoise path. # --skip-full records the skip explicitly; tile fixtures (the parity # inputs) are still captured and full-photo gates remain pending. from rapidraw_denoise.inference import denoise - full_predictions: dict[str, np.ndarray] = {} - full_infos: dict[str, dict] = {} + full_prediction: np.ndarray | None = None + full_info: dict | None = None full_skipped_reason: str | None = None if args.skip_full: full_skipped_reason = ( @@ -293,15 +293,14 @@ def hook(_m, inputs, _out): "impractical); full-photo assembly gates remain pending" ) else: - for ensemble in (1, 4): - out, _dis, info = denoise( - packed, model, profile, NATIVE_TILE, NATIVE_HALO, - ensemble, args.noise_scale, - ) - if out.shape != shape or not np.isfinite(out).all(): - raise ValueError(f"Invalid full prediction for ensemble {ensemble}") - full_predictions[str(ensemble)] = out.astype(np.float32) - full_infos[str(ensemble)] = info + out, info = denoise( + packed, model, profile, NATIVE_TILE, NATIVE_HALO, + args.noise_scale, + ) + if out.shape != shape or not np.isfinite(out).all(): + raise ValueError("Invalid full prediction") + full_prediction = out.astype(np.float32) + full_info = info # Publish to a temp dir then atomically rename (never overwrite). import tempfile, shutil @@ -317,8 +316,8 @@ def hook(_m, inputs, _out): origin_y=np.int64(y), origin_x=np.int64(x), ) - for ensemble, arr in full_predictions.items(): - np.save(tmp / f"full-ensemble-{ensemble}.npy", arr, allow_pickle=False) + if full_prediction is not None: + np.save(tmp / "full-prediction.npy", full_prediction, allow_pickle=False) if captured_blocks: (tmp / "blocks").mkdir() for key, val in captured_blocks.items(): @@ -356,7 +355,7 @@ def hook(_m, inputs, _out): "tile_input_dtype": "float32", "tile_output_shape": [1, 4, NATIVE_TILE, NATIVE_TILE], "tile_output_dtype": "float32", - "full_infos": full_infos, + "full_info": full_info, "full_skipped_reason": full_skipped_reason, "block_io_shapes": { k: {"features": v["features_shape"], "offsets": v["offsets_shape"]} diff --git a/denoise/rapidraw_denoise/cli.py b/denoise/rapidraw_denoise/cli.py index 9e498c2faf..4ab710cdfa 100644 --- a/denoise/rapidraw_denoise/cli.py +++ b/denoise/rapidraw_denoise/cli.py @@ -18,8 +18,7 @@ def main(): p.add_argument("--checkpoint",required=True,type=Path) p.add_argument("--device",choices=["cuda","cpu"],default="cuda") p.add_argument("--tile",type=int,default=256) - p.add_argument("--halo",type=int,default=64) - p.add_argument("--ensemble",type=int,choices=[1,4,8],default=1) + p.add_argument("--halo",type=int,default=40) p.add_argument("--pilot",action="store_true") p.add_argument("--noise-scale",type=float,default=1.) p.add_argument("--noise-profile",type=Path) @@ -40,7 +39,7 @@ def main(): if args.noise_profile: profile=NoiseProfile(**json.loads(args.noise_profile.read_text())) model=load_model(args.checkpoint,args.device) - candidate,disagreement,info=denoise(frame.packed,model,profile,args.tile,args.halo,args.ensemble,args.noise_scale,args.pilot, + candidate,info=denoise(frame.packed,model,profile,args.tile,args.halo,args.noise_scale,args.pilot, progress=lambda x: print(json.dumps(x),flush=True)) info.update(source_metadata) info.update({"checkpoint_sha256":sha256(args.checkpoint),"status":"experimental; commercial parity unverified", @@ -49,7 +48,6 @@ def main(): temporary=Path(tempfile.mkdtemp(prefix=".denoise-",dir=args.output.parent)) try: np.save(temporary/"denoised-raw.npy",candidate,allow_pickle=False) - np.save(temporary/"transform-disagreement.npy",disagreement,allow_pickle=False) for name,planes in [("original",None),("denoised",candidate)]: rgb=frame.render(planes) save_render(temporary/f"{name}.tiff",rgb) diff --git a/denoise/rapidraw_denoise/compare_full.py b/denoise/rapidraw_denoise/compare_full.py index 7da16e23e4..5b654d4c00 100644 --- a/denoise/rapidraw_denoise/compare_full.py +++ b/denoise/rapidraw_denoise/compare_full.py @@ -7,11 +7,6 @@ errors and exit codes: 0 = all pass, 2 = gate failures, 1 = tool/invalid error (saved as JSON). -For ensemble-4 comparisons (``--ensemble 4``), an additional diagnostic -``regions_e4`` entry unions the inverse-mapped tile-boundary lines of all -four rotation passes (diagnostic attribution only; tiling, thresholds and -pass/fail gates are unchanged), plus stats over the newly attributed seam -pixels (``seam_e4_extra``). """ import hashlib import json @@ -22,8 +17,10 @@ from .metrics import full_metrics, load_float32_array, require_case_set -CORE = 192 -OUTER_BAND = 64 +TILE = 320 +HALO = 40 +CORE = TILE - 2 * HALO +OUTER_BAND = HALO SEAM_HALF_WIDTH = 8 @@ -64,44 +61,7 @@ def region_masks(h, w): return {"outer": outer, "seam": seam, "interior": interior} -def region_masks_e4(h, w): - """Diagnostic seam mask unioning all four rotation-pass boundaries. - - Pass ``i`` tiles ``rot90(conditioned, i)``; stitch lines at multiples of - CORE in pass coordinates map back through the inverse rotation. For odd - dimensions the mirrored lines (H-1-p, W-1-q) differ from the unrotated - grid, which the plain mask would misclassify as interior. Diagnostic - only: same outer band, same construction rule per pass. - """ - rows = set(_grid_lines(h)) | {h - 1 - p for p in _grid_lines(h)} - cols = set(_grid_lines(w)) | {w - 1 - q for q in _grid_lines(w)} - rows = sorted(r for r in rows if 0 <= r < h) - cols = sorted(c for c in cols if 0 <= c < w) - yy, xx = np.mgrid[0:h, 0:w] - outer = ((yy < OUTER_BAND) | (yy >= h - OUTER_BAND) - | (xx < OUTER_BAND) | (xx >= w - OUTER_BAND)) - seam_y = np.zeros(h, bool) - seam_y[rows] = True - seam_x = np.zeros(w, bool) - seam_x[cols] = True - seam = (seam_y[:, None] | seam_x[None, :]) & ~outer - interior = ~(outer | seam) - return {"outer": outer, "seam": seam, "interior": interior} - - -def _region_stats(ae, masks): - stats = {} - for rname, mask in masks.items(): - sub = ae[:, mask] - stats[rname] = { - "max": float(sub.max()), - "mean": float(sub.mean()), - "pixels": int(sub.size), - } - return stats - - -def compare_files(name, actual_path, ref_path, ensemble=1): +def compare_files(name, actual_path, ref_path): actual_path, ref_path = Path(actual_path), Path(ref_path) record = {"comparison": name, "actual": {"path": str(actual_path), @@ -122,39 +82,19 @@ def compare_files(name, actual_path, ref_path, ensemble=1): entry.update(record) # Regions form a disjoint complete partition; order for readability. entry["regions"] = {k: entry["regions"][k] for k in ("outer", "seam", "interior")} - if ensemble == 4: - masks4 = region_masks_e4(h, w) - ae = np.abs(actual.astype("float64") - ref.astype("float64")) - entry["regions_e4"] = _region_stats(ae, masks4) - extra = masks4["seam"] & ~masks["seam"] - sub = ae[:, extra] - entry["seam_e4_extra"] = { - "note": "pixels attributed to seam only by the 4-pass diagnostic " - "mask (inverse-mapped rotation boundaries); diagnostic, " - "no gate effect", - "max": float(sub.max()), - "mean": float(sub.mean()), - "pixels": int(sub.size), - } - elif ensemble != 1: - raise ValueError(f"ensemble must be 1 or 4, got {ensemble}") return entry def main(argv=None): argv = list(sys.argv[1:] if argv is None else argv) # Usage: compare_full.py PAIRS_JSON OUTPUT_JSON [--expect-cases a,b,...] - # [--ensemble {1,4}] - positional, expected, ensemble = [], None, 1 + positional, expected = [], None i = 0 while i < len(argv): tok = argv[i] if tok == "--expect-cases" and i + 1 < len(argv): expected = [c for c in argv[i + 1].split(",") if c] i += 2 - elif tok == "--ensemble" and i + 1 < len(argv): - ensemble = int(argv[i + 1]) - i += 2 elif tok.startswith("--"): print(json.dumps({"status": "error", "error": f"unknown flag: {tok}"}), @@ -166,7 +106,7 @@ def main(argv=None): if len(positional) != 2: print(json.dumps({"status": "error", "error": "usage: compare_full.py PAIRS_JSON OUTPUT_JSON " - "[--expect-cases a,b,...] [--ensemble {1,4}]"}), + "[--expect-cases a,b,...]"}), file=sys.stderr) return 1 pairs_path, output_path = Path(positional[0]), Path(positional[1]) @@ -178,7 +118,7 @@ def main(argv=None): try: pairs = require_case_set(json.loads(pairs_path.read_text()), expected, label=str(pairs_path)) - out = [compare_files(name, a, r, ensemble) for name, a, r in pairs] + out = [compare_files(name, a, r) for name, a, r in pairs] except Exception as exc: # invalid inputs reject the report error = {"status": "error", "error": f"{type(exc).__name__}: {exc}", "pairs": str(pairs_path)} diff --git a/denoise/rapidraw_denoise/inference.py b/denoise/rapidraw_denoise/inference.py index fdd31a2197..38cfff28a2 100644 --- a/denoise/rapidraw_denoise/inference.py +++ b/denoise/rapidraw_denoise/inference.py @@ -1,10 +1,9 @@ """Memory-bounded RAW inference with a frozen photometric noise profile.""" import torch -from .pipeline import TilePredictor, denoise_with_predictor, tiled_apply, transform, inverse_transform +from .pipeline import TilePredictor, denoise_with_predictor, tiled_apply from .vendor.nonlocalmf.network import SimpleBlockMatchingUNet -__all__ = ["load_model", "tiled_apply", "transform", "inverse_transform", - "denoise", "TorchTilePredictor", "SimpleBlockMatchingUNet"] +__all__ = ["load_model", "tiled_apply", "denoise", "TorchTilePredictor", "SimpleBlockMatchingUNet"] def load_model(checkpoint, device="cuda"): @@ -45,5 +44,5 @@ def execution_info(self): @torch.inference_mode() -def denoise(packed, model, profile=None, tile=256, halo=64, ensemble=1, noise_scale=1., pilot=False, progress=None, row_correction=False): - return denoise_with_predictor(packed, TorchTilePredictor(model), profile, tile, halo, ensemble, noise_scale, pilot, progress, row_correction) +def denoise(packed, model, profile=None, tile=256, halo=40, noise_scale=1., pilot=False, progress=None, row_correction=False): + return denoise_with_predictor(packed, TorchTilePredictor(model), profile, tile, halo, noise_scale, pilot, progress, row_correction) diff --git a/denoise/rapidraw_denoise/onnx_backend.py b/denoise/rapidraw_denoise/onnx_backend.py index 7f10875531..e599175397 100644 --- a/denoise/rapidraw_denoise/onnx_backend.py +++ b/denoise/rapidraw_denoise/onnx_backend.py @@ -1,7 +1,7 @@ """ONNX Runtime tile predictor for the exported Nonlocal reference bundle. Loads and validates the bundle once per job and reuses one InferenceSession -across all tiles and ensemble passes. Fail-closed validation order: +across all tiles. Fail-closed validation order: 1. Bundle files present; manifest version / tile / precision / I/O contract. 2. Pinned checkpoint lineage (non-zero, pinned-match, equals the pinned SHA). diff --git a/denoise/rapidraw_denoise/onnx_export.py b/denoise/rapidraw_denoise/onnx_export.py index 6f95cfef5e..871645e520 100644 --- a/denoise/rapidraw_denoise/onnx_export.py +++ b/denoise/rapidraw_denoise/onnx_export.py @@ -30,6 +30,7 @@ PINNED_CHECKPOINT_SHA256 = "c16747d852b93a95908792cdbac901f89cca98e35b91ea7db6214de42fbd3cad" PRODUCTION_TILE = 320 +HALO = 40 INPUT_NAME = "raw_with_noise" OUTPUT_NAME = "denoised_raw" @@ -301,8 +302,8 @@ def visit(graph) -> None: "signal clipped [0,1]; noise sqrt(max(var,1e-12))", "pipeline_revision": "inference.py@tiled_apply+denoise (shared; unmodified)", "tile": args.tile, - "halo": 64, - "retained_core": 192, + "halo": HALO, + "retained_core": args.tile - 2 * HALO, "sampler_implementation": sampler_label, "precision": "fp32", "opset_imports": opsets, diff --git a/denoise/rapidraw_denoise/onnx_rewrite.py b/denoise/rapidraw_denoise/onnx_rewrite.py new file mode 100644 index 0000000000..c95796c5a6 --- /dev/null +++ b/denoise/rapidraw_denoise/onnx_rewrite.py @@ -0,0 +1,350 @@ +"""Throughput-only graph rewrites for the packed Nonlocal ONNX bundle. + +The accepted packed graph spends 92.6 % of its FLOPs in twenty dense 1x1 +convolutions (800-channel and up) and ONNX Runtime's CUDA FP32 path runs them +through a legacy SIMT convolution kernel at about a third of cuBLAS SGEMM speed. +Its ten grouped 1x1 convolutions carry almost no arithmetic but run through a +cuDNN direct-convolution kernel that takes a third of the remaining time. The +rewrites remove those costs without changing a weight, a dtype or the set of +products summed for any output: + +* drop provably all-zero Conv biases (the exporter materialises them as + ``Expand(CastLike(0))`` nodes feeding three-input Convs); +* express every dense batch-1 1x1 Conv with at least ``--min-channels`` input + channels as ``Reshape -> MatMul -> Reshape`` so the product runs as one GEMM; +* express every grouped batch-1 1x1 Conv with at least ``--min-channels`` input + channels the same way, as one batched GEMM with the group as the batch. + +The tool derives a new, separately hashed bundle from an accepted bundle and +records the source hash in the manifest; it never modifies or overwrites the +source. A derived bundle is not an accepted bundle until ``onnx_validate`` passes +on every fixture tile, CPU and CUDA, under the frozen gates. + +Usage: + python -m rapidraw_denoise.onnx_rewrite --bundle SRC --output NEW [--fixtures DIR] +""" +from __future__ import annotations + +import argparse +import hashlib +import json +import sys +import tempfile +from pathlib import Path + +import numpy as np + +MIN_IN_CHANNELS = 256 +MODEL_NAME = "model.onnx" +REPORT_NAME = "rewrite-report.json" +REWRITE_VERSION = "gemm-1x1-v2" + + +def sha256_file(path: Path) -> str: + digest = hashlib.sha256() + with open(path, "rb") as stream: + for block in iter(lambda: stream.read(4 * 1024 * 1024), b""): + digest.update(block) + return digest.hexdigest() + + +def _is_zero_tensor(tensor) -> bool: + from onnx import numpy_helper + + return not np.any(numpy_helper.to_array(tensor)) + + +def _make_zero_test(model): + """Return ``is_zero(name)``: True only when the tensor is provably all zero.""" + inits = {i.name: i for i in model.graph.initializer} + producers = {o: n for n in model.graph.node for o in n.output} + + def is_zero(name: str, depth: int = 0) -> bool: + if depth > 8: + return False + if name in inits: + return _is_zero_tensor(inits[name]) + node = producers.get(name) + if node is None: + return False + if node.op_type == "Constant": + for attr in node.attribute: + if attr.name == "value" and attr.HasField("t"): + return _is_zero_tensor(attr.t) + return False + if node.op_type == "ConstantOfShape": + for attr in node.attribute: + if attr.name == "value" and attr.HasField("t"): + return _is_zero_tensor(attr.t) + return True # ONNX default fill value is float 0 + if node.op_type in ("Expand", "CastLike", "Cast"): + return is_zero(node.input[0], depth + 1) + return False + + return is_zero + + +def _attributes(node) -> dict: + from onnx import helper + + return {a.name: helper.get_attribute_value(a) for a in node.attribute} + + +def _is_pointwise(node, inits) -> bool: + """A 1x1, stride-1, unpadded, undilated Conv with a constant weight (any group).""" + a = _attributes(node) + return ( + list(a.get("kernel_shape", [])) == [1, 1] + and list(a.get("strides", [1, 1])) == [1, 1] + and list(a.get("pads", [0, 0, 0, 0])) == [0, 0, 0, 0] + and list(a.get("dilations", [1, 1])) == [1, 1] + and a.get("auto_pad", b"NOTSET") in (b"NOTSET", "NOTSET") + and len(node.input) >= 2 + and node.input[1] in inits + ) + + +def _referenced_names(graph) -> set[str]: + names = {o.name for o in graph.output} + + def visit(g) -> None: + for n in g.node: + names.update(n.input) + for attr in n.attribute: + if attr.HasField("g"): + visit(attr.g) + for sub in attr.graphs: + visit(sub) + + visit(graph) + return names + + +def rewrite_model(model, min_in_channels: int = MIN_IN_CHANNELS, drop_zero_bias: bool = True, + dense_as_matmul: bool = True, grouped_as_matmul: bool = True): + """Rewrite ``model`` in place and return a report dictionary.""" + import onnx + from onnx import helper, numpy_helper, shape_inference + + inferred = shape_inference.infer_shapes(model) + shapes: dict[str, list[int]] = {} + for info in list(inferred.graph.value_info) + list(inferred.graph.input) + list(inferred.graph.output): + dims = info.type.tensor_type.shape.dim + shapes[info.name] = [d.dim_value if d.HasField("dim_value") else -1 for d in dims] + + graph = model.graph + inits = {i.name: i for i in graph.initializer} + is_zero = _make_zero_test(model) + used_names = {n.name for n in graph.node if n.name} | set(inits) + + def unique(base: str) -> str: + name, k = base, 0 + while name in used_names: + k += 1 + name = f"{base}_{k}" + used_names.add(name) + return name + + new_nodes, new_inits = [], [] + biases_removed = converted = converted_grouped = 0 + converted_shapes: list[list[int]] = [] + grouped_shapes: list[list[int]] = [] + for node in graph.node: + if node.op_type != "Conv": + new_nodes.append(node) + continue + zero_bias = len(node.input) == 3 and is_zero(node.input[2]) + shape_in = shapes.get(node.input[0]) + groups = int(_attributes(node).get("group", 1)) + enabled = dense_as_matmul if groups == 1 else grouped_as_matmul + if (enabled and _is_pointwise(node, inits) and shape_in is not None + and len(shape_in) == 4 and shape_in[0] == 1 and min(shape_in) > 0 + and (len(node.input) == 2 or zero_bias)): + weight = numpy_helper.to_array(inits[node.input[1]]) + out_ch, group_in = int(weight.shape[0]), int(weight.shape[1]) + in_ch = group_in * groups + if (in_ch >= min_in_channels and weight.shape[2:] == (1, 1) and shape_in[1] == in_ch + and out_ch % groups == 0): + _, _, height, width = shape_in + tag = unique(f"{node.name or node.output[0]}_gemm") + # [O, C/G, 1, 1] -> [G, O/G, C/G]; output channels of group g are contiguous. + w_array = np.ascontiguousarray(weight.reshape(groups, out_ch // groups, group_in)) + if groups == 1: + w_array = w_array[0] + w_init = numpy_helper.from_array(w_array, f"{tag}_w") + flat = [groups, group_in, height * width] if groups > 1 else [1, in_ch, height * width] + s_in = numpy_helper.from_array(np.array(flat, np.int64), f"{tag}_shape_in") + s_out = numpy_helper.from_array(np.array([1, out_ch, height, width], np.int64), f"{tag}_shape_out") + new_inits += [w_init, s_in, s_out] + new_nodes += [ + helper.make_node("Reshape", [node.input[0], s_in.name], [f"{tag}_in"], name=f"{tag}_reshape_in"), + helper.make_node("MatMul", [w_init.name, f"{tag}_in"], [f"{tag}_out"], name=f"{tag}_matmul"), + helper.make_node("Reshape", [f"{tag}_out", s_out.name], [node.output[0]], name=f"{tag}_reshape_out"), + ] + biases_removed += int(zero_bias) + if groups == 1: + converted += 1 + converted_shapes.append([out_ch, in_ch, height, width]) + else: + converted_grouped += 1 + grouped_shapes.append([out_ch, in_ch, groups, height, width]) + continue + if drop_zero_bias and zero_bias: + node = onnx.NodeProto.FromString(node.SerializeToString()) + del node.input[2] + biases_removed += 1 + new_nodes.append(node) + + del graph.node[:] + graph.node.extend(new_nodes) + graph.initializer.extend(new_inits) + + # Prune nodes and initializers that only fed the removed biases / converted weights. + while True: + used = _referenced_names(graph) + live = [n for n in graph.node if any(o in used for o in n.output)] + if len(live) == len(graph.node): + break + del graph.node[:] + graph.node.extend(live) + used = _referenced_names(graph) + keep = [i for i in graph.initializer if i.name in used] + pruned = len(graph.initializer) - len(keep) + del graph.initializer[:] + graph.initializer.extend(keep) + return { + "rewrite_version": REWRITE_VERSION, + "zero_biases_removed": biases_removed, + "dense_1x1_conv_to_matmul": converted, + "grouped_1x1_conv_to_matmul": converted_grouped, + "min_in_channels": min_in_channels, + "converted_shapes_oc_ic_h_w": converted_shapes, + "grouped_shapes_oc_ic_groups_h_w": grouped_shapes, + "initializers_pruned": pruned, + } + + +def _op_counts(model) -> dict[str, int]: + counts: dict[str, int] = {} + for n in model.graph.node: + key = f"{n.domain or 'ai.onnx'}::{n.op_type}" + counts[key] = counts.get(key, 0) + 1 + return dict(sorted(counts.items())) + + +def build_arg_parser() -> argparse.ArgumentParser: + p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + p.add_argument("--bundle", type=Path, required=True, help="Accepted source bundle (model.onnx + manifest.json)") + p.add_argument("--output", type=Path, required=True, help="New bundle directory to create (never overwritten)") + p.add_argument("--min-channels", type=int, default=MIN_IN_CHANNELS) + p.add_argument("--halo", type=int, default=None, + help="Record this tiling halo (and the matching retained core) in the derived " + "manifest; the graph itself is tile-size specific, not halo specific. " + "Default keeps the source manifest's geometry.") + p.add_argument("--fixtures", type=Path, default=None, + help="Optional fixture dir; the rewritten graph must reproduce the source " + "graph's output on the first tile within 1e-5 on CPU before the bundle is written") + return p + + +def main(argv: list[str] | None = None) -> int: + args = build_arg_parser().parse_args(argv) + try: + import onnx + + source = args.bundle.resolve() + output = args.output.resolve() + if output.exists(): + raise ValueError(f"--output already exists (never overwrite): {output}") + manifest = json.loads((source / "manifest.json").read_text()) + source_model = source / MODEL_NAME + source_sha = sha256_file(source_model) + if manifest.get("model_sha256") != source_sha: + raise ValueError("Source manifest model hash does not match model.onnx") + if manifest.get("precision") != "fp32" or manifest.get("diagnostic_shape"): + raise ValueError("Only production FP32 bundles can be rewritten") + + if args.halo is not None: + tile = manifest.get("tile") + if type(tile) is not int or args.halo < 0 or args.halo % 4 or tile <= 2 * args.halo: + raise ValueError(f"--halo must be a non-negative multiple of 4 below half the tile ({tile})") + + model = onnx.load(str(source_model)) + before_ops = _op_counts(model) + report = rewrite_model(model, args.min_channels) + if not (report["dense_1x1_conv_to_matmul"] or report["grouped_1x1_conv_to_matmul"] + or report["zero_biases_removed"]): + raise ValueError("No rewrite applied; refusing to create an identical bundle") + onnx.checker.check_model(model, full_check=True) + + output.parent.mkdir(parents=True, exist_ok=True) + tmp = Path(tempfile.mkdtemp(prefix=".onnx-rewrite-", dir=str(output.parent))) + try: + new_model = tmp / MODEL_NAME + onnx.save(model, str(new_model)) + new_sha = sha256_file(new_model) + check: dict = {"status": "not run (no --fixtures)"} + if args.fixtures is not None: + import onnxruntime as ort + + tiles = sorted((args.fixtures.resolve() / "tiles").glob("tile-*.npz")) + if not tiles: + raise ValueError(f"No fixture tiles in {args.fixtures / 'tiles'}") + with np.load(tiles[0], allow_pickle=False) as f: + x = np.ascontiguousarray(f["input"]) + options = ort.SessionOptions() + options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_DISABLE_ALL + outs = [] + for path in (source_model, new_model): + session = ort.InferenceSession(str(path), options, providers=["CPUExecutionProvider"]) + outs.append(session.run(None, {manifest["input_name"]: x})[0]) + worst = float(np.abs(outs[0].astype(np.float64) - outs[1]).max()) + if worst > 1e-5: + raise ValueError(f"Rewritten graph differs from the source graph by {worst:.3e} (> 1e-5)") + check = {"status": "CPU, first fixture tile, rewritten vs source graph", "max_abs_diff": worst} + summary = { + **report, + "source_model_sha256": source_sha, + "model_sha256": new_sha, + "operators_before": before_ops, + "operators_after": _op_counts(model), + "equivalence_check": check, + "note": "Derived graph; accepted only after onnx_validate passes on every fixture tile " + "(CPU and CUDA) under the frozen gates.", + } + (tmp / REPORT_NAME).write_text(json.dumps(summary, indent=2) + "\n") + derived = dict(manifest) + derived["model_sha256"] = new_sha + if args.halo is not None: + derived["halo"] = args.halo + derived["retained_core"] = manifest["tile"] - 2 * args.halo + derived["graph_rewrites"] = { + "version": REWRITE_VERSION, + "source_model_sha256": source_sha, + "zero_biases_removed": report["zero_biases_removed"], + "dense_1x1_conv_to_matmul": report["dense_1x1_conv_to_matmul"], + "grouped_1x1_conv_to_matmul": report["grouped_1x1_conv_to_matmul"], + "min_in_channels": report["min_in_channels"], + } + derived["validation_report_file"] = REPORT_NAME + derived["validation_report_sha256"] = sha256_file(tmp / REPORT_NAME) + (tmp / "manifest.json").write_text(json.dumps(derived, indent=2) + "\n") + sums = {n: sha256_file(tmp / n) for n in (MODEL_NAME, REPORT_NAME, "manifest.json")} + (tmp / "SHA256SUMS.txt").write_text("".join(f"{h} {n}\n" for n, h in sorted(sums.items()))) + tmp.rename(output) + except BaseException: + import shutil + + shutil.rmtree(tmp, ignore_errors=True) + raise + print(json.dumps({"status": "rewritten", "output": str(output), "model_sha256": new_sha, + "source_model_sha256": source_sha, **{k: report[k] for k in ( + "zero_biases_removed", "dense_1x1_conv_to_matmul", "grouped_1x1_conv_to_matmul")}}, indent=2)) + return 0 + except Exception as exc: + print(json.dumps({"status": "error", "error": str(exc)}), file=sys.stderr) + return 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/denoise/rapidraw_denoise/paired.py b/denoise/rapidraw_denoise/paired.py index 493c5b6ef7..6fbd065c6f 100644 --- a/denoise/rapidraw_denoise/paired.py +++ b/denoise/rapidraw_denoise/paired.py @@ -90,7 +90,7 @@ def main(): size=384;y=int((h-size)*fy)//4*4;x=int((w-size)*fx)//4*4 ref=clean[:,y:y+size,x:x+size].copy() obs=noisy[:,y:y+size,x:x+size].copy() - out,_,info=denoise(obs,model,profile,tile=320,halo=64) + out,info=denoise(obs,model,profile,tile=320,halo=40) normalized=(out-offset)/gain;input_normalized=(obs-offset)/gain record={"scene":entry["scene"],"crop":i,"alignment":alignment, "photometric_gain":gain.ravel().tolist(),"photometric_offset":offset.ravel().tolist(), diff --git a/denoise/rapidraw_denoise/pipeline.py b/denoise/rapidraw_denoise/pipeline.py index b22bf2d274..2a64cdde7d 100644 --- a/denoise/rapidraw_denoise/pipeline.py +++ b/denoise/rapidraw_denoise/pipeline.py @@ -1,9 +1,8 @@ """Backend-neutral NumPy orchestration for Nonlocal RAW inference. -Moved verbatim from inference.py: conditioning, tiling, transforms, Welford -ensemble accumulation, and pilot/row-correction orchestration keep identical -arithmetic, order of operations, clipping locations, reflection padding, and -progress semantics. This module must not import Torch or ONNX Runtime; tensor +Moved verbatim from inference.py: conditioning, tiling and +pilot/row-correction orchestration keep identical arithmetic, order of +operations, clipping locations, reflection padding, and progress semantics. This module must not import Torch or ONNX Runtime; tensor prediction is supplied through the small TilePredictor contract. """ import time @@ -22,7 +21,7 @@ def execution_info(self) -> dict: raise NotImplementedError -def tiled_apply(image, predict, tile=256, halo=64, progress=None): +def tiled_apply(image, predict, tile=256, halo=40, progress=None): if tile % 4 or halo % 4 or halo < 0 or tile <= 2 * halo: raise ValueError("Tile and halo must be multiples of four, with tile > 2*halo") _, h, w = image.shape @@ -44,19 +43,9 @@ def tiled_apply(image, predict, tile=256, halo=64, progress=None): return out -def transform(x, i): - x = np.rot90(x, i % 4, axes=(-2,-1)) - return x[..., ::-1] if i >= 4 else x - - -def inverse_transform(x, i): - x = x[..., ::-1] if i >= 4 else x - return np.rot90(x, -(i % 4), axes=(-2,-1)) - - -def denoise_with_predictor(packed, predictor, profile=None, tile=256, halo=64, ensemble=1, noise_scale=1., pilot=False, progress=None, row_correction=False): - if ensemble not in (1, 4, 8) or not np.isfinite(noise_scale) or noise_scale <= 0: - raise ValueError("Ensemble must be 1, 4, or 8; noise scale must be positive") +def denoise_with_predictor(packed, predictor, profile=None, tile=256, halo=40, noise_scale=1., pilot=False, progress=None, row_correction=False): + if not np.isfinite(noise_scale) or noise_scale <= 0: + raise ValueError("Noise scale must be positive") start=time.monotonic() profile = profile or estimate_noise(packed) image=np.clip(packed,0,1) @@ -86,23 +75,12 @@ def predict(tile_array): image=np.clip(packed-row,0,1) row_info={"rms":float(np.sqrt(np.mean(row**2))),"shrinkage":shrink.ravel().tolist()} conditioned=np.concatenate([image,np.sqrt(np.maximum(variance,1e-12))*noise_scale]) - mean=np.zeros_like(image); m2=np.zeros_like(image) - for i in range(ensemble): - tile_progress = (lambda done, total: progress({"completed_passes":i,"total_passes":ensemble, - "completed_tiles":done,"total_tiles":total})) if progress else None - candidate=inverse_transform(tiled_apply(transform(conditioned,i),predict,tile,halo,tile_progress),i) - delta=candidate-mean - mean+=delta/(i+1) - m2+=delta*(candidate-mean) - if progress: - progress({"completed_passes":i+1,"total_passes":ensemble}) - disagreement=m2/max(ensemble-1,1) - if not np.isfinite(mean).all(): + tile_progress = (lambda done, total: progress({"completed_tiles":done,"total_tiles":total})) if progress else None + prediction=tiled_apply(conditioned,predict,tile,halo,tile_progress) + if not np.isfinite(prediction).all(): raise RuntimeError("Nonfinite inference result") - info={"elapsed_seconds":time.monotonic()-start,"tile":tile,"halo":halo,"ensemble":ensemble, + info={"elapsed_seconds":time.monotonic()-start,"tile":tile,"halo":halo, "pilot_conditioning":pilot,"row_correction":row_info,"noise_scale":noise_scale,"noise_profile":profile.to_dict(), - "parameters":None, - "disagreement_mean_variance":float(disagreement.mean()), - "disagreement_interpretation":"Transform sensitivity, not calibrated epistemic uncertainty"} + "parameters":None} info.update(predictor.execution_info()) - return mean,disagreement,info + return prediction,info diff --git a/denoise/rapidraw_denoise/study/candidates.py b/denoise/rapidraw_denoise/study/candidates.py index e7f9acaae9..e7336b8178 100644 --- a/denoise/rapidraw_denoise/study/candidates.py +++ b/denoise/rapidraw_denoise/study/candidates.py @@ -4,7 +4,8 @@ import json, time import numpy as np import torch -from rapidraw_denoise.inference import load_model, transform, inverse_transform +from rapidraw_denoise.inference import load_model +from rapidraw_denoise.study.transforms import transform, inverse_transform from rapidraw_denoise.raw import sha256 from .evaluate import extra_metrics from .common import ROOT as R, STUDY as S diff --git a/denoise/rapidraw_denoise/study/evaluate.py b/denoise/rapidraw_denoise/study/evaluate.py index 0d9eac7a58..de394fe603 100644 --- a/denoise/rapidraw_denoise/study/evaluate.py +++ b/denoise/rapidraw_denoise/study/evaluate.py @@ -7,7 +7,8 @@ import torch from scipy.ndimage import gaussian_filter from rapidraw_denoise.raw import sha256 -from rapidraw_denoise.inference import load_model, transform, inverse_transform +from rapidraw_denoise.inference import load_model +from rapidraw_denoise.study.transforms import transform, inverse_transform from rapidraw_denoise.benchmark import metrics from .common import ROOT, STUDY diff --git a/denoise/rapidraw_denoise/study/phase.py b/denoise/rapidraw_denoise/study/phase.py index a919a98e9f..81e2d21bd7 100644 --- a/denoise/rapidraw_denoise/study/phase.py +++ b/denoise/rapidraw_denoise/study/phase.py @@ -3,7 +3,8 @@ require_space() import json, time, numpy as np, torch from .evaluate import extra_metrics -from rapidraw_denoise.inference import load_model, transform, inverse_transform +from rapidraw_denoise.inference import load_model +from rapidraw_denoise.study.transforms import transform, inverse_transform from .common import ROOT, STUDY as S out = S / "validation-phase" diff --git a/denoise/rapidraw_denoise/study/test.py b/denoise/rapidraw_denoise/study/test.py index bf87db2895..f6258bbf68 100644 --- a/denoise/rapidraw_denoise/study/test.py +++ b/denoise/rapidraw_denoise/study/test.py @@ -3,7 +3,8 @@ require_space() from pathlib import Path import json, time, numpy as np, torch -from rapidraw_denoise.inference import load_model, transform, inverse_transform +from rapidraw_denoise.inference import load_model +from rapidraw_denoise.study.transforms import transform, inverse_transform from rapidraw_denoise.raw import sha256 from .evaluate import extra_metrics from .common import ROOT as R, STUDY as S diff --git a/denoise/rapidraw_denoise/study/transforms.py b/denoise/rapidraw_denoise/study/transforms.py new file mode 100644 index 0000000000..5f4be1b198 --- /dev/null +++ b/denoise/rapidraw_denoise/study/transforms.py @@ -0,0 +1,12 @@ +"""Rotation/flip test-time transforms used by the checkpoint-selection study.""" +import numpy as np + + +def transform(x, i): + x = np.rot90(x, i % 4, axes=(-2,-1)) + return x[..., ::-1] if i >= 4 else x + + +def inverse_transform(x, i): + x = x[..., ::-1] if i >= 4 else x + return np.rot90(x, -(i % 4), axes=(-2,-1)) diff --git a/denoise/rapidraw_denoise/worker.py b/denoise/rapidraw_denoise/worker.py index a28445721c..dd906c9038 100644 --- a/denoise/rapidraw_denoise/worker.py +++ b/denoise/rapidraw_denoise/worker.py @@ -21,7 +21,7 @@ def validate_request(request): raise ValueError("Invalid packed RGBG shape") if request.get("model_sha256") != CHECKPOINT_SHA256: raise ValueError("Unrecognized Nonlocal checkpoint") - if type(request.get("ensemble")) is not int or request["ensemble"] not in (1, 4) or request.get("tile") != 320 or request.get("halo") != 64: + if request.get("tile") != 320 or request.get("halo") != 40: raise ValueError("Unsupported native inference configuration") return tuple(shape) @@ -46,13 +46,11 @@ def run(directory, checkpoint): from .inference import denoise, load_model torch.set_num_threads(4) def progress(event): - fraction = (event["completed_passes"] + event.get("completed_tiles", 0) - / max(event.get("total_tiles", 1), 1)) / event["total_passes"] + fraction = event["completed_tiles"] / max(event["total_tiles"], 1) print(json.dumps({"progress": min(fraction, 1), "stage": "Nonlocal RAW tiles"}), flush=True) print(json.dumps({"progress": 0, "stage": "Loading Nonlocal CUDA model"}), flush=True) model = load_model(checkpoint, "cuda") - output, _, info = denoise(packed, model, tile=320, halo=64, - ensemble=request["ensemble"], progress=progress) + output, info = denoise(packed, model, tile=320, halo=40, progress=progress) if output.shape != shape or not np.isfinite(output).all(): raise ValueError("Invalid RAW prediction") temporary = directory / "prediction.f32.partial" diff --git a/denoise/scripts/bench_candidate.py b/denoise/scripts/bench_candidate.py index 7c4b345f35..fe37a75c5c 100644 --- a/denoise/scripts/bench_candidate.py +++ b/denoise/scripts/bench_candidate.py @@ -2,8 +2,7 @@ Benchmarks one ONNX bundle with NumPy-in/NumPy-out timing only: session init, 3 warm-up + 10 measured tile runs (np.median / np.quantile-linear -P95), ensemble-1 xN full assemblies + one ensemble-4 on the representative -photo. Imports numpy/scipy/onnxruntime only; asserts Torch was never +P95), then xN full-photo assemblies on the representative photo. Imports numpy/scipy/onnxruntime only; asserts Torch was never imported. One process per configuration; env fixed at launch. Optional benchmark-only ORT profiling (--profile) records a per-node @@ -15,8 +14,8 @@ opt-in and off by default; it must not change inference numerics. Usage: bench_candidate.py --bundle DIR --fixtures DIR --out PATH - [--provider cuda] [--optimize] [--photo portrait] [--repeats-e1 3] - [--skip-e4] [--profile] [--profile-dir DIR] + [--provider cuda] [--optimize] [--photo portrait] [--repeats 3] + [--profile] [--profile-dir DIR] [--prod-cuda] [--cuda-arena ...] [--cuda-cudnn-search ...] [--cuda-device-id 0] [--maps-out PATH] """ @@ -98,8 +97,7 @@ def main(): help="FAST_EXPERIMENTAL only: CUDA TF32 on (labels the run " "experimental; never a strict pass)") ap.add_argument("--photo", default="portrait") - ap.add_argument("--repeats-e1", type=int, default=3) - ap.add_argument("--skip-e4", action="store_true") + ap.add_argument("--repeats", type=int, default=3) ap.add_argument("--torch-prime", action="store_true", help="Diagnostic only: import torch and touch the CUDA " "context/allocator before session creation (mimics the " @@ -318,21 +316,14 @@ def main(): f"nodes={audit['node_events']} " f"by_provider={audit['node_events_by_provider']}", flush=True) - rep["full_e1_wall_s_runs"] = [] - for r in range(ns.repeats_e1): + rep["full_wall_s_runs"] = [] + for r in range(ns.repeats): t1 = time.monotonic() - out, _, _ = denoise_with_predictor(photo, pred, profile, 320, 64, 1) + out, _ = denoise_with_predictor(photo, pred, profile, 320, 40) wall = time.monotonic() - t1 assert out.shape == photo.shape and np.isfinite(out).all() - rep["full_e1_wall_s_runs"].append(wall) - print(f"e1 run{r}: wall={wall:.1f}s", flush=True) - if not ns.skip_e4: - t1 = time.monotonic() - out, _, _ = denoise_with_predictor(photo, pred, profile, 320, 64, 4) - wall = time.monotonic() - t1 - assert out.shape == photo.shape and np.isfinite(out).all() - rep["full_e4_wall_s"] = wall - print(f"e4: wall={wall:.1f}s", flush=True) + rep["full_wall_s_runs"].append(wall) + print(f"run{r}: wall={wall:.1f}s", flush=True) rep["torch_imported_at_end"] = "torch" in sys.modules if not ns.torch_prime: assert "torch" not in sys.modules, "torch must not be imported on this path" diff --git a/denoise/scripts/bench_v2.py b/denoise/scripts/bench_v2.py index 07ac92569f..2ac90f4a59 100644 --- a/denoise/scripts/bench_v2.py +++ b/denoise/scripts/bench_v2.py @@ -3,7 +3,7 @@ Usage: bench_v2.py {onnx-cuda,torch-native,torch-native-tf32off,torch-ref, onnx-strict,torch-default,torch-orig,torch-strict, torch-packed,torch-refstrict} - [--repeats-e1 N] [--skip-e4] + [--repeats N] [--bundle DIR] [--fixtures DIR] [--photo NAME] [--ckpt PATH] [--out PATH] @@ -12,7 +12,7 @@ hard-coded paths grew optional CLI overrides with identical defaults. New reproducible modes (one fresh process per mode, same portrait -[4,2327,3491] workload, same 320/64 e1 path): +[4,2327,3491] workload, same 320/40 path): - onnx-strict: ONNX production provider policy (packed bundle, graph opt off, device 0, SameAsRequested, HEURISTIC, TF32 off, no fallback) in a TORCH-PRELOADED library context (torch is imported before ORT session @@ -31,8 +31,8 @@ in-process (both vendor.sampling and vendor.network entry points), both TF32 paths off. Never mutates source/default runtime behavior. - torch-refstrict (T_REF): reference sampler rebound explicitly - in-process, both TF32 paths off. Tile section always runs; the single e1 - is skipped with a labeled reason when the tile median is prohibitive. + in-process, both TF32 paths off. Tile section always runs; the full + photo run is skipped with a labeled reason when the tile median is prohibitive. Sampler identity is proven behaviorally (native must differ from the reference loop bitwise; reference must reproduce it bitwise; packed reports @@ -49,7 +49,7 @@ native binding must differ from it; missing ext fails loud). Saves ordered warm-up + measured samples (median = np.median; P95 = np.quantile linear, both documented), session init separately, full-assembly NumPy-in -> -NumPy-out walls (3x ensemble-1 plus one ensemble-4 unless skipped), and +NumPy-out walls (3 repeats by default), and scoped memory notes (Torch allocator counters on torch runs only; ORT VRAM recorded null, never as Torch counters). """ @@ -87,7 +87,7 @@ "torch-refstrict": "reference", } -# torch-refstrict skips its single e1 above this tile-median threshold. +# torch-refstrict skips its full-photo run above this tile-median threshold. REFSTRICT_TILE_SKIP_MS = 1500.0 @@ -154,8 +154,7 @@ def expected_probe_equal(binding): def main(): ap = argparse.ArgumentParser() ap.add_argument("backend", choices=LEGACY_BACKENDS + NEW_BACKENDS) - ap.add_argument("--repeats-e1", type=int, default=3) - ap.add_argument("--skip-e4", action="store_true") + ap.add_argument("--repeats", type=int, default=3) ap.add_argument("--bundle", default=None, help="ONNX bundle dir (default: packed for onnx-strict, " "pinned legacy otherwise)") @@ -168,9 +167,9 @@ def main(): help="Report path (default: /tmp/nlx-bench-v2-BACKEND.json)") ns = ap.parse_args() backend = ns.backend - if backend in NEW_BACKENDS and ns.repeats_e1 < 0: + if backend in NEW_BACKENDS and ns.repeats < 0: print(json.dumps({"status": "error", - "error": "--repeats-e1 must be >= 0"}), + "error": "--repeats must be >= 0"}), file=sys.stderr) return 1 @@ -309,21 +308,14 @@ def counting_ref(*a, **k): rep["torch_vram_reserved_B"] = torch.cuda.max_memory_reserved() rep["rss_mb"] = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 1024 - rep["full_e1_wall_s_runs"] = [] - for r in range(ns.repeats_e1): + rep["full_wall_s_runs"] = [] + for r in range(ns.repeats): t1 = time.monotonic() - out, _, _ = denoise_with_predictor(portrait, pred, profile, 320, 64, 1) + out, _ = denoise_with_predictor(portrait, pred, profile, 320, 40) wall = time.monotonic() - t1 assert out.shape == portrait.shape and np.isfinite(out).all() - rep["full_e1_wall_s_runs"].append(wall) - print(f"{backend} e1 run{r}: wall={wall:.1f}s", flush=True) - if not ns.skip_e4: - t1 = time.monotonic() - out, _, _ = denoise_with_predictor(portrait, pred, profile, 320, 64, 4) - wall = time.monotonic() - t1 - assert out.shape == portrait.shape and np.isfinite(out).all() - rep["full_e4_wall_s"] = wall - print(f"{backend} e4: wall={wall:.1f}s", flush=True) + rep["full_wall_s_runs"].append(wall) + print(f"{backend} run{r}: wall={wall:.1f}s", flush=True) rep["smi_after"] = smi() with open(out_name, "w") as f: json.dump(rep, f, indent=2) @@ -554,21 +546,21 @@ def counting(*a, **k): rep["torch_vram_reserved_B"] = torch.cuda.max_memory_reserved() rep["rss_mb"] = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 1024 - rep["full_e1_wall_s_runs"] = [] + rep["full_wall_s_runs"] = [] if backend == "torch-refstrict" and float(np.median(meas)) > \ REFSTRICT_TILE_SKIP_MS: - rep["full_e1_skipped"] = ( + rep["full_skipped"] = ( f"tile median {float(np.median(meas)):.0f} ms exceeds " f"{REFSTRICT_TILE_SKIP_MS:.0f} ms; tile-only comparison") else: - for r in range(ns.repeats_e1): + for r in range(ns.repeats): t1 = time.monotonic() - out, _, _ = denoise_with_predictor(photo, pred, profile, - 320, 64, 1) + out, _ = denoise_with_predictor(photo, pred, profile, + 320, 40) wall = time.monotonic() - t1 assert out.shape == photo.shape and np.isfinite(out).all() - rep["full_e1_wall_s_runs"].append(wall) - print(f"{backend} e1 run{r}: wall={wall:.1f}s", flush=True) + rep["full_wall_s_runs"].append(wall) + print(f"{backend} run{r}: wall={wall:.1f}s", flush=True) if backend != "onnx-strict": network_mod.deform_neighbourhood = orig_network sampling_mod.deform_neighbourhood = orig_sampling diff --git a/denoise/scripts/compare_controlled_full.py b/denoise/scripts/compare_controlled_full.py index a00f6788a4..1acf41dcd8 100644 --- a/denoise/scripts/compare_controlled_full.py +++ b/denoise/scripts/compare_controlled_full.py @@ -1,11 +1,10 @@ """Phase-C closeout driver: ONNX-vs-controlled full-photo comparison. Compares assembled ONNX full predictions against controlled-fp32 reference -full predictions for the six required photo/ensemble cases -(portrait/landscape/phone x e1/e4) with the repaired comparator, frozen -gates, exact expected case identities, and ensemble-4 diagnostic seam masks. -Saved arrays are re-measured (no inference rerun). Exit codes: 0 = 6/6 -pass, 2 = gate failures (report still written), 1 = tool/invalid error +full predictions for the three required photo cases +(portrait/landscape/phone) with the repaired comparator, frozen gates and +exact expected case identities. Saved arrays are re-measured (no inference +rerun). Exit codes: 0 = 3/3 pass, 2 = gate failures (report still written), 1 = tool/invalid error (such as a missing case or a job that did not actually finish). Usage: compare_controlled_full.py [--fixtures DIR] [--onnx-dir DIR] @@ -21,16 +20,15 @@ from rapidraw_denoise.compare_full import compare_files # noqa: E402 PHOTOS = ("portrait", "landscape", "phone") -ENSEMBLES = (1, 4) -EXPECTED = [f"{p}-e{e}" for p in PHOTOS for e in ENSEMBLES] +EXPECTED = list(PHOTOS) def main(argv=None): ap = argparse.ArgumentParser(description=__doc__) ap.add_argument("--fixtures", default="/tmp/nlx-fixtures", - help="Dir with -controlled-full/full-ensemble-.npy") + help="Dir with -controlled-full/full-prediction.npy") ap.add_argument("--onnx-dir", default="/tmp/nlx-fixtures", - help="Dir with --full-e.npy") + help="Dir with --full.npy") ap.add_argument("--tag", default="onnx") ap.add_argument("--report", default="/tmp/nlx-controlled-full-report.json") args = ap.parse_args(argv) @@ -43,20 +41,19 @@ def main(argv=None): try: out = [] for photo in PHOTOS: - for ensemble in ENSEMBLES: - name = f"{photo}-e{ensemble}" - ref = (Path(args.fixtures) / f"{photo}-controlled-full" - / f"full-ensemble-{ensemble}.npy") - act = (Path(args.onnx_dir) / f"{photo}-{args.tag}-full-e{ensemble}.npy") - for p, what in ((ref, "controlled reference"), (act, "onnx assembly")): - if not p.is_file() or p.stat().st_size == 0: - raise FileNotFoundError( - f"{name}: {what} missing or empty: {p} " - "(capture may still be running; not a pass)") - out.append(compare_files(name, act, ref, ensemble=ensemble)) + name = photo + ref = (Path(args.fixtures) / f"{photo}-controlled-full" + / "full-prediction.npy") + act = Path(args.onnx_dir) / f"{photo}-{args.tag}-full.npy" + for p, what in ((ref, "controlled reference"), (act, "onnx assembly")): + if not p.is_file() or p.stat().st_size == 0: + raise FileNotFoundError( + f"{name}: {what} missing or empty: {p} " + "(capture may still be running; not a pass)") + out.append(compare_files(name, act, ref)) names = [e["comparison"] for e in out] if sorted(names) != sorted(EXPECTED): - raise ValueError(f"case identities != expected six: {names}") + raise ValueError(f"case identities != expected three: {names}") report = {"status": "pass" if all(e["pass"] for e in out) else "fail", "expected_cases": EXPECTED, "comparisons": out} diff --git a/denoise/scripts/compare_production_full.py b/denoise/scripts/compare_production_full.py index 534ef7c361..adb79ff716 100644 --- a/denoise/scripts/compare_production_full.py +++ b/denoise/scripts/compare_production_full.py @@ -1,9 +1,8 @@ """Production-policy recompute: refactored vs original production full-photo. -Re-measures the six saved production full predictions with the hardened +Re-measures the three saved production full predictions with the hardened comparator (no inference rerun): refactored-assembly vs original-CUDA -production reference, frozen gates, exact expected identities, ensemble-4 -diagnostics. Production TF32 policy failures are preserved as failures; +production reference, frozen gates and exact expected identities. Production TF32 policy failures are preserved as failures; this driver never loosens gates or promotes defaults. Usage: compare_production_full.py [--fixtures DIR] [--report PATH] @@ -18,8 +17,7 @@ from rapidraw_denoise.compare_full import compare_files # noqa: E402 PHOTOS = ("portrait", "landscape", "phone") -ENSEMBLES = (1, 4) -EXPECTED = [f"{p}-newprod-e{e}" for p in PHOTOS for e in ENSEMBLES] +EXPECTED = [f"{p}-newprod" for p in PHOTOS] def main(argv=None): @@ -36,18 +34,17 @@ def main(argv=None): try: out = [] for photo in PHOTOS: - for ensemble in ENSEMBLES: - name = f"{photo}-newprod-e{ensemble}" - act = (Path(args.fixtures) / f"{photo}-newprod-full-e{ensemble}.npy") - ref = (Path(args.fixtures) / f"{photo}-production-cuda" - / f"full-ensemble-{ensemble}.npy") - for p, what in ((act, "refactored assembly"), (ref, "production reference")): - if not p.is_file() or p.stat().st_size == 0: - raise FileNotFoundError(f"{name}: {what} missing: {p}") - out.append(compare_files(name, act, ref, ensemble=ensemble)) + name = f"{photo}-newprod" + act = Path(args.fixtures) / f"{photo}-newprod-full.npy" + ref = (Path(args.fixtures) / f"{photo}-production-cuda" + / "full-prediction.npy") + for p, what in ((act, "refactored assembly"), (ref, "production reference")): + if not p.is_file() or p.stat().st_size == 0: + raise FileNotFoundError(f"{name}: {what} missing: {p}") + out.append(compare_files(name, act, ref)) names = [e["comparison"] for e in out] if sorted(names) != sorted(EXPECTED): - raise ValueError(f"case identities != expected six: {names}") + raise ValueError(f"case identities != expected three: {names}") report = {"status": "pass" if all(e["pass"] for e in out) else "fail", "expected_cases": EXPECTED, "note": "production-policy comparison (CUDA TF32 production " diff --git a/denoise/scripts/coreml_fullphoto.py b/denoise/scripts/coreml_fullphoto.py index 4aa7a5a002..e01f9455ea 100644 --- a/denoise/scripts/coreml_fullphoto.py +++ b/denoise/scripts/coreml_fullphoto.py @@ -4,10 +4,10 @@ denoise_with_predictor pipeline (same TILE/HALO/CORE contract as the ONNX full-photo runs), saves the output, and scores it against a reference full-photo array with the frozen full_metrics gates via compare_files. -One photo/ensemble per invocation; all paths are argv (no hardcodes). +One photo per invocation; all paths are argv (no hardcodes). Usage: coreml_fullphoto.py --package DIR --native-dir DIR --fixtures-root DIR - --photo portrait --ensemble 1 --reference REF.npy --out OUT.npy [--report JSON] + --photo portrait --reference REF.npy --out OUT.npy [--report JSON] """ import argparse import json @@ -15,7 +15,7 @@ import time from pathlib import Path -TILE, HALO = 320, 64 +TILE, HALO = 320, 40 def main(): @@ -24,7 +24,6 @@ def main(): ap.add_argument("--native-dir", required=True) ap.add_argument("--fixtures-root", required=True) ap.add_argument("--photo", required=True) - ap.add_argument("--ensemble", type=int, choices=(1, 4), default=1) ap.add_argument("--reference", required=True) ap.add_argument("--out", required=True) ap.add_argument("--report", default=None) @@ -76,13 +75,11 @@ def counting(tile): pred.predict = counting t1 = time.monotonic() - out, _, _ = denoise_with_predictor(packed, pred, profile, TILE, HALO, - ns.ensemble) + out, _ = denoise_with_predictor(packed, pred, profile, TILE, HALO) wall = time.monotonic() - t1 np.save(ns.out, out) - cmp_rec = compare_files(f"coreml-{ns.photo}-e{ns.ensemble}", ns.out, - ns.reference, ensemble=ns.ensemble) - rep = {"photo": ns.photo, "ensemble": ns.ensemble, + cmp_rec = compare_files(f"coreml-{ns.photo}", ns.out, ns.reference) + rep = {"photo": ns.photo, "package_sha256": manifest["package_sha256"], "load_s": load_s, "wall_s": wall, "tile_calls": calls[0], "output_shape": list(out.shape), "comparison": cmp_rec} @@ -91,7 +88,7 @@ def counting(tile): m = {k: cmp_rec.get(k) for k in ("pass", "max_abs_err", "mae", "p99_abs_err", "elementwise_violations", "consistency_problems")} - print(json.dumps({"photo": ns.photo, "ensemble": ns.ensemble, + print(json.dumps({"photo": ns.photo, "wall_s": round(wall, 1), "tiles": calls[0], "pass": m.get("pass"), "max_abs_err": m.get("max_abs_err")}, indent=2)) diff --git a/denoise/scripts/fullphoto_newprod.py b/denoise/scripts/fullphoto_newprod.py index e650cc71fe..a833911038 100644 --- a/denoise/scripts/fullphoto_newprod.py +++ b/denoise/scripts/fullphoto_newprod.py @@ -24,24 +24,22 @@ open(f"/tmp/nlx-fixtures/{photo}-controlled-fp32/metadata.json") ) profile = NoiseProfile(**meta["noise_profile"]) - pentry = {"photoshape": list(packed.shape), "passes": {}} - for ensemble in (1, 4): - events = [] - t1 = time.monotonic() - out, _, info = denoise( - packed, model, profile, 320, 64, ensemble, - progress=events.append, - ) - wall = time.monotonic() - t1 - tile_events = [e for e in events if "completed_tiles" in e] - np.save(f"/tmp/nlx-fixtures/{photo}-newprod-full-e{ensemble}.npy", out) - pentry["passes"][str(ensemble)] = { - "wall_seconds": wall, - "tile_events": len(tile_events), - "inference_elapsed": info["elapsed_seconds"], - } - print(f"{photo} e{ensemble}: {wall:.1f}s tile_events={len(tile_events)}", - flush=True) + events = [] + t1 = time.monotonic() + out, info = denoise( + packed, model, profile, 320, 40, + progress=events.append, + ) + wall = time.monotonic() - t1 + tile_events = [e for e in events if "completed_tiles" in e] + np.save(f"/tmp/nlx-fixtures/{photo}-newprod-full.npy", out) + pentry = { + "photoshape": list(packed.shape), + "wall_seconds": wall, + "tile_events": len(tile_events), + "inference_elapsed": info["elapsed_seconds"], + } + print(f"{photo}: {wall:.1f}s tile_events={len(tile_events)}", flush=True) report["photos"][photo] = pentry with open("/tmp/nlx-newprod-full-report.json", "w") as f: json.dump(report, f, indent=2) diff --git a/denoise/scripts/fullphoto_onnx.py b/denoise/scripts/fullphoto_onnx.py index 6d211448bd..079141d639 100644 --- a/denoise/scripts/fullphoto_onnx.py +++ b/denoise/scripts/fullphoto_onnx.py @@ -16,7 +16,7 @@ from rapidraw_denoise.onnx_backend import OnnxTilePredictor from rapidraw_denoise.pipeline import denoise_with_predictor -TILE, HALO, CORE = 320, 64, 192 +TILE, HALO, CORE = 320, 40, 240 PHOTOS = ["portrait", "landscape", "phone"] @@ -36,7 +36,7 @@ def main(argv=None): ap = argparse.ArgumentParser() ap.add_argument("--bundle", default="/tmp/nlx-model/bundle-pinned-320-fp32") ap.add_argument("--tag", default="onnx", - help="output filename tag: {photo}-{tag}-full-e{1,4}.npy") + help="output filename tag: {photo}-{tag}-full.npy") ap.add_argument("--optimize", action="store_true") ap.add_argument("--report", default=None) ns = ap.parse_args(argv) @@ -62,39 +62,34 @@ def main(argv=None): nx = (w + CORE - 1) // CORE pentry = { "packed_shape": [4, h, w], - "expected_tiles_per_pass": ny * nx, + "expected_tiles": ny * nx, "source_sha256": meta["source_sha256"], - "passes": {}, } - for ensemble in (1, 4): - calls = [0] + calls = [0] - orig_predict = predictor.predict + orig_predict = predictor.predict - def counting_predict(tile, _o=orig_predict, _c=calls): - _c[0] += 1 - return _o(tile) + def counting_predict(tile, _o=orig_predict, _c=calls): + _c[0] += 1 + return _o(tile) - predictor.predict = counting_predict - t1 = time.monotonic() - out, _, info = denoise_with_predictor( - packed, predictor, profile, TILE, HALO, ensemble - ) - wall = time.monotonic() - t1 - predictor.predict = orig_predict - np.save(f"/tmp/nlx-fixtures/{photo}-{ns.tag}-full-e{ensemble}.npy", out) - pentry["passes"][str(ensemble)] = { - "wall_seconds": wall, - "actual_tile_calls": calls[0], - "expected_tile_calls": ny * nx * ensemble, - "output_shape": list(out.shape), - } - assert calls[0] == ny * nx * ensemble, (calls[0], ny * nx * ensemble) - assert out.shape == packed.shape - print( - f"{photo} e{ensemble}: {wall:.1f}s tiles={calls[0]}", - flush=True, - ) + predictor.predict = counting_predict + t1 = time.monotonic() + out, info = denoise_with_predictor( + packed, predictor, profile, TILE, HALO + ) + wall = time.monotonic() - t1 + predictor.predict = orig_predict + np.save(f"/tmp/nlx-fixtures/{photo}-{ns.tag}-full.npy", out) + pentry.update({ + "wall_seconds": wall, + "actual_tile_calls": calls[0], + "expected_tile_calls": ny * nx, + "output_shape": list(out.shape), + }) + assert calls[0] == ny * nx, (calls[0], ny * nx) + assert out.shape == packed.shape + print(f"{photo}: {wall:.1f}s tiles={calls[0]}", flush=True) report["photos"][photo] = pentry report_path = ns.report or f"/tmp/nlx-{ns.tag}-full-report.json" with open(report_path, "w") as f: diff --git a/denoise/scripts/host_refactor_check.py b/denoise/scripts/host_refactor_check.py index 1764bd7a6f..d5115f9189 100644 --- a/denoise/scripts/host_refactor_check.py +++ b/denoise/scripts/host_refactor_check.py @@ -24,11 +24,11 @@ fm = frozen.load_model(CKPT, device="cpu") nm = new.load_model(CKPT, device="cpu") -fo, _, fi = frozen.denoise(packed, fm, profile, tile=64, halo=16, ensemble=4) -no, _, ni = new.denoise(packed, nm, profile, tile=64, halo=16, ensemble=4) +fo, _, fi = frozen.denoise(packed, fm, profile, tile=64, halo=16) +no, ni = new.denoise(packed, nm, profile, tile=64, halo=16) print("bitwise equal:", np.array_equal(no, fo)) print("max abs diff:", float(np.abs(no.astype("float64") - fo.astype("float64")).max())) -for key in ("tile", "halo", "ensemble", "noise_scale", "parameters", "device"): +for key in ("tile", "halo", "noise_scale", "parameters", "device"): assert ni[key] == fi[key], key assert ni["noise_profile"] == fi["noise_profile"] print("info keys match") diff --git a/denoise/tests/test_backend_parity.py b/denoise/tests/test_backend_parity.py index 619581ee2e..06fbfea865 100644 --- a/denoise/tests/test_backend_parity.py +++ b/denoise/tests/test_backend_parity.py @@ -59,14 +59,13 @@ def test_refactored_torch_matches_frozen_original(): profile = _profile() frozen_model = frozen.load_model(LOCAL_CKPT, device="cpu") new_model = new_inference.load_model(LOCAL_CKPT, device="cpu") - frozen_out, frozen_dis, frozen_info = frozen.denoise( - packed, frozen_model, profile, tile=64, halo=16, ensemble=4) - new_out, new_dis, new_info = new_inference.denoise( - packed, new_model, profile, tile=64, halo=16, ensemble=4) + frozen_out, _, frozen_info = frozen.denoise( + packed, frozen_model, profile, tile=64, halo=16) + new_out, new_info = new_inference.denoise( + packed, new_model, profile, tile=64, halo=16) assert np.array_equal(new_out, frozen_out), float( np.abs(new_out - frozen_out).max()) - assert np.array_equal(new_dis, frozen_dis) - for key in ("tile", "halo", "ensemble", "noise_scale", "parameters"): + for key in ("tile", "halo", "noise_scale", "parameters"): assert new_info[key] == frozen_info[key], key assert new_info["noise_profile"] == frozen_info["noise_profile"] assert new_info["device"] == frozen_info["device"] == "cpu" @@ -75,13 +74,13 @@ def test_refactored_torch_matches_frozen_original(): @requires_ckpt def test_wrappers_preserve_helper_defaults_and_contracts(): import inspect - for fn in ("tiled_apply", "transform", "inverse_transform"): + for fn in ("tiled_apply",): assert getattr(new_inference, fn) is getattr( __import__("rapidraw_denoise.pipeline", fromlist=[fn]), fn) sig = inspect.signature(new_inference.denoise) assert sig.parameters["tile"].default == 256 - assert sig.parameters["halo"].default == 64 - assert sig.parameters["ensemble"].default == 1 + assert sig.parameters["halo"].default == 40 + assert "ensemble" not in sig.parameters @requires_bundle diff --git a/denoise/tests/test_bench_v2_modes.py b/denoise/tests/test_bench_v2_modes.py index 562dbc7c6e..993f84a568 100644 --- a/denoise/tests/test_bench_v2_modes.py +++ b/denoise/tests/test_bench_v2_modes.py @@ -85,7 +85,7 @@ def test_cli_lists_new_modes_and_overrides(): "torch-strict", "torch-packed", "torch-refstrict"): assert mode in proc.stdout for flag in ("--bundle", "--fixtures", "--photo", "--ckpt", "--out", - "--repeats-e1", "--skip-e4"): + "--repeats"): assert flag in proc.stdout diff --git a/denoise/tests/test_capture_reference.py b/denoise/tests/test_capture_reference.py index dac0564b1e..53d0d623f4 100644 --- a/denoise/tests/test_capture_reference.py +++ b/denoise/tests/test_capture_reference.py @@ -36,8 +36,7 @@ def _write_native_dir(tmp_path: Path, h: int = 64, w: int = 64) -> Path: "source_sha256": "source", "shape": [4, h, w], "tile": 320, - "halo": 64, - "ensemble": 1, + "halo": 40, } (worker_dir / "request.json").write_text(json.dumps(request)) return worker_dir @@ -139,7 +138,7 @@ def test_native_transport_roundtrip_is_little_endian(tmp_path): raw = (worker_dir / "input.f32").read_bytes() request = json.loads((worker_dir / "request.json").read_text()) assert request["shape"] == [4, 16, 32] - assert request["tile"] == 320 and request["halo"] == 64 + assert request["tile"] == 320 and request["halo"] == 40 arr = np.fromfile(worker_dir / "input.f32", dtype=" 0 and masks["seam"].sum() > 0 assert masks["interior"].sum() > 0 @@ -490,27 +490,6 @@ def test_region_masks_grid_lines_are_exact(): assert total == 400 * 400 # disjoint complete partition -def test_e4_diagnostic_supersets_unrotated_seam_and_labels_extra(tmp_path): - from rapidraw_denoise.compare_full import compare_files, region_masks, region_masks_e4 - rng = np.random.default_rng(7) - ref = rng.normal(size=(4, 400, 401)).astype(np.float32) - out = (ref.astype("float64") + 1e-7).astype(np.float32) - a, r = tmp_path / "a.npy", tmp_path / "r.npy" - np.save(a, out) - np.save(r, ref) - e1 = compare_files("e1", a, r, ensemble=1) - e4 = compare_files("e4", a, r, ensemble=4) - assert "regions_e4" not in e1 - assert set(e4["regions_e4"]) == {"outer", "seam", "interior"} - m1, m4 = region_masks(400, 401), region_masks_e4(400, 401) - assert (m4["seam"] & ~m1["seam"]).any() # mirrored lines add pixels - assert not (m1["seam"] & ~m4["seam"]).any() # nothing lost - assert e4["seam_e4_extra"]["pixels"] > 0 - assert e4["pass"] == e1["pass"] # diagnostic has no gate effect - with pytest.raises(ValueError, match="ensemble"): - compare_files("bad", a, r, ensemble=8) - - def test_repro_source_imports_no_torch(): # Static half of the standalone contract (runtime half is proven by # fresh-process reconstruction runs, where pytest has no Torch loaded). diff --git a/denoise/tests/test_contracts.py b/denoise/tests/test_contracts.py index 128a99bb38..a405a4b464 100644 --- a/denoise/tests/test_contracts.py +++ b/denoise/tests/test_contracts.py @@ -2,7 +2,8 @@ import pytest import torch from rapidraw_denoise.raw import cfa_positions, pack, unpack -from rapidraw_denoise.inference import tiled_apply, transform, inverse_transform +from rapidraw_denoise.inference import tiled_apply +from rapidraw_denoise.study.transforms import transform, inverse_transform from rapidraw_denoise.noise import estimate_noise, synthesize, NoiseProfile from rapidraw_denoise.vendor.nonlocalmf.sampling import reference_sampling diff --git a/denoise/tests/test_onnx_pipeline.py b/denoise/tests/test_onnx_pipeline.py index d4d9e36812..16e393a9fe 100644 --- a/denoise/tests/test_onnx_pipeline.py +++ b/denoise/tests/test_onnx_pipeline.py @@ -42,10 +42,7 @@ def test_tiling_identity_odd_rectangular(): assert np.array_equal(pipeline.tiled_apply(a, lambda x: x, tile=32, halo=8), a) -def test_transform_inverses_rectangular_and_negative_stride(): - x = np.arange(4 * 13 * 17).reshape(4, 13, 17) - for i in range(8): - assert np.array_equal(pipeline.inverse_transform(pipeline.transform(x, i), i), x) +def test_tiling_handles_negative_stride_views(): view = np.ascontiguousarray( np.random.default_rng(3).normal(size=(8, 40, 40)).astype(np.float32) )[:, ::-1, :] @@ -54,22 +51,19 @@ def test_transform_inverses_rectangular_and_negative_stride(): assert np.array_equal(out, np.ascontiguousarray(view)) -def test_four_rotation_ensemble_averages_transforms(): +def test_single_pass_returns_the_tiled_prediction(): rng = np.random.default_rng(5) packed = rng.uniform(-0.05, 0.9, size=(4, 96, 96)).astype(np.float32) profile = _profile() - mean, _, info = pipeline.denoise_with_predictor( - packed, Identity(), profile, tile=64, halo=16, ensemble=4) - assert info["ensemble"] == 4 - # Manual Welford mean of the four inverse-transformed tiled singles. + prediction, info = pipeline.denoise_with_predictor( + packed, Identity(), profile, tile=64, halo=16) + assert "ensemble" not in info and "disagreement_mean_variance" not in info + assert info["tile"] == 64 and info["halo"] == 16 image = np.clip(packed, 0, 1) variance = profile.variance(packed) conditioned = np.concatenate([image, np.sqrt(np.maximum(variance, 1e-12))]) - singles = [pipeline.inverse_transform( - pipeline.tiled_apply(pipeline.transform(conditioned, i), - lambda t: t[:4], tile=64, halo=16), i) - for i in range(4)] - np.testing.assert_array_equal(mean, sum(singles) / 4) + expected = pipeline.tiled_apply(conditioned, lambda t: t[:4], tile=64, halo=16) + np.testing.assert_array_equal(prediction, expected) def test_invalid_inputs_fail_closed(): @@ -78,7 +72,7 @@ def test_invalid_inputs_fail_closed(): tile=32, halo=16) with pytest.raises(ValueError): pipeline.denoise_with_predictor( - np.ones((4, 96, 96), np.float32), Identity(), _profile(), ensemble=3) + np.ones((4, 96, 96), np.float32), Identity(), _profile(), noise_scale=0) with pytest.raises(RuntimeError): pipeline.tiled_apply(np.ones((4, 64, 64), np.float32), lambda x: x * np.nan, tile=32, halo=8) @@ -95,9 +89,7 @@ def test_progress_semantics_preserved(): events = [] pipeline.denoise_with_predictor( np.random.default_rng(6).uniform(-0.05, 0.9, size=(4, 96, 96)).astype(np.float32), - Identity(), _profile(), tile=64, halo=16, ensemble=4, - progress=events.append) - passes = [e for e in events if set(e) == {"completed_passes", "total_passes"}] - assert [e["completed_passes"] for e in passes] == [1, 2, 3, 4] - tiles = [e for e in events if "completed_tiles" in e] - assert tiles and all(e["total_passes"] == 4 for e in tiles) + Identity(), _profile(), tile=64, halo=16, progress=events.append) + assert events and all(set(e) == {"completed_tiles", "total_tiles"} for e in events) + assert [e["completed_tiles"] for e in events] == list(range(1, len(events) + 1)) + assert all(e["total_tiles"] == len(events) for e in events) diff --git a/denoise/tests/test_onnx_rewrite.py b/denoise/tests/test_onnx_rewrite.py new file mode 100644 index 0000000000..2f66221fbe --- /dev/null +++ b/denoise/tests/test_onnx_rewrite.py @@ -0,0 +1,198 @@ +"""Checks for the throughput-only ONNX graph rewrites (zero-bias removal, dense 1x1 Conv -> MatMul). + +Synthetic graphs only; the accepted-bundle parity gates are recorded evidence, not part of this suite. +""" +import hashlib +import json + +import numpy as np +import pytest + +onnx = pytest.importorskip("onnx") +ort = pytest.importorskip("onnxruntime") +from onnx import TensorProto, helper, numpy_helper # noqa: E402 + +from rapidraw_denoise import onnx_rewrite # noqa: E402 + +RNG = np.random.default_rng(0) + + +def _initializer(name, array): + return numpy_helper.from_array(np.asarray(array, np.float32), name) + + +def _zero_bias_nodes(prefix, channels): + """The exporter's materialised zero bias: Constant -> CastLike -> Expand.""" + return [ + helper.make_node("Constant", [], [f"{prefix}_c"], name=f"{prefix}_c", + value=numpy_helper.from_array(np.array(0.0, np.float32))), + helper.make_node("CastLike", [f"{prefix}_c", f"{prefix}_w"], [f"{prefix}_cl"], name=f"{prefix}_cl"), + helper.make_node("Expand", [f"{prefix}_cl", f"{prefix}_shape"], [f"{prefix}_b"], name=f"{prefix}_e"), + ] + + +def _model(nodes, inits, in_shape, out_shape, out_name): + graph = helper.make_graph( + nodes, "g", + [helper.make_tensor_value_info("x", TensorProto.FLOAT, in_shape)], + [helper.make_tensor_value_info(out_name, TensorProto.FLOAT, out_shape)], inits) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 18)]) + model.ir_version = 10 + onnx.checker.check_model(model, full_check=True) + return model + + +def _conv_model(*, in_ch=8, out_ch=6, hw=5, batch=1, kernel=1, group=1, zero_bias=True, bias_value=None): + weight = RNG.standard_normal((out_ch, in_ch // group, kernel, kernel)).astype(np.float32) + inits = [_initializer("conv_w", weight)] + nodes, inputs = [], ["x", "conv_w"] + if zero_bias: + inits.append(numpy_helper.from_array(np.array([out_ch], np.int64), "conv_shape")) + nodes += _zero_bias_nodes("conv", out_ch) + inputs.append("conv_b") + elif bias_value is not None: + inits.append(_initializer("conv_b", np.full(out_ch, bias_value))) + inputs.append("conv_b") + pad = kernel // 2 + nodes.append(helper.make_node("Conv", inputs, ["y"], name="conv", kernel_shape=[kernel, kernel], + pads=[pad] * 4, group=group)) + return _model(nodes, inits, [batch, in_ch, hw, hw], [batch, out_ch, hw, hw], "y") + + +def _run(model, x): + options = ort.SessionOptions() + options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_DISABLE_ALL + session = ort.InferenceSession(model.SerializeToString(), options, providers=["CPUExecutionProvider"]) + return session.run(None, {"x": x})[0] + + +def _ops(model): + return [n.op_type for n in model.graph.node] + + +def _check_equivalent(model, report_expectation, **kwargs): + x = RNG.standard_normal([d.dim_value for d in model.graph.input[0].type.tensor_type.shape.dim]).astype(np.float32) + before = _run(model, x) + rewritten = onnx.ModelProto.FromString(model.SerializeToString()) + report = onnx_rewrite.rewrite_model(rewritten, **kwargs) + onnx.checker.check_model(rewritten, full_check=True) + np.testing.assert_allclose(_run(rewritten, x), before, atol=1e-5, rtol=1e-5) + for key, value in report_expectation.items(): + assert report[key] == value, (key, report) + return rewritten, report + + +def test_dense_1x1_with_zero_bias_becomes_matmul_and_stays_equivalent(): + model = _conv_model() + rewritten, report = _check_equivalent( + model, {"dense_1x1_conv_to_matmul": 1, "zero_biases_removed": 1}, min_in_channels=8) + assert _ops(rewritten) == ["Reshape", "MatMul", "Reshape"] + # The now-unused conv weight and the zero-bias chain are gone, not carried along. + names = {i.name for i in rewritten.graph.initializer} + assert "conv_w" not in names and "conv_shape" not in names + assert report["initializers_pruned"] >= 2 + + +def test_channel_threshold_leaves_small_1x1_as_conv_but_drops_its_zero_bias(): + rewritten, _ = _check_equivalent( + _conv_model(in_ch=8), {"dense_1x1_conv_to_matmul": 0, "zero_biases_removed": 1}, min_in_channels=16) + assert _ops(rewritten) == ["Conv"] + assert len(rewritten.graph.node[0].input) == 2 + + +def test_nonzero_bias_is_never_removed_or_converted(): + rewritten, _ = _check_equivalent( + _conv_model(zero_bias=False, bias_value=0.5), + {"dense_1x1_conv_to_matmul": 0, "zero_biases_removed": 0}, min_in_channels=8) + assert _ops(rewritten) == ["Conv"] and len(rewritten.graph.node[0].input) == 3 + + +def test_spatial_and_batched_convs_are_untouched(): + for kwargs in ({"kernel": 3}, {"batch": 2}, {"kernel": 3, "group": 8, "out_ch": 8}): + rewritten, _ = _check_equivalent( + _conv_model(**kwargs), {"dense_1x1_conv_to_matmul": 0, "grouped_1x1_conv_to_matmul": 0}, + min_in_channels=8) + assert _ops(rewritten) == ["Conv"], kwargs + + +def test_grouped_1x1_becomes_batched_matmul_and_keeps_group_channel_order(): + for group, out_ch in ((2, 6), (4, 8), (8, 16)): + model = _conv_model(in_ch=8, out_ch=out_ch, group=group, zero_bias=group != 4) + rewritten, report = _check_equivalent( + model, {"dense_1x1_conv_to_matmul": 0, "grouped_1x1_conv_to_matmul": 1, + "zero_biases_removed": int(group != 4)}, min_in_channels=8) + assert _ops(rewritten) == ["Reshape", "MatMul", "Reshape"], group + assert report["grouped_shapes_oc_ic_groups_h_w"] == [[out_ch, 8, group, 5, 5]] + + +def test_grouped_rewrite_can_be_disabled_independently(): + model = _conv_model(in_ch=8, out_ch=6, group=2) + rewritten, _ = _check_equivalent(model, {"grouped_1x1_conv_to_matmul": 0, "zero_biases_removed": 1}, + min_in_channels=8, grouped_as_matmul=False) + assert _ops(rewritten) == ["Conv"] + rewritten, _ = _check_equivalent(_conv_model(in_ch=8, out_ch=6), {"dense_1x1_conv_to_matmul": 0}, + min_in_channels=8, dense_as_matmul=False) + assert _ops(rewritten) == ["Conv"] + + +def test_bias_free_dense_1x1_converts_without_touching_bias_logic(): + rewritten, _ = _check_equivalent( + _conv_model(zero_bias=False), {"dense_1x1_conv_to_matmul": 1, "zero_biases_removed": 0}, min_in_channels=8) + assert _ops(rewritten) == ["Reshape", "MatMul", "Reshape"] + + +def _bundle(tmp_path, model, **manifest_overrides): + bundle = tmp_path / "source" + bundle.mkdir() + onnx.save(model, str(bundle / "model.onnx")) + sha = hashlib.sha256((bundle / "model.onnx").read_bytes()).hexdigest() + manifest = {"manifest_version": 1, "precision": "fp32", "diagnostic_shape": False, + "model_sha256": sha, "input_name": "x", "output_name": "y", **manifest_overrides} + (bundle / "manifest.json").write_text(json.dumps(manifest)) + return bundle, sha + + +def test_derive_bundle_records_lineage_and_never_overwrites(tmp_path, capsys): + bundle, source_sha = _bundle(tmp_path, _conv_model(in_ch=256, out_ch=8)) + out = tmp_path / "derived" + assert onnx_rewrite.main(["--bundle", str(bundle), "--output", str(out), "--min-channels", "256"]) == 0 + manifest = json.loads((out / "manifest.json").read_text()) + new_sha = hashlib.sha256((out / "model.onnx").read_bytes()).hexdigest() + assert manifest["model_sha256"] == new_sha != source_sha + assert manifest["graph_rewrites"]["source_model_sha256"] == source_sha + assert manifest["graph_rewrites"]["dense_1x1_conv_to_matmul"] == 1 + assert manifest["precision"] == "fp32" + sums = dict(line.split(" ")[::-1] for line in (out / "SHA256SUMS.txt").read_text().splitlines()) + assert sums["model.onnx"] == new_sha + # Source untouched, second run refuses to overwrite. + assert hashlib.sha256((bundle / "model.onnx").read_bytes()).hexdigest() == source_sha + assert onnx_rewrite.main(["--bundle", str(bundle), "--output", str(out)]) == 1 + assert not list(tmp_path.glob(".onnx-rewrite-*")) + + +def test_derive_can_record_a_new_tiling_halo(tmp_path): + bundle, _ = _bundle(tmp_path, _conv_model(in_ch=256, out_ch=8), + tile=320, halo=64, retained_core=192) + out = tmp_path / "derived" + assert onnx_rewrite.main(["--bundle", str(bundle), "--output", str(out), "--halo", "40"]) == 0 + manifest = json.loads((out / "manifest.json").read_text()) + assert (manifest["tile"], manifest["halo"], manifest["retained_core"]) == (320, 40, 240) + for bad in ("42", "160", "-4"): + assert onnx_rewrite.main(["--bundle", str(bundle), "--output", str(tmp_path / f"bad{bad}"), + "--halo", bad]) == 1 + + +def test_derive_rejects_hash_mismatch_nothing_to_rewrite_and_non_fp32(tmp_path): + bundle, _ = _bundle(tmp_path, _conv_model(in_ch=256, out_ch=8), model_sha256="0" * 64) + assert onnx_rewrite.main(["--bundle", str(bundle), "--output", str(tmp_path / "a")]) == 1 + assert not (tmp_path / "a").exists() + + small = tmp_path / "small" + small.mkdir() + bundle, _ = _bundle(small, _conv_model(in_ch=8, out_ch=6, zero_bias=False)) + assert onnx_rewrite.main(["--bundle", str(bundle), "--output", str(tmp_path / "b")]) == 1 + + half = tmp_path / "half" + half.mkdir() + bundle, _ = _bundle(half, _conv_model(in_ch=256, out_ch=8), precision="fp16") + assert onnx_rewrite.main(["--bundle", str(bundle), "--output", str(tmp_path / "c")]) == 1 diff --git a/denoise/tests/test_worker.py b/denoise/tests/test_worker.py index 708f5c6440..42ed5568e6 100644 --- a/denoise/tests/test_worker.py +++ b/denoise/tests/test_worker.py @@ -6,13 +6,13 @@ def request(**overrides): return dict(protocol=1, algorithm=ALGORITHM, model_sha256=CHECKPOINT_SHA256, - shape=[4, 128, 256], tile=320, halo=64, ensemble=1) | overrides + shape=[4, 128, 256], tile=320, halo=40) | overrides @pytest.mark.parametrize("override", [ {"shape": [4, 2**64, 16]}, {"shape": [4, 0, 32]}, {"shape": [3, 16, 16]}, - {"shape": [4, 16., 16]}, {"shape": [4, 7, 32]}, {"ensemble": True}, - {"ensemble": 8}, {"model_sha256": "wrong"}, {"protocol": 2}, {"halo": 0}, + {"shape": [4, 16., 16]}, {"shape": [4, 7, 32]}, + {"model_sha256": "wrong"}, {"protocol": 2}, {"halo": 0}, {"halo": 64}, ]) def test_invalid_protocol_fails_before_loading_model(override): with pytest.raises(ValueError): diff --git a/mcp/scripts/nonlocal-e2e.mjs b/mcp/scripts/nonlocal-e2e.mjs index 7e61d7c30c..ed8d767eac 100644 --- a/mcp/scripts/nonlocal-e2e.mjs +++ b/mcp/scripts/nonlocal-e2e.mjs @@ -106,12 +106,12 @@ try { await render(zeroJob.result_session_id, 'zero-strength'); const rejected = await call( 'start_denoise', - { session_id: parent.session_id, method: 'bm3d', quality: 'maximum' }, + { session_id: parent.session_id, method: 'nonlocal', quality: 'maximum' }, true, ); assert.match(JSON.stringify(rejected), /quality/i); - // A different quality has a distinct prediction cache. Cancel after CUDA - // begins, then restart the same immutable snapshot with resume_job. + // Cancel after CUDA begins, then restart the same immutable snapshot with + // resume_job. let job = await data('start_denoise', { session_id: parent.session_id, method: 'nonlocal', diff --git a/mcp/scripts/nonlocal-lib-bench.mjs b/mcp/scripts/nonlocal-lib-bench.mjs index 2a22dc2edb..b09cd3145a 100644 --- a/mcp/scripts/nonlocal-lib-bench.mjs +++ b/mcp/scripts/nonlocal-lib-bench.mjs @@ -1,7 +1,7 @@ /** Benchmark-only native Nonlocal library A/B runner (one fresh job). * * Launches one fresh native engine + isolated workspace, opens a fixed RAW, - * runs a single balanced (e1) nonlocal denoise at a fixed intensity, times + * runs a single-pass nonlocal denoise at a fixed intensity, times * start→success, captures backend provenance, renders a deterministic PNG, * and records the engine's actually-mapped CUDA libraries from /proc maps. * Intended for N_S vs N_P comparisons where only process-local library @@ -121,7 +121,6 @@ try { const startedJob = await data('start_denoise', { session_id: opened.session_id, method: 'nonlocal', - quality: 'balanced', intensity, }); let state; diff --git a/mcp/src/tools.ts b/mcp/src/tools.ts index 23d547d35a..7bf1d6dfd1 100644 --- a/mcp/src/tools.ts +++ b/mcp/src/tools.ts @@ -414,12 +414,11 @@ export const toolDefinitions: ToolDefinition[] = [ ), tool( 'start_denoise', - 'Start background denoise with captured source and edits. ai (default) is lightweight NIND; bm3d is native CPU; nonlocal is Bayer RAW-to-RAW using the provider-selected installed bundle (RAPIDRAW_NONLOCAL_BUNDLE override when set) with native ONNX on CPU or Linux CUDA, or direct CoreML.framework on macOS. Use install_model kind=nonlocal first when the bundle is absent; start_denoise never downloads. Nonlocal creates a float Bayer DNG developed by the normal RAW pipeline; quality balanced uses one pass, maximum four rotations. Intensity defaults to 100 for nonlocal, 50 otherwise. Returns job_id; get_job yields result_session_id. One worker per workspace; editing remains available. Original session is preserved.', + 'Start background denoise with captured source and edits. ai (default) is lightweight NIND; bm3d is native CPU; nonlocal is Bayer RAW-to-RAW using the provider-selected installed bundle (RAPIDRAW_NONLOCAL_BUNDLE override when set) with native ONNX on CPU or Linux CUDA, or direct CoreML.framework on macOS. Use install_model kind=nonlocal first when the bundle is absent; start_denoise never downloads. Nonlocal creates a float Bayer DNG developed by the normal RAW pipeline. Intensity defaults to 100 for nonlocal, 50 otherwise. Returns job_id; get_job yields result_session_id. One worker per workspace; editing remains available. Original session is preserved.', { ...mutation, intensity: percent.optional(), method: z.enum(['ai', 'bm3d', 'nonlocal']).optional(), - quality: z.enum(['balanced', 'maximum']).optional(), }, ), tool( diff --git a/skills/rapidraw-mcp/SKILL.md b/skills/rapidraw-mcp/SKILL.md index 837f3eb862..61d2ac4f78 100644 --- a/skills/rapidraw-mcp/SKILL.md +++ b/skills/rapidraw-mcp/SKILL.md @@ -41,7 +41,7 @@ For a collection, inspect each source at overview and relevant native detail bef When noise blocks the intended detail or tonal separation, read [dedicated denoising](references/advanced-editing.md#dedicated-denoising). Ordinary noise-reduction sliders do not exercise the dedicated AI, BM3D or Nonlocal operations; an unsuccessful slider treatment does not establish the engine's limit. -Choose the denoiser from the source and the host. For Bayer RAW on a configured native host (ONNX CPU or Linux CUDA, or macOS CoreML with its pinned `.mlpackage` bundle), prefer `method: "nonlocal"`: it denoises before demosaicing and returns a Bayer DNG that keeps RapidRAW's native RAW pipeline. Start with `quality: "balanced"` and `intensity: 100` for high-ISO stills; start at `intensity: 50` when faint stars or subtle texture need extra protection, comparing both candidates from the same undenoised parent. Reserve `quality: "maximum"` for a final keeper when roughly four inference passes are acceptable, since the measured gain over balanced was small. Use `method: "ai"` (NIND) for a quick low-resource workflow, and `method: "bm3d"` as a CPU baseline. Nonlocal requires Bayer RAW and a pinned native ONNX bundle; there is no automatic Mac-to-server offload. Always inspect the result session at overview and native detail before accepting it. +Choose the denoiser from the source and the host. For Bayer RAW on a configured native host (ONNX CPU or Linux CUDA, or macOS CoreML with its pinned `.mlpackage` bundle), prefer `method: "nonlocal"`: it denoises before demosaicing and returns a Bayer DNG that keeps RapidRAW's native RAW pipeline. Start with `intensity: 100` for high-ISO stills; start at `intensity: 50` when faint stars or subtle texture need extra protection, comparing both candidates from the same undenoised parent. Use `method: "ai"` (NIND) for a quick low-resource workflow, and `method: "bm3d"` as a CPU baseline. Nonlocal requires Bayer RAW and a pinned native ONNX bundle; there is no automatic Mac-to-server offload. Always inspect the result session at overview and native detail before accepting it. For a stitched full-sphere image or a flat view selected from it, read [spherical photo handoff](references/spherical-photos.md). Projection, colour assumptions, periodic seams and final metadata need their own checks before a normal image export qualifies as a 360 deliverable. diff --git a/skills/rapidraw-mcp/references/advanced-editing.md b/skills/rapidraw-mcp/references/advanced-editing.md index 9379ad1056..8fab6b1a4b 100644 --- a/skills/rapidraw-mcp/references/advanced-editing.md +++ b/skills/rapidraw-mcp/references/advanced-editing.md @@ -154,13 +154,13 @@ Use this path when noise is a consequential obstacle to the chosen edit. `colorN ### Recommended method selection -- **Configured native host (ONNX CPU or Linux CUDA, or macOS CoreML) + Bayer RAW:** use `method: "nonlocal"`. Start with `quality: "balanced"`, `intensity: 100` for general high-ISO stills, and compare `intensity: 50` when faint stars or fine texture may be softened. Use `quality: "maximum"` only for a keeper where roughly four inference passes are acceptable; its measured improvement over balanced was small. +- **Configured native host (ONNX CPU or Linux CUDA, or macOS CoreML) + Bayer RAW:** use `method: "nonlocal"`. Start with `intensity: 100` for general high-ISO stills, and compare `intensity: 50` when faint stars or fine texture may be softened. - **Local Mac, CPU, or quick low-resource work:** use `method: "ai"` (NIND), the lightweight option. It can use the existing configured ONNX path where available. - **CPU reference or no AI assets:** use `method: "bm3d"` as a conventional baseline. For fair comparisons, derive every candidate from the same undenoised parent and hold the crop, grade and review regions constant. Nonlocal is Bayer RAW-only and runs in-process on the configured native ONNX host; there is no automatic Mac-to-server offload. Inspect the returned session at overview and native detail before accepting any method. -On a configured native host, `start_denoise` also accepts `method: "nonlocal"` for Bayer RAW. Inspect `models.groups.nonlocal`: bundle/provider verification does not guarantee CUDA availability until the job starts. On macOS, `RAPIDRAW_NONLOCAL_PROVIDER=coreml` with the pinned `.mlpackage` bundle runs direct CoreML.framework inference with no Python/coremltools dependency. Use `quality: "balanced"` (one pass) first; `"maximum"` averages four rotations. Intensity defaults to 100 and blends in sensor space; cached full-strength predictions can be reused for other intensities. The result is a float Bayer DNG using RapidRAW's native demosaic/colour pipeline and captured edits in a new session. Open its DNG independently for normal RAW development, or continue the returned session to retain its exact edits. This is a background MCP option, not the desktop denoise dialog or automatic Mac-to-server offload. See [GPU setup](../../../denoise/README.md#native-raw-denoise-through-mcp). Do not route Nonlocal through the synchronous `denoise` compatibility tool. +On a configured native host, `start_denoise` also accepts `method: "nonlocal"` for Bayer RAW. Inspect `models.groups.nonlocal`: bundle/provider verification does not guarantee CUDA availability until the job starts. On macOS, `RAPIDRAW_NONLOCAL_PROVIDER=coreml` with the pinned `.mlpackage` bundle runs direct CoreML.framework inference with no Python/coremltools dependency. Nonlocal has no quality setting; do not pass `quality`. Model time scales with photograph size (about 23 s for 32.5 MP on an RTX 5060 Ti, estimated), so poll the job. Intensity defaults to 100 and blends in sensor space; cached full-strength predictions can be reused for other intensities. The result is a float Bayer DNG using RapidRAW's native demosaic/colour pipeline and captured edits in a new session. Open its DNG independently for normal RAW development, or continue the returned session to retain its exact edits. This is a background MCP option, not the desktop denoise dialog or automatic Mac-to-server offload. See [GPU setup](../../../denoise/README.md#native-raw-denoise-through-mcp). Do not route Nonlocal through the synchronous `denoise` compatibility tool. For a method or strength comparison, derive each candidate from the same undenoised parent and hold the grading and geometry constant. Denoising an already derived candidate compounds the filtering. AI intensity blends the filtered result with the source; BM3D intensity also changes filter parameters, so equal percentages are not equal treatments. Zero disables filtering; the synchronous tool is a no-op, while a background job produces a separate unchanged-source session. NIND uses CPU inference by default; a configured Linux deployment can use the [optional CUDA path](https://github.com/sheldonxxxx/RapidRAW/blob/main/docs/mcp/onnx-cuda.md). Inspect `models.onnx_execution` for the selected provider. Lowering nonzero AI intensity does not reduce the inference workload. Full-resolution processing can be slow. Prefer [background jobs](review-and-jobs.md#background-denoise) so polling, editing and cancellation remain available. For an older synchronous engine, the fallback client's native timeout and request wait are separate, as explained in [connection and recovery](connection.md#persistent-fallback-client). diff --git a/skills/rapidraw-mcp/references/review-and-jobs.md b/skills/rapidraw-mcp/references/review-and-jobs.md index 6f5d38f5f5..054086d6b5 100644 --- a/skills/rapidraw-mcp/references/review-and-jobs.md +++ b/skills/rapidraw-mcp/references/review-and-jobs.md @@ -47,7 +47,7 @@ The EV statistics and row profile describe the selected full-resolution region. ## Background denoise -1. Inspect `models` for AI and `models.groups.nonlocal` for the GPU backend; install missing assets only when needed and authorized. Call `start_denoise` with parent session, latest revision, `method: "ai"`, `"bm3d"`, or `"nonlocal"`, and intensity. Nonlocal requires Bayer RAW and a pinned native bundle installed once via `install_model kind="nonlocal"` for the selected provider (`cpu`, Linux CUDA, or macOS CoreML); `RAPIDRAW_NONLOCAL_BUNDLE` is an explicit manual-bundle override. Its optional quality is `balanced` or `maximum`. Source capture and asset verification precede the job response; Nonlocal inference runs in-process through native ONNX Runtime or directly through CoreML.framework. +1. Inspect `models` for AI and `models.groups.nonlocal` for the GPU backend; install missing assets only when needed and authorized. Call `start_denoise` with parent session, latest revision, `method: "ai"`, `"bm3d"`, or `"nonlocal"`, and intensity. Nonlocal requires Bayer RAW and a pinned native bundle installed once via `install_model kind="nonlocal"` for the selected provider (`cpu`, Linux CUDA, or macOS CoreML); `RAPIDRAW_NONLOCAL_BUNDLE` is an explicit manual-bundle override. Source capture and asset verification precede the job response; Nonlocal inference runs in-process through native ONNX Runtime or directly through CoreML.framework. 2. Keep `job_id`. Poll `get_job` at a reasonable interval; the session remains available for editing and rendering. One denoise job runs per workspace. `progress_percent` covers processing and result publication; 100 means the result was saved. 3. Success returns a separate `result_session_id`. Inspect it and matching parent detail. Its source pixels, RAW interpretation and inherited edits come from the captured input, even if the parent was edited while filtering ran. Save that session before delivery. 4. To stop, call `cancel_job`, then poll until terminal. Cancellation is cooperative between BM3D patches/NIND tiles; Nonlocal cancellation lands on tile boundaries and cannot abort an active ORT or CoreML call. It does not kill the editing engine. If completion won the race, use the returned completed status. diff --git a/src-tauri/src/denoising.rs b/src-tauri/src/denoising.rs index 7070ddc2ea..2f61978498 100644 --- a/src-tauri/src/denoising.rs +++ b/src-tauri/src/denoising.rs @@ -451,8 +451,8 @@ fn restore_denoise_source( /// float Bayer DNG is developed through the normal RAW pipeline below, so the /// caller must still apply the default RAW processing for RAW sources. /// Strength blends inside the RAW domain; no sRGB blending happens here. -/// Uses balanced quality; the model bundle must already be installed -/// (missing bundles fail with `MODEL_NOT_INSTALLED`). +/// The model bundle must already be installed (missing bundles fail with +/// `MODEL_NOT_INSTALLED`). fn run_nonlocal_denoise( path_str: &str, file_bytes: &[u8], @@ -478,15 +478,9 @@ fn run_nonlocal_denoise( ); }), }; - let output = crate::raw_denoise::denoise( - Path::new(path_str), - &source_sha, - &root, - intensity, - "balanced", - &control, - ) - .map_err(|e| format!("{e:#}"))?; + let output = + crate::raw_denoise::denoise(Path::new(path_str), &source_sha, &root, intensity, &control) + .map_err(|e| format!("{e:#}"))?; let dng_path = output.path(); let dng_bytes = fs::read(&dng_path).map_err(|e| e.to_string())?; load_base_image_from_bytes( diff --git a/src-tauri/src/mcp_bridge/jobs.rs b/src-tauri/src/mcp_bridge/jobs.rs index c08d2a4fbc..10e527b21a 100644 --- a/src-tauri/src/mcp_bridge/jobs.rs +++ b/src-tauri/src/mcp_bridge/jobs.rs @@ -47,8 +47,6 @@ struct Job { stage: String, method: String, intensity: f32, - #[serde(default = "balanced_quality")] - quality: String, /// Backend generation that owns unfinished work. Absent on manifests /// written before the native ONNX switch; serde default keeps those /// manifests listable, and resume rejects them explicitly below. @@ -59,9 +57,6 @@ struct Job { error: Option, attempt: u32, } -fn balanced_quality() -> String { - "balanced".into() -} fn legacy_backend_generation() -> String { String::new() } @@ -88,7 +83,7 @@ impl Job { "Completed results persist without polling. After process interruption, resume_job restarts computation from the captured input; partial tiles are not checkpoints." }; json!({"job_id":self.job_id,"status":self.status,"progress_percent":self.progress_percent,"stage":self.stage, - "method":self.method,"intensity":self.intensity,"quality":self.quality,"backend":self.backend,"parent_session_id":self.input.id,"parent_revision":self.input.revision, + "method":self.method,"intensity":self.intensity,"backend":self.backend,"parent_session_id":self.input.id,"parent_revision":self.input.revision, "result_session_id":if self.status==Status::Succeeded { Some(&self.result_session_id) } else { None }, "error":self.error,"attempt":self.attempt,"recoverable":recoverable, "recovery":recovery}) @@ -182,8 +177,6 @@ impl Jobs { job.input .validate_restored(&root.join("sessions").join(&job.input.id))?; if !["ai", "bm3d", "nonlocal"].contains(&job.method.as_str()) - || !["balanced", "maximum"].contains(&job.quality.as_str()) - || (job.method != "nonlocal" && job.quality != "balanced") || !job.intensity.is_finite() || !(0.0..=100.0).contains(&job.intensity) { @@ -382,7 +375,7 @@ impl Jobs { control.check()?; atomic_write( &directory.join("session.json"), - &serde_json::to_vec_pretty(&result).map_err(|e| e.to_string())?, + &super::sessions::session_manifest_bytes(&result, true)?, false, )?; entry.job.status = Status::Succeeded; @@ -416,7 +409,7 @@ impl Jobs { object.remove("mcpSourceDomain"); } metadata["derivedFrom"] = json!({"session_id":job.input.id,"revision":job.input.revision,"job_id":job.job_id, - "operation":"denoise","method":"nonlocal","intensity":job.intensity,"quality":job.quality, + "operation":"denoise","method":"nonlocal","intensity":job.intensity, "source_domain_preserved":true,"nonlocal":output.provenance}); let source_sha256 = crate::raw_denoise::hash_file(&derived).map_err(|e| e.to_string())?; let result = Session { @@ -442,7 +435,7 @@ impl Jobs { control.check()?; atomic_write( &directory.join("session.json"), - &serde_json::to_vec_pretty(&result).map_err(|e| e.to_string())?, + &super::sessions::session_manifest_bytes(&result, true)?, false, )?; entry.job.status = Status::Succeeded; @@ -523,22 +516,6 @@ impl Bridge { if !["ai", "bm3d", "nonlocal"].contains(&method) { return Err("INVALID_ARGUMENT: method must be ai, bm3d or nonlocal".into()); } - let quality = params - .get("quality") - .map(|v| { - v.as_str() - .ok_or("INVALID_ARGUMENT: quality must be a string") - }) - .transpose()? - .unwrap_or("balanced"); - if !["balanced", "maximum"].contains(&quality) - || (method != "nonlocal" && params.get("quality").is_some()) - { - return Err( - "INVALID_ARGUMENT: quality is balanced or maximum and applies only to nonlocal" - .into(), - ); - } Job { job_id: uuid::Uuid::new_v4().to_string(), status: Status::Running, @@ -552,7 +529,6 @@ impl Bridge { 0.0, 100.0, )? as f32, - quality: quality.into(), backend: if method == "nonlocal" { nonlocal_generation_for_new_job() } else { @@ -641,7 +617,6 @@ impl Bridge { &job.input.source_sha256, &jobs.root, job.intensity / 100., - &job.quality, &control, ) .map_err(|e| format!("{e:#}"))?; @@ -743,7 +718,6 @@ mod tests { stage: "AI tiles".into(), method: "bm3d".into(), intensity: 50.0, - quality: balanced_quality(), backend: String::new(), input, result_session_id: uuid::Uuid::new_v4().to_string(), diff --git a/src-tauri/src/mcp_bridge/render.rs b/src-tauri/src/mcp_bridge/render.rs index 047abea4e3..3b547eba41 100644 --- a/src-tauri/src/mcp_bridge/render.rs +++ b/src-tauri/src/mcp_bridge/render.rs @@ -679,7 +679,7 @@ impl Bridge { if !manifest.is_file() { continue; } - let session: Session = serde_json::from_slice( + let session = super::sessions::session_from_manifest_bytes( &fs::read(&manifest).map_err(|e| e.to_string())?, ) .map_err(|e| format!("INVALID_SESSION: Cannot verify protected source paths: {e}"))?; diff --git a/src-tauri/src/mcp_bridge/sessions.rs b/src-tauri/src/mcp_bridge/sessions.rs index 1c912a355d..cefbb279ba 100644 --- a/src-tauri/src/mcp_bridge/sessions.rs +++ b/src-tauri/src/mcp_bridge/sessions.rs @@ -161,11 +161,7 @@ impl Bridge { .join("sessions") .join(&session.id) .join("session.json"); - atomic_write( - &path, - &serde_json::to_vec_pretty(session).map_err(|e| e.to_string())?, - true, - ) + atomic_write(&path, &session_manifest_bytes(session, true)?, true) } pub fn commit( &mut self, @@ -460,6 +456,82 @@ pub(super) fn validate_metadata(metadata: &Value) -> Result<()> { Ok(()) } +/// Embedded images (AI mask bitmaps) repeat in every history snapshot. The +/// manifest stores each distinct image once in a hash-keyed table and refers to +/// it from the snapshots; restore expands them again, so the engine and the +/// MCP surface still see the original adjustments. +const BLOB_MIN_LEN: usize = 4096; +const BLOB_PREFIX: &str = "blob:sha256:"; + +fn externalize_blobs(value: &mut Value, blobs: &mut BTreeMap) { + match value { + Value::String(text) if text.len() >= BLOB_MIN_LEN && text.starts_with("data:image/") => { + use sha2::{Digest, Sha256}; + let key: String = Sha256::digest(text.as_bytes()) + .iter() + .map(|byte| format!("{byte:02x}")) + .collect(); + let reference = format!("{BLOB_PREFIX}{key}"); + let original = std::mem::replace(text, reference); + blobs.entry(key).or_insert(original); + } + Value::Array(items) => items.iter_mut().for_each(|v| externalize_blobs(v, blobs)), + Value::Object(map) => map.values_mut().for_each(|v| externalize_blobs(v, blobs)), + _ => {} + } +} + +fn internalize_blobs(value: &mut Value, blobs: &BTreeMap) -> Result<()> { + match value { + Value::String(text) if text.starts_with(BLOB_PREFIX) => { + let key = &text[BLOB_PREFIX.len()..]; + let original = blobs.get(key).ok_or_else(|| { + format!("INVALID_SESSION: Session manifest is missing embedded image {key}") + })?; + *text = original.clone(); + } + Value::Array(items) => { + for item in items { + internalize_blobs(item, blobs)?; + } + } + Value::Object(map) => { + for item in map.values_mut() { + internalize_blobs(item, blobs)?; + } + } + _ => {} + } + Ok(()) +} + +/// Serialize a session manifest, storing repeated embedded images once. +pub(super) fn session_manifest_bytes(session: &Session, pretty: bool) -> Result> { + let mut value = serde_json::to_value(session).map_err(|e| e.to_string())?; + let mut blobs = BTreeMap::new(); + externalize_blobs(&mut value, &mut blobs); + if !blobs.is_empty() { + value["blobs"] = serde_json::to_value(blobs).map_err(|e| e.to_string())?; + } + if pretty { + serde_json::to_vec_pretty(&value) + } else { + serde_json::to_vec(&value) + } + .map_err(|e| e.to_string()) +} + +/// Parse a manifest written with or without the shared image table. +pub(super) fn session_from_manifest_bytes(bytes: &[u8]) -> std::result::Result { + let mut value: Value = serde_json::from_slice(bytes).map_err(|e| e.to_string())?; + if let Some(table) = value.as_object_mut().and_then(|map| map.remove("blobs")) { + let blobs: BTreeMap = + serde_json::from_value(table).map_err(|e| format!("invalid image table: {e}"))?; + internalize_blobs(&mut value, &blobs)?; + } + serde_json::from_value(value).map_err(|e| e.to_string()) +} + pub(super) fn restore_session(directory: &Path) -> Result { let manifest = directory.join("session.json"); let file_type = fs::symlink_metadata(&manifest) @@ -470,7 +542,7 @@ pub(super) fn restore_session(directory: &Path) -> Result { "INVALID_SESSION: Session manifest must be a regular file, not a symlink".into(), ); } - let session: Session = serde_json::from_slice(&fs::read(&manifest).map_err(|e| e.to_string())?) + let session = session_from_manifest_bytes(&fs::read(&manifest).map_err(|e| e.to_string())?) .map_err(|e| format!("INVALID_SESSION: {}: {e}", manifest.display()))?; session.validate_restored(directory)?; Ok(session) @@ -540,6 +612,44 @@ mod tests { (fixture, directory) } + #[test] + fn manifest_stores_repeated_embedded_images_once_and_restores_them() { + let bitmap = format!("data:image/png;base64,{}", "A".repeat(BLOB_MIN_LEN)); + let other = format!("data:image/png;base64,{}", "B".repeat(BLOB_MIN_LEN)); + let mut before = session(); + for (index, snapshot) in before.history.iter_mut().enumerate() { + snapshot.adjustments = json!({"masks":[{"subMasks":[ + {"parameters":{"maskDataBase64": bitmap}}, + {"parameters":{"maskDataBase64": if index == 0 { &bitmap } else { &other }}} + ]}], "note": "short data:image/ text stays inline"}); + } + let bytes = session_manifest_bytes(&before, false).unwrap(); + let text = String::from_utf8(bytes.clone()).unwrap(); + assert_eq!(text.matches(&bitmap).count(), 1); + assert_eq!(text.matches(&other).count(), 1); + assert!(text.contains("blob:sha256:")); + let after = session_from_manifest_bytes(&bytes).unwrap(); + for (a, b) in before.history.iter().zip(&after.history) { + assert_eq!(a.adjustments, b.adjustments); + } + } + + #[test] + fn manifest_without_image_table_still_restores_and_missing_images_are_rejected() { + let before = session(); + let legacy = serde_json::to_vec(&before).unwrap(); + assert_eq!( + session_from_manifest_bytes(&legacy).unwrap().history.len(), + 2 + ); + let mut value = serde_json::to_value(&before).unwrap(); + value["history"][0]["adjustments"] = + json!({"x": format!("{BLOB_PREFIX}{}", "0".repeat(64))}); + value["blobs"] = json!({}); + let error = session_from_manifest_bytes(&serde_json::to_vec(&value).unwrap()).unwrap_err(); + assert!(error.contains("missing embedded image")); + } + #[test] fn restore_checks_all_history_including_inactive_snapshots() { let root = tempfile::tempdir().unwrap(); diff --git a/src-tauri/src/nonlocal_coreml.rs b/src-tauri/src/nonlocal_coreml.rs index c849824d38..82d4492d4d 100644 --- a/src-tauri/src/nonlocal_coreml.rs +++ b/src-tauri/src/nonlocal_coreml.rs @@ -20,7 +20,7 @@ //! which natively compiles the `.mlpackage` to `.mlmodelc` and loads one //! `MLModel` per denoise job with `MLComputeUnitsAll` and //! `allowLowPrecisionAccumulationOnGPU = NO`. The shared Rust -//! noise/conditioning/tiling/ensemble driver (`nonlocal_onnx::denoise_packed`) +//! noise/conditioning/tiling driver (`nonlocal_onnx::denoise_packed`) //! is reused through the [`nonlocal_onnx::TilePredictor`] contract; no //! pipeline arithmetic is duplicated here. @@ -302,8 +302,7 @@ mod ffi { } } -/// One compiled `MLModel` per denoise job, reused for all tiles and ensemble -/// passes. The handle is owned by the job worker thread. +/// One compiled `MLModel` per denoise job, reused for all tiles. The handle is owned by the job worker thread. #[cfg(target_os = "macos")] pub(crate) struct Backend { handle: *mut c_void, @@ -920,17 +919,16 @@ mod tests { } /// Native Rust/CoreML full-photo parity through the shared - /// noise/conditioning/tiling/Welford driver. Requires the exact accepted - /// package plus a packed `[4,H,W]` little-endian float32 input and the - /// frozen torch reference outputs: - /// `RAPIDRAW_NONLOCAL_COREML_BUNDLE`, `RAPIDRAW_NONLOCAL_COREML_PHOTO_INPUT`, - /// `RAPIDRAW_NONLOCAL_COREML_PHOTO_REF_E1` and - /// `RAPIDRAW_NONLOCAL_COREML_PHOTO_REF_E4` (single-array float32 `.npy`). + /// noise/conditioning/tiling driver. Requires the exact accepted package + /// plus a packed `[4,H,W]` little-endian float32 input and the frozen + /// torch reference output computed with the production 320/40 tiling: + /// `RAPIDRAW_NONLOCAL_COREML_BUNDLE`, `RAPIDRAW_NONLOCAL_COREML_PHOTO_INPUT` + /// and `RAPIDRAW_NONLOCAL_COREML_PHOTO_REF` (single-array float32 `.npy`). /// Gates match the tile family (zero elementwise violations at /// 1e-4 atol/rtol, MAE <= 1e-5, p99 <= 1e-4). #[test] #[cfg(target_os = "macos")] - #[ignore = "Requires exact CoreML bundle + full-photo packed input and torch references; runs e1+e4 end to end"] + #[ignore = "Requires exact CoreML bundle + full-photo packed input and torch reference; runs end to end"] fn accepted_package_full_photo_parity() { let bundle = std::path::PathBuf::from( std::env::var_os("RAPIDRAW_NONLOCAL_COREML_BUNDLE") @@ -953,75 +951,65 @@ mod tests { cancelled: std::sync::Arc::new(std::sync::atomic::AtomicBool::new(false)), progress: std::sync::Arc::new(|_, _| {}), }; - for (ensemble, var) in [ - (1u32, "RAPIDRAW_NONLOCAL_COREML_PHOTO_REF_E1"), - (4u32, "RAPIDRAW_NONLOCAL_COREML_PHOTO_REF_E4"), - ] { - let ref_path = std::path::PathBuf::from( - std::env::var_os(var) - .unwrap_or_else(|| panic!("Set {var} to the torch reference .npy")), - ); - let (ref_shape, reference) = - frozen_fixtures::read_npy(&std::fs::read(&ref_path).expect("reference readable")) - .expect("reference parses"); - assert_eq!(ref_shape.len(), 3, "reference must be [4,H,W]"); - let (height, width) = (ref_shape[1], ref_shape[2]); - assert_eq!(ref_shape[0], 4); - let raw = std::fs::read(&input_path).expect("photo input readable"); - assert_eq!( - raw.len(), - 4 * height * width * 4, - "input size matches reference" - ); - let packed: Vec = raw - .as_chunks::<4>() - .0 - .iter() - .map(|b| f32::from_le_bytes(*b)) - .collect(); - let started = std::time::Instant::now(); - let result = crate::nonlocal_onnx::denoise_packed( - &mut backend, - &packed, - height, - width, - ensemble, - &control, - ) - .expect("full-photo inference"); - let wall_s = started.elapsed().as_secs_f64(); - assert_eq!(result.prediction.len(), reference.len()); - let label = format!("coreml-full-e{ensemble}"); - // Full-array gates, same family as the tile gates. - let mut max_abs = 0f64; - let mut sum_abs = 0f64; - let mut violations = 0usize; - let mut abs_dev = Vec::with_capacity(reference.len()); - for (a, r) in result.prediction.iter().zip(reference.iter()) { - assert!(a.is_finite(), "{label}: finite"); - let err = (*a as f64 - *r as f64).abs(); - max_abs = max_abs.max(err); - sum_abs += err; - if err > 1e-4 + 1e-4 * (*r as f64).abs() { - violations += 1; - } - abs_dev.push(err); + let var = "RAPIDRAW_NONLOCAL_COREML_PHOTO_REF"; + let ref_path = std::path::PathBuf::from( + std::env::var_os(var) + .unwrap_or_else(|| panic!("Set {var} to the torch reference .npy")), + ); + let (ref_shape, reference) = + frozen_fixtures::read_npy(&std::fs::read(&ref_path).expect("reference readable")) + .expect("reference parses"); + assert_eq!(ref_shape.len(), 3, "reference must be [4,H,W]"); + let (height, width) = (ref_shape[1], ref_shape[2]); + assert_eq!(ref_shape[0], 4); + let raw = std::fs::read(&input_path).expect("photo input readable"); + assert_eq!( + raw.len(), + 4 * height * width * 4, + "input size matches reference" + ); + let packed: Vec = raw + .as_chunks::<4>() + .0 + .iter() + .map(|b| f32::from_le_bytes(*b)) + .collect(); + let started = std::time::Instant::now(); + let result = + crate::nonlocal_onnx::denoise_packed(&mut backend, &packed, height, width, &control) + .expect("full-photo inference"); + let wall_s = started.elapsed().as_secs_f64(); + assert_eq!(result.prediction.len(), reference.len()); + let label = "coreml-full"; + // Full-array gates, same family as the tile gates. + let mut max_abs = 0f64; + let mut sum_abs = 0f64; + let mut violations = 0usize; + let mut abs_dev = Vec::with_capacity(reference.len()); + for (a, r) in result.prediction.iter().zip(reference.iter()) { + assert!(a.is_finite(), "{label}: finite"); + let err = (*a as f64 - *r as f64).abs(); + max_abs = max_abs.max(err); + sum_abs += err; + if err > 1e-4 + 1e-4 * (*r as f64).abs() { + violations += 1; } - let mae = sum_abs / reference.len() as f64; - abs_dev.sort_by(|a, b| a.partial_cmp(b).expect("finite")); - let pos = 0.99 * (reference.len() as f64 - 1.0); - let lo = pos.floor() as usize; - let hi = pos.ceil() as usize; - let p99 = abs_dev[lo] + (pos - lo as f64) * (abs_dev[hi] - abs_dev[lo]); - eprintln!( - "{label}: max={max_abs:e} mae={mae:e} p99={p99:e} violations={violations} \ - inference_s={:.1} load_s={:.3} disagreement={:e}", - result.elapsed_secs, backend.load_secs, result.disagreement_mean_variance - ); - assert_eq!(violations, 0, "{label}: elementwise violations"); - assert!(mae <= 1e-5, "{label}: mae {mae:e}"); - assert!(p99 <= 1e-4, "{label}: p99 {p99:e}"); - eprintln!("{label}: WALL_S={wall_s:.1}"); + abs_dev.push(err); } + let mae = sum_abs / reference.len() as f64; + abs_dev.sort_by(|a, b| a.partial_cmp(b).expect("finite")); + let pos = 0.99 * (reference.len() as f64 - 1.0); + let lo = pos.floor() as usize; + let hi = pos.ceil() as usize; + let p99 = abs_dev[lo] + (pos - lo as f64) * (abs_dev[hi] - abs_dev[lo]); + eprintln!( + "{label}: max={max_abs:e} mae={mae:e} p99={p99:e} violations={violations} \ + inference_s={:.1} load_s={:.3}", + result.elapsed_secs, backend.load_secs + ); + assert_eq!(violations, 0, "{label}: elementwise violations"); + assert!(mae <= 1e-5, "{label}: mae {mae:e}"); + assert!(p99 <= 1e-4, "{label}: p99 {p99:e}"); + eprintln!("{label}: WALL_S={wall_s:.1}"); } } diff --git a/src-tauri/src/nonlocal_install.rs b/src-tauri/src/nonlocal_install.rs index 0a1ddf21f7..07947adc26 100644 --- a/src-tauri/src/nonlocal_install.rs +++ b/src-tauri/src/nonlocal_install.rs @@ -26,12 +26,12 @@ use std::{ /// Public Hugging Face model repository carrying the runtime artifacts. pub(crate) const HUB_REPO: &str = "sheldonxxxx/RapidRAW-Nonlocal-Denoise"; /// Immutable Hub commit the application downloads from (never `main`). -pub(crate) const HUB_REVISION: &str = "2c17892faf67b0511fd9debe2e59cb3cd75a850f"; +pub(crate) const HUB_REVISION: &str = "af6bf36ca55591d103e465200efc4f13b7112a7b"; /// Human release tag pointing at [`HUB_REVISION`]; informational only. -pub(crate) const HUB_TAG: &str = "nonlocal-v1"; +pub(crate) const HUB_TAG: &str = "nonlocal-v2"; /// Hard-pinned SHA-256 of `distribution.json` at [`HUB_REVISION`]. pub(crate) const DISTRIBUTION_SHA256: &str = - "de9dddc3e4d2dfbcb55227479da4cb073430f707445fc00cc38b822970600bb5"; + "91311f274f504ad98e26aad2178349e922bb44d9b793816029340411f1648bb1"; pub(crate) const DISTRIBUTION_PATH: &str = "distribution.json"; pub(crate) const MAX_DISTRIBUTION_BYTES: u64 = 1024 * 1024; /// Per-file sanity cap (~6x the largest accepted artifact). @@ -988,6 +988,30 @@ mod tests { } } + /// End-to-end install of the ONNX bundle from the pinned Hub revision into + /// a temporary workspace, then validate and status from the installed path. + /// Requires public network access to huggingface.co. + #[test] + #[ignore = "Downloads the ~86MB ONNX bundle from the pinned Hub revision"] + fn install_onnx_from_hub_into_temp_workspace() { + if std::env::var_os("RAPIDRAW_NONLOCAL_BUNDLE").is_some() { + eprintln!("SKIP: RAPIDRAW_NONLOCAL_BUNDLE is set in this environment"); + return; + } + let models = tempfile::tempdir().unwrap(); + let report = block_on(install(models.path(), Provider::Cpu)).unwrap(); + assert_eq!(report["installed"], json!(true)); + assert_eq!(report["backend"], json!("native-onnx-v1")); + assert_eq!(report["hub_revision"], json!(HUB_REVISION)); + let (bundle, via_env) = resolve_bundle(Some(models.path()), Provider::Cpu).unwrap(); + assert!(!via_env); + assert_eq!(bundle, install_dir(models.path(), Provider::Cpu)); + validate_resolved_bundle(Provider::Cpu, &bundle).unwrap(); + let status = status_for(models.path(), Provider::Cpu); + assert_eq!(status["ready"], json!(true)); + assert_eq!(status["provider_requested"], json!("cpu")); + } + #[test] fn failed_install_leaves_previous_bundle_untouched() { // A "valid" previous install is simulated by a marker tree; the diff --git a/src-tauri/src/nonlocal_onnx.rs b/src-tauri/src/nonlocal_onnx.rs index d0a49b39b8..de977923c7 100644 --- a/src-tauri/src/nonlocal_onnx.rs +++ b/src-tauri/src/nonlocal_onnx.rs @@ -8,10 +8,9 @@ //! //! The host-side conditioning pipeline ports //! `denoise/rapidraw_denoise/pipeline.py` + `noise.py` -//! (stratified-haar-irls-v1, 320/64 tiling, RGBG order, balanced=e1, -//! maximum=e4) into Rust. The Python package remains only as a -//! development/reference oracle and research tooling, never as a -//! production runtime dependency. +//! (stratified-haar-irls-v1, 320/40 tiling, RGBG order, one pass per tile) +//! into Rust. The Python package remains only as a development/reference +//! oracle and research tooling, never as a production runtime dependency. use crate::denoising::DenoiseControl; use anyhow::{Context, Result, bail, ensure}; @@ -30,9 +29,14 @@ use std::{ time::Instant, }; -/// Accepted ONNX bundle model hash (model.onnx, packed sampler, FP32). +/// Accepted ONNX bundle model hash (`model.onnx`, packed sampler, FP32). It +/// is the reference export with the zero-bias nodes dropped and every dense +/// and grouped 1x1 convolution expressed as a MatMul +/// (`rapidraw_denoise.onnx_rewrite`, version `gemm-1x1-v2`), which passes the +/// frozen strict-FP32 gates against the reference export. The validated hash +/// feeds cache identity and provenance. pub(crate) const ONNX_MODEL_SHA: &str = - "df7f21ddbdfecd0984896a902623f75aa883d25b5f862c49d8c2b3f63a7dd2d9"; + "ba788a97f247058ea03d14bbf2f289adc8f4666039f27d0d530e6f67d23dd73c"; /// Pinned source checkpoint lineage the bundle must record. pub(crate) const SOURCE_CHECKPOINT_SHA: &str = "c16747d852b93a95908792cdbac901f89cca98e35b91ea7db6214de42fbd3cad"; @@ -41,11 +45,9 @@ pub(crate) const BACKEND_GENERATION: &str = "native-onnx-v1"; pub(crate) const ALGORITHM: &str = "nonlocal-raw-v1"; pub(crate) const TILE: usize = 320; -pub(crate) const HALO: usize = 64; +pub(crate) const HALO: usize = 40; pub(crate) const CORE: usize = TILE - 2 * HALO; const BLOCK: usize = 32; -const ENSEMBLE_BALANCED: u32 = 1; -const ENSEMBLE_MAXIMUM: u32 = 4; /// Between-tile wall-clock guard, matching the retired worker's 30 minutes. /// An active ORT call is not abortable; cancellation lands on tile boundaries. const JOB_WALL_LIMIT_SECS: u64 = 1800; @@ -298,7 +300,7 @@ pub(crate) fn validate_bundle(bundle: &Path) -> Result { manifest_i64(&manifest, "tile")? == TILE as i64 && manifest_i64(&manifest, "halo")? == HALO as i64 && manifest_i64(&manifest, "retained_core")? == CORE as i64, - "NONLOCAL_BUNDLE_INVALID: tile/halo/retained_core must be 320/64/192" + "NONLOCAL_BUNDLE_INVALID: tile/halo/retained_core must be 320/40/240" ); ensure!( manifest.get("cfa_order").and_then(|v| v.as_str()) == Some("RGBG"), @@ -341,11 +343,12 @@ pub(crate) fn validate_bundle(bundle: &Path) -> Result { ); let actual = hash_file(&model_path)?; ensure!( - actual == ONNX_MODEL_SHA, + actual == recorded, "NONLOCAL_BUNDLE_INVALID: model file hash mismatch (modified graph rejected)" ); Ok(BundleContract { model_path, + model_sha256: actual, input_name, output_name, }) @@ -354,6 +357,8 @@ pub(crate) fn validate_bundle(bundle: &Path) -> Result { #[derive(Debug)] pub(crate) struct BundleContract { pub model_path: PathBuf, + /// Hash of the validated model file; always `ONNX_MODEL_SHA`. + pub model_sha256: String, pub input_name: String, pub output_name: String, } @@ -375,7 +380,7 @@ fn expect_io_type(kind: &str, name: &str, ty: &ValueType, dims: &[usize]) -> Res } } -/// One ORT session per denoise job, reused for all tiles and ensemble passes. +/// One ORT session per denoise job, reused for all tiles. pub(crate) struct Backend { session: Session, pub input_name: String, @@ -880,91 +885,6 @@ pub(crate) fn conditioned_input( out } -/// Exact Welford update, mirroring pipeline.py: -/// delta = candidate - mean; mean += delta/(i+1); m2 += delta*(candidate-mean). -fn welford_update(mean: &mut [f32], m2: &mut [f32], candidate: &[f32], pass: usize) { - let denom = pass as f32 + 1.0; - for i in 0..mean.len() { - let delta = candidate[i] - mean[i]; - mean[i] += delta / denom; - m2[i] += delta * (candidate[i] - mean[i]); - } -} -/// Rotate CHW counter-clockwise `k` quarter turns, then flip width if asked. -/// Mirrors `pipeline.transform` (rot90 over (-2,-1), optional `[..., ::-1]`). -pub(crate) fn transform_chw( - src: &[f32], - channels: usize, - height: usize, - width: usize, - pass: usize, -) -> (Vec, usize, usize) { - let mut data = src.to_vec(); - let (mut h, mut w) = (height, width); - for _ in 0..(pass % 4) { - let mut next = vec![0f32; data.len()]; - for c in 0..channels { - for y in 0..w { - for x in 0..h { - // Counter-clockwise quarter turn: out[y,x] = in[x, W-1-y]. - next[(c * w + y) * h + x] = data[(c * h + x) * w + (w - 1 - y)]; - } - } - } - data = next; - std::mem::swap(&mut h, &mut w); - } - if pass >= 4 { - let mut next = vec![0f32; data.len()]; - for c in 0..channels { - for y in 0..h { - for x in 0..w { - next[(c * h + y) * w + x] = data[(c * h + y) * w + (w - 1 - x)]; - } - } - } - data = next; - } - (data, h, w) -} - -/// Inverse of [`transform_chw`]: unflip, then rotate clockwise. -pub(crate) fn inverse_transform_chw( - src: &[f32], - channels: usize, - height: usize, - width: usize, - pass: usize, -) -> (Vec, usize, usize) { - let mut data = src.to_vec(); - let (mut h, mut w) = (height, width); - if pass >= 4 { - let mut next = vec![0f32; data.len()]; - for c in 0..channels { - for y in 0..h { - for x in 0..w { - next[(c * h + y) * w + x] = data[(c * h + y) * w + (w - 1 - x)]; - } - } - } - data = next; - } - for _ in 0..(pass % 4) { - let mut next = vec![0f32; data.len()]; - for c in 0..channels { - for y in 0..w { - for x in 0..h { - // Clockwise quarter turn: out[y,x] = in[H-1-x, y]. - next[(c * w + y) * h + x] = data[(c * h + (h - 1 - x)) * w + y]; - } - } - } - data = next; - std::mem::swap(&mut h, &mut w); - } - (data, h, w) -} - fn reflect_index(p: usize, halo: usize, n: usize, pad_after: usize) -> usize { let total = halo + n + pad_after; debug_assert!(p < total); @@ -1010,27 +930,20 @@ pub(crate) struct PackResult { /// Denoised [4,H,W] plane-major prediction. pub prediction: Vec, pub profile: NoiseProfile, - pub disagreement_mean_variance: f64, pub elapsed_secs: f64, } -/// Full conditioning + tiled native inference + Welford ensemble driver, -/// shared by the ONNX and direct CoreML backends through [`TilePredictor`]. -/// `packed` is plane-major [4,H,W]; ensemble is 1 (balanced) or 4 (maximum). -/// Cancellation and the 30-minute wall guard land on tile boundaries: an -/// active backend prediction is not abortable. +/// Full conditioning + tiled native inference, shared by the ONNX and direct +/// CoreML backends through [`TilePredictor`]. `packed` is plane-major +/// [4,H,W]. Cancellation and the 30-minute wall guard land on tile +/// boundaries: an active backend prediction is not abortable. pub(crate) fn denoise_packed( backend: &mut P, packed: &[f32], height: usize, width: usize, - ensemble: u32, control: &DenoiseControl, ) -> Result { - ensure!( - ensemble == ENSEMBLE_BALANCED || ensemble == ENSEMBLE_MAXIMUM, - "NONLOCAL_ENSEMBLE_INVALID: ensemble must be 1 or 4" - ); ensure!( packed.len() == 4 * height * width, "NONLOCAL_INPUT_INVALID: packed shape mismatch" @@ -1039,70 +952,43 @@ pub(crate) fn denoise_packed( control.check().map_err(anyhow::Error::msg)?; let profile = estimate_noise(packed, height, width)?; let conditioned = conditioned_input(packed, height, width, &profile); - let plane = height * width; - let mut mean = vec![0f32; 4 * plane]; - let mut m2 = vec![0f32; 4 * plane]; - for pass in 0..ensemble as usize { - control.check().map_err(anyhow::Error::msg)?; - let (rotated, rh, rw) = transform_chw(&conditioned, 8, height, width, pass); - let (padded, ph, pw, ny, nx) = reflect_pad(&rotated, 8, rh, rw, TILE, HALO); - let total = ny * nx; - let mut out = vec![0f32; 4 * rh * rw]; - for ty in 0..ny { - for tx in 0..nx { - control.check().map_err(anyhow::Error::msg)?; - if started.elapsed().as_secs() > JOB_WALL_LIMIT_SECS { - bail!("NONLOCAL_TIMEOUT: inference exceeded 30 minutes"); - } - let mut tile = vec![0f32; 8 * TILE * TILE]; - for c in 0..8 { - for y in 0..TILE { - let src_row = &padded[(c * ph + ty * CORE + y) * pw + tx * CORE..][..TILE]; - tile[(c * TILE + y) * TILE..(c * TILE + y) * TILE + TILE] - .copy_from_slice(src_row); - } + let (padded, ph, pw, ny, nx) = reflect_pad(&conditioned, 8, height, width, TILE, HALO); + let mut out = vec![0f32; 4 * height * width]; + for ty in 0..ny { + for tx in 0..nx { + control.check().map_err(anyhow::Error::msg)?; + if started.elapsed().as_secs() > JOB_WALL_LIMIT_SECS { + bail!("NONLOCAL_TIMEOUT: inference exceeded 30 minutes"); + } + let mut tile = vec![0f32; 8 * TILE * TILE]; + for c in 0..8 { + for y in 0..TILE { + let src_row = &padded[(c * ph + ty * CORE + y) * pw + tx * CORE..][..TILE]; + tile[(c * TILE + y) * TILE..(c * TILE + y) * TILE + TILE] + .copy_from_slice(src_row); } - let pred = backend.predict_tile(&tile)?; - let (cy, cx) = (rh - ty * CORE, rw - tx * CORE); - let (cy, cx) = (cy.min(CORE), cx.min(CORE)); - for c in 0..4 { - for y in 0..cy { - let src_row = &pred[(c * TILE + HALO + y) * TILE + HALO..][..cx]; - out[((c * rh + ty * CORE + y) * rw + tx * CORE)..][..cx] - .copy_from_slice(src_row); - } + } + let pred = backend.predict_tile(&tile)?; + let (cy, cx) = (height - ty * CORE, width - tx * CORE); + let (cy, cx) = (cy.min(CORE), cx.min(CORE)); + for c in 0..4 { + for y in 0..cy { + let src_row = &pred[(c * TILE + HALO + y) * TILE + HALO..][..cx]; + out[((c * height + ty * CORE + y) * width + tx * CORE)..][..cx] + .copy_from_slice(src_row); } - let done = (pass * total + ty * nx + tx + 1) as f32; - let all = (ensemble as usize * total) as f32; - control.report(0.1 + done / all * 0.8, "Nonlocal inference"); } + let done = (ty * nx + tx + 1) as f32; + control.report(0.1 + done / (ny * nx) as f32 * 0.8, "Nonlocal inference"); } - let (candidate, ch, cw) = inverse_transform_chw(&out, 4, rh, rw, pass); - ensure!( - ch == height && cw == width, - "NONLOCAL_SHAPE_INVALID: ensemble pass changed image geometry" - ); - // Exact Welford update (see welford_update). - welford_update(&mut mean, &mut m2, &candidate, pass); - control.check().map_err(anyhow::Error::msg)?; - } - let denom = (ensemble as usize - 1).max(1) as f32; - let mut disag_sum = 0f64; - for v in &m2 { - ensure!( - v.is_finite(), - "NONLOCAL_OUTPUT_INVALID: nonfinite ensemble result" - ); - disag_sum += (*v / denom) as f64; } ensure!( - mean.iter().all(|v| v.is_finite()), + out.iter().all(|v| v.is_finite()), "NONLOCAL_OUTPUT_INVALID: nonfinite inference result" ); Ok(PackResult { - prediction: mean, + prediction: out, profile, - disagreement_mean_variance: disag_sum / m2.len() as f64, elapsed_secs: started.elapsed().as_secs_f64(), }) } @@ -1246,53 +1132,6 @@ mod tests { assert!(estimate_noise(&vec![1.0f32; 4 * 128 * 128], 128, 128).is_err()); } - #[test] - fn rotations_match_numpy_and_roundtrip() { - // Oracle values from rapidraw_denoise.pipeline on a seeded (2,5,6). - // Here we verify shape preservation, exact spot values from a - // fixed input, and lossless roundtrips for all 8 passes. - let src: Vec = (0..2 * 5 * 6).map(|i| i as f32 * 0.25 + 0.125).collect(); - for pass in 0..8 { - let (t, th, tw) = transform_chw(&src, 2, 5, 6, pass); - let (b, bh, bw) = inverse_transform_chw(&t, 2, th, tw, pass); - assert_eq!((bh, bw), (5, 6)); - assert_eq!(b, src, "pass {pass} roundtrip must be lossless"); - } - // CCW quarter turn: out[y,x] = in[x, W-1-y]. - let at = |c: usize, y: usize, x: usize, h: usize, w: usize| (c * h + y) * w + x; - let (t1, h1, w1) = transform_chw(&src, 2, 5, 6, 1); - assert_eq!((h1, w1), (6, 5)); - assert_eq!(t1[at(0, 0, 0, 6, 5)], src[at(0, 0, 5, 5, 6)]); - assert_eq!(t1[at(1, 5, 4, 6, 5)], src[at(1, 4, 0, 5, 6)]); - // Flip pass: out = in reversed along width. - let (t4, _, _) = transform_chw(&src, 2, 5, 6, 4); - assert_eq!(t4[at(0, 2, 0, 5, 6)], src[at(0, 2, 5, 5, 6)]); - } - - #[test] - fn welford_matches_reference_arithmetic() { - // Oracle: pipeline.py Welford over four candidates. - let cands = [ - vec![1f32, 2., 3., 4.], - vec![2., 3., 4., 5.], - vec![0., 1., 5., 3.], - vec![3., 3., 3., 3.], - ]; - let mut mean = vec![0f32; 4]; - let mut m2 = vec![0f32; 4]; - for (i, c) in cands.iter().enumerate() { - welford_update(&mut mean, &mut m2, c, i); - } - for (m, e) in mean.iter().zip([1.5f32, 2.25, 3.75, 3.75]) { - assert!((m - e).abs() < 1e-6, "{m} vs {e}"); - } - for (v, e) in m2.iter().zip([5.0f32, 2.75, 2.75, 2.75]) { - assert!((v - e).abs() < 1e-5, "{v} vs {e}"); - } - let disag: f64 = m2.iter().map(|v| (*v / 3.0) as f64).sum::() / 4.0; - assert!((disag - 1.1041666269302368).abs() < 1e-6); - } - #[test] fn reflect_pad_matches_numpy() { // Oracle: np.pad 1ch 5x6 with tile=8 halo=2 -> 12x12. @@ -1413,7 +1252,7 @@ mod tests { "manifest_version": 1, "checkpoint_pinned_match": true, "source_checkpoint_sha256": SOURCE_CHECKPOINT_SHA, - "precision": "fp32", "tile": 320, "halo": 64, "retained_core": 192, + "precision": "fp32", "tile": 320, "halo": 40, "retained_core": 240, "cfa_order": "RGBG", "input_name": "raw_with_noise", "input_shape": [1, 8, 320, 320], "input_dtype": "float32", @@ -1433,6 +1272,37 @@ mod tests { assert!(err.contains("lineage"), "{err}"); } + #[test] + fn bundle_validation_binds_file_to_the_pinned_hash() { + let root = tempfile::tempdir().unwrap(); + let manifest = |sha: &str| { + json!({ + "manifest_version": 1, + "checkpoint_pinned_match": true, + "source_checkpoint_sha256": SOURCE_CHECKPOINT_SHA, + "precision": "fp32", "tile": 320, "halo": 40, "retained_core": 240, + "cfa_order": "RGBG", + "input_name": "raw_with_noise", "input_shape": [1, 8, 320, 320], + "input_dtype": "float32", + "output_name": "denoised_raw", "output_shape": [1, 4, 320, 320], + "output_dtype": "float32", + "model_sha256": sha, + }) + .to_string() + }; + std::fs::write(root.path().join("model.onnx"), b"decoy").unwrap(); + // A file that is not the pinned bytes never passes, even when the + // manifest records the pinned hash. + std::fs::write(root.path().join("manifest.json"), manifest(ONNX_MODEL_SHA)).unwrap(); + let err = validate_bundle(root.path()).unwrap_err().to_string(); + assert!(err.contains("model file hash mismatch"), "{err}"); + // A manifest recording the decoy's own hash is still not an accepted model. + let decoy = hex::encode(Sha256::digest(b"decoy")); + std::fs::write(root.path().join("manifest.json"), manifest(&decoy)).unwrap(); + let err = validate_bundle(root.path()).unwrap_err().to_string(); + assert!(err.contains("not the accepted bundle"), "{err}"); + } + #[test] fn nnls_matches_scipy_on_reference_subproblem() { // First IRLS weighted subproblem of golden channel 0, captured @@ -1541,4 +1411,63 @@ mod tests { } let _ = test_control(); } + + /// End-to-end tiling check at the production 320/40 geometry: the whole + /// conditioning + tiling + stitching driver against a reference produced + /// by `rapidraw_denoise.pipeline.denoise_with_predictor` (tile 320, halo + /// 40) on the same input through the same bundle on CPU. Input and + /// reference are plane-major little-endian f32 `[4,H,W]`; the geometry + /// comes from `RAPIDRAW_NONLOCAL_PHOTO_HEIGHT` and `..._WIDTH`. Gates match + /// the tile family: no elementwise violation at 1e-4 atol/rtol, MAE <= 1e-5, + /// p99 <= 1e-4. + #[test] + #[ignore = "Requires RAPIDRAW_NONLOCAL_BUNDLE, RAPIDRAW_NONLOCAL_PHOTO_{INPUT,REF,HEIGHT,WIDTH}; runs a full image on CPU"] + fn rust_cpu_full_photo_matches_reference_tiling() { + let var = |name: &str| std::env::var(name).unwrap_or_else(|_| panic!("Set {name}")); + let read_f32 = |path: String| -> Vec { + std::fs::read(path) + .unwrap() + .as_chunks::<4>() + .0 + .iter() + .map(|b| f32::from_le_bytes(*b)) + .collect() + }; + let bundle = PathBuf::from(var("RAPIDRAW_NONLOCAL_BUNDLE")); + let height: usize = var("RAPIDRAW_NONLOCAL_PHOTO_HEIGHT").parse().unwrap(); + let width: usize = var("RAPIDRAW_NONLOCAL_PHOTO_WIDTH").parse().unwrap(); + let packed = read_f32(var("RAPIDRAW_NONLOCAL_PHOTO_INPUT")); + let reference = read_f32(var("RAPIDRAW_NONLOCAL_PHOTO_REF")); + assert_eq!(packed.len(), 4 * height * width); + assert_eq!(reference.len(), packed.len()); + let contract = validate_bundle(&bundle).unwrap(); + let config = NativeConfig { + bundle, + provider: Provider::Cpu, + device_id: 0, + mem_limit_mb: None, + }; + let mut backend = open_backend(&config, &contract).unwrap(); + let result = denoise_packed(&mut backend, &packed, height, width, &test_control()).unwrap(); + let mut errors: Vec = Vec::with_capacity(reference.len()); + let mut violations = 0usize; + for (got, want) in result.prediction.iter().zip(&reference) { + let err = (*got as f64 - *want as f64).abs(); + if err > 1e-4 + 1e-4 * (*want as f64).abs() { + violations += 1; + } + errors.push(err); + } + let mae = errors.iter().sum::() / errors.len() as f64; + errors.sort_by(|a, b| a.partial_cmp(b).unwrap()); + let p99 = errors[(0.99 * (errors.len() - 1) as f64) as usize]; + eprintln!( + "full-photo parity: max={:e} mae={mae:e} p99={p99:e} violations={violations} tiles_seconds={:.1}", + errors[errors.len() - 1], + result.elapsed_secs + ); + assert_eq!(violations, 0, "elementwise violations"); + assert!(mae <= 1e-5, "mae {mae:e}"); + assert!(p99 <= 1e-4, "p99 {p99:e}"); + } } diff --git a/src-tauri/src/raw_denoise.rs b/src-tauri/src/raw_denoise.rs index a0d0b25976..daecc58a71 100644 --- a/src-tauri/src/raw_denoise.rs +++ b/src-tauri/src/raw_denoise.rs @@ -8,7 +8,7 @@ use crate::{ denoising::DenoiseControl, nonlocal_coreml::{self, COREML_BACKEND_GENERATION, COREML_PACKAGE_SHA}, nonlocal_onnx::{ - self, ALGORITHM, BACKEND_GENERATION, NativeConfig, ONNX_MODEL_SHA, SOURCE_CHECKPOINT_SHA, + self, ALGORITHM, BACKEND_GENERATION, HALO, NativeConfig, SOURCE_CHECKPOINT_SHA, TILE, }, }; use anyhow::{Context, Result, bail, ensure}; @@ -293,17 +293,16 @@ fn write_packed_input(path: &Path, values: &[f32]) -> Result<()> { /// different hash preimage, so they can never be accepted as native results. fn native_onnx_request( config: &NativeConfig, - packed_height: usize, - packed_width: usize, + model_sha256: &str, + (packed_height, packed_width): (usize, usize), input_sha: &str, source_sha: &str, - ensemble: u32, ort_build: &str, ) -> Value { json!({"protocol":2,"algorithm":ALGORITHM,"backend":BACKEND_GENERATION, - "model_sha256":ONNX_MODEL_SHA,"source_checkpoint_sha256":SOURCE_CHECKPOINT_SHA, + "model_sha256":model_sha256,"source_checkpoint_sha256":SOURCE_CHECKPOINT_SHA, "input_sha256":input_sha,"source_sha256":source_sha, - "shape":[4,packed_height,packed_width],"tile":320,"halo":64,"ensemble":ensemble, + "shape":[4,packed_height,packed_width],"tile":TILE,"halo":HALO, "provider":config.provider_as_str(), "device_id":config.device_id_or_null(), "graph_optimization":false,"tf32":false, @@ -377,13 +376,12 @@ fn native_coreml_request( packed_width: usize, input_sha: &str, source_sha: &str, - ensemble: u32, ) -> Value { json!({"protocol":3,"algorithm":ALGORITHM,"backend":COREML_BACKEND_GENERATION, "package_sha256":COREML_PACKAGE_SHA,"package_tree_sha256":tree_sha256, "source_checkpoint_sha256":SOURCE_CHECKPOINT_SHA, "input_sha256":input_sha,"source_sha256":source_sha, - "shape":[4,packed_height,packed_width],"tile":320,"halo":64,"ensemble":ensemble, + "shape":[4,packed_height,packed_width],"tile":TILE,"halo":HALO, "provider":"coreml","compute_units":nonlocal_coreml::COREML_COMPUTE_UNITS, "allow_low_precision_accumulation":false, "os":nonlocal_coreml::runtime_identity(), @@ -455,8 +453,8 @@ fn prediction_coreml( /// generation provider comes only from the hash-validated receipt (which /// `prediction()` already required to equal the fresh request); the noise /// profile is recomputed deterministically from the current hash-verified -/// packed input. No ONNX inference runs here, and timing/disagreement fields -/// are left absent because they were never persisted. +/// packed input. No ONNX inference runs here, and the timing field is left +/// absent because it was never persisted. fn cache_hit_provenance( request: &Value, receipt: &Value, @@ -511,7 +509,6 @@ fn cache_hit_fragment( let profile = nonlocal_onnx::estimate_noise(packed, height, width)?; Ok(json!({ "provider_actual": provider, - "ensemble": request["ensemble"].clone(), "noise_profile": profile.to_json(), })) } @@ -557,10 +554,9 @@ impl Output { /// Truthful provenance for a strength/intensity-zero job: no bundle or /// provider configuration was consulted and no inference ran, so no backend, /// runtime, model/package or cache identity may be claimed. -fn skipped_provenance(source_sha: &str, quality: &str, packed: &Packed) -> Value { +fn skipped_provenance(source_sha: &str, packed: &Packed) -> Value { json!({"algorithm":ALGORITHM, "inference_executed":false,"reason":"intensity-zero", - "quality":quality, "source_sha256":source_sha,"packed_shape":[4,packed.height,packed.width],"output":"float32-bayer-dng", "normalization":"per-CFA black/white; signed noise estimation; unclipped highlights preserved"}) } @@ -570,13 +566,8 @@ pub(crate) fn denoise( source_sha: &str, root: &Path, strength: f32, - quality: &str, control: &DenoiseControl, ) -> Result { - ensure!( - ["balanced", "maximum"].contains(&quality), - "Invalid Nonlocal quality" - ); ensure!( strength.is_finite() && (0.0..=1.0).contains(&strength), "Invalid Nonlocal strength" @@ -639,19 +630,19 @@ pub(crate) fn denoise( "compute_units":nonlocal_coreml::COREML_COMPUTE_UNITS, "allow_low_precision_accumulation":false, "os":nonlocal_coreml::runtime_identity(), - "quality":quality,"cache_hit":false, + "cache_hit":false, "source_sha256":source_sha,"packed_shape":[4,packed.height,packed.width],"output":"float32-bayer-dng", "normalization":"per-CFA black/white; signed noise estimation; unclipped highlights preserved"}) } _ => json!({"algorithm":ALGORITHM,"backend":BACKEND_GENERATION, - "runtime":"onnxruntime","model_sha256":ONNX_MODEL_SHA, + "runtime":"onnxruntime", "source_checkpoint_sha256":SOURCE_CHECKPOINT_SHA, - "quality":quality,"cache_hit":false, + "cache_hit":false, "source_sha256":source_sha,"packed_shape":[4,packed.height,packed.width],"output":"float32-bayer-dng", "normalization":"per-CFA black/white; signed noise estimation; unclipped highlights preserved"}), } } else { - skipped_provenance(source_sha, quality, &packed) + skipped_provenance(source_sha, &packed) }; if strength > 0. { let config = configuration_with_models(Some(&models_dir))?; @@ -667,7 +658,6 @@ pub(crate) fn denoise( packed.width, &hash_file(&input_path)?, source_sha, - if quality == "maximum" { 4 } else { 1 }, ); let request_bytes = serde_json::to_vec(&request)?; let request_sha = hex::encode(Sha256::digest(&request_bytes)); @@ -693,14 +683,12 @@ pub(crate) fn denoise( "NONLOCAL_CACHE_INVALID: Remove this corrupt cache entry before retrying", )?; provenance["provider_actual"] = fragment["provider_actual"].clone(); - provenance["ensemble"] = fragment["ensemble"].clone(); provenance["noise_profile"] = fragment["noise_profile"].clone(); hit } else { fs::write(directory.path().join("request.json"), &request_bytes)?; control.check().map_err(anyhow::Error::msg)?; - // One compiled MLModel per job, reused for all tiles and - // passes. macOS only; other platforms fail closed in + // One compiled MLModel per job, reused for all tiles. macOS only; other platforms fail closed in // configuration() before reaching here. #[cfg(target_os = "macos")] let outcome: Result<(Vec, Value)> = (|| { @@ -713,7 +701,6 @@ pub(crate) fn denoise( &packed.values, packed.height, packed.width, - if quality == "maximum" { 4 } else { 1 }, control, )?; control.check().map_err(anyhow::Error::msg)?; @@ -739,9 +726,6 @@ pub(crate) fn denoise( )?; control.check().map_err(anyhow::Error::msg)?; provenance["noise_profile"] = result.profile.to_json(); - provenance["ensemble"] = request["ensemble"].clone(); - provenance["disagreement_mean_variance"] = - json!(result.disagreement_mean_variance); provenance["inference_seconds"] = json!(result.elapsed_secs); let staging = tempfile::Builder::new() .prefix("cache-") @@ -766,6 +750,7 @@ pub(crate) fn denoise( blend(&mut raw, &packed, &values, strength)?; } else { let contract = nonlocal_onnx::validate_bundle(&config.bundle)?; + provenance["model_sha256"] = json!(contract.model_sha256); provenance["provider_requested"] = json!(config.provider_as_str()); provenance["graph_optimization"] = json!(false); provenance["tf32"] = json!(false); @@ -774,11 +759,10 @@ pub(crate) fn denoise( write_packed_input(&input_path, &packed.values)?; let request = native_onnx_request( &config, - packed.height, - packed.width, + &contract.model_sha256, + (packed.height, packed.width), &hash_file(&input_path)?, source_sha, - if quality == "maximum" { 4 } else { 1 }, &nonlocal_onnx::ort_identity(), ); let request_bytes = serde_json::to_vec(&request)?; @@ -805,13 +789,12 @@ pub(crate) fn denoise( "NONLOCAL_CACHE_INVALID: Remove this corrupt cache entry before retrying", )?; provenance["provider_actual"] = fragment["provider_actual"].clone(); - provenance["ensemble"] = fragment["ensemble"].clone(); provenance["noise_profile"] = fragment["noise_profile"].clone(); hit } else { fs::write(directory.path().join("request.json"), &request_bytes)?; control.check().map_err(anyhow::Error::msg)?; - // One ORT session per job, reused for all tiles and passes. + // One ORT session per job, reused for all tiles. let mut backend = nonlocal_onnx::open_backend(&config, &contract)?; provenance["provider_actual"] = json!(backend.provider); let result = nonlocal_onnx::denoise_packed( @@ -819,7 +802,6 @@ pub(crate) fn denoise( &packed.values, packed.height, packed.width, - if quality == "maximum" { 4 } else { 1 }, control, )?; control.check().map_err(anyhow::Error::msg)?; @@ -827,7 +809,7 @@ pub(crate) fn denoise( write_packed_input(&prediction_path, &result.prediction)?; let receipt = json!({"protocol":2,"algorithm":ALGORITHM,"backend":BACKEND_GENERATION, "shape":[4,packed.height,packed.width], - "input_sha256":request["input_sha256"],"model_sha256":ONNX_MODEL_SHA, + "input_sha256":request["input_sha256"],"model_sha256":request["model_sha256"], "source_checkpoint_sha256":SOURCE_CHECKPOINT_SHA, "request_sha256":request_sha,"prediction_sha256":hash_file(&prediction_path)?, "provider":request["provider"],"device_id":request["device_id"], @@ -844,8 +826,6 @@ pub(crate) fn denoise( )?; control.check().map_err(anyhow::Error::msg)?; provenance["noise_profile"] = result.profile.to_json(); - provenance["ensemble"] = request["ensemble"].clone(); - provenance["disagreement_mean_variance"] = json!(result.disagreement_mean_variance); provenance["inference_seconds"] = json!(result.elapsed_secs); let staging = tempfile::Builder::new() .prefix("cache-") @@ -877,6 +857,7 @@ pub(crate) fn denoise( #[cfg(test)] mod tests { use super::*; + use crate::nonlocal_onnx::ONNX_MODEL_SHA; use rawler::{ CFA, decoders::Camera, @@ -1071,7 +1052,14 @@ mod tests { device_id: 0, mem_limit_mb: None, }; - let request = native_onnx_request(&config, H, W, "input-sha", "source-sha", 1, "test"); + let request = native_onnx_request( + &config, + ONNX_MODEL_SHA, + (H, W), + "input-sha", + "source-sha", + "test", + ); let request_sha = hex::encode(Sha256::digest(serde_json::to_vec(&request).unwrap())); let root = tempfile::tempdir().unwrap(); let entry = root.path().join(&request_sha); @@ -1093,7 +1081,6 @@ mod tests { let validated = prediction(&entry, &request, &request_sha, packed.len()).unwrap(); let fragment = cache_hit_provenance(&request, &validated.1, &packed, H, W).unwrap(); assert_eq!(fragment["provider_actual"], json!("cpu")); - assert_eq!(fragment["ensemble"], json!(1)); assert_eq!( fragment["noise_profile"]["method"], json!("stratified-haar-irls-v1") @@ -1114,7 +1101,23 @@ mod tests { json!(SOURCE_CHECKPOINT_SHA) ); assert!(!fragment.get("inference_seconds").is_some()); - assert!(!fragment.get("disagreement_mean_variance").is_some()); + } + #[test] + fn model_hash_is_part_of_native_cache_identity() { + // A different graph is arithmetically equivalent at best, never + // bit-identical, so its predictions must not share a cache entry. + let config = NativeConfig { + bundle: PathBuf::from("/bundle"), + provider: crate::nonlocal_onnx::Provider::Cuda, + device_id: 0, + mem_limit_mb: None, + }; + let sha = |model: &str| { + let request = native_onnx_request(&config, model, (8, 16), "input", "source", "ort"); + assert_eq!(request["model_sha256"], json!(model)); + hex::encode(Sha256::digest(serde_json::to_vec(&request).unwrap())) + }; + assert_ne!(sha(nonlocal_onnx::ONNX_MODEL_SHA), sha("other-model-sha")); } #[test] fn coreml_cache_identity_is_isolated_from_onnx() { @@ -1126,8 +1129,15 @@ mod tests { device_id: 0, mem_limit_mb: None, }; - let onnx = native_onnx_request(&config, 8, 16, "input-sha", "source-sha", 1, "test-ort"); - let coreml = native_coreml_request("tree-sha", 8, 16, "input-sha", "source-sha", 1); + let onnx = native_onnx_request( + &config, + ONNX_MODEL_SHA, + (8, 16), + "input-sha", + "source-sha", + "test-ort", + ); + let coreml = native_coreml_request("tree-sha", 8, 16, "input-sha", "source-sha"); assert_eq!(onnx["protocol"], json!(2)); assert_eq!(coreml["protocol"], json!(3)); assert_eq!(onnx["backend"], json!(BACKEND_GENERATION)); @@ -1144,7 +1154,8 @@ mod tests { assert_eq!(coreml["compute_units"], json!("all")); assert_eq!(coreml["allow_low_precision_accumulation"], json!(false)); assert!(coreml["os"].as_str().unwrap().starts_with("coreml/")); - assert_eq!(coreml["ensemble"], json!(1)); + assert_eq!(coreml["tile"], json!(320)); + assert_eq!(coreml["halo"], json!(40)); // Request bytes (and therefore cache keys) differ across families. assert_ne!( hex::encode(Sha256::digest(serde_json::to_vec(&onnx).unwrap())), @@ -1154,7 +1165,7 @@ mod tests { #[test] fn coreml_cache_rejects_cross_family_and_corrupt_receipts() { let root = tempfile::tempdir().unwrap(); - let request = native_coreml_request("tree-sha", 8, 8, "input-sha", "source-sha", 1); + let request = native_coreml_request("tree-sha", 8, 8, "input-sha", "source-sha"); let request_sha = hex::encode(Sha256::digest(serde_json::to_vec(&request).unwrap())); let pixels = root.path().join("prediction.f32"); fs::write(&pixels, vec![0_u8; 4 * 8 * 8 * 4]).unwrap(); @@ -1235,7 +1246,7 @@ mod tests { } } } - let request = native_coreml_request("tree-sha", H, W, "input-sha", "source-sha", 4); + let request = native_coreml_request("tree-sha", H, W, "input-sha", "source-sha"); let request_sha = hex::encode(Sha256::digest(serde_json::to_vec(&request).unwrap())); let root = tempfile::tempdir().unwrap(); let entry = root.path().join(&request_sha); @@ -1252,7 +1263,6 @@ mod tests { let validated = prediction_coreml(&entry, &request, &request_sha, packed.len()).unwrap(); let fragment = cache_hit_provenance_coreml(&request, &validated.1, &packed, H, W).unwrap(); assert_eq!(fragment["provider_actual"], json!("coreml")); - assert_eq!(fragment["ensemble"], json!(4)); assert_eq!( fragment["noise_profile"]["method"], json!("stratified-haar-irls-v1") @@ -1274,11 +1284,10 @@ mod tests { // zero-strength provenance that cannot claim a backend that never ran. let mut raw = fixture("RGGB", 0); let packed = prepare(&mut raw).unwrap(); - let value = skipped_provenance("source-sha", "balanced", &packed); + let value = skipped_provenance("source-sha", &packed); assert_eq!(value["algorithm"], json!(ALGORITHM)); assert_eq!(value["inference_executed"], json!(false)); assert_eq!(value["reason"], json!("intensity-zero")); - assert_eq!(value["quality"], json!("balanced")); assert_eq!(value["source_sha256"], json!("source-sha")); assert_eq!( value["packed_shape"], @@ -1365,7 +1374,14 @@ mod tests { device_id: 0, mem_limit_mb: None, }; - let request = native_onnx_request(&config, 8, 16, "input-sha", "source-sha", 1, "test"); + let request = native_onnx_request( + &config, + ONNX_MODEL_SHA, + (8, 16), + "input-sha", + "source-sha", + "test", + ); assert_eq!(request["protocol"], json!(2)); assert_eq!(request["backend"], json!(BACKEND_GENERATION)); assert_eq!(request["model_sha256"], json!(ONNX_MODEL_SHA)); @@ -1374,12 +1390,12 @@ mod tests { json!(SOURCE_CHECKPOINT_SHA) ); assert_eq!(request["shape"], json!([4, 8, 16])); + assert_eq!(request["tile"], json!(320)); + assert_eq!(request["halo"], json!(40)); assert_eq!(request["provider"], json!("cpu")); assert_eq!(request["device_id"], Value::Null); assert_eq!(request["graph_optimization"], json!(false)); assert_eq!(request["tf32"], json!(false)); - let max_request = native_onnx_request(&config, 8, 16, "input-sha", "source-sha", 4, "test"); - assert_eq!(max_request["ensemble"], json!(4)); // Request bytes must be valid JSON with stable key order for hashing. let bytes = serde_json::to_vec(&request).unwrap(); let parsed: Value = serde_json::from_slice(&bytes).unwrap(); @@ -1432,13 +1448,15 @@ mod tests { let input_sha = hash_file(&input_path).unwrap(); let config = configuration_with_models(None) .expect("Set RAPIDRAW_NONLOCAL_BUNDLE (+ provider) to export the production request"); + let model_sha256 = nonlocal_onnx::validate_bundle(&config.bundle) + .unwrap() + .model_sha256; let request = native_onnx_request( &config, - packed.height, - packed.width, + &model_sha256, + (packed.height, packed.width), &input_sha, &source_sha, - 1, &nonlocal_onnx::ort_identity(), ); let request_bytes = serde_json::to_vec(&request).unwrap(); @@ -1462,9 +1480,9 @@ mod tests { "white": packed.white, "crop_area": crop_before.map(|r| json!({"x": r.p.x, "y": r.p.y, "w": r.d.w, "h": r.d.h})), "active_area": active_before.map(|r| json!({"x": r.p.x, "y": r.p.y, "w": r.d.w, "h": r.d.h})), - "tile": 320, "halo": 64, "ensemble": 1, + "tile": TILE, "halo": HALO, "algorithm": ALGORITHM, "backend": BACKEND_GENERATION, - "model_sha256": ONNX_MODEL_SHA, + "model_sha256": model_sha256, "source_checkpoint_sha256": SOURCE_CHECKPOINT_SHA, "protocol": 2, }); fs::write(