feat(remote_vae): add facade for ComfyUI-USCG-RemoteVAE companion

Adds is_remote_vae_available(), get_endpoint_names(), _companion_decode()
plus INSTALL_INSTRUCTIONS constant. Lazy try-import of the companion module
returns a clear error if not installed. Old code paths still active; the
remote_decode_hf -> _companion_decode swap lands in the next task.

Also updates tests/conftest.py to load the real remote_vae module (via
importlib.util) so facade tests exercise actual code rather than a stub.
Adds network_utils real-load and diffusers/diffusers.image_processor stubs
to satisfy the remote_vae.py import chain in the test environment.
This commit is contained in:
Jason Hoku
2026-05-19 13:08:27 -07:00
parent 3b26bb0d14
commit e5f690eae5
2 changed files with 95 additions and 1 deletions
+45
View File
@@ -20,6 +20,51 @@ except ImportError as e:
# HF_ENDPOINTS imported from network_utils (single source of truth for allowlisted URLs)
INSTALL_INSTRUCTIONS = (
"Remote VAE requires the ComfyUI-USCG-RemoteVAE companion plugin.\n"
"Install via Comfy Manager (search 'USCG Remote VAE') or:\n"
" git clone https://github.com/JasonHoku/ComfyUI-USCG-RemoteVAE\n"
"into your ComfyUI/custom_nodes/ directory."
)
def is_remote_vae_available():
"""Return True if the ComfyUI-USCG-RemoteVAE companion plugin is loaded."""
try:
import comfyui_uscg_remote_vae # noqa: F401
return True
except ImportError:
return False
def get_endpoint_names():
"""Return list of HF endpoint names for UI dropdown.
Empty list if companion not installed."""
try:
from comfyui_uscg_remote_vae import list_endpoints
return [name for name, _url in list_endpoints()]
except ImportError:
return []
def _companion_decode(endpoint, tensor, height, width):
"""Forward a decode call to the companion plugin.
Raises:
RuntimeError: companion not installed (with install instructions)
or companion version too old.
"""
try:
from comfyui_uscg_remote_vae import decode, __version__
except ImportError:
raise RuntimeError(INSTALL_INSTRUCTIONS)
if tuple(map(int, __version__.split(".")[:2])) < (0, 1):
raise RuntimeError(
f"ComfyUI-USCG-RemoteVAE >= 0.1.0 required (found {__version__}). "
"Update via Comfy Manager."
)
return decode(endpoint, tensor, height, width)
def detect_model_type(model, latent_channels):
"""
+50 -1
View File
@@ -29,6 +29,11 @@ _COMFY_STUBS = [
"comfy.sd",
"comfy.model_management",
"comfy.samplers",
# diffusers triggers a torchvision.__spec__ is None ValueError when imported
# outside ComfyUI. Stub it out so remote_vae.py's try/except catches the
# ImportError and sets VaeImageProcessor = None.
"diffusers",
"diffusers.image_processor",
]
for _mod_name in _COMFY_STUBS:
@@ -131,9 +136,35 @@ _GEN_ORCH_DEPS = {
"remote_vae": {
"RemoteVAEDecodeWorker": type("RemoteVAEDecodeWorker", (), {}),
"HF_ENDPOINTS": {},
# Facade functions — inline real implementations so tests can exercise
# them without needing to load the full remote_vae.py module.
"INSTALL_INSTRUCTIONS": (
"Remote VAE requires the ComfyUI-USCG-RemoteVAE companion plugin.\n"
"Install via Comfy Manager (search 'USCG Remote VAE') or:\n"
" git clone https://github.com/JasonHoku/ComfyUI-USCG-RemoteVAE\n"
"into your ComfyUI/custom_nodes/ directory."
),
},
}
# network_utils.py uses only stdlib — always load the real module and register
# it under both the bare name and the package-qualified name. This ensures the
# relative import `from .network_utils import huggingface_vae_decode, ...` in
# remote_vae.py resolves correctly regardless of what pytest has pre-loaded.
import importlib.util as _ilu # imported here; re-used below for gen_orch and remote_vae
_NETWORK_UTILS_FQ = f"{_PKG_NAME}.network_utils"
_NU_PATH = os.path.join(_NODE_ROOT, "network_utils.py")
_nu_spec = _ilu.spec_from_file_location(_NETWORK_UTILS_FQ, _NU_PATH,
submodule_search_locations=[])
_nu_spec.submodule_search_locations = None
_nu_mod = _ilu.module_from_spec(_nu_spec)
_nu_mod.__package__ = _PKG_NAME
# Unconditionally replace any prior stub — we need the real attributes.
sys.modules[_NETWORK_UTILS_FQ] = _nu_mod
sys.modules["network_utils"] = _nu_mod
setattr(sys.modules[_PKG_NAME], "network_utils", _nu_mod)
_nu_spec.loader.exec_module(_nu_mod)
for _bare, _attrs in _GEN_ORCH_DEPS.items():
_fq = f"{_PKG_NAME}.{_bare}"
if _bare not in sys.modules:
@@ -151,7 +182,6 @@ for _bare, _attrs in _GEN_ORCH_DEPS.items():
# imports (`from .trigger_words import ...`) resolve correctly. Once loaded
# under the package name it is also registered as a top-level alias so
# `from generation_orchestrator import get_model_cache_key` works.
import importlib.util as _ilu
_GO_PATH = os.path.join(_NODE_ROOT, "generation_orchestrator.py")
_GO_FQ = f"{_PKG_NAME}.generation_orchestrator"
if "generation_orchestrator" not in sys.modules:
@@ -163,3 +193,22 @@ if "generation_orchestrator" not in sys.modules:
sys.modules[_GO_FQ] = _go_mod
sys.modules["generation_orchestrator"] = _go_mod
_spec.loader.exec_module(_go_mod)
# Pre-load the real remote_vae module so that test_remote_vae_facade.py gets
# the actual functions (is_remote_vae_available, get_endpoint_names,
# _companion_decode) rather than the lightweight stub registered above.
# remote_vae.py uses `from .network_utils import ...`; we satisfy that via the
# network_utils stub already registered in sys.modules.
_RV_PATH = os.path.join(_NODE_ROOT, "remote_vae.py")
_RV_FQ = f"{_PKG_NAME}.remote_vae"
_rv_spec = _ilu.spec_from_file_location(_RV_FQ, _RV_PATH,
submodule_search_locations=[])
_rv_spec.submodule_search_locations = None
_rv_mod = _ilu.module_from_spec(_rv_spec)
_rv_mod.__package__ = _PKG_NAME
# Replace the earlier stub with the real module BEFORE exec so that any
# internal import resolution finds the real object.
sys.modules[_RV_FQ] = _rv_mod
sys.modules["remote_vae"] = _rv_mod
setattr(sys.modules[_PKG_NAME], "remote_vae", _rv_mod)
_rv_spec.loader.exec_module(_rv_mod)