Run rect_proj, combined_proj (BatchNorm), and the transformer encoder
in fixed chunk_size batches so every kernel invocation sees the same
tensor shape. This prevents cuDNN from re-benchmarking on every new
batch size, eliminating ~460ms latency spikes.
Also compute cdist per-image on actual regions instead of padding all
images to k_max, reducing wasted compute on the quadratic distance
matrix.
Made-with: Cursor
After loading weights, replaces the classifier Linear layer with one
whose output dimension is rounded up to the next multiple of 64.
Extra positions get bias=-inf so they never affect argmax or softmax.
- English (858->896): neutral on crops, ~2% on pages
- Multilingual (14247->14272): +1.5% crops/s, +6.5% pages/s
Vectorize post-processing format step and default to white padding
Ryan Chesler
0a7079c
Replace the per-sentence Python loop in _postprocess_batch_inner
(parse_relational_results + reorder_boxes) with batched numpy
operations: pad group indices to equal length, batch min/max across
all sentences at once, and inline the reading-order sort.
Reduces post_fmt from ~7ms to ~1ms per batch of 8 on OmniDocBench.
Also change default pad color from ImageNet mean [0.485, 0.456, 0.406]
to white [1.0, 1.0, 1.0] to match the v1 pipeline and training data.
improve install docs and process
Ryan Chesler
65abd8a
The .gitignore excludes *.so files, which caused hatchling to silently
drop the compiled CUDA extension from the wheel.
pyproject.toml: add artifacts glob so hatchling includes the .so
despite gitignore; add ninja to build hook deps for faster compilation.
hatch_build.py: set a platform-specific wheel tag (cpXY-cpXY-linux_*)
instead of py3-none-any when the extension is built.
README.md + quickstart.md: rewrite install instructions to recommend
installing torch first with the correct CUDA index URL, then using
--no-build-isolation to avoid CUDA version mismatches during the
extension build.
Remove unused code: training-only methods, dead model blocks, and test files
Strip inference-only package of ~1,500 lines of dead code:
- Remove tests/ (hardcoded paths, not useful for end users)
- Remove 11 unused nn.Module classes from blocks.py (ResBlock, Inception, GC/MAGC attention, mCReLU, etc.), keep only conv2d_block
- Remove training-only methods from encoders (base, recognizer, relational)
- Remove worker_messages.py (training message protocol, unused at inference)
- Remove rrect_to_quads.py (unused Python wrapper; C++ function imported directly)
- Remove NaNHook and GradReversal from fots_detector.py
- Remove find_checkpoint, tensor_all_reduce, tensor_all_gather from utils.py
- Simplify data_container.py to minimal base class
- Update quickstart.md for NemotronOCRV2
Made-with: Cursor
- Remove unused _SUPPORTS_PAD_COLOR and import inspect
- Rename _HF_* constants to _DEFAULT_*
- Update module and class docstrings for current features
- Gate all CUDA syncs and per-batch logging behind verbose_post flag
- Update example.py to use NemotronOCRV2 with --detector-only and
--skip-relational flags
- Add inference modes section to README.md and quickstart.md
Make per-batch CUDA syncs and logging optional via verbose_post
Timing syncs and per-batch logger.info calls are now gated behind
the verbose_post flag (default False). In production mode, no
torch.cuda.synchronize() calls are made between pipeline phases,
eliminating unnecessary GPU pipeline stalls.
Enable profiling with verbose_post=True to restore full per-phase
timing and detailed logging.
Add detector_only and skip_relational inference modes
Two new constructor flags for NemotronOCRV2:
- detector_only=True: loads only the detector, skips recognizer and
relational model entirely. Returns bounding boxes with confidence
scores but no text. ~20% faster crops, ~50% faster pages.
- skip_relational=True: loads detector + recognizer, skips the
relational model. Returns per-word text without reading order
grouping. ~8% faster crops
Both modes reduce GPU memory and model load time proportionally.
Lower default detector batch size to 8 for 16GB GPU compatibility
Computes argmax + softmax probability of the winning token inside
the recognizer chunk loop and discards the full [N, T, num_tokens]
logits tensor immediately, reducing peak VRAM.
Also converts region vertices to numpy during graph_build to avoid
per-region .cpu() calls in post-processing.
Add pre-NMS centerness + peak filter for faster inference
