Pressing code changes on `CorridorKeyModule`

Author: gtkaczCreated May 24, 2026Updated May 24, 2026

Hey guys, I did some quick CR on the CorridorKeyModule and found some relevant stuff:

  1. CHECKPOINT_DIR is clobbered to a CWD-relative path: backend.py:74 (line 22 has the correct absolute definition; line 74 reassigns to os.path.join("CorridorKeyModule", "checkpoints")). Breaks the installed package
  2. MLX auto-download branch is unreachable dead code: backend.py:85–119 is indented inside an except ImportError clause that returns "torch" on line 89, so lines 91–119 (cache lookup + urllib.request.urlretrieve + os.rename + final return "mlx") are never executed. Apple-Silicon users silently get the slower Torch path
  3. Autocast hardcoded dtype=torch.float16: inference_engine.py:472. The _is_rocm branch elsewhere in the file has no companion in autocast. fp16 in attention on ROCm/HIP is documented to produce NaN/Inf; bf16 is both more stable and 10–25% faster than the fp32 fallback users hit after the first NaN
  4. _clamp(x, min) silently ignores its min argument: color_utils.py:42–48. Both branches hard-code 0.0. Current call sites pass 0.0 by coincidence; any future caller gets silent wrong output
  5. _postprocess_torch docstring documents a non-existent async API: inference_engine.py:328–334. Describes a sync=True/False parameter and a PendingTransfer return type. Neither exists. The function always blocks and returns list[dict[str, np.ndarray]]
  6. process_frame return contract is undocumented: inference_engine.py:434–435. The Returns block lists {alpha, fg, comp} but omits processed (the primary linear-premultiplied-RGBA output), the dict | list[dict] polymorphism, and the comp is None when generate_comp=False rule
  7. apply_garbage_matte tensor branch is a pass stub: color_utils.py:200–202. Numpy path handles broadcasting; tensor path silently does nothing or produces wrong shapes. The GPU postprocess path therefore ignores garbage-matte input
  8. macos-latest runner is Intel x86_64: every pytest.mark.mlx test is structurally skipped. The MLX backend is dead in CI
  9. get_checkerboard_linear_torch @lru_cache(maxsize=4) holds GPU tensors strongly: ~95 MB / unique (w,h,device); unbounded VRAM growth in a long-running server with variable sizes
  10. connected_components torch.randperm(H*W, device=device) per call: 63 MB transient GPU alloc at 4K for a tiebreaker that does not need randomness
  11. clean_matte_torch runs up to 12 consecutive max_pool2d dilation iterations: collapsible to one dilated conv. ~5–8 ms savings per 4K frame
  12. tensor.to(device, non_blocking=True) on unpinned numpy silently synchronizes: non_blocking=True is only honored on pinned memory
  13. DecoderHead.forward four-way flatten/transpose/linear/transpose/view: ~3–5 ms per frame. Collapsible to a 1x1 nn.Conv2d
  14. cv2.INTER_LANCZOS4 in CPU postprocess: ~30 ms per 4K frame; INTER_LINEAR is 6x faster, visually indistinguishable
  15. Per-channel torch.stack-then-redistribute in despill: small but recurring

I think most (if not all) of these should be addressed ASAP, specially the faulty code and the bottom easy performance gains. If the maintainers agree with my assessment I can handle implementation myself.

Source: nikopueringer/CorridorKey