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:
CHECKPOINT_DIRis clobbered to a CWD-relative path:backend.py:74(line 22 has the correct absolute definition; line 74 reassigns toos.path.join("CorridorKeyModule", "checkpoints")). Breaks the installed package- MLX auto-download branch is unreachable dead code:
backend.py:85–119is indented inside anexcept ImportErrorclause that returns"torch"on line 89, so lines 91–119 (cache lookup +urllib.request.urlretrieve+os.rename+ finalreturn "mlx") are never executed. Apple-Silicon users silently get the slower Torch path - Autocast hardcoded
dtype=torch.float16:inference_engine.py:472. The_is_rocmbranch elsewhere in the file has no companion in autocast.fp16in attention on ROCm/HIP is documented to produceNaN/Inf;bf16is both more stable and 10–25% faster than thefp32fallback users hit after the firstNaN _clamp(x, min)silently ignores itsminargument:color_utils.py:42–48. Both branches hard-code0.0. Current call sites pass0.0by coincidence; any future caller gets silent wrong output_postprocess_torchdocstring documents a non-existent async API:inference_engine.py:328–334. Describes async=True/Falseparameter and aPendingTransferreturn type. Neither exists. The function always blocks and returnslist[dict[str, np.ndarray]]process_framereturn contract is undocumented:inference_engine.py:434–435. The Returns block lists{alpha, fg, comp}but omitsprocessed(the primary linear-premultiplied-RGBA output), thedict | list[dict]polymorphism, and thecomp is None when generate_comp=Falseruleapply_garbage_mattetensor branch is apassstub: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 inputmacos-latestrunner is Intel x86_64: everypytest.mark.mlxtest is structurally skipped. The MLX backend is dead in CIget_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 sizesconnected_componentstorch.randperm(H*W, device=device)per call: 63 MB transient GPU alloc at 4K for a tiebreaker that does not need randomnessclean_matte_torchruns up to 12 consecutivemax_pool2ddilation iterations: collapsible to one dilated conv. ~5–8 ms savings per 4K frametensor.to(device, non_blocking=True)on unpinned numpy silently synchronizes:non_blocking=Trueis only honored on pinned memoryDecoderHead.forwardfour-wayflatten/transpose/linear/transpose/view: ~3–5 ms per frame. Collapsible to a 1x1nn.Conv2dcv2.INTER_LANCZOS4in CPU postprocess: ~30 ms per 4K frame;INTER_LINEARis 6x faster, visually indistinguishable- 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