Scope the torch.load map_location wrapper instead of replacing it globally#52
Open
deepme987 wants to merge 1 commit into
Open
Scope the torch.load map_location wrapper instead of replacing it globally#52deepme987 wants to merge 1 commit into
deepme987 wants to merge 1 commit into
Conversation
…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.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Problem
chatterbox_node.pyinstallspatched_torch_loadby replacingtorch.loadprocess-wide at import time:Because this runs at import, it changes
torch.loadfor 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):
torch.load. Whichever pack imports last ownstorch.load, so the process-wide behavior depends on custom-node import order, which is not deterministic.map_locationforced onto callers that never opted in. After import, everytorch.loadcall anywhere in the process getsmap_locationinjected 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_loadexactly as-is, but install it only for the duration of this pack's own model loads, via a context manager:All four
Chatterbox*.from_local(...)call sites (load_turbo_model,load_tts_model,load_multilingual_model,load_vc_model) are wrapped withwith default_map_location():, so the device-defaulting still applies to every Chatterbox model load.torch.loadis 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 OKpatched_torch_loadvia the context manager