Skip to content

Scope the torch.load map_location wrapper instead of replacing it globally#52

Open
deepme987 wants to merge 1 commit into
filliptm:mainfrom
deepme987:deepme987/fix/scope-torch-load-patch
Open

Scope the torch.load map_location wrapper instead of replacing it globally#52
deepme987 wants to merge 1 commit into
filliptm:mainfrom
deepme987:deepme987/fix/scope-torch-load-patch

Conversation

@deepme987
Copy link
Copy Markdown

Problem

chatterbox_node.py installs patched_torch_load by replacing torch.load process-wide at import time:

torch.load = patched_torch_load

Because this runs at import, it changes torch.load for the entire Python process — ComfyUI core and every other installed custom node pack, not just this one.

Two concrete issues in a shared environment (and especially on multi-tenant cloud runtimes where one ComfyUI process serves many users' jobs back to back):

  1. Cross-pack clobbering — last import wins. Other custom node packs also wrap torch.load. Whichever pack imports last owns torch.load, so the process-wide behavior depends on custom-node import order, which is not deterministic.
  2. map_location forced onto callers that never opted in. After import, every torch.load call anywhere in the process gets map_location injected when the caller omitted it — including ComfyUI core's own checkpoint loading and other packs that deliberately relied on the default device placement.

Fix

Keep patched_torch_load exactly as-is, but install it only for the duration of this pack's own model loads, via a context manager:

@contextlib.contextmanager
def default_map_location():
    previous = torch.load
    torch.load = patched_torch_load
    try:
        yield
    finally:
        torch.load = previous

All four Chatterbox*.from_local(...) call sites (load_turbo_model, load_tts_model, load_multilingual_model, load_vc_model) are wrapped with with default_map_location():, so the device-defaulting still applies to every Chatterbox model load. torch.load is restored immediately afterward and left untouched for ComfyUI core and other packs.

No change to how this pack loads models or to the MPS/CUDA/CPU defaulting behavior — only the blast radius of the patch is reduced from process-global to call-scoped.

Testing

  • python -c "import ast; ast.parse(open('chatterbox_node.py').read())" — syntax OK
  • All four model-load functions still route through patched_torch_load via the context manager

…t globally

`chatterbox_node.py` did `torch.load = patched_torch_load` at import time,
which replaces torch.load for the entire Python process — ComfyUI core and
every other custom node pack included.

Two problems in a shared environment (and on multi-tenant cloud runtimes
where one process serves many users' jobs back to back):

1. Cross-pack clobbering. Other packs also wrap `torch.load`. Whichever
   imports last wins, so process-wide torch.load behavior depends on
   custom-node import order, which is not deterministic.
2. Behavior imposed on unrelated callers. After import, every `torch.load`
   in the process gets `map_location` forced onto it — including ComfyUI
   core's checkpoint loading and other packs that explicitly wanted default
   device placement.

Fix: keep `patched_torch_load` exactly as-is, but install it only for the
duration of this pack's own Chatterbox model loads via a
`default_map_location()` context manager, restoring the previous torch.load
in `finally`. All four `Chatterbox*.from_local(...)` call sites are wrapped,
so device defaulting still works for every Chatterbox load; torch.load is
left untouched for everyone else.

No functional change to how this pack loads models; only the blast radius of
the patch is reduced from process-global to call-scoped.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant