mirror of
https://github.com/comfyanonymous/ComfyUI.git
synced 2026-07-21 23:41:28 +08:00
Compare commits
26
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
f9bb70a2ea | ||
|
|
a2b60dea17 | ||
|
|
7288264c7e | ||
|
|
defb663b94 | ||
|
|
d0f0b15cf5 | ||
|
|
b5bb83c964 | ||
|
|
f6d5068ac0 | ||
|
|
be95871adc | ||
|
|
f756d801a1 | ||
|
|
1d23a875ed | ||
|
|
ef6722f6be | ||
|
|
783782d5d7 | ||
|
|
3e3ed8cc2a | ||
|
|
67f6cb3527 | ||
|
|
0230e0e7cc | ||
|
|
b5921c8ac2 | ||
|
|
63103d519e | ||
|
|
cf758bd256 | ||
|
|
10b45a71cd | ||
|
|
fa7553138e | ||
|
|
cf9cbec596 | ||
|
|
96f1cee9f5 | ||
|
|
97f58baaaf | ||
|
|
e8e8fee224 | ||
|
|
e9c311b245 | ||
|
|
e6e0936128 |
+1
-1
@@ -1,2 +1,2 @@
|
||||
.\python_embeded\python.exe -s ComfyUI\main.py --windows-standalone-build --disable-smart-memory
|
||||
.\python_embeded\python.exe -s ComfyUI\main.py --windows-standalone-build --enable-dynamic-vram
|
||||
pause
|
||||
@@ -1,2 +1,2 @@
|
||||
.\python_embeded\python.exe -s ComfyUI\main.py --windows-standalone-build
|
||||
pause
|
||||
.\python_embeded\python.exe -s ComfyUI\main.py --windows-standalone-build
|
||||
pause
|
||||
|
||||
@@ -23,4 +23,3 @@ web_custom_versions/
|
||||
.DS_Store
|
||||
filtered-openapi.yaml
|
||||
uv.lock
|
||||
.pyisolate_venvs/
|
||||
|
||||
+1
-1
@@ -1,2 +1,2 @@
|
||||
# Admins
|
||||
* @comfyanonymous @kosinkadink @guill
|
||||
* @comfyanonymous @kosinkadink @guill @alexisrolland @rattus128 @kijai
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
<div align="center">
|
||||
|
||||
# ComfyUI
|
||||
**The most powerful and modular visual AI engine and application.**
|
||||
**The most powerful and modular AI engine for content creation.**
|
||||
|
||||
|
||||
[![Website][website-shield]][website-url]
|
||||
@@ -31,10 +31,16 @@
|
||||
[github-downloads-latest-shield]: https://img.shields.io/github/downloads/comfyanonymous/ComfyUI/latest/total?style=flat&label=downloads%40latest
|
||||
[github-downloads-link]: https://github.com/comfyanonymous/ComfyUI/releases
|
||||
|
||||

|
||||
<img width="1590" height="795" alt="ComfyUI Screenshot" src="https://github.com/user-attachments/assets/36e065e0-bfae-4456-8c7f-8369d5ea48a2" />
|
||||
<br>
|
||||
</div>
|
||||
|
||||
ComfyUI lets you design and execute advanced stable diffusion pipelines using a graph/nodes/flowchart based interface. Available on Windows, Linux, and macOS.
|
||||
ComfyUI is the AI creation engine for visual professionals who demand control over every model, every parameter, and every output. Its powerful and modular node graph interface empowers creatives to generate images, videos, 3D models, audio, and more...
|
||||
- ComfyUI natively supports the latest open-source state of the art models.
|
||||
- API nodes provide access to the best closed source models such as Nano Banana, Seedance, Hunyuan3D, etc.
|
||||
- It is available on Windows, Linux, and macOS, locally with our desktop application or on our cloud.
|
||||
- The most sophisticated workflows can be exposed through a simple UI thanks to App Mode.
|
||||
- It integrates seamlessly into production pipelines with our API endpoints.
|
||||
|
||||
## Get Started
|
||||
|
||||
@@ -77,6 +83,7 @@ See what ComfyUI can do with the [newer template workflows](https://comfy.org/wo
|
||||
- [Hunyuan Image 2.1](https://comfyanonymous.github.io/ComfyUI_examples/hunyuan_image/)
|
||||
- [Flux 2](https://comfyanonymous.github.io/ComfyUI_examples/flux2/)
|
||||
- [Z Image](https://comfyanonymous.github.io/ComfyUI_examples/z_image/)
|
||||
- Ernie Image
|
||||
- Image Editing Models
|
||||
- [Omnigen 2](https://comfyanonymous.github.io/ComfyUI_examples/omnigen/)
|
||||
- [Flux Kontext](https://comfyanonymous.github.io/ComfyUI_examples/flux/#flux-kontext-image-editing-model)
|
||||
@@ -193,13 +200,15 @@ If you have trouble extracting it, right click the file -> properties -> unblock
|
||||
|
||||
The portable above currently comes with python 3.13 and pytorch cuda 13.0. Update your Nvidia drivers if it doesn't start.
|
||||
|
||||
#### Alternative Downloads:
|
||||
#### All Official Portable Downloads:
|
||||
|
||||
[Portable for AMD GPUs](https://github.com/comfyanonymous/ComfyUI/releases/latest/download/ComfyUI_windows_portable_amd.7z)
|
||||
|
||||
[Experimental portable for Intel GPUs](https://github.com/comfyanonymous/ComfyUI/releases/latest/download/ComfyUI_windows_portable_intel.7z)
|
||||
[Portable for Intel GPUs](https://github.com/comfyanonymous/ComfyUI/releases/latest/download/ComfyUI_windows_portable_intel.7z)
|
||||
|
||||
[Portable with pytorch cuda 12.6 and python 3.12](https://github.com/comfyanonymous/ComfyUI/releases/latest/download/ComfyUI_windows_portable_nvidia_cu126.7z) (Supports Nvidia 10 series and older GPUs).
|
||||
[Portable for Nvidia GPUs](https://github.com/comfyanonymous/ComfyUI/releases/latest/download/ComfyUI_windows_portable_nvidia.7z) (supports 20 series and above).
|
||||
|
||||
[Portable for Nvidia GPUs with pytorch cuda 12.6 and python 3.12](https://github.com/comfyanonymous/ComfyUI/releases/latest/download/ComfyUI_windows_portable_nvidia_cu126.7z) (Supports Nvidia 10 series and older GPUs).
|
||||
|
||||
#### How do I share models between another UI and ComfyUI?
|
||||
|
||||
|
||||
@@ -90,7 +90,6 @@ parser.add_argument("--force-channels-last", action="store_true", help="Force ch
|
||||
parser.add_argument("--directml", type=int, nargs="?", metavar="DIRECTML_DEVICE", const=-1, help="Use torch-directml.")
|
||||
|
||||
parser.add_argument("--oneapi-device-selector", type=str, default=None, metavar="SELECTOR_STRING", help="Sets the oneAPI device(s) this instance will use.")
|
||||
parser.add_argument("--disable-ipex-optimize", action="store_true", help="Disables ipex.optimize default when loading models with Intel's Extension for Pytorch.")
|
||||
parser.add_argument("--supports-fp8-compute", action="store_true", help="ComfyUI will act like if the device supports fp8 compute.")
|
||||
|
||||
class LatentPreviewMethod(enum.Enum):
|
||||
@@ -184,8 +183,6 @@ parser.add_argument("--disable-api-nodes", action="store_true", help="Disable lo
|
||||
|
||||
parser.add_argument("--multi-user", action="store_true", help="Enables per-user storage.")
|
||||
|
||||
parser.add_argument("--use-process-isolation", action="store_true", help="Enable process isolation for custom nodes with pyproject.toml manifests containing a [tool.comfy.isolation] section.")
|
||||
|
||||
parser.add_argument("--verbose", default='INFO', const='DEBUG', nargs="?", choices=['DEBUG', 'INFO', 'WARNING', 'ERROR', 'CRITICAL'], help='Set the logging level')
|
||||
parser.add_argument("--log-stdout", action="store_true", help="Send normal process output to stdout instead of stderr (default).")
|
||||
|
||||
|
||||
+1
-35
@@ -14,9 +14,6 @@ if TYPE_CHECKING:
|
||||
import comfy.lora
|
||||
import comfy.model_management
|
||||
import comfy.patcher_extension
|
||||
from comfy.cli_args import args
|
||||
import uuid
|
||||
import os
|
||||
from node_helpers import conditioning_set_values
|
||||
|
||||
# #######################################################################################################
|
||||
@@ -64,37 +61,8 @@ class EnumHookScope(enum.Enum):
|
||||
HookedOnly = "hooked_only"
|
||||
|
||||
|
||||
_ISOLATION_HOOKREF_MODE = args.use_process_isolation or os.environ.get("PYISOLATE_CHILD") == "1"
|
||||
|
||||
|
||||
class _HookRef:
|
||||
def __init__(self):
|
||||
if _ISOLATION_HOOKREF_MODE:
|
||||
self._pyisolate_id = str(uuid.uuid4())
|
||||
|
||||
def _ensure_pyisolate_id(self):
|
||||
pyisolate_id = getattr(self, "_pyisolate_id", None)
|
||||
if pyisolate_id is None:
|
||||
pyisolate_id = str(uuid.uuid4())
|
||||
self._pyisolate_id = pyisolate_id
|
||||
return pyisolate_id
|
||||
|
||||
def __eq__(self, other):
|
||||
if not _ISOLATION_HOOKREF_MODE:
|
||||
return self is other
|
||||
if not isinstance(other, _HookRef):
|
||||
return False
|
||||
return self._ensure_pyisolate_id() == other._ensure_pyisolate_id()
|
||||
|
||||
def __hash__(self):
|
||||
if not _ISOLATION_HOOKREF_MODE:
|
||||
return id(self)
|
||||
return hash(self._ensure_pyisolate_id())
|
||||
|
||||
def __str__(self):
|
||||
if not _ISOLATION_HOOKREF_MODE:
|
||||
return super().__str__()
|
||||
return f"PYISOLATE_HOOKREF:{self._ensure_pyisolate_id()}"
|
||||
pass
|
||||
|
||||
|
||||
def default_should_register(hook: Hook, model: ModelPatcher, model_options: dict, target_dict: dict[str], registered: HookGroup):
|
||||
@@ -200,8 +168,6 @@ class WeightHook(Hook):
|
||||
key_map = comfy.lora.model_lora_keys_clip(model.model, key_map)
|
||||
else:
|
||||
key_map = comfy.lora.model_lora_keys_unet(model.model, key_map)
|
||||
if self.weights is None:
|
||||
self.weights = {}
|
||||
weights = comfy.lora.load_lora(self.weights, key_map, log_missing=False)
|
||||
else:
|
||||
if target == EnumWeightTarget.Clip:
|
||||
|
||||
@@ -1,436 +0,0 @@
|
||||
# pylint: disable=consider-using-from-import,cyclic-import,global-statement,global-variable-not-assigned,import-outside-toplevel,logging-fstring-interpolation
|
||||
from __future__ import annotations
|
||||
import asyncio
|
||||
import inspect
|
||||
import logging
|
||||
import os
|
||||
import time
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Dict, List, Optional, Set, TYPE_CHECKING
|
||||
_IMPORT_TORCH = os.environ.get("PYISOLATE_IMPORT_TORCH", "1") == "1"
|
||||
|
||||
load_isolated_node = None
|
||||
find_manifest_directories = None
|
||||
build_stub_class = None
|
||||
get_class_types_for_extension = None
|
||||
scan_shm_forensics = None
|
||||
start_shm_forensics = None
|
||||
|
||||
if _IMPORT_TORCH:
|
||||
import folder_paths
|
||||
from .extension_loader import load_isolated_node
|
||||
from .manifest_loader import find_manifest_directories
|
||||
from .runtime_helpers import build_stub_class, get_class_types_for_extension
|
||||
from .shm_forensics import scan_shm_forensics, start_shm_forensics
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from pyisolate import ExtensionManager
|
||||
from .extension_wrapper import ComfyNodeExtension
|
||||
|
||||
LOG_PREFIX = "]["
|
||||
isolated_node_timings: List[tuple[float, Path, int]] = []
|
||||
|
||||
if _IMPORT_TORCH:
|
||||
PYISOLATE_VENV_ROOT = Path(folder_paths.base_path) / ".pyisolate_venvs"
|
||||
PYISOLATE_VENV_ROOT.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
_WORKFLOW_BOUNDARY_MIN_FREE_VRAM_BYTES = 2 * 1024 * 1024 * 1024
|
||||
_MODEL_PATCHER_IDLE_TIMEOUT_MS = 120000
|
||||
|
||||
|
||||
def initialize_proxies() -> None:
|
||||
from .child_hooks import is_child_process
|
||||
|
||||
is_child = is_child_process()
|
||||
|
||||
if is_child:
|
||||
from .child_hooks import initialize_child_process
|
||||
|
||||
initialize_child_process()
|
||||
else:
|
||||
from .host_hooks import initialize_host_process
|
||||
|
||||
initialize_host_process()
|
||||
if start_shm_forensics is not None:
|
||||
start_shm_forensics()
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class IsolatedNodeSpec:
|
||||
node_name: str
|
||||
display_name: str
|
||||
stub_class: type
|
||||
module_path: Path
|
||||
|
||||
|
||||
_ISOLATED_NODE_SPECS: List[IsolatedNodeSpec] = []
|
||||
_CLAIMED_PATHS: Set[Path] = set()
|
||||
_ISOLATION_SCAN_ATTEMPTED = False
|
||||
_EXTENSION_MANAGERS: List["ExtensionManager"] = []
|
||||
_RUNNING_EXTENSIONS: Dict[str, "ComfyNodeExtension"] = {}
|
||||
_ISOLATION_BACKGROUND_TASK: Optional["asyncio.Task[List[IsolatedNodeSpec]]"] = None
|
||||
_EARLY_START_TIME: Optional[float] = None
|
||||
|
||||
|
||||
def start_isolation_loading_early(loop: "asyncio.AbstractEventLoop") -> None:
|
||||
global _ISOLATION_BACKGROUND_TASK, _EARLY_START_TIME
|
||||
if _ISOLATION_BACKGROUND_TASK is not None:
|
||||
return
|
||||
_EARLY_START_TIME = time.perf_counter()
|
||||
_ISOLATION_BACKGROUND_TASK = loop.create_task(initialize_isolation_nodes())
|
||||
|
||||
|
||||
async def await_isolation_loading() -> List[IsolatedNodeSpec]:
|
||||
global _ISOLATION_BACKGROUND_TASK, _EARLY_START_TIME
|
||||
if _ISOLATION_BACKGROUND_TASK is not None:
|
||||
specs = await _ISOLATION_BACKGROUND_TASK
|
||||
return specs
|
||||
return await initialize_isolation_nodes()
|
||||
|
||||
|
||||
async def initialize_isolation_nodes() -> List[IsolatedNodeSpec]:
|
||||
global _ISOLATED_NODE_SPECS, _ISOLATION_SCAN_ATTEMPTED, _CLAIMED_PATHS
|
||||
|
||||
if _ISOLATED_NODE_SPECS:
|
||||
return _ISOLATED_NODE_SPECS
|
||||
|
||||
if _ISOLATION_SCAN_ATTEMPTED:
|
||||
return []
|
||||
|
||||
_ISOLATION_SCAN_ATTEMPTED = True
|
||||
if find_manifest_directories is None or load_isolated_node is None or build_stub_class is None:
|
||||
return []
|
||||
manifest_entries = find_manifest_directories()
|
||||
_CLAIMED_PATHS = {entry[0].resolve() for entry in manifest_entries}
|
||||
|
||||
if not manifest_entries:
|
||||
return []
|
||||
|
||||
os.environ["PYISOLATE_ISOLATION_ACTIVE"] = "1"
|
||||
concurrency_limit = max(1, (os.cpu_count() or 4) // 2)
|
||||
semaphore = asyncio.Semaphore(concurrency_limit)
|
||||
|
||||
async def load_with_semaphore(
|
||||
node_dir: Path, manifest: Path
|
||||
) -> List[IsolatedNodeSpec]:
|
||||
async with semaphore:
|
||||
load_start = time.perf_counter()
|
||||
spec_list = await load_isolated_node(
|
||||
node_dir,
|
||||
manifest,
|
||||
logger,
|
||||
lambda name, info, extension: build_stub_class(
|
||||
name,
|
||||
info,
|
||||
extension,
|
||||
_RUNNING_EXTENSIONS,
|
||||
logger,
|
||||
),
|
||||
PYISOLATE_VENV_ROOT,
|
||||
_EXTENSION_MANAGERS,
|
||||
)
|
||||
spec_list = [
|
||||
IsolatedNodeSpec(
|
||||
node_name=node_name,
|
||||
display_name=display_name,
|
||||
stub_class=stub_cls,
|
||||
module_path=node_dir,
|
||||
)
|
||||
for node_name, display_name, stub_cls in spec_list
|
||||
]
|
||||
isolated_node_timings.append(
|
||||
(time.perf_counter() - load_start, node_dir, len(spec_list))
|
||||
)
|
||||
return spec_list
|
||||
|
||||
tasks = [
|
||||
load_with_semaphore(node_dir, manifest)
|
||||
for node_dir, manifest in manifest_entries
|
||||
]
|
||||
results = await asyncio.gather(*tasks, return_exceptions=True)
|
||||
|
||||
specs: List[IsolatedNodeSpec] = []
|
||||
for result in results:
|
||||
if isinstance(result, Exception):
|
||||
logger.error(
|
||||
"%s Isolated node failed during startup; continuing: %s",
|
||||
LOG_PREFIX,
|
||||
result,
|
||||
)
|
||||
continue
|
||||
specs.extend(result)
|
||||
|
||||
_ISOLATED_NODE_SPECS = specs
|
||||
return list(_ISOLATED_NODE_SPECS)
|
||||
|
||||
|
||||
def _get_class_types_for_extension(extension_name: str) -> Set[str]:
|
||||
"""Get all node class types (node names) belonging to an extension."""
|
||||
extension = _RUNNING_EXTENSIONS.get(extension_name)
|
||||
if not extension:
|
||||
return set()
|
||||
|
||||
ext_path = Path(extension.module_path)
|
||||
class_types = set()
|
||||
for spec in _ISOLATED_NODE_SPECS:
|
||||
if spec.module_path.resolve() == ext_path.resolve():
|
||||
class_types.add(spec.node_name)
|
||||
|
||||
return class_types
|
||||
|
||||
|
||||
async def notify_execution_graph(needed_class_types: Set[str], caches: list | None = None) -> None:
|
||||
"""Evict running extensions not needed for current execution.
|
||||
|
||||
When *caches* is provided, cache entries for evicted extensions' node
|
||||
class_types are invalidated to prevent stale ``RemoteObjectHandle``
|
||||
references from surviving in the output cache.
|
||||
"""
|
||||
await wait_for_model_patcher_quiescence(
|
||||
timeout_ms=_MODEL_PATCHER_IDLE_TIMEOUT_MS,
|
||||
fail_loud=True,
|
||||
marker="ISO:notify_graph_wait_idle",
|
||||
)
|
||||
|
||||
evicted_class_types: Set[str] = set()
|
||||
|
||||
async def _stop_extension(
|
||||
ext_name: str, extension: "ComfyNodeExtension", reason: str
|
||||
) -> None:
|
||||
# Collect class_types BEFORE stopping so we can invalidate cache entries.
|
||||
ext_class_types = _get_class_types_for_extension(ext_name)
|
||||
evicted_class_types.update(ext_class_types)
|
||||
logger.info("%s ISO:eject_start ext=%s reason=%s", LOG_PREFIX, ext_name, reason)
|
||||
logger.debug("%s ISO:stop_start ext=%s", LOG_PREFIX, ext_name)
|
||||
stop_result = extension.stop()
|
||||
if inspect.isawaitable(stop_result):
|
||||
await stop_result
|
||||
_RUNNING_EXTENSIONS.pop(ext_name, None)
|
||||
logger.debug("%s ISO:stop_done ext=%s", LOG_PREFIX, ext_name)
|
||||
if scan_shm_forensics is not None:
|
||||
scan_shm_forensics("ISO:stop_extension", refresh_model_context=True)
|
||||
|
||||
if scan_shm_forensics is not None:
|
||||
scan_shm_forensics("ISO:notify_graph_start", refresh_model_context=True)
|
||||
isolated_class_types_in_graph = needed_class_types.intersection(
|
||||
{spec.node_name for spec in _ISOLATED_NODE_SPECS}
|
||||
)
|
||||
graph_uses_isolation = bool(isolated_class_types_in_graph)
|
||||
logger.debug(
|
||||
"%s ISO:notify_graph_start running=%d needed=%d",
|
||||
LOG_PREFIX,
|
||||
len(_RUNNING_EXTENSIONS),
|
||||
len(needed_class_types),
|
||||
)
|
||||
if graph_uses_isolation:
|
||||
for ext_name, extension in list(_RUNNING_EXTENSIONS.items()):
|
||||
ext_class_types = _get_class_types_for_extension(ext_name)
|
||||
|
||||
# If NONE of this extension's nodes are in the execution graph -> evict.
|
||||
if not ext_class_types.intersection(needed_class_types):
|
||||
await _stop_extension(
|
||||
ext_name,
|
||||
extension,
|
||||
"isolated custom_node not in execution graph, evicting",
|
||||
)
|
||||
else:
|
||||
logger.debug(
|
||||
"%s ISO:notify_graph_skip_evict running=%d reason=no isolated nodes in graph",
|
||||
LOG_PREFIX,
|
||||
len(_RUNNING_EXTENSIONS),
|
||||
)
|
||||
|
||||
# Isolated child processes add steady VRAM pressure; reclaim host-side models
|
||||
# at workflow boundaries so subsequent host nodes (e.g. CLIP encode) keep headroom.
|
||||
try:
|
||||
import comfy.model_management as model_management
|
||||
|
||||
device = model_management.get_torch_device()
|
||||
if getattr(device, "type", None) == "cuda":
|
||||
required = max(
|
||||
model_management.minimum_inference_memory(),
|
||||
_WORKFLOW_BOUNDARY_MIN_FREE_VRAM_BYTES,
|
||||
)
|
||||
free_before = model_management.get_free_memory(device)
|
||||
if free_before < required and _RUNNING_EXTENSIONS and graph_uses_isolation:
|
||||
for ext_name, extension in list(_RUNNING_EXTENSIONS.items()):
|
||||
await _stop_extension(
|
||||
ext_name,
|
||||
extension,
|
||||
f"boundary low-vram restart (free={int(free_before)} target={int(required)})",
|
||||
)
|
||||
if model_management.get_free_memory(device) < required:
|
||||
model_management.unload_all_models()
|
||||
model_management.cleanup_models_gc()
|
||||
model_management.cleanup_models()
|
||||
if model_management.get_free_memory(device) < required:
|
||||
model_management.free_memory(required, device, for_dynamic=False)
|
||||
model_management.soft_empty_cache()
|
||||
except Exception:
|
||||
logger.debug(
|
||||
"%s workflow-boundary host VRAM relief failed", LOG_PREFIX, exc_info=True
|
||||
)
|
||||
finally:
|
||||
# Invalidate cached outputs for evicted extensions so stale
|
||||
# RemoteObjectHandle references are not served from cache.
|
||||
if evicted_class_types and caches:
|
||||
total_invalidated = 0
|
||||
for cache in caches:
|
||||
if hasattr(cache, "invalidate_by_class_types"):
|
||||
total_invalidated += cache.invalidate_by_class_types(
|
||||
evicted_class_types
|
||||
)
|
||||
if total_invalidated > 0:
|
||||
logger.info(
|
||||
"%s ISO:cache_invalidated count=%d class_types=%s",
|
||||
LOG_PREFIX,
|
||||
total_invalidated,
|
||||
evicted_class_types,
|
||||
)
|
||||
scan_shm_forensics("ISO:notify_graph_done", refresh_model_context=True)
|
||||
logger.debug(
|
||||
"%s ISO:notify_graph_done running=%d", LOG_PREFIX, len(_RUNNING_EXTENSIONS)
|
||||
)
|
||||
|
||||
|
||||
async def flush_running_extensions_transport_state() -> int:
|
||||
await wait_for_model_patcher_quiescence(
|
||||
timeout_ms=_MODEL_PATCHER_IDLE_TIMEOUT_MS,
|
||||
fail_loud=True,
|
||||
marker="ISO:flush_transport_wait_idle",
|
||||
)
|
||||
total_flushed = 0
|
||||
for ext_name, extension in list(_RUNNING_EXTENSIONS.items()):
|
||||
flush_fn = getattr(extension, "flush_transport_state", None)
|
||||
if not callable(flush_fn):
|
||||
continue
|
||||
try:
|
||||
flushed = await flush_fn()
|
||||
if isinstance(flushed, int):
|
||||
total_flushed += flushed
|
||||
if flushed > 0:
|
||||
logger.debug(
|
||||
"%s %s workflow-end flush released=%d",
|
||||
LOG_PREFIX,
|
||||
ext_name,
|
||||
flushed,
|
||||
)
|
||||
except Exception:
|
||||
logger.debug(
|
||||
"%s %s workflow-end flush failed", LOG_PREFIX, ext_name, exc_info=True
|
||||
)
|
||||
scan_shm_forensics(
|
||||
"ISO:flush_running_extensions_transport_state", refresh_model_context=True
|
||||
)
|
||||
return total_flushed
|
||||
|
||||
|
||||
async def wait_for_model_patcher_quiescence(
|
||||
timeout_ms: int = _MODEL_PATCHER_IDLE_TIMEOUT_MS,
|
||||
*,
|
||||
fail_loud: bool = False,
|
||||
marker: str = "ISO:wait_model_patcher_idle",
|
||||
) -> bool:
|
||||
try:
|
||||
from comfy.isolation.model_patcher_proxy_registry import ModelPatcherRegistry
|
||||
|
||||
registry = ModelPatcherRegistry()
|
||||
start = time.perf_counter()
|
||||
idle = await registry.wait_all_idle(timeout_ms)
|
||||
elapsed_ms = (time.perf_counter() - start) * 1000.0
|
||||
if idle:
|
||||
logger.debug(
|
||||
"%s %s idle=1 timeout_ms=%d elapsed_ms=%.3f",
|
||||
LOG_PREFIX,
|
||||
marker,
|
||||
timeout_ms,
|
||||
elapsed_ms,
|
||||
)
|
||||
return True
|
||||
|
||||
states = await registry.get_all_operation_states()
|
||||
logger.error(
|
||||
"%s %s idle_timeout timeout_ms=%d elapsed_ms=%.3f states=%s",
|
||||
LOG_PREFIX,
|
||||
marker,
|
||||
timeout_ms,
|
||||
elapsed_ms,
|
||||
states,
|
||||
)
|
||||
if fail_loud:
|
||||
raise TimeoutError(
|
||||
f"ModelPatcherRegistry did not quiesce within {timeout_ms} ms"
|
||||
)
|
||||
return False
|
||||
except Exception:
|
||||
if fail_loud:
|
||||
raise
|
||||
logger.debug("%s %s failed", LOG_PREFIX, marker, exc_info=True)
|
||||
return False
|
||||
|
||||
|
||||
def get_claimed_paths() -> Set[Path]:
|
||||
return _CLAIMED_PATHS
|
||||
|
||||
|
||||
def update_rpc_event_loops(loop: "asyncio.AbstractEventLoop | None" = None) -> None:
|
||||
"""Update all active RPC instances with the current event loop.
|
||||
|
||||
This MUST be called at the start of each workflow execution to ensure
|
||||
RPC calls are scheduled on the correct event loop. This handles the case
|
||||
where asyncio.run() creates a new event loop for each workflow.
|
||||
|
||||
Args:
|
||||
loop: The event loop to use. If None, uses asyncio.get_running_loop().
|
||||
"""
|
||||
if loop is None:
|
||||
try:
|
||||
loop = asyncio.get_running_loop()
|
||||
except RuntimeError:
|
||||
loop = asyncio.get_event_loop()
|
||||
|
||||
update_count = 0
|
||||
|
||||
# Update RPCs from ExtensionManagers
|
||||
for manager in _EXTENSION_MANAGERS:
|
||||
if not hasattr(manager, "extensions"):
|
||||
continue
|
||||
for name, extension in manager.extensions.items():
|
||||
if hasattr(extension, "rpc") and extension.rpc is not None:
|
||||
if hasattr(extension.rpc, "update_event_loop"):
|
||||
extension.rpc.update_event_loop(loop)
|
||||
update_count += 1
|
||||
logger.debug(f"{LOG_PREFIX}Updated loop on extension '{name}'")
|
||||
|
||||
# Also update RPCs from running extensions (they may have direct RPC refs)
|
||||
for name, extension in _RUNNING_EXTENSIONS.items():
|
||||
if hasattr(extension, "rpc") and extension.rpc is not None:
|
||||
if hasattr(extension.rpc, "update_event_loop"):
|
||||
extension.rpc.update_event_loop(loop)
|
||||
update_count += 1
|
||||
logger.debug(f"{LOG_PREFIX}Updated loop on running extension '{name}'")
|
||||
|
||||
if update_count > 0:
|
||||
logger.debug(f"{LOG_PREFIX}Updated event loop on {update_count} RPC instances")
|
||||
else:
|
||||
logger.debug(
|
||||
f"{LOG_PREFIX}No RPC instances found to update (managers={len(_EXTENSION_MANAGERS)}, running={len(_RUNNING_EXTENSIONS)})"
|
||||
)
|
||||
|
||||
|
||||
__all__ = [
|
||||
"LOG_PREFIX",
|
||||
"initialize_proxies",
|
||||
"initialize_isolation_nodes",
|
||||
"start_isolation_loading_early",
|
||||
"await_isolation_loading",
|
||||
"notify_execution_graph",
|
||||
"flush_running_extensions_transport_state",
|
||||
"wait_for_model_patcher_quiescence",
|
||||
"get_claimed_paths",
|
||||
"update_rpc_event_loops",
|
||||
"IsolatedNodeSpec",
|
||||
"get_class_types_for_extension",
|
||||
]
|
||||
@@ -1,864 +0,0 @@
|
||||
# pylint: disable=import-outside-toplevel,logging-fstring-interpolation,protected-access,raise-missing-from,useless-return,wrong-import-position
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import os
|
||||
import inspect
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, List, Optional, cast
|
||||
|
||||
from pyisolate.interfaces import IsolationAdapter, SerializerRegistryProtocol # type: ignore[import-untyped]
|
||||
from pyisolate._internal.rpc_protocol import AsyncRPC, ProxiedSingleton # type: ignore[import-untyped]
|
||||
|
||||
_IMPORT_TORCH = os.environ.get("PYISOLATE_IMPORT_TORCH", "1") == "1"
|
||||
|
||||
# Singleton proxies that do NOT transitively import torch/PIL/psutil/aiohttp.
|
||||
# Safe to import in sealed workers without host framework modules.
|
||||
from comfy.isolation.proxies.folder_paths_proxy import FolderPathsProxy
|
||||
from comfy.isolation.proxies.helper_proxies import HelperProxiesService
|
||||
from comfy.isolation.proxies.web_directory_proxy import WebDirectoryProxy
|
||||
|
||||
# Singleton proxies that transitively import torch, PIL, or heavy host modules.
|
||||
# Only available when torch/host framework is present.
|
||||
CLIPProxy = None
|
||||
CLIPRegistry = None
|
||||
ModelPatcherProxy = None
|
||||
ModelPatcherRegistry = None
|
||||
ModelSamplingProxy = None
|
||||
ModelSamplingRegistry = None
|
||||
VAEProxy = None
|
||||
VAERegistry = None
|
||||
FirstStageModelRegistry = None
|
||||
ModelManagementProxy = None
|
||||
PromptServerService = None
|
||||
ProgressProxy = None
|
||||
UtilsProxy = None
|
||||
_HAS_TORCH_PROXIES = False
|
||||
if _IMPORT_TORCH:
|
||||
from comfy.isolation.clip_proxy import CLIPProxy, CLIPRegistry
|
||||
from comfy.isolation.model_patcher_proxy import (
|
||||
ModelPatcherProxy,
|
||||
ModelPatcherRegistry,
|
||||
)
|
||||
from comfy.isolation.model_sampling_proxy import (
|
||||
ModelSamplingProxy,
|
||||
ModelSamplingRegistry,
|
||||
)
|
||||
from comfy.isolation.vae_proxy import VAEProxy, VAERegistry, FirstStageModelRegistry
|
||||
from comfy.isolation.proxies.model_management_proxy import ModelManagementProxy
|
||||
from comfy.isolation.proxies.prompt_server_impl import PromptServerService
|
||||
from comfy.isolation.proxies.progress_proxy import ProgressProxy
|
||||
from comfy.isolation.proxies.utils_proxy import UtilsProxy
|
||||
_HAS_TORCH_PROXIES = True
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Force /dev/shm for shared memory (bwrap makes /tmp private)
|
||||
import tempfile
|
||||
|
||||
if os.path.exists("/dev/shm"):
|
||||
# Only override if not already set or if default is not /dev/shm
|
||||
current_tmp = tempfile.gettempdir()
|
||||
if not current_tmp.startswith("/dev/shm"):
|
||||
logger.debug(
|
||||
f"Configuring shared memory: Changing TMPDIR from {current_tmp} to /dev/shm"
|
||||
)
|
||||
os.environ["TMPDIR"] = "/dev/shm"
|
||||
tempfile.tempdir = None # Clear cache to force re-evaluation
|
||||
|
||||
|
||||
class ComfyUIAdapter(IsolationAdapter):
|
||||
# ComfyUI-specific IsolationAdapter implementation
|
||||
|
||||
@property
|
||||
def identifier(self) -> str:
|
||||
return "comfyui"
|
||||
|
||||
def get_path_config(self, module_path: str) -> Optional[Dict[str, Any]]:
|
||||
if "ComfyUI" in module_path and "custom_nodes" in module_path:
|
||||
parts = module_path.split("ComfyUI")
|
||||
if len(parts) > 1:
|
||||
comfy_root = parts[0] + "ComfyUI"
|
||||
return {
|
||||
"preferred_root": comfy_root,
|
||||
"additional_paths": [
|
||||
os.path.join(comfy_root, "custom_nodes"),
|
||||
os.path.join(comfy_root, "comfy"),
|
||||
],
|
||||
"filtered_subdirs": ["comfy", "app", "comfy_execution", "utils"],
|
||||
}
|
||||
return None
|
||||
|
||||
def get_sandbox_system_paths(self) -> Optional[List[str]]:
|
||||
"""Returns required application paths to mount in the sandbox."""
|
||||
# By inspecting where our adapter is loaded from, we can determine the comfy root
|
||||
adapter_file = inspect.getfile(self.__class__)
|
||||
# adapter_file = /home/johnj/ComfyUI/comfy/isolation/adapter.py
|
||||
comfy_root = os.path.dirname(os.path.dirname(os.path.dirname(adapter_file)))
|
||||
if os.path.exists(comfy_root):
|
||||
return [comfy_root]
|
||||
return None
|
||||
|
||||
def setup_child_environment(self, snapshot: Dict[str, Any]) -> None:
|
||||
comfy_root = snapshot.get("preferred_root")
|
||||
if not comfy_root:
|
||||
return
|
||||
|
||||
requirements_path = Path(comfy_root) / "requirements.txt"
|
||||
if requirements_path.exists():
|
||||
import re
|
||||
|
||||
for line in requirements_path.read_text().splitlines():
|
||||
line = line.strip()
|
||||
if not line or line.startswith("#"):
|
||||
continue
|
||||
pkg_name = re.split(r"[<>=!~\[]", line)[0].strip()
|
||||
if pkg_name:
|
||||
logging.getLogger(pkg_name).setLevel(logging.ERROR)
|
||||
|
||||
def register_serializers(self, registry: SerializerRegistryProtocol) -> None:
|
||||
if not _IMPORT_TORCH:
|
||||
# Sealed worker without torch — register torch-free TensorValue handler
|
||||
# so IMAGE/MASK/LATENT tensors arrive as numpy arrays, not raw dicts.
|
||||
import numpy as np
|
||||
|
||||
_TORCH_DTYPE_TO_NUMPY = {
|
||||
"torch.float32": np.float32,
|
||||
"torch.float64": np.float64,
|
||||
"torch.float16": np.float16,
|
||||
"torch.bfloat16": np.float32, # numpy has no bfloat16; upcast
|
||||
"torch.int32": np.int32,
|
||||
"torch.int64": np.int64,
|
||||
"torch.int16": np.int16,
|
||||
"torch.int8": np.int8,
|
||||
"torch.uint8": np.uint8,
|
||||
"torch.bool": np.bool_,
|
||||
}
|
||||
|
||||
def _deserialize_tensor_value(data: Dict[str, Any]) -> Any:
|
||||
dtype_str = data["dtype"]
|
||||
np_dtype = _TORCH_DTYPE_TO_NUMPY.get(dtype_str, np.float32)
|
||||
shape = tuple(data["tensor_size"])
|
||||
arr = np.array(data["data"], dtype=np_dtype).reshape(shape)
|
||||
return arr
|
||||
|
||||
_NUMPY_TO_TORCH_DTYPE = {
|
||||
np.float32: "torch.float32",
|
||||
np.float64: "torch.float64",
|
||||
np.float16: "torch.float16",
|
||||
np.int32: "torch.int32",
|
||||
np.int64: "torch.int64",
|
||||
np.int16: "torch.int16",
|
||||
np.int8: "torch.int8",
|
||||
np.uint8: "torch.uint8",
|
||||
np.bool_: "torch.bool",
|
||||
}
|
||||
|
||||
def _serialize_tensor_value(obj: Any) -> Dict[str, Any]:
|
||||
arr = np.asarray(obj, dtype=np.float32) if obj.dtype not in _NUMPY_TO_TORCH_DTYPE else np.asarray(obj)
|
||||
dtype_str = _NUMPY_TO_TORCH_DTYPE.get(arr.dtype.type, "torch.float32")
|
||||
return {
|
||||
"__type__": "TensorValue",
|
||||
"dtype": dtype_str,
|
||||
"tensor_size": list(arr.shape),
|
||||
"requires_grad": False,
|
||||
"data": arr.tolist(),
|
||||
}
|
||||
|
||||
registry.register("TensorValue", _serialize_tensor_value, _deserialize_tensor_value, data_type=True)
|
||||
# ndarray output from sealed workers serializes as TensorValue for host torch reconstruction
|
||||
registry.register("ndarray", _serialize_tensor_value, _deserialize_tensor_value, data_type=True)
|
||||
return
|
||||
|
||||
import torch
|
||||
|
||||
def serialize_device(obj: Any) -> Dict[str, Any]:
|
||||
return {"__type__": "device", "device_str": str(obj)}
|
||||
|
||||
def deserialize_device(data: Dict[str, Any]) -> Any:
|
||||
return torch.device(data["device_str"])
|
||||
|
||||
registry.register("device", serialize_device, deserialize_device)
|
||||
|
||||
_VALID_DTYPES = {
|
||||
"float16", "float32", "float64", "bfloat16",
|
||||
"int8", "int16", "int32", "int64",
|
||||
"uint8", "bool",
|
||||
}
|
||||
|
||||
def serialize_dtype(obj: Any) -> Dict[str, Any]:
|
||||
return {"__type__": "dtype", "dtype_str": str(obj)}
|
||||
|
||||
def deserialize_dtype(data: Dict[str, Any]) -> Any:
|
||||
dtype_name = data["dtype_str"].replace("torch.", "")
|
||||
if dtype_name not in _VALID_DTYPES:
|
||||
raise ValueError(f"Invalid dtype: {data['dtype_str']}")
|
||||
return getattr(torch, dtype_name)
|
||||
|
||||
registry.register("dtype", serialize_dtype, deserialize_dtype)
|
||||
|
||||
from comfy_api.latest._io import FolderType
|
||||
from comfy_api.latest._ui import SavedImages, SavedResult
|
||||
|
||||
def serialize_saved_result(obj: Any) -> Dict[str, Any]:
|
||||
return {
|
||||
"__type__": "SavedResult",
|
||||
"filename": obj.filename,
|
||||
"subfolder": obj.subfolder,
|
||||
"folder_type": obj.type.value,
|
||||
}
|
||||
|
||||
def deserialize_saved_result(data: Dict[str, Any]) -> Any:
|
||||
if isinstance(data, SavedResult):
|
||||
return data
|
||||
folder_type = data["folder_type"] if "folder_type" in data else data["type"]
|
||||
return SavedResult(
|
||||
filename=data["filename"],
|
||||
subfolder=data["subfolder"],
|
||||
type=FolderType(folder_type),
|
||||
)
|
||||
|
||||
registry.register(
|
||||
"SavedResult",
|
||||
serialize_saved_result,
|
||||
deserialize_saved_result,
|
||||
data_type=True,
|
||||
)
|
||||
|
||||
def serialize_saved_images(obj: Any) -> Dict[str, Any]:
|
||||
return {
|
||||
"__type__": "SavedImages",
|
||||
"results": [serialize_saved_result(result) for result in obj.results],
|
||||
"is_animated": obj.is_animated,
|
||||
}
|
||||
|
||||
def deserialize_saved_images(data: Dict[str, Any]) -> Any:
|
||||
return SavedImages(
|
||||
results=[deserialize_saved_result(result) for result in data["results"]],
|
||||
is_animated=data.get("is_animated", False),
|
||||
)
|
||||
|
||||
registry.register(
|
||||
"SavedImages",
|
||||
serialize_saved_images,
|
||||
deserialize_saved_images,
|
||||
data_type=True,
|
||||
)
|
||||
|
||||
def serialize_model_patcher(obj: Any) -> Dict[str, Any]:
|
||||
# Child-side: must already have _instance_id (proxy)
|
||||
if os.environ.get("PYISOLATE_CHILD") == "1":
|
||||
if hasattr(obj, "_instance_id"):
|
||||
return {"__type__": "ModelPatcherRef", "model_id": obj._instance_id}
|
||||
raise RuntimeError(
|
||||
f"ModelPatcher in child lacks _instance_id: "
|
||||
f"{type(obj).__module__}.{type(obj).__name__}"
|
||||
)
|
||||
# Host-side: register with registry
|
||||
if hasattr(obj, "_instance_id"):
|
||||
return {"__type__": "ModelPatcherRef", "model_id": obj._instance_id}
|
||||
model_id = ModelPatcherRegistry().register(obj)
|
||||
return {"__type__": "ModelPatcherRef", "model_id": model_id}
|
||||
|
||||
def deserialize_model_patcher(data: Any) -> Any:
|
||||
"""Deserialize ModelPatcher refs; pass through already-materialized objects."""
|
||||
if isinstance(data, dict):
|
||||
return ModelPatcherProxy(
|
||||
data["model_id"], registry=None, manage_lifecycle=False
|
||||
)
|
||||
return data
|
||||
|
||||
def deserialize_model_patcher_ref(data: Dict[str, Any]) -> Any:
|
||||
"""Context-aware ModelPatcherRef deserializer for both host and child."""
|
||||
is_child = os.environ.get("PYISOLATE_CHILD") == "1"
|
||||
if is_child:
|
||||
return ModelPatcherProxy(
|
||||
data["model_id"], registry=None, manage_lifecycle=False
|
||||
)
|
||||
else:
|
||||
return ModelPatcherRegistry()._get_instance(data["model_id"])
|
||||
|
||||
# Register ModelPatcher type for serialization
|
||||
registry.register(
|
||||
"ModelPatcher", serialize_model_patcher, deserialize_model_patcher
|
||||
)
|
||||
# Register ModelPatcherProxy type (already a proxy, just return ref)
|
||||
registry.register(
|
||||
"ModelPatcherProxy", serialize_model_patcher, deserialize_model_patcher
|
||||
)
|
||||
# Register ModelPatcherRef for deserialization (context-aware: host or child)
|
||||
registry.register("ModelPatcherRef", None, deserialize_model_patcher_ref)
|
||||
|
||||
def serialize_clip(obj: Any) -> Dict[str, Any]:
|
||||
if hasattr(obj, "_instance_id"):
|
||||
return {"__type__": "CLIPRef", "clip_id": obj._instance_id}
|
||||
clip_id = CLIPRegistry().register(obj)
|
||||
return {"__type__": "CLIPRef", "clip_id": clip_id}
|
||||
|
||||
def deserialize_clip(data: Any) -> Any:
|
||||
if isinstance(data, dict):
|
||||
return CLIPProxy(data["clip_id"], registry=None, manage_lifecycle=False)
|
||||
return data
|
||||
|
||||
def deserialize_clip_ref(data: Dict[str, Any]) -> Any:
|
||||
"""Context-aware CLIPRef deserializer for both host and child."""
|
||||
is_child = os.environ.get("PYISOLATE_CHILD") == "1"
|
||||
if is_child:
|
||||
return CLIPProxy(data["clip_id"], registry=None, manage_lifecycle=False)
|
||||
else:
|
||||
return CLIPRegistry()._get_instance(data["clip_id"])
|
||||
|
||||
# Register CLIP type for serialization
|
||||
registry.register("CLIP", serialize_clip, deserialize_clip)
|
||||
# Register CLIPProxy type (already a proxy, just return ref)
|
||||
registry.register("CLIPProxy", serialize_clip, deserialize_clip)
|
||||
# Register CLIPRef for deserialization (context-aware: host or child)
|
||||
registry.register("CLIPRef", None, deserialize_clip_ref)
|
||||
|
||||
def serialize_vae(obj: Any) -> Dict[str, Any]:
|
||||
if hasattr(obj, "_instance_id"):
|
||||
return {"__type__": "VAERef", "vae_id": obj._instance_id}
|
||||
vae_id = VAERegistry().register(obj)
|
||||
return {"__type__": "VAERef", "vae_id": vae_id}
|
||||
|
||||
def deserialize_vae(data: Any) -> Any:
|
||||
if isinstance(data, dict):
|
||||
return VAEProxy(data["vae_id"])
|
||||
return data
|
||||
|
||||
def deserialize_vae_ref(data: Dict[str, Any]) -> Any:
|
||||
"""Context-aware VAERef deserializer for both host and child."""
|
||||
is_child = os.environ.get("PYISOLATE_CHILD") == "1"
|
||||
if is_child:
|
||||
# Child: create a proxy
|
||||
return VAEProxy(data["vae_id"])
|
||||
else:
|
||||
# Host: lookup real VAE from registry
|
||||
return VAERegistry()._get_instance(data["vae_id"])
|
||||
|
||||
# Register VAE type for serialization
|
||||
registry.register("VAE", serialize_vae, deserialize_vae)
|
||||
# Register VAEProxy type (already a proxy, just return ref)
|
||||
registry.register("VAEProxy", serialize_vae, deserialize_vae)
|
||||
# Register VAERef for deserialization (context-aware: host or child)
|
||||
registry.register("VAERef", None, deserialize_vae_ref)
|
||||
|
||||
# ModelSampling serialization - handles ModelSampling* types
|
||||
# copyreg removed - no pickle fallback allowed
|
||||
|
||||
def serialize_model_sampling(obj: Any) -> Dict[str, Any]:
|
||||
# Proxy with _instance_id — return ref (works from both host and child)
|
||||
if hasattr(obj, "_instance_id"):
|
||||
return {"__type__": "ModelSamplingRef", "ms_id": obj._instance_id}
|
||||
# Child-side: object created locally in child (e.g. ModelSamplingAdvanced
|
||||
# in nodes_z_image_turbo.py). Serialize as inline data so the host can
|
||||
# reconstruct the real torch.nn.Module.
|
||||
if os.environ.get("PYISOLATE_CHILD") == "1":
|
||||
import base64
|
||||
import io as _io
|
||||
|
||||
# Identify base classes from comfy.model_sampling
|
||||
bases = []
|
||||
for base in type(obj).__mro__:
|
||||
if base.__module__ == "comfy.model_sampling" and base.__name__ != "object":
|
||||
bases.append(base.__name__)
|
||||
# Serialize state_dict as base64 safetensors-like
|
||||
sd = obj.state_dict()
|
||||
sd_serialized = {}
|
||||
for k, v in sd.items():
|
||||
buf = _io.BytesIO()
|
||||
torch.save(v, buf)
|
||||
sd_serialized[k] = base64.b64encode(buf.getvalue()).decode("ascii")
|
||||
# Capture plain attrs (shift, multiplier, sigma_data, etc.)
|
||||
plain_attrs = {}
|
||||
for k, v in obj.__dict__.items():
|
||||
if k.startswith("_"):
|
||||
continue
|
||||
if isinstance(v, (bool, int, float, str)):
|
||||
plain_attrs[k] = v
|
||||
return {
|
||||
"__type__": "ModelSamplingInline",
|
||||
"bases": bases,
|
||||
"state_dict": sd_serialized,
|
||||
"attrs": plain_attrs,
|
||||
}
|
||||
# Host-side: register with ModelSamplingRegistry and return JSON-safe dict
|
||||
ms_id = ModelSamplingRegistry().register(obj)
|
||||
return {"__type__": "ModelSamplingRef", "ms_id": ms_id}
|
||||
|
||||
def deserialize_model_sampling(data: Any) -> Any:
|
||||
"""Deserialize ModelSampling refs or inline data."""
|
||||
if isinstance(data, dict):
|
||||
if data.get("__type__") == "ModelSamplingInline":
|
||||
return _reconstruct_model_sampling_inline(data)
|
||||
return ModelSamplingProxy(data["ms_id"])
|
||||
return data
|
||||
|
||||
def _reconstruct_model_sampling_inline(data: Dict[str, Any]) -> Any:
|
||||
"""Reconstruct a ModelSampling object on the host from inline child data."""
|
||||
import comfy.model_sampling as _ms
|
||||
import base64
|
||||
import io as _io
|
||||
|
||||
# Resolve base classes
|
||||
base_classes = []
|
||||
for name in data["bases"]:
|
||||
cls = getattr(_ms, name, None)
|
||||
if cls is not None:
|
||||
base_classes.append(cls)
|
||||
if not base_classes:
|
||||
raise RuntimeError(
|
||||
f"Cannot reconstruct ModelSampling: no known bases in {data['bases']}"
|
||||
)
|
||||
# Create dynamic class matching the child's class hierarchy
|
||||
ReconstructedSampling = type("ReconstructedSampling", tuple(base_classes), {})
|
||||
obj = ReconstructedSampling.__new__(ReconstructedSampling)
|
||||
torch.nn.Module.__init__(obj)
|
||||
# Restore plain attributes first
|
||||
for k, v in data.get("attrs", {}).items():
|
||||
setattr(obj, k, v)
|
||||
# Restore state_dict (buffers like sigmas)
|
||||
for k, v_b64 in data.get("state_dict", {}).items():
|
||||
buf = _io.BytesIO(base64.b64decode(v_b64))
|
||||
tensor = torch.load(buf, weights_only=True)
|
||||
# Register as buffer so it's part of state_dict
|
||||
parts = k.split(".")
|
||||
if len(parts) == 1:
|
||||
cast(Any, obj).register_buffer(parts[0], tensor) # pylint: disable=no-member
|
||||
else:
|
||||
setattr(obj, parts[0], tensor)
|
||||
# Register on host so future references use proxy pattern.
|
||||
# Skip in child process — register() is async RPC and cannot be
|
||||
# called synchronously during deserialization.
|
||||
if os.environ.get("PYISOLATE_CHILD") != "1":
|
||||
ModelSamplingRegistry().register(obj)
|
||||
return obj
|
||||
|
||||
def deserialize_model_sampling_ref(data: Dict[str, Any]) -> Any:
|
||||
"""Context-aware ModelSamplingRef deserializer for both host and child."""
|
||||
is_child = os.environ.get("PYISOLATE_CHILD") == "1"
|
||||
if is_child:
|
||||
return ModelSamplingProxy(data["ms_id"])
|
||||
else:
|
||||
return ModelSamplingRegistry()._get_instance(data["ms_id"])
|
||||
|
||||
# Register all ModelSampling* and StableCascadeSampling classes dynamically
|
||||
import comfy.model_sampling
|
||||
|
||||
for ms_cls in vars(comfy.model_sampling).values():
|
||||
if not isinstance(ms_cls, type):
|
||||
continue
|
||||
if not issubclass(ms_cls, torch.nn.Module):
|
||||
continue
|
||||
if not (ms_cls.__name__.startswith("ModelSampling") or ms_cls.__name__ == "StableCascadeSampling"):
|
||||
continue
|
||||
registry.register(
|
||||
ms_cls.__name__,
|
||||
serialize_model_sampling,
|
||||
deserialize_model_sampling,
|
||||
)
|
||||
registry.register(
|
||||
"ModelSamplingProxy", serialize_model_sampling, deserialize_model_sampling
|
||||
)
|
||||
# Register ModelSamplingRef for deserialization (context-aware: host or child)
|
||||
registry.register("ModelSamplingRef", None, deserialize_model_sampling_ref)
|
||||
# Register ModelSamplingInline for deserialization (child→host inline transfer)
|
||||
registry.register(
|
||||
"ModelSamplingInline", None, lambda data: _reconstruct_model_sampling_inline(data)
|
||||
)
|
||||
|
||||
def serialize_cond(obj: Any) -> Dict[str, Any]:
|
||||
type_key = f"{type(obj).__module__}.{type(obj).__name__}"
|
||||
return {
|
||||
"__type__": type_key,
|
||||
"cond": obj.cond,
|
||||
}
|
||||
|
||||
def deserialize_cond(data: Dict[str, Any]) -> Any:
|
||||
import importlib
|
||||
|
||||
type_key = data["__type__"]
|
||||
module_name, class_name = type_key.rsplit(".", 1)
|
||||
module = importlib.import_module(module_name)
|
||||
cls = getattr(module, class_name)
|
||||
return cls(data["cond"])
|
||||
|
||||
def _serialize_public_state(obj: Any) -> Dict[str, Any]:
|
||||
state: Dict[str, Any] = {}
|
||||
for key, value in obj.__dict__.items():
|
||||
if key.startswith("_"):
|
||||
continue
|
||||
if callable(value):
|
||||
continue
|
||||
state[key] = value
|
||||
return state
|
||||
|
||||
def serialize_latent_format(obj: Any) -> Dict[str, Any]:
|
||||
type_key = f"{type(obj).__module__}.{type(obj).__name__}"
|
||||
return {
|
||||
"__type__": type_key,
|
||||
"state": _serialize_public_state(obj),
|
||||
}
|
||||
|
||||
def deserialize_latent_format(data: Dict[str, Any]) -> Any:
|
||||
import importlib
|
||||
|
||||
type_key = data["__type__"]
|
||||
module_name, class_name = type_key.rsplit(".", 1)
|
||||
module = importlib.import_module(module_name)
|
||||
cls = getattr(module, class_name)
|
||||
obj = cls()
|
||||
for key, value in data.get("state", {}).items():
|
||||
prop = getattr(type(obj), key, None)
|
||||
if isinstance(prop, property) and prop.fset is None:
|
||||
continue
|
||||
setattr(obj, key, value)
|
||||
return obj
|
||||
|
||||
import comfy.conds
|
||||
|
||||
for cond_cls in vars(comfy.conds).values():
|
||||
if not isinstance(cond_cls, type):
|
||||
continue
|
||||
if not issubclass(cond_cls, comfy.conds.CONDRegular):
|
||||
continue
|
||||
type_key = f"{cond_cls.__module__}.{cond_cls.__name__}"
|
||||
registry.register(type_key, serialize_cond, deserialize_cond)
|
||||
registry.register(cond_cls.__name__, serialize_cond, deserialize_cond)
|
||||
|
||||
import comfy.latent_formats
|
||||
|
||||
for latent_cls in vars(comfy.latent_formats).values():
|
||||
if not isinstance(latent_cls, type):
|
||||
continue
|
||||
if not issubclass(latent_cls, comfy.latent_formats.LatentFormat):
|
||||
continue
|
||||
type_key = f"{latent_cls.__module__}.{latent_cls.__name__}"
|
||||
registry.register(
|
||||
type_key, serialize_latent_format, deserialize_latent_format
|
||||
)
|
||||
registry.register(
|
||||
latent_cls.__name__, serialize_latent_format, deserialize_latent_format
|
||||
)
|
||||
|
||||
# V3 API: unwrap NodeOutput.args
|
||||
def deserialize_node_output(data: Any) -> Any:
|
||||
return getattr(data, "args", data)
|
||||
|
||||
registry.register("NodeOutput", None, deserialize_node_output)
|
||||
|
||||
# KSAMPLER serializer: stores sampler name instead of function object
|
||||
# sampler_function is a callable which gets filtered out by JSONSocketTransport
|
||||
def serialize_ksampler(obj: Any) -> Dict[str, Any]:
|
||||
func_name = obj.sampler_function.__name__
|
||||
# Map function name back to sampler name
|
||||
if func_name == "sample_unipc":
|
||||
sampler_name = "uni_pc"
|
||||
elif func_name == "sample_unipc_bh2":
|
||||
sampler_name = "uni_pc_bh2"
|
||||
elif func_name == "dpm_fast_function":
|
||||
sampler_name = "dpm_fast"
|
||||
elif func_name == "dpm_adaptive_function":
|
||||
sampler_name = "dpm_adaptive"
|
||||
elif func_name.startswith("sample_"):
|
||||
sampler_name = func_name[7:] # Remove "sample_" prefix
|
||||
else:
|
||||
sampler_name = func_name
|
||||
return {
|
||||
"__type__": "KSAMPLER",
|
||||
"sampler_name": sampler_name,
|
||||
"extra_options": obj.extra_options,
|
||||
"inpaint_options": obj.inpaint_options,
|
||||
}
|
||||
|
||||
def deserialize_ksampler(data: Dict[str, Any]) -> Any:
|
||||
import comfy.samplers
|
||||
|
||||
return comfy.samplers.ksampler(
|
||||
data["sampler_name"],
|
||||
data.get("extra_options", {}),
|
||||
data.get("inpaint_options", {}),
|
||||
)
|
||||
|
||||
registry.register("KSAMPLER", serialize_ksampler, deserialize_ksampler)
|
||||
|
||||
from comfy.isolation.model_patcher_proxy_utils import register_hooks_serializers
|
||||
|
||||
register_hooks_serializers(registry)
|
||||
|
||||
# -- File3D (comfy_api.latest._util.geometry_types) ---------------------
|
||||
# Origin: comfy_api by ComfyOrg (Alexander Piskun), PR #12129
|
||||
|
||||
def serialize_file3d(obj: Any) -> Dict[str, Any]:
|
||||
import base64
|
||||
return {
|
||||
"__type__": "File3D",
|
||||
"format": obj.format,
|
||||
"data": base64.b64encode(obj.get_bytes()).decode("ascii"),
|
||||
}
|
||||
|
||||
def deserialize_file3d(data: Any) -> Any:
|
||||
import base64
|
||||
from io import BytesIO
|
||||
from comfy_api.latest._util.geometry_types import File3D
|
||||
return File3D(BytesIO(base64.b64decode(data["data"])), file_format=data["format"])
|
||||
|
||||
registry.register("File3D", serialize_file3d, deserialize_file3d, data_type=True)
|
||||
|
||||
# -- VIDEO (comfy_api.latest._input_impl.video_types) -------------------
|
||||
# Origin: ComfyAPI Core v0.0.2 by ComfyOrg (guill), PR #8962
|
||||
|
||||
def serialize_video(obj: Any) -> Dict[str, Any]:
|
||||
components = obj.get_components()
|
||||
images = components.images.detach() if components.images.requires_grad else components.images
|
||||
result: Dict[str, Any] = {
|
||||
"__type__": "VIDEO",
|
||||
"images": images,
|
||||
"frame_rate_num": components.frame_rate.numerator,
|
||||
"frame_rate_den": components.frame_rate.denominator,
|
||||
}
|
||||
if components.audio is not None:
|
||||
waveform = components.audio["waveform"]
|
||||
if waveform.requires_grad:
|
||||
waveform = waveform.detach()
|
||||
result["audio_waveform"] = waveform
|
||||
result["audio_sample_rate"] = components.audio["sample_rate"]
|
||||
if components.metadata is not None:
|
||||
result["metadata"] = components.metadata
|
||||
return result
|
||||
|
||||
def deserialize_video(data: Any) -> Any:
|
||||
from fractions import Fraction
|
||||
from comfy_api.latest._input_impl.video_types import VideoFromComponents
|
||||
from comfy_api.latest._util.video_types import VideoComponents
|
||||
audio = None
|
||||
if "audio_waveform" in data:
|
||||
audio = {"waveform": data["audio_waveform"], "sample_rate": data["audio_sample_rate"]}
|
||||
components = VideoComponents(
|
||||
images=data["images"],
|
||||
frame_rate=Fraction(data["frame_rate_num"], data["frame_rate_den"]),
|
||||
audio=audio,
|
||||
metadata=data.get("metadata"),
|
||||
)
|
||||
return VideoFromComponents(components)
|
||||
|
||||
registry.register("VIDEO", serialize_video, deserialize_video, data_type=True)
|
||||
registry.register("VideoFromFile", serialize_video, deserialize_video, data_type=True)
|
||||
registry.register("VideoFromComponents", serialize_video, deserialize_video, data_type=True)
|
||||
|
||||
def setup_web_directory(self, module: Any) -> None:
|
||||
"""Detect WEB_DIRECTORY on a module and populate/register it.
|
||||
|
||||
Called by the sealed worker after loading the node module.
|
||||
Mirrors extension_wrapper.py:216-227 for host-coupled nodes.
|
||||
Does NOT import extension_wrapper.py (it has `import torch` at module level).
|
||||
"""
|
||||
import shutil
|
||||
|
||||
web_dir_attr = getattr(module, "WEB_DIRECTORY", None)
|
||||
if web_dir_attr is None:
|
||||
return
|
||||
|
||||
module_dir = os.path.dirname(os.path.abspath(module.__file__))
|
||||
web_dir_path = os.path.abspath(os.path.join(module_dir, web_dir_attr))
|
||||
|
||||
# Read extension name from pyproject.toml
|
||||
ext_name = os.path.basename(module_dir)
|
||||
pyproject = os.path.join(module_dir, "pyproject.toml")
|
||||
if os.path.exists(pyproject):
|
||||
try:
|
||||
import tomllib
|
||||
except ImportError:
|
||||
import tomli as tomllib # type: ignore[no-redef]
|
||||
try:
|
||||
with open(pyproject, "rb") as f:
|
||||
data = tomllib.load(f)
|
||||
name = data.get("project", {}).get("name")
|
||||
if name:
|
||||
ext_name = name
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# Populate web dir if empty (mirrors _run_prestartup_web_copy)
|
||||
if not (os.path.isdir(web_dir_path) and any(os.scandir(web_dir_path))):
|
||||
os.makedirs(web_dir_path, exist_ok=True)
|
||||
|
||||
# Module-defined copy spec
|
||||
copy_spec = getattr(module, "_PRESTARTUP_WEB_COPY", None)
|
||||
if copy_spec is not None and callable(copy_spec):
|
||||
try:
|
||||
copy_spec(web_dir_path)
|
||||
except Exception as e:
|
||||
logger.warning("][ _PRESTARTUP_WEB_COPY failed: %s", e)
|
||||
|
||||
# Fallback: comfy_3d_viewers
|
||||
try:
|
||||
from comfy_3d_viewers import copy_viewer, VIEWER_FILES
|
||||
for viewer in VIEWER_FILES:
|
||||
try:
|
||||
copy_viewer(viewer, web_dir_path)
|
||||
except Exception:
|
||||
pass
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
# Fallback: comfy_dynamic_widgets
|
||||
try:
|
||||
from comfy_dynamic_widgets import get_js_path
|
||||
src = os.path.realpath(get_js_path())
|
||||
if os.path.exists(src):
|
||||
dst_dir = os.path.join(web_dir_path, "js")
|
||||
os.makedirs(dst_dir, exist_ok=True)
|
||||
shutil.copy2(src, os.path.join(dst_dir, "dynamic_widgets.js"))
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
if os.path.isdir(web_dir_path) and any(os.scandir(web_dir_path)):
|
||||
WebDirectoryProxy.register_web_dir(ext_name, web_dir_path)
|
||||
logger.info(
|
||||
"][ Adapter: registered web dir for %s (%d files)",
|
||||
ext_name,
|
||||
sum(1 for _ in Path(web_dir_path).rglob("*") if _.is_file()),
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def register_host_event_handlers(extension: Any) -> None:
|
||||
"""Register host-side event handlers for an isolated extension.
|
||||
|
||||
Wires ``"progress"`` events from the child to ``comfy.utils.PROGRESS_BAR_HOOK``
|
||||
so the ComfyUI frontend receives progress bar updates.
|
||||
"""
|
||||
register_event_handler = inspect.getattr_static(
|
||||
extension, "register_event_handler", None
|
||||
)
|
||||
if not callable(register_event_handler):
|
||||
return
|
||||
|
||||
def _host_progress_handler(payload: dict) -> None:
|
||||
import comfy.utils
|
||||
|
||||
hook = comfy.utils.PROGRESS_BAR_HOOK
|
||||
if hook is not None:
|
||||
hook(
|
||||
payload.get("value", 0),
|
||||
payload.get("total", 0),
|
||||
payload.get("preview"),
|
||||
payload.get("node_id"),
|
||||
)
|
||||
|
||||
extension.register_event_handler("progress", _host_progress_handler)
|
||||
|
||||
def setup_child_event_hooks(self, extension: Any) -> None:
|
||||
"""Wire PROGRESS_BAR_HOOK in the child to emit_event on the extension.
|
||||
|
||||
Host-coupled only — sealed workers do not have comfy.utils (torch).
|
||||
"""
|
||||
is_child = os.environ.get("PYISOLATE_CHILD") == "1"
|
||||
logger.info("][ ISO:setup_child_event_hooks called, PYISOLATE_CHILD=%s", is_child)
|
||||
if not is_child:
|
||||
return
|
||||
|
||||
if not _IMPORT_TORCH:
|
||||
logger.info("][ ISO:setup_child_event_hooks skipped — sealed worker (no torch)")
|
||||
return
|
||||
|
||||
import comfy.utils
|
||||
|
||||
def _event_progress_hook(value, total, preview=None, node_id=None):
|
||||
logger.debug("][ ISO:event_progress value=%s/%s node_id=%s", value, total, node_id)
|
||||
extension.emit_event("progress", {
|
||||
"value": value,
|
||||
"total": total,
|
||||
"node_id": node_id,
|
||||
})
|
||||
|
||||
comfy.utils.PROGRESS_BAR_HOOK = _event_progress_hook
|
||||
logger.info("][ ISO:PROGRESS_BAR_HOOK wired to event channel")
|
||||
|
||||
def provide_rpc_services(self) -> List[type[ProxiedSingleton]]:
|
||||
# Always available — no torch/PIL dependency
|
||||
services: List[type[ProxiedSingleton]] = [
|
||||
FolderPathsProxy,
|
||||
HelperProxiesService,
|
||||
WebDirectoryProxy,
|
||||
]
|
||||
# Torch/PIL-dependent proxies
|
||||
if _HAS_TORCH_PROXIES:
|
||||
services.extend([
|
||||
PromptServerService,
|
||||
ModelManagementProxy,
|
||||
UtilsProxy,
|
||||
ProgressProxy,
|
||||
VAERegistry,
|
||||
CLIPRegistry,
|
||||
ModelPatcherRegistry,
|
||||
ModelSamplingRegistry,
|
||||
FirstStageModelRegistry,
|
||||
])
|
||||
return services
|
||||
|
||||
def handle_api_registration(self, api: ProxiedSingleton, rpc: AsyncRPC) -> None:
|
||||
# Resolve the real name whether it's an instance or the Singleton class itself
|
||||
api_name = api.__name__ if isinstance(api, type) else api.__class__.__name__
|
||||
|
||||
if api_name == "FolderPathsProxy":
|
||||
import folder_paths
|
||||
|
||||
# Replace module-level functions with proxy methods
|
||||
# This is aggressive but necessary for transparent proxying
|
||||
# Handle both instance and class cases
|
||||
instance = api() if isinstance(api, type) else api
|
||||
for name in dir(instance):
|
||||
if not name.startswith("_"):
|
||||
setattr(folder_paths, name, getattr(instance, name))
|
||||
|
||||
# Fence: isolated children get writable temp inside sandbox
|
||||
if os.environ.get("PYISOLATE_CHILD") == "1":
|
||||
import tempfile
|
||||
_child_temp = os.path.join(tempfile.gettempdir(), "comfyui_temp")
|
||||
os.makedirs(_child_temp, exist_ok=True)
|
||||
folder_paths.temp_directory = _child_temp
|
||||
|
||||
return
|
||||
|
||||
if api_name == "ModelManagementProxy":
|
||||
if _IMPORT_TORCH:
|
||||
import comfy.model_management
|
||||
|
||||
instance = api() if isinstance(api, type) else api
|
||||
# Replace module-level functions with proxy methods
|
||||
for name in dir(instance):
|
||||
if not name.startswith("_"):
|
||||
setattr(comfy.model_management, name, getattr(instance, name))
|
||||
return
|
||||
|
||||
if api_name == "UtilsProxy":
|
||||
if not _IMPORT_TORCH:
|
||||
logger.info("][ ISO:UtilsProxy handle_api_registration skipped — sealed worker (no torch)")
|
||||
return
|
||||
|
||||
import comfy.utils
|
||||
|
||||
# Static Injection of RPC mechanism to ensure Child can access it
|
||||
# independent of instance lifecycle.
|
||||
api.set_rpc(rpc)
|
||||
|
||||
is_child = os.environ.get("PYISOLATE_CHILD") == "1"
|
||||
logger.info("][ ISO:UtilsProxy handle_api_registration PYISOLATE_CHILD=%s", is_child)
|
||||
|
||||
# Progress hook wiring moved to setup_child_event_hooks via event channel
|
||||
|
||||
return
|
||||
|
||||
if api_name == "PromptServerService":
|
||||
if not _IMPORT_TORCH:
|
||||
return
|
||||
import server
|
||||
from comfy.isolation.proxies.prompt_server_impl import PromptServerStub
|
||||
|
||||
stub = PromptServerStub()
|
||||
if (
|
||||
hasattr(server, "PromptServer")
|
||||
and getattr(server.PromptServer, "instance", None) is not stub
|
||||
):
|
||||
server.PromptServer.instance = stub
|
||||
@@ -1,122 +0,0 @@
|
||||
# pylint: disable=import-outside-toplevel,logging-fstring-interpolation
|
||||
# Child process initialization for PyIsolate
|
||||
import logging
|
||||
import os
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def is_child_process() -> bool:
|
||||
return os.environ.get("PYISOLATE_CHILD") == "1"
|
||||
|
||||
|
||||
def _load_extra_model_paths() -> None:
|
||||
"""Load extra_model_paths.yaml so the child's folder_paths has the same search paths as the host.
|
||||
|
||||
The host loads this in main.py:143-145. The child is spawned by
|
||||
pyisolate's uds_client.py and never runs main.py, so folder_paths
|
||||
only has the base model directories. Any isolated node calling
|
||||
folder_paths.get_filename_list() in define_schema() would get empty
|
||||
results for folders whose files live in extra_model_paths locations.
|
||||
"""
|
||||
import folder_paths # noqa: F401 — side-effect import; load_extra_path_config writes to folder_paths internals
|
||||
from utils.extra_config import load_extra_path_config
|
||||
|
||||
extra_config_path = os.path.join(
|
||||
os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))),
|
||||
"extra_model_paths.yaml",
|
||||
)
|
||||
if os.path.isfile(extra_config_path):
|
||||
load_extra_path_config(extra_config_path)
|
||||
|
||||
|
||||
def initialize_child_process() -> None:
|
||||
if os.environ.get("PYISOLATE_IMPORT_TORCH", "1") != "0":
|
||||
_load_extra_model_paths()
|
||||
_setup_child_loop_bridge()
|
||||
|
||||
# Manual RPC injection
|
||||
try:
|
||||
from pyisolate._internal.rpc_protocol import get_child_rpc_instance
|
||||
|
||||
rpc = get_child_rpc_instance()
|
||||
if rpc:
|
||||
_setup_proxy_callers(rpc)
|
||||
else:
|
||||
_setup_proxy_callers()
|
||||
except Exception as e:
|
||||
logger.error(f"][ child_hooks Manual RPC Injection failed: {e}")
|
||||
_setup_proxy_callers()
|
||||
|
||||
_setup_logging()
|
||||
|
||||
|
||||
def _setup_child_loop_bridge() -> None:
|
||||
import asyncio
|
||||
|
||||
main_loop = None
|
||||
try:
|
||||
main_loop = asyncio.get_running_loop()
|
||||
except RuntimeError:
|
||||
try:
|
||||
main_loop = asyncio.get_event_loop()
|
||||
except RuntimeError:
|
||||
pass
|
||||
|
||||
if main_loop is None:
|
||||
return
|
||||
|
||||
try:
|
||||
from .proxies.base import set_global_loop
|
||||
|
||||
set_global_loop(main_loop)
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
|
||||
def _setup_prompt_server_stub(rpc=None) -> None:
|
||||
try:
|
||||
from .proxies.prompt_server_impl import PromptServerStub
|
||||
|
||||
if rpc:
|
||||
PromptServerStub.set_rpc(rpc)
|
||||
elif hasattr(PromptServerStub, "clear_rpc"):
|
||||
PromptServerStub.clear_rpc()
|
||||
else:
|
||||
PromptServerStub._rpc = None # type: ignore[attr-defined]
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to setup PromptServerStub: {e}")
|
||||
|
||||
|
||||
def _setup_proxy_callers(rpc=None) -> None:
|
||||
try:
|
||||
from .proxies.folder_paths_proxy import FolderPathsProxy
|
||||
from .proxies.helper_proxies import HelperProxiesService
|
||||
from .proxies.model_management_proxy import ModelManagementProxy
|
||||
from .proxies.progress_proxy import ProgressProxy
|
||||
from .proxies.prompt_server_impl import PromptServerStub
|
||||
from .proxies.utils_proxy import UtilsProxy
|
||||
|
||||
if rpc is None:
|
||||
FolderPathsProxy.clear_rpc()
|
||||
HelperProxiesService.clear_rpc()
|
||||
ModelManagementProxy.clear_rpc()
|
||||
ProgressProxy.clear_rpc()
|
||||
PromptServerStub.clear_rpc()
|
||||
UtilsProxy.clear_rpc()
|
||||
return
|
||||
|
||||
FolderPathsProxy.set_rpc(rpc)
|
||||
HelperProxiesService.set_rpc(rpc)
|
||||
ModelManagementProxy.set_rpc(rpc)
|
||||
ProgressProxy.set_rpc(rpc)
|
||||
PromptServerStub.set_rpc(rpc)
|
||||
UtilsProxy.set_rpc(rpc)
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to setup child singleton proxy callers: {e}")
|
||||
|
||||
|
||||
def _setup_logging() -> None:
|
||||
logging.getLogger().setLevel(logging.INFO)
|
||||
@@ -1,327 +0,0 @@
|
||||
# pylint: disable=attribute-defined-outside-init,import-outside-toplevel,logging-fstring-interpolation
|
||||
# CLIP Proxy implementation
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from typing import TYPE_CHECKING, Any, Optional
|
||||
|
||||
from comfy.isolation.proxies.base import (
|
||||
IS_CHILD_PROCESS,
|
||||
BaseProxy,
|
||||
BaseRegistry,
|
||||
detach_if_grad,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from comfy.isolation.model_patcher_proxy import ModelPatcherProxy
|
||||
|
||||
|
||||
class CondStageModelRegistry(BaseRegistry[Any]):
|
||||
_type_prefix = "cond_stage_model"
|
||||
|
||||
async def get_property(self, instance_id: str, name: str) -> Any:
|
||||
obj = self._get_instance(instance_id)
|
||||
return getattr(obj, name)
|
||||
|
||||
|
||||
class CondStageModelProxy(BaseProxy[CondStageModelRegistry]):
|
||||
_registry_class = CondStageModelRegistry
|
||||
__module__ = "comfy.sd"
|
||||
|
||||
def __getattr__(self, name: str) -> Any:
|
||||
try:
|
||||
return self._call_rpc("get_property", name)
|
||||
except Exception as e:
|
||||
raise AttributeError(
|
||||
f"'{self.__class__.__name__}' object has no attribute '{name}'"
|
||||
) from e
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return f"<CondStageModelProxy {self._instance_id}>"
|
||||
|
||||
|
||||
class TokenizerRegistry(BaseRegistry[Any]):
|
||||
_type_prefix = "tokenizer"
|
||||
|
||||
async def get_property(self, instance_id: str, name: str) -> Any:
|
||||
obj = self._get_instance(instance_id)
|
||||
return getattr(obj, name)
|
||||
|
||||
|
||||
class TokenizerProxy(BaseProxy[TokenizerRegistry]):
|
||||
_registry_class = TokenizerRegistry
|
||||
__module__ = "comfy.sd"
|
||||
|
||||
def __getattr__(self, name: str) -> Any:
|
||||
try:
|
||||
return self._call_rpc("get_property", name)
|
||||
except Exception as e:
|
||||
raise AttributeError(
|
||||
f"'{self.__class__.__name__}' object has no attribute '{name}'"
|
||||
) from e
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return f"<TokenizerProxy {self._instance_id}>"
|
||||
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class CLIPRegistry(BaseRegistry[Any]):
|
||||
_type_prefix = "clip"
|
||||
_allowed_setters = {
|
||||
"layer_idx",
|
||||
"tokenizer_options",
|
||||
"use_clip_schedule",
|
||||
"apply_hooks_to_conds",
|
||||
}
|
||||
|
||||
async def get_ram_usage(self, instance_id: str) -> int:
|
||||
return self._get_instance(instance_id).get_ram_usage()
|
||||
|
||||
async def get_patcher_id(self, instance_id: str) -> str:
|
||||
from comfy.isolation.model_patcher_proxy import ModelPatcherRegistry
|
||||
|
||||
return ModelPatcherRegistry().register(self._get_instance(instance_id).patcher)
|
||||
|
||||
async def get_cond_stage_model_id(self, instance_id: str) -> str:
|
||||
return CondStageModelRegistry().register(
|
||||
self._get_instance(instance_id).cond_stage_model
|
||||
)
|
||||
|
||||
async def get_tokenizer_id(self, instance_id: str) -> str:
|
||||
return TokenizerRegistry().register(self._get_instance(instance_id).tokenizer)
|
||||
|
||||
async def load_model(self, instance_id: str) -> None:
|
||||
self._get_instance(instance_id).load_model()
|
||||
|
||||
async def clip_layer(self, instance_id: str, layer_idx: int) -> None:
|
||||
self._get_instance(instance_id).clip_layer(layer_idx)
|
||||
|
||||
async def set_tokenizer_option(
|
||||
self, instance_id: str, option_name: str, value: Any
|
||||
) -> None:
|
||||
self._get_instance(instance_id).set_tokenizer_option(option_name, value)
|
||||
|
||||
async def get_property(self, instance_id: str, name: str) -> Any:
|
||||
return getattr(self._get_instance(instance_id), name)
|
||||
|
||||
async def set_property(self, instance_id: str, name: str, value: Any) -> None:
|
||||
if name not in self._allowed_setters:
|
||||
raise PermissionError(f"Setting '{name}' is not allowed via RPC")
|
||||
setattr(self._get_instance(instance_id), name, value)
|
||||
|
||||
async def tokenize(
|
||||
self, instance_id: str, text: str, return_word_ids: bool = False, **kwargs: Any
|
||||
) -> Any:
|
||||
return self._get_instance(instance_id).tokenize(
|
||||
text, return_word_ids=return_word_ids, **kwargs
|
||||
)
|
||||
|
||||
async def encode(self, instance_id: str, text: str) -> Any:
|
||||
return detach_if_grad(self._get_instance(instance_id).encode(text))
|
||||
|
||||
async def encode_from_tokens(
|
||||
self,
|
||||
instance_id: str,
|
||||
tokens: Any,
|
||||
return_pooled: bool = False,
|
||||
return_dict: bool = False,
|
||||
) -> Any:
|
||||
return detach_if_grad(
|
||||
self._get_instance(instance_id).encode_from_tokens(
|
||||
tokens, return_pooled=return_pooled, return_dict=return_dict
|
||||
)
|
||||
)
|
||||
|
||||
async def encode_from_tokens_scheduled(
|
||||
self,
|
||||
instance_id: str,
|
||||
tokens: Any,
|
||||
unprojected: bool = False,
|
||||
add_dict: Optional[dict] = None,
|
||||
show_pbar: bool = True,
|
||||
) -> Any:
|
||||
add_dict = add_dict or {}
|
||||
return detach_if_grad(
|
||||
self._get_instance(instance_id).encode_from_tokens_scheduled(
|
||||
tokens, unprojected=unprojected, add_dict=add_dict, show_pbar=show_pbar
|
||||
)
|
||||
)
|
||||
|
||||
async def add_patches(
|
||||
self,
|
||||
instance_id: str,
|
||||
patches: Any,
|
||||
strength_patch: float = 1.0,
|
||||
strength_model: float = 1.0,
|
||||
) -> Any:
|
||||
return self._get_instance(instance_id).add_patches(
|
||||
patches, strength_patch=strength_patch, strength_model=strength_model
|
||||
)
|
||||
|
||||
async def get_key_patches(self, instance_id: str) -> Any:
|
||||
return self._get_instance(instance_id).get_key_patches()
|
||||
|
||||
async def load_sd(
|
||||
self, instance_id: str, sd: dict, full_model: bool = False
|
||||
) -> Any:
|
||||
return self._get_instance(instance_id).load_sd(sd, full_model=full_model)
|
||||
|
||||
async def get_sd(self, instance_id: str) -> Any:
|
||||
return self._get_instance(instance_id).get_sd()
|
||||
|
||||
async def clone(self, instance_id: str) -> str:
|
||||
return self.register(self._get_instance(instance_id).clone())
|
||||
|
||||
|
||||
class CLIPProxy(BaseProxy[CLIPRegistry]):
|
||||
_registry_class = CLIPRegistry
|
||||
__module__ = "comfy.sd"
|
||||
|
||||
def get_ram_usage(self) -> int:
|
||||
return self._call_rpc("get_ram_usage")
|
||||
|
||||
@property
|
||||
def patcher(self) -> "ModelPatcherProxy":
|
||||
from comfy.isolation.model_patcher_proxy import ModelPatcherProxy
|
||||
|
||||
if not hasattr(self, "_patcher_proxy"):
|
||||
patcher_id = self._call_rpc("get_patcher_id")
|
||||
self._patcher_proxy = ModelPatcherProxy(patcher_id, manage_lifecycle=False)
|
||||
return self._patcher_proxy
|
||||
|
||||
@patcher.setter
|
||||
def patcher(self, value: Any) -> None:
|
||||
from comfy.isolation.model_patcher_proxy import ModelPatcherProxy
|
||||
|
||||
if isinstance(value, ModelPatcherProxy):
|
||||
self._patcher_proxy = value
|
||||
else:
|
||||
logger.warning(
|
||||
f"Attempted to set CLIPProxy.patcher to non-proxy object: {value}"
|
||||
)
|
||||
|
||||
@property
|
||||
def cond_stage_model(self) -> CondStageModelProxy:
|
||||
if not hasattr(self, "_cond_stage_model_proxy"):
|
||||
csm_id = self._call_rpc("get_cond_stage_model_id")
|
||||
self._cond_stage_model_proxy = CondStageModelProxy(
|
||||
csm_id, manage_lifecycle=False
|
||||
)
|
||||
return self._cond_stage_model_proxy
|
||||
|
||||
@property
|
||||
def tokenizer(self) -> TokenizerProxy:
|
||||
if not hasattr(self, "_tokenizer_proxy"):
|
||||
tok_id = self._call_rpc("get_tokenizer_id")
|
||||
self._tokenizer_proxy = TokenizerProxy(tok_id, manage_lifecycle=False)
|
||||
return self._tokenizer_proxy
|
||||
|
||||
def load_model(self) -> ModelPatcherProxy:
|
||||
self._call_rpc("load_model")
|
||||
return self.patcher
|
||||
|
||||
@property
|
||||
def layer_idx(self) -> Optional[int]:
|
||||
return self._call_rpc("get_property", "layer_idx")
|
||||
|
||||
@layer_idx.setter
|
||||
def layer_idx(self, value: Optional[int]) -> None:
|
||||
self._call_rpc("set_property", "layer_idx", value)
|
||||
|
||||
@property
|
||||
def tokenizer_options(self) -> dict:
|
||||
return self._call_rpc("get_property", "tokenizer_options")
|
||||
|
||||
@tokenizer_options.setter
|
||||
def tokenizer_options(self, value: dict) -> None:
|
||||
self._call_rpc("set_property", "tokenizer_options", value)
|
||||
|
||||
@property
|
||||
def use_clip_schedule(self) -> bool:
|
||||
return self._call_rpc("get_property", "use_clip_schedule")
|
||||
|
||||
@use_clip_schedule.setter
|
||||
def use_clip_schedule(self, value: bool) -> None:
|
||||
self._call_rpc("set_property", "use_clip_schedule", value)
|
||||
|
||||
@property
|
||||
def apply_hooks_to_conds(self) -> Any:
|
||||
return self._call_rpc("get_property", "apply_hooks_to_conds")
|
||||
|
||||
@apply_hooks_to_conds.setter
|
||||
def apply_hooks_to_conds(self, value: Any) -> None:
|
||||
self._call_rpc("set_property", "apply_hooks_to_conds", value)
|
||||
|
||||
def clip_layer(self, layer_idx: int) -> None:
|
||||
return self._call_rpc("clip_layer", layer_idx)
|
||||
|
||||
def set_tokenizer_option(self, option_name: str, value: Any) -> None:
|
||||
return self._call_rpc("set_tokenizer_option", option_name, value)
|
||||
|
||||
def tokenize(self, text: str, return_word_ids: bool = False, **kwargs: Any) -> Any:
|
||||
return self._call_rpc(
|
||||
"tokenize", text, return_word_ids=return_word_ids, **kwargs
|
||||
)
|
||||
|
||||
def encode(self, text: str) -> Any:
|
||||
return self._call_rpc("encode", text)
|
||||
|
||||
def encode_from_tokens(
|
||||
self, tokens: Any, return_pooled: bool = False, return_dict: bool = False
|
||||
) -> Any:
|
||||
res = self._call_rpc(
|
||||
"encode_from_tokens",
|
||||
tokens,
|
||||
return_pooled=return_pooled,
|
||||
return_dict=return_dict,
|
||||
)
|
||||
if return_pooled and isinstance(res, list) and not return_dict:
|
||||
return tuple(res)
|
||||
return res
|
||||
|
||||
def encode_from_tokens_scheduled(
|
||||
self,
|
||||
tokens: Any,
|
||||
unprojected: bool = False,
|
||||
add_dict: Optional[dict] = None,
|
||||
show_pbar: bool = True,
|
||||
) -> Any:
|
||||
add_dict = add_dict or {}
|
||||
return self._call_rpc(
|
||||
"encode_from_tokens_scheduled",
|
||||
tokens,
|
||||
unprojected=unprojected,
|
||||
add_dict=add_dict,
|
||||
show_pbar=show_pbar,
|
||||
)
|
||||
|
||||
def add_patches(
|
||||
self, patches: Any, strength_patch: float = 1.0, strength_model: float = 1.0
|
||||
) -> Any:
|
||||
return self._call_rpc(
|
||||
"add_patches",
|
||||
patches,
|
||||
strength_patch=strength_patch,
|
||||
strength_model=strength_model,
|
||||
)
|
||||
|
||||
def get_key_patches(self) -> Any:
|
||||
return self._call_rpc("get_key_patches")
|
||||
|
||||
def load_sd(self, sd: dict, full_model: bool = False) -> Any:
|
||||
return self._call_rpc("load_sd", sd, full_model=full_model)
|
||||
|
||||
def get_sd(self) -> Any:
|
||||
return self._call_rpc("get_sd")
|
||||
|
||||
def clone(self) -> CLIPProxy:
|
||||
new_id = self._call_rpc("clone")
|
||||
return CLIPProxy(new_id, self._registry, manage_lifecycle=not IS_CHILD_PROCESS)
|
||||
|
||||
|
||||
if not IS_CHILD_PROCESS:
|
||||
_CLIP_REGISTRY_SINGLETON = CLIPRegistry()
|
||||
_COND_STAGE_MODEL_REGISTRY_SINGLETON = CondStageModelRegistry()
|
||||
_TOKENIZER_REGISTRY_SINGLETON = TokenizerRegistry()
|
||||
@@ -1,16 +0,0 @@
|
||||
"""Compatibility shim for the indexed serializer path."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
|
||||
def register_custom_node_serializers(_registry: Any) -> None:
|
||||
"""Legacy no-op shim.
|
||||
|
||||
Serializer registration now lives directly in the active isolation adapter.
|
||||
This module remains importable because the isolation index still references it.
|
||||
"""
|
||||
return None
|
||||
|
||||
__all__ = ["register_custom_node_serializers"]
|
||||
@@ -1,540 +0,0 @@
|
||||
# pylint: disable=cyclic-import,import-outside-toplevel,redefined-outer-name
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import os
|
||||
import inspect
|
||||
import sys
|
||||
import types
|
||||
import platform
|
||||
from pathlib import Path
|
||||
from typing import Any, Callable, Dict, List, Tuple
|
||||
|
||||
import pyisolate
|
||||
from pyisolate import ExtensionManager, ExtensionManagerConfig
|
||||
from packaging.requirements import InvalidRequirement, Requirement
|
||||
from packaging.utils import canonicalize_name
|
||||
|
||||
from .manifest_loader import is_cache_valid, load_from_cache, save_to_cache
|
||||
from .host_policy import load_host_policy
|
||||
|
||||
try:
|
||||
import tomllib
|
||||
except ImportError:
|
||||
import tomli as tomllib # type: ignore[no-redef]
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _register_web_directory(extension_name: str, node_dir: Path) -> None:
|
||||
"""Register an isolated extension's web directory on the host side."""
|
||||
import nodes
|
||||
|
||||
# Method 1: pyproject.toml [tool.comfy] web field
|
||||
pyproject = node_dir / "pyproject.toml"
|
||||
if pyproject.exists():
|
||||
try:
|
||||
with pyproject.open("rb") as f:
|
||||
data = tomllib.load(f)
|
||||
web_dir_name = data.get("tool", {}).get("comfy", {}).get("web")
|
||||
if web_dir_name:
|
||||
web_dir_path = str(node_dir / web_dir_name)
|
||||
if os.path.isdir(web_dir_path):
|
||||
nodes.EXTENSION_WEB_DIRS[extension_name] = web_dir_path
|
||||
logger.debug(
|
||||
"][ Registered web dir for isolated %s: %s",
|
||||
extension_name,
|
||||
web_dir_path,
|
||||
)
|
||||
return
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# Method 2: __init__.py WEB_DIRECTORY constant (parse without importing)
|
||||
init_file = node_dir / "__init__.py"
|
||||
if init_file.exists():
|
||||
try:
|
||||
source = init_file.read_text()
|
||||
for line in source.splitlines():
|
||||
stripped = line.strip()
|
||||
if stripped.startswith("WEB_DIRECTORY"):
|
||||
# Parse: WEB_DIRECTORY = "./web" or WEB_DIRECTORY = "web"
|
||||
_, _, value = stripped.partition("=")
|
||||
value = value.strip().strip("\"'")
|
||||
if value:
|
||||
web_dir_path = str((node_dir / value).resolve())
|
||||
if os.path.isdir(web_dir_path):
|
||||
nodes.EXTENSION_WEB_DIRS[extension_name] = web_dir_path
|
||||
logger.debug(
|
||||
"][ Registered web dir for isolated %s: %s",
|
||||
extension_name,
|
||||
web_dir_path,
|
||||
)
|
||||
return
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
def _get_extension_type(execution_model: str) -> type[Any]:
|
||||
if execution_model == "sealed_worker":
|
||||
return pyisolate.SealedNodeExtension
|
||||
|
||||
from .extension_wrapper import ComfyNodeExtension
|
||||
|
||||
return ComfyNodeExtension
|
||||
|
||||
|
||||
async def _stop_extension_safe(extension: Any, extension_name: str) -> None:
|
||||
try:
|
||||
stop_result = extension.stop()
|
||||
if inspect.isawaitable(stop_result):
|
||||
await stop_result
|
||||
except Exception:
|
||||
logger.debug("][ %s stop failed", extension_name, exc_info=True)
|
||||
|
||||
|
||||
def _normalize_dependency_spec(dep: str, base_paths: list[Path]) -> str:
|
||||
req, sep, marker = dep.partition(";")
|
||||
req = req.strip()
|
||||
marker_suffix = f";{marker}" if sep else ""
|
||||
|
||||
def _resolve_local_path(local_path: str) -> Path | None:
|
||||
for base in base_paths:
|
||||
candidate = (base / local_path).resolve()
|
||||
if candidate.exists():
|
||||
return candidate
|
||||
return None
|
||||
|
||||
if req.startswith("./") or req.startswith("../"):
|
||||
resolved = _resolve_local_path(req)
|
||||
if resolved is not None:
|
||||
return f"{resolved}{marker_suffix}"
|
||||
|
||||
if req.startswith("file://"):
|
||||
raw = req[len("file://") :]
|
||||
if raw.startswith("./") or raw.startswith("../"):
|
||||
resolved = _resolve_local_path(raw)
|
||||
if resolved is not None:
|
||||
return f"file://{resolved}{marker_suffix}"
|
||||
|
||||
return dep
|
||||
|
||||
|
||||
def _dependency_name_from_spec(dep: str) -> str | None:
|
||||
stripped = dep.strip()
|
||||
if not stripped or stripped == "-e" or stripped.startswith("-e "):
|
||||
return None
|
||||
if stripped.startswith(("/", "./", "../", "file://")):
|
||||
return None
|
||||
|
||||
try:
|
||||
return canonicalize_name(Requirement(stripped).name)
|
||||
except InvalidRequirement:
|
||||
return None
|
||||
|
||||
|
||||
def _parse_cuda_wheels_config(
|
||||
tool_config: dict[str, object], dependencies: list[str]
|
||||
) -> dict[str, object] | None:
|
||||
raw_config = tool_config.get("cuda_wheels")
|
||||
if raw_config is None:
|
||||
return None
|
||||
if not isinstance(raw_config, dict):
|
||||
raise ExtensionLoadError("[tool.comfy.isolation.cuda_wheels] must be a table")
|
||||
|
||||
index_url = raw_config.get("index_url")
|
||||
index_urls = raw_config.get("index_urls")
|
||||
if index_urls is not None:
|
||||
if not isinstance(index_urls, list) or not all(
|
||||
isinstance(u, str) and u.strip() for u in index_urls
|
||||
):
|
||||
raise ExtensionLoadError(
|
||||
"[tool.comfy.isolation.cuda_wheels.index_urls] must be a list of non-empty strings"
|
||||
)
|
||||
elif not isinstance(index_url, str) or not index_url.strip():
|
||||
raise ExtensionLoadError(
|
||||
"[tool.comfy.isolation.cuda_wheels.index_url] must be a non-empty string"
|
||||
)
|
||||
|
||||
packages = raw_config.get("packages")
|
||||
if not isinstance(packages, list) or not all(
|
||||
isinstance(package_name, str) and package_name.strip()
|
||||
for package_name in packages
|
||||
):
|
||||
raise ExtensionLoadError(
|
||||
"[tool.comfy.isolation.cuda_wheels.packages] must be a list of non-empty strings"
|
||||
)
|
||||
|
||||
declared_dependencies = {
|
||||
dependency_name
|
||||
for dep in dependencies
|
||||
if (dependency_name := _dependency_name_from_spec(dep)) is not None
|
||||
}
|
||||
normalized_packages = [canonicalize_name(package_name) for package_name in packages]
|
||||
missing = [
|
||||
package_name
|
||||
for package_name in normalized_packages
|
||||
if package_name not in declared_dependencies
|
||||
]
|
||||
if missing:
|
||||
missing_joined = ", ".join(sorted(missing))
|
||||
raise ExtensionLoadError(
|
||||
"[tool.comfy.isolation.cuda_wheels.packages] references undeclared dependencies: "
|
||||
f"{missing_joined}"
|
||||
)
|
||||
|
||||
package_map = raw_config.get("package_map", {})
|
||||
if not isinstance(package_map, dict):
|
||||
raise ExtensionLoadError(
|
||||
"[tool.comfy.isolation.cuda_wheels.package_map] must be a table"
|
||||
)
|
||||
|
||||
normalized_package_map: dict[str, str] = {}
|
||||
for dependency_name, index_package_name in package_map.items():
|
||||
if not isinstance(dependency_name, str) or not dependency_name.strip():
|
||||
raise ExtensionLoadError(
|
||||
"[tool.comfy.isolation.cuda_wheels.package_map] keys must be non-empty strings"
|
||||
)
|
||||
if not isinstance(index_package_name, str) or not index_package_name.strip():
|
||||
raise ExtensionLoadError(
|
||||
"[tool.comfy.isolation.cuda_wheels.package_map] values must be non-empty strings"
|
||||
)
|
||||
canonical_dependency_name = canonicalize_name(dependency_name)
|
||||
if canonical_dependency_name not in normalized_packages:
|
||||
raise ExtensionLoadError(
|
||||
"[tool.comfy.isolation.cuda_wheels.package_map] can only override packages listed in "
|
||||
"[tool.comfy.isolation.cuda_wheels.packages]"
|
||||
)
|
||||
normalized_package_map[canonical_dependency_name] = index_package_name.strip()
|
||||
|
||||
result: dict = {
|
||||
"packages": normalized_packages,
|
||||
"package_map": normalized_package_map,
|
||||
}
|
||||
if index_urls is not None:
|
||||
result["index_urls"] = [u.rstrip("/") + "/" for u in index_urls]
|
||||
else:
|
||||
result["index_url"] = index_url.rstrip("/") + "/"
|
||||
return result
|
||||
|
||||
|
||||
def get_enforcement_policy() -> Dict[str, bool]:
|
||||
return {
|
||||
"force_isolated": os.environ.get("PYISOLATE_ENFORCE_ISOLATED") == "1",
|
||||
"force_sandbox": os.environ.get("PYISOLATE_ENFORCE_SANDBOX") == "1",
|
||||
}
|
||||
|
||||
|
||||
class ExtensionLoadError(RuntimeError):
|
||||
pass
|
||||
|
||||
|
||||
def register_dummy_module(extension_name: str, node_dir: Path) -> None:
|
||||
normalized_name = extension_name.replace("-", "_").replace(".", "_")
|
||||
if normalized_name not in sys.modules:
|
||||
dummy_module = types.ModuleType(normalized_name)
|
||||
dummy_module.__file__ = str(node_dir / "__init__.py")
|
||||
dummy_module.__path__ = [str(node_dir)]
|
||||
dummy_module.__package__ = normalized_name
|
||||
sys.modules[normalized_name] = dummy_module
|
||||
|
||||
|
||||
def _is_stale_node_cache(cached_data: Dict[str, Dict]) -> bool:
|
||||
for details in cached_data.values():
|
||||
if not isinstance(details, dict):
|
||||
return True
|
||||
if details.get("is_v3") and "schema_v1" not in details:
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
async def load_isolated_node(
|
||||
node_dir: Path,
|
||||
manifest_path: Path,
|
||||
logger: logging.Logger,
|
||||
build_stub_class: Callable[[str, Dict[str, object], Any], type],
|
||||
venv_root: Path,
|
||||
extension_managers: List[ExtensionManager],
|
||||
) -> List[Tuple[str, str, type]]:
|
||||
try:
|
||||
with manifest_path.open("rb") as handle:
|
||||
manifest_data = tomllib.load(handle)
|
||||
except Exception as e:
|
||||
logger.warning(f"][ Failed to parse {manifest_path}: {e}")
|
||||
return []
|
||||
|
||||
# Parse [tool.comfy.isolation]
|
||||
tool_config = manifest_data.get("tool", {}).get("comfy", {}).get("isolation", {})
|
||||
can_isolate = tool_config.get("can_isolate", False)
|
||||
share_torch = tool_config.get("share_torch", False)
|
||||
package_manager = tool_config.get("package_manager", "uv")
|
||||
is_conda = package_manager == "conda"
|
||||
execution_model = tool_config.get("execution_model")
|
||||
if execution_model is None:
|
||||
execution_model = "sealed_worker" if is_conda else "host-coupled"
|
||||
|
||||
if "sealed_host_ro_paths" in tool_config:
|
||||
raise ValueError(
|
||||
"Manifest field 'sealed_host_ro_paths' is not allowed. "
|
||||
"Configure [tool.comfy.host].sealed_worker_ro_import_paths in host policy."
|
||||
)
|
||||
|
||||
# Conda-specific manifest fields
|
||||
conda_channels: list[str] = (
|
||||
tool_config.get("conda_channels", []) if is_conda else []
|
||||
)
|
||||
conda_dependencies: list[str] = (
|
||||
tool_config.get("conda_dependencies", []) if is_conda else []
|
||||
)
|
||||
conda_platforms: list[str] = (
|
||||
tool_config.get("conda_platforms", []) if is_conda else []
|
||||
)
|
||||
conda_python: str = (
|
||||
tool_config.get("conda_python", "*") if is_conda else "*"
|
||||
)
|
||||
|
||||
# Parse [project] dependencies
|
||||
project_config = manifest_data.get("project", {})
|
||||
dependencies = project_config.get("dependencies", [])
|
||||
if not isinstance(dependencies, list):
|
||||
dependencies = []
|
||||
|
||||
# Get extension name (default to folder name if not in project.name)
|
||||
extension_name = project_config.get("name", node_dir.name)
|
||||
|
||||
# LOGIC: Isolation Decision
|
||||
policy = get_enforcement_policy()
|
||||
isolated = can_isolate or policy["force_isolated"]
|
||||
|
||||
if not isolated:
|
||||
return []
|
||||
|
||||
import folder_paths
|
||||
|
||||
base_paths = [Path(folder_paths.base_path), node_dir]
|
||||
dependencies = [
|
||||
_normalize_dependency_spec(dep, base_paths) if isinstance(dep, str) else dep
|
||||
for dep in dependencies
|
||||
]
|
||||
cuda_wheels = _parse_cuda_wheels_config(tool_config, dependencies)
|
||||
|
||||
manager_config = ExtensionManagerConfig(venv_root_path=str(venv_root))
|
||||
extension_type = _get_extension_type(execution_model)
|
||||
manager: ExtensionManager = pyisolate.ExtensionManager(
|
||||
extension_type, manager_config
|
||||
)
|
||||
extension_managers.append(manager)
|
||||
|
||||
host_policy = load_host_policy(Path(folder_paths.base_path))
|
||||
|
||||
sandbox_config = {}
|
||||
is_linux = platform.system() == "Linux"
|
||||
|
||||
if is_conda:
|
||||
share_torch = False
|
||||
share_cuda_ipc = False
|
||||
else:
|
||||
share_cuda_ipc = share_torch and is_linux
|
||||
|
||||
if is_linux and isolated:
|
||||
sandbox_config = {
|
||||
"network": host_policy["allow_network"],
|
||||
"writable_paths": host_policy["writable_paths"],
|
||||
"readonly_paths": host_policy["readonly_paths"],
|
||||
}
|
||||
|
||||
extension_config: dict = {
|
||||
"name": extension_name,
|
||||
"module_path": str(node_dir),
|
||||
"isolated": True,
|
||||
"dependencies": dependencies,
|
||||
"share_torch": share_torch,
|
||||
"share_cuda_ipc": share_cuda_ipc,
|
||||
"sandbox_mode": host_policy["sandbox_mode"],
|
||||
"sandbox": sandbox_config,
|
||||
}
|
||||
|
||||
share_torch_no_deps = tool_config.get("share_torch_no_deps", [])
|
||||
if share_torch_no_deps:
|
||||
if not isinstance(share_torch_no_deps, list) or not all(
|
||||
isinstance(dep, str) and dep.strip() for dep in share_torch_no_deps
|
||||
):
|
||||
raise ExtensionLoadError(
|
||||
"[tool.comfy.isolation.share_torch_no_deps] must be a list of non-empty strings"
|
||||
)
|
||||
extension_config["share_torch_no_deps"] = share_torch_no_deps
|
||||
|
||||
_is_sealed = execution_model == "sealed_worker"
|
||||
_is_sandboxed = host_policy["sandbox_mode"] != "disabled" and is_linux
|
||||
logger.info(
|
||||
"][ Loading isolated node: %s (torch_share [%s], sealed [%s], sandboxed [%s])",
|
||||
extension_name,
|
||||
"x" if share_torch else " ",
|
||||
"x" if _is_sealed else " ",
|
||||
"x" if _is_sandboxed else " ",
|
||||
)
|
||||
|
||||
if cuda_wheels is not None:
|
||||
extension_config["cuda_wheels"] = cuda_wheels
|
||||
|
||||
extra_index_urls = tool_config.get("extra_index_urls", [])
|
||||
if extra_index_urls:
|
||||
if not isinstance(extra_index_urls, list) or not all(
|
||||
isinstance(u, str) and u.strip() for u in extra_index_urls
|
||||
):
|
||||
raise ExtensionLoadError(
|
||||
"[tool.comfy.isolation.extra_index_urls] must be a list of non-empty strings"
|
||||
)
|
||||
extension_config["extra_index_urls"] = extra_index_urls
|
||||
|
||||
# Conda-specific keys
|
||||
if is_conda:
|
||||
extension_config["package_manager"] = "conda"
|
||||
extension_config["conda_channels"] = conda_channels
|
||||
extension_config["conda_dependencies"] = conda_dependencies
|
||||
extension_config["conda_python"] = conda_python
|
||||
find_links = tool_config.get("find_links", [])
|
||||
if find_links:
|
||||
extension_config["find_links"] = find_links
|
||||
if conda_platforms:
|
||||
extension_config["conda_platforms"] = conda_platforms
|
||||
|
||||
if execution_model != "host-coupled":
|
||||
extension_config["execution_model"] = execution_model
|
||||
if execution_model == "sealed_worker":
|
||||
policy_ro_paths = host_policy.get("sealed_worker_ro_import_paths", [])
|
||||
if isinstance(policy_ro_paths, list) and policy_ro_paths:
|
||||
extension_config["sealed_host_ro_paths"] = list(policy_ro_paths)
|
||||
# Sealed workers keep the host RPC service inventory even when the
|
||||
# child resolves no API classes locally.
|
||||
|
||||
extension = manager.load_extension(extension_config)
|
||||
register_dummy_module(extension_name, node_dir)
|
||||
|
||||
# Register host-side event handlers via adapter
|
||||
from .adapter import ComfyUIAdapter
|
||||
ComfyUIAdapter.register_host_event_handlers(extension)
|
||||
|
||||
# Register web directory on the host — only when sandbox is disabled.
|
||||
# In sandbox mode, serving untrusted JS to the browser is not safe.
|
||||
if host_policy["sandbox_mode"] == "disabled":
|
||||
_register_web_directory(extension_name, node_dir)
|
||||
|
||||
# Register for proxied web serving — the child's web dir may have
|
||||
# content that doesn't exist on the host (e.g., pip-installed viewer
|
||||
# bundles). The WebDirectoryCache will lazily fetch via RPC.
|
||||
from .proxies.web_directory_proxy import WebDirectoryProxy, get_web_directory_cache
|
||||
|
||||
class ChildWebDirectoryProxy:
|
||||
def __init__(self, host_extension):
|
||||
self._host_extension = host_extension
|
||||
self._caller = None
|
||||
|
||||
def _get_caller(self):
|
||||
self._host_extension.proxy
|
||||
rpc = self._host_extension._extension.rpc
|
||||
caller = rpc.create_caller(WebDirectoryProxy, WebDirectoryProxy.get_remote_id())
|
||||
if self._caller is not caller:
|
||||
self._caller = caller
|
||||
return self._caller
|
||||
|
||||
def list_web_files(self, extension_name: str):
|
||||
from .proxies.base import run_sync_rpc_coro
|
||||
return run_sync_rpc_coro(self._get_caller().list_web_files(extension_name))
|
||||
|
||||
def get_web_file(self, extension_name: str, relative_path: str):
|
||||
from .proxies.base import run_sync_rpc_coro
|
||||
return run_sync_rpc_coro(
|
||||
self._get_caller().get_web_file(extension_name, relative_path)
|
||||
)
|
||||
|
||||
cache = get_web_directory_cache()
|
||||
cache.register_proxy(extension_name, ChildWebDirectoryProxy(extension))
|
||||
|
||||
# Try cache first (lazy spawn)
|
||||
if is_cache_valid(node_dir, manifest_path, venv_root):
|
||||
cached_data = load_from_cache(node_dir, venv_root)
|
||||
if cached_data:
|
||||
if _is_stale_node_cache(cached_data):
|
||||
pass
|
||||
else:
|
||||
try:
|
||||
flushed = await extension.flush_pending_routes()
|
||||
logger.info("][ %s flushed %d routes", extension_name, flushed)
|
||||
except Exception as exc:
|
||||
logger.warning("][ %s route flush failed: %s", extension_name, exc)
|
||||
specs: List[Tuple[str, str, type]] = []
|
||||
for node_name, details in cached_data.items():
|
||||
stub_cls = build_stub_class(node_name, details, extension)
|
||||
specs.append(
|
||||
(node_name, details.get("display_name", node_name), stub_cls)
|
||||
)
|
||||
return specs
|
||||
# Cache miss - spawn process and get metadata
|
||||
|
||||
try:
|
||||
remote_nodes: Dict[str, str] = await extension.list_nodes()
|
||||
except Exception as exc:
|
||||
logger.warning(
|
||||
"][ %s metadata discovery failed, skipping isolated load: %s",
|
||||
extension_name,
|
||||
exc,
|
||||
)
|
||||
await _stop_extension_safe(extension, extension_name)
|
||||
return []
|
||||
|
||||
if not remote_nodes:
|
||||
logger.debug("][ %s exposed no isolated nodes; skipping", extension_name)
|
||||
await _stop_extension_safe(extension, extension_name)
|
||||
return []
|
||||
|
||||
specs: List[Tuple[str, str, type]] = []
|
||||
cache_data: Dict[str, Dict] = {}
|
||||
|
||||
for node_name, display_name in remote_nodes.items():
|
||||
try:
|
||||
details = await extension.get_node_details(node_name)
|
||||
except Exception as exc:
|
||||
logger.warning(
|
||||
"][ %s failed to load metadata for %s, skipping node: %s",
|
||||
extension_name,
|
||||
node_name,
|
||||
exc,
|
||||
)
|
||||
continue
|
||||
details["display_name"] = display_name
|
||||
cache_data[node_name] = details
|
||||
stub_cls = build_stub_class(node_name, details, extension)
|
||||
specs.append((node_name, display_name, stub_cls))
|
||||
|
||||
if not specs:
|
||||
logger.warning(
|
||||
"][ %s produced no usable nodes after metadata scan; skipping",
|
||||
extension_name,
|
||||
)
|
||||
await _stop_extension_safe(extension, extension_name)
|
||||
return []
|
||||
|
||||
# Save metadata to cache for future runs
|
||||
save_to_cache(node_dir, venv_root, cache_data, manifest_path)
|
||||
logger.debug(f"][ {extension_name} metadata cached")
|
||||
|
||||
# Re-check web directory AFTER child has populated it
|
||||
if host_policy["sandbox_mode"] == "disabled":
|
||||
_register_web_directory(extension_name, node_dir)
|
||||
|
||||
# Flush any routes the child buffered during module import — must happen
|
||||
# before router freeze and before we kill the child process.
|
||||
try:
|
||||
flushed = await extension.flush_pending_routes()
|
||||
logger.info("][ %s flushed %d routes", extension_name, flushed)
|
||||
except Exception as exc:
|
||||
logger.warning("][ %s route flush failed: %s", extension_name, exc)
|
||||
|
||||
# EJECT: Kill process after getting metadata (will respawn on first execution)
|
||||
await _stop_extension_safe(extension, extension_name)
|
||||
|
||||
return specs
|
||||
|
||||
|
||||
__all__ = ["ExtensionLoadError", "register_dummy_module", "load_isolated_node"]
|
||||
@@ -1,942 +0,0 @@
|
||||
# pylint: disable=consider-using-from-import,cyclic-import,import-outside-toplevel,logging-fstring-interpolation,protected-access,wrong-import-position
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import torch
|
||||
|
||||
|
||||
class AttrDict(dict):
|
||||
def __getattr__(self, item):
|
||||
try:
|
||||
return self[item]
|
||||
except KeyError as e:
|
||||
raise AttributeError(item) from e
|
||||
|
||||
def copy(self):
|
||||
return AttrDict(super().copy())
|
||||
|
||||
|
||||
import importlib
|
||||
import inspect
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import sys
|
||||
import uuid
|
||||
from dataclasses import asdict
|
||||
from typing import Any, Dict, List, Tuple
|
||||
|
||||
from pyisolate import ExtensionBase
|
||||
|
||||
from comfy_api.internal import _ComfyNodeInternal
|
||||
|
||||
LOG_PREFIX = "]["
|
||||
V3_DISCOVERY_TIMEOUT = 30
|
||||
_PRE_EXEC_MIN_FREE_VRAM_BYTES = 2 * 1024 * 1024 * 1024
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _run_prestartup_web_copy(module: Any, module_dir: str, web_dir_path: str) -> None:
|
||||
"""Run the web asset copy step that prestartup_script.py used to do.
|
||||
|
||||
If the module's web/ directory is empty and the module had a
|
||||
prestartup_script.py that copied assets from pip packages, this
|
||||
function replicates that work inside the child process.
|
||||
|
||||
Generic pattern: reads _PRESTARTUP_WEB_COPY from the module if
|
||||
defined, otherwise falls back to detecting common asset packages.
|
||||
"""
|
||||
import shutil
|
||||
|
||||
# Already populated — nothing to do
|
||||
if os.path.isdir(web_dir_path) and any(os.scandir(web_dir_path)):
|
||||
return
|
||||
|
||||
os.makedirs(web_dir_path, exist_ok=True)
|
||||
|
||||
# Try module-defined copy spec first (generic hook for any node pack)
|
||||
copy_spec = getattr(module, "_PRESTARTUP_WEB_COPY", None)
|
||||
if copy_spec is not None and callable(copy_spec):
|
||||
try:
|
||||
copy_spec(web_dir_path)
|
||||
logger.info(
|
||||
"%s Ran _PRESTARTUP_WEB_COPY for %s", LOG_PREFIX, module_dir
|
||||
)
|
||||
return
|
||||
except Exception as e:
|
||||
logger.warning(
|
||||
"%s _PRESTARTUP_WEB_COPY failed for %s: %s",
|
||||
LOG_PREFIX, module_dir, e,
|
||||
)
|
||||
|
||||
# Fallback: detect comfy_3d_viewers and run copy_viewer()
|
||||
try:
|
||||
from comfy_3d_viewers import copy_viewer, VIEWER_FILES
|
||||
viewers = list(VIEWER_FILES.keys())
|
||||
for viewer in viewers:
|
||||
try:
|
||||
copy_viewer(viewer, web_dir_path)
|
||||
except Exception:
|
||||
pass
|
||||
if any(os.scandir(web_dir_path)):
|
||||
logger.info(
|
||||
"%s Copied %d viewer types from comfy_3d_viewers to %s",
|
||||
LOG_PREFIX, len(viewers), web_dir_path,
|
||||
)
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
# Fallback: detect comfy_dynamic_widgets
|
||||
try:
|
||||
from comfy_dynamic_widgets import get_js_path
|
||||
src = os.path.realpath(get_js_path())
|
||||
if os.path.exists(src):
|
||||
dst_dir = os.path.join(web_dir_path, "js")
|
||||
os.makedirs(dst_dir, exist_ok=True)
|
||||
dst = os.path.join(dst_dir, "dynamic_widgets.js")
|
||||
shutil.copy2(src, dst)
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
|
||||
def _read_extension_name(module_dir: str) -> str:
|
||||
"""Read extension name from pyproject.toml, falling back to directory name."""
|
||||
pyproject = os.path.join(module_dir, "pyproject.toml")
|
||||
if os.path.exists(pyproject):
|
||||
try:
|
||||
import tomllib
|
||||
except ImportError:
|
||||
import tomli as tomllib # type: ignore[no-redef]
|
||||
try:
|
||||
with open(pyproject, "rb") as f:
|
||||
data = tomllib.load(f)
|
||||
name = data.get("project", {}).get("name")
|
||||
if name:
|
||||
return name
|
||||
except Exception:
|
||||
pass
|
||||
return os.path.basename(module_dir)
|
||||
|
||||
|
||||
def _flush_tensor_transport_state(marker: str) -> int:
|
||||
try:
|
||||
from pyisolate import flush_tensor_keeper # type: ignore[attr-defined]
|
||||
except Exception:
|
||||
return 0
|
||||
if not callable(flush_tensor_keeper):
|
||||
return 0
|
||||
flushed = flush_tensor_keeper()
|
||||
if flushed > 0:
|
||||
logger.debug(
|
||||
"%s %s flush_tensor_keeper released=%d", LOG_PREFIX, marker, flushed
|
||||
)
|
||||
return flushed
|
||||
|
||||
|
||||
def _relieve_child_vram_pressure(marker: str) -> None:
|
||||
import comfy.model_management as model_management
|
||||
|
||||
model_management.cleanup_models_gc()
|
||||
model_management.cleanup_models()
|
||||
|
||||
device = model_management.get_torch_device()
|
||||
if not hasattr(device, "type") or device.type == "cpu":
|
||||
return
|
||||
|
||||
required = max(
|
||||
model_management.minimum_inference_memory(),
|
||||
_PRE_EXEC_MIN_FREE_VRAM_BYTES,
|
||||
)
|
||||
if model_management.get_free_memory(device) < required:
|
||||
model_management.free_memory(required, device, for_dynamic=True)
|
||||
if model_management.get_free_memory(device) < required:
|
||||
model_management.free_memory(required, device, for_dynamic=False)
|
||||
model_management.cleanup_models()
|
||||
model_management.soft_empty_cache()
|
||||
logger.debug("%s %s free_memory target=%d", LOG_PREFIX, marker, required)
|
||||
|
||||
|
||||
def _sanitize_for_transport(value):
|
||||
primitives = (str, int, float, bool, type(None))
|
||||
if isinstance(value, primitives):
|
||||
return value
|
||||
|
||||
cls_name = value.__class__.__name__
|
||||
if cls_name == "FlexibleOptionalInputType":
|
||||
return {
|
||||
"__pyisolate_flexible_optional__": True,
|
||||
"type": _sanitize_for_transport(getattr(value, "type", "*")),
|
||||
}
|
||||
if cls_name == "AnyType":
|
||||
return {"__pyisolate_any_type__": True, "value": str(value)}
|
||||
if cls_name == "ByPassTypeTuple":
|
||||
return {
|
||||
"__pyisolate_bypass_tuple__": [
|
||||
_sanitize_for_transport(v) for v in tuple(value)
|
||||
]
|
||||
}
|
||||
|
||||
if isinstance(value, dict):
|
||||
return {k: _sanitize_for_transport(v) for k, v in value.items()}
|
||||
if isinstance(value, tuple):
|
||||
return {"__pyisolate_tuple__": [_sanitize_for_transport(v) for v in value]}
|
||||
if isinstance(value, list):
|
||||
return [_sanitize_for_transport(v) for v in value]
|
||||
|
||||
return str(value)
|
||||
|
||||
|
||||
# Re-export RemoteObjectHandle from pyisolate for backward compatibility
|
||||
# The canonical definition is now in pyisolate._internal.remote_handle
|
||||
from pyisolate._internal.remote_handle import RemoteObjectHandle # noqa: E402,F401
|
||||
|
||||
|
||||
class ComfyNodeExtension(ExtensionBase):
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
self.node_classes: Dict[str, type] = {}
|
||||
self.display_names: Dict[str, str] = {}
|
||||
self.node_instances: Dict[str, Any] = {}
|
||||
self.remote_objects: Dict[str, Any] = {}
|
||||
self._route_handlers: Dict[str, Any] = {}
|
||||
self._module: Any = None
|
||||
self._metadata_ready = asyncio.Event()
|
||||
|
||||
async def on_module_loaded(self, module: Any) -> None:
|
||||
try:
|
||||
self._module = module
|
||||
|
||||
# Registries are initialized in host_hooks.py initialize_host_process()
|
||||
# They auto-register via ProxiedSingleton when instantiated
|
||||
# NO additional setup required here - if a registry is missing from host_hooks, it WILL fail
|
||||
|
||||
self.node_classes = getattr(module, "NODE_CLASS_MAPPINGS", {}) or {}
|
||||
self.display_names = getattr(module, "NODE_DISPLAY_NAME_MAPPINGS", {}) or {}
|
||||
self._register_module_routes(module)
|
||||
|
||||
# Register web directory with WebDirectoryProxy (child-side)
|
||||
web_dir_attr = getattr(module, "WEB_DIRECTORY", None)
|
||||
if web_dir_attr is not None:
|
||||
module_dir = os.path.dirname(os.path.abspath(module.__file__))
|
||||
web_dir_path = os.path.abspath(os.path.join(module_dir, web_dir_attr))
|
||||
ext_name = _read_extension_name(module_dir)
|
||||
|
||||
# If web dir is empty, run the copy step that prestartup_script.py did
|
||||
_run_prestartup_web_copy(module, module_dir, web_dir_path)
|
||||
|
||||
if os.path.isdir(web_dir_path) and any(os.scandir(web_dir_path)):
|
||||
from comfy.isolation.proxies.web_directory_proxy import WebDirectoryProxy
|
||||
WebDirectoryProxy.register_web_dir(ext_name, web_dir_path)
|
||||
|
||||
try:
|
||||
from comfy_api.latest import ComfyExtension
|
||||
|
||||
for name, obj in inspect.getmembers(module):
|
||||
if not (
|
||||
inspect.isclass(obj)
|
||||
and issubclass(obj, ComfyExtension)
|
||||
and obj is not ComfyExtension
|
||||
):
|
||||
continue
|
||||
if not obj.__module__.startswith(module.__name__):
|
||||
continue
|
||||
try:
|
||||
ext_instance = obj()
|
||||
try:
|
||||
await asyncio.wait_for(
|
||||
ext_instance.on_load(), timeout=V3_DISCOVERY_TIMEOUT
|
||||
)
|
||||
except asyncio.TimeoutError:
|
||||
logger.error(
|
||||
"%s V3 Extension %s timed out in on_load()",
|
||||
LOG_PREFIX,
|
||||
name,
|
||||
)
|
||||
continue
|
||||
try:
|
||||
v3_nodes = await asyncio.wait_for(
|
||||
ext_instance.get_node_list(), timeout=V3_DISCOVERY_TIMEOUT
|
||||
)
|
||||
except asyncio.TimeoutError:
|
||||
logger.error(
|
||||
"%s V3 Extension %s timed out in get_node_list()",
|
||||
LOG_PREFIX,
|
||||
name,
|
||||
)
|
||||
continue
|
||||
for node_cls in v3_nodes:
|
||||
if hasattr(node_cls, "GET_SCHEMA"):
|
||||
schema = node_cls.GET_SCHEMA()
|
||||
self.node_classes[schema.node_id] = node_cls
|
||||
if schema.display_name:
|
||||
self.display_names[schema.node_id] = schema.display_name
|
||||
except Exception as e:
|
||||
logger.error("%s V3 Extension %s failed: %s", LOG_PREFIX, name, e)
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
module_name = getattr(module, "__name__", "isolated_nodes")
|
||||
for node_cls in self.node_classes.values():
|
||||
if hasattr(node_cls, "__module__") and "/" in str(node_cls.__module__):
|
||||
node_cls.__module__ = module_name
|
||||
|
||||
self.node_instances = {}
|
||||
finally:
|
||||
self._metadata_ready.set()
|
||||
|
||||
def _register_module_routes(self, module: Any) -> None:
|
||||
"""Bridge legacy module-level ROUTES declarations into isolated routing."""
|
||||
routes = getattr(module, "ROUTES", None) or []
|
||||
if not routes:
|
||||
return
|
||||
|
||||
from comfy.isolation.proxies.prompt_server_impl import PromptServerStub
|
||||
|
||||
prompt_server = PromptServerStub()
|
||||
route_table = getattr(prompt_server, "routes", None)
|
||||
if route_table is None:
|
||||
logger.warning("%s Route registration unavailable for %s", LOG_PREFIX, module)
|
||||
return
|
||||
|
||||
for route_spec in routes:
|
||||
if not isinstance(route_spec, dict):
|
||||
logger.warning("%s Ignoring non-dict ROUTES entry: %r", LOG_PREFIX, route_spec)
|
||||
continue
|
||||
|
||||
method = str(route_spec.get("method", "")).strip().upper()
|
||||
path = str(route_spec.get("path", "")).strip()
|
||||
handler_ref = route_spec.get("handler")
|
||||
if not method or not path:
|
||||
logger.warning("%s Ignoring incomplete route spec: %r", LOG_PREFIX, route_spec)
|
||||
continue
|
||||
|
||||
if isinstance(handler_ref, str):
|
||||
handler = getattr(module, handler_ref, None)
|
||||
else:
|
||||
handler = handler_ref
|
||||
if not callable(handler):
|
||||
logger.warning(
|
||||
"%s Ignoring route with missing handler %r for %s %s",
|
||||
LOG_PREFIX,
|
||||
handler_ref,
|
||||
method,
|
||||
path,
|
||||
)
|
||||
continue
|
||||
|
||||
decorator = getattr(route_table, method.lower(), None)
|
||||
if not callable(decorator):
|
||||
logger.warning("%s Unsupported route method %s for %s", LOG_PREFIX, method, path)
|
||||
continue
|
||||
|
||||
decorator(path)(handler)
|
||||
self._route_handlers[f"{method} {path}"] = handler
|
||||
logger.info("%s buffered legacy route %s %s", LOG_PREFIX, method, path)
|
||||
|
||||
async def list_nodes(self) -> Dict[str, str]:
|
||||
await asyncio.wait_for(
|
||||
self._metadata_ready.wait(), timeout=V3_DISCOVERY_TIMEOUT
|
||||
)
|
||||
return {name: self.display_names.get(name, name) for name in self.node_classes}
|
||||
|
||||
async def get_node_info(self, node_name: str) -> Dict[str, Any]:
|
||||
return await self.get_node_details(node_name)
|
||||
|
||||
async def get_node_details(self, node_name: str) -> Dict[str, Any]:
|
||||
await asyncio.wait_for(
|
||||
self._metadata_ready.wait(), timeout=V3_DISCOVERY_TIMEOUT
|
||||
)
|
||||
node_cls = self._get_node_class(node_name)
|
||||
is_v3 = issubclass(node_cls, _ComfyNodeInternal)
|
||||
|
||||
input_types_raw = (
|
||||
node_cls.INPUT_TYPES() if hasattr(node_cls, "INPUT_TYPES") else {}
|
||||
)
|
||||
output_is_list = getattr(node_cls, "OUTPUT_IS_LIST", None)
|
||||
if output_is_list is not None:
|
||||
output_is_list = tuple(bool(x) for x in output_is_list)
|
||||
|
||||
details: Dict[str, Any] = {
|
||||
"input_types": _sanitize_for_transport(input_types_raw),
|
||||
"return_types": tuple(
|
||||
str(t) for t in getattr(node_cls, "RETURN_TYPES", ())
|
||||
),
|
||||
"return_names": getattr(node_cls, "RETURN_NAMES", None),
|
||||
"function": str(getattr(node_cls, "FUNCTION", "execute")),
|
||||
"category": str(getattr(node_cls, "CATEGORY", "")),
|
||||
"output_node": bool(getattr(node_cls, "OUTPUT_NODE", False)),
|
||||
"output_is_list": output_is_list,
|
||||
"is_v3": is_v3,
|
||||
}
|
||||
|
||||
if is_v3:
|
||||
try:
|
||||
schema = node_cls.GET_SCHEMA()
|
||||
schema_v1 = asdict(schema.get_v1_info(node_cls))
|
||||
try:
|
||||
schema_v3 = asdict(schema.get_v3_info(node_cls))
|
||||
except (AttributeError, TypeError):
|
||||
schema_v3 = self._build_schema_v3_fallback(schema)
|
||||
details.update(
|
||||
{
|
||||
"schema_v1": schema_v1,
|
||||
"schema_v3": schema_v3,
|
||||
"hidden": [h.value for h in (schema.hidden or [])],
|
||||
"description": getattr(schema, "description", ""),
|
||||
"deprecated": bool(getattr(node_cls, "DEPRECATED", False)),
|
||||
"experimental": bool(getattr(node_cls, "EXPERIMENTAL", False)),
|
||||
"api_node": bool(getattr(node_cls, "API_NODE", False)),
|
||||
"input_is_list": bool(
|
||||
getattr(node_cls, "INPUT_IS_LIST", False)
|
||||
),
|
||||
"not_idempotent": bool(
|
||||
getattr(node_cls, "NOT_IDEMPOTENT", False)
|
||||
),
|
||||
"accept_all_inputs": bool(
|
||||
getattr(node_cls, "ACCEPT_ALL_INPUTS", False)
|
||||
),
|
||||
}
|
||||
)
|
||||
except Exception as exc:
|
||||
logger.warning(
|
||||
"%s V3 schema serialization failed for %s: %s",
|
||||
LOG_PREFIX,
|
||||
node_name,
|
||||
exc,
|
||||
)
|
||||
return details
|
||||
|
||||
def _build_schema_v3_fallback(self, schema) -> Dict[str, Any]:
|
||||
input_dict: Dict[str, Any] = {}
|
||||
output_dict: Dict[str, Any] = {}
|
||||
hidden_list: List[str] = []
|
||||
|
||||
if getattr(schema, "inputs", None):
|
||||
for inp in schema.inputs:
|
||||
self._add_schema_io_v3(inp, input_dict)
|
||||
if getattr(schema, "outputs", None):
|
||||
for out in schema.outputs:
|
||||
self._add_schema_io_v3(out, output_dict)
|
||||
if getattr(schema, "hidden", None):
|
||||
for h in schema.hidden:
|
||||
hidden_list.append(getattr(h, "value", str(h)))
|
||||
|
||||
return {
|
||||
"input": input_dict,
|
||||
"output": output_dict,
|
||||
"hidden": hidden_list,
|
||||
"name": getattr(schema, "node_id", None),
|
||||
"display_name": getattr(schema, "display_name", None),
|
||||
"description": getattr(schema, "description", None),
|
||||
"category": getattr(schema, "category", None),
|
||||
"output_node": getattr(schema, "is_output_node", False),
|
||||
"deprecated": getattr(schema, "is_deprecated", False),
|
||||
"experimental": getattr(schema, "is_experimental", False),
|
||||
"api_node": getattr(schema, "is_api_node", False),
|
||||
}
|
||||
|
||||
def _add_schema_io_v3(self, io_obj: Any, target: Dict[str, Any]) -> None:
|
||||
io_id = getattr(io_obj, "id", None)
|
||||
if io_id is None:
|
||||
return
|
||||
|
||||
io_type_fn = getattr(io_obj, "get_io_type", None)
|
||||
io_type = (
|
||||
io_type_fn() if callable(io_type_fn) else getattr(io_obj, "io_type", None)
|
||||
)
|
||||
|
||||
as_dict_fn = getattr(io_obj, "as_dict", None)
|
||||
payload = as_dict_fn() if callable(as_dict_fn) else {}
|
||||
|
||||
target[str(io_id)] = (io_type, payload)
|
||||
|
||||
async def get_input_types(self, node_name: str) -> Dict[str, Any]:
|
||||
node_cls = self._get_node_class(node_name)
|
||||
if hasattr(node_cls, "INPUT_TYPES"):
|
||||
return node_cls.INPUT_TYPES()
|
||||
return {}
|
||||
|
||||
async def execute_node(self, node_name: str, **inputs: Any) -> Tuple[Any, ...]:
|
||||
logger.debug(
|
||||
"%s ISO:child_execute_start ext=%s node=%s input_keys=%d",
|
||||
LOG_PREFIX,
|
||||
getattr(self, "name", "?"),
|
||||
node_name,
|
||||
len(inputs),
|
||||
)
|
||||
if os.environ.get("PYISOLATE_CHILD") == "1":
|
||||
_relieve_child_vram_pressure("EXT:pre_execute")
|
||||
|
||||
resolved_inputs = self._resolve_remote_objects(inputs)
|
||||
|
||||
instance = self._get_node_instance(node_name)
|
||||
node_cls = self._get_node_class(node_name)
|
||||
|
||||
# V3 API nodes expect hidden parameters in cls.hidden, not as kwargs
|
||||
# Hidden params come through RPC as string keys like "Hidden.prompt"
|
||||
from comfy_api.latest._io import Hidden, HiddenHolder
|
||||
|
||||
# Map string representations back to Hidden enum keys
|
||||
hidden_string_map = {
|
||||
"Hidden.unique_id": Hidden.unique_id,
|
||||
"Hidden.prompt": Hidden.prompt,
|
||||
"Hidden.extra_pnginfo": Hidden.extra_pnginfo,
|
||||
"Hidden.dynprompt": Hidden.dynprompt,
|
||||
"Hidden.auth_token_comfy_org": Hidden.auth_token_comfy_org,
|
||||
"Hidden.api_key_comfy_org": Hidden.api_key_comfy_org,
|
||||
# Uppercase enum VALUE forms — V3 execution engine passes these
|
||||
"UNIQUE_ID": Hidden.unique_id,
|
||||
"PROMPT": Hidden.prompt,
|
||||
"EXTRA_PNGINFO": Hidden.extra_pnginfo,
|
||||
"DYNPROMPT": Hidden.dynprompt,
|
||||
"AUTH_TOKEN_COMFY_ORG": Hidden.auth_token_comfy_org,
|
||||
"API_KEY_COMFY_ORG": Hidden.api_key_comfy_org,
|
||||
}
|
||||
|
||||
# Find and extract hidden parameters (both enum and string form)
|
||||
hidden_found = {}
|
||||
keys_to_remove = []
|
||||
|
||||
for key in list(resolved_inputs.keys()):
|
||||
# Check string form first (from RPC serialization)
|
||||
if key in hidden_string_map:
|
||||
hidden_found[hidden_string_map[key]] = resolved_inputs[key]
|
||||
keys_to_remove.append(key)
|
||||
# Also check enum form (direct calls)
|
||||
elif isinstance(key, Hidden):
|
||||
hidden_found[key] = resolved_inputs[key]
|
||||
keys_to_remove.append(key)
|
||||
|
||||
# Remove hidden params from kwargs
|
||||
for key in keys_to_remove:
|
||||
resolved_inputs.pop(key)
|
||||
|
||||
# Set hidden on node class if any hidden params found
|
||||
if hidden_found:
|
||||
if not hasattr(node_cls, "hidden") or node_cls.hidden is None:
|
||||
node_cls.hidden = HiddenHolder.from_dict(hidden_found)
|
||||
else:
|
||||
# Update existing hidden holder
|
||||
for key, value in hidden_found.items():
|
||||
setattr(node_cls.hidden, key.value.lower(), value)
|
||||
|
||||
# INPUT_IS_LIST: ComfyUI's executor passes all inputs as lists when this
|
||||
# flag is set. The isolation RPC delivers unwrapped values, so we must
|
||||
# wrap each input in a single-element list to match the contract.
|
||||
if getattr(node_cls, "INPUT_IS_LIST", False):
|
||||
resolved_inputs = {k: [v] for k, v in resolved_inputs.items()}
|
||||
|
||||
function_name = getattr(node_cls, "FUNCTION", "execute")
|
||||
if not hasattr(instance, function_name):
|
||||
raise AttributeError(f"Node {node_name} missing callable '{function_name}'")
|
||||
|
||||
handler = getattr(instance, function_name)
|
||||
|
||||
try:
|
||||
import torch
|
||||
if asyncio.iscoroutinefunction(handler):
|
||||
with torch.inference_mode():
|
||||
result = await handler(**resolved_inputs)
|
||||
else:
|
||||
import functools
|
||||
|
||||
def _run_with_inference_mode(**kwargs):
|
||||
with torch.inference_mode():
|
||||
return handler(**kwargs)
|
||||
|
||||
loop = asyncio.get_running_loop()
|
||||
result = await loop.run_in_executor(
|
||||
None, functools.partial(_run_with_inference_mode, **resolved_inputs)
|
||||
)
|
||||
except Exception:
|
||||
logger.exception(
|
||||
"%s ISO:child_execute_error ext=%s node=%s",
|
||||
LOG_PREFIX,
|
||||
getattr(self, "name", "?"),
|
||||
node_name,
|
||||
)
|
||||
raise
|
||||
|
||||
if type(result).__name__ == "NodeOutput":
|
||||
node_output_dict = {
|
||||
"__node_output__": True,
|
||||
"args": self._wrap_unpicklable_objects(result.args),
|
||||
}
|
||||
if result.ui is not None:
|
||||
node_output_dict["ui"] = self._wrap_unpicklable_objects(result.ui)
|
||||
if getattr(result, "expand", None) is not None:
|
||||
node_output_dict["expand"] = result.expand
|
||||
if getattr(result, "block_execution", None) is not None:
|
||||
node_output_dict["block_execution"] = result.block_execution
|
||||
return node_output_dict
|
||||
if self._is_comfy_protocol_return(result):
|
||||
wrapped = self._wrap_unpicklable_objects(result)
|
||||
return wrapped
|
||||
|
||||
if not isinstance(result, tuple):
|
||||
result = (result,)
|
||||
wrapped = self._wrap_unpicklable_objects(result)
|
||||
return wrapped
|
||||
|
||||
async def flush_pending_routes(self) -> int:
|
||||
"""Flush buffered route registrations to host via RPC. Called by host after node discovery."""
|
||||
from comfy.isolation.proxies.prompt_server_impl import PromptServerStub
|
||||
return await PromptServerStub.flush_child_routes()
|
||||
|
||||
async def flush_transport_state(self) -> int:
|
||||
if os.environ.get("PYISOLATE_CHILD") != "1":
|
||||
return 0
|
||||
logger.debug(
|
||||
"%s ISO:child_flush_start ext=%s", LOG_PREFIX, getattr(self, "name", "?")
|
||||
)
|
||||
flushed = _flush_tensor_transport_state("EXT:workflow_end")
|
||||
try:
|
||||
from comfy.isolation.model_patcher_proxy_registry import (
|
||||
ModelPatcherRegistry,
|
||||
)
|
||||
|
||||
registry = ModelPatcherRegistry()
|
||||
removed = registry.sweep_pending_cleanup()
|
||||
if removed > 0:
|
||||
logger.debug(
|
||||
"%s EXT:workflow_end registry sweep removed=%d", LOG_PREFIX, removed
|
||||
)
|
||||
except Exception:
|
||||
logger.debug(
|
||||
"%s EXT:workflow_end registry sweep failed", LOG_PREFIX, exc_info=True
|
||||
)
|
||||
logger.debug(
|
||||
"%s ISO:child_flush_done ext=%s flushed=%d",
|
||||
LOG_PREFIX,
|
||||
getattr(self, "name", "?"),
|
||||
flushed,
|
||||
)
|
||||
return flushed
|
||||
|
||||
async def get_remote_object(self, object_id: str) -> Any:
|
||||
"""Retrieve a remote object by ID for host-side deserialization."""
|
||||
if object_id not in self.remote_objects:
|
||||
raise KeyError(f"Remote object {object_id} not found")
|
||||
|
||||
return self.remote_objects[object_id]
|
||||
|
||||
def _store_remote_object_handle(self, obj: Any) -> RemoteObjectHandle:
|
||||
object_id = str(uuid.uuid4())
|
||||
self.remote_objects[object_id] = obj
|
||||
return RemoteObjectHandle(object_id, type(obj).__name__)
|
||||
|
||||
async def call_remote_object_method(
|
||||
self,
|
||||
object_id: str,
|
||||
method_name: str,
|
||||
*args: Any,
|
||||
**kwargs: Any,
|
||||
) -> Any:
|
||||
"""Invoke a method or attribute-backed accessor on a child-owned object."""
|
||||
obj = await self.get_remote_object(object_id)
|
||||
|
||||
if method_name == "get_patcher_attr":
|
||||
return getattr(obj, args[0])
|
||||
if method_name == "get_model_options":
|
||||
return getattr(obj, "model_options")
|
||||
if method_name == "set_model_options":
|
||||
setattr(obj, "model_options", args[0])
|
||||
return None
|
||||
if method_name == "get_object_patches":
|
||||
return getattr(obj, "object_patches")
|
||||
if method_name == "get_patches":
|
||||
return getattr(obj, "patches")
|
||||
if method_name == "get_wrappers":
|
||||
return getattr(obj, "wrappers")
|
||||
if method_name == "get_callbacks":
|
||||
return getattr(obj, "callbacks")
|
||||
if method_name == "get_load_device":
|
||||
return getattr(obj, "load_device")
|
||||
if method_name == "get_offload_device":
|
||||
return getattr(obj, "offload_device")
|
||||
if method_name == "get_hook_mode":
|
||||
return getattr(obj, "hook_mode")
|
||||
if method_name == "get_parent":
|
||||
parent = getattr(obj, "parent", None)
|
||||
if parent is None:
|
||||
return None
|
||||
return self._store_remote_object_handle(parent)
|
||||
if method_name == "get_inner_model_attr":
|
||||
attr_name = args[0]
|
||||
if hasattr(obj.model, attr_name):
|
||||
return getattr(obj.model, attr_name)
|
||||
if hasattr(obj, attr_name):
|
||||
return getattr(obj, attr_name)
|
||||
return None
|
||||
if method_name == "inner_model_apply_model":
|
||||
return obj.model.apply_model(*args[0], **args[1])
|
||||
if method_name == "inner_model_extra_conds_shapes":
|
||||
return obj.model.extra_conds_shapes(*args[0], **args[1])
|
||||
if method_name == "inner_model_extra_conds":
|
||||
return obj.model.extra_conds(*args[0], **args[1])
|
||||
if method_name == "inner_model_memory_required":
|
||||
return obj.model.memory_required(*args[0], **args[1])
|
||||
if method_name == "process_latent_in":
|
||||
return obj.model.process_latent_in(*args[0], **args[1])
|
||||
if method_name == "process_latent_out":
|
||||
return obj.model.process_latent_out(*args[0], **args[1])
|
||||
if method_name == "scale_latent_inpaint":
|
||||
return obj.model.scale_latent_inpaint(*args[0], **args[1])
|
||||
if method_name.startswith("get_"):
|
||||
attr_name = method_name[4:]
|
||||
if hasattr(obj, attr_name):
|
||||
return getattr(obj, attr_name)
|
||||
|
||||
target = getattr(obj, method_name)
|
||||
if callable(target):
|
||||
result = target(*args, **kwargs)
|
||||
if inspect.isawaitable(result):
|
||||
result = await result
|
||||
if type(result).__name__ == "ModelPatcher":
|
||||
return self._store_remote_object_handle(result)
|
||||
return result
|
||||
if args or kwargs:
|
||||
raise TypeError(f"{method_name} is not callable on remote object {object_id}")
|
||||
return target
|
||||
|
||||
def _wrap_unpicklable_objects(self, data: Any) -> Any:
|
||||
if isinstance(data, (str, int, float, bool, type(None))):
|
||||
return data
|
||||
if isinstance(data, torch.Tensor):
|
||||
tensor = data.detach() if data.requires_grad else data
|
||||
if os.environ.get("PYISOLATE_CHILD") == "1" and tensor.device.type != "cpu":
|
||||
return tensor.cpu()
|
||||
return tensor
|
||||
|
||||
# Special-case clip vision outputs: preserve attribute access by packing fields
|
||||
if hasattr(data, "penultimate_hidden_states") or hasattr(
|
||||
data, "last_hidden_state"
|
||||
):
|
||||
fields = {}
|
||||
for attr in (
|
||||
"penultimate_hidden_states",
|
||||
"last_hidden_state",
|
||||
"image_embeds",
|
||||
"text_embeds",
|
||||
):
|
||||
if hasattr(data, attr):
|
||||
try:
|
||||
fields[attr] = self._wrap_unpicklable_objects(
|
||||
getattr(data, attr)
|
||||
)
|
||||
except Exception:
|
||||
pass
|
||||
if fields:
|
||||
return {"__pyisolate_attribute_container__": True, "data": fields}
|
||||
|
||||
# Avoid converting arbitrary objects with stateful methods (models, etc.)
|
||||
# They will be handled via RemoteObjectHandle below.
|
||||
|
||||
type_name = type(data).__name__
|
||||
if type_name == "ModelPatcherProxy":
|
||||
return {"__type__": "ModelPatcherRef", "model_id": data._instance_id}
|
||||
if type_name == "CLIPProxy":
|
||||
return {"__type__": "CLIPRef", "clip_id": data._instance_id}
|
||||
if type_name == "VAEProxy":
|
||||
return {"__type__": "VAERef", "vae_id": data._instance_id}
|
||||
if type_name == "ModelSamplingProxy":
|
||||
return {"__type__": "ModelSamplingRef", "ms_id": data._instance_id}
|
||||
|
||||
if isinstance(data, (list, tuple)):
|
||||
wrapped = [self._wrap_unpicklable_objects(item) for item in data]
|
||||
return tuple(wrapped) if isinstance(data, tuple) else wrapped
|
||||
if isinstance(data, dict):
|
||||
converted_dict = {
|
||||
k: self._wrap_unpicklable_objects(v) for k, v in data.items()
|
||||
}
|
||||
return {"__pyisolate_attrdict__": True, "data": converted_dict}
|
||||
|
||||
from pyisolate._internal.serialization_registry import SerializerRegistry
|
||||
|
||||
registry = SerializerRegistry.get_instance()
|
||||
if registry.is_data_type(type_name):
|
||||
serializer = registry.get_serializer(type_name)
|
||||
if serializer:
|
||||
return serializer(data)
|
||||
|
||||
return self._store_remote_object_handle(data)
|
||||
|
||||
def _resolve_remote_objects(self, data: Any) -> Any:
|
||||
if isinstance(data, RemoteObjectHandle):
|
||||
if data.object_id not in self.remote_objects:
|
||||
raise KeyError(f"Remote object {data.object_id} not found")
|
||||
return self.remote_objects[data.object_id]
|
||||
|
||||
if isinstance(data, dict):
|
||||
ref_type = data.get("__type__")
|
||||
if ref_type in ("CLIPRef", "ModelPatcherRef", "VAERef"):
|
||||
from pyisolate._internal.model_serialization import (
|
||||
deserialize_proxy_result,
|
||||
)
|
||||
|
||||
return deserialize_proxy_result(data)
|
||||
if ref_type == "ModelSamplingRef":
|
||||
from pyisolate._internal.model_serialization import (
|
||||
deserialize_proxy_result,
|
||||
)
|
||||
|
||||
return deserialize_proxy_result(data)
|
||||
return {k: self._resolve_remote_objects(v) for k, v in data.items()}
|
||||
|
||||
if isinstance(data, (list, tuple)):
|
||||
resolved = [self._resolve_remote_objects(item) for item in data]
|
||||
return tuple(resolved) if isinstance(data, tuple) else resolved
|
||||
return data
|
||||
|
||||
def _get_node_class(self, node_name: str) -> type:
|
||||
if node_name not in self.node_classes:
|
||||
raise KeyError(f"Unknown node: {node_name}")
|
||||
return self.node_classes[node_name]
|
||||
|
||||
def _get_node_instance(self, node_name: str) -> Any:
|
||||
if node_name not in self.node_instances:
|
||||
if node_name not in self.node_classes:
|
||||
raise KeyError(f"Unknown node: {node_name}")
|
||||
self.node_instances[node_name] = self.node_classes[node_name]()
|
||||
return self.node_instances[node_name]
|
||||
|
||||
async def before_module_loaded(self) -> None:
|
||||
try:
|
||||
from comfy.isolation import initialize_proxies
|
||||
|
||||
initialize_proxies()
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
"%s before_module_loaded initialize_proxies FAILED: %s", LOG_PREFIX, e
|
||||
)
|
||||
|
||||
await super().before_module_loaded()
|
||||
try:
|
||||
from comfy_api.latest import ComfyAPI_latest
|
||||
from .proxies.progress_proxy import ProgressProxy
|
||||
|
||||
ComfyAPI_latest.Execution = ProgressProxy
|
||||
# ComfyAPI_latest.execution = ProgressProxy() # Eliminated to avoid Singleton collision
|
||||
# fp_proxy = FolderPathsProxy() # Eliminated to avoid Singleton collision
|
||||
# latest_ui.folder_paths = fp_proxy
|
||||
# latest_resources.folder_paths = fp_proxy
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
async def call_route_handler(
|
||||
self,
|
||||
handler_module: str,
|
||||
handler_func: str,
|
||||
request_data: Dict[str, Any],
|
||||
) -> Any:
|
||||
cache_key = f"{handler_module}.{handler_func}"
|
||||
if cache_key not in self._route_handlers:
|
||||
if self._module is not None and hasattr(self._module, "__file__"):
|
||||
node_dir = os.path.dirname(self._module.__file__)
|
||||
if node_dir not in sys.path:
|
||||
sys.path.insert(0, node_dir)
|
||||
try:
|
||||
module = importlib.import_module(handler_module)
|
||||
self._route_handlers[cache_key] = getattr(module, handler_func)
|
||||
except (ImportError, AttributeError) as e:
|
||||
raise ValueError(f"Route handler not found: {cache_key}") from e
|
||||
|
||||
handler = self._route_handlers[cache_key]
|
||||
mock_request = MockRequest(request_data)
|
||||
|
||||
if asyncio.iscoroutinefunction(handler):
|
||||
result = await handler(mock_request)
|
||||
else:
|
||||
result = handler(mock_request)
|
||||
return self._serialize_response(result)
|
||||
|
||||
def _is_comfy_protocol_return(self, result: Any) -> bool:
|
||||
"""
|
||||
Check if the result matches the ComfyUI 'Protocol Return' schema.
|
||||
|
||||
A Protocol Return is a dictionary containing specific reserved keys that
|
||||
ComfyUI's execution engine interprets as instructions (UI updates,
|
||||
Workflow expansion, etc.) rather than purely data outputs.
|
||||
|
||||
Schema:
|
||||
- Must be a dict
|
||||
- Must contain at least one of: 'ui', 'result', 'expand'
|
||||
"""
|
||||
if not isinstance(result, dict):
|
||||
return False
|
||||
return any(key in result for key in ("ui", "result", "expand"))
|
||||
|
||||
def _serialize_response(self, response: Any) -> Dict[str, Any]:
|
||||
if response is None:
|
||||
return {"type": "text", "body": "", "status": 204}
|
||||
if isinstance(response, dict):
|
||||
return {"type": "json", "body": response, "status": 200}
|
||||
if isinstance(response, str):
|
||||
return {"type": "text", "body": response, "status": 200}
|
||||
if hasattr(response, "text") and hasattr(response, "status"):
|
||||
return {
|
||||
"type": "text",
|
||||
"body": response.text
|
||||
if hasattr(response, "text")
|
||||
else str(response.body),
|
||||
"status": response.status,
|
||||
"headers": dict(response.headers)
|
||||
if hasattr(response, "headers")
|
||||
else {},
|
||||
}
|
||||
if hasattr(response, "body") and hasattr(response, "status"):
|
||||
body = response.body
|
||||
if isinstance(body, bytes):
|
||||
try:
|
||||
return {
|
||||
"type": "text",
|
||||
"body": body.decode("utf-8"),
|
||||
"status": response.status,
|
||||
}
|
||||
except UnicodeDecodeError:
|
||||
return {
|
||||
"type": "binary",
|
||||
"body": body.hex(),
|
||||
"status": response.status,
|
||||
}
|
||||
return {"type": "json", "body": body, "status": response.status}
|
||||
return {"type": "text", "body": str(response), "status": 200}
|
||||
|
||||
|
||||
class MockRequest:
|
||||
def __init__(self, data: Dict[str, Any]):
|
||||
self.method = data.get("method", "GET")
|
||||
self.path = data.get("path", "/")
|
||||
self.query = data.get("query", {})
|
||||
self._body = data.get("body", {})
|
||||
self._text = data.get("text", "")
|
||||
self.headers = data.get("headers", {})
|
||||
self.content_type = data.get(
|
||||
"content_type", self.headers.get("Content-Type", "application/json")
|
||||
)
|
||||
self.match_info = data.get("match_info", {})
|
||||
|
||||
async def json(self) -> Any:
|
||||
if isinstance(self._body, dict):
|
||||
return self._body
|
||||
if isinstance(self._body, str):
|
||||
return json.loads(self._body)
|
||||
return {}
|
||||
|
||||
async def post(self) -> Dict[str, Any]:
|
||||
if isinstance(self._body, dict):
|
||||
return self._body
|
||||
return {}
|
||||
|
||||
async def text(self) -> str:
|
||||
if self._text:
|
||||
return self._text
|
||||
if isinstance(self._body, str):
|
||||
return self._body
|
||||
if isinstance(self._body, dict):
|
||||
return json.dumps(self._body)
|
||||
return ""
|
||||
|
||||
async def read(self) -> bytes:
|
||||
return (await self.text()).encode("utf-8")
|
||||
@@ -1,25 +0,0 @@
|
||||
# pylint: disable=import-outside-toplevel
|
||||
# Host process initialization for PyIsolate
|
||||
import logging
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def initialize_host_process() -> None:
|
||||
from .proxies.folder_paths_proxy import FolderPathsProxy
|
||||
from .proxies.helper_proxies import HelperProxiesService
|
||||
from .proxies.model_management_proxy import ModelManagementProxy
|
||||
from .proxies.progress_proxy import ProgressProxy
|
||||
from .proxies.prompt_server_impl import PromptServerService
|
||||
from .proxies.utils_proxy import UtilsProxy
|
||||
from .proxies.web_directory_proxy import WebDirectoryProxy
|
||||
from .vae_proxy import VAERegistry
|
||||
|
||||
FolderPathsProxy()
|
||||
HelperProxiesService()
|
||||
ModelManagementProxy()
|
||||
ProgressProxy()
|
||||
PromptServerService()
|
||||
UtilsProxy()
|
||||
WebDirectoryProxy()
|
||||
VAERegistry()
|
||||
@@ -1,180 +0,0 @@
|
||||
# pylint: disable=logging-fstring-interpolation
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import os
|
||||
from pathlib import Path
|
||||
from pathlib import PurePosixPath
|
||||
from typing import Dict, List, TypedDict
|
||||
|
||||
try:
|
||||
import tomllib
|
||||
except ImportError:
|
||||
import tomli as tomllib # type: ignore[no-redef]
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
HOST_POLICY_PATH_ENV = "COMFY_HOST_POLICY_PATH"
|
||||
VALID_SANDBOX_MODES = frozenset({"required", "disabled"})
|
||||
FORBIDDEN_WRITABLE_PATHS = frozenset({"/tmp"})
|
||||
|
||||
|
||||
class HostSecurityPolicy(TypedDict):
|
||||
sandbox_mode: str
|
||||
allow_network: bool
|
||||
writable_paths: List[str]
|
||||
readonly_paths: List[str]
|
||||
sealed_worker_ro_import_paths: List[str]
|
||||
whitelist: Dict[str, str]
|
||||
|
||||
|
||||
DEFAULT_POLICY: HostSecurityPolicy = {
|
||||
"sandbox_mode": "required",
|
||||
"allow_network": False,
|
||||
"writable_paths": ["/dev/shm"],
|
||||
"readonly_paths": [],
|
||||
"sealed_worker_ro_import_paths": [],
|
||||
"whitelist": {},
|
||||
}
|
||||
|
||||
|
||||
def _default_policy() -> HostSecurityPolicy:
|
||||
return {
|
||||
"sandbox_mode": DEFAULT_POLICY["sandbox_mode"],
|
||||
"allow_network": DEFAULT_POLICY["allow_network"],
|
||||
"writable_paths": list(DEFAULT_POLICY["writable_paths"]),
|
||||
"readonly_paths": list(DEFAULT_POLICY["readonly_paths"]),
|
||||
"sealed_worker_ro_import_paths": list(DEFAULT_POLICY["sealed_worker_ro_import_paths"]),
|
||||
"whitelist": dict(DEFAULT_POLICY["whitelist"]),
|
||||
}
|
||||
|
||||
|
||||
def _normalize_writable_paths(paths: list[object]) -> list[str]:
|
||||
normalized_paths: list[str] = []
|
||||
for raw_path in paths:
|
||||
# Host-policy paths are contract-style POSIX paths; keep representation
|
||||
# stable across Windows/Linux so tests and config behavior stay consistent.
|
||||
normalized_path = str(PurePosixPath(str(raw_path).replace("\\", "/")))
|
||||
if normalized_path in FORBIDDEN_WRITABLE_PATHS:
|
||||
continue
|
||||
normalized_paths.append(normalized_path)
|
||||
return normalized_paths
|
||||
|
||||
|
||||
def _load_whitelist_file(file_path: Path, config_path: Path) -> Dict[str, str]:
|
||||
if not file_path.is_absolute():
|
||||
file_path = config_path.parent / file_path
|
||||
if not file_path.exists():
|
||||
logger.warning("whitelist_file %s not found, skipping.", file_path)
|
||||
return {}
|
||||
entries: Dict[str, str] = {}
|
||||
for line in file_path.read_text().splitlines():
|
||||
line = line.strip()
|
||||
if not line or line.startswith("#"):
|
||||
continue
|
||||
entries[line] = "*"
|
||||
logger.debug("Loaded %d whitelist entries from %s", len(entries), file_path)
|
||||
return entries
|
||||
|
||||
|
||||
def _normalize_sealed_worker_ro_import_paths(raw_paths: object) -> list[str]:
|
||||
if not isinstance(raw_paths, list):
|
||||
raise ValueError(
|
||||
"tool.comfy.host.sealed_worker_ro_import_paths must be a list of absolute paths."
|
||||
)
|
||||
|
||||
normalized_paths: list[str] = []
|
||||
seen: set[str] = set()
|
||||
for raw_path in raw_paths:
|
||||
if not isinstance(raw_path, str) or not raw_path.strip():
|
||||
raise ValueError(
|
||||
"tool.comfy.host.sealed_worker_ro_import_paths entries must be non-empty strings."
|
||||
)
|
||||
normalized_path = str(PurePosixPath(raw_path.replace("\\", "/")))
|
||||
# Accept both POSIX absolute paths (/home/...) and Windows drive-letter paths (D:/...)
|
||||
is_absolute = normalized_path.startswith("/") or (
|
||||
len(normalized_path) >= 3 and normalized_path[1] == ":" and normalized_path[2] == "/"
|
||||
)
|
||||
if not is_absolute:
|
||||
raise ValueError(
|
||||
"tool.comfy.host.sealed_worker_ro_import_paths entries must be absolute paths."
|
||||
)
|
||||
if normalized_path not in seen:
|
||||
seen.add(normalized_path)
|
||||
normalized_paths.append(normalized_path)
|
||||
|
||||
return normalized_paths
|
||||
|
||||
|
||||
def load_host_policy(comfy_root: Path) -> HostSecurityPolicy:
|
||||
config_override = os.environ.get(HOST_POLICY_PATH_ENV)
|
||||
config_path = Path(config_override) if config_override else comfy_root / "pyproject.toml"
|
||||
policy = _default_policy()
|
||||
|
||||
if not config_path.exists():
|
||||
logger.debug("Host policy file missing at %s, using defaults.", config_path)
|
||||
return policy
|
||||
|
||||
try:
|
||||
with config_path.open("rb") as f:
|
||||
data = tomllib.load(f)
|
||||
except Exception:
|
||||
logger.warning(
|
||||
"Failed to parse host policy from %s, using defaults.",
|
||||
config_path,
|
||||
exc_info=True,
|
||||
)
|
||||
return policy
|
||||
|
||||
tool_config = data.get("tool", {}).get("comfy", {}).get("host", {})
|
||||
if not isinstance(tool_config, dict):
|
||||
logger.debug("No [tool.comfy.host] section found, using defaults.")
|
||||
return policy
|
||||
|
||||
sandbox_mode = tool_config.get("sandbox_mode")
|
||||
if isinstance(sandbox_mode, str):
|
||||
normalized_sandbox_mode = sandbox_mode.strip().lower()
|
||||
if normalized_sandbox_mode in VALID_SANDBOX_MODES:
|
||||
policy["sandbox_mode"] = normalized_sandbox_mode
|
||||
else:
|
||||
logger.warning(
|
||||
"Invalid host sandbox_mode %r in %s, using default %r.",
|
||||
sandbox_mode,
|
||||
config_path,
|
||||
DEFAULT_POLICY["sandbox_mode"],
|
||||
)
|
||||
|
||||
if "allow_network" in tool_config:
|
||||
policy["allow_network"] = bool(tool_config["allow_network"])
|
||||
|
||||
if "writable_paths" in tool_config:
|
||||
policy["writable_paths"] = _normalize_writable_paths(tool_config["writable_paths"])
|
||||
|
||||
if "readonly_paths" in tool_config:
|
||||
policy["readonly_paths"] = [str(p) for p in tool_config["readonly_paths"]]
|
||||
|
||||
if "sealed_worker_ro_import_paths" in tool_config:
|
||||
policy["sealed_worker_ro_import_paths"] = _normalize_sealed_worker_ro_import_paths(
|
||||
tool_config["sealed_worker_ro_import_paths"]
|
||||
)
|
||||
|
||||
whitelist_file = tool_config.get("whitelist_file")
|
||||
if isinstance(whitelist_file, str):
|
||||
policy["whitelist"].update(_load_whitelist_file(Path(whitelist_file), config_path))
|
||||
|
||||
whitelist_raw = tool_config.get("whitelist")
|
||||
if isinstance(whitelist_raw, dict):
|
||||
policy["whitelist"].update({str(k): str(v) for k, v in whitelist_raw.items()})
|
||||
|
||||
os.environ["PYISOLATE_SANDBOX_MODE"] = policy["sandbox_mode"]
|
||||
|
||||
logger.debug(
|
||||
"Loaded Host Policy: %d whitelisted nodes, Sandbox=%s, Network=%s",
|
||||
len(policy["whitelist"]),
|
||||
policy["sandbox_mode"],
|
||||
policy["allow_network"],
|
||||
)
|
||||
return policy
|
||||
|
||||
|
||||
__all__ = ["HostSecurityPolicy", "load_host_policy", "DEFAULT_POLICY"]
|
||||
@@ -1,221 +0,0 @@
|
||||
# pylint: disable=import-outside-toplevel
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import sys
|
||||
import tempfile
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, List, Optional, Tuple
|
||||
|
||||
import folder_paths
|
||||
|
||||
try:
|
||||
import tomllib
|
||||
except ImportError:
|
||||
import tomli as tomllib # type: ignore[no-redef]
|
||||
|
||||
LOG_PREFIX = "]["
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
CACHE_SUBDIR = "cache"
|
||||
CACHE_KEY_FILE = "cache_key"
|
||||
CACHE_DATA_FILE = "node_info.json"
|
||||
CACHE_KEY_LENGTH = 16
|
||||
_NESTED_SCAN_ROOT = "packages"
|
||||
_IGNORED_MANIFEST_DIRS = {".git", ".venv", "__pycache__"}
|
||||
|
||||
|
||||
def _read_manifest(manifest_path: Path) -> dict[str, Any] | None:
|
||||
try:
|
||||
with manifest_path.open("rb") as f:
|
||||
data = tomllib.load(f)
|
||||
if isinstance(data, dict):
|
||||
return data
|
||||
except Exception:
|
||||
return None
|
||||
return None
|
||||
|
||||
|
||||
def _is_isolation_manifest(data: dict[str, Any]) -> bool:
|
||||
return (
|
||||
"tool" in data
|
||||
and "comfy" in data["tool"]
|
||||
and "isolation" in data["tool"]["comfy"]
|
||||
)
|
||||
|
||||
|
||||
def _discover_nested_manifests(entry: Path) -> List[Tuple[Path, Path]]:
|
||||
packages_root = entry / _NESTED_SCAN_ROOT
|
||||
if not packages_root.exists() or not packages_root.is_dir():
|
||||
return []
|
||||
|
||||
nested: List[Tuple[Path, Path]] = []
|
||||
for manifest in sorted(packages_root.rglob("pyproject.toml")):
|
||||
node_dir = manifest.parent
|
||||
if any(part in _IGNORED_MANIFEST_DIRS for part in node_dir.parts):
|
||||
continue
|
||||
|
||||
data = _read_manifest(manifest)
|
||||
if not data or not _is_isolation_manifest(data):
|
||||
continue
|
||||
|
||||
isolation = data["tool"]["comfy"]["isolation"]
|
||||
if isolation.get("standalone") is True:
|
||||
nested.append((node_dir, manifest))
|
||||
|
||||
return nested
|
||||
|
||||
|
||||
def find_manifest_directories() -> List[Tuple[Path, Path]]:
|
||||
"""Find custom node directories containing a valid pyproject.toml with [tool.comfy.isolation]."""
|
||||
manifest_dirs: List[Tuple[Path, Path]] = []
|
||||
|
||||
# Standard custom_nodes paths
|
||||
for base_path in folder_paths.get_folder_paths("custom_nodes"):
|
||||
base = Path(base_path)
|
||||
if not base.exists() or not base.is_dir():
|
||||
continue
|
||||
|
||||
for entry in base.iterdir():
|
||||
if not entry.is_dir():
|
||||
continue
|
||||
|
||||
# Look for pyproject.toml
|
||||
manifest = entry / "pyproject.toml"
|
||||
if not manifest.exists():
|
||||
continue
|
||||
|
||||
data = _read_manifest(manifest)
|
||||
if not data or not _is_isolation_manifest(data):
|
||||
continue
|
||||
|
||||
manifest_dirs.append((entry, manifest))
|
||||
manifest_dirs.extend(_discover_nested_manifests(entry))
|
||||
|
||||
return manifest_dirs
|
||||
|
||||
|
||||
def compute_cache_key(node_dir: Path, manifest_path: Path) -> str:
|
||||
"""Hash manifest + .py mtimes + Python version + PyIsolate version."""
|
||||
hasher = hashlib.sha256()
|
||||
|
||||
try:
|
||||
# Hashing the manifest content ensures config changes invalidate cache
|
||||
hasher.update(manifest_path.read_bytes())
|
||||
except OSError:
|
||||
hasher.update(b"__manifest_read_error__")
|
||||
|
||||
try:
|
||||
py_files = sorted(node_dir.rglob("*.py"))
|
||||
for py_file in py_files:
|
||||
rel_path = py_file.relative_to(node_dir)
|
||||
if "__pycache__" in str(rel_path) or ".venv" in str(rel_path):
|
||||
continue
|
||||
hasher.update(str(rel_path).encode("utf-8"))
|
||||
try:
|
||||
hasher.update(str(py_file.stat().st_mtime).encode("utf-8"))
|
||||
except OSError:
|
||||
hasher.update(b"__file_stat_error__")
|
||||
except OSError:
|
||||
hasher.update(b"__dir_scan_error__")
|
||||
|
||||
hasher.update(sys.version.encode("utf-8"))
|
||||
|
||||
try:
|
||||
import pyisolate
|
||||
|
||||
hasher.update(pyisolate.__version__.encode("utf-8"))
|
||||
except (ImportError, AttributeError):
|
||||
hasher.update(b"__pyisolate_unknown__")
|
||||
|
||||
return hasher.hexdigest()[:CACHE_KEY_LENGTH]
|
||||
|
||||
|
||||
def get_cache_path(node_dir: Path, venv_root: Path) -> Tuple[Path, Path]:
|
||||
"""Return (cache_key_file, cache_data_file) in venv_root/{node}/cache/."""
|
||||
cache_dir = venv_root / node_dir.name / CACHE_SUBDIR
|
||||
return (cache_dir / CACHE_KEY_FILE, cache_dir / CACHE_DATA_FILE)
|
||||
|
||||
|
||||
def is_cache_valid(node_dir: Path, manifest_path: Path, venv_root: Path) -> bool:
|
||||
"""Return True only if stored cache key matches current computed key."""
|
||||
try:
|
||||
cache_key_file, cache_data_file = get_cache_path(node_dir, venv_root)
|
||||
if not cache_key_file.exists() or not cache_data_file.exists():
|
||||
return False
|
||||
current_key = compute_cache_key(node_dir, manifest_path)
|
||||
stored_key = cache_key_file.read_text(encoding="utf-8").strip()
|
||||
return current_key == stored_key
|
||||
except Exception as e:
|
||||
logger.debug(
|
||||
"%s Cache validation error for %s: %s", LOG_PREFIX, node_dir.name, e
|
||||
)
|
||||
return False
|
||||
|
||||
|
||||
def load_from_cache(node_dir: Path, venv_root: Path) -> Optional[Dict[str, Any]]:
|
||||
"""Load node metadata from cache, return None on any error."""
|
||||
try:
|
||||
_, cache_data_file = get_cache_path(node_dir, venv_root)
|
||||
if not cache_data_file.exists():
|
||||
return None
|
||||
data = json.loads(cache_data_file.read_text(encoding="utf-8"))
|
||||
if not isinstance(data, dict):
|
||||
return None
|
||||
return data
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
|
||||
def save_to_cache(
|
||||
node_dir: Path, venv_root: Path, node_data: Dict[str, Any], manifest_path: Path
|
||||
) -> None:
|
||||
"""Save node metadata and cache key atomically."""
|
||||
try:
|
||||
cache_key_file, cache_data_file = get_cache_path(node_dir, venv_root)
|
||||
cache_dir = cache_key_file.parent
|
||||
cache_dir.mkdir(parents=True, exist_ok=True)
|
||||
cache_key = compute_cache_key(node_dir, manifest_path)
|
||||
|
||||
# Atomic write: data
|
||||
tmp_data_fd, tmp_data_path = tempfile.mkstemp(dir=str(cache_dir), suffix=".tmp")
|
||||
try:
|
||||
with os.fdopen(tmp_data_fd, "w", encoding="utf-8") as f:
|
||||
json.dump(node_data, f, indent=2)
|
||||
os.replace(tmp_data_path, cache_data_file)
|
||||
except Exception:
|
||||
try:
|
||||
os.unlink(tmp_data_path)
|
||||
except OSError:
|
||||
pass
|
||||
raise
|
||||
|
||||
# Atomic write: key
|
||||
tmp_key_fd, tmp_key_path = tempfile.mkstemp(dir=str(cache_dir), suffix=".tmp")
|
||||
try:
|
||||
with os.fdopen(tmp_key_fd, "w", encoding="utf-8") as f:
|
||||
f.write(cache_key)
|
||||
os.replace(tmp_key_path, cache_key_file)
|
||||
except Exception:
|
||||
try:
|
||||
os.unlink(tmp_key_path)
|
||||
except OSError:
|
||||
pass
|
||||
raise
|
||||
|
||||
except Exception as e:
|
||||
logger.warning("%s Cache save failed for %s: %s", LOG_PREFIX, node_dir.name, e)
|
||||
|
||||
|
||||
__all__ = [
|
||||
"LOG_PREFIX",
|
||||
"find_manifest_directories",
|
||||
"compute_cache_key",
|
||||
"get_cache_path",
|
||||
"is_cache_valid",
|
||||
"load_from_cache",
|
||||
"save_to_cache",
|
||||
]
|
||||
@@ -1,891 +0,0 @@
|
||||
# pylint: disable=bare-except,consider-using-from-import,import-outside-toplevel,protected-access
|
||||
# RPC proxy for ModelPatcher (parent process)
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from typing import Any, Optional, List, Set, Dict, Callable
|
||||
|
||||
from comfy.isolation.proxies.base import (
|
||||
IS_CHILD_PROCESS,
|
||||
BaseProxy,
|
||||
)
|
||||
from comfy.isolation.model_patcher_proxy_registry import (
|
||||
ModelPatcherRegistry,
|
||||
AutoPatcherEjector,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class ModelPatcherProxy(BaseProxy[ModelPatcherRegistry]):
|
||||
_registry_class = ModelPatcherRegistry
|
||||
__module__ = "comfy.model_patcher"
|
||||
_APPLY_MODEL_GUARD_PADDING_BYTES = 32 * 1024 * 1024
|
||||
|
||||
def _spawn_related_proxy(self, instance_id: str) -> "ModelPatcherProxy":
|
||||
proxy = ModelPatcherProxy(
|
||||
instance_id,
|
||||
self._registry,
|
||||
manage_lifecycle=not IS_CHILD_PROCESS,
|
||||
)
|
||||
if getattr(self, "_rpc_caller", None) is not None:
|
||||
proxy._rpc_caller = self._rpc_caller
|
||||
return proxy
|
||||
|
||||
def _get_rpc(self) -> Any:
|
||||
if self._rpc_caller is None:
|
||||
from pyisolate._internal.rpc_protocol import get_child_rpc_instance
|
||||
|
||||
rpc = get_child_rpc_instance()
|
||||
if rpc is not None:
|
||||
self._rpc_caller = rpc.create_caller(
|
||||
self._registry_class, self._registry_class.get_remote_id()
|
||||
)
|
||||
else:
|
||||
self._rpc_caller = self._registry
|
||||
return self._rpc_caller
|
||||
|
||||
def get_all_callbacks(self, call_type: str = None) -> Any:
|
||||
return self._call_rpc("get_all_callbacks", call_type)
|
||||
|
||||
def get_all_wrappers(self, wrapper_type: str = None) -> Any:
|
||||
return self._call_rpc("get_all_wrappers", wrapper_type)
|
||||
|
||||
def _load_list(self, *args, **kwargs) -> Any:
|
||||
return self._call_rpc("load_list_internal", *args, **kwargs)
|
||||
|
||||
def prepare_hook_patches_current_keyframe(
|
||||
self, t: Any, hook_group: Any, model_options: Any
|
||||
) -> None:
|
||||
self._call_rpc(
|
||||
"prepare_hook_patches_current_keyframe", t, hook_group, model_options
|
||||
)
|
||||
|
||||
def add_hook_patches(
|
||||
self,
|
||||
hook: Any,
|
||||
patches: Any,
|
||||
strength_patch: float = 1.0,
|
||||
strength_model: float = 1.0,
|
||||
) -> None:
|
||||
self._call_rpc(
|
||||
"add_hook_patches", hook, patches, strength_patch, strength_model
|
||||
)
|
||||
|
||||
def clear_cached_hook_weights(self) -> None:
|
||||
self._call_rpc("clear_cached_hook_weights")
|
||||
|
||||
def get_combined_hook_patches(self, hooks: Any) -> Any:
|
||||
return self._call_rpc("get_combined_hook_patches", hooks)
|
||||
|
||||
def get_additional_models_with_key(self, key: str) -> Any:
|
||||
return self._call_rpc("get_additional_models_with_key", key)
|
||||
|
||||
@property
|
||||
def object_patches(self) -> Any:
|
||||
return self._call_rpc("get_object_patches")
|
||||
|
||||
@property
|
||||
def patches(self) -> Any:
|
||||
res = self._call_rpc("get_patches")
|
||||
if isinstance(res, dict):
|
||||
new_res = {}
|
||||
for k, v in res.items():
|
||||
new_list = []
|
||||
for item in v:
|
||||
if isinstance(item, list):
|
||||
new_list.append(tuple(item))
|
||||
else:
|
||||
new_list.append(item)
|
||||
new_res[k] = new_list
|
||||
return new_res
|
||||
return res
|
||||
|
||||
@property
|
||||
def pinned(self) -> Set:
|
||||
val = self._call_rpc("get_patcher_attr", "pinned")
|
||||
return set(val) if val is not None else set()
|
||||
|
||||
@property
|
||||
def hook_patches(self) -> Dict:
|
||||
val = self._call_rpc("get_patcher_attr", "hook_patches")
|
||||
if val is None:
|
||||
return {}
|
||||
try:
|
||||
from comfy.hooks import _HookRef
|
||||
import json
|
||||
|
||||
new_val = {}
|
||||
for k, v in val.items():
|
||||
if isinstance(k, str):
|
||||
if k.startswith("PYISOLATE_HOOKREF:"):
|
||||
ref_id = k.split(":", 1)[1]
|
||||
h = _HookRef()
|
||||
h._pyisolate_id = ref_id
|
||||
new_val[h] = v
|
||||
elif k.startswith("__pyisolate_key__"):
|
||||
try:
|
||||
json_str = k[len("__pyisolate_key__") :]
|
||||
data = json.loads(json_str)
|
||||
ref_id = None
|
||||
if isinstance(data, list):
|
||||
for item in data:
|
||||
if (
|
||||
isinstance(item, list)
|
||||
and len(item) == 2
|
||||
and item[0] == "id"
|
||||
):
|
||||
ref_id = item[1]
|
||||
break
|
||||
if ref_id:
|
||||
h = _HookRef()
|
||||
h._pyisolate_id = ref_id
|
||||
new_val[h] = v
|
||||
else:
|
||||
new_val[k] = v
|
||||
except Exception:
|
||||
new_val[k] = v
|
||||
else:
|
||||
new_val[k] = v
|
||||
else:
|
||||
new_val[k] = v
|
||||
return new_val
|
||||
except ImportError:
|
||||
return val
|
||||
|
||||
def set_hook_mode(self, hook_mode: Any) -> None:
|
||||
self._call_rpc("set_hook_mode", hook_mode)
|
||||
|
||||
def register_all_hook_patches(
|
||||
self,
|
||||
hooks: Any,
|
||||
target_dict: Any,
|
||||
model_options: Any = None,
|
||||
registered: Any = None,
|
||||
) -> None:
|
||||
self._call_rpc(
|
||||
"register_all_hook_patches", hooks, target_dict, model_options, registered
|
||||
)
|
||||
|
||||
def is_clone(self, other: Any) -> bool:
|
||||
if isinstance(other, ModelPatcherProxy):
|
||||
return self._call_rpc("is_clone_by_id", other._instance_id)
|
||||
return False
|
||||
|
||||
def clone(self) -> ModelPatcherProxy:
|
||||
new_id = self._call_rpc("clone")
|
||||
return self._spawn_related_proxy(new_id)
|
||||
|
||||
def clone_has_same_weights(self, clone: Any) -> bool:
|
||||
if isinstance(clone, ModelPatcherProxy):
|
||||
return self._call_rpc("clone_has_same_weights_by_id", clone._instance_id)
|
||||
if not IS_CHILD_PROCESS:
|
||||
return self._call_rpc("is_clone", clone)
|
||||
return False
|
||||
|
||||
def get_model_object(self, name: str) -> Any:
|
||||
return self._call_rpc("get_model_object", name)
|
||||
|
||||
@property
|
||||
def model_options(self) -> dict:
|
||||
data = self._call_rpc("get_model_options")
|
||||
import json
|
||||
|
||||
def _decode_keys(obj):
|
||||
if isinstance(obj, dict):
|
||||
new_d = {}
|
||||
for k, v in obj.items():
|
||||
if isinstance(k, str) and k.startswith("__pyisolate_key__"):
|
||||
try:
|
||||
json_str = k[17:]
|
||||
val = json.loads(json_str)
|
||||
if isinstance(val, list):
|
||||
val = tuple(val)
|
||||
new_d[val] = _decode_keys(v)
|
||||
except:
|
||||
new_d[k] = _decode_keys(v)
|
||||
else:
|
||||
new_d[k] = _decode_keys(v)
|
||||
return new_d
|
||||
if isinstance(obj, list):
|
||||
return [_decode_keys(x) for x in obj]
|
||||
return obj
|
||||
|
||||
return _decode_keys(data)
|
||||
|
||||
@model_options.setter
|
||||
def model_options(self, value: dict) -> None:
|
||||
self._call_rpc("set_model_options", value)
|
||||
|
||||
def apply_hooks(self, hooks: Any) -> Any:
|
||||
return self._call_rpc("apply_hooks", hooks)
|
||||
|
||||
def prepare_state(self, timestep: Any) -> Any:
|
||||
return self._call_rpc("prepare_state", timestep)
|
||||
|
||||
def restore_hook_patches(self) -> None:
|
||||
self._call_rpc("restore_hook_patches")
|
||||
|
||||
def unpatch_hooks(self, whitelist_keys_set: Optional[Set[str]] = None) -> None:
|
||||
self._call_rpc("unpatch_hooks", whitelist_keys_set)
|
||||
|
||||
def model_patches_to(self, device: Any) -> Any:
|
||||
return self._call_rpc("model_patches_to", device)
|
||||
|
||||
def partially_load(
|
||||
self, device: Any, extra_memory: Any, force_patch_weights: bool = False
|
||||
) -> Any:
|
||||
return self._call_rpc(
|
||||
"partially_load", device, extra_memory, force_patch_weights
|
||||
)
|
||||
|
||||
def partially_unload(
|
||||
self, device_to: Any, memory_to_free: int = 0, force_patch_weights: bool = False
|
||||
) -> int:
|
||||
return self._call_rpc(
|
||||
"partially_unload", device_to, memory_to_free, force_patch_weights
|
||||
)
|
||||
|
||||
def load(
|
||||
self,
|
||||
device_to: Any = None,
|
||||
lowvram_model_memory: int = 0,
|
||||
force_patch_weights: bool = False,
|
||||
full_load: bool = False,
|
||||
) -> None:
|
||||
self._call_rpc(
|
||||
"load", device_to, lowvram_model_memory, force_patch_weights, full_load
|
||||
)
|
||||
|
||||
def patch_model(
|
||||
self,
|
||||
device_to: Any = None,
|
||||
lowvram_model_memory: int = 0,
|
||||
load_weights: bool = True,
|
||||
force_patch_weights: bool = False,
|
||||
) -> Any:
|
||||
self._call_rpc(
|
||||
"patch_model",
|
||||
device_to,
|
||||
lowvram_model_memory,
|
||||
load_weights,
|
||||
force_patch_weights,
|
||||
)
|
||||
return self
|
||||
|
||||
def unpatch_model(
|
||||
self, device_to: Any = None, unpatch_weights: bool = True
|
||||
) -> None:
|
||||
self._call_rpc("unpatch_model", device_to, unpatch_weights)
|
||||
|
||||
def detach(self, unpatch_all: bool = True) -> Any:
|
||||
self._call_rpc("detach", unpatch_all)
|
||||
return self.model
|
||||
|
||||
def _cpu_tensor_bytes(self, obj: Any) -> int:
|
||||
import torch
|
||||
|
||||
if isinstance(obj, torch.Tensor):
|
||||
if obj.device.type == "cpu":
|
||||
return obj.nbytes
|
||||
return 0
|
||||
if isinstance(obj, dict):
|
||||
return sum(self._cpu_tensor_bytes(v) for v in obj.values())
|
||||
if isinstance(obj, (list, tuple)):
|
||||
return sum(self._cpu_tensor_bytes(v) for v in obj)
|
||||
return 0
|
||||
|
||||
def _ensure_apply_model_headroom(self, required_bytes: int) -> bool:
|
||||
if required_bytes <= 0:
|
||||
return True
|
||||
|
||||
import torch
|
||||
import comfy.model_management as model_management
|
||||
|
||||
target_raw = self.load_device
|
||||
try:
|
||||
if isinstance(target_raw, torch.device):
|
||||
target = target_raw
|
||||
elif isinstance(target_raw, str):
|
||||
target = torch.device(target_raw)
|
||||
elif isinstance(target_raw, int):
|
||||
target = torch.device(f"cuda:{target_raw}")
|
||||
else:
|
||||
target = torch.device(target_raw)
|
||||
except Exception:
|
||||
return True
|
||||
|
||||
if target.type != "cuda":
|
||||
return True
|
||||
|
||||
required = required_bytes + self._APPLY_MODEL_GUARD_PADDING_BYTES
|
||||
if model_management.get_free_memory(target) >= required:
|
||||
return True
|
||||
|
||||
model_management.cleanup_models_gc()
|
||||
model_management.cleanup_models()
|
||||
model_management.soft_empty_cache()
|
||||
|
||||
if model_management.get_free_memory(target) < required:
|
||||
model_management.free_memory(required, target, for_dynamic=True)
|
||||
model_management.soft_empty_cache()
|
||||
|
||||
if model_management.get_free_memory(target) < required:
|
||||
# Escalate to non-dynamic unloading before dispatching CUDA transfer.
|
||||
model_management.free_memory(required, target, for_dynamic=False)
|
||||
model_management.soft_empty_cache()
|
||||
|
||||
if model_management.get_free_memory(target) < required:
|
||||
model_management.load_models_gpu(
|
||||
[self],
|
||||
minimum_memory_required=required,
|
||||
)
|
||||
|
||||
return model_management.get_free_memory(target) >= required
|
||||
|
||||
def apply_model(self, *args, **kwargs) -> Any:
|
||||
import torch
|
||||
|
||||
def _preferred_device() -> Any:
|
||||
for value in args:
|
||||
if isinstance(value, torch.Tensor):
|
||||
return value.device
|
||||
for value in kwargs.values():
|
||||
if isinstance(value, torch.Tensor):
|
||||
return value.device
|
||||
return None
|
||||
|
||||
def _move_result_to_device(obj: Any, device: Any) -> Any:
|
||||
if device is None:
|
||||
return obj
|
||||
if isinstance(obj, torch.Tensor):
|
||||
return obj.to(device) if obj.device != device else obj
|
||||
if isinstance(obj, dict):
|
||||
return {k: _move_result_to_device(v, device) for k, v in obj.items()}
|
||||
if isinstance(obj, list):
|
||||
return [_move_result_to_device(v, device) for v in obj]
|
||||
if isinstance(obj, tuple):
|
||||
return tuple(_move_result_to_device(v, device) for v in obj)
|
||||
return obj
|
||||
|
||||
# DynamicVRAM models must keep load/offload decisions in host process.
|
||||
# Child-side CUDA staging here can deadlock before first inference RPC.
|
||||
if self.is_dynamic():
|
||||
out = self._call_rpc("inner_model_apply_model", args, kwargs)
|
||||
return _move_result_to_device(out, _preferred_device())
|
||||
|
||||
required_bytes = self._cpu_tensor_bytes(args) + self._cpu_tensor_bytes(kwargs)
|
||||
self._ensure_apply_model_headroom(required_bytes)
|
||||
target_device = self.load_device
|
||||
|
||||
def _to_cuda(obj: Any) -> Any:
|
||||
if isinstance(obj, torch.Tensor) and obj.device.type == "cpu":
|
||||
return obj.to(target_device)
|
||||
if isinstance(obj, dict):
|
||||
return {k: _to_cuda(v) for k, v in obj.items()}
|
||||
if isinstance(obj, list):
|
||||
return [_to_cuda(v) for v in obj]
|
||||
if isinstance(obj, tuple):
|
||||
return tuple(_to_cuda(v) for v in obj)
|
||||
return obj
|
||||
|
||||
try:
|
||||
args_cuda = _to_cuda(args)
|
||||
kwargs_cuda = _to_cuda(kwargs)
|
||||
except torch.OutOfMemoryError:
|
||||
self._ensure_apply_model_headroom(required_bytes)
|
||||
args_cuda = _to_cuda(args)
|
||||
kwargs_cuda = _to_cuda(kwargs)
|
||||
|
||||
out = self._call_rpc("inner_model_apply_model", args_cuda, kwargs_cuda)
|
||||
return _move_result_to_device(out, _preferred_device())
|
||||
|
||||
def model_state_dict(self, filter_prefix: Optional[str] = None) -> Any:
|
||||
keys = self._call_rpc("model_state_dict", filter_prefix)
|
||||
return dict.fromkeys(keys, None)
|
||||
|
||||
def add_patches(self, *args: Any, **kwargs: Any) -> Any:
|
||||
res = self._call_rpc("add_patches", *args, **kwargs)
|
||||
if isinstance(res, list):
|
||||
return [tuple(x) if isinstance(x, list) else x for x in res]
|
||||
return res
|
||||
|
||||
def get_key_patches(self, filter_prefix: Optional[str] = None) -> Any:
|
||||
return self._call_rpc("get_key_patches", filter_prefix)
|
||||
|
||||
def patch_weight_to_device(self, key, device_to=None, inplace_update=False):
|
||||
self._call_rpc("patch_weight_to_device", key, device_to, inplace_update)
|
||||
|
||||
def pin_weight_to_device(self, key):
|
||||
self._call_rpc("pin_weight_to_device", key)
|
||||
|
||||
def unpin_weight(self, key):
|
||||
self._call_rpc("unpin_weight", key)
|
||||
|
||||
def unpin_all_weights(self):
|
||||
self._call_rpc("unpin_all_weights")
|
||||
|
||||
def calculate_weight(self, patches, weight, key, intermediate_dtype=None):
|
||||
return self._call_rpc(
|
||||
"calculate_weight", patches, weight, key, intermediate_dtype
|
||||
)
|
||||
|
||||
def inject_model(self) -> None:
|
||||
self._call_rpc("inject_model")
|
||||
|
||||
def eject_model(self) -> None:
|
||||
self._call_rpc("eject_model")
|
||||
|
||||
def use_ejected(self, skip_and_inject_on_exit_only: bool = False) -> Any:
|
||||
return AutoPatcherEjector(
|
||||
self, skip_and_inject_on_exit_only=skip_and_inject_on_exit_only
|
||||
)
|
||||
|
||||
@property
|
||||
def is_injected(self) -> bool:
|
||||
return self._call_rpc("get_is_injected")
|
||||
|
||||
@property
|
||||
def skip_injection(self) -> bool:
|
||||
return self._call_rpc("get_skip_injection")
|
||||
|
||||
@skip_injection.setter
|
||||
def skip_injection(self, value: bool) -> None:
|
||||
self._call_rpc("set_skip_injection", value)
|
||||
|
||||
def clean_hooks(self) -> None:
|
||||
self._call_rpc("clean_hooks")
|
||||
|
||||
def pre_run(self) -> None:
|
||||
self._call_rpc("pre_run")
|
||||
|
||||
def cleanup(self) -> None:
|
||||
try:
|
||||
self._call_rpc("cleanup")
|
||||
except Exception:
|
||||
logger.debug(
|
||||
"ModelPatcherProxy cleanup RPC failed for %s",
|
||||
self._instance_id,
|
||||
exc_info=True,
|
||||
)
|
||||
finally:
|
||||
super().cleanup()
|
||||
|
||||
@property
|
||||
def model(self) -> _InnerModelProxy:
|
||||
return _InnerModelProxy(self)
|
||||
|
||||
def __getattr__(self, name: str) -> Any:
|
||||
_whitelisted_attrs = {
|
||||
"hook_patches_backup",
|
||||
"hook_backup",
|
||||
"cached_hook_patches",
|
||||
"current_hooks",
|
||||
"forced_hooks",
|
||||
"is_clip",
|
||||
"patches_uuid",
|
||||
"pinned",
|
||||
"attachments",
|
||||
"additional_models",
|
||||
"injections",
|
||||
"hook_patches",
|
||||
"model_lowvram",
|
||||
"model_loaded_weight_memory",
|
||||
"backup",
|
||||
"object_patches_backup",
|
||||
"weight_wrapper_patches",
|
||||
"weight_inplace_update",
|
||||
"force_cast_weights",
|
||||
}
|
||||
if name in _whitelisted_attrs:
|
||||
return self._call_rpc("get_patcher_attr", name)
|
||||
raise AttributeError(
|
||||
f"'{type(self).__name__}' object has no attribute '{name}'"
|
||||
)
|
||||
|
||||
def load_lora(
|
||||
self,
|
||||
lora_path: str,
|
||||
strength_model: float,
|
||||
clip: Optional[Any] = None,
|
||||
strength_clip: float = 1.0,
|
||||
) -> tuple:
|
||||
clip_id = None
|
||||
if clip is not None:
|
||||
clip_id = getattr(clip, "_instance_id", getattr(clip, "_clip_id", None))
|
||||
result = self._call_rpc(
|
||||
"load_lora", lora_path, strength_model, clip_id, strength_clip
|
||||
)
|
||||
new_model = None
|
||||
if result.get("model_id"):
|
||||
new_model = self._spawn_related_proxy(result["model_id"])
|
||||
new_clip = None
|
||||
if result.get("clip_id"):
|
||||
from comfy.isolation.clip_proxy import CLIPProxy
|
||||
|
||||
new_clip = CLIPProxy(result["clip_id"])
|
||||
return (new_model, new_clip)
|
||||
|
||||
@property
|
||||
def load_device(self) -> Any:
|
||||
return self._call_rpc("get_load_device")
|
||||
|
||||
@property
|
||||
def offload_device(self) -> Any:
|
||||
return self._call_rpc("get_offload_device")
|
||||
|
||||
@property
|
||||
def device(self) -> Any:
|
||||
return self.load_device
|
||||
|
||||
def current_loaded_device(self) -> Any:
|
||||
return self._call_rpc("current_loaded_device")
|
||||
|
||||
@property
|
||||
def size(self) -> int:
|
||||
return self._call_rpc("get_size")
|
||||
|
||||
def model_size(self) -> Any:
|
||||
return self._call_rpc("model_size")
|
||||
|
||||
def loaded_size(self) -> Any:
|
||||
return self._call_rpc("loaded_size")
|
||||
|
||||
def get_ram_usage(self) -> int:
|
||||
return self._call_rpc("get_ram_usage")
|
||||
|
||||
def lowvram_patch_counter(self) -> int:
|
||||
return self._call_rpc("lowvram_patch_counter")
|
||||
|
||||
def memory_required(self, input_shape: Any) -> Any:
|
||||
return self._call_rpc("memory_required", input_shape)
|
||||
|
||||
def get_operation_state(self) -> Dict[str, Any]:
|
||||
state = self._call_rpc("get_operation_state")
|
||||
return state if isinstance(state, dict) else {}
|
||||
|
||||
def wait_for_idle(self, timeout_ms: int = 0) -> bool:
|
||||
return bool(self._call_rpc("wait_for_idle", timeout_ms))
|
||||
|
||||
def is_dynamic(self) -> bool:
|
||||
return bool(self._call_rpc("is_dynamic"))
|
||||
|
||||
def get_free_memory(self, device: Any) -> Any:
|
||||
return self._call_rpc("get_free_memory", device)
|
||||
|
||||
def partially_unload_ram(self, ram_to_unload: int) -> Any:
|
||||
return self._call_rpc("partially_unload_ram", ram_to_unload)
|
||||
|
||||
def model_dtype(self) -> Any:
|
||||
res = self._call_rpc("model_dtype")
|
||||
if isinstance(res, str) and res.startswith("torch."):
|
||||
try:
|
||||
import torch
|
||||
|
||||
attr = res.split(".")[-1]
|
||||
if hasattr(torch, attr):
|
||||
return getattr(torch, attr)
|
||||
except ImportError:
|
||||
pass
|
||||
return res
|
||||
|
||||
@property
|
||||
def hook_mode(self) -> Any:
|
||||
return self._call_rpc("get_hook_mode")
|
||||
|
||||
@hook_mode.setter
|
||||
def hook_mode(self, value: Any) -> None:
|
||||
self._call_rpc("set_hook_mode", value)
|
||||
|
||||
def set_model_sampler_cfg_function(
|
||||
self, sampler_cfg_function: Any, disable_cfg1_optimization: bool = False
|
||||
) -> None:
|
||||
self._call_rpc(
|
||||
"set_model_sampler_cfg_function",
|
||||
sampler_cfg_function,
|
||||
disable_cfg1_optimization,
|
||||
)
|
||||
|
||||
def set_model_sampler_post_cfg_function(
|
||||
self, post_cfg_function: Any, disable_cfg1_optimization: bool = False
|
||||
) -> None:
|
||||
self._call_rpc(
|
||||
"set_model_sampler_post_cfg_function",
|
||||
post_cfg_function,
|
||||
disable_cfg1_optimization,
|
||||
)
|
||||
|
||||
def set_model_sampler_pre_cfg_function(
|
||||
self, pre_cfg_function: Any, disable_cfg1_optimization: bool = False
|
||||
) -> None:
|
||||
self._call_rpc(
|
||||
"set_model_sampler_pre_cfg_function",
|
||||
pre_cfg_function,
|
||||
disable_cfg1_optimization,
|
||||
)
|
||||
|
||||
def set_model_sampler_calc_cond_batch_function(self, fn: Any) -> None:
|
||||
self._call_rpc("set_model_sampler_calc_cond_batch_function", fn)
|
||||
|
||||
def set_model_unet_function_wrapper(self, unet_wrapper_function: Any) -> None:
|
||||
self._call_rpc("set_model_unet_function_wrapper", unet_wrapper_function)
|
||||
|
||||
def set_model_denoise_mask_function(self, denoise_mask_function: Any) -> None:
|
||||
self._call_rpc("set_model_denoise_mask_function", denoise_mask_function)
|
||||
|
||||
def set_model_patch(self, patch: Any, name: str) -> None:
|
||||
self._call_rpc("set_model_patch", patch, name)
|
||||
|
||||
def set_model_patch_replace(
|
||||
self,
|
||||
patch: Any,
|
||||
name: str,
|
||||
block_name: str,
|
||||
number: int,
|
||||
transformer_index: Optional[int] = None,
|
||||
) -> None:
|
||||
self._call_rpc(
|
||||
"set_model_patch_replace",
|
||||
patch,
|
||||
name,
|
||||
block_name,
|
||||
number,
|
||||
transformer_index,
|
||||
)
|
||||
|
||||
def set_model_attn1_patch(self, patch: Any) -> None:
|
||||
self.set_model_patch(patch, "attn1_patch")
|
||||
|
||||
def set_model_attn2_patch(self, patch: Any) -> None:
|
||||
self.set_model_patch(patch, "attn2_patch")
|
||||
|
||||
def set_model_attn1_replace(
|
||||
self,
|
||||
patch: Any,
|
||||
block_name: str,
|
||||
number: int,
|
||||
transformer_index: Optional[int] = None,
|
||||
) -> None:
|
||||
self.set_model_patch_replace(
|
||||
patch, "attn1", block_name, number, transformer_index
|
||||
)
|
||||
|
||||
def set_model_attn2_replace(
|
||||
self,
|
||||
patch: Any,
|
||||
block_name: str,
|
||||
number: int,
|
||||
transformer_index: Optional[int] = None,
|
||||
) -> None:
|
||||
self.set_model_patch_replace(
|
||||
patch, "attn2", block_name, number, transformer_index
|
||||
)
|
||||
|
||||
def set_model_attn1_output_patch(self, patch: Any) -> None:
|
||||
self.set_model_patch(patch, "attn1_output_patch")
|
||||
|
||||
def set_model_attn2_output_patch(self, patch: Any) -> None:
|
||||
self.set_model_patch(patch, "attn2_output_patch")
|
||||
|
||||
def set_model_input_block_patch(self, patch: Any) -> None:
|
||||
self.set_model_patch(patch, "input_block_patch")
|
||||
|
||||
def set_model_input_block_patch_after_skip(self, patch: Any) -> None:
|
||||
self.set_model_patch(patch, "input_block_patch_after_skip")
|
||||
|
||||
def set_model_output_block_patch(self, patch: Any) -> None:
|
||||
self.set_model_patch(patch, "output_block_patch")
|
||||
|
||||
def set_model_emb_patch(self, patch: Any) -> None:
|
||||
self.set_model_patch(patch, "emb_patch")
|
||||
|
||||
def set_model_forward_timestep_embed_patch(self, patch: Any) -> None:
|
||||
self.set_model_patch(patch, "forward_timestep_embed_patch")
|
||||
|
||||
def set_model_double_block_patch(self, patch: Any) -> None:
|
||||
self.set_model_patch(patch, "double_block")
|
||||
|
||||
def set_model_post_input_patch(self, patch: Any) -> None:
|
||||
self.set_model_patch(patch, "post_input")
|
||||
|
||||
def set_model_rope_options(
|
||||
self,
|
||||
scale_x=1.0,
|
||||
shift_x=0.0,
|
||||
scale_y=1.0,
|
||||
shift_y=0.0,
|
||||
scale_t=1.0,
|
||||
shift_t=0.0,
|
||||
**kwargs: Any,
|
||||
) -> None:
|
||||
options = {
|
||||
"scale_x": scale_x,
|
||||
"shift_x": shift_x,
|
||||
"scale_y": scale_y,
|
||||
"shift_y": shift_y,
|
||||
"scale_t": scale_t,
|
||||
"shift_t": shift_t,
|
||||
}
|
||||
options.update(kwargs)
|
||||
self._call_rpc("set_model_rope_options", options)
|
||||
|
||||
def set_model_compute_dtype(self, dtype: Any) -> None:
|
||||
self._call_rpc("set_model_compute_dtype", dtype)
|
||||
|
||||
def add_object_patch(self, name: str, obj: Any) -> None:
|
||||
self._call_rpc("add_object_patch", name, obj)
|
||||
|
||||
def add_weight_wrapper(self, name: str, function: Any) -> None:
|
||||
self._call_rpc("add_weight_wrapper", name, function)
|
||||
|
||||
def add_wrapper_with_key(self, wrapper_type: Any, key: str, fn: Any) -> None:
|
||||
self._call_rpc("add_wrapper_with_key", wrapper_type, key, fn)
|
||||
|
||||
def add_wrapper(self, wrapper_type: str, wrapper: Callable) -> None:
|
||||
self.add_wrapper_with_key(wrapper_type, None, wrapper)
|
||||
|
||||
def remove_wrappers_with_key(self, wrapper_type: str, key: str) -> None:
|
||||
self._call_rpc("remove_wrappers_with_key", wrapper_type, key)
|
||||
|
||||
@property
|
||||
def wrappers(self) -> Any:
|
||||
return self._call_rpc("get_wrappers")
|
||||
|
||||
def add_callback_with_key(self, call_type: str, key: str, callback: Any) -> None:
|
||||
self._call_rpc("add_callback_with_key", call_type, key, callback)
|
||||
|
||||
def add_callback(self, call_type: str, callback: Any) -> None:
|
||||
self.add_callback_with_key(call_type, None, callback)
|
||||
|
||||
def remove_callbacks_with_key(self, call_type: str, key: str) -> None:
|
||||
self._call_rpc("remove_callbacks_with_key", call_type, key)
|
||||
|
||||
@property
|
||||
def callbacks(self) -> Any:
|
||||
return self._call_rpc("get_callbacks")
|
||||
|
||||
def set_attachments(self, key: str, attachment: Any) -> None:
|
||||
self._call_rpc("set_attachments", key, attachment)
|
||||
|
||||
def get_attachment(self, key: str) -> Any:
|
||||
return self._call_rpc("get_attachment", key)
|
||||
|
||||
def remove_attachments(self, key: str) -> None:
|
||||
self._call_rpc("remove_attachments", key)
|
||||
|
||||
def set_injections(self, key: str, injections: Any) -> None:
|
||||
self._call_rpc("set_injections", key, injections)
|
||||
|
||||
def get_injections(self, key: str) -> Any:
|
||||
return self._call_rpc("get_injections", key)
|
||||
|
||||
def remove_injections(self, key: str) -> None:
|
||||
self._call_rpc("remove_injections", key)
|
||||
|
||||
def set_additional_models(self, key: str, models: Any) -> None:
|
||||
ids = [m._instance_id for m in models]
|
||||
self._call_rpc("set_additional_models", key, ids)
|
||||
|
||||
def remove_additional_models(self, key: str) -> None:
|
||||
self._call_rpc("remove_additional_models", key)
|
||||
|
||||
def get_nested_additional_models(self) -> Any:
|
||||
return self._call_rpc("get_nested_additional_models")
|
||||
|
||||
def get_additional_models(self) -> List[ModelPatcherProxy]:
|
||||
ids = self._call_rpc("get_additional_models")
|
||||
return [self._spawn_related_proxy(mid) for mid in ids]
|
||||
|
||||
def model_patches_models(self) -> Any:
|
||||
return self._call_rpc("model_patches_models")
|
||||
|
||||
@property
|
||||
def parent(self) -> Any:
|
||||
return self._call_rpc("get_parent")
|
||||
|
||||
def model_mmap_residency(self, free: bool = False) -> tuple:
|
||||
result = self._call_rpc("model_mmap_residency", free)
|
||||
if isinstance(result, list):
|
||||
return tuple(result)
|
||||
return result
|
||||
|
||||
def pinned_memory_size(self) -> int:
|
||||
return self._call_rpc("pinned_memory_size")
|
||||
|
||||
def get_non_dynamic_delegate(self) -> ModelPatcherProxy:
|
||||
new_id = self._call_rpc("get_non_dynamic_delegate")
|
||||
return self._spawn_related_proxy(new_id)
|
||||
|
||||
def disable_model_cfg1_optimization(self) -> None:
|
||||
self._call_rpc("disable_model_cfg1_optimization")
|
||||
|
||||
def set_model_noise_refiner_patch(self, patch: Any) -> None:
|
||||
self.set_model_patch(patch, "noise_refiner")
|
||||
|
||||
|
||||
class _InnerModelProxy:
|
||||
def __init__(self, parent: ModelPatcherProxy):
|
||||
self._parent = parent
|
||||
self._model_sampling = None
|
||||
|
||||
def __getattr__(self, name: str) -> Any:
|
||||
if name.startswith("_"):
|
||||
raise AttributeError(name)
|
||||
if name == "model_config":
|
||||
from types import SimpleNamespace
|
||||
|
||||
data = self._parent._call_rpc("get_inner_model_attr", name)
|
||||
if isinstance(data, dict):
|
||||
return SimpleNamespace(**data)
|
||||
return data
|
||||
if name in (
|
||||
"latent_format",
|
||||
"model_type",
|
||||
"current_weight_patches_uuid",
|
||||
):
|
||||
return self._parent._call_rpc("get_inner_model_attr", name)
|
||||
if name == "load_device":
|
||||
return self._parent._call_rpc("get_inner_model_attr", "load_device")
|
||||
if name == "device":
|
||||
return self._parent._call_rpc("get_inner_model_attr", "device")
|
||||
if name == "current_patcher":
|
||||
proxy = ModelPatcherProxy(
|
||||
self._parent._instance_id,
|
||||
self._parent._registry,
|
||||
manage_lifecycle=False,
|
||||
)
|
||||
if getattr(self._parent, "_rpc_caller", None) is not None:
|
||||
proxy._rpc_caller = self._parent._rpc_caller
|
||||
return proxy
|
||||
if name == "model_sampling":
|
||||
if self._model_sampling is None:
|
||||
self._model_sampling = self._parent._call_rpc(
|
||||
"get_model_object", "model_sampling"
|
||||
)
|
||||
return self._model_sampling
|
||||
if name == "extra_conds_shapes":
|
||||
return lambda *a, **k: self._parent._call_rpc(
|
||||
"inner_model_extra_conds_shapes", a, k
|
||||
)
|
||||
if name == "extra_conds":
|
||||
return lambda *a, **k: self._parent._call_rpc(
|
||||
"inner_model_extra_conds", a, k
|
||||
)
|
||||
if name == "memory_required":
|
||||
return lambda *a, **k: self._parent._call_rpc(
|
||||
"inner_model_memory_required", a, k
|
||||
)
|
||||
if name == "apply_model":
|
||||
# Delegate to parent's method to get the CPU->CUDA optimization
|
||||
return self._parent.apply_model
|
||||
if name == "process_latent_in":
|
||||
return lambda *a, **k: self._parent._call_rpc("process_latent_in", a, k)
|
||||
if name == "process_latent_out":
|
||||
return lambda *a, **k: self._parent._call_rpc("process_latent_out", a, k)
|
||||
if name == "scale_latent_inpaint":
|
||||
return lambda *a, **k: self._parent._call_rpc("scale_latent_inpaint", a, k)
|
||||
if name == "diffusion_model":
|
||||
return self._parent._call_rpc("get_inner_model_attr", "diffusion_model")
|
||||
if name == "state_dict":
|
||||
return lambda: self._parent.model_state_dict()
|
||||
raise AttributeError(f"'{name}' not supported on isolated InnerModel")
|
||||
File diff suppressed because it is too large
Load Diff
@@ -1,150 +0,0 @@
|
||||
# pylint: disable=import-outside-toplevel,logging-fstring-interpolation,protected-access
|
||||
# Isolation utilities and serializers for ModelPatcherProxy
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import os
|
||||
from typing import Any
|
||||
|
||||
from comfy.cli_args import args
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def maybe_wrap_model_for_isolation(model_patcher: Any) -> Any:
|
||||
from comfy.isolation.model_patcher_proxy_registry import ModelPatcherRegistry
|
||||
from comfy.isolation.model_patcher_proxy import ModelPatcherProxy
|
||||
|
||||
is_child = os.environ.get("PYISOLATE_CHILD") == "1"
|
||||
isolation_active = args.use_process_isolation or is_child
|
||||
|
||||
if not isolation_active:
|
||||
return model_patcher
|
||||
if is_child:
|
||||
return model_patcher
|
||||
if isinstance(model_patcher, ModelPatcherProxy):
|
||||
return model_patcher
|
||||
|
||||
registry = ModelPatcherRegistry()
|
||||
model_id = registry.register(model_patcher)
|
||||
logger.debug(f"Isolated ModelPatcher: {model_id}")
|
||||
return ModelPatcherProxy(model_id, registry, manage_lifecycle=True)
|
||||
|
||||
|
||||
def register_hooks_serializers(registry=None):
|
||||
from pyisolate._internal.serialization_registry import SerializerRegistry
|
||||
import comfy.hooks
|
||||
|
||||
if registry is None:
|
||||
registry = SerializerRegistry.get_instance()
|
||||
|
||||
def serialize_enum(obj):
|
||||
return {"__enum__": f"{type(obj).__name__}.{obj.name}"}
|
||||
|
||||
def deserialize_enum(data):
|
||||
cls_name, val_name = data["__enum__"].split(".")
|
||||
cls = getattr(comfy.hooks, cls_name)
|
||||
return cls[val_name]
|
||||
|
||||
registry.register("EnumHookType", serialize_enum, deserialize_enum)
|
||||
registry.register("EnumHookScope", serialize_enum, deserialize_enum)
|
||||
registry.register("EnumHookMode", serialize_enum, deserialize_enum)
|
||||
registry.register("EnumWeightTarget", serialize_enum, deserialize_enum)
|
||||
|
||||
def serialize_hook_group(obj):
|
||||
return {"__type__": "HookGroup", "hooks": obj.hooks}
|
||||
|
||||
def deserialize_hook_group(data):
|
||||
hg = comfy.hooks.HookGroup()
|
||||
for h in data["hooks"]:
|
||||
hg.add(h)
|
||||
return hg
|
||||
|
||||
registry.register("HookGroup", serialize_hook_group, deserialize_hook_group)
|
||||
|
||||
def serialize_dict_state(obj):
|
||||
d = obj.__dict__.copy()
|
||||
d["__type__"] = type(obj).__name__
|
||||
if "custom_should_register" in d:
|
||||
del d["custom_should_register"]
|
||||
return d
|
||||
|
||||
def deserialize_dict_state_generic(cls):
|
||||
def _deserialize(data):
|
||||
h = cls()
|
||||
h.__dict__.update(data)
|
||||
return h
|
||||
|
||||
return _deserialize
|
||||
|
||||
def deserialize_hook_keyframe(data):
|
||||
h = comfy.hooks.HookKeyframe(strength=data.get("strength", 1.0))
|
||||
h.__dict__.update(data)
|
||||
return h
|
||||
|
||||
registry.register("HookKeyframe", serialize_dict_state, deserialize_hook_keyframe)
|
||||
|
||||
def deserialize_hook_keyframe_group(data):
|
||||
h = comfy.hooks.HookKeyframeGroup()
|
||||
h.__dict__.update(data)
|
||||
return h
|
||||
|
||||
registry.register(
|
||||
"HookKeyframeGroup", serialize_dict_state, deserialize_hook_keyframe_group
|
||||
)
|
||||
|
||||
def deserialize_hook(data):
|
||||
h = comfy.hooks.Hook()
|
||||
h.__dict__.update(data)
|
||||
return h
|
||||
|
||||
registry.register("Hook", serialize_dict_state, deserialize_hook)
|
||||
|
||||
def deserialize_weight_hook(data):
|
||||
h = comfy.hooks.WeightHook()
|
||||
h.__dict__.update(data)
|
||||
return h
|
||||
|
||||
registry.register("WeightHook", serialize_dict_state, deserialize_weight_hook)
|
||||
|
||||
def serialize_set(obj):
|
||||
return {"__set__": list(obj)}
|
||||
|
||||
def deserialize_set(data):
|
||||
return set(data["__set__"])
|
||||
|
||||
registry.register("set", serialize_set, deserialize_set)
|
||||
|
||||
try:
|
||||
from comfy.weight_adapter.lora import LoRAAdapter
|
||||
|
||||
def serialize_lora(obj):
|
||||
return {"weights": {}, "loaded_keys": list(obj.loaded_keys)}
|
||||
|
||||
def deserialize_lora(data):
|
||||
return LoRAAdapter(set(data["loaded_keys"]), data["weights"])
|
||||
|
||||
registry.register("LoRAAdapter", serialize_lora, deserialize_lora)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
try:
|
||||
from comfy.hooks import _HookRef
|
||||
import uuid
|
||||
|
||||
def serialize_hook_ref(obj):
|
||||
return {
|
||||
"__hook_ref__": True,
|
||||
"id": getattr(obj, "_pyisolate_id", str(uuid.uuid4())),
|
||||
}
|
||||
|
||||
def deserialize_hook_ref(data):
|
||||
h = _HookRef()
|
||||
h._pyisolate_id = data.get("id", str(uuid.uuid4()))
|
||||
return h
|
||||
|
||||
registry.register("_HookRef", serialize_hook_ref, deserialize_hook_ref)
|
||||
except ImportError:
|
||||
pass
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to register _HookRef: {e}")
|
||||
@@ -1,360 +0,0 @@
|
||||
# pylint: disable=import-outside-toplevel
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
import os
|
||||
import threading
|
||||
import time
|
||||
from typing import Any
|
||||
|
||||
from comfy.isolation.proxies.base import (
|
||||
BaseProxy,
|
||||
BaseRegistry,
|
||||
detach_if_grad,
|
||||
get_thread_loop,
|
||||
run_coro_in_new_loop,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _describe_value(obj: Any) -> str:
|
||||
try:
|
||||
import torch
|
||||
except Exception:
|
||||
torch = None
|
||||
try:
|
||||
if torch is not None and isinstance(obj, torch.Tensor):
|
||||
return (
|
||||
"Tensor(shape=%s,dtype=%s,device=%s,id=%s)"
|
||||
% (tuple(obj.shape), obj.dtype, obj.device, id(obj))
|
||||
)
|
||||
except Exception:
|
||||
pass
|
||||
return "%s(id=%s)" % (type(obj).__name__, id(obj))
|
||||
|
||||
|
||||
def _prefer_device(*tensors: Any) -> Any:
|
||||
try:
|
||||
import torch
|
||||
except Exception:
|
||||
return None
|
||||
for t in tensors:
|
||||
if isinstance(t, torch.Tensor) and t.is_cuda:
|
||||
return t.device
|
||||
for t in tensors:
|
||||
if isinstance(t, torch.Tensor):
|
||||
return t.device
|
||||
return None
|
||||
|
||||
|
||||
def _to_device(obj: Any, device: Any) -> Any:
|
||||
try:
|
||||
import torch
|
||||
except Exception:
|
||||
return obj
|
||||
if device is None:
|
||||
return obj
|
||||
if isinstance(obj, torch.Tensor):
|
||||
if obj.device != device:
|
||||
return obj.to(device)
|
||||
return obj
|
||||
if isinstance(obj, (list, tuple)):
|
||||
converted = [_to_device(x, device) for x in obj]
|
||||
return type(obj)(converted) if isinstance(obj, tuple) else converted
|
||||
if isinstance(obj, dict):
|
||||
return {k: _to_device(v, device) for k, v in obj.items()}
|
||||
return obj
|
||||
|
||||
|
||||
def _to_cpu_for_rpc(obj: Any) -> Any:
|
||||
try:
|
||||
import torch
|
||||
except Exception:
|
||||
return obj
|
||||
if isinstance(obj, torch.Tensor):
|
||||
t = obj.detach() if obj.requires_grad else obj
|
||||
if t.is_cuda:
|
||||
return t.to("cpu")
|
||||
return t
|
||||
if isinstance(obj, (list, tuple)):
|
||||
converted = [_to_cpu_for_rpc(x) for x in obj]
|
||||
return type(obj)(converted) if isinstance(obj, tuple) else converted
|
||||
if isinstance(obj, dict):
|
||||
return {k: _to_cpu_for_rpc(v) for k, v in obj.items()}
|
||||
return obj
|
||||
|
||||
|
||||
class ModelSamplingRegistry(BaseRegistry[Any]):
|
||||
_type_prefix = "modelsampling"
|
||||
|
||||
async def calculate_input(self, instance_id: str, sigma: Any, noise: Any) -> Any:
|
||||
sampling = self._get_instance(instance_id)
|
||||
return detach_if_grad(sampling.calculate_input(sigma, noise))
|
||||
|
||||
async def calculate_denoised(
|
||||
self, instance_id: str, sigma: Any, model_output: Any, model_input: Any
|
||||
) -> Any:
|
||||
sampling = self._get_instance(instance_id)
|
||||
return detach_if_grad(
|
||||
sampling.calculate_denoised(sigma, model_output, model_input)
|
||||
)
|
||||
|
||||
async def noise_scaling(
|
||||
self,
|
||||
instance_id: str,
|
||||
sigma: Any,
|
||||
noise: Any,
|
||||
latent_image: Any,
|
||||
max_denoise: bool = False,
|
||||
) -> Any:
|
||||
sampling = self._get_instance(instance_id)
|
||||
return detach_if_grad(
|
||||
sampling.noise_scaling(sigma, noise, latent_image, max_denoise=max_denoise)
|
||||
)
|
||||
|
||||
async def inverse_noise_scaling(
|
||||
self, instance_id: str, sigma: Any, latent: Any
|
||||
) -> Any:
|
||||
sampling = self._get_instance(instance_id)
|
||||
return detach_if_grad(sampling.inverse_noise_scaling(sigma, latent))
|
||||
|
||||
async def timestep(self, instance_id: str, sigma: Any) -> Any:
|
||||
sampling = self._get_instance(instance_id)
|
||||
return sampling.timestep(sigma)
|
||||
|
||||
async def sigma(self, instance_id: str, timestep: Any) -> Any:
|
||||
sampling = self._get_instance(instance_id)
|
||||
return sampling.sigma(timestep)
|
||||
|
||||
async def percent_to_sigma(self, instance_id: str, percent: float) -> Any:
|
||||
sampling = self._get_instance(instance_id)
|
||||
return sampling.percent_to_sigma(percent)
|
||||
|
||||
async def get_sigma_min(self, instance_id: str) -> Any:
|
||||
sampling = self._get_instance(instance_id)
|
||||
return detach_if_grad(sampling.sigma_min)
|
||||
|
||||
async def get_sigma_max(self, instance_id: str) -> Any:
|
||||
sampling = self._get_instance(instance_id)
|
||||
return detach_if_grad(sampling.sigma_max)
|
||||
|
||||
async def get_sigma_data(self, instance_id: str) -> Any:
|
||||
sampling = self._get_instance(instance_id)
|
||||
return detach_if_grad(sampling.sigma_data)
|
||||
|
||||
async def get_sigmas(self, instance_id: str) -> Any:
|
||||
sampling = self._get_instance(instance_id)
|
||||
return detach_if_grad(sampling.sigmas)
|
||||
|
||||
async def set_sigmas(self, instance_id: str, sigmas: Any) -> None:
|
||||
sampling = self._get_instance(instance_id)
|
||||
sampling.set_sigmas(sigmas)
|
||||
|
||||
|
||||
class ModelSamplingProxy(BaseProxy[ModelSamplingRegistry]):
|
||||
_registry_class = ModelSamplingRegistry
|
||||
__module__ = "comfy.isolation.model_sampling_proxy"
|
||||
|
||||
def _get_rpc(self) -> Any:
|
||||
if self._rpc_caller is None:
|
||||
from pyisolate._internal.rpc_protocol import get_child_rpc_instance
|
||||
|
||||
rpc = get_child_rpc_instance()
|
||||
if rpc is not None:
|
||||
self._rpc_caller = rpc.create_caller(
|
||||
ModelSamplingRegistry, ModelSamplingRegistry.get_remote_id()
|
||||
)
|
||||
else:
|
||||
registry = ModelSamplingRegistry()
|
||||
|
||||
class _LocalCaller:
|
||||
def calculate_input(
|
||||
self, instance_id: str, sigma: Any, noise: Any
|
||||
) -> Any:
|
||||
return registry.calculate_input(instance_id, sigma, noise)
|
||||
|
||||
def calculate_denoised(
|
||||
self,
|
||||
instance_id: str,
|
||||
sigma: Any,
|
||||
model_output: Any,
|
||||
model_input: Any,
|
||||
) -> Any:
|
||||
return registry.calculate_denoised(
|
||||
instance_id, sigma, model_output, model_input
|
||||
)
|
||||
|
||||
def noise_scaling(
|
||||
self,
|
||||
instance_id: str,
|
||||
sigma: Any,
|
||||
noise: Any,
|
||||
latent_image: Any,
|
||||
max_denoise: bool = False,
|
||||
) -> Any:
|
||||
return registry.noise_scaling(
|
||||
instance_id, sigma, noise, latent_image, max_denoise
|
||||
)
|
||||
|
||||
def inverse_noise_scaling(
|
||||
self, instance_id: str, sigma: Any, latent: Any
|
||||
) -> Any:
|
||||
return registry.inverse_noise_scaling(
|
||||
instance_id, sigma, latent
|
||||
)
|
||||
|
||||
def timestep(self, instance_id: str, sigma: Any) -> Any:
|
||||
return registry.timestep(instance_id, sigma)
|
||||
|
||||
def sigma(self, instance_id: str, timestep: Any) -> Any:
|
||||
return registry.sigma(instance_id, timestep)
|
||||
|
||||
def percent_to_sigma(self, instance_id: str, percent: float) -> Any:
|
||||
return registry.percent_to_sigma(instance_id, percent)
|
||||
|
||||
def get_sigma_min(self, instance_id: str) -> Any:
|
||||
return registry.get_sigma_min(instance_id)
|
||||
|
||||
def get_sigma_max(self, instance_id: str) -> Any:
|
||||
return registry.get_sigma_max(instance_id)
|
||||
|
||||
def get_sigma_data(self, instance_id: str) -> Any:
|
||||
return registry.get_sigma_data(instance_id)
|
||||
|
||||
def get_sigmas(self, instance_id: str) -> Any:
|
||||
return registry.get_sigmas(instance_id)
|
||||
|
||||
def set_sigmas(self, instance_id: str, sigmas: Any) -> None:
|
||||
return registry.set_sigmas(instance_id, sigmas)
|
||||
|
||||
self._rpc_caller = _LocalCaller()
|
||||
return self._rpc_caller
|
||||
|
||||
def _call(self, method_name: str, *args: Any) -> Any:
|
||||
rpc = self._get_rpc()
|
||||
method = getattr(rpc, method_name)
|
||||
result = method(self._instance_id, *args)
|
||||
timeout_ms = self._rpc_timeout_ms()
|
||||
start_epoch = time.time()
|
||||
start_perf = time.perf_counter()
|
||||
thread_id = threading.get_ident()
|
||||
call_id = "%s:%s:%s:%.6f" % (
|
||||
self._instance_id,
|
||||
method_name,
|
||||
thread_id,
|
||||
start_perf,
|
||||
)
|
||||
logger.debug(
|
||||
"ISO:modelsampling_rpc_start method=%s instance_id=%s call_id=%s start_ts=%.6f thread=%s timeout_ms=%s",
|
||||
method_name,
|
||||
self._instance_id,
|
||||
call_id,
|
||||
start_epoch,
|
||||
thread_id,
|
||||
timeout_ms,
|
||||
)
|
||||
if asyncio.iscoroutine(result):
|
||||
result = asyncio.wait_for(result, timeout=timeout_ms / 1000.0)
|
||||
try:
|
||||
asyncio.get_running_loop()
|
||||
out = run_coro_in_new_loop(result)
|
||||
except RuntimeError:
|
||||
loop = get_thread_loop()
|
||||
out = loop.run_until_complete(result)
|
||||
else:
|
||||
out = result
|
||||
logger.debug(
|
||||
"ISO:modelsampling_rpc_after_await method=%s instance_id=%s call_id=%s out=%s",
|
||||
method_name,
|
||||
self._instance_id,
|
||||
call_id,
|
||||
_describe_value(out),
|
||||
)
|
||||
elapsed_ms = (time.perf_counter() - start_perf) * 1000.0
|
||||
logger.debug(
|
||||
"ISO:modelsampling_rpc_end method=%s instance_id=%s call_id=%s elapsed_ms=%.3f thread=%s",
|
||||
method_name,
|
||||
self._instance_id,
|
||||
call_id,
|
||||
elapsed_ms,
|
||||
thread_id,
|
||||
)
|
||||
logger.debug(
|
||||
"ISO:modelsampling_rpc_return method=%s instance_id=%s call_id=%s",
|
||||
method_name,
|
||||
self._instance_id,
|
||||
call_id,
|
||||
)
|
||||
return out
|
||||
|
||||
@staticmethod
|
||||
def _rpc_timeout_ms() -> int:
|
||||
raw = os.environ.get(
|
||||
"COMFY_ISOLATION_MODEL_SAMPLING_RPC_TIMEOUT_MS",
|
||||
os.environ.get("COMFY_ISOLATION_LOAD_RPC_TIMEOUT_MS", "30000"),
|
||||
)
|
||||
try:
|
||||
timeout_ms = int(raw)
|
||||
except ValueError:
|
||||
timeout_ms = 30000
|
||||
return max(1, timeout_ms)
|
||||
|
||||
@property
|
||||
def sigma_min(self) -> Any:
|
||||
return self._call("get_sigma_min")
|
||||
|
||||
@property
|
||||
def sigma_max(self) -> Any:
|
||||
return self._call("get_sigma_max")
|
||||
|
||||
@property
|
||||
def sigma_data(self) -> Any:
|
||||
return self._call("get_sigma_data")
|
||||
|
||||
@property
|
||||
def sigmas(self) -> Any:
|
||||
return self._call("get_sigmas")
|
||||
|
||||
def calculate_input(self, sigma: Any, noise: Any) -> Any:
|
||||
return self._call("calculate_input", sigma, noise)
|
||||
|
||||
def calculate_denoised(
|
||||
self, sigma: Any, model_output: Any, model_input: Any
|
||||
) -> Any:
|
||||
return self._call("calculate_denoised", sigma, model_output, model_input)
|
||||
|
||||
def noise_scaling(
|
||||
self, sigma: Any, noise: Any, latent_image: Any, max_denoise: bool = False
|
||||
) -> Any:
|
||||
preferred_device = _prefer_device(noise, latent_image)
|
||||
out = self._call(
|
||||
"noise_scaling",
|
||||
_to_cpu_for_rpc(sigma),
|
||||
_to_cpu_for_rpc(noise),
|
||||
_to_cpu_for_rpc(latent_image),
|
||||
max_denoise,
|
||||
)
|
||||
return _to_device(out, preferred_device)
|
||||
|
||||
def inverse_noise_scaling(self, sigma: Any, latent: Any) -> Any:
|
||||
preferred_device = _prefer_device(latent)
|
||||
out = self._call(
|
||||
"inverse_noise_scaling",
|
||||
_to_cpu_for_rpc(sigma),
|
||||
_to_cpu_for_rpc(latent),
|
||||
)
|
||||
return _to_device(out, preferred_device)
|
||||
|
||||
def timestep(self, sigma: Any) -> Any:
|
||||
return self._call("timestep", sigma)
|
||||
|
||||
def sigma(self, timestep: Any) -> Any:
|
||||
return self._call("sigma", timestep)
|
||||
|
||||
def percent_to_sigma(self, percent: float) -> Any:
|
||||
return self._call("percent_to_sigma", percent)
|
||||
|
||||
def set_sigmas(self, sigmas: Any) -> None:
|
||||
return self._call("set_sigmas", sigmas)
|
||||
@@ -1,17 +0,0 @@
|
||||
from .base import (
|
||||
IS_CHILD_PROCESS,
|
||||
BaseProxy,
|
||||
BaseRegistry,
|
||||
detach_if_grad,
|
||||
get_thread_loop,
|
||||
run_coro_in_new_loop,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"IS_CHILD_PROCESS",
|
||||
"BaseRegistry",
|
||||
"BaseProxy",
|
||||
"get_thread_loop",
|
||||
"run_coro_in_new_loop",
|
||||
"detach_if_grad",
|
||||
]
|
||||
@@ -1,301 +0,0 @@
|
||||
# pylint: disable=global-statement,import-outside-toplevel,protected-access
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import concurrent.futures
|
||||
import logging
|
||||
import os
|
||||
import threading
|
||||
import time
|
||||
import weakref
|
||||
from typing import Any, Callable, Dict, Generic, Optional, TypeVar
|
||||
|
||||
try:
|
||||
from pyisolate import ProxiedSingleton
|
||||
except ImportError:
|
||||
|
||||
class ProxiedSingleton: # type: ignore[no-redef]
|
||||
pass
|
||||
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
IS_CHILD_PROCESS = os.environ.get("PYISOLATE_CHILD") == "1"
|
||||
_thread_local = threading.local()
|
||||
T = TypeVar("T")
|
||||
|
||||
|
||||
def get_thread_loop() -> asyncio.AbstractEventLoop:
|
||||
loop = getattr(_thread_local, "loop", None)
|
||||
if loop is None or loop.is_closed():
|
||||
loop = asyncio.new_event_loop()
|
||||
_thread_local.loop = loop
|
||||
return loop
|
||||
|
||||
|
||||
def run_coro_in_new_loop(coro: Any) -> Any:
|
||||
result_box: Dict[str, Any] = {}
|
||||
exc_box: Dict[str, BaseException] = {}
|
||||
|
||||
def runner() -> None:
|
||||
loop = asyncio.new_event_loop()
|
||||
asyncio.set_event_loop(loop)
|
||||
try:
|
||||
result_box["value"] = loop.run_until_complete(coro)
|
||||
except Exception as exc: # noqa: BLE001
|
||||
exc_box["exc"] = exc
|
||||
finally:
|
||||
loop.close()
|
||||
|
||||
t = threading.Thread(target=runner, daemon=True)
|
||||
t.start()
|
||||
t.join()
|
||||
if "exc" in exc_box:
|
||||
raise exc_box["exc"]
|
||||
return result_box.get("value")
|
||||
|
||||
|
||||
def detach_if_grad(obj: Any) -> Any:
|
||||
try:
|
||||
import torch
|
||||
except Exception:
|
||||
return obj
|
||||
|
||||
if isinstance(obj, torch.Tensor):
|
||||
return obj.detach() if obj.requires_grad else obj
|
||||
if isinstance(obj, (list, tuple)):
|
||||
return type(obj)(detach_if_grad(x) for x in obj)
|
||||
if isinstance(obj, dict):
|
||||
return {k: detach_if_grad(v) for k, v in obj.items()}
|
||||
return obj
|
||||
|
||||
|
||||
class BaseRegistry(ProxiedSingleton, Generic[T]):
|
||||
_type_prefix: str = "base"
|
||||
|
||||
def __init__(self) -> None:
|
||||
if hasattr(ProxiedSingleton, "__init__") and ProxiedSingleton is not object:
|
||||
super().__init__()
|
||||
self._registry: Dict[str, T] = {}
|
||||
self._id_map: Dict[int, str] = {}
|
||||
self._counter = 0
|
||||
self._lock = threading.Lock()
|
||||
|
||||
def register(self, instance: T) -> str:
|
||||
with self._lock:
|
||||
obj_id = id(instance)
|
||||
if obj_id in self._id_map:
|
||||
return self._id_map[obj_id]
|
||||
instance_id = f"{self._type_prefix}_{self._counter}"
|
||||
self._counter += 1
|
||||
self._registry[instance_id] = instance
|
||||
self._id_map[obj_id] = instance_id
|
||||
return instance_id
|
||||
|
||||
def unregister_sync(self, instance_id: str) -> None:
|
||||
with self._lock:
|
||||
instance = self._registry.pop(instance_id, None)
|
||||
if instance:
|
||||
self._id_map.pop(id(instance), None)
|
||||
|
||||
def _get_instance(self, instance_id: str) -> T:
|
||||
if IS_CHILD_PROCESS:
|
||||
raise RuntimeError(
|
||||
f"[{self.__class__.__name__}] _get_instance called in child"
|
||||
)
|
||||
with self._lock:
|
||||
instance = self._registry.get(instance_id)
|
||||
if instance is None:
|
||||
raise ValueError(f"{instance_id} not found")
|
||||
return instance
|
||||
|
||||
|
||||
_GLOBAL_LOOP: Optional[asyncio.AbstractEventLoop] = None
|
||||
|
||||
|
||||
def set_global_loop(loop: asyncio.AbstractEventLoop) -> None:
|
||||
global _GLOBAL_LOOP
|
||||
_GLOBAL_LOOP = loop
|
||||
|
||||
|
||||
def run_sync_rpc_coro(coro: Any, timeout_ms: Optional[int] = None) -> Any:
|
||||
if timeout_ms is not None:
|
||||
coro = asyncio.wait_for(coro, timeout=timeout_ms / 1000.0)
|
||||
|
||||
try:
|
||||
if _GLOBAL_LOOP is not None and _GLOBAL_LOOP.is_running():
|
||||
try:
|
||||
curr_loop = asyncio.get_running_loop()
|
||||
if curr_loop is _GLOBAL_LOOP:
|
||||
pass
|
||||
except RuntimeError:
|
||||
future = asyncio.run_coroutine_threadsafe(coro, _GLOBAL_LOOP)
|
||||
return future.result(
|
||||
timeout=(timeout_ms / 1000.0) if timeout_ms is not None else None
|
||||
)
|
||||
|
||||
try:
|
||||
asyncio.get_running_loop()
|
||||
return run_coro_in_new_loop(coro)
|
||||
except RuntimeError:
|
||||
loop = get_thread_loop()
|
||||
return loop.run_until_complete(coro)
|
||||
except asyncio.TimeoutError as exc:
|
||||
raise TimeoutError(f"Isolation RPC timeout (timeout_ms={timeout_ms})") from exc
|
||||
except concurrent.futures.TimeoutError as exc:
|
||||
raise TimeoutError(f"Isolation RPC timeout (timeout_ms={timeout_ms})") from exc
|
||||
|
||||
|
||||
def call_singleton_rpc(
|
||||
caller: Any,
|
||||
method_name: str,
|
||||
*args: Any,
|
||||
timeout_ms: Optional[int] = None,
|
||||
**kwargs: Any,
|
||||
) -> Any:
|
||||
if caller is None:
|
||||
raise RuntimeError(f"No RPC caller available for {method_name}")
|
||||
method = getattr(caller, method_name)
|
||||
return run_sync_rpc_coro(method(*args, **kwargs), timeout_ms=timeout_ms)
|
||||
|
||||
|
||||
class BaseProxy(Generic[T]):
|
||||
_registry_class: type = BaseRegistry # type: ignore[type-arg]
|
||||
__module__: str = "comfy.isolation.proxies.base"
|
||||
_TIMEOUT_RPC_METHODS = frozenset(
|
||||
{
|
||||
"partially_load",
|
||||
"partially_unload",
|
||||
"load",
|
||||
"patch_model",
|
||||
"unpatch_model",
|
||||
"inner_model_apply_model",
|
||||
"memory_required",
|
||||
"model_dtype",
|
||||
"inner_model_memory_required",
|
||||
"inner_model_extra_conds_shapes",
|
||||
"inner_model_extra_conds",
|
||||
"process_latent_in",
|
||||
"process_latent_out",
|
||||
"scale_latent_inpaint",
|
||||
}
|
||||
)
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
instance_id: str,
|
||||
registry: Optional[Any] = None,
|
||||
manage_lifecycle: bool = False,
|
||||
) -> None:
|
||||
self._instance_id = instance_id
|
||||
self._rpc_caller: Optional[Any] = None
|
||||
self._registry = registry if registry is not None else self._registry_class()
|
||||
self._manage_lifecycle = manage_lifecycle
|
||||
self._cleaned_up = False
|
||||
if manage_lifecycle and not IS_CHILD_PROCESS:
|
||||
self._finalizer = weakref.finalize(
|
||||
self, self._registry.unregister_sync, instance_id
|
||||
)
|
||||
|
||||
def _get_rpc(self) -> Any:
|
||||
if self._rpc_caller is None:
|
||||
from pyisolate._internal.rpc_protocol import get_child_rpc_instance
|
||||
|
||||
rpc = get_child_rpc_instance()
|
||||
if rpc is None:
|
||||
raise RuntimeError(f"[{self.__class__.__name__}] No RPC in child")
|
||||
self._rpc_caller = rpc.create_caller(
|
||||
self._registry_class, self._registry_class.get_remote_id()
|
||||
)
|
||||
return self._rpc_caller
|
||||
|
||||
def _rpc_timeout_ms_for_method(self, method_name: str) -> Optional[int]:
|
||||
if method_name not in self._TIMEOUT_RPC_METHODS:
|
||||
return None
|
||||
try:
|
||||
timeout_ms = int(
|
||||
os.environ.get("COMFY_ISOLATION_LOAD_RPC_TIMEOUT_MS", "120000")
|
||||
)
|
||||
except ValueError:
|
||||
timeout_ms = 120000
|
||||
return max(1, timeout_ms)
|
||||
|
||||
def _call_rpc(self, method_name: str, *args: Any, **kwargs: Any) -> Any:
|
||||
rpc = self._get_rpc()
|
||||
method = getattr(rpc, method_name)
|
||||
timeout_ms = self._rpc_timeout_ms_for_method(method_name)
|
||||
coro = method(self._instance_id, *args, **kwargs)
|
||||
if timeout_ms is not None:
|
||||
coro = asyncio.wait_for(coro, timeout=timeout_ms / 1000.0)
|
||||
|
||||
start_epoch = time.time()
|
||||
start_perf = time.perf_counter()
|
||||
thread_id = threading.get_ident()
|
||||
try:
|
||||
running_loop = asyncio.get_running_loop()
|
||||
loop_id: Optional[int] = id(running_loop)
|
||||
except RuntimeError:
|
||||
loop_id = None
|
||||
logger.debug(
|
||||
"ISO:rpc_start proxy=%s method=%s instance_id=%s start_ts=%.6f "
|
||||
"thread=%s loop=%s timeout_ms=%s",
|
||||
self.__class__.__name__,
|
||||
method_name,
|
||||
self._instance_id,
|
||||
start_epoch,
|
||||
thread_id,
|
||||
loop_id,
|
||||
timeout_ms,
|
||||
)
|
||||
|
||||
try:
|
||||
return run_sync_rpc_coro(coro, timeout_ms=timeout_ms)
|
||||
except TimeoutError as exc:
|
||||
raise TimeoutError(
|
||||
f"Isolation RPC timeout in {self.__class__.__name__}.{method_name} "
|
||||
f"(instance_id={self._instance_id}, timeout_ms={timeout_ms})"
|
||||
) from exc
|
||||
finally:
|
||||
end_epoch = time.time()
|
||||
elapsed_ms = (time.perf_counter() - start_perf) * 1000.0
|
||||
logger.debug(
|
||||
"ISO:rpc_end proxy=%s method=%s instance_id=%s end_ts=%.6f "
|
||||
"elapsed_ms=%.3f thread=%s loop=%s",
|
||||
self.__class__.__name__,
|
||||
method_name,
|
||||
self._instance_id,
|
||||
end_epoch,
|
||||
elapsed_ms,
|
||||
thread_id,
|
||||
loop_id,
|
||||
)
|
||||
|
||||
def __getstate__(self) -> Dict[str, Any]:
|
||||
return {"_instance_id": self._instance_id}
|
||||
|
||||
def __setstate__(self, state: Dict[str, Any]) -> None:
|
||||
self._instance_id = state["_instance_id"]
|
||||
self._rpc_caller = None
|
||||
self._registry = self._registry_class()
|
||||
self._manage_lifecycle = False
|
||||
self._cleaned_up = False
|
||||
|
||||
def cleanup(self) -> None:
|
||||
if self._cleaned_up or IS_CHILD_PROCESS:
|
||||
return
|
||||
self._cleaned_up = True
|
||||
finalizer = getattr(self, "_finalizer", None)
|
||||
if finalizer is not None:
|
||||
finalizer.detach()
|
||||
self._registry.unregister_sync(self._instance_id)
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return f"<{self.__class__.__name__} {self._instance_id}>"
|
||||
|
||||
|
||||
def create_rpc_method(method_name: str) -> Callable[..., Any]:
|
||||
def method(self: BaseProxy[Any], *args: Any, **kwargs: Any) -> Any:
|
||||
return self._call_rpc(method_name, *args, **kwargs)
|
||||
|
||||
method.__name__ = method_name
|
||||
return method
|
||||
@@ -1,206 +0,0 @@
|
||||
from __future__ import annotations
|
||||
import logging
|
||||
import os
|
||||
|
||||
from typing import Any, Dict, Optional
|
||||
|
||||
from pyisolate import ProxiedSingleton
|
||||
|
||||
from .base import call_singleton_rpc
|
||||
|
||||
_fp_logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _folder_paths():
|
||||
import folder_paths
|
||||
|
||||
return folder_paths
|
||||
|
||||
|
||||
def _is_child_process() -> bool:
|
||||
return os.environ.get("PYISOLATE_CHILD") == "1"
|
||||
|
||||
|
||||
def _serialize_folder_names_and_paths(data: dict[str, tuple[list[str], set[str]]]) -> dict[str, dict[str, list[str]]]:
|
||||
return {
|
||||
key: {"paths": list(paths), "extensions": sorted(list(extensions))}
|
||||
for key, (paths, extensions) in data.items()
|
||||
}
|
||||
|
||||
|
||||
def _deserialize_folder_names_and_paths(data: dict[str, dict[str, list[str]]]) -> dict[str, tuple[list[str], set[str]]]:
|
||||
return {
|
||||
key: (list(value.get("paths", [])), set(value.get("extensions", [])))
|
||||
for key, value in data.items()
|
||||
}
|
||||
|
||||
|
||||
class FolderPathsProxy(ProxiedSingleton):
|
||||
"""
|
||||
Dynamic proxy for folder_paths.
|
||||
Uses __getattr__ for most lookups, with explicit handling for
|
||||
mutable collections to ensure efficient by-value transfer.
|
||||
"""
|
||||
|
||||
_rpc: Optional[Any] = None
|
||||
|
||||
@classmethod
|
||||
def set_rpc(cls, rpc: Any) -> None:
|
||||
cls._rpc = rpc.create_caller(cls, cls.get_remote_id())
|
||||
|
||||
@classmethod
|
||||
def clear_rpc(cls) -> None:
|
||||
cls._rpc = None
|
||||
|
||||
@classmethod
|
||||
def _get_caller(cls) -> Any:
|
||||
if cls._rpc is None:
|
||||
raise RuntimeError("FolderPathsProxy RPC caller is not configured")
|
||||
return cls._rpc
|
||||
|
||||
def __getattr__(self, name):
|
||||
if _is_child_process():
|
||||
property_rpc = {
|
||||
"models_dir": "rpc_get_models_dir",
|
||||
"folder_names_and_paths": "rpc_get_folder_names_and_paths",
|
||||
"extension_mimetypes_cache": "rpc_get_extension_mimetypes_cache",
|
||||
"filename_list_cache": "rpc_get_filename_list_cache",
|
||||
}
|
||||
rpc_name = property_rpc.get(name)
|
||||
if rpc_name is not None:
|
||||
return call_singleton_rpc(self._get_caller(), rpc_name)
|
||||
raise AttributeError(name)
|
||||
return getattr(_folder_paths(), name)
|
||||
|
||||
@property
|
||||
def folder_names_and_paths(self) -> Dict:
|
||||
if _is_child_process():
|
||||
payload = call_singleton_rpc(self._get_caller(), "rpc_get_folder_names_and_paths")
|
||||
return _deserialize_folder_names_and_paths(payload)
|
||||
return _folder_paths().folder_names_and_paths
|
||||
|
||||
@property
|
||||
def extension_mimetypes_cache(self) -> Dict:
|
||||
if _is_child_process():
|
||||
return dict(call_singleton_rpc(self._get_caller(), "rpc_get_extension_mimetypes_cache"))
|
||||
return dict(_folder_paths().extension_mimetypes_cache)
|
||||
|
||||
@property
|
||||
def filename_list_cache(self) -> Dict:
|
||||
if _is_child_process():
|
||||
return dict(call_singleton_rpc(self._get_caller(), "rpc_get_filename_list_cache"))
|
||||
return dict(_folder_paths().filename_list_cache)
|
||||
|
||||
@property
|
||||
def models_dir(self) -> str:
|
||||
if _is_child_process():
|
||||
return str(call_singleton_rpc(self._get_caller(), "rpc_get_models_dir"))
|
||||
return _folder_paths().models_dir
|
||||
|
||||
def get_temp_directory(self) -> str:
|
||||
if _is_child_process():
|
||||
return call_singleton_rpc(self._get_caller(), "rpc_get_temp_directory")
|
||||
return _folder_paths().get_temp_directory()
|
||||
|
||||
def get_input_directory(self) -> str:
|
||||
if _is_child_process():
|
||||
return call_singleton_rpc(self._get_caller(), "rpc_get_input_directory")
|
||||
return _folder_paths().get_input_directory()
|
||||
|
||||
def get_output_directory(self) -> str:
|
||||
if _is_child_process():
|
||||
return call_singleton_rpc(self._get_caller(), "rpc_get_output_directory")
|
||||
return _folder_paths().get_output_directory()
|
||||
|
||||
def get_user_directory(self) -> str:
|
||||
if _is_child_process():
|
||||
return call_singleton_rpc(self._get_caller(), "rpc_get_user_directory")
|
||||
return _folder_paths().get_user_directory()
|
||||
|
||||
def get_annotated_filepath(self, name: str, default_dir: str | None = None) -> str:
|
||||
if _is_child_process():
|
||||
return call_singleton_rpc(
|
||||
self._get_caller(), "rpc_get_annotated_filepath", name, default_dir
|
||||
)
|
||||
return _folder_paths().get_annotated_filepath(name, default_dir)
|
||||
|
||||
def exists_annotated_filepath(self, name: str) -> bool:
|
||||
if _is_child_process():
|
||||
return bool(
|
||||
call_singleton_rpc(self._get_caller(), "rpc_exists_annotated_filepath", name)
|
||||
)
|
||||
return bool(_folder_paths().exists_annotated_filepath(name))
|
||||
|
||||
def add_model_folder_path(
|
||||
self, folder_name: str, full_folder_path: str, is_default: bool = False
|
||||
) -> None:
|
||||
if _is_child_process():
|
||||
call_singleton_rpc(
|
||||
self._get_caller(),
|
||||
"rpc_add_model_folder_path",
|
||||
folder_name,
|
||||
full_folder_path,
|
||||
is_default,
|
||||
)
|
||||
return None
|
||||
_folder_paths().add_model_folder_path(folder_name, full_folder_path, is_default)
|
||||
return None
|
||||
|
||||
def get_folder_paths(self, folder_name: str) -> list[str]:
|
||||
if _is_child_process():
|
||||
return list(call_singleton_rpc(self._get_caller(), "rpc_get_folder_paths", folder_name))
|
||||
return list(_folder_paths().get_folder_paths(folder_name))
|
||||
|
||||
def get_filename_list(self, folder_name: str) -> list[str]:
|
||||
if _is_child_process():
|
||||
return list(call_singleton_rpc(self._get_caller(), "rpc_get_filename_list", folder_name))
|
||||
return list(_folder_paths().get_filename_list(folder_name))
|
||||
|
||||
def get_full_path(self, folder_name: str, filename: str) -> str | None:
|
||||
if _is_child_process():
|
||||
return call_singleton_rpc(self._get_caller(), "rpc_get_full_path", folder_name, filename)
|
||||
return _folder_paths().get_full_path(folder_name, filename)
|
||||
|
||||
async def rpc_get_models_dir(self) -> str:
|
||||
return _folder_paths().models_dir
|
||||
|
||||
async def rpc_get_folder_names_and_paths(self) -> dict[str, dict[str, list[str]]]:
|
||||
return _serialize_folder_names_and_paths(_folder_paths().folder_names_and_paths)
|
||||
|
||||
async def rpc_get_extension_mimetypes_cache(self) -> dict[str, Any]:
|
||||
return dict(_folder_paths().extension_mimetypes_cache)
|
||||
|
||||
async def rpc_get_filename_list_cache(self) -> dict[str, Any]:
|
||||
return dict(_folder_paths().filename_list_cache)
|
||||
|
||||
async def rpc_get_temp_directory(self) -> str:
|
||||
return _folder_paths().get_temp_directory()
|
||||
|
||||
async def rpc_get_input_directory(self) -> str:
|
||||
return _folder_paths().get_input_directory()
|
||||
|
||||
async def rpc_get_output_directory(self) -> str:
|
||||
return _folder_paths().get_output_directory()
|
||||
|
||||
async def rpc_get_user_directory(self) -> str:
|
||||
return _folder_paths().get_user_directory()
|
||||
|
||||
async def rpc_get_annotated_filepath(self, name: str, default_dir: str | None = None) -> str:
|
||||
return _folder_paths().get_annotated_filepath(name, default_dir)
|
||||
|
||||
async def rpc_exists_annotated_filepath(self, name: str) -> bool:
|
||||
return _folder_paths().exists_annotated_filepath(name)
|
||||
|
||||
async def rpc_add_model_folder_path(
|
||||
self, folder_name: str, full_folder_path: str, is_default: bool = False
|
||||
) -> None:
|
||||
_folder_paths().add_model_folder_path(folder_name, full_folder_path, is_default)
|
||||
|
||||
async def rpc_get_folder_paths(self, folder_name: str) -> list[str]:
|
||||
return _folder_paths().get_folder_paths(folder_name)
|
||||
|
||||
async def rpc_get_filename_list(self, folder_name: str) -> list[str]:
|
||||
return _folder_paths().get_filename_list(folder_name)
|
||||
|
||||
async def rpc_get_full_path(self, folder_name: str, filename: str) -> str | None:
|
||||
return _folder_paths().get_full_path(folder_name, filename)
|
||||
@@ -1,158 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from typing import Any, Dict, Optional
|
||||
|
||||
from pyisolate import ProxiedSingleton
|
||||
|
||||
from .base import call_singleton_rpc
|
||||
|
||||
|
||||
class AnyTypeProxy(str):
|
||||
"""Replacement for custom AnyType objects used by some nodes."""
|
||||
|
||||
def __new__(cls, value: str = "*"):
|
||||
return super().__new__(cls, value)
|
||||
|
||||
def __ne__(self, other): # type: ignore[override]
|
||||
return False
|
||||
|
||||
|
||||
class FlexibleOptionalInputProxy(dict):
|
||||
"""Replacement for FlexibleOptionalInputType to allow dynamic inputs."""
|
||||
|
||||
def __init__(self, flex_type, data: Optional[Dict[str, object]] = None):
|
||||
super().__init__()
|
||||
self.type = flex_type
|
||||
if data:
|
||||
self.update(data)
|
||||
|
||||
def __getitem__(self, key): # type: ignore[override]
|
||||
return (self.type,)
|
||||
|
||||
def __contains__(self, key): # type: ignore[override]
|
||||
return True
|
||||
|
||||
|
||||
class ByPassTypeTupleProxy(tuple):
|
||||
"""Replacement for ByPassTypeTuple to mirror wildcard fallback behavior."""
|
||||
|
||||
def __new__(cls, values):
|
||||
return super().__new__(cls, values)
|
||||
|
||||
def __getitem__(self, index): # type: ignore[override]
|
||||
if index >= len(self):
|
||||
return AnyTypeProxy("*")
|
||||
return super().__getitem__(index)
|
||||
|
||||
|
||||
def _restore_special_value(value: Any) -> Any:
|
||||
if isinstance(value, dict):
|
||||
if value.get("__pyisolate_any_type__"):
|
||||
return AnyTypeProxy(value.get("value", "*"))
|
||||
if value.get("__pyisolate_flexible_optional__"):
|
||||
flex_type = _restore_special_value(value.get("type"))
|
||||
data_raw = value.get("data")
|
||||
data = (
|
||||
{k: _restore_special_value(v) for k, v in data_raw.items()}
|
||||
if isinstance(data_raw, dict)
|
||||
else {}
|
||||
)
|
||||
return FlexibleOptionalInputProxy(flex_type, data)
|
||||
if value.get("__pyisolate_tuple__") is not None:
|
||||
return tuple(
|
||||
_restore_special_value(v) for v in value["__pyisolate_tuple__"]
|
||||
)
|
||||
if value.get("__pyisolate_bypass_tuple__") is not None:
|
||||
return ByPassTypeTupleProxy(
|
||||
tuple(
|
||||
_restore_special_value(v)
|
||||
for v in value["__pyisolate_bypass_tuple__"]
|
||||
)
|
||||
)
|
||||
return {k: _restore_special_value(v) for k, v in value.items()}
|
||||
if isinstance(value, list):
|
||||
return [_restore_special_value(v) for v in value]
|
||||
return value
|
||||
|
||||
|
||||
def _serialize_special_value(value: Any) -> Any:
|
||||
if isinstance(value, AnyTypeProxy):
|
||||
return {"__pyisolate_any_type__": True, "value": str(value)}
|
||||
if isinstance(value, FlexibleOptionalInputProxy):
|
||||
return {
|
||||
"__pyisolate_flexible_optional__": True,
|
||||
"type": _serialize_special_value(value.type),
|
||||
"data": {k: _serialize_special_value(v) for k, v in value.items()},
|
||||
}
|
||||
if isinstance(value, ByPassTypeTupleProxy):
|
||||
return {
|
||||
"__pyisolate_bypass_tuple__": [_serialize_special_value(v) for v in value]
|
||||
}
|
||||
if isinstance(value, tuple):
|
||||
return {"__pyisolate_tuple__": [_serialize_special_value(v) for v in value]}
|
||||
if isinstance(value, list):
|
||||
return [_serialize_special_value(v) for v in value]
|
||||
if isinstance(value, dict):
|
||||
return {k: _serialize_special_value(v) for k, v in value.items()}
|
||||
return value
|
||||
|
||||
|
||||
def _restore_input_types_local(raw: Dict[str, object]) -> Dict[str, object]:
|
||||
if not isinstance(raw, dict):
|
||||
return raw # type: ignore[return-value]
|
||||
|
||||
restored: Dict[str, object] = {}
|
||||
for section, entries in raw.items():
|
||||
if isinstance(entries, dict) and entries.get("__pyisolate_flexible_optional__"):
|
||||
restored[section] = _restore_special_value(entries)
|
||||
elif isinstance(entries, dict):
|
||||
restored[section] = {
|
||||
k: _restore_special_value(v) for k, v in entries.items()
|
||||
}
|
||||
else:
|
||||
restored[section] = _restore_special_value(entries)
|
||||
return restored
|
||||
|
||||
|
||||
class HelperProxiesService(ProxiedSingleton):
|
||||
_rpc: Optional[Any] = None
|
||||
|
||||
@classmethod
|
||||
def set_rpc(cls, rpc: Any) -> None:
|
||||
cls._rpc = rpc.create_caller(cls, cls.get_remote_id())
|
||||
|
||||
@classmethod
|
||||
def clear_rpc(cls) -> None:
|
||||
cls._rpc = None
|
||||
|
||||
@classmethod
|
||||
def _get_caller(cls) -> Any:
|
||||
if cls._rpc is None:
|
||||
raise RuntimeError("HelperProxiesService RPC caller is not configured")
|
||||
return cls._rpc
|
||||
|
||||
async def rpc_restore_input_types(self, raw: Dict[str, object]) -> Dict[str, object]:
|
||||
restored = _restore_input_types_local(raw)
|
||||
return _serialize_special_value(restored)
|
||||
|
||||
|
||||
def restore_input_types(raw: Dict[str, object]) -> Dict[str, object]:
|
||||
"""Restore serialized INPUT_TYPES payload back into ComfyUI-compatible objects."""
|
||||
if os.environ.get("PYISOLATE_CHILD") == "1":
|
||||
payload = call_singleton_rpc(
|
||||
HelperProxiesService._get_caller(),
|
||||
"rpc_restore_input_types",
|
||||
raw,
|
||||
)
|
||||
return _restore_input_types_local(payload)
|
||||
return _restore_input_types_local(raw)
|
||||
|
||||
|
||||
__all__ = [
|
||||
"AnyTypeProxy",
|
||||
"FlexibleOptionalInputProxy",
|
||||
"ByPassTypeTupleProxy",
|
||||
"HelperProxiesService",
|
||||
"restore_input_types",
|
||||
]
|
||||
@@ -1,142 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from typing import Any, Optional
|
||||
|
||||
from pyisolate import ProxiedSingleton
|
||||
|
||||
from .base import call_singleton_rpc
|
||||
|
||||
|
||||
def _mm():
|
||||
import comfy.model_management
|
||||
|
||||
return comfy.model_management
|
||||
|
||||
|
||||
def _is_child_process() -> bool:
|
||||
return os.environ.get("PYISOLATE_CHILD") == "1"
|
||||
|
||||
|
||||
class TorchDeviceProxy:
|
||||
def __init__(self, device_str: str):
|
||||
self._device_str = device_str
|
||||
if ":" in device_str:
|
||||
device_type, index = device_str.split(":", 1)
|
||||
self.type = device_type
|
||||
self.index = int(index)
|
||||
else:
|
||||
self.type = device_str
|
||||
self.index = None
|
||||
|
||||
def __str__(self) -> str:
|
||||
return self._device_str
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return f"TorchDeviceProxy({self._device_str!r})"
|
||||
|
||||
|
||||
def _serialize_value(value: Any) -> Any:
|
||||
value_type = type(value)
|
||||
if value_type.__module__ == "torch" and value_type.__name__ == "device":
|
||||
return {"__pyisolate_torch_device__": str(value)}
|
||||
if isinstance(value, TorchDeviceProxy):
|
||||
return {"__pyisolate_torch_device__": str(value)}
|
||||
if isinstance(value, tuple):
|
||||
return {"__pyisolate_tuple__": [_serialize_value(item) for item in value]}
|
||||
if isinstance(value, list):
|
||||
return [_serialize_value(item) for item in value]
|
||||
if isinstance(value, dict):
|
||||
return {key: _serialize_value(inner) for key, inner in value.items()}
|
||||
return value
|
||||
|
||||
|
||||
def _deserialize_value(value: Any) -> Any:
|
||||
if isinstance(value, dict):
|
||||
if "__pyisolate_torch_device__" in value:
|
||||
return TorchDeviceProxy(value["__pyisolate_torch_device__"])
|
||||
if "__pyisolate_tuple__" in value:
|
||||
return tuple(_deserialize_value(item) for item in value["__pyisolate_tuple__"])
|
||||
return {key: _deserialize_value(inner) for key, inner in value.items()}
|
||||
if isinstance(value, list):
|
||||
return [_deserialize_value(item) for item in value]
|
||||
return value
|
||||
|
||||
|
||||
def _normalize_argument(value: Any) -> Any:
|
||||
if isinstance(value, TorchDeviceProxy):
|
||||
import torch
|
||||
|
||||
return torch.device(str(value))
|
||||
if isinstance(value, dict):
|
||||
if "__pyisolate_torch_device__" in value:
|
||||
import torch
|
||||
|
||||
return torch.device(value["__pyisolate_torch_device__"])
|
||||
if "__pyisolate_tuple__" in value:
|
||||
return tuple(_normalize_argument(item) for item in value["__pyisolate_tuple__"])
|
||||
return {key: _normalize_argument(inner) for key, inner in value.items()}
|
||||
if isinstance(value, list):
|
||||
return [_normalize_argument(item) for item in value]
|
||||
return value
|
||||
|
||||
|
||||
class ModelManagementProxy(ProxiedSingleton):
|
||||
"""
|
||||
Exact-relay proxy for comfy.model_management.
|
||||
Child calls never import comfy.model_management directly; they serialize
|
||||
arguments, relay to host, and deserialize the host result back.
|
||||
"""
|
||||
|
||||
_rpc: Optional[Any] = None
|
||||
|
||||
@classmethod
|
||||
def set_rpc(cls, rpc: Any) -> None:
|
||||
cls._rpc = rpc.create_caller(cls, cls.get_remote_id())
|
||||
|
||||
@classmethod
|
||||
def clear_rpc(cls) -> None:
|
||||
cls._rpc = None
|
||||
|
||||
@classmethod
|
||||
def _get_caller(cls) -> Any:
|
||||
if cls._rpc is None:
|
||||
raise RuntimeError("ModelManagementProxy RPC caller is not configured")
|
||||
return cls._rpc
|
||||
|
||||
def _relay_call(self, method_name: str, *args: Any, **kwargs: Any) -> Any:
|
||||
payload = call_singleton_rpc(
|
||||
self._get_caller(),
|
||||
"rpc_call",
|
||||
method_name,
|
||||
_serialize_value(args),
|
||||
_serialize_value(kwargs),
|
||||
)
|
||||
return _deserialize_value(payload)
|
||||
|
||||
@property
|
||||
def VRAMState(self):
|
||||
return _mm().VRAMState
|
||||
|
||||
@property
|
||||
def CPUState(self):
|
||||
return _mm().CPUState
|
||||
|
||||
@property
|
||||
def OOM_EXCEPTION(self):
|
||||
return _mm().OOM_EXCEPTION
|
||||
|
||||
def __getattr__(self, name: str):
|
||||
if _is_child_process():
|
||||
def child_method(*args: Any, **kwargs: Any) -> Any:
|
||||
return self._relay_call(name, *args, **kwargs)
|
||||
|
||||
return child_method
|
||||
return getattr(_mm(), name)
|
||||
|
||||
async def rpc_call(self, method_name: str, args: Any, kwargs: Any) -> Any:
|
||||
normalized_args = _normalize_argument(_deserialize_value(args))
|
||||
normalized_kwargs = _normalize_argument(_deserialize_value(kwargs))
|
||||
method = getattr(_mm(), method_name)
|
||||
result = method(*normalized_args, **normalized_kwargs)
|
||||
return _serialize_value(result)
|
||||
@@ -1,87 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import os
|
||||
from typing import Any, Optional
|
||||
|
||||
try:
|
||||
from pyisolate import ProxiedSingleton
|
||||
except ImportError:
|
||||
|
||||
class ProxiedSingleton:
|
||||
pass
|
||||
|
||||
from .base import call_singleton_rpc
|
||||
|
||||
|
||||
def _get_progress_state():
|
||||
from comfy_execution.progress import get_progress_state
|
||||
|
||||
return get_progress_state()
|
||||
|
||||
|
||||
def _is_child_process() -> bool:
|
||||
return os.environ.get("PYISOLATE_CHILD") == "1"
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class ProgressProxy(ProxiedSingleton):
|
||||
_rpc: Optional[Any] = None
|
||||
|
||||
@classmethod
|
||||
def set_rpc(cls, rpc: Any) -> None:
|
||||
cls._rpc = rpc.create_caller(cls, cls.get_remote_id())
|
||||
|
||||
@classmethod
|
||||
def clear_rpc(cls) -> None:
|
||||
cls._rpc = None
|
||||
|
||||
@classmethod
|
||||
def _get_caller(cls) -> Any:
|
||||
if cls._rpc is None:
|
||||
raise RuntimeError("ProgressProxy RPC caller is not configured")
|
||||
return cls._rpc
|
||||
|
||||
def set_progress(
|
||||
self,
|
||||
value: float,
|
||||
max_value: float,
|
||||
node_id: Optional[str] = None,
|
||||
image: Any = None,
|
||||
) -> None:
|
||||
if _is_child_process():
|
||||
call_singleton_rpc(
|
||||
self._get_caller(),
|
||||
"rpc_set_progress",
|
||||
value,
|
||||
max_value,
|
||||
node_id,
|
||||
image,
|
||||
)
|
||||
return None
|
||||
|
||||
_get_progress_state().update_progress(
|
||||
node_id=node_id,
|
||||
value=value,
|
||||
max_value=max_value,
|
||||
image=image,
|
||||
)
|
||||
return None
|
||||
|
||||
async def rpc_set_progress(
|
||||
self,
|
||||
value: float,
|
||||
max_value: float,
|
||||
node_id: Optional[str] = None,
|
||||
image: Any = None,
|
||||
) -> None:
|
||||
_get_progress_state().update_progress(
|
||||
node_id=node_id,
|
||||
value=value,
|
||||
max_value=max_value,
|
||||
image=image,
|
||||
)
|
||||
|
||||
|
||||
__all__ = ["ProgressProxy"]
|
||||
@@ -1,306 +0,0 @@
|
||||
# pylint: disable=import-outside-toplevel,logging-fstring-interpolation,redefined-outer-name,reimported,super-init-not-called
|
||||
"""Stateless RPC Implementation for PromptServer.
|
||||
|
||||
Replaces the legacy PromptServerProxy (Singleton) with a clean Service/Stub architecture.
|
||||
- Host: PromptServerService (RPC Handler)
|
||||
- Child: PromptServerStub (Interface Implementation)
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from typing import Any, Dict, Optional, Callable
|
||||
|
||||
import logging
|
||||
|
||||
# IMPORTS
|
||||
from pyisolate import ProxiedSingleton
|
||||
from .base import call_singleton_rpc
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
LOG_PREFIX = "[Isolation:C<->H]"
|
||||
|
||||
# ...
|
||||
|
||||
# =============================================================================
|
||||
# CHILD SIDE: PromptServerStub
|
||||
# =============================================================================
|
||||
|
||||
|
||||
class PromptServerStub:
|
||||
"""Stateless Stub for PromptServer."""
|
||||
|
||||
# Masquerade as the real server module
|
||||
__module__ = "server"
|
||||
|
||||
_instance: Optional["PromptServerStub"] = None
|
||||
_rpc: Optional[Any] = None # This will be the Caller object
|
||||
_source_file: Optional[str] = None
|
||||
|
||||
def __init__(self):
|
||||
self.routes = RouteStub(self)
|
||||
|
||||
@classmethod
|
||||
def set_rpc(cls, rpc: Any) -> None:
|
||||
"""Inject RPC client (called by adapter.py or manually)."""
|
||||
# Create caller for HOST Service
|
||||
# Assuming Host Service is registered as "PromptServerService" (class name)
|
||||
# We target the Host Service Class
|
||||
target_id = "PromptServerService"
|
||||
# We need to pass a class to create_caller? Usually yes.
|
||||
# But we don't have the Service class imported here necessarily (if running on child).
|
||||
# pyisolate check verify_service type?
|
||||
# If we pass PromptServerStub as the 'class', it might mismatch if checking types.
|
||||
# But we can try passing PromptServerStub if it mirrors the service name? No, stub is PromptServerStub.
|
||||
# We need a dummy class with right name?
|
||||
# Or just rely on string ID if create_caller supports it?
|
||||
# Standard: rpc.create_caller(PromptServerStub, target_id)
|
||||
# But wait, PromptServerStub is the *Local* class.
|
||||
# We want to call *Remote* class.
|
||||
# If we use PromptServerStub as the type, returning object will be typed as PromptServerStub?
|
||||
# The first arg is 'service_cls'.
|
||||
cls._rpc = rpc.create_caller(
|
||||
PromptServerService, target_id
|
||||
) # We import Service below?
|
||||
|
||||
@classmethod
|
||||
def clear_rpc(cls) -> None:
|
||||
cls._rpc = None
|
||||
|
||||
# We need PromptServerService available for the create_caller call?
|
||||
# Or just use the Stub class if ID matches?
|
||||
# prompt_server_impl.py defines BOTH. So PromptServerService IS available!
|
||||
|
||||
@property
|
||||
def instance(self) -> "PromptServerStub":
|
||||
return self
|
||||
|
||||
# ... Compatibility ...
|
||||
@classmethod
|
||||
def _get_source_file(cls) -> str:
|
||||
if cls._source_file is None:
|
||||
import folder_paths
|
||||
|
||||
cls._source_file = os.path.join(folder_paths.base_path, "server.py")
|
||||
return cls._source_file
|
||||
|
||||
@property
|
||||
def __file__(self) -> str:
|
||||
return self._get_source_file()
|
||||
|
||||
# --- Properties ---
|
||||
@property
|
||||
def client_id(self) -> Optional[str]:
|
||||
return "isolated_client"
|
||||
|
||||
@property
|
||||
def supports(self) -> set:
|
||||
return {"custom_nodes_from_web"}
|
||||
|
||||
@property
|
||||
def app(self):
|
||||
return _AppStub(self)
|
||||
|
||||
@property
|
||||
def prompt_queue(self):
|
||||
raise RuntimeError(
|
||||
"PromptServer.prompt_queue is not accessible in isolated nodes."
|
||||
)
|
||||
|
||||
# --- UI Communication (RPC Delegates) ---
|
||||
async def send_sync(
|
||||
self, event: str, data: Dict[str, Any], sid: Optional[str] = None
|
||||
) -> None:
|
||||
if self._rpc:
|
||||
await self._rpc.ui_send_sync(event, data, sid)
|
||||
|
||||
async def send(
|
||||
self, event: str, data: Dict[str, Any], sid: Optional[str] = None
|
||||
) -> None:
|
||||
if self._rpc:
|
||||
await self._rpc.ui_send(event, data, sid)
|
||||
|
||||
def send_progress_text(self, text: str, node_id: str, sid=None) -> None:
|
||||
if self._rpc:
|
||||
# Fire and forget likely needed. If method is async on host, caller invocation returns coroutine.
|
||||
# We must schedule it?
|
||||
# Or use fire_remote equivalent?
|
||||
# Caller object usually proxies calls. If host method is async, it returns coro.
|
||||
# If we are sync here (send_progress_text checks imply sync usage), we must background it.
|
||||
# But UtilsProxy hook wrapper creates task.
|
||||
# Does send_progress_text need to be sync? Yes, node code calls it sync.
|
||||
import asyncio
|
||||
|
||||
try:
|
||||
loop = asyncio.get_running_loop()
|
||||
loop.create_task(self._rpc.ui_send_progress_text(text, node_id, sid))
|
||||
except RuntimeError:
|
||||
call_singleton_rpc(self._rpc, "ui_send_progress_text", text, node_id, sid)
|
||||
|
||||
# --- Route Registration Logic ---
|
||||
_pending_child_routes: list = []
|
||||
|
||||
def register_route(self, method: str, path: str, handler: Callable):
|
||||
"""Buffer route registration. Routes are flushed via flush_child_routes()."""
|
||||
PromptServerStub._pending_child_routes.append((method, path, handler))
|
||||
logger.info("%s Buffered isolated route %s %s", LOG_PREFIX, method, path)
|
||||
|
||||
@classmethod
|
||||
async def flush_child_routes(cls):
|
||||
"""Send all buffered route registrations to host via RPC. Call from on_module_loaded."""
|
||||
if not cls._rpc:
|
||||
return 0
|
||||
flushed = 0
|
||||
for method, path, handler in cls._pending_child_routes:
|
||||
try:
|
||||
await cls._rpc.register_route_rpc(method, path, handler)
|
||||
flushed += 1
|
||||
except Exception as e:
|
||||
logger.error("%s Child route flush failed %s %s: %s", LOG_PREFIX, method, path, e)
|
||||
cls._pending_child_routes = []
|
||||
return flushed
|
||||
|
||||
|
||||
class RouteStub:
|
||||
"""Simulates aiohttp.web.RouteTableDef."""
|
||||
|
||||
def __init__(self, stub: PromptServerStub):
|
||||
self._stub = stub
|
||||
|
||||
def get(self, path: str):
|
||||
def decorator(handler):
|
||||
self._stub.register_route("GET", path, handler)
|
||||
return handler
|
||||
|
||||
return decorator
|
||||
|
||||
def post(self, path: str):
|
||||
def decorator(handler):
|
||||
self._stub.register_route("POST", path, handler)
|
||||
return handler
|
||||
|
||||
return decorator
|
||||
|
||||
def patch(self, path: str):
|
||||
def decorator(handler):
|
||||
self._stub.register_route("PATCH", path, handler)
|
||||
return handler
|
||||
|
||||
return decorator
|
||||
|
||||
def put(self, path: str):
|
||||
def decorator(handler):
|
||||
self._stub.register_route("PUT", path, handler)
|
||||
return handler
|
||||
|
||||
return decorator
|
||||
|
||||
def delete(self, path: str):
|
||||
def decorator(handler):
|
||||
self._stub.register_route("DELETE", path, handler)
|
||||
return handler
|
||||
|
||||
return decorator
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# HOST SIDE: PromptServerService
|
||||
# =============================================================================
|
||||
|
||||
|
||||
class PromptServerService(ProxiedSingleton):
|
||||
"""Host-side RPC Service for PromptServer."""
|
||||
|
||||
def __init__(self):
|
||||
pass
|
||||
|
||||
@property
|
||||
def server(self):
|
||||
from server import PromptServer
|
||||
|
||||
return PromptServer.instance
|
||||
|
||||
async def ui_send_sync(
|
||||
self, event: str, data: Dict[str, Any], sid: Optional[str] = None
|
||||
):
|
||||
await self.server.send_sync(event, data, sid)
|
||||
|
||||
async def ui_send(
|
||||
self, event: str, data: Dict[str, Any], sid: Optional[str] = None
|
||||
):
|
||||
await self.server.send(event, data, sid)
|
||||
|
||||
async def ui_send_progress_text(self, text: str, node_id: str, sid=None):
|
||||
# Made async to be awaitable by RPC layer
|
||||
self.server.send_progress_text(text, node_id, sid)
|
||||
|
||||
async def register_route_rpc(self, method: str, path: str, child_handler_proxy):
|
||||
"""RPC Target: Register a route that forwards to the Child."""
|
||||
from aiohttp import web
|
||||
logger.info("%s Registering isolated route %s %s", LOG_PREFIX, method, path)
|
||||
|
||||
async def route_wrapper(request: web.Request) -> web.Response:
|
||||
# 1. Capture request data
|
||||
req_data = {
|
||||
"method": request.method,
|
||||
"path": request.path,
|
||||
"query": dict(request.query),
|
||||
}
|
||||
if request.can_read_body:
|
||||
req_data["text"] = await request.text()
|
||||
|
||||
try:
|
||||
# 2. Call Child Handler via RPC (child_handler_proxy is async callable)
|
||||
result = await child_handler_proxy(req_data)
|
||||
|
||||
# 3. Serialize Response
|
||||
return self._serialize_response(result)
|
||||
except Exception as e:
|
||||
logger.error(f"{LOG_PREFIX} Isolated Route Error: {e}")
|
||||
return web.Response(status=500, text=str(e))
|
||||
|
||||
self.server.app.router.add_route(method, path, route_wrapper)
|
||||
logger.info("%s Registered isolated route %s %s", LOG_PREFIX, method, path)
|
||||
|
||||
def _serialize_response(self, result: Any) -> Any:
|
||||
"""Helper to convert Child result -> web.Response"""
|
||||
from aiohttp import web
|
||||
if isinstance(result, web.Response):
|
||||
return result
|
||||
# Handle dict (json)
|
||||
if isinstance(result, dict):
|
||||
return web.json_response(result)
|
||||
# Handle string
|
||||
if isinstance(result, str):
|
||||
return web.Response(text=result)
|
||||
# Fallback
|
||||
return web.Response(text=str(result))
|
||||
|
||||
|
||||
class _RouterStub:
|
||||
"""Captures router.add_route and router.add_static calls in isolation child."""
|
||||
|
||||
def __init__(self, stub):
|
||||
self._stub = stub
|
||||
|
||||
def add_route(self, method, path, handler, **kwargs):
|
||||
self._stub.register_route(method, path, handler)
|
||||
|
||||
def add_static(self, prefix, path, **kwargs):
|
||||
# Static file serving not supported in isolation — silently skip
|
||||
pass
|
||||
|
||||
|
||||
class _AppStub:
|
||||
"""Captures PromptServer.app access patterns in isolation child."""
|
||||
|
||||
def __init__(self, stub):
|
||||
self.router = _RouterStub(stub)
|
||||
self.frozen = False
|
||||
|
||||
def add_routes(self, routes):
|
||||
# aiohttp route table — iterate and register each
|
||||
for route in routes:
|
||||
if hasattr(route, "method") and hasattr(route, "handler"):
|
||||
self.router.add_route(route.method, route.path, route.handler)
|
||||
# StaticDef and other non-method routes — silently skip
|
||||
@@ -1,64 +0,0 @@
|
||||
# pylint: disable=cyclic-import,import-outside-toplevel
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Optional, Any
|
||||
from pyisolate import ProxiedSingleton
|
||||
|
||||
import os
|
||||
|
||||
|
||||
def _comfy_utils():
|
||||
import comfy.utils
|
||||
return comfy.utils
|
||||
|
||||
|
||||
class UtilsProxy(ProxiedSingleton):
|
||||
"""
|
||||
Proxy for comfy.utils.
|
||||
Primarily handles the PROGRESS_BAR_HOOK to ensure progress updates
|
||||
from isolated nodes reach the host.
|
||||
"""
|
||||
|
||||
# _instance and __new__ removed to rely on SingletonMetaclass
|
||||
_rpc: Optional[Any] = None
|
||||
|
||||
@classmethod
|
||||
def set_rpc(cls, rpc: Any) -> None:
|
||||
# Create caller using class name as ID (standard for Singletons)
|
||||
cls._rpc = rpc.create_caller(cls, "UtilsProxy")
|
||||
|
||||
@classmethod
|
||||
def clear_rpc(cls) -> None:
|
||||
cls._rpc = None
|
||||
|
||||
async def progress_bar_hook(
|
||||
self,
|
||||
value: int,
|
||||
total: int,
|
||||
preview: Optional[bytes] = None,
|
||||
node_id: Optional[str] = None,
|
||||
) -> Any:
|
||||
"""
|
||||
Host-side implementation: forwards the call to the real global hook.
|
||||
Child-side: this method call is intercepted by RPC and sent to host.
|
||||
"""
|
||||
if os.environ.get("PYISOLATE_CHILD") == "1":
|
||||
if UtilsProxy._rpc is None:
|
||||
raise RuntimeError("UtilsProxy RPC caller is not configured")
|
||||
return await UtilsProxy._rpc.progress_bar_hook(
|
||||
value, total, preview, node_id
|
||||
)
|
||||
|
||||
# Host Execution
|
||||
utils = _comfy_utils()
|
||||
if utils.PROGRESS_BAR_HOOK is not None:
|
||||
return utils.PROGRESS_BAR_HOOK(value, total, preview, node_id)
|
||||
return None
|
||||
|
||||
def set_progress_bar_global_hook(self, hook: Any) -> None:
|
||||
"""Forward hook registration (though usually not needed from child)."""
|
||||
if os.environ.get("PYISOLATE_CHILD") == "1":
|
||||
raise RuntimeError(
|
||||
"UtilsProxy.set_progress_bar_global_hook is not available in child without exact relay support"
|
||||
)
|
||||
_comfy_utils().set_progress_bar_global_hook(hook)
|
||||
@@ -1,229 +0,0 @@
|
||||
"""WebDirectoryProxy — serves isolated node web assets via RPC.
|
||||
|
||||
Child side: enumerates and reads files from the extension's web/ directory.
|
||||
Host side: gets an RPC proxy that fetches file listings and contents on demand.
|
||||
|
||||
Only files with allowed extensions (.js, .html, .css) are served.
|
||||
Directory traversal is rejected. File contents are base64-encoded for
|
||||
safe JSON-RPC transport.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
import binascii
|
||||
import logging
|
||||
import os
|
||||
from pathlib import Path, PurePosixPath
|
||||
from typing import Any, Dict, List
|
||||
|
||||
from pyisolate import ProxiedSingleton
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
ALLOWED_EXTENSIONS = frozenset({".js", ".html", ".css"})
|
||||
|
||||
MIME_TYPES = {
|
||||
".js": "application/javascript",
|
||||
".html": "text/html",
|
||||
".css": "text/css",
|
||||
}
|
||||
|
||||
|
||||
class WebDirectoryProxy(ProxiedSingleton):
|
||||
"""Proxy for serving isolated extension web directories.
|
||||
|
||||
On the child side, this class has direct filesystem access to the
|
||||
extension's web/ directory. On the host side, callers get an RPC
|
||||
proxy whose method calls are forwarded to the child.
|
||||
"""
|
||||
|
||||
# {extension_name: absolute_path_to_web_dir}
|
||||
_web_dirs: dict[str, str] = {}
|
||||
|
||||
@classmethod
|
||||
def register_web_dir(cls, extension_name: str, web_dir_path: str) -> None:
|
||||
"""Register an extension's web directory (child-side only)."""
|
||||
cls._web_dirs[extension_name] = web_dir_path
|
||||
logger.info(
|
||||
"][ WebDirectoryProxy: registered %s -> %s",
|
||||
extension_name,
|
||||
web_dir_path,
|
||||
)
|
||||
|
||||
def list_web_files(self, extension_name: str) -> List[Dict[str, str]]:
|
||||
"""Return a list of servable files in the extension's web directory.
|
||||
|
||||
Each entry is {"relative_path": "js/foo.js", "content_type": "application/javascript"}.
|
||||
Only files with allowed extensions are included.
|
||||
"""
|
||||
web_dir = self._web_dirs.get(extension_name)
|
||||
if not web_dir:
|
||||
return []
|
||||
|
||||
root = Path(web_dir)
|
||||
if not root.is_dir():
|
||||
return []
|
||||
|
||||
result: List[Dict[str, str]] = []
|
||||
for path in sorted(root.rglob("*")):
|
||||
if not path.is_file():
|
||||
continue
|
||||
ext = path.suffix.lower()
|
||||
if ext not in ALLOWED_EXTENSIONS:
|
||||
continue
|
||||
rel = path.relative_to(root)
|
||||
result.append({
|
||||
"relative_path": str(PurePosixPath(rel)),
|
||||
"content_type": MIME_TYPES[ext],
|
||||
})
|
||||
return result
|
||||
|
||||
def get_web_file(
|
||||
self, extension_name: str, relative_path: str
|
||||
) -> Dict[str, Any]:
|
||||
"""Return the contents of a single web file as base64.
|
||||
|
||||
Raises ValueError for traversal attempts or disallowed file types.
|
||||
Returns {"content": <base64 str>, "content_type": <MIME str>}.
|
||||
"""
|
||||
_validate_path(relative_path)
|
||||
|
||||
web_dir = self._web_dirs.get(extension_name)
|
||||
if not web_dir:
|
||||
raise FileNotFoundError(
|
||||
f"No web directory registered for {extension_name}"
|
||||
)
|
||||
|
||||
root = Path(web_dir).resolve()
|
||||
target = (root / relative_path).resolve()
|
||||
|
||||
# Ensure resolved path is under the web directory
|
||||
if os.path.commonpath([str(root), str(target)]) != str(root):
|
||||
raise ValueError(f"Path escapes web directory: {relative_path}")
|
||||
|
||||
if not target.is_file():
|
||||
raise FileNotFoundError(f"File not found: {relative_path}")
|
||||
|
||||
ext = target.suffix.lower()
|
||||
if ext not in ALLOWED_EXTENSIONS:
|
||||
raise ValueError(f"Disallowed file type: {ext}")
|
||||
|
||||
content_type = MIME_TYPES[ext]
|
||||
raw = target.read_bytes()
|
||||
|
||||
return {
|
||||
"content": base64.b64encode(raw).decode("ascii"),
|
||||
"content_type": content_type,
|
||||
}
|
||||
|
||||
|
||||
def _validate_path(relative_path: str) -> None:
|
||||
"""Reject directory traversal and absolute paths."""
|
||||
if os.path.isabs(relative_path):
|
||||
raise ValueError(f"Absolute paths are not allowed: {relative_path}")
|
||||
if ".." in PurePosixPath(relative_path).parts:
|
||||
raise ValueError(f"Directory traversal is not allowed: {relative_path}")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Host-side cache and aiohttp handler
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class WebDirectoryCache:
|
||||
"""Host-side in-memory cache for proxied web directory contents.
|
||||
|
||||
Populated lazily via RPC calls to the child's WebDirectoryProxy.
|
||||
Once a file is cached, subsequent requests are served from memory.
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
# {extension_name: {relative_path: {"content": bytes, "content_type": str}}}
|
||||
self._file_cache: dict[str, dict[str, dict[str, Any]]] = {}
|
||||
# {extension_name: [{"relative_path": str, "content_type": str}, ...]}
|
||||
self._listing_cache: dict[str, list[dict[str, str]]] = {}
|
||||
# {extension_name: WebDirectoryProxy (RPC proxy instance)}
|
||||
self._proxies: dict[str, Any] = {}
|
||||
|
||||
def register_proxy(self, extension_name: str, proxy: Any) -> None:
|
||||
"""Register an RPC proxy for an extension's web directory."""
|
||||
self._proxies[extension_name] = proxy
|
||||
logger.info(
|
||||
"][ WebDirectoryCache: registered proxy for %s", extension_name
|
||||
)
|
||||
|
||||
@property
|
||||
def extension_names(self) -> list[str]:
|
||||
return list(self._proxies.keys())
|
||||
|
||||
def list_files(self, extension_name: str) -> list[dict[str, str]]:
|
||||
"""List servable files for an extension (cached after first call)."""
|
||||
if extension_name not in self._listing_cache:
|
||||
proxy = self._proxies.get(extension_name)
|
||||
if proxy is None:
|
||||
return []
|
||||
try:
|
||||
self._listing_cache[extension_name] = proxy.list_web_files(
|
||||
extension_name
|
||||
)
|
||||
except Exception:
|
||||
logger.warning(
|
||||
"][ WebDirectoryCache: failed to list files for %s",
|
||||
extension_name,
|
||||
exc_info=True,
|
||||
)
|
||||
return []
|
||||
return self._listing_cache[extension_name]
|
||||
|
||||
def get_file(
|
||||
self, extension_name: str, relative_path: str
|
||||
) -> dict[str, Any] | None:
|
||||
"""Get file content (cached after first fetch). Returns None on miss."""
|
||||
ext_cache = self._file_cache.get(extension_name)
|
||||
if ext_cache and relative_path in ext_cache:
|
||||
return ext_cache[relative_path]
|
||||
|
||||
proxy = self._proxies.get(extension_name)
|
||||
if proxy is None:
|
||||
return None
|
||||
|
||||
try:
|
||||
result = proxy.get_web_file(extension_name, relative_path)
|
||||
except (FileNotFoundError, ValueError):
|
||||
return None
|
||||
except Exception:
|
||||
logger.warning(
|
||||
"][ WebDirectoryCache: failed to fetch %s/%s",
|
||||
extension_name,
|
||||
relative_path,
|
||||
exc_info=True,
|
||||
)
|
||||
return None
|
||||
|
||||
try:
|
||||
decoded = {
|
||||
"content": base64.b64decode(result["content"], validate=True),
|
||||
"content_type": result["content_type"],
|
||||
}
|
||||
except (binascii.Error, KeyError, TypeError):
|
||||
logger.warning(
|
||||
"][ WebDirectoryCache: invalid payload for %s/%s",
|
||||
extension_name,
|
||||
relative_path,
|
||||
exc_info=True,
|
||||
)
|
||||
return None
|
||||
|
||||
if extension_name not in self._file_cache:
|
||||
self._file_cache[extension_name] = {}
|
||||
self._file_cache[extension_name][relative_path] = decoded
|
||||
return decoded
|
||||
|
||||
|
||||
# Global cache instance — populated during isolation loading
|
||||
_web_directory_cache = WebDirectoryCache()
|
||||
|
||||
|
||||
def get_web_directory_cache() -> WebDirectoryCache:
|
||||
return _web_directory_cache
|
||||
@@ -1,49 +0,0 @@
|
||||
import asyncio
|
||||
import logging
|
||||
import threading
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class RpcBridge:
|
||||
"""Minimal helper to run coroutines synchronously inside isolated processes.
|
||||
|
||||
If an event loop is already running, the coroutine is executed on a fresh
|
||||
thread with its own loop to avoid nested run_until_complete errors.
|
||||
"""
|
||||
|
||||
def run_sync(self, maybe_coro):
|
||||
if not asyncio.iscoroutine(maybe_coro):
|
||||
return maybe_coro
|
||||
|
||||
try:
|
||||
loop = asyncio.get_running_loop()
|
||||
except RuntimeError:
|
||||
loop = None
|
||||
|
||||
if loop and loop.is_running():
|
||||
result_container = {}
|
||||
exc_container = {}
|
||||
|
||||
def _runner():
|
||||
try:
|
||||
new_loop = asyncio.new_event_loop()
|
||||
asyncio.set_event_loop(new_loop)
|
||||
result_container["value"] = new_loop.run_until_complete(maybe_coro)
|
||||
except Exception as exc: # pragma: no cover
|
||||
exc_container["error"] = exc
|
||||
finally:
|
||||
try:
|
||||
new_loop.close()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
t = threading.Thread(target=_runner, daemon=True)
|
||||
t.start()
|
||||
t.join()
|
||||
|
||||
if "error" in exc_container:
|
||||
raise exc_container["error"]
|
||||
return result_container.get("value")
|
||||
|
||||
return asyncio.run(maybe_coro)
|
||||
@@ -1,471 +0,0 @@
|
||||
# pylint: disable=consider-using-from-import,import-outside-toplevel,no-member
|
||||
from __future__ import annotations
|
||||
|
||||
import copy
|
||||
import logging
|
||||
import os
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, List, Set, TYPE_CHECKING
|
||||
|
||||
from .proxies.helper_proxies import restore_input_types
|
||||
from .shm_forensics import scan_shm_forensics
|
||||
|
||||
_IMPORT_TORCH = os.environ.get("PYISOLATE_IMPORT_TORCH", "1") == "1"
|
||||
|
||||
_ComfyNodeInternal = object
|
||||
latest_io = None
|
||||
|
||||
if _IMPORT_TORCH:
|
||||
from comfy_api.internal import _ComfyNodeInternal
|
||||
from comfy_api.latest import _io as latest_io
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from .extension_wrapper import ComfyNodeExtension
|
||||
|
||||
LOG_PREFIX = "]["
|
||||
_PRE_EXEC_MIN_FREE_VRAM_BYTES = 2 * 1024 * 1024 * 1024
|
||||
|
||||
|
||||
class _RemoteObjectRegistryCaller:
|
||||
def __init__(self, extension: Any) -> None:
|
||||
self._extension = extension
|
||||
|
||||
def __getattr__(self, method_name: str) -> Any:
|
||||
async def _call(instance_id: str, *args: Any, **kwargs: Any) -> Any:
|
||||
return await self._extension.call_remote_object_method(
|
||||
instance_id,
|
||||
method_name,
|
||||
*args,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
return _call
|
||||
|
||||
|
||||
def _wrap_remote_handles_as_host_proxies(value: Any, extension: Any) -> Any:
|
||||
from pyisolate._internal.remote_handle import RemoteObjectHandle
|
||||
|
||||
if isinstance(value, RemoteObjectHandle):
|
||||
if value.type_name == "ModelPatcher":
|
||||
from comfy.isolation.model_patcher_proxy import ModelPatcherProxy
|
||||
|
||||
proxy = ModelPatcherProxy(value.object_id, manage_lifecycle=False)
|
||||
proxy._rpc_caller = _RemoteObjectRegistryCaller(extension) # type: ignore[attr-defined]
|
||||
proxy._pyisolate_remote_handle = value # type: ignore[attr-defined]
|
||||
return proxy
|
||||
if value.type_name == "VAE":
|
||||
from comfy.isolation.vae_proxy import VAEProxy
|
||||
|
||||
proxy = VAEProxy(value.object_id, manage_lifecycle=False)
|
||||
proxy._rpc_caller = _RemoteObjectRegistryCaller(extension) # type: ignore[attr-defined]
|
||||
proxy._pyisolate_remote_handle = value # type: ignore[attr-defined]
|
||||
return proxy
|
||||
if value.type_name == "CLIP":
|
||||
from comfy.isolation.clip_proxy import CLIPProxy
|
||||
|
||||
proxy = CLIPProxy(value.object_id, manage_lifecycle=False)
|
||||
proxy._rpc_caller = _RemoteObjectRegistryCaller(extension) # type: ignore[attr-defined]
|
||||
proxy._pyisolate_remote_handle = value # type: ignore[attr-defined]
|
||||
return proxy
|
||||
if value.type_name == "ModelSampling":
|
||||
from comfy.isolation.model_sampling_proxy import ModelSamplingProxy
|
||||
|
||||
proxy = ModelSamplingProxy(value.object_id, manage_lifecycle=False)
|
||||
proxy._rpc_caller = _RemoteObjectRegistryCaller(extension) # type: ignore[attr-defined]
|
||||
proxy._pyisolate_remote_handle = value # type: ignore[attr-defined]
|
||||
return proxy
|
||||
return value
|
||||
|
||||
if isinstance(value, dict):
|
||||
return {
|
||||
k: _wrap_remote_handles_as_host_proxies(v, extension) for k, v in value.items()
|
||||
}
|
||||
|
||||
if isinstance(value, (list, tuple)):
|
||||
wrapped = [_wrap_remote_handles_as_host_proxies(item, extension) for item in value]
|
||||
return type(value)(wrapped)
|
||||
|
||||
return value
|
||||
|
||||
|
||||
def _resource_snapshot() -> Dict[str, int]:
|
||||
fd_count = -1
|
||||
shm_sender_files = 0
|
||||
try:
|
||||
fd_count = len(os.listdir("/proc/self/fd"))
|
||||
except Exception:
|
||||
pass
|
||||
try:
|
||||
shm_root = Path("/dev/shm")
|
||||
if shm_root.exists():
|
||||
prefix = f"torch_{os.getpid()}_"
|
||||
shm_sender_files = sum(1 for _ in shm_root.glob(f"{prefix}*"))
|
||||
except Exception:
|
||||
pass
|
||||
return {"fd_count": fd_count, "shm_sender_files": shm_sender_files}
|
||||
|
||||
|
||||
def _tensor_transport_summary(value: Any) -> Dict[str, int]:
|
||||
summary: Dict[str, int] = {
|
||||
"tensor_count": 0,
|
||||
"cpu_tensors": 0,
|
||||
"cuda_tensors": 0,
|
||||
"shared_cpu_tensors": 0,
|
||||
"tensor_bytes": 0,
|
||||
}
|
||||
try:
|
||||
import torch
|
||||
except Exception:
|
||||
return summary
|
||||
|
||||
def visit(node: Any) -> None:
|
||||
if isinstance(node, torch.Tensor):
|
||||
summary["tensor_count"] += 1
|
||||
summary["tensor_bytes"] += int(node.numel() * node.element_size())
|
||||
if node.device.type == "cpu":
|
||||
summary["cpu_tensors"] += 1
|
||||
if node.is_shared():
|
||||
summary["shared_cpu_tensors"] += 1
|
||||
elif node.device.type == "cuda":
|
||||
summary["cuda_tensors"] += 1
|
||||
return
|
||||
if isinstance(node, dict):
|
||||
for v in node.values():
|
||||
visit(v)
|
||||
return
|
||||
if isinstance(node, (list, tuple)):
|
||||
for v in node:
|
||||
visit(v)
|
||||
|
||||
visit(value)
|
||||
return summary
|
||||
|
||||
|
||||
def _extract_hidden_unique_id(inputs: Dict[str, Any]) -> str | None:
|
||||
for key, value in inputs.items():
|
||||
key_text = str(key)
|
||||
if "unique_id" in key_text:
|
||||
return str(value)
|
||||
return None
|
||||
|
||||
|
||||
def _flush_tensor_transport_state(marker: str, logger: logging.Logger) -> None:
|
||||
try:
|
||||
from pyisolate import flush_tensor_keeper # type: ignore[attr-defined]
|
||||
except Exception:
|
||||
return
|
||||
if not callable(flush_tensor_keeper):
|
||||
return
|
||||
flushed = flush_tensor_keeper()
|
||||
if flushed > 0:
|
||||
logger.debug(
|
||||
"%s %s flush_tensor_keeper released=%d", LOG_PREFIX, marker, flushed
|
||||
)
|
||||
|
||||
|
||||
def _relieve_host_vram_pressure(marker: str, logger: logging.Logger) -> None:
|
||||
import comfy.model_management as model_management
|
||||
|
||||
model_management.cleanup_models_gc()
|
||||
model_management.cleanup_models()
|
||||
|
||||
device = model_management.get_torch_device()
|
||||
if not hasattr(device, "type") or device.type == "cpu":
|
||||
return
|
||||
|
||||
required = max(
|
||||
model_management.minimum_inference_memory(),
|
||||
_PRE_EXEC_MIN_FREE_VRAM_BYTES,
|
||||
)
|
||||
if model_management.get_free_memory(device) < required:
|
||||
model_management.free_memory(required, device, for_dynamic=True)
|
||||
if model_management.get_free_memory(device) < required:
|
||||
model_management.free_memory(required, device, for_dynamic=False)
|
||||
model_management.cleanup_models()
|
||||
model_management.soft_empty_cache()
|
||||
logger.debug("%s %s free_memory target=%d", LOG_PREFIX, marker, required)
|
||||
|
||||
|
||||
def _detach_shared_cpu_tensors(value: Any) -> Any:
|
||||
try:
|
||||
import torch
|
||||
except Exception:
|
||||
return value
|
||||
|
||||
if isinstance(value, torch.Tensor):
|
||||
if value.device.type == "cpu" and value.is_shared():
|
||||
clone = value.clone()
|
||||
if value.requires_grad:
|
||||
clone.requires_grad_(True)
|
||||
return clone
|
||||
return value
|
||||
if isinstance(value, list):
|
||||
return [_detach_shared_cpu_tensors(v) for v in value]
|
||||
if isinstance(value, tuple):
|
||||
return tuple(_detach_shared_cpu_tensors(v) for v in value)
|
||||
if isinstance(value, dict):
|
||||
return {k: _detach_shared_cpu_tensors(v) for k, v in value.items()}
|
||||
return value
|
||||
|
||||
|
||||
def build_stub_class(
|
||||
node_name: str,
|
||||
info: Dict[str, object],
|
||||
extension: "ComfyNodeExtension",
|
||||
running_extensions: Dict[str, "ComfyNodeExtension"],
|
||||
logger: logging.Logger,
|
||||
) -> type:
|
||||
if latest_io is None:
|
||||
raise RuntimeError("comfy_api.latest._io is required to build isolation stubs")
|
||||
is_v3 = bool(info.get("is_v3", False))
|
||||
function_name = "_pyisolate_execute"
|
||||
restored_input_types = restore_input_types(info.get("input_types", {}))
|
||||
|
||||
async def _execute(self, **inputs):
|
||||
from comfy.isolation import _RUNNING_EXTENSIONS
|
||||
|
||||
# Update BOTH the local dict AND the module-level dict
|
||||
running_extensions[extension.name] = extension
|
||||
_RUNNING_EXTENSIONS[extension.name] = extension
|
||||
prev_child = None
|
||||
node_unique_id = _extract_hidden_unique_id(inputs)
|
||||
summary = _tensor_transport_summary(inputs)
|
||||
resources = _resource_snapshot()
|
||||
logger.debug(
|
||||
"%s ISO:execute_start ext=%s node=%s uid=%s",
|
||||
LOG_PREFIX,
|
||||
extension.name,
|
||||
node_name,
|
||||
node_unique_id or "-",
|
||||
)
|
||||
logger.debug(
|
||||
"%s ISO:execute_start ext=%s node=%s uid=%s tensors=%d cpu=%d cuda=%d shared_cpu=%d bytes=%d fds=%d sender_shm=%d",
|
||||
LOG_PREFIX,
|
||||
extension.name,
|
||||
node_name,
|
||||
node_unique_id or "-",
|
||||
summary["tensor_count"],
|
||||
summary["cpu_tensors"],
|
||||
summary["cuda_tensors"],
|
||||
summary["shared_cpu_tensors"],
|
||||
summary["tensor_bytes"],
|
||||
resources["fd_count"],
|
||||
resources["shm_sender_files"],
|
||||
)
|
||||
scan_shm_forensics("RUNTIME:execute_start", refresh_model_context=True)
|
||||
try:
|
||||
if os.environ.get("PYISOLATE_CHILD") != "1":
|
||||
_relieve_host_vram_pressure("RUNTIME:pre_execute", logger)
|
||||
scan_shm_forensics("RUNTIME:pre_execute", refresh_model_context=True)
|
||||
from pyisolate._internal.model_serialization import (
|
||||
serialize_for_isolation,
|
||||
deserialize_from_isolation,
|
||||
)
|
||||
|
||||
prev_child = os.environ.pop("PYISOLATE_CHILD", None)
|
||||
logger.debug(
|
||||
"%s ISO:serialize_start ext=%s node=%s uid=%s",
|
||||
LOG_PREFIX,
|
||||
extension.name,
|
||||
node_name,
|
||||
node_unique_id or "-",
|
||||
)
|
||||
# Unwrap NodeOutput-like dicts before serialization.
|
||||
# OUTPUT_NODE nodes return {"ui": {...}, "result": (outputs...)}
|
||||
# and the executor may pass this dict as input to downstream nodes.
|
||||
unwrapped_inputs = {}
|
||||
for k, v in inputs.items():
|
||||
if isinstance(v, dict) and "result" in v and ("ui" in v or "__node_output__" in v):
|
||||
result = v.get("result")
|
||||
if isinstance(result, (tuple, list)) and len(result) > 0:
|
||||
unwrapped_inputs[k] = result[0]
|
||||
else:
|
||||
unwrapped_inputs[k] = result
|
||||
else:
|
||||
unwrapped_inputs[k] = v
|
||||
serialized = serialize_for_isolation(unwrapped_inputs)
|
||||
logger.debug(
|
||||
"%s ISO:serialize_done ext=%s node=%s uid=%s",
|
||||
LOG_PREFIX,
|
||||
extension.name,
|
||||
node_name,
|
||||
node_unique_id or "-",
|
||||
)
|
||||
logger.debug(
|
||||
"%s ISO:dispatch_start ext=%s node=%s uid=%s",
|
||||
LOG_PREFIX,
|
||||
extension.name,
|
||||
node_name,
|
||||
node_unique_id or "-",
|
||||
)
|
||||
result = await extension.execute_node(node_name, **serialized)
|
||||
logger.debug(
|
||||
"%s ISO:dispatch_done ext=%s node=%s uid=%s",
|
||||
LOG_PREFIX,
|
||||
extension.name,
|
||||
node_name,
|
||||
node_unique_id or "-",
|
||||
)
|
||||
# Reconstruct NodeOutput if the child serialized one
|
||||
if isinstance(result, dict) and result.get("__node_output__"):
|
||||
from comfy_api.latest import io as latest_io
|
||||
args_raw = result.get("args", ())
|
||||
deserialized_args = await deserialize_from_isolation(args_raw, extension)
|
||||
deserialized_args = _wrap_remote_handles_as_host_proxies(
|
||||
deserialized_args, extension
|
||||
)
|
||||
deserialized_args = _detach_shared_cpu_tensors(deserialized_args)
|
||||
ui_raw = result.get("ui")
|
||||
deserialized_ui = None
|
||||
if ui_raw is not None:
|
||||
deserialized_ui = await deserialize_from_isolation(ui_raw, extension)
|
||||
deserialized_ui = _wrap_remote_handles_as_host_proxies(
|
||||
deserialized_ui, extension
|
||||
)
|
||||
deserialized_ui = _detach_shared_cpu_tensors(deserialized_ui)
|
||||
scan_shm_forensics("RUNTIME:post_execute", refresh_model_context=True)
|
||||
return latest_io.NodeOutput(
|
||||
*deserialized_args,
|
||||
ui=deserialized_ui,
|
||||
expand=result.get("expand"),
|
||||
block_execution=result.get("block_execution"),
|
||||
)
|
||||
# OUTPUT_NODE: if sealed worker returned a tuple/list whose first
|
||||
# element is a {"ui": ...} dict, unwrap it for the executor.
|
||||
if (isinstance(result, (tuple, list)) and len(result) == 1
|
||||
and isinstance(result[0], dict) and "ui" in result[0]):
|
||||
return result[0]
|
||||
deserialized = await deserialize_from_isolation(result, extension)
|
||||
deserialized = _wrap_remote_handles_as_host_proxies(deserialized, extension)
|
||||
scan_shm_forensics("RUNTIME:post_execute", refresh_model_context=True)
|
||||
return _detach_shared_cpu_tensors(deserialized)
|
||||
except ImportError:
|
||||
return await extension.execute_node(node_name, **inputs)
|
||||
except Exception:
|
||||
logger.exception(
|
||||
"%s ISO:execute_error ext=%s node=%s uid=%s",
|
||||
LOG_PREFIX,
|
||||
extension.name,
|
||||
node_name,
|
||||
node_unique_id or "-",
|
||||
)
|
||||
raise
|
||||
finally:
|
||||
if prev_child is not None:
|
||||
os.environ["PYISOLATE_CHILD"] = prev_child
|
||||
logger.debug(
|
||||
"%s ISO:execute_end ext=%s node=%s uid=%s",
|
||||
LOG_PREFIX,
|
||||
extension.name,
|
||||
node_name,
|
||||
node_unique_id or "-",
|
||||
)
|
||||
scan_shm_forensics("RUNTIME:execute_end", refresh_model_context=True)
|
||||
|
||||
def _input_types(
|
||||
cls,
|
||||
include_hidden: bool = True,
|
||||
return_schema: bool = False,
|
||||
live_inputs: Any = None,
|
||||
):
|
||||
if not is_v3:
|
||||
return restored_input_types
|
||||
|
||||
inputs_copy = copy.deepcopy(restored_input_types)
|
||||
if not include_hidden:
|
||||
inputs_copy.pop("hidden", None)
|
||||
|
||||
v3_data: Dict[str, Any] = {"hidden_inputs": {}}
|
||||
dynamic = inputs_copy.pop("dynamic_paths", None)
|
||||
if dynamic is not None:
|
||||
v3_data["dynamic_paths"] = dynamic
|
||||
|
||||
if return_schema:
|
||||
hidden_vals = info.get("hidden", []) or []
|
||||
hidden_enums = []
|
||||
for h in hidden_vals:
|
||||
try:
|
||||
hidden_enums.append(latest_io.Hidden(h))
|
||||
except Exception:
|
||||
hidden_enums.append(h)
|
||||
|
||||
class SchemaProxy:
|
||||
hidden = hidden_enums
|
||||
|
||||
return inputs_copy, SchemaProxy, v3_data
|
||||
return inputs_copy
|
||||
|
||||
def _validate_class(cls):
|
||||
return True
|
||||
|
||||
def _get_node_info_v1(cls):
|
||||
node_info = copy.deepcopy(info.get("schema_v1", {}))
|
||||
relative_python_module = node_info.get("python_module")
|
||||
if not isinstance(relative_python_module, str) or not relative_python_module:
|
||||
relative_python_module = f"custom_nodes.{extension.name}"
|
||||
node_info["python_module"] = relative_python_module
|
||||
return node_info
|
||||
|
||||
def _get_base_class(cls):
|
||||
return latest_io.ComfyNode
|
||||
|
||||
attributes: Dict[str, object] = {
|
||||
"FUNCTION": function_name,
|
||||
"CATEGORY": info.get("category", ""),
|
||||
"OUTPUT_NODE": info.get("output_node", False),
|
||||
"RETURN_TYPES": tuple(info.get("return_types", ()) or ()),
|
||||
"RETURN_NAMES": info.get("return_names"),
|
||||
function_name: _execute,
|
||||
"_pyisolate_extension": extension,
|
||||
"_pyisolate_node_name": node_name,
|
||||
"INPUT_TYPES": classmethod(_input_types),
|
||||
}
|
||||
|
||||
output_is_list = info.get("output_is_list")
|
||||
if output_is_list is not None:
|
||||
attributes["OUTPUT_IS_LIST"] = tuple(output_is_list)
|
||||
|
||||
if is_v3:
|
||||
attributes["VALIDATE_CLASS"] = classmethod(_validate_class)
|
||||
attributes["GET_NODE_INFO_V1"] = classmethod(_get_node_info_v1)
|
||||
attributes["GET_BASE_CLASS"] = classmethod(_get_base_class)
|
||||
attributes["DESCRIPTION"] = info.get("description", "")
|
||||
attributes["EXPERIMENTAL"] = info.get("experimental", False)
|
||||
attributes["DEPRECATED"] = info.get("deprecated", False)
|
||||
attributes["API_NODE"] = info.get("api_node", False)
|
||||
attributes["NOT_IDEMPOTENT"] = info.get("not_idempotent", False)
|
||||
attributes["ACCEPT_ALL_INPUTS"] = info.get("accept_all_inputs", False)
|
||||
attributes["_ACCEPT_ALL_INPUTS"] = info.get("accept_all_inputs", False)
|
||||
attributes["INPUT_IS_LIST"] = info.get("input_is_list", False)
|
||||
|
||||
class_name = f"PyIsolate_{node_name}".replace(" ", "_")
|
||||
bases = (_ComfyNodeInternal,) if is_v3 else ()
|
||||
stub_cls = type(class_name, bases, attributes)
|
||||
|
||||
if is_v3:
|
||||
try:
|
||||
stub_cls.VALIDATE_CLASS()
|
||||
except Exception as e:
|
||||
logger.error("%s VALIDATE_CLASS failed: %s - %s", LOG_PREFIX, node_name, e)
|
||||
|
||||
return stub_cls
|
||||
|
||||
|
||||
def get_class_types_for_extension(
|
||||
extension_name: str,
|
||||
running_extensions: Dict[str, "ComfyNodeExtension"],
|
||||
specs: List[Any],
|
||||
) -> Set[str]:
|
||||
extension = running_extensions.get(extension_name)
|
||||
if not extension:
|
||||
return set()
|
||||
|
||||
ext_path = Path(extension.module_path)
|
||||
class_types = set()
|
||||
for spec in specs:
|
||||
if spec.module_path.resolve() == ext_path.resolve():
|
||||
class_types.add(spec.node_name)
|
||||
return class_types
|
||||
|
||||
|
||||
__all__ = ["build_stub_class", "get_class_types_for_extension"]
|
||||
@@ -1,217 +0,0 @@
|
||||
# pylint: disable=consider-using-from-import,import-outside-toplevel
|
||||
from __future__ import annotations
|
||||
|
||||
import atexit
|
||||
import hashlib
|
||||
import logging
|
||||
import os
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, List, Set
|
||||
|
||||
LOG_PREFIX = "]["
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _shm_debug_enabled() -> bool:
|
||||
return os.environ.get("COMFY_ISO_SHM_DEBUG") == "1"
|
||||
|
||||
|
||||
class _SHMForensicsTracker:
|
||||
def __init__(self) -> None:
|
||||
self._started = False
|
||||
self._tracked_files: Set[str] = set()
|
||||
self._current_model_context: Dict[str, str] = {
|
||||
"id": "unknown",
|
||||
"name": "unknown",
|
||||
"hash": "????",
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
def _snapshot_shm() -> Set[str]:
|
||||
shm_path = Path("/dev/shm")
|
||||
if not shm_path.exists():
|
||||
return set()
|
||||
return {f.name for f in shm_path.glob("torch_*")}
|
||||
|
||||
def start(self) -> None:
|
||||
if self._started or not _shm_debug_enabled():
|
||||
return
|
||||
self._tracked_files = self._snapshot_shm()
|
||||
self._started = True
|
||||
logger.debug(
|
||||
"%s SHM:forensics_enabled tracked=%d", LOG_PREFIX, len(self._tracked_files)
|
||||
)
|
||||
|
||||
def stop(self) -> None:
|
||||
if not self._started:
|
||||
return
|
||||
self.scan("shutdown", refresh_model_context=True)
|
||||
self._started = False
|
||||
logger.debug("%s SHM:forensics_disabled", LOG_PREFIX)
|
||||
|
||||
def _compute_model_hash(self, model_patcher: Any) -> str:
|
||||
try:
|
||||
model_instance_id = getattr(model_patcher, "_instance_id", None)
|
||||
if model_instance_id is not None:
|
||||
model_id_text = str(model_instance_id)
|
||||
return model_id_text[-4:] if len(model_id_text) >= 4 else model_id_text
|
||||
|
||||
import torch
|
||||
|
||||
real_model = (
|
||||
model_patcher.model
|
||||
if hasattr(model_patcher, "model")
|
||||
else model_patcher
|
||||
)
|
||||
tensor = None
|
||||
if hasattr(real_model, "parameters"):
|
||||
for p in real_model.parameters():
|
||||
if torch.is_tensor(p) and p.numel() > 0:
|
||||
tensor = p
|
||||
break
|
||||
|
||||
if tensor is None:
|
||||
return "0000"
|
||||
|
||||
flat = tensor.flatten()
|
||||
values = []
|
||||
indices = [0, flat.shape[0] // 2, flat.shape[0] - 1]
|
||||
for i in indices:
|
||||
if i < flat.shape[0]:
|
||||
values.append(flat[i].item())
|
||||
|
||||
size = 0
|
||||
if hasattr(model_patcher, "model_size"):
|
||||
size = model_patcher.model_size()
|
||||
sample_str = f"{values}_{id(model_patcher):016x}_{size}"
|
||||
return hashlib.sha256(sample_str.encode()).hexdigest()[-4:]
|
||||
except Exception:
|
||||
return "err!"
|
||||
|
||||
def _get_models_snapshot(self) -> List[Dict[str, Any]]:
|
||||
try:
|
||||
import comfy.model_management as model_management
|
||||
except Exception:
|
||||
return []
|
||||
|
||||
snapshot: List[Dict[str, Any]] = []
|
||||
try:
|
||||
for loaded_model in model_management.current_loaded_models:
|
||||
model = loaded_model.model
|
||||
if model is None:
|
||||
continue
|
||||
if str(getattr(loaded_model, "device", "")) != "cuda:0":
|
||||
continue
|
||||
|
||||
name = (
|
||||
model.model.__class__.__name__
|
||||
if hasattr(model, "model")
|
||||
else type(model).__name__
|
||||
)
|
||||
model_hash = self._compute_model_hash(model)
|
||||
model_instance_id = getattr(model, "_instance_id", None)
|
||||
if model_instance_id is None:
|
||||
model_instance_id = model_hash
|
||||
snapshot.append(
|
||||
{
|
||||
"name": str(name),
|
||||
"id": str(model_instance_id),
|
||||
"hash": str(model_hash or "????"),
|
||||
"used": bool(getattr(loaded_model, "currently_used", False)),
|
||||
}
|
||||
)
|
||||
except Exception:
|
||||
return []
|
||||
|
||||
return snapshot
|
||||
|
||||
def _update_model_context(self) -> None:
|
||||
snapshot = self._get_models_snapshot()
|
||||
selected = None
|
||||
|
||||
used_models = [m for m in snapshot if m.get("used") and m.get("id")]
|
||||
if used_models:
|
||||
selected = used_models[-1]
|
||||
else:
|
||||
live_models = [m for m in snapshot if m.get("id")]
|
||||
if live_models:
|
||||
selected = live_models[-1]
|
||||
|
||||
if selected is None:
|
||||
self._current_model_context = {
|
||||
"id": "unknown",
|
||||
"name": "unknown",
|
||||
"hash": "????",
|
||||
}
|
||||
return
|
||||
|
||||
self._current_model_context = {
|
||||
"id": str(selected.get("id", "unknown")),
|
||||
"name": str(selected.get("name", "unknown")),
|
||||
"hash": str(selected.get("hash", "????") or "????"),
|
||||
}
|
||||
|
||||
def scan(self, marker: str, refresh_model_context: bool = True) -> None:
|
||||
if not self._started or not _shm_debug_enabled():
|
||||
return
|
||||
|
||||
if refresh_model_context:
|
||||
self._update_model_context()
|
||||
|
||||
current = self._snapshot_shm()
|
||||
added = current - self._tracked_files
|
||||
removed = self._tracked_files - current
|
||||
self._tracked_files = current
|
||||
|
||||
if not added and not removed:
|
||||
logger.debug("%s SHM:scan marker=%s changes=0", LOG_PREFIX, marker)
|
||||
return
|
||||
|
||||
for filename in sorted(added):
|
||||
logger.info("%s SHM:created | %s", LOG_PREFIX, filename)
|
||||
model_id = self._current_model_context["id"]
|
||||
if model_id == "unknown":
|
||||
logger.error(
|
||||
"%s SHM:model_association_missing | file=%s | reason=no_active_model_context",
|
||||
LOG_PREFIX,
|
||||
filename,
|
||||
)
|
||||
else:
|
||||
logger.info(
|
||||
"%s SHM:model_association | model=%s | file=%s | name=%s | hash=%s",
|
||||
LOG_PREFIX,
|
||||
model_id,
|
||||
filename,
|
||||
self._current_model_context["name"],
|
||||
self._current_model_context["hash"],
|
||||
)
|
||||
|
||||
for filename in sorted(removed):
|
||||
logger.info("%s SHM:deleted | %s", LOG_PREFIX, filename)
|
||||
|
||||
logger.debug(
|
||||
"%s SHM:scan marker=%s created=%d deleted=%d active=%d",
|
||||
LOG_PREFIX,
|
||||
marker,
|
||||
len(added),
|
||||
len(removed),
|
||||
len(self._tracked_files),
|
||||
)
|
||||
|
||||
|
||||
_TRACKER = _SHMForensicsTracker()
|
||||
|
||||
|
||||
def start_shm_forensics() -> None:
|
||||
_TRACKER.start()
|
||||
|
||||
|
||||
def scan_shm_forensics(marker: str, refresh_model_context: bool = True) -> None:
|
||||
_TRACKER.scan(marker, refresh_model_context=refresh_model_context)
|
||||
|
||||
|
||||
def stop_shm_forensics() -> None:
|
||||
_TRACKER.stop()
|
||||
|
||||
|
||||
atexit.register(stop_shm_forensics)
|
||||
@@ -1,214 +0,0 @@
|
||||
# pylint: disable=attribute-defined-outside-init
|
||||
import logging
|
||||
from typing import Any
|
||||
|
||||
from comfy.isolation.proxies.base import (
|
||||
IS_CHILD_PROCESS,
|
||||
BaseProxy,
|
||||
BaseRegistry,
|
||||
detach_if_grad,
|
||||
)
|
||||
from comfy.isolation.model_patcher_proxy import ModelPatcherProxy, ModelPatcherRegistry
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class FirstStageModelRegistry(BaseRegistry[Any]):
|
||||
_type_prefix = "first_stage_model"
|
||||
|
||||
async def get_property(self, instance_id: str, name: str) -> Any:
|
||||
obj = self._get_instance(instance_id)
|
||||
return getattr(obj, name)
|
||||
|
||||
async def has_property(self, instance_id: str, name: str) -> bool:
|
||||
obj = self._get_instance(instance_id)
|
||||
return hasattr(obj, name)
|
||||
|
||||
|
||||
class FirstStageModelProxy(BaseProxy[FirstStageModelRegistry]):
|
||||
_registry_class = FirstStageModelRegistry
|
||||
__module__ = "comfy.ldm.models.autoencoder"
|
||||
|
||||
def __getattr__(self, name: str) -> Any:
|
||||
try:
|
||||
return self._call_rpc("get_property", name)
|
||||
except Exception as e:
|
||||
raise AttributeError(
|
||||
f"'{self.__class__.__name__}' object has no attribute '{name}'"
|
||||
) from e
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return f"<FirstStageModelProxy {self._instance_id}>"
|
||||
|
||||
|
||||
class VAERegistry(BaseRegistry[Any]):
|
||||
_type_prefix = "vae"
|
||||
|
||||
async def get_patcher_id(self, instance_id: str) -> str:
|
||||
vae = self._get_instance(instance_id)
|
||||
return ModelPatcherRegistry().register(vae.patcher)
|
||||
|
||||
async def get_first_stage_model_id(self, instance_id: str) -> str:
|
||||
vae = self._get_instance(instance_id)
|
||||
return FirstStageModelRegistry().register(vae.first_stage_model)
|
||||
|
||||
async def encode(self, instance_id: str, pixels: Any) -> Any:
|
||||
return detach_if_grad(self._get_instance(instance_id).encode(pixels))
|
||||
|
||||
async def encode_tiled(
|
||||
self,
|
||||
instance_id: str,
|
||||
pixels: Any,
|
||||
tile_x: int = 512,
|
||||
tile_y: int = 512,
|
||||
overlap: int = 64,
|
||||
) -> Any:
|
||||
return detach_if_grad(
|
||||
self._get_instance(instance_id).encode_tiled(
|
||||
pixels, tile_x=tile_x, tile_y=tile_y, overlap=overlap
|
||||
)
|
||||
)
|
||||
|
||||
async def decode(self, instance_id: str, samples: Any, **kwargs: Any) -> Any:
|
||||
return detach_if_grad(self._get_instance(instance_id).decode(samples, **kwargs))
|
||||
|
||||
async def decode_tiled(
|
||||
self,
|
||||
instance_id: str,
|
||||
samples: Any,
|
||||
tile_x: int = 64,
|
||||
tile_y: int = 64,
|
||||
overlap: int = 16,
|
||||
**kwargs: Any,
|
||||
) -> Any:
|
||||
return detach_if_grad(
|
||||
self._get_instance(instance_id).decode_tiled(
|
||||
samples, tile_x=tile_x, tile_y=tile_y, overlap=overlap, **kwargs
|
||||
)
|
||||
)
|
||||
|
||||
async def get_property(self, instance_id: str, name: str) -> Any:
|
||||
return getattr(self._get_instance(instance_id), name)
|
||||
|
||||
async def memory_used_encode(self, instance_id: str, shape: Any, dtype: Any) -> int:
|
||||
return self._get_instance(instance_id).memory_used_encode(shape, dtype)
|
||||
|
||||
async def memory_used_decode(self, instance_id: str, shape: Any, dtype: Any) -> int:
|
||||
return self._get_instance(instance_id).memory_used_decode(shape, dtype)
|
||||
|
||||
async def process_input(self, instance_id: str, image: Any) -> Any:
|
||||
return detach_if_grad(self._get_instance(instance_id).process_input(image))
|
||||
|
||||
async def process_output(self, instance_id: str, image: Any) -> Any:
|
||||
return detach_if_grad(self._get_instance(instance_id).process_output(image))
|
||||
|
||||
|
||||
class VAEProxy(BaseProxy[VAERegistry]):
|
||||
_registry_class = VAERegistry
|
||||
__module__ = "comfy.sd"
|
||||
|
||||
@property
|
||||
def patcher(self) -> ModelPatcherProxy:
|
||||
if not hasattr(self, "_patcher_proxy"):
|
||||
patcher_id = self._call_rpc("get_patcher_id")
|
||||
self._patcher_proxy = ModelPatcherProxy(patcher_id, manage_lifecycle=False)
|
||||
return self._patcher_proxy
|
||||
|
||||
@property
|
||||
def first_stage_model(self) -> FirstStageModelProxy:
|
||||
if not hasattr(self, "_first_stage_model_proxy"):
|
||||
fsm_id = self._call_rpc("get_first_stage_model_id")
|
||||
self._first_stage_model_proxy = FirstStageModelProxy(
|
||||
fsm_id, manage_lifecycle=False
|
||||
)
|
||||
return self._first_stage_model_proxy
|
||||
|
||||
@property
|
||||
def vae_dtype(self) -> Any:
|
||||
return self._get_property("vae_dtype")
|
||||
|
||||
def encode(self, pixels: Any) -> Any:
|
||||
return self._call_rpc("encode", pixels)
|
||||
|
||||
def encode_tiled(
|
||||
self, pixels: Any, tile_x: int = 512, tile_y: int = 512, overlap: int = 64
|
||||
) -> Any:
|
||||
return self._call_rpc("encode_tiled", pixels, tile_x, tile_y, overlap)
|
||||
|
||||
def decode(self, samples: Any, **kwargs: Any) -> Any:
|
||||
return self._call_rpc("decode", samples, **kwargs)
|
||||
|
||||
def decode_tiled(
|
||||
self,
|
||||
samples: Any,
|
||||
tile_x: int = 64,
|
||||
tile_y: int = 64,
|
||||
overlap: int = 16,
|
||||
**kwargs: Any,
|
||||
) -> Any:
|
||||
return self._call_rpc(
|
||||
"decode_tiled", samples, tile_x, tile_y, overlap, **kwargs
|
||||
)
|
||||
|
||||
def get_sd(self) -> Any:
|
||||
return self._call_rpc("get_sd")
|
||||
|
||||
def _get_property(self, name: str) -> Any:
|
||||
return self._call_rpc("get_property", name)
|
||||
|
||||
@property
|
||||
def latent_dim(self) -> int:
|
||||
return self._get_property("latent_dim")
|
||||
|
||||
@property
|
||||
def latent_channels(self) -> int:
|
||||
return self._get_property("latent_channels")
|
||||
|
||||
@property
|
||||
def downscale_ratio(self) -> Any:
|
||||
return self._get_property("downscale_ratio")
|
||||
|
||||
@property
|
||||
def upscale_ratio(self) -> Any:
|
||||
return self._get_property("upscale_ratio")
|
||||
|
||||
@property
|
||||
def output_channels(self) -> int:
|
||||
return self._get_property("output_channels")
|
||||
|
||||
@property
|
||||
def not_video(self) -> bool:
|
||||
return self._get_property("not_video")
|
||||
|
||||
@property
|
||||
def device(self) -> Any:
|
||||
return self._get_property("device")
|
||||
|
||||
@property
|
||||
def working_dtypes(self) -> Any:
|
||||
return self._get_property("working_dtypes")
|
||||
|
||||
@property
|
||||
def disable_offload(self) -> bool:
|
||||
return self._get_property("disable_offload")
|
||||
|
||||
@property
|
||||
def size(self) -> Any:
|
||||
return self._get_property("size")
|
||||
|
||||
def memory_used_encode(self, shape: Any, dtype: Any) -> int:
|
||||
return self._call_rpc("memory_used_encode", shape, dtype)
|
||||
|
||||
def memory_used_decode(self, shape: Any, dtype: Any) -> int:
|
||||
return self._call_rpc("memory_used_decode", shape, dtype)
|
||||
|
||||
def process_input(self, image: Any) -> Any:
|
||||
return self._call_rpc("process_input", image)
|
||||
|
||||
def process_output(self, image: Any) -> Any:
|
||||
return self._call_rpc("process_output", image)
|
||||
|
||||
|
||||
if not IS_CHILD_PROCESS:
|
||||
_VAE_REGISTRY_SINGLETON = VAERegistry()
|
||||
_FIRST_STAGE_MODEL_REGISTRY_SINGLETON = FirstStageModelRegistry()
|
||||
@@ -1,5 +1,4 @@
|
||||
import math
|
||||
import os
|
||||
from functools import partial
|
||||
|
||||
from scipy import integrate
|
||||
@@ -13,8 +12,8 @@ from . import deis
|
||||
from . import sa_solver
|
||||
import comfy.model_patcher
|
||||
import comfy.model_sampling
|
||||
|
||||
import comfy.memory_management
|
||||
from comfy.cli_args import args
|
||||
from comfy.utils import model_trange as trange
|
||||
|
||||
def append_zero(x):
|
||||
@@ -192,13 +191,6 @@ def sample_euler(model, x, sigmas, extra_args=None, callback=None, disable=None,
|
||||
"""Implements Algorithm 2 (Euler steps) from Karras et al. (2022)."""
|
||||
extra_args = {} if extra_args is None else extra_args
|
||||
s_in = x.new_ones([x.shape[0]])
|
||||
isolation_active = args.use_process_isolation or os.environ.get("PYISOLATE_CHILD") == "1"
|
||||
if isolation_active:
|
||||
target_device = sigmas.device
|
||||
if x.device != target_device:
|
||||
x = x.to(target_device)
|
||||
s_in = s_in.to(target_device)
|
||||
|
||||
for i in trange(len(sigmas) - 1, disable=disable):
|
||||
if s_churn > 0:
|
||||
gamma = min(s_churn / (len(sigmas) - 1), 2 ** 0.5 - 1) if s_tmin <= sigmas[i] <= s_tmax else 0.
|
||||
|
||||
@@ -16,6 +16,7 @@ from comfy.ldm.lightricks.model import (
|
||||
from comfy.ldm.lightricks.symmetric_patchifier import AudioPatchifier
|
||||
from comfy.ldm.lightricks.embeddings_connector import Embeddings1DConnector
|
||||
import comfy.ldm.common_dit
|
||||
import comfy.model_prefetch
|
||||
|
||||
class CompressedTimestep:
|
||||
"""Store video timestep embeddings in compressed form using per-frame indexing."""
|
||||
@@ -907,9 +908,11 @@ class LTXAVModel(LTXVModel):
|
||||
"""Process transformer blocks for LTXAV."""
|
||||
patches_replace = transformer_options.get("patches_replace", {})
|
||||
blocks_replace = patches_replace.get("dit", {})
|
||||
prefetch_queue = comfy.model_prefetch.make_prefetch_queue(list(self.transformer_blocks), vx.device, transformer_options)
|
||||
|
||||
# Process transformer blocks
|
||||
for i, block in enumerate(self.transformer_blocks):
|
||||
comfy.model_prefetch.prefetch_queue_pop(prefetch_queue, vx.device, block)
|
||||
if ("double_block", i) in blocks_replace:
|
||||
|
||||
def block_wrap(args):
|
||||
@@ -982,6 +985,8 @@ class LTXAVModel(LTXVModel):
|
||||
a_prompt_timestep=a_prompt_timestep,
|
||||
)
|
||||
|
||||
comfy.model_prefetch.prefetch_queue_pop(prefetch_queue, vx.device, None)
|
||||
|
||||
return [vx, ax]
|
||||
|
||||
def _process_output(self, x, embedded_timestep, keyframe_idxs, **kwargs):
|
||||
|
||||
@@ -14,6 +14,8 @@ from .sub_quadratic_attention import efficient_dot_product_attention
|
||||
|
||||
from comfy import model_management
|
||||
|
||||
TORCH_HAS_GQA = model_management.torch_version_numeric >= (2, 5)
|
||||
|
||||
if model_management.xformers_enabled():
|
||||
import xformers
|
||||
import xformers.ops
|
||||
@@ -150,7 +152,12 @@ def attention_basic(q, k, v, heads, mask=None, attn_precision=None, skip_reshape
|
||||
b, _, dim_head = q.shape
|
||||
dim_head //= heads
|
||||
|
||||
scale = dim_head ** -0.5
|
||||
if kwargs.get("enable_gqa", False) and q.shape[-3] != k.shape[-3]:
|
||||
n_rep = q.shape[-3] // k.shape[-3]
|
||||
k = k.repeat_interleave(n_rep, dim=-3)
|
||||
v = v.repeat_interleave(n_rep, dim=-3)
|
||||
|
||||
scale = kwargs.get("scale", dim_head ** -0.5)
|
||||
|
||||
h = heads
|
||||
if skip_reshape:
|
||||
@@ -219,6 +226,10 @@ def attention_sub_quad(query, key, value, heads, mask=None, attn_precision=None,
|
||||
b, _, dim_head = query.shape
|
||||
dim_head //= heads
|
||||
|
||||
if "scale" in kwargs:
|
||||
# Pre-scale query to match requested scale (cancels internal 1/sqrt(dim_head))
|
||||
query = query * (kwargs["scale"] * dim_head ** 0.5)
|
||||
|
||||
if skip_reshape:
|
||||
query = query.reshape(b * heads, -1, dim_head)
|
||||
value = value.reshape(b * heads, -1, dim_head)
|
||||
@@ -290,7 +301,7 @@ def attention_split(q, k, v, heads, mask=None, attn_precision=None, skip_reshape
|
||||
b, _, dim_head = q.shape
|
||||
dim_head //= heads
|
||||
|
||||
scale = dim_head ** -0.5
|
||||
scale = kwargs.get("scale", dim_head ** -0.5)
|
||||
|
||||
if skip_reshape:
|
||||
q, k, v = map(
|
||||
@@ -500,8 +511,13 @@ def attention_pytorch(q, k, v, heads, mask=None, attn_precision=None, skip_resha
|
||||
if mask.ndim == 3:
|
||||
mask = mask.unsqueeze(1)
|
||||
|
||||
# Pass through extra SDPA kwargs (scale, enable_gqa) if provided
|
||||
# enable_gqa requires PyTorch 2.5+; older versions use manual KV expansion above
|
||||
sdpa_keys = ("scale", "enable_gqa") if TORCH_HAS_GQA else ("scale",)
|
||||
sdpa_extra = {k: v for k, v in kwargs.items() if k in sdpa_keys}
|
||||
|
||||
if SDP_BATCH_LIMIT >= b:
|
||||
out = comfy.ops.scaled_dot_product_attention(q, k, v, attn_mask=mask, dropout_p=0.0, is_causal=False)
|
||||
out = comfy.ops.scaled_dot_product_attention(q, k, v, attn_mask=mask, dropout_p=0.0, is_causal=False, **sdpa_extra)
|
||||
if not skip_output_reshape:
|
||||
out = (
|
||||
out.transpose(1, 2).reshape(b, -1, heads * dim_head)
|
||||
@@ -519,7 +535,7 @@ def attention_pytorch(q, k, v, heads, mask=None, attn_precision=None, skip_resha
|
||||
k[i : i + SDP_BATCH_LIMIT],
|
||||
v[i : i + SDP_BATCH_LIMIT],
|
||||
attn_mask=m,
|
||||
dropout_p=0.0, is_causal=False
|
||||
dropout_p=0.0, is_causal=False, **sdpa_extra
|
||||
).transpose(1, 2).reshape(-1, q.shape[2], heads * dim_head)
|
||||
return out
|
||||
|
||||
|
||||
@@ -17,6 +17,7 @@
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
import comfy.memory_management
|
||||
import comfy.utils
|
||||
import comfy.model_management
|
||||
import comfy.model_base
|
||||
@@ -342,6 +343,12 @@ def model_lora_keys_unet(model, key_map={}):
|
||||
key_map["base_model.model.{}".format(key_lora)] = k # Official base model loras
|
||||
key_map["lycoris_{}".format(key_lora.replace(".", "_"))] = k # LyCORIS/LoKR format
|
||||
|
||||
if isinstance(model, comfy.model_base.ErnieImage):
|
||||
for k in sdk:
|
||||
if k.startswith("diffusion_model.") and k.endswith(".weight"):
|
||||
key_lora = k[len("diffusion_model."):-len(".weight")]
|
||||
key_map["transformer.{}".format(key_lora)] = k
|
||||
|
||||
return key_map
|
||||
|
||||
|
||||
@@ -467,3 +474,17 @@ def calculate_weight(patches, weight, key, intermediate_dtype=torch.float32, ori
|
||||
weight = old_weight
|
||||
|
||||
return weight
|
||||
|
||||
def prefetch_prepared_value(value, allocate_buffer, stream):
|
||||
if isinstance(value, torch.Tensor):
|
||||
dest = allocate_buffer(comfy.memory_management.vram_aligned_size(value))
|
||||
comfy.model_management.cast_to_gathered([value], dest, non_blocking=True, stream=stream)
|
||||
return comfy.memory_management.interpret_gathered_like([value], dest)[0]
|
||||
elif isinstance(value, weight_adapter.WeightAdapterBase):
|
||||
return type(value)(value.loaded_keys, prefetch_prepared_value(value.weights, allocate_buffer, stream))
|
||||
elif isinstance(value, tuple):
|
||||
return tuple(prefetch_prepared_value(item, allocate_buffer, stream) for item in value)
|
||||
elif isinstance(value, list):
|
||||
return [prefetch_prepared_value(item, allocate_buffer, stream) for item in value]
|
||||
|
||||
return value
|
||||
|
||||
+6
-14
@@ -20,7 +20,6 @@ import comfy.ldm.hunyuan3dv2_1
|
||||
import comfy.ldm.hunyuan3dv2_1.hunyuandit
|
||||
import torch
|
||||
import logging
|
||||
import os
|
||||
import comfy.ldm.lightricks.av_model
|
||||
import comfy.context_windows
|
||||
from comfy.ldm.modules.diffusionmodules.openaimodel import UNetModel, Timestep
|
||||
@@ -121,20 +120,8 @@ def model_sampling(model_config, model_type):
|
||||
elif model_type == ModelType.V_PREDICTION_DDPM:
|
||||
c = comfy.model_sampling.V_PREDICTION_DDPM
|
||||
|
||||
from comfy.cli_args import args
|
||||
isolation_runtime_enabled = args.use_process_isolation or os.environ.get("PYISOLATE_CHILD") == "1"
|
||||
|
||||
class ModelSampling(s, c):
|
||||
if isolation_runtime_enabled:
|
||||
def __reduce__(self):
|
||||
"""Ensure pickling yields a proxy instead of failing on local class."""
|
||||
try:
|
||||
from comfy.isolation.model_sampling_proxy import ModelSamplingRegistry, ModelSamplingProxy
|
||||
registry = ModelSamplingRegistry()
|
||||
ms_id = registry.register(self)
|
||||
return (ModelSamplingProxy, (ms_id,))
|
||||
except Exception as exc:
|
||||
raise RuntimeError("Failed to serialize ModelSampling for isolation.") from exc
|
||||
pass
|
||||
|
||||
return ModelSampling(model_config)
|
||||
|
||||
@@ -227,6 +214,11 @@ class BaseModel(torch.nn.Module):
|
||||
if "latent_shapes" in extra_conds:
|
||||
xc = utils.unpack_latents(xc, extra_conds.pop("latent_shapes"))
|
||||
|
||||
transformer_options = transformer_options.copy()
|
||||
transformer_options["prefetch_dynamic_vbars"] = (
|
||||
self.current_patcher is not None and self.current_patcher.is_dynamic()
|
||||
)
|
||||
|
||||
model_output = self.diffusion_model(xc, t, context=context, control=control, transformer_options=transformer_options, **extra_conds)
|
||||
if len(model_output) > 1 and not torch.is_tensor(model_output):
|
||||
model_output, _ = utils.pack_latents(model_output)
|
||||
|
||||
+40
-111
@@ -31,6 +31,7 @@ from contextlib import nullcontext
|
||||
import comfy.memory_management
|
||||
import comfy.utils
|
||||
import comfy.quant_ops
|
||||
import comfy_aimdo.vram_buffer
|
||||
|
||||
class VRAMState(Enum):
|
||||
DISABLED = 0 #No vram present: no need to move models to vram
|
||||
@@ -112,10 +113,6 @@ if args.directml is not None:
|
||||
# torch_directml.disable_tiled_resources(True)
|
||||
lowvram_available = False #TODO: need to find a way to get free memory in directml before this can be enabled by default.
|
||||
|
||||
try:
|
||||
import intel_extension_for_pytorch as ipex # noqa: F401
|
||||
except:
|
||||
pass
|
||||
|
||||
try:
|
||||
_ = torch.xpu.device_count()
|
||||
@@ -498,9 +495,6 @@ except:
|
||||
|
||||
current_loaded_models = []
|
||||
|
||||
def _isolation_mode_enabled():
|
||||
return args.use_process_isolation or os.environ.get("PYISOLATE_CHILD") == "1"
|
||||
|
||||
def module_size(module):
|
||||
module_mem = 0
|
||||
sd = module.state_dict()
|
||||
@@ -586,9 +580,6 @@ class LoadedModel:
|
||||
|
||||
real_model = self.model.model
|
||||
|
||||
if is_intel_xpu() and not args.disable_ipex_optimize and 'ipex' in globals() and real_model is not None:
|
||||
with torch.no_grad():
|
||||
real_model = ipex.optimize(real_model.eval(), inplace=True, graph_mode=True, concat_linear=True)
|
||||
|
||||
self.real_model = weakref.ref(real_model)
|
||||
self.model_finalizer = weakref.finalize(real_model, cleanup_models)
|
||||
@@ -607,9 +598,8 @@ class LoadedModel:
|
||||
if freed >= memory_to_free:
|
||||
return False
|
||||
self.model.detach(unpatch_weights)
|
||||
if self.model_finalizer is not None:
|
||||
self.model_finalizer.detach()
|
||||
self.model_finalizer = None
|
||||
self.model_finalizer.detach()
|
||||
self.model_finalizer = None
|
||||
self.real_model = None
|
||||
return True
|
||||
|
||||
@@ -623,15 +613,8 @@ class LoadedModel:
|
||||
if self._patcher_finalizer is not None:
|
||||
self._patcher_finalizer.detach()
|
||||
|
||||
def dead_state(self):
|
||||
model_ref_gone = self.model is None
|
||||
real_model_ref = self.real_model
|
||||
real_model_ref_gone = callable(real_model_ref) and real_model_ref() is None
|
||||
return model_ref_gone, real_model_ref_gone
|
||||
|
||||
def is_dead(self):
|
||||
model_ref_gone, real_model_ref_gone = self.dead_state()
|
||||
return model_ref_gone or real_model_ref_gone
|
||||
return self.real_model() is not None and self.model is None
|
||||
|
||||
|
||||
def use_more_memory(extra_memory, loaded_models, device):
|
||||
@@ -678,7 +661,6 @@ def free_memory(memory_required, device, keep_loaded=[], for_dynamic=False, pins
|
||||
unloaded_model = []
|
||||
can_unload = []
|
||||
unloaded_models = []
|
||||
isolation_active = _isolation_mode_enabled()
|
||||
|
||||
for i in range(len(current_loaded_models) -1, -1, -1):
|
||||
shift_model = current_loaded_models[i]
|
||||
@@ -687,17 +669,6 @@ def free_memory(memory_required, device, keep_loaded=[], for_dynamic=False, pins
|
||||
can_unload.append((-shift_model.model_offloaded_memory(), sys.getrefcount(shift_model.model), shift_model.model_memory(), i))
|
||||
shift_model.currently_used = False
|
||||
|
||||
if can_unload and isolation_active:
|
||||
try:
|
||||
from pyisolate import flush_tensor_keeper # type: ignore[attr-defined]
|
||||
except Exception:
|
||||
flush_tensor_keeper = None
|
||||
if callable(flush_tensor_keeper):
|
||||
flushed = flush_tensor_keeper()
|
||||
if flushed > 0:
|
||||
logging.debug("][ MM:tensor_keeper_flush | released=%d", flushed)
|
||||
gc.collect()
|
||||
|
||||
can_unload_sorted = sorted(can_unload)
|
||||
for x in can_unload_sorted:
|
||||
i = x[-1]
|
||||
@@ -728,13 +699,7 @@ def free_memory(memory_required, device, keep_loaded=[], for_dynamic=False, pins
|
||||
logging.debug(f"RAM Unloading {current_loaded_models[i].model.model.__class__.__name__}")
|
||||
|
||||
for i in sorted(unloaded_model, reverse=True):
|
||||
unloaded = current_loaded_models.pop(i)
|
||||
model_obj = unloaded.model
|
||||
if model_obj is not None:
|
||||
cleanup = getattr(model_obj, "cleanup", None)
|
||||
if callable(cleanup):
|
||||
cleanup()
|
||||
unloaded_models.append(unloaded)
|
||||
unloaded_models.append(current_loaded_models.pop(i))
|
||||
|
||||
if len(unloaded_model) > 0:
|
||||
soft_empty_cache()
|
||||
@@ -793,9 +758,7 @@ def load_models_gpu(models, memory_required=0, force_patch_weights=False, minimu
|
||||
for i in to_unload:
|
||||
model_to_unload = current_loaded_models.pop(i)
|
||||
model_to_unload.model.detach(unpatch_all=False)
|
||||
if model_to_unload.model_finalizer is not None:
|
||||
model_to_unload.model_finalizer.detach()
|
||||
model_to_unload.model_finalizer = None
|
||||
model_to_unload.model_finalizer.detach()
|
||||
|
||||
|
||||
total_memory_required = {}
|
||||
@@ -868,62 +831,25 @@ def loaded_models(only_currently_used=False):
|
||||
|
||||
|
||||
def cleanup_models_gc():
|
||||
do_gc = False
|
||||
|
||||
reset_cast_buffers()
|
||||
if not _isolation_mode_enabled():
|
||||
dead_found = False
|
||||
for i in range(len(current_loaded_models)):
|
||||
if current_loaded_models[i].is_dead():
|
||||
dead_found = True
|
||||
break
|
||||
|
||||
if dead_found:
|
||||
logging.info("Potential memory leak detected with model NoneType, doing a full garbage collect, for maximum performance avoid circular references in the model code.")
|
||||
gc.collect()
|
||||
soft_empty_cache()
|
||||
|
||||
for i in range(len(current_loaded_models) - 1, -1, -1):
|
||||
cur = current_loaded_models[i]
|
||||
if cur.is_dead():
|
||||
logging.warning("WARNING, memory leak with model NoneType. Please make sure it is not being referenced from somewhere.")
|
||||
leaked = current_loaded_models.pop(i)
|
||||
model_obj = getattr(leaked, "model", None)
|
||||
if model_obj is not None:
|
||||
cleanup = getattr(model_obj, "cleanup", None)
|
||||
if callable(cleanup):
|
||||
cleanup()
|
||||
return
|
||||
|
||||
dead_found = False
|
||||
has_real_model_leak = False
|
||||
for i in range(len(current_loaded_models)):
|
||||
model_ref_gone, real_model_ref_gone = current_loaded_models[i].dead_state()
|
||||
if model_ref_gone or real_model_ref_gone:
|
||||
dead_found = True
|
||||
if real_model_ref_gone and not model_ref_gone:
|
||||
has_real_model_leak = True
|
||||
cur = current_loaded_models[i]
|
||||
if cur.is_dead():
|
||||
logging.info("Potential memory leak detected with model {}, doing a full garbage collect, for maximum performance avoid circular references in the model code.".format(cur.real_model().__class__.__name__))
|
||||
do_gc = True
|
||||
break
|
||||
|
||||
if dead_found:
|
||||
if has_real_model_leak:
|
||||
logging.info("Potential memory leak detected with model NoneType, doing a full garbage collect, for maximum performance avoid circular references in the model code.")
|
||||
else:
|
||||
logging.debug("Cleaning stale loaded-model entries with released patcher references.")
|
||||
if do_gc:
|
||||
gc.collect()
|
||||
soft_empty_cache()
|
||||
|
||||
for i in range(len(current_loaded_models) - 1, -1, -1):
|
||||
for i in range(len(current_loaded_models)):
|
||||
cur = current_loaded_models[i]
|
||||
model_ref_gone, real_model_ref_gone = cur.dead_state()
|
||||
if model_ref_gone or real_model_ref_gone:
|
||||
if real_model_ref_gone and not model_ref_gone:
|
||||
logging.warning("WARNING, memory leak with model NoneType. Please make sure it is not being referenced from somewhere.")
|
||||
else:
|
||||
logging.debug("Cleaning stale loaded-model entry with released patcher reference.")
|
||||
leaked = current_loaded_models.pop(i)
|
||||
model_obj = getattr(leaked, "model", None)
|
||||
if model_obj is not None:
|
||||
cleanup = getattr(model_obj, "cleanup", None)
|
||||
if callable(cleanup):
|
||||
cleanup()
|
||||
if cur.is_dead():
|
||||
logging.warning("WARNING, memory leak with model {}. Please make sure it is not being referenced from somewhere.".format(cur.real_model().__class__.__name__))
|
||||
|
||||
|
||||
def archive_model_dtypes(model):
|
||||
@@ -937,20 +863,11 @@ def archive_model_dtypes(model):
|
||||
def cleanup_models():
|
||||
to_delete = []
|
||||
for i in range(len(current_loaded_models)):
|
||||
real_model_ref = current_loaded_models[i].real_model
|
||||
if real_model_ref is None:
|
||||
to_delete = [i] + to_delete
|
||||
continue
|
||||
if callable(real_model_ref) and real_model_ref() is None:
|
||||
if current_loaded_models[i].real_model() is None:
|
||||
to_delete = [i] + to_delete
|
||||
|
||||
for i in to_delete:
|
||||
x = current_loaded_models.pop(i)
|
||||
model_obj = getattr(x, "model", None)
|
||||
if model_obj is not None:
|
||||
cleanup = getattr(model_obj, "cleanup", None)
|
||||
if callable(cleanup):
|
||||
cleanup()
|
||||
del x
|
||||
|
||||
def dtype_size(dtype):
|
||||
@@ -1259,6 +1176,10 @@ stream_counters = {}
|
||||
|
||||
STREAM_CAST_BUFFERS = {}
|
||||
LARGEST_CASTED_WEIGHT = (None, 0)
|
||||
STREAM_AIMDO_CAST_BUFFERS = {}
|
||||
LARGEST_AIMDO_CASTED_WEIGHT = (None, 0)
|
||||
|
||||
DEFAULT_AIMDO_CAST_BUFFER_RESERVATION_SIZE = 16 * 1024 ** 3
|
||||
|
||||
def get_cast_buffer(offload_stream, device, size, ref):
|
||||
global LARGEST_CASTED_WEIGHT
|
||||
@@ -1292,13 +1213,26 @@ def get_cast_buffer(offload_stream, device, size, ref):
|
||||
|
||||
return cast_buffer
|
||||
|
||||
def get_aimdo_cast_buffer(offload_stream, device):
|
||||
cast_buffer = STREAM_AIMDO_CAST_BUFFERS.get(offload_stream, None)
|
||||
if cast_buffer is None:
|
||||
cast_buffer = comfy_aimdo.vram_buffer.VRAMBuffer(DEFAULT_AIMDO_CAST_BUFFER_RESERVATION_SIZE, device.index)
|
||||
STREAM_AIMDO_CAST_BUFFERS[offload_stream] = cast_buffer
|
||||
|
||||
return cast_buffer
|
||||
def reset_cast_buffers():
|
||||
global LARGEST_CASTED_WEIGHT
|
||||
global LARGEST_AIMDO_CASTED_WEIGHT
|
||||
|
||||
LARGEST_CASTED_WEIGHT = (None, 0)
|
||||
for offload_stream in STREAM_CAST_BUFFERS:
|
||||
offload_stream.synchronize()
|
||||
LARGEST_AIMDO_CASTED_WEIGHT = (None, 0)
|
||||
for offload_stream in set(STREAM_CAST_BUFFERS) | set(STREAM_AIMDO_CAST_BUFFERS):
|
||||
if offload_stream is not None:
|
||||
offload_stream.synchronize()
|
||||
synchronize()
|
||||
|
||||
STREAM_CAST_BUFFERS.clear()
|
||||
STREAM_AIMDO_CAST_BUFFERS.clear()
|
||||
soft_empty_cache()
|
||||
|
||||
def get_offload_stream(device):
|
||||
@@ -1658,10 +1592,7 @@ def should_use_fp16(device=None, model_params=0, prioritize_performance=True, ma
|
||||
return False
|
||||
|
||||
if is_intel_xpu():
|
||||
if torch_version_numeric < (2, 3):
|
||||
return True
|
||||
else:
|
||||
return torch.xpu.get_device_properties(device).has_fp16
|
||||
return torch.xpu.get_device_properties(device).has_fp16
|
||||
|
||||
if is_ascend_npu():
|
||||
return True
|
||||
@@ -1727,10 +1658,7 @@ def should_use_bf16(device=None, model_params=0, prioritize_performance=True, ma
|
||||
return False
|
||||
|
||||
if is_intel_xpu():
|
||||
if torch_version_numeric < (2, 3):
|
||||
return True
|
||||
else:
|
||||
return torch.xpu.is_bf16_supported()
|
||||
return torch.xpu.is_bf16_supported()
|
||||
|
||||
if is_ascend_npu():
|
||||
return True
|
||||
@@ -1861,6 +1789,7 @@ def soft_empty_cache(force=False):
|
||||
if cpu_state == CPUState.MPS:
|
||||
torch.mps.empty_cache()
|
||||
elif is_intel_xpu():
|
||||
torch.xpu.synchronize()
|
||||
torch.xpu.empty_cache()
|
||||
elif is_ascend_npu():
|
||||
torch.npu.empty_cache()
|
||||
|
||||
+12
-1
@@ -121,9 +121,20 @@ class LowVramPatch:
|
||||
self.patches = patches
|
||||
self.convert_func = convert_func # TODO: remove
|
||||
self.set_func = set_func
|
||||
self.prepared_patches = None
|
||||
|
||||
def prepare(self, allocate_buffer, stream):
|
||||
self.prepared_patches = [
|
||||
(patch[0], comfy.lora.prefetch_prepared_value(patch[1], allocate_buffer, stream), patch[2], patch[3], patch[4])
|
||||
for patch in self.patches[self.key]
|
||||
]
|
||||
|
||||
def clear_prepared(self):
|
||||
self.prepared_patches = None
|
||||
|
||||
def __call__(self, weight):
|
||||
return comfy.lora.calculate_weight(self.patches[self.key], weight, self.key, intermediate_dtype=weight.dtype)
|
||||
patches = self.prepared_patches if self.prepared_patches is not None else self.patches[self.key]
|
||||
return comfy.lora.calculate_weight(patches, weight, self.key, intermediate_dtype=weight.dtype)
|
||||
|
||||
LOWVRAM_PATCH_ESTIMATE_MATH_FACTOR = 2
|
||||
|
||||
|
||||
@@ -0,0 +1,65 @@
|
||||
import comfy_aimdo.model_vbar
|
||||
import comfy.model_management
|
||||
import comfy.ops
|
||||
|
||||
PREFETCH_QUEUES = []
|
||||
|
||||
def cleanup_prefetched_modules(comfy_modules):
|
||||
for s in comfy_modules:
|
||||
prefetch = getattr(s, "_prefetch", None)
|
||||
if prefetch is None:
|
||||
continue
|
||||
for param_key in ("weight", "bias"):
|
||||
lowvram_fn = getattr(s, param_key + "_lowvram_function", None)
|
||||
if lowvram_fn is not None:
|
||||
lowvram_fn.clear_prepared()
|
||||
if prefetch["signature"] is not None:
|
||||
comfy_aimdo.model_vbar.vbar_unpin(s._v)
|
||||
delattr(s, "_prefetch")
|
||||
|
||||
def cleanup_prefetch_queues():
|
||||
global PREFETCH_QUEUES
|
||||
|
||||
for queue in PREFETCH_QUEUES:
|
||||
for entry in queue:
|
||||
if entry is None or not isinstance(entry, tuple):
|
||||
continue
|
||||
_, prefetch_state = entry
|
||||
comfy_modules = prefetch_state[1]
|
||||
if comfy_modules is not None:
|
||||
cleanup_prefetched_modules(comfy_modules)
|
||||
PREFETCH_QUEUES = []
|
||||
|
||||
def prefetch_queue_pop(queue, device, module):
|
||||
if queue is None:
|
||||
return
|
||||
|
||||
consumed = queue.pop(0)
|
||||
if consumed is not None:
|
||||
offload_stream, prefetch_state = consumed
|
||||
offload_stream.wait_stream(comfy.model_management.current_stream(device))
|
||||
_, comfy_modules = prefetch_state
|
||||
if comfy_modules is not None:
|
||||
cleanup_prefetched_modules(comfy_modules)
|
||||
|
||||
prefetch = queue[0]
|
||||
if prefetch is not None:
|
||||
comfy_modules = []
|
||||
for s in prefetch.modules():
|
||||
if hasattr(s, "_v"):
|
||||
comfy_modules.append(s)
|
||||
|
||||
offload_stream = comfy.ops.cast_modules_with_vbar(comfy_modules, None, device, None, True)
|
||||
comfy.model_management.sync_stream(device, offload_stream)
|
||||
queue[0] = (offload_stream, (prefetch, comfy_modules))
|
||||
|
||||
def make_prefetch_queue(queue, device, transformer_options):
|
||||
if (not transformer_options.get("prefetch_dynamic_vbars", False)
|
||||
or comfy.model_management.NUM_STREAMS == 0
|
||||
or comfy.model_management.is_device_cpu(device)
|
||||
or not comfy.model_management.device_supports_non_blocking(device)):
|
||||
return None
|
||||
|
||||
queue = [None] + queue + [None]
|
||||
PREFETCH_QUEUES.append(queue)
|
||||
return queue
|
||||
+214
-54
@@ -86,38 +86,61 @@ def materialize_meta_param(s, param_keys):
|
||||
setattr(s, param_key, torch.nn.Parameter(torch.zeros(param.shape, dtype=param.dtype), requires_grad=param.requires_grad))
|
||||
|
||||
|
||||
def cast_bias_weight_with_vbar(s, dtype, device, bias_dtype, non_blocking, compute_dtype, want_requant):
|
||||
#vbar doesn't support CPU weights, but some custom nodes have weird paths
|
||||
#that might switch the layer to the CPU and expect it to work. We have to take
|
||||
#a clone conservatively as we are mmapped and some SFT files are packed misaligned
|
||||
#If you are a custom node author reading this, please move your layer to the GPU
|
||||
#or declare your ModelPatcher as CPU in the first place.
|
||||
if comfy.model_management.is_device_cpu(device):
|
||||
materialize_meta_param(s, ["weight", "bias"])
|
||||
weight = s.weight.to(dtype=dtype, copy=True)
|
||||
if isinstance(weight, QuantizedTensor):
|
||||
weight = weight.dequantize()
|
||||
bias = None
|
||||
if s.bias is not None:
|
||||
bias = s.bias.to(dtype=bias_dtype, copy=True)
|
||||
return weight, bias, (None, None, None)
|
||||
|
||||
# FIXME: add n=1 cache hit fast path
|
||||
def cast_modules_with_vbar(comfy_modules, dtype, device, bias_dtype, non_blocking):
|
||||
offload_stream = None
|
||||
xfer_dest = None
|
||||
cast_buffer = None
|
||||
cast_buffer_offset = 0
|
||||
|
||||
def ensure_offload_stream(module, required_size, check_largest):
|
||||
nonlocal offload_stream
|
||||
nonlocal cast_buffer
|
||||
|
||||
if offload_stream is None:
|
||||
offload_stream = comfy.model_management.get_offload_stream(device)
|
||||
if offload_stream is None or not check_largest or len(comfy_modules) != 1:
|
||||
return
|
||||
|
||||
current_size = 0 if cast_buffer is None else cast_buffer.size()
|
||||
if current_size < required_size and module is comfy.model_management.LARGEST_AIMDO_CASTED_WEIGHT[0]:
|
||||
offload_stream = comfy.model_management.get_offload_stream(device)
|
||||
cast_buffer = None
|
||||
if required_size > comfy.model_management.LARGEST_AIMDO_CASTED_WEIGHT[1]:
|
||||
comfy.model_management.LARGEST_AIMDO_CASTED_WEIGHT = (module, required_size)
|
||||
|
||||
def get_cast_buffer(buffer_size):
|
||||
nonlocal offload_stream
|
||||
nonlocal cast_buffer
|
||||
nonlocal cast_buffer_offset
|
||||
|
||||
if buffer_size == 0:
|
||||
return None
|
||||
|
||||
if offload_stream is None:
|
||||
return torch.empty((buffer_size,), dtype=torch.uint8, device=device)
|
||||
|
||||
cast_buffer = comfy.model_management.get_aimdo_cast_buffer(offload_stream, device)
|
||||
buffer = comfy_aimdo.torch.aimdo_to_tensor(cast_buffer.get(buffer_size, cast_buffer_offset), device)
|
||||
cast_buffer_offset += buffer_size
|
||||
return buffer
|
||||
|
||||
for s in comfy_modules:
|
||||
signature = comfy_aimdo.model_vbar.vbar_fault(s._v)
|
||||
resident = comfy_aimdo.model_vbar.vbar_signature_compare(signature, s._v_signature)
|
||||
prefetch = {
|
||||
"signature": signature,
|
||||
"resident": resident,
|
||||
}
|
||||
|
||||
signature = comfy_aimdo.model_vbar.vbar_fault(s._v)
|
||||
resident = comfy_aimdo.model_vbar.vbar_signature_compare(signature, s._v_signature)
|
||||
if signature is not None:
|
||||
if resident:
|
||||
weight = s._v_weight
|
||||
bias = s._v_bias
|
||||
else:
|
||||
xfer_dest = comfy_aimdo.torch.aimdo_to_tensor(s._v, device)
|
||||
s._prefetch = prefetch
|
||||
continue
|
||||
|
||||
if not resident:
|
||||
materialize_meta_param(s, ["weight", "bias"])
|
||||
xfer_dest = comfy_aimdo.torch.aimdo_to_tensor(s._v, device) if signature is not None else None
|
||||
cast_geometry = comfy.memory_management.tensors_to_geometries([ s.weight, s.bias ])
|
||||
cast_dest = None
|
||||
needs_cast = False
|
||||
|
||||
xfer_source = [ s.weight, s.bias ]
|
||||
|
||||
@@ -129,22 +152,15 @@ def cast_bias_weight_with_vbar(s, dtype, device, bias_dtype, non_blocking, compu
|
||||
if data is None:
|
||||
continue
|
||||
if data.dtype != geometry.dtype:
|
||||
needs_cast = True
|
||||
cast_dest = xfer_dest
|
||||
if cast_dest is None:
|
||||
cast_dest = torch.empty((comfy.memory_management.vram_aligned_size(cast_geometry),), dtype=torch.uint8, device=device)
|
||||
xfer_dest = None
|
||||
break
|
||||
|
||||
dest_size = comfy.memory_management.vram_aligned_size(xfer_source)
|
||||
offload_stream = comfy.model_management.get_offload_stream(device)
|
||||
if xfer_dest is None and offload_stream is not None:
|
||||
xfer_dest = comfy.model_management.get_cast_buffer(offload_stream, device, dest_size, s)
|
||||
if xfer_dest is None:
|
||||
offload_stream = comfy.model_management.get_offload_stream(device)
|
||||
xfer_dest = comfy.model_management.get_cast_buffer(offload_stream, device, dest_size, s)
|
||||
ensure_offload_stream(s, dest_size if xfer_dest is None else 0, True)
|
||||
if xfer_dest is None:
|
||||
xfer_dest = torch.empty((dest_size,), dtype=torch.uint8, device=device)
|
||||
offload_stream = None
|
||||
xfer_dest = get_cast_buffer(dest_size)
|
||||
|
||||
if signature is None and pin is None:
|
||||
comfy.pinned_memory.pin_memory(s)
|
||||
@@ -157,27 +173,54 @@ def cast_bias_weight_with_vbar(s, dtype, device, bias_dtype, non_blocking, compu
|
||||
xfer_source = [ pin ]
|
||||
#send it over
|
||||
comfy.model_management.cast_to_gathered(xfer_source, xfer_dest, non_blocking=non_blocking, stream=offload_stream)
|
||||
comfy.model_management.sync_stream(device, offload_stream)
|
||||
|
||||
if cast_dest is not None:
|
||||
for param_key in ("weight", "bias"):
|
||||
lowvram_fn = getattr(s, param_key + "_lowvram_function", None)
|
||||
if lowvram_fn is not None:
|
||||
ensure_offload_stream(s, cast_buffer_offset, False)
|
||||
lowvram_fn.prepare(lambda size: get_cast_buffer(size), offload_stream)
|
||||
|
||||
prefetch["xfer_dest"] = xfer_dest
|
||||
prefetch["cast_dest"] = cast_dest
|
||||
prefetch["cast_geometry"] = cast_geometry
|
||||
prefetch["needs_cast"] = needs_cast
|
||||
s._prefetch = prefetch
|
||||
|
||||
return offload_stream
|
||||
|
||||
|
||||
def resolve_cast_module_with_vbar(s, dtype, device, bias_dtype, compute_dtype, want_requant):
|
||||
|
||||
prefetch = getattr(s, "_prefetch", None)
|
||||
|
||||
if prefetch["resident"]:
|
||||
weight = s._v_weight
|
||||
bias = s._v_bias
|
||||
else:
|
||||
xfer_dest = prefetch["xfer_dest"]
|
||||
if prefetch["needs_cast"]:
|
||||
cast_dest = prefetch["cast_dest"] if prefetch["cast_dest"] is not None else torch.empty((comfy.memory_management.vram_aligned_size(prefetch["cast_geometry"]),), dtype=torch.uint8, device=device)
|
||||
for pre_cast, post_cast in zip(comfy.memory_management.interpret_gathered_like([s.weight, s.bias ], xfer_dest),
|
||||
comfy.memory_management.interpret_gathered_like(cast_geometry, cast_dest)):
|
||||
comfy.memory_management.interpret_gathered_like(prefetch["cast_geometry"], cast_dest)):
|
||||
if post_cast is not None:
|
||||
post_cast.copy_(pre_cast)
|
||||
xfer_dest = cast_dest
|
||||
|
||||
params = comfy.memory_management.interpret_gathered_like(cast_geometry, xfer_dest)
|
||||
params = comfy.memory_management.interpret_gathered_like(prefetch["cast_geometry"], xfer_dest)
|
||||
weight = params[0]
|
||||
bias = params[1]
|
||||
if signature is not None:
|
||||
if prefetch["signature"] is not None:
|
||||
s._v_weight = weight
|
||||
s._v_bias = bias
|
||||
s._v_signature=signature
|
||||
s._v_signature = prefetch["signature"]
|
||||
|
||||
def post_cast(s, param_key, x, dtype, resident, update_weight):
|
||||
lowvram_fn = getattr(s, param_key + "_lowvram_function", None)
|
||||
fns = getattr(s, param_key + "_function", [])
|
||||
|
||||
if x is None:
|
||||
return None
|
||||
|
||||
orig = x
|
||||
|
||||
def to_dequant(tensor, dtype):
|
||||
@@ -205,14 +248,12 @@ def cast_bias_weight_with_vbar(s, dtype, device, bias_dtype, non_blocking, compu
|
||||
x = f(x)
|
||||
return x
|
||||
|
||||
update_weight = signature is not None
|
||||
update_weight = prefetch["signature"] is not None
|
||||
weight = post_cast(s, "weight", weight, dtype, prefetch["resident"], update_weight)
|
||||
if bias is not None:
|
||||
bias = post_cast(s, "bias", bias, bias_dtype, prefetch["resident"], update_weight)
|
||||
|
||||
weight = post_cast(s, "weight", weight, dtype, resident, update_weight)
|
||||
if s.bias is not None:
|
||||
bias = post_cast(s, "bias", bias, bias_dtype, resident, update_weight)
|
||||
|
||||
#FIXME: weird offload return protocol
|
||||
return weight, bias, (offload_stream, device if signature is not None else None, None)
|
||||
return weight, bias
|
||||
|
||||
|
||||
def cast_bias_weight(s, input=None, dtype=None, device=None, bias_dtype=None, offloadable=False, compute_dtype=None, want_requant=False):
|
||||
@@ -230,10 +271,46 @@ def cast_bias_weight(s, input=None, dtype=None, device=None, bias_dtype=None, of
|
||||
if device is None:
|
||||
device = input.device
|
||||
|
||||
def format_return(result, offloadable):
|
||||
weight, bias, offload_stream = result
|
||||
return (weight, bias, offload_stream) if offloadable else (weight, bias)
|
||||
|
||||
non_blocking = comfy.model_management.device_supports_non_blocking(device)
|
||||
|
||||
if hasattr(s, "_v"):
|
||||
return cast_bias_weight_with_vbar(s, dtype, device, bias_dtype, non_blocking, compute_dtype, want_requant)
|
||||
|
||||
#vbar doesn't support CPU weights, but some custom nodes have weird paths
|
||||
#that might switch the layer to the CPU and expect it to work. We have to take
|
||||
#a clone conservatively as we are mmapped and some SFT files are packed misaligned
|
||||
#If you are a custom node author reading this, please move your layer to the GPU
|
||||
#or declare your ModelPatcher as CPU in the first place.
|
||||
if comfy.model_management.is_device_cpu(device):
|
||||
materialize_meta_param(s, ["weight", "bias"])
|
||||
weight = s.weight.to(dtype=dtype, copy=True)
|
||||
if isinstance(weight, QuantizedTensor):
|
||||
weight = weight.dequantize()
|
||||
bias = s.bias.to(dtype=bias_dtype, copy=True) if s.bias is not None else None
|
||||
return format_return((weight, bias, (None, None, None)), offloadable)
|
||||
|
||||
prefetched = hasattr(s, "_prefetch")
|
||||
offload_stream = None
|
||||
offload_device = None
|
||||
if not prefetched:
|
||||
offload_stream = cast_modules_with_vbar([s], dtype, device, bias_dtype, non_blocking)
|
||||
comfy.model_management.sync_stream(device, offload_stream)
|
||||
|
||||
weight, bias = resolve_cast_module_with_vbar(s, dtype, device, bias_dtype, compute_dtype, want_requant)
|
||||
|
||||
if not prefetched:
|
||||
if getattr(s, "_prefetch")["signature"] is not None:
|
||||
offload_device = device
|
||||
for param_key in ("weight", "bias"):
|
||||
lowvram_fn = getattr(s, param_key + "_lowvram_function", None)
|
||||
if lowvram_fn is not None:
|
||||
lowvram_fn.clear_prepared()
|
||||
delattr(s, "_prefetch")
|
||||
return format_return((weight, bias, (offload_stream, offload_device, None)), offloadable)
|
||||
|
||||
|
||||
if offloadable and (device != s.weight.device or
|
||||
(s.bias is not None and device != s.bias.device)):
|
||||
@@ -280,11 +357,7 @@ def cast_bias_weight(s, input=None, dtype=None, device=None, bias_dtype=None, of
|
||||
for f in s.weight_function:
|
||||
weight = f(weight)
|
||||
|
||||
if offloadable:
|
||||
return weight, bias, (offload_stream, weight_a, bias_a)
|
||||
else:
|
||||
#Legacy function signature
|
||||
return weight, bias
|
||||
return format_return((weight, bias, (offload_stream, weight_a, bias_a)), offloadable)
|
||||
|
||||
|
||||
def uncast_bias_weight(s, weight, bias, offload_stream):
|
||||
@@ -1173,6 +1246,93 @@ def mixed_precision_ops(quant_config={}, compute_dtype=torch.bfloat16, full_prec
|
||||
self._buffers[key] = fn(buf)
|
||||
return self
|
||||
|
||||
class Embedding(manual_cast.Embedding):
|
||||
def _load_from_state_dict(self, state_dict, prefix, local_metadata,
|
||||
strict, missing_keys, unexpected_keys, error_msgs):
|
||||
weight_key = f"{prefix}weight"
|
||||
layer_conf = state_dict.pop(f"{prefix}comfy_quant", None)
|
||||
if layer_conf is not None:
|
||||
layer_conf = json.loads(layer_conf.numpy().tobytes())
|
||||
|
||||
# Only fp8 makes sense for embeddings (per-row dequant via index select).
|
||||
# Block-scaled formats (NVFP4, MXFP8) can't do per-row lookup efficiently.
|
||||
quant_format = layer_conf.get("format", None) if layer_conf is not None else None
|
||||
if quant_format in ["float8_e4m3fn", "float8_e5m2"] and weight_key in state_dict:
|
||||
self.quant_format = quant_format
|
||||
qconfig = QUANT_ALGOS[quant_format]
|
||||
layout_cls = get_layout_class(qconfig["comfy_tensor_layout"])
|
||||
weight = state_dict.pop(weight_key)
|
||||
manually_loaded_keys = [weight_key]
|
||||
|
||||
scale_key = f"{prefix}weight_scale"
|
||||
scale = state_dict.pop(scale_key, None)
|
||||
if scale is not None:
|
||||
scale = scale.float()
|
||||
manually_loaded_keys.append(scale_key)
|
||||
|
||||
params = layout_cls.Params(
|
||||
scale=scale if scale is not None else torch.ones((), dtype=torch.float32),
|
||||
orig_dtype=MixedPrecisionOps._compute_dtype,
|
||||
orig_shape=(self.num_embeddings, self.embedding_dim),
|
||||
)
|
||||
self.weight = torch.nn.Parameter(
|
||||
QuantizedTensor(weight.to(dtype=qconfig["storage_t"]), qconfig["comfy_tensor_layout"], params),
|
||||
requires_grad=False)
|
||||
|
||||
super()._load_from_state_dict(state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs)
|
||||
for k in manually_loaded_keys:
|
||||
if k in missing_keys:
|
||||
missing_keys.remove(k)
|
||||
else:
|
||||
if layer_conf is not None:
|
||||
state_dict[f"{prefix}comfy_quant"] = torch.tensor(list(json.dumps(layer_conf).encode('utf-8')), dtype=torch.uint8)
|
||||
super()._load_from_state_dict(state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs)
|
||||
|
||||
def state_dict(self, *args, destination=None, prefix="", **kwargs):
|
||||
if destination is not None:
|
||||
sd = destination
|
||||
else:
|
||||
sd = {}
|
||||
|
||||
if not hasattr(self, 'weight') or self.weight is None:
|
||||
return sd
|
||||
|
||||
if isinstance(self.weight, QuantizedTensor):
|
||||
sd_out = self.weight.state_dict("{}weight".format(prefix))
|
||||
for k in sd_out:
|
||||
sd[k] = sd_out[k]
|
||||
|
||||
quant_conf = {"format": self.quant_format}
|
||||
sd["{}comfy_quant".format(prefix)] = torch.tensor(list(json.dumps(quant_conf).encode('utf-8')), dtype=torch.uint8)
|
||||
else:
|
||||
sd["{}weight".format(prefix)] = self.weight
|
||||
return sd
|
||||
|
||||
def forward_comfy_cast_weights(self, input, out_dtype=None):
|
||||
weight = self.weight
|
||||
|
||||
# Optimized path: lookup in fp8, dequantize only the selected rows.
|
||||
if isinstance(weight, QuantizedTensor) and len(self.weight_function) == 0:
|
||||
qdata, _, offload_stream = cast_bias_weight(self, device=input.device, dtype=weight.dtype, offloadable=True)
|
||||
if isinstance(qdata, QuantizedTensor):
|
||||
scale = qdata._params.scale
|
||||
qdata = qdata._qdata
|
||||
else:
|
||||
scale = None
|
||||
|
||||
x = torch.nn.functional.embedding(
|
||||
input, qdata, self.padding_idx, self.max_norm,
|
||||
self.norm_type, self.scale_grad_by_freq, self.sparse)
|
||||
uncast_bias_weight(self, qdata, None, offload_stream)
|
||||
target_dtype = out_dtype if out_dtype is not None else weight._params.orig_dtype
|
||||
x = x.to(dtype=target_dtype)
|
||||
if scale is not None and scale != 1.0:
|
||||
x = x * scale.to(dtype=target_dtype)
|
||||
return x
|
||||
|
||||
# Fallback for non-quantized or weight_function (LoRA) case
|
||||
return super().forward_comfy_cast_weights(input, out_dtype=out_dtype)
|
||||
|
||||
return MixedPrecisionOps
|
||||
|
||||
def pick_operations(weight_dtype, compute_dtype, load_device=None, disable_fast_fp8=False, fp8_optimizations=False, model_config=None):
|
||||
|
||||
@@ -3,6 +3,7 @@ import comfy.model_management
|
||||
|
||||
RMSNorm = torch.nn.RMSNorm
|
||||
|
||||
# Note: torch's fused F.rms_norm is faster but produces slightly different output than manual implementations (rsqrt/reduction rounding).
|
||||
def rms_norm(x, weight=None, eps=1e-6):
|
||||
if weight is None:
|
||||
return torch.nn.functional.rms_norm(x, (x.shape[-1],), eps=eps)
|
||||
|
||||
+11
-53
@@ -11,14 +11,12 @@ from functools import partial
|
||||
import collections
|
||||
import math
|
||||
import logging
|
||||
import os
|
||||
import comfy.sampler_helpers
|
||||
import comfy.model_patcher
|
||||
import comfy.patcher_extension
|
||||
import comfy.hooks
|
||||
import comfy.context_windows
|
||||
import comfy.utils
|
||||
from comfy.cli_args import args
|
||||
import scipy.stats
|
||||
import numpy
|
||||
|
||||
@@ -212,11 +210,9 @@ def _calc_cond_batch_outer(model: BaseModel, conds: list[list[dict]], x_in: torc
|
||||
_calc_cond_batch,
|
||||
comfy.patcher_extension.get_all_wrappers(comfy.patcher_extension.WrappersMP.CALC_COND_BATCH, model_options, is_model_options=True)
|
||||
)
|
||||
result = executor.execute(model, conds, x_in, timestep, model_options)
|
||||
return result
|
||||
return executor.execute(model, conds, x_in, timestep, model_options)
|
||||
|
||||
def _calc_cond_batch(model: BaseModel, conds: list[list[dict]], x_in: torch.Tensor, timestep, model_options):
|
||||
isolation_active = args.use_process_isolation or os.environ.get("PYISOLATE_CHILD") == "1"
|
||||
out_conds = []
|
||||
out_counts = []
|
||||
# separate conds by matching hooks
|
||||
@@ -273,8 +269,7 @@ def _calc_cond_batch(model: BaseModel, conds: list[list[dict]], x_in: torch.Tens
|
||||
for k, v in to_run[tt][0].conditioning.items():
|
||||
cond_shapes[k].append(v.size())
|
||||
|
||||
memory_required = model.memory_required(input_shape, cond_shapes=cond_shapes)
|
||||
if memory_required * 1.5 < free_memory:
|
||||
if model.memory_required(input_shape, cond_shapes=cond_shapes) * 1.5 < free_memory:
|
||||
to_batch = batch_amount
|
||||
break
|
||||
|
||||
@@ -299,17 +294,9 @@ def _calc_cond_batch(model: BaseModel, conds: list[list[dict]], x_in: torch.Tens
|
||||
patches = p.patches
|
||||
|
||||
batch_chunks = len(cond_or_uncond)
|
||||
if isolation_active:
|
||||
target_device = model.load_device if hasattr(model, "load_device") else input_x[0].device
|
||||
input_x = torch.cat(input_x).to(target_device)
|
||||
else:
|
||||
input_x = torch.cat(input_x)
|
||||
input_x = torch.cat(input_x)
|
||||
c = cond_cat(c)
|
||||
if isolation_active:
|
||||
timestep_ = torch.cat([timestep] * batch_chunks).to(target_device)
|
||||
mult = [m.to(target_device) if hasattr(m, "to") else m for m in mult]
|
||||
else:
|
||||
timestep_ = torch.cat([timestep] * batch_chunks)
|
||||
timestep_ = torch.cat([timestep] * batch_chunks)
|
||||
|
||||
transformer_options = model.current_patcher.apply_hooks(hooks=hooks)
|
||||
if 'transformer_options' in model_options:
|
||||
@@ -340,17 +327,9 @@ def _calc_cond_batch(model: BaseModel, conds: list[list[dict]], x_in: torch.Tens
|
||||
for o in range(batch_chunks):
|
||||
cond_index = cond_or_uncond[o]
|
||||
a = area[o]
|
||||
out_t = output[o]
|
||||
mult_t = mult[o]
|
||||
if isolation_active:
|
||||
target_dev = out_conds[cond_index].device
|
||||
if hasattr(out_t, "device") and out_t.device != target_dev:
|
||||
out_t = out_t.to(target_dev)
|
||||
if hasattr(mult_t, "device") and mult_t.device != target_dev:
|
||||
mult_t = mult_t.to(target_dev)
|
||||
if a is None:
|
||||
out_conds[cond_index] += out_t * mult_t
|
||||
out_counts[cond_index] += mult_t
|
||||
out_conds[cond_index] += output[o] * mult[o]
|
||||
out_counts[cond_index] += mult[o]
|
||||
else:
|
||||
out_c = out_conds[cond_index]
|
||||
out_cts = out_counts[cond_index]
|
||||
@@ -358,8 +337,8 @@ def _calc_cond_batch(model: BaseModel, conds: list[list[dict]], x_in: torch.Tens
|
||||
for i in range(dims):
|
||||
out_c = out_c.narrow(i + 2, a[i + dims], a[i])
|
||||
out_cts = out_cts.narrow(i + 2, a[i + dims], a[i])
|
||||
out_c += out_t * mult_t
|
||||
out_cts += mult_t
|
||||
out_c += output[o] * mult[o]
|
||||
out_cts += mult[o]
|
||||
|
||||
for i in range(len(out_conds)):
|
||||
out_conds[i] /= out_counts[i]
|
||||
@@ -413,31 +392,14 @@ class KSamplerX0Inpaint:
|
||||
self.inner_model = model
|
||||
self.sigmas = sigmas
|
||||
def __call__(self, x, sigma, denoise_mask, model_options={}, seed=None):
|
||||
isolation_active = args.use_process_isolation or os.environ.get("PYISOLATE_CHILD") == "1"
|
||||
if denoise_mask is not None:
|
||||
if isolation_active and denoise_mask.device != x.device:
|
||||
denoise_mask = denoise_mask.to(x.device)
|
||||
if "denoise_mask_function" in model_options:
|
||||
denoise_mask = model_options["denoise_mask_function"](sigma, denoise_mask, extra_options={"model": self.inner_model, "sigmas": self.sigmas})
|
||||
latent_mask = 1. - denoise_mask
|
||||
if isolation_active:
|
||||
latent_image = self.latent_image
|
||||
if hasattr(latent_image, "device") and latent_image.device != x.device:
|
||||
latent_image = latent_image.to(x.device)
|
||||
scaled = self.inner_model.inner_model.scale_latent_inpaint(x=x, sigma=sigma, noise=self.noise, latent_image=latent_image)
|
||||
if hasattr(scaled, "device") and scaled.device != x.device:
|
||||
scaled = scaled.to(x.device)
|
||||
else:
|
||||
scaled = self.inner_model.inner_model.scale_latent_inpaint(
|
||||
x=x, sigma=sigma, noise=self.noise, latent_image=self.latent_image
|
||||
)
|
||||
x = x * denoise_mask + scaled * latent_mask
|
||||
x = x * denoise_mask + self.inner_model.inner_model.scale_latent_inpaint(x=x, sigma=sigma, noise=self.noise, latent_image=self.latent_image) * latent_mask
|
||||
out = self.inner_model(x, sigma, model_options=model_options, seed=seed)
|
||||
if denoise_mask is not None:
|
||||
latent_image = self.latent_image
|
||||
if isolation_active and hasattr(latent_image, "device") and latent_image.device != out.device:
|
||||
latent_image = latent_image.to(out.device)
|
||||
out = out * denoise_mask + latent_image * latent_mask
|
||||
out = out * denoise_mask + self.latent_image * latent_mask
|
||||
return out
|
||||
|
||||
def simple_scheduler(model_sampling, steps):
|
||||
@@ -779,11 +741,7 @@ class KSAMPLER(Sampler):
|
||||
else:
|
||||
model_k.noise = noise
|
||||
|
||||
max_denoise = self.max_denoise(model_wrap, sigmas)
|
||||
model_sampling = model_wrap.inner_model.model_sampling
|
||||
noise = model_sampling.noise_scaling(
|
||||
sigmas[0], noise, latent_image, max_denoise
|
||||
)
|
||||
noise = model_wrap.inner_model.model_sampling.noise_scaling(sigmas[0], noise, latent_image, self.max_denoise(model_wrap, sigmas))
|
||||
|
||||
k_callback = None
|
||||
total_steps = len(sigmas) - 1
|
||||
|
||||
+17
@@ -65,6 +65,7 @@ import comfy.text_encoders.ace15
|
||||
import comfy.text_encoders.longcat_image
|
||||
import comfy.text_encoders.qwen35
|
||||
import comfy.text_encoders.ernie
|
||||
import comfy.text_encoders.gemma4
|
||||
|
||||
import comfy.model_patcher
|
||||
import comfy.lora
|
||||
@@ -1271,6 +1272,9 @@ class TEModel(Enum):
|
||||
QWEN35_9B = 26
|
||||
QWEN35_27B = 27
|
||||
MINISTRAL_3_3B = 28
|
||||
GEMMA_4_E4B = 29
|
||||
GEMMA_4_E2B = 30
|
||||
GEMMA_4_31B = 31
|
||||
|
||||
|
||||
def detect_te_model(sd):
|
||||
@@ -1296,6 +1300,12 @@ def detect_te_model(sd):
|
||||
return TEModel.BYT5_SMALL_GLYPH
|
||||
return TEModel.T5_BASE
|
||||
if 'model.layers.0.post_feedforward_layernorm.weight' in sd:
|
||||
if 'model.layers.59.self_attn.q_norm.weight' in sd:
|
||||
return TEModel.GEMMA_4_31B
|
||||
if 'model.layers.41.self_attn.q_norm.weight' in sd and 'model.layers.47.self_attn.q_norm.weight' not in sd:
|
||||
return TEModel.GEMMA_4_E4B
|
||||
if 'model.layers.34.self_attn.q_norm.weight' in sd and 'model.layers.41.self_attn.q_norm.weight' not in sd:
|
||||
return TEModel.GEMMA_4_E2B
|
||||
if 'model.layers.47.self_attn.q_norm.weight' in sd:
|
||||
return TEModel.GEMMA_3_12B
|
||||
if 'model.layers.0.self_attn.q_norm.weight' in sd:
|
||||
@@ -1435,6 +1445,13 @@ def load_text_encoder_state_dicts(state_dicts=[], embedding_directory=None, clip
|
||||
else:
|
||||
clip_target.clip = comfy.text_encoders.sa_t5.SAT5Model
|
||||
clip_target.tokenizer = comfy.text_encoders.sa_t5.SAT5Tokenizer
|
||||
elif te_model in (TEModel.GEMMA_4_E4B, TEModel.GEMMA_4_E2B, TEModel.GEMMA_4_31B):
|
||||
variant = {TEModel.GEMMA_4_E4B: comfy.text_encoders.gemma4.Gemma4_E4B,
|
||||
TEModel.GEMMA_4_E2B: comfy.text_encoders.gemma4.Gemma4_E2B,
|
||||
TEModel.GEMMA_4_31B: comfy.text_encoders.gemma4.Gemma4_31B}[te_model]
|
||||
clip_target.clip = comfy.text_encoders.gemma4.gemma4_te(**llama_detect(clip_data), model_class=variant)
|
||||
clip_target.tokenizer = variant.tokenizer
|
||||
tokenizer_data["tokenizer_json"] = clip_data[0].get("tokenizer_json", None)
|
||||
elif te_model == TEModel.GEMMA_2_2B:
|
||||
clip_target.clip = comfy.text_encoders.lumina2.te(**llama_detect(clip_data))
|
||||
clip_target.tokenizer = comfy.text_encoders.lumina2.LuminaTokenizer
|
||||
|
||||
@@ -1879,6 +1879,86 @@ class CogVideoX_I2V(CogVideoX_T2V):
|
||||
out = model_base.CogVideoX(self, image_to_video=True, device=device)
|
||||
return out
|
||||
|
||||
models = [LotusD, Stable_Zero123, SD15_instructpix2pix, SD15, SD20, SD21UnclipL, SD21UnclipH, SDXL_instructpix2pix, SDXLRefiner, SDXL, SSD1B, KOALA_700M, KOALA_1B, Segmind_Vega, SD_X4Upscaler, Stable_Cascade_C, Stable_Cascade_B, SV3D_u, SV3D_p, SD3, StableAudio, AuraFlow, PixArtAlpha, PixArtSigma, HunyuanDiT, HunyuanDiT1, FluxInpaint, Flux, LongCatImage, FluxSchnell, GenmoMochi, LTXV, LTXAV, HunyuanVideo15_SR_Distilled, HunyuanVideo15, HunyuanImage21Refiner, HunyuanImage21, HunyuanVideoSkyreelsI2V, HunyuanVideoI2V, HunyuanVideo, CosmosT2V, CosmosI2V, CosmosT2IPredict2, CosmosI2VPredict2, ZImagePixelSpace, ZImage, Lumina2, WAN22_T2V, WAN21_T2V, WAN21_I2V, WAN21_FunControl2V, WAN21_Vace, WAN21_Camera, WAN22_Camera, WAN22_S2V, WAN21_HuMo, WAN22_Animate, WAN21_FlowRVS, WAN21_SCAIL, Hunyuan3Dv2mini, Hunyuan3Dv2, Hunyuan3Dv2_1, HiDream, Chroma, ChromaRadiance, ACEStep, ACEStep15, Omnigen2, QwenImage, Flux2, Kandinsky5Image, Kandinsky5, Anima, RT_DETR_v4, ErnieImage, SAM3, SAM31, CogVideoX_I2V, CogVideoX_T2V]
|
||||
|
||||
models += [SVD_img2vid]
|
||||
models = [
|
||||
LotusD,
|
||||
Stable_Zero123,
|
||||
SD15_instructpix2pix,
|
||||
SD15,
|
||||
SD20,
|
||||
SD21UnclipL,
|
||||
SD21UnclipH,
|
||||
SDXL_instructpix2pix,
|
||||
SDXLRefiner,
|
||||
SDXL,
|
||||
SSD1B,
|
||||
KOALA_700M,
|
||||
KOALA_1B,
|
||||
Segmind_Vega,
|
||||
SD_X4Upscaler,
|
||||
Stable_Cascade_C,
|
||||
Stable_Cascade_B,
|
||||
SV3D_u,
|
||||
SV3D_p,
|
||||
SD3,
|
||||
StableAudio,
|
||||
AuraFlow,
|
||||
PixArtAlpha,
|
||||
PixArtSigma,
|
||||
HunyuanDiT,
|
||||
HunyuanDiT1,
|
||||
FluxInpaint,
|
||||
Flux,
|
||||
LongCatImage,
|
||||
FluxSchnell,
|
||||
GenmoMochi,
|
||||
LTXV,
|
||||
LTXAV,
|
||||
HunyuanVideo15_SR_Distilled,
|
||||
HunyuanVideo15,
|
||||
HunyuanImage21Refiner,
|
||||
HunyuanImage21,
|
||||
HunyuanVideoSkyreelsI2V,
|
||||
HunyuanVideoI2V,
|
||||
HunyuanVideo,
|
||||
CosmosT2V,
|
||||
CosmosI2V,
|
||||
CosmosT2IPredict2,
|
||||
CosmosI2VPredict2,
|
||||
ZImagePixelSpace,
|
||||
ZImage,
|
||||
Lumina2,
|
||||
WAN22_T2V,
|
||||
WAN21_T2V,
|
||||
WAN21_I2V,
|
||||
WAN21_FunControl2V,
|
||||
WAN21_Vace,
|
||||
WAN21_Camera,
|
||||
WAN22_Camera,
|
||||
WAN22_S2V,
|
||||
WAN21_HuMo,
|
||||
WAN22_Animate,
|
||||
WAN21_FlowRVS,
|
||||
WAN21_SCAIL,
|
||||
Hunyuan3Dv2mini,
|
||||
Hunyuan3Dv2,
|
||||
Hunyuan3Dv2_1,
|
||||
HiDream,
|
||||
Chroma,
|
||||
ChromaRadiance,
|
||||
ACEStep,
|
||||
ACEStep15,
|
||||
Omnigen2,
|
||||
QwenImage,
|
||||
Flux2,
|
||||
Kandinsky5Image,
|
||||
Kandinsky5,
|
||||
Anima,
|
||||
RT_DETR_v4,
|
||||
ErnieImage,
|
||||
SAM3,
|
||||
SAM31,
|
||||
CogVideoX_I2V,
|
||||
CogVideoX_T2V,
|
||||
SVD_img2vid,
|
||||
]
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -521,7 +521,7 @@ class Attention(nn.Module):
|
||||
else:
|
||||
present_key_value = (xk, xv, index + num_tokens)
|
||||
|
||||
if sliding_window is not None and xk.shape[2] > sliding_window:
|
||||
if sliding_window is not None and xk.shape[2] > sliding_window and seq_length == 1:
|
||||
xk = xk[:, :, -sliding_window:]
|
||||
xv = xv[:, :, -sliding_window:]
|
||||
attention_mask = attention_mask[..., -sliding_window:] if attention_mask is not None else None
|
||||
@@ -533,12 +533,12 @@ class Attention(nn.Module):
|
||||
return self.o_proj(output), present_key_value
|
||||
|
||||
class MLP(nn.Module):
|
||||
def __init__(self, config: Llama2Config, device=None, dtype=None, ops: Any = None):
|
||||
def __init__(self, config: Llama2Config, device=None, dtype=None, ops: Any = None, intermediate_size=None):
|
||||
super().__init__()
|
||||
ops = ops or nn
|
||||
self.gate_proj = ops.Linear(config.hidden_size, config.intermediate_size, bias=False, device=device, dtype=dtype)
|
||||
self.up_proj = ops.Linear(config.hidden_size, config.intermediate_size, bias=False, device=device, dtype=dtype)
|
||||
self.down_proj = ops.Linear(config.intermediate_size, config.hidden_size, bias=False, device=device, dtype=dtype)
|
||||
intermediate_size = intermediate_size or config.intermediate_size
|
||||
self.gate_proj = ops.Linear(config.hidden_size, intermediate_size, bias=False, device=device, dtype=dtype)
|
||||
self.up_proj = ops.Linear(config.hidden_size, intermediate_size, bias=False, device=device, dtype=dtype)
|
||||
self.down_proj = ops.Linear(intermediate_size, config.hidden_size, bias=False, device=device, dtype=dtype)
|
||||
if config.mlp_activation == "silu":
|
||||
self.activation = torch.nn.functional.silu
|
||||
elif config.mlp_activation == "gelu_pytorch_tanh":
|
||||
@@ -647,24 +647,25 @@ class TransformerBlockGemma2(nn.Module):
|
||||
|
||||
return x, present_key_value
|
||||
|
||||
def _make_scaled_embedding(ops, vocab_size, hidden_size, scale, device, dtype):
|
||||
class ScaledEmbedding(ops.Embedding):
|
||||
def forward(self, input_ids, out_dtype=None):
|
||||
return super().forward(input_ids, out_dtype=out_dtype) * scale
|
||||
return ScaledEmbedding(vocab_size, hidden_size, device=device, dtype=dtype)
|
||||
|
||||
|
||||
class Llama2_(nn.Module):
|
||||
def __init__(self, config, device=None, dtype=None, ops=None):
|
||||
super().__init__()
|
||||
self.config = config
|
||||
self.vocab_size = config.vocab_size
|
||||
|
||||
self.embed_tokens = ops.Embedding(
|
||||
config.vocab_size,
|
||||
config.hidden_size,
|
||||
device=device,
|
||||
dtype=dtype
|
||||
)
|
||||
if self.config.transformer_type == "gemma2" or self.config.transformer_type == "gemma3":
|
||||
transformer = TransformerBlockGemma2
|
||||
self.normalize_in = True
|
||||
self.embed_tokens = _make_scaled_embedding(ops, config.vocab_size, config.hidden_size, config.hidden_size ** 0.5, device, dtype)
|
||||
else:
|
||||
transformer = TransformerBlock
|
||||
self.normalize_in = False
|
||||
self.embed_tokens = ops.Embedding(config.vocab_size, config.hidden_size, device=device, dtype=dtype)
|
||||
|
||||
self.layers = nn.ModuleList([
|
||||
transformer(config, index=i, device=device, dtype=dtype, ops=ops)
|
||||
@@ -690,15 +691,12 @@ class Llama2_(nn.Module):
|
||||
self.config.rope_dims,
|
||||
device=device)
|
||||
|
||||
def forward(self, x, attention_mask=None, embeds=None, num_tokens=None, intermediate_output=None, final_layer_norm_intermediate=True, dtype=None, position_ids=None, embeds_info=[], past_key_values=None):
|
||||
def forward(self, x, attention_mask=None, embeds=None, num_tokens=None, intermediate_output=None, final_layer_norm_intermediate=True, dtype=None, position_ids=None, embeds_info=[], past_key_values=None, input_ids=None):
|
||||
if embeds is not None:
|
||||
x = embeds
|
||||
else:
|
||||
x = self.embed_tokens(x, out_dtype=dtype)
|
||||
|
||||
if self.normalize_in:
|
||||
x *= self.config.hidden_size ** 0.5
|
||||
|
||||
seq_len = x.shape[1]
|
||||
past_len = 0
|
||||
if past_key_values is not None and len(past_key_values) > 0:
|
||||
@@ -850,7 +848,7 @@ class BaseGenerate:
|
||||
torch.empty([batch, model_config.num_key_value_heads, max_cache_len, model_config.head_dim], device=device, dtype=execution_dtype), 0))
|
||||
return past_key_values
|
||||
|
||||
def generate(self, embeds=None, do_sample=True, max_length=256, temperature=1.0, top_k=50, top_p=0.9, min_p=0.0, repetition_penalty=1.0, seed=42, stop_tokens=None, initial_tokens=[], execution_dtype=None, min_tokens=0, presence_penalty=0.0):
|
||||
def generate(self, embeds=None, do_sample=True, max_length=256, temperature=1.0, top_k=50, top_p=0.9, min_p=0.0, repetition_penalty=1.0, seed=42, stop_tokens=None, initial_tokens=[], execution_dtype=None, min_tokens=0, presence_penalty=0.0, initial_input_ids=None):
|
||||
device = embeds.device
|
||||
|
||||
if stop_tokens is None:
|
||||
@@ -875,14 +873,16 @@ class BaseGenerate:
|
||||
pbar = comfy.utils.ProgressBar(max_length)
|
||||
|
||||
# Generation loop
|
||||
current_input_ids = initial_input_ids
|
||||
for step in tqdm(range(max_length), desc="Generating tokens"):
|
||||
x, _, past_key_values = self.model.forward(None, embeds=embeds, attention_mask=None, past_key_values=past_key_values)
|
||||
x, _, past_key_values = self.model.forward(None, embeds=embeds, attention_mask=None, past_key_values=past_key_values, input_ids=current_input_ids)
|
||||
logits = self.logits(x)[:, -1]
|
||||
next_token = self.sample_token(logits, temperature, top_k, top_p, min_p, repetition_penalty, initial_tokens + generated_token_ids, generator, do_sample=do_sample, presence_penalty=presence_penalty)
|
||||
token_id = next_token[0].item()
|
||||
generated_token_ids.append(token_id)
|
||||
|
||||
embeds = self.model.embed_tokens(next_token).to(execution_dtype)
|
||||
current_input_ids = next_token if initial_input_ids is not None else None
|
||||
pbar.update(1)
|
||||
|
||||
if token_id in stop_tokens:
|
||||
|
||||
@@ -93,8 +93,7 @@ class Gemma3_12BModel(sd1_clip.SDClipModel):
|
||||
|
||||
def generate(self, tokens, do_sample, max_length, temperature, top_k, top_p, min_p, repetition_penalty, seed, presence_penalty):
|
||||
tokens_only = [[t[0] for t in b] for b in tokens]
|
||||
embeds, _, _, embeds_info = self.process_tokens(tokens_only, self.execution_device)
|
||||
comfy.utils.normalize_image_embeddings(embeds, embeds_info, self.transformer.model.config.hidden_size ** 0.5)
|
||||
embeds, _, _, _ = self.process_tokens(tokens_only, self.execution_device)
|
||||
return self.transformer.generate(embeds, do_sample, max_length, temperature, top_k, top_p, min_p, repetition_penalty, seed, stop_tokens=[106], presence_penalty=presence_penalty) # 106 is <end_of_turn>
|
||||
|
||||
class DualLinearProjection(torch.nn.Module):
|
||||
|
||||
@@ -50,8 +50,7 @@ class Gemma3_4B_Vision_Model(sd1_clip.SDClipModel):
|
||||
super().__init__(device=device, layer=layer, layer_idx=layer_idx, textmodel_json_config={}, dtype=dtype, special_tokens={"start": 2, "pad": 0}, layer_norm_hidden_state=False, model_class=comfy.text_encoders.llama.Gemma3_4B_Vision, enable_attention_masks=attention_mask, return_attention_masks=attention_mask, model_options=model_options)
|
||||
|
||||
def process_tokens(self, tokens, device):
|
||||
embeds, _, _, embeds_info = super().process_tokens(tokens, device)
|
||||
comfy.utils.normalize_image_embeddings(embeds, embeds_info, self.transformer.model.config.hidden_size ** 0.5)
|
||||
embeds, _, _, _ = super().process_tokens(tokens, device)
|
||||
return embeds
|
||||
|
||||
class LuminaModel(sd1_clip.SD1ClipModel):
|
||||
|
||||
@@ -408,8 +408,6 @@ class Qwen35Transformer(Llama2_):
|
||||
nn.Module.__init__(self)
|
||||
self.config = config
|
||||
self.vocab_size = config.vocab_size
|
||||
self.normalize_in = False
|
||||
|
||||
self.embed_tokens = ops.Embedding(config.vocab_size, config.hidden_size, device=device, dtype=dtype)
|
||||
self.layers = nn.ModuleList([
|
||||
Qwen35TransformerBlock(config, index=i, device=device, dtype=dtype, ops=ops)
|
||||
|
||||
@@ -1446,10 +1446,3 @@ def deepcopy_list_dict(obj, memo=None):
|
||||
memo[obj_id] = res
|
||||
return res
|
||||
|
||||
def normalize_image_embeddings(embeds, embeds_info, scale_factor):
|
||||
"""Normalize image embeddings to match text embedding scale"""
|
||||
for info in embeds_info:
|
||||
if info.get("type") == "image":
|
||||
start_idx = info["index"]
|
||||
end_idx = start_idx + info["size"]
|
||||
embeds[:, start_idx:end_idx, :] /= scale_factor
|
||||
|
||||
@@ -290,7 +290,7 @@ class VideoFromFile(VideoInput):
|
||||
alphas = []
|
||||
alpha_channel = True
|
||||
break
|
||||
if frame.format.name in ("yuvj420p", "rgb24", "rgba", "pal8"):
|
||||
if frame.format.name in ("yuvj420p", "yuvj422p", "yuvj444p", "rgb24", "rgba", "pal8"):
|
||||
process_image_format = lambda a: a.float() / 255.0
|
||||
if alpha_channel:
|
||||
image_format = 'rgba'
|
||||
|
||||
@@ -43,7 +43,67 @@ class UploadType(str, Enum):
|
||||
model = "file_upload"
|
||||
|
||||
|
||||
class RemoteItemSchema:
|
||||
"""Describes how to map API response objects to rich dropdown items.
|
||||
|
||||
All *_field parameters use dot-path notation (e.g. ``"labels.gender"``).
|
||||
``label_field`` and ``description_field`` additionally support template strings
|
||||
with ``{field}`` placeholders (e.g. ``"{name} ({labels.accent})"``).
|
||||
"""
|
||||
def __init__(
|
||||
self,
|
||||
value_field: str,
|
||||
label_field: str,
|
||||
preview_url_field: str | None = None,
|
||||
preview_type: Literal["image", "video", "audio"] = "image",
|
||||
description_field: str | None = None,
|
||||
search_fields: list[str] | None = None,
|
||||
):
|
||||
if preview_type not in ("image", "video", "audio"):
|
||||
raise ValueError(
|
||||
f"RemoteItemSchema: 'preview_type' must be 'image', 'video', or 'audio'; got {preview_type!r}."
|
||||
)
|
||||
if search_fields is not None:
|
||||
for f in search_fields:
|
||||
if "{" in f or "}" in f:
|
||||
raise ValueError(
|
||||
f"RemoteItemSchema: 'search_fields' must be dot-paths, not template strings (got {f!r})."
|
||||
)
|
||||
self.value_field = value_field
|
||||
"""Dot-path to the unique identifier within each item.
|
||||
This value is stored in the widget and passed to execute()."""
|
||||
self.label_field = label_field
|
||||
"""Dot-path to the display name, or a template string with {field} placeholders."""
|
||||
self.preview_url_field = preview_url_field
|
||||
"""Dot-path to a preview media URL. If None, no preview is shown."""
|
||||
self.preview_type = preview_type
|
||||
"""How to render the preview: "image", "video", or "audio"."""
|
||||
self.description_field = description_field
|
||||
"""Optional dot-path or template for a subtitle line shown below the label."""
|
||||
self.search_fields = search_fields
|
||||
"""Dot-paths to fields included in the search index. When unset, search falls back to
|
||||
the resolved label (i.e. ``label_field`` after template substitution). Note that template
|
||||
label strings (e.g. ``"{first} {last}"``) are not valid path entries here — list the
|
||||
underlying paths (``["first", "last"]``) instead."""
|
||||
|
||||
def as_dict(self):
|
||||
return prune_dict({
|
||||
"value_field": self.value_field,
|
||||
"label_field": self.label_field,
|
||||
"preview_url_field": self.preview_url_field,
|
||||
"preview_type": self.preview_type,
|
||||
"description_field": self.description_field,
|
||||
"search_fields": self.search_fields,
|
||||
})
|
||||
|
||||
|
||||
class RemoteOptions:
|
||||
"""Plain remote combo: fetches a list of strings/objects and populates a standard dropdown.
|
||||
|
||||
Use this for lightweight lists from endpoints that return a bare array (or an array under
|
||||
``response_key``). For rich dropdowns with previews, search, filtering, or pagination,
|
||||
use :class:`RemoteComboOptions` and the ``remote_combo=`` parameter on ``Combo.Input``.
|
||||
"""
|
||||
def __init__(self, route: str, refresh_button: bool, control_after_refresh: Literal["first", "last"]="first",
|
||||
timeout: int=None, max_retries: int=None, refresh: int=None):
|
||||
self.route = route
|
||||
@@ -70,6 +130,80 @@ class RemoteOptions:
|
||||
})
|
||||
|
||||
|
||||
class RemoteComboOptions:
|
||||
"""Rich remote combo: populates a Vue dropdown with previews, search, and filtering.
|
||||
|
||||
Attached to a :class:`Combo.Input` via ``remote_combo=`` (not ``remote=``). Requires an
|
||||
``item_schema`` describing how to map API response objects to dropdown items.
|
||||
|
||||
Response-shape contract: the endpoint returns the full items array in a single response
|
||||
(either at the top level, or at the dot-path given by ``response_key``). Backing endpoints
|
||||
that paginate upstream are expected to aggregate and cache server-side.
|
||||
"""
|
||||
def __init__(
|
||||
self,
|
||||
route: str,
|
||||
item_schema: RemoteItemSchema,
|
||||
refresh_button: bool = True,
|
||||
auto_select: Literal["first", "last"] | None = None,
|
||||
timeout: int | None = None,
|
||||
max_retries: int | None = None,
|
||||
refresh: int | None = None,
|
||||
response_key: str | None = None,
|
||||
):
|
||||
if auto_select is not None and auto_select not in ("first", "last"):
|
||||
raise ValueError(
|
||||
f"RemoteComboOptions: 'auto_select' must be 'first', 'last', or None; got {auto_select!r}."
|
||||
)
|
||||
if refresh is not None and 0 < refresh < 128:
|
||||
raise ValueError(
|
||||
f"RemoteComboOptions: 'refresh' must be >= 128 (ms TTL) or <= 0 (cache never expires); got {refresh}."
|
||||
)
|
||||
if timeout is not None and timeout < 0:
|
||||
raise ValueError(
|
||||
f"RemoteComboOptions: 'timeout' must be >= 0 (got {timeout})."
|
||||
)
|
||||
if max_retries is not None and max_retries < 0:
|
||||
raise ValueError(
|
||||
f"RemoteComboOptions: 'max_retries' must be >= 0 (got {max_retries})."
|
||||
)
|
||||
if not route.startswith("/"):
|
||||
raise ValueError(
|
||||
f"RemoteComboOptions: 'route' must be a relative path starting with '/'; got {route!r}."
|
||||
)
|
||||
self.route = route
|
||||
"""Relative path to the remote source (must start with ``/``). The frontend resolves this
|
||||
against the comfy-api base URL and injects auth headers; absolute URLs are rejected."""
|
||||
self.item_schema = item_schema
|
||||
"""Required: describes how each API response object maps to a dropdown item."""
|
||||
self.refresh_button = refresh_button
|
||||
"""Specifies whether to show a refresh button next to the widget."""
|
||||
self.auto_select = auto_select
|
||||
"""Fallback item to select when the widget's value is empty. Never overrides an existing
|
||||
selection. Default None means no fallback."""
|
||||
self.timeout = timeout
|
||||
"""Maximum time to wait for a response, in milliseconds."""
|
||||
self.max_retries = max_retries
|
||||
"""Maximum number of retries before aborting the request. Default None uses the frontend's built-in limit."""
|
||||
self.refresh = refresh
|
||||
"""TTL of the cached value in milliseconds. Must be >= 128 (ms TTL) or <= 0 (cache never expires,
|
||||
re-fetched only via the refresh button). Default None uses the frontend's built-in behavior."""
|
||||
self.response_key = response_key
|
||||
"""Dot-path to the items array within the response (when not at the top level)."""
|
||||
|
||||
def as_dict(self):
|
||||
return prune_dict({
|
||||
"route": self.route,
|
||||
"item_schema": self.item_schema.as_dict(),
|
||||
"refresh_button": self.refresh_button,
|
||||
"auto_select": self.auto_select,
|
||||
"timeout": self.timeout,
|
||||
"max_retries": self.max_retries,
|
||||
"refresh": self.refresh,
|
||||
"response_key": self.response_key,
|
||||
})
|
||||
|
||||
|
||||
class NumberDisplay(str, Enum):
|
||||
number = "number"
|
||||
slider = "slider"
|
||||
@@ -359,11 +493,16 @@ class Combo(ComfyTypeIO):
|
||||
upload: UploadType=None,
|
||||
image_folder: FolderType=None,
|
||||
remote: RemoteOptions=None,
|
||||
remote_combo: RemoteComboOptions=None,
|
||||
socketless: bool=None,
|
||||
extra_dict=None,
|
||||
raw_link: bool=None,
|
||||
advanced: bool=None,
|
||||
):
|
||||
if remote is not None and remote_combo is not None:
|
||||
raise ValueError("Combo.Input: pass either 'remote' or 'remote_combo', not both.")
|
||||
if options is not None and remote_combo is not None:
|
||||
raise ValueError("Combo.Input: pass either 'options' or 'remote_combo', not both.")
|
||||
if isinstance(options, type) and issubclass(options, Enum):
|
||||
options = [v.value for v in options]
|
||||
if isinstance(default, Enum):
|
||||
@@ -375,6 +514,7 @@ class Combo(ComfyTypeIO):
|
||||
self.upload = upload
|
||||
self.image_folder = image_folder
|
||||
self.remote = remote
|
||||
self.remote_combo = remote_combo
|
||||
self.default: str
|
||||
|
||||
def as_dict(self):
|
||||
@@ -385,6 +525,7 @@ class Combo(ComfyTypeIO):
|
||||
**({self.upload.value: True} if self.upload is not None else {}),
|
||||
"image_folder": self.image_folder.value if self.image_folder else None,
|
||||
"remote": self.remote.as_dict() if self.remote else None,
|
||||
"remote_combo": self.remote_combo.as_dict() if self.remote_combo else None,
|
||||
})
|
||||
|
||||
class Output(Output):
|
||||
@@ -2221,7 +2362,9 @@ class NodeReplace:
|
||||
__all__ = [
|
||||
"FolderType",
|
||||
"UploadType",
|
||||
"RemoteItemSchema",
|
||||
"RemoteOptions",
|
||||
"RemoteComboOptions",
|
||||
"NumberDisplay",
|
||||
"ControlAfterGenerate",
|
||||
|
||||
|
||||
+5
-24
@@ -65,22 +65,6 @@ class SavedAudios(_UIOutput):
|
||||
return {"audio": self.results}
|
||||
|
||||
|
||||
def _is_isolated_child() -> bool:
|
||||
return os.environ.get("PYISOLATE_CHILD") == "1"
|
||||
|
||||
|
||||
def _get_preview_folder_type() -> FolderType:
|
||||
if _is_isolated_child():
|
||||
return FolderType.output
|
||||
return FolderType.temp
|
||||
|
||||
|
||||
def _get_preview_route_prefix(folder_type: FolderType) -> str:
|
||||
if folder_type == FolderType.output:
|
||||
return "output"
|
||||
return "temp"
|
||||
|
||||
|
||||
def _get_directory_by_folder_type(folder_type: FolderType) -> str:
|
||||
if folder_type == FolderType.input:
|
||||
return folder_paths.get_input_directory()
|
||||
@@ -404,11 +388,10 @@ class AudioSaveHelper:
|
||||
|
||||
class PreviewImage(_UIOutput):
|
||||
def __init__(self, image: Image.Type, animated: bool = False, cls: type[ComfyNode] = None, **kwargs):
|
||||
folder_type = _get_preview_folder_type()
|
||||
self.values = ImageSaveHelper.save_images(
|
||||
image,
|
||||
filename_prefix="ComfyUI_temp_" + ''.join(random.choice("abcdefghijklmnopqrstupvxyz") for _ in range(5)),
|
||||
folder_type=folder_type,
|
||||
folder_type=FolderType.temp,
|
||||
cls=cls,
|
||||
compress_level=1,
|
||||
)
|
||||
@@ -429,11 +412,10 @@ class PreviewMask(PreviewImage):
|
||||
|
||||
class PreviewAudio(_UIOutput):
|
||||
def __init__(self, audio: dict, cls: type[ComfyNode] = None, **kwargs):
|
||||
folder_type = _get_preview_folder_type()
|
||||
self.values = AudioSaveHelper.save_audio(
|
||||
audio,
|
||||
filename_prefix="ComfyUI_temp_" + "".join(random.choice("abcdefghijklmnopqrstuvwxyz") for _ in range(5)),
|
||||
folder_type=folder_type,
|
||||
folder_type=FolderType.temp,
|
||||
cls=cls,
|
||||
format="flac",
|
||||
quality="128k",
|
||||
@@ -456,16 +438,15 @@ class PreviewUI3D(_UIOutput):
|
||||
self.model_file = model_file
|
||||
self.camera_info = camera_info
|
||||
self.bg_image_path = None
|
||||
folder_type = _get_preview_folder_type()
|
||||
bg_image = kwargs.get("bg_image", None)
|
||||
if bg_image is not None:
|
||||
img_array = (bg_image[0].cpu().numpy() * 255).astype(np.uint8)
|
||||
img = PILImage.fromarray(img_array)
|
||||
preview_dir = _get_directory_by_folder_type(folder_type)
|
||||
temp_dir = folder_paths.get_temp_directory()
|
||||
filename = f"bg_{uuid.uuid4().hex}.png"
|
||||
bg_image_path = os.path.join(preview_dir, filename)
|
||||
bg_image_path = os.path.join(temp_dir, filename)
|
||||
img.save(bg_image_path, compress_level=1)
|
||||
self.bg_image_path = f"{_get_preview_route_prefix(folder_type)}/{filename}"
|
||||
self.bg_image_path = f"temp/{filename}"
|
||||
|
||||
def as_dict(self):
|
||||
return {"result": [self.model_file, self.camera_info, self.bg_image_path]}
|
||||
|
||||
@@ -132,6 +132,10 @@ class GetAssetResponse(BaseModel):
|
||||
error: TaskStatusError | None = Field(None)
|
||||
|
||||
|
||||
class SeedanceCreateVisualValidateSessionRequest(BaseModel):
|
||||
name: str | None = Field(None, max_length=64)
|
||||
|
||||
|
||||
class SeedanceCreateVisualValidateSessionResponse(BaseModel):
|
||||
session_id: str = Field(...)
|
||||
h5_link: str = Field(...)
|
||||
@@ -141,6 +145,7 @@ class SeedanceGetVisualValidateSessionResponse(BaseModel):
|
||||
session_id: str = Field(...)
|
||||
status: str = Field(...)
|
||||
group_id: str | None = Field(None)
|
||||
name: str | None = Field(None)
|
||||
error_code: str | None = Field(None)
|
||||
error_message: str | None = Field(None)
|
||||
|
||||
|
||||
@@ -1,152 +0,0 @@
|
||||
from enum import Enum
|
||||
from typing import Optional, Dict, Any
|
||||
|
||||
from pydantic import BaseModel, Field, StrictBytes
|
||||
|
||||
|
||||
class MoonvalleyPromptResponse(BaseModel):
|
||||
error: Optional[Dict[str, Any]] = None
|
||||
frame_conditioning: Optional[Dict[str, Any]] = None
|
||||
id: Optional[str] = None
|
||||
inference_params: Optional[Dict[str, Any]] = None
|
||||
meta: Optional[Dict[str, Any]] = None
|
||||
model_params: Optional[Dict[str, Any]] = None
|
||||
output_url: Optional[str] = None
|
||||
prompt_text: Optional[str] = None
|
||||
status: Optional[str] = None
|
||||
|
||||
|
||||
class MoonvalleyTextToVideoInferenceParams(BaseModel):
|
||||
add_quality_guidance: Optional[bool] = Field(
|
||||
True, description='Whether to add quality guidance'
|
||||
)
|
||||
caching_coefficient: Optional[float] = Field(
|
||||
0.3, description='Caching coefficient for optimization'
|
||||
)
|
||||
caching_cooldown: Optional[int] = Field(
|
||||
3, description='Number of caching cooldown steps'
|
||||
)
|
||||
caching_warmup: Optional[int] = Field(
|
||||
3, description='Number of caching warmup steps'
|
||||
)
|
||||
clip_value: Optional[float] = Field(
|
||||
3, description='CLIP value for generation control'
|
||||
)
|
||||
conditioning_frame_index: Optional[int] = Field(
|
||||
0, description='Index of the conditioning frame'
|
||||
)
|
||||
cooldown_steps: Optional[int] = Field(
|
||||
75, description='Number of cooldown steps (calculated based on num_frames)'
|
||||
)
|
||||
fps: Optional[int] = Field(
|
||||
24, description='Frames per second of the generated video'
|
||||
)
|
||||
guidance_scale: Optional[float] = Field(
|
||||
10, description='Guidance scale for generation control'
|
||||
)
|
||||
height: Optional[int] = Field(
|
||||
1080, description='Height of the generated video in pixels'
|
||||
)
|
||||
negative_prompt: Optional[str] = Field(None, description='Negative prompt text')
|
||||
num_frames: Optional[int] = Field(64, description='Number of frames to generate')
|
||||
seed: Optional[int] = Field(
|
||||
None, description='Random seed for generation (default: random)'
|
||||
)
|
||||
shift_value: Optional[float] = Field(
|
||||
3, description='Shift value for generation control'
|
||||
)
|
||||
steps: Optional[int] = Field(80, description='Number of denoising steps')
|
||||
use_guidance_schedule: Optional[bool] = Field(
|
||||
True, description='Whether to use guidance scheduling'
|
||||
)
|
||||
use_negative_prompts: Optional[bool] = Field(
|
||||
False, description='Whether to use negative prompts'
|
||||
)
|
||||
use_timestep_transform: Optional[bool] = Field(
|
||||
True, description='Whether to use timestep transformation'
|
||||
)
|
||||
warmup_steps: Optional[int] = Field(
|
||||
0, description='Number of warmup steps (calculated based on num_frames)'
|
||||
)
|
||||
width: Optional[int] = Field(
|
||||
1920, description='Width of the generated video in pixels'
|
||||
)
|
||||
|
||||
|
||||
class MoonvalleyTextToVideoRequest(BaseModel):
|
||||
image_url: Optional[str] = None
|
||||
inference_params: Optional[MoonvalleyTextToVideoInferenceParams] = None
|
||||
prompt_text: Optional[str] = None
|
||||
webhook_url: Optional[str] = None
|
||||
|
||||
|
||||
class MoonvalleyUploadFileRequest(BaseModel):
|
||||
file: Optional[StrictBytes] = None
|
||||
|
||||
|
||||
class MoonvalleyUploadFileResponse(BaseModel):
|
||||
access_url: Optional[str] = None
|
||||
|
||||
|
||||
class MoonvalleyVideoToVideoInferenceParams(BaseModel):
|
||||
add_quality_guidance: Optional[bool] = Field(
|
||||
True, description='Whether to add quality guidance'
|
||||
)
|
||||
caching_coefficient: Optional[float] = Field(
|
||||
0.3, description='Caching coefficient for optimization'
|
||||
)
|
||||
caching_cooldown: Optional[int] = Field(
|
||||
3, description='Number of caching cooldown steps'
|
||||
)
|
||||
caching_warmup: Optional[int] = Field(
|
||||
3, description='Number of caching warmup steps'
|
||||
)
|
||||
clip_value: Optional[float] = Field(
|
||||
3, description='CLIP value for generation control'
|
||||
)
|
||||
conditioning_frame_index: Optional[int] = Field(
|
||||
0, description='Index of the conditioning frame'
|
||||
)
|
||||
cooldown_steps: Optional[int] = Field(
|
||||
36, description='Number of cooldown steps (calculated based on num_frames)'
|
||||
)
|
||||
guidance_scale: Optional[float] = Field(
|
||||
15, description='Guidance scale for generation control'
|
||||
)
|
||||
negative_prompt: Optional[str] = Field(None, description='Negative prompt text')
|
||||
seed: Optional[int] = Field(
|
||||
None, description='Random seed for generation (default: random)'
|
||||
)
|
||||
shift_value: Optional[float] = Field(
|
||||
3, description='Shift value for generation control'
|
||||
)
|
||||
steps: Optional[int] = Field(80, description='Number of denoising steps')
|
||||
use_guidance_schedule: Optional[bool] = Field(
|
||||
True, description='Whether to use guidance scheduling'
|
||||
)
|
||||
use_negative_prompts: Optional[bool] = Field(
|
||||
False, description='Whether to use negative prompts'
|
||||
)
|
||||
use_timestep_transform: Optional[bool] = Field(
|
||||
True, description='Whether to use timestep transformation'
|
||||
)
|
||||
warmup_steps: Optional[int] = Field(
|
||||
24, description='Number of warmup steps (calculated based on num_frames)'
|
||||
)
|
||||
|
||||
|
||||
class ControlType(str, Enum):
|
||||
motion_control = 'motion_control'
|
||||
pose_control = 'pose_control'
|
||||
|
||||
|
||||
class MoonvalleyVideoToVideoRequest(BaseModel):
|
||||
control_type: ControlType = Field(
|
||||
..., description='Supported types for video control'
|
||||
)
|
||||
inference_params: Optional[MoonvalleyVideoToVideoInferenceParams] = None
|
||||
prompt_text: str = Field(..., description='Describes the video to generate')
|
||||
video_url: str = Field(..., description='Url to control video')
|
||||
webhook_url: Optional[str] = Field(
|
||||
None, description='Optional webhook URL for notifications'
|
||||
)
|
||||
@@ -1,4 +1,4 @@
|
||||
from typing import Optional, Union
|
||||
from typing import Optional
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
@@ -72,8 +72,11 @@ class VideoEnhancementFilter(BaseModel):
|
||||
grain: Optional[float] = Field(None, description="Grain after AI model processing")
|
||||
grainSize: Optional[float] = Field(None, description="Size of generated grain")
|
||||
recoverOriginalDetailValue: Optional[float] = Field(None, description="Source details into the output video")
|
||||
creativity: Optional[str] = Field(None, description="Creativity level(high, low) for slc-1 only")
|
||||
creativity: float | str | None = Field(None, description="slc-1/slp-2.5: enum (low/middle/high). ast-2: decimal 0.0-1.0.")
|
||||
isOptimizedMode: Optional[bool] = Field(None, description="Set to true for Starlight Creative (slc-1) only")
|
||||
prompt: str | None = Field(None, description="Descriptive scene prompt (ast-2 only)")
|
||||
sharp: float | None = Field(None, description="ast-2 pre-enhance sharpness")
|
||||
realism: float | None = Field(None, description="ast-2 realism control")
|
||||
|
||||
|
||||
class OutputInformationVideo(BaseModel):
|
||||
@@ -90,7 +93,7 @@ class Overrides(BaseModel):
|
||||
|
||||
class CreateVideoRequest(BaseModel):
|
||||
source: CreateVideoRequestSource = Field(...)
|
||||
filters: list[Union[VideoFrameInterpolationFilter, VideoEnhancementFilter]] = Field(...)
|
||||
filters: list[VideoFrameInterpolationFilter | VideoEnhancementFilter] = Field(...)
|
||||
output: OutputInformationVideo = Field(...)
|
||||
overrides: Overrides = Field(Overrides(isPaidDiffusion=True))
|
||||
|
||||
|
||||
@@ -19,6 +19,7 @@ from comfy_api_nodes.apis.bytedance import (
|
||||
Seedance2TaskCreationRequest,
|
||||
SeedanceCreateAssetRequest,
|
||||
SeedanceCreateAssetResponse,
|
||||
SeedanceCreateVisualValidateSessionRequest,
|
||||
SeedanceCreateVisualValidateSessionResponse,
|
||||
SeedanceGetVisualValidateSessionResponse,
|
||||
SeedanceVirtualLibraryCreateAssetRequest,
|
||||
@@ -196,11 +197,16 @@ def _rewrite_asset_refs(prompt: str, labels: dict[int, str]) -> str:
|
||||
return _ASSET_REF_RE.sub(_sub, prompt)
|
||||
|
||||
|
||||
async def _obtain_group_id_via_h5_auth(cls: type[IO.ComfyNode]) -> str:
|
||||
async def _obtain_group_id_via_h5_auth(
|
||||
cls: type[IO.ComfyNode],
|
||||
group_name: str | None = None,
|
||||
) -> str:
|
||||
payload = SeedanceCreateVisualValidateSessionRequest(name=group_name)
|
||||
session = await sync_op(
|
||||
cls,
|
||||
ApiEndpoint(path="/proxy/seedance/visual-validate/sessions", method="POST"),
|
||||
response_model=SeedanceCreateVisualValidateSessionResponse,
|
||||
data=payload,
|
||||
)
|
||||
logger.warning("Seedance authentication required. Open link: %s", session.h5_link)
|
||||
|
||||
@@ -229,10 +235,15 @@ async def _obtain_group_id_via_h5_auth(cls: type[IO.ComfyNode]) -> str:
|
||||
return result.group_id
|
||||
|
||||
|
||||
async def _resolve_group_id(cls: type[IO.ComfyNode], group_id: str) -> str:
|
||||
async def _resolve_group_id(
|
||||
cls: type[IO.ComfyNode],
|
||||
group_id: str,
|
||||
group_name: str | None = None,
|
||||
) -> str:
|
||||
if group_id and group_id.strip():
|
||||
return group_id.strip()
|
||||
return await _obtain_group_id_via_h5_auth(cls)
|
||||
label = (group_name or "").strip() or None
|
||||
return await _obtain_group_id_via_h5_auth(cls, group_name=label)
|
||||
|
||||
|
||||
async def _create_seedance_asset(
|
||||
@@ -1403,7 +1414,6 @@ class ByteDance2TextToVideoNode(IO.ComfyNode):
|
||||
status_extractor=lambda r: r.status,
|
||||
price_extractor=_seedance2_price_extractor(model_id, has_video_input=False),
|
||||
poll_interval=9,
|
||||
max_poll_attempts=180,
|
||||
)
|
||||
return IO.NodeOutput(await download_url_to_video_output(response.content.video_url))
|
||||
|
||||
@@ -1585,7 +1595,6 @@ class ByteDance2FirstLastFrameNode(IO.ComfyNode):
|
||||
status_extractor=lambda r: r.status,
|
||||
price_extractor=_seedance2_price_extractor(model_id, has_video_input=False),
|
||||
poll_interval=9,
|
||||
max_poll_attempts=180,
|
||||
)
|
||||
return IO.NodeOutput(await download_url_to_video_output(response.content.video_url))
|
||||
|
||||
@@ -1907,7 +1916,6 @@ class ByteDance2ReferenceNode(IO.ComfyNode):
|
||||
status_extractor=lambda r: r.status,
|
||||
price_extractor=_seedance2_price_extractor(model_id, has_video_input=has_video_input),
|
||||
poll_interval=9,
|
||||
max_poll_attempts=180,
|
||||
)
|
||||
return IO.NodeOutput(await download_url_to_video_output(response.content.video_url))
|
||||
|
||||
@@ -1939,6 +1947,55 @@ async def process_video_task(
|
||||
return IO.NodeOutput(await download_url_to_video_output(response.content.video_url))
|
||||
|
||||
|
||||
def _seedance_group_picker_input() -> IO.Combo.Input:
|
||||
"""Combo populated from /proxy/seedance/visual-validate/groups. Empty selection triggers H5 enrollment."""
|
||||
return IO.Combo.Input(
|
||||
"group_id",
|
||||
default="",
|
||||
tooltip=(
|
||||
"Pick an existing verified group, or leave empty to run real-person H5 "
|
||||
"authentication and create a new group."
|
||||
),
|
||||
remote_combo=IO.RemoteComboOptions(
|
||||
route="/proxy/seedance/visual-validate/groups",
|
||||
response_key="groups",
|
||||
item_schema=IO.RemoteItemSchema(
|
||||
value_field="group_id",
|
||||
label_field="name",
|
||||
description_field="created_at",
|
||||
search_fields=["name", "group_id"],
|
||||
),
|
||||
refresh=60_000,
|
||||
),
|
||||
optional=True,
|
||||
)
|
||||
|
||||
|
||||
def _seedance_group_name_input() -> IO.String.Input:
|
||||
return IO.String.Input(
|
||||
"group_name",
|
||||
default="",
|
||||
tooltip=(
|
||||
"Optional label for a new group. Used only when group_id is empty; the label is "
|
||||
"shown later in the group picker so you can identify this group at a glance. "
|
||||
"Up to 64 characters."
|
||||
),
|
||||
optional=True,
|
||||
)
|
||||
|
||||
|
||||
def _seedance_asset_name_input() -> IO.String.Input:
|
||||
return IO.String.Input(
|
||||
"asset_name",
|
||||
default="",
|
||||
tooltip=(
|
||||
"Optional label for the asset, shown in the asset selector dropdown. "
|
||||
"Up to 64 characters. Leave empty to identify the asset by its id."
|
||||
),
|
||||
optional=True,
|
||||
)
|
||||
|
||||
|
||||
class ByteDanceCreateImageAsset(IO.ComfyNode):
|
||||
|
||||
@classmethod
|
||||
@@ -1949,22 +2006,15 @@ class ByteDanceCreateImageAsset(IO.ComfyNode):
|
||||
category="api node/image/ByteDance",
|
||||
description=(
|
||||
"Create a Seedance 2.0 personal image asset. Uploads the input image and "
|
||||
"registers it in the given asset group. If group_id is empty, runs a real-person "
|
||||
"H5 authentication flow to create a new group before adding the asset."
|
||||
"registers it in the selected asset group. Leave group_id empty to run a "
|
||||
"real-person H5 authentication flow and create a new group; provide group_name "
|
||||
"to label the new group."
|
||||
),
|
||||
inputs=[
|
||||
IO.Image.Input("image", tooltip="Image to register as a personal asset."),
|
||||
IO.String.Input(
|
||||
"group_id",
|
||||
default="",
|
||||
tooltip="Reuse an existing Seedance asset group ID to skip repeated human verification for the "
|
||||
"same person. Leave empty to run real-person authentication in the browser and create a new group.",
|
||||
),
|
||||
# IO.String.Input(
|
||||
# "name",
|
||||
# default="",
|
||||
# tooltip="Asset name (up to 64 characters).",
|
||||
# ),
|
||||
_seedance_group_picker_input(),
|
||||
_seedance_group_name_input(),
|
||||
_seedance_asset_name_input(),
|
||||
],
|
||||
outputs=[
|
||||
IO.String.Output(display_name="asset_id"),
|
||||
@@ -1983,18 +2033,17 @@ class ByteDanceCreateImageAsset(IO.ComfyNode):
|
||||
cls,
|
||||
image: Input.Image,
|
||||
group_id: str = "",
|
||||
# name: str = "",
|
||||
group_name: str = "",
|
||||
asset_name: str = "",
|
||||
) -> IO.NodeOutput:
|
||||
# if len(name) > 64:
|
||||
# raise ValueError("Name of asset can not be greater then 64 symbols")
|
||||
validate_image_dimensions(image, min_width=300, max_width=6000, min_height=300, max_height=6000)
|
||||
validate_image_aspect_ratio(image, min_ratio=(0.4, 1), max_ratio=(2.5, 1))
|
||||
resolved_group = await _resolve_group_id(cls, group_id)
|
||||
resolved_group = await _resolve_group_id(cls, group_id, group_name=group_name)
|
||||
asset_id = await _create_seedance_asset(
|
||||
cls,
|
||||
group_id=resolved_group,
|
||||
url=await upload_image_to_comfyapi(cls, image),
|
||||
name="",
|
||||
name=asset_name.strip()[:64],
|
||||
asset_type="Image",
|
||||
)
|
||||
await _wait_for_asset_active(cls, asset_id, resolved_group)
|
||||
@@ -2016,22 +2065,15 @@ class ByteDanceCreateVideoAsset(IO.ComfyNode):
|
||||
category="api node/video/ByteDance",
|
||||
description=(
|
||||
"Create a Seedance 2.0 personal video asset. Uploads the input video and "
|
||||
"registers it in the given asset group. If group_id is empty, runs a real-person "
|
||||
"H5 authentication flow to create a new group before adding the asset."
|
||||
"registers it in the selected asset group. Leave group_id empty to run a "
|
||||
"real-person H5 authentication flow and create a new group; provide group_name "
|
||||
"to label the new group."
|
||||
),
|
||||
inputs=[
|
||||
IO.Video.Input("video", tooltip="Video to register as a personal asset."),
|
||||
IO.String.Input(
|
||||
"group_id",
|
||||
default="",
|
||||
tooltip="Reuse an existing Seedance asset group ID to skip repeated human verification for the "
|
||||
"same person. Leave empty to run real-person authentication in the browser and create a new group.",
|
||||
),
|
||||
# IO.String.Input(
|
||||
# "name",
|
||||
# default="",
|
||||
# tooltip="Asset name (up to 64 characters).",
|
||||
# ),
|
||||
_seedance_group_picker_input(),
|
||||
_seedance_group_name_input(),
|
||||
_seedance_asset_name_input(),
|
||||
],
|
||||
outputs=[
|
||||
IO.String.Output(display_name="asset_id"),
|
||||
@@ -2050,10 +2092,9 @@ class ByteDanceCreateVideoAsset(IO.ComfyNode):
|
||||
cls,
|
||||
video: Input.Video,
|
||||
group_id: str = "",
|
||||
# name: str = "",
|
||||
group_name: str = "",
|
||||
asset_name: str = "",
|
||||
) -> IO.NodeOutput:
|
||||
# if len(name) > 64:
|
||||
# raise ValueError("Name of asset can not be greater then 64 symbols")
|
||||
validate_video_duration(video, min_duration=2, max_duration=15)
|
||||
validate_video_dimensions(video, min_width=300, max_width=6000, min_height=300, max_height=6000)
|
||||
|
||||
@@ -2072,12 +2113,12 @@ class ByteDanceCreateVideoAsset(IO.ComfyNode):
|
||||
if not (24 <= fps <= 60):
|
||||
raise ValueError(f"Asset video FPS must be in [24, 60], got {fps:.2f}.")
|
||||
|
||||
resolved_group = await _resolve_group_id(cls, group_id)
|
||||
resolved_group = await _resolve_group_id(cls, group_id, group_name=group_name)
|
||||
asset_id = await _create_seedance_asset(
|
||||
cls,
|
||||
group_id=resolved_group,
|
||||
url=await upload_video_to_comfyapi(cls, video),
|
||||
name="",
|
||||
name=asset_name.strip()[:64],
|
||||
asset_type="Video",
|
||||
)
|
||||
await _wait_for_asset_active(cls, asset_id, resolved_group)
|
||||
@@ -2089,6 +2130,92 @@ class ByteDanceCreateVideoAsset(IO.ComfyNode):
|
||||
return IO.NodeOutput(asset_id, resolved_group)
|
||||
|
||||
|
||||
def _seedance_asset_picker_input(asset_type: str, preview_type: str) -> IO.Combo.Input:
|
||||
"""Combo populated from /proxy/seedance/assets, scoped to one asset_type."""
|
||||
return IO.Combo.Input(
|
||||
"asset_id",
|
||||
tooltip=(
|
||||
f"Pick a previously-created Seedance {asset_type.lower()} asset. The dropdown shows "
|
||||
"your assets across all your verified groups; type a group name to filter."
|
||||
),
|
||||
remote_combo=IO.RemoteComboOptions(
|
||||
route=f"/proxy/seedance/assets?asset_type={asset_type}",
|
||||
response_key="assets",
|
||||
item_schema=IO.RemoteItemSchema(
|
||||
value_field="asset_id",
|
||||
label_field="name",
|
||||
description_field="group_name",
|
||||
preview_url_field="url",
|
||||
preview_type=preview_type,
|
||||
search_fields=["name", "asset_id", "group_name", "group_id"],
|
||||
),
|
||||
refresh=60_000,
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
class ByteDanceSelectImageAsset(IO.ComfyNode):
|
||||
|
||||
@classmethod
|
||||
def define_schema(cls) -> IO.Schema:
|
||||
return IO.Schema(
|
||||
node_id="ByteDanceSelectImageAsset",
|
||||
display_name="ByteDance Select Image Asset",
|
||||
category="api node/image/ByteDance",
|
||||
description=(
|
||||
"Pick a previously-created Seedance image asset. Outputs the selected asset_id "
|
||||
"for use with downstream Seedance 2.0 reference/first-last-frame nodes."
|
||||
),
|
||||
inputs=[
|
||||
_seedance_asset_picker_input("Image", "image"),
|
||||
],
|
||||
outputs=[IO.String.Output(display_name="asset_id")],
|
||||
hidden=[
|
||||
IO.Hidden.auth_token_comfy_org,
|
||||
IO.Hidden.api_key_comfy_org,
|
||||
IO.Hidden.unique_id,
|
||||
],
|
||||
# is_api_node=True,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
async def execute(cls, asset_id: str) -> IO.NodeOutput:
|
||||
if not asset_id or not asset_id.strip():
|
||||
raise ValueError("asset_id is required. Pick an asset from the dropdown.")
|
||||
return IO.NodeOutput(asset_id.strip())
|
||||
|
||||
|
||||
class ByteDanceSelectVideoAsset(IO.ComfyNode):
|
||||
|
||||
@classmethod
|
||||
def define_schema(cls) -> IO.Schema:
|
||||
return IO.Schema(
|
||||
node_id="ByteDanceSelectVideoAsset",
|
||||
display_name="ByteDance Select Video Asset",
|
||||
category="api node/video/ByteDance",
|
||||
description=(
|
||||
"Pick a previously-created Seedance video asset. Outputs the selected asset_id "
|
||||
"for use with downstream Seedance 2.0 reference/first-last-frame nodes."
|
||||
),
|
||||
inputs=[
|
||||
_seedance_asset_picker_input("Video", "video"),
|
||||
],
|
||||
outputs=[IO.String.Output(display_name="asset_id")],
|
||||
hidden=[
|
||||
IO.Hidden.auth_token_comfy_org,
|
||||
IO.Hidden.api_key_comfy_org,
|
||||
IO.Hidden.unique_id,
|
||||
],
|
||||
# is_api_node=True,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
async def execute(cls, asset_id: str) -> IO.NodeOutput:
|
||||
if not asset_id or not asset_id.strip():
|
||||
raise ValueError("asset_id is required. Pick an asset from the dropdown.")
|
||||
return IO.NodeOutput(asset_id.strip())
|
||||
|
||||
|
||||
class ByteDanceExtension(ComfyExtension):
|
||||
@override
|
||||
async def get_node_list(self) -> list[type[IO.ComfyNode]]:
|
||||
@@ -2104,6 +2231,8 @@ class ByteDanceExtension(ComfyExtension):
|
||||
ByteDance2ReferenceNode,
|
||||
ByteDanceCreateImageAsset,
|
||||
ByteDanceCreateVideoAsset,
|
||||
ByteDanceSelectImageAsset,
|
||||
ByteDanceSelectVideoAsset,
|
||||
]
|
||||
|
||||
|
||||
|
||||
@@ -233,6 +233,44 @@ class ElevenLabsVoiceSelector(IO.ComfyNode):
|
||||
return IO.NodeOutput(voice_id)
|
||||
|
||||
|
||||
class ElevenLabsRichVoiceSelector(IO.ComfyNode):
|
||||
@classmethod
|
||||
def define_schema(cls) -> IO.Schema:
|
||||
return IO.Schema(
|
||||
node_id="ElevenLabsRichVoiceSelector",
|
||||
display_name="ElevenLabs Voice Selector (Rich)",
|
||||
category="api node/audio/ElevenLabs",
|
||||
description="Select an ElevenLabs voice with audio preview and rich metadata.",
|
||||
inputs=[
|
||||
IO.Combo.Input(
|
||||
"voice",
|
||||
remote_combo=IO.RemoteComboOptions(
|
||||
route="/proxy/elevenlabs/v2/voices?page_size=100",
|
||||
response_key="items",
|
||||
refresh_button=True,
|
||||
refresh=43200000,
|
||||
item_schema=IO.RemoteItemSchema(
|
||||
value_field="voice_id",
|
||||
label_field="name",
|
||||
preview_url_field="preview_url",
|
||||
preview_type="audio",
|
||||
search_fields=["name", "labels.gender", "labels.accent", "labels.use_case"],
|
||||
),
|
||||
),
|
||||
tooltip="Choose a voice with audio preview.",
|
||||
),
|
||||
],
|
||||
outputs=[
|
||||
IO.Custom(ELEVENLABS_VOICE).Output(display_name="voice"),
|
||||
],
|
||||
is_api_node=False,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, voice: str) -> IO.NodeOutput:
|
||||
return IO.NodeOutput(voice) # voice is already the voice_id from item_schema.value_field
|
||||
|
||||
|
||||
class ElevenLabsTextToSpeech(IO.ComfyNode):
|
||||
@classmethod
|
||||
def define_schema(cls) -> IO.Schema:
|
||||
@@ -911,6 +949,7 @@ class ElevenLabsExtension(ComfyExtension):
|
||||
return [
|
||||
ElevenLabsSpeechToText,
|
||||
ElevenLabsVoiceSelector,
|
||||
ElevenLabsRichVoiceSelector,
|
||||
ElevenLabsTextToSpeech,
|
||||
ElevenLabsAudioIsolation,
|
||||
ElevenLabsTextToSoundEffects,
|
||||
|
||||
@@ -178,7 +178,6 @@ class HitPawGeneralImageEnhance(IO.ComfyNode):
|
||||
status_extractor=lambda x: x.data.status,
|
||||
price_extractor=lambda x: request_price,
|
||||
poll_interval=10.0,
|
||||
max_poll_attempts=480,
|
||||
)
|
||||
return IO.NodeOutput(await download_url_to_image_tensor(final_response.data.res_url))
|
||||
|
||||
@@ -324,7 +323,6 @@ class HitPawVideoEnhance(IO.ComfyNode):
|
||||
status_extractor=lambda x: x.data.status,
|
||||
price_extractor=lambda x: request_price,
|
||||
poll_interval=10.0,
|
||||
max_poll_attempts=320,
|
||||
)
|
||||
return IO.NodeOutput(await download_url_to_video_output(final_response.data.res_url))
|
||||
|
||||
|
||||
@@ -276,7 +276,6 @@ async def finish_omni_video_task(cls: type[IO.ComfyNode], response: TaskStatusRe
|
||||
cls,
|
||||
ApiEndpoint(path=f"/proxy/kling/v1/videos/omni-video/{response.data.task_id}"),
|
||||
response_model=TaskStatusResponse,
|
||||
max_poll_attempts=280,
|
||||
status_extractor=lambda r: (r.data.task_status if r.data else None),
|
||||
)
|
||||
return IO.NodeOutput(await download_url_to_video_output(final_response.data.task_result.videos[0].url))
|
||||
@@ -3062,7 +3061,6 @@ class KlingVideoNode(IO.ComfyNode):
|
||||
cls,
|
||||
ApiEndpoint(path=poll_path),
|
||||
response_model=TaskStatusResponse,
|
||||
max_poll_attempts=280,
|
||||
status_extractor=lambda r: (r.data.task_status if r.data else None),
|
||||
)
|
||||
return IO.NodeOutput(await download_url_to_video_output(final_response.data.task_result.videos[0].url))
|
||||
@@ -3188,7 +3186,6 @@ class KlingFirstLastFrameNode(IO.ComfyNode):
|
||||
cls,
|
||||
ApiEndpoint(path=f"/proxy/kling/v1/videos/image2video/{response.data.task_id}"),
|
||||
response_model=TaskStatusResponse,
|
||||
max_poll_attempts=280,
|
||||
status_extractor=lambda r: (r.data.task_status if r.data else None),
|
||||
)
|
||||
return IO.NodeOutput(await download_url_to_video_output(final_response.data.task_result.videos[0].url))
|
||||
@@ -3291,6 +3288,53 @@ class KlingAvatarNode(IO.ComfyNode):
|
||||
return IO.NodeOutput(await download_url_to_video_output(final_response.data.task_result.videos[0].url))
|
||||
|
||||
|
||||
KLING_ELEMENT_ID = "KLING_ELEMENT_ID"
|
||||
|
||||
|
||||
class KlingElementSelector(IO.ComfyNode):
|
||||
"""Select a Kling preset element (character, scene, effect, etc.) for use in video generation."""
|
||||
|
||||
@classmethod
|
||||
def define_schema(cls) -> IO.Schema:
|
||||
return IO.Schema(
|
||||
node_id="KlingElementSelector",
|
||||
display_name="Kling Element Selector",
|
||||
category="api node/video/Kling",
|
||||
description="Browse and select a Kling preset element with image preview. Elements provide consistent characters, scenes, costumes, and effects for video generation.",
|
||||
inputs=[
|
||||
IO.Combo.Input(
|
||||
"element",
|
||||
remote_combo=IO.RemoteComboOptions(
|
||||
route="/proxy/kling/v1/general/advanced-presets-elements",
|
||||
refresh_button=True,
|
||||
refresh=43200000,
|
||||
response_key="data",
|
||||
item_schema=IO.RemoteItemSchema(
|
||||
value_field="task_result.elements.0.element_id",
|
||||
label_field="task_result.elements.0.element_name",
|
||||
preview_url_field="task_result.elements.0.element_image_list.frontal_image",
|
||||
preview_type="image",
|
||||
description_field="task_result.elements.0.element_description",
|
||||
search_fields=["task_result.elements.0.element_name", "task_result.elements.0.element_description"],
|
||||
),
|
||||
),
|
||||
tooltip="Select a preset element to use in video generation.",
|
||||
),
|
||||
],
|
||||
outputs=[IO.Custom(KLING_ELEMENT_ID).Output(display_name="element_id")],
|
||||
hidden=[
|
||||
IO.Hidden.auth_token_comfy_org,
|
||||
IO.Hidden.api_key_comfy_org,
|
||||
IO.Hidden.unique_id,
|
||||
],
|
||||
is_api_node=False,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
async def execute(cls, element: str) -> IO.NodeOutput:
|
||||
return IO.NodeOutput(element)
|
||||
|
||||
|
||||
class KlingExtension(ComfyExtension):
|
||||
@override
|
||||
async def get_node_list(self) -> list[type[IO.ComfyNode]]:
|
||||
@@ -3320,6 +3364,7 @@ class KlingExtension(ComfyExtension):
|
||||
KlingVideoNode,
|
||||
KlingFirstLastFrameNode,
|
||||
KlingAvatarNode,
|
||||
KlingElementSelector,
|
||||
]
|
||||
|
||||
|
||||
|
||||
@@ -230,7 +230,6 @@ class MagnificImageUpscalerCreativeNode(IO.ComfyNode):
|
||||
status_extractor=lambda x: x.status,
|
||||
price_extractor=lambda _: price_usd,
|
||||
poll_interval=10.0,
|
||||
max_poll_attempts=480,
|
||||
)
|
||||
return IO.NodeOutput(await download_url_to_image_tensor(final_response.generated[0]))
|
||||
|
||||
@@ -391,7 +390,6 @@ class MagnificImageUpscalerPreciseV2Node(IO.ComfyNode):
|
||||
status_extractor=lambda x: x.status,
|
||||
price_extractor=lambda _: price_usd,
|
||||
poll_interval=10.0,
|
||||
max_poll_attempts=480,
|
||||
)
|
||||
return IO.NodeOutput(await download_url_to_image_tensor(final_response.generated[0]))
|
||||
|
||||
@@ -541,7 +539,6 @@ class MagnificImageStyleTransferNode(IO.ComfyNode):
|
||||
response_model=TaskResponse,
|
||||
status_extractor=lambda x: x.status,
|
||||
poll_interval=10.0,
|
||||
max_poll_attempts=480,
|
||||
)
|
||||
return IO.NodeOutput(await download_url_to_image_tensor(final_response.generated[0]))
|
||||
|
||||
@@ -782,7 +779,6 @@ class MagnificImageRelightNode(IO.ComfyNode):
|
||||
response_model=TaskResponse,
|
||||
status_extractor=lambda x: x.status,
|
||||
poll_interval=10.0,
|
||||
max_poll_attempts=480,
|
||||
)
|
||||
return IO.NodeOutput(await download_url_to_image_tensor(final_response.generated[0]))
|
||||
|
||||
@@ -924,7 +920,6 @@ class MagnificImageSkinEnhancerNode(IO.ComfyNode):
|
||||
response_model=TaskResponse,
|
||||
status_extractor=lambda x: x.status,
|
||||
poll_interval=10.0,
|
||||
max_poll_attempts=480,
|
||||
)
|
||||
return IO.NodeOutput(await download_url_to_image_tensor(final_response.generated[0]))
|
||||
|
||||
|
||||
@@ -1,534 +0,0 @@
|
||||
import logging
|
||||
|
||||
from typing_extensions import override
|
||||
|
||||
from comfy_api.latest import IO, ComfyExtension, Input
|
||||
from comfy_api_nodes.apis.moonvalley import (
|
||||
MoonvalleyPromptResponse,
|
||||
MoonvalleyTextToVideoInferenceParams,
|
||||
MoonvalleyTextToVideoRequest,
|
||||
MoonvalleyVideoToVideoInferenceParams,
|
||||
MoonvalleyVideoToVideoRequest,
|
||||
)
|
||||
from comfy_api_nodes.util import (
|
||||
ApiEndpoint,
|
||||
download_url_to_video_output,
|
||||
poll_op,
|
||||
sync_op,
|
||||
trim_video,
|
||||
upload_images_to_comfyapi,
|
||||
upload_video_to_comfyapi,
|
||||
validate_container_format_is_mp4,
|
||||
validate_image_dimensions,
|
||||
validate_string,
|
||||
)
|
||||
|
||||
API_UPLOADS_ENDPOINT = "/proxy/moonvalley/uploads"
|
||||
API_PROMPTS_ENDPOINT = "/proxy/moonvalley/prompts"
|
||||
API_VIDEO2VIDEO_ENDPOINT = "/proxy/moonvalley/prompts/video-to-video"
|
||||
API_TXT2VIDEO_ENDPOINT = "/proxy/moonvalley/prompts/text-to-video"
|
||||
API_IMG2VIDEO_ENDPOINT = "/proxy/moonvalley/prompts/image-to-video"
|
||||
|
||||
MIN_WIDTH = 300
|
||||
MIN_HEIGHT = 300
|
||||
|
||||
MAX_WIDTH = 10000
|
||||
MAX_HEIGHT = 10000
|
||||
|
||||
MIN_VID_WIDTH = 300
|
||||
MIN_VID_HEIGHT = 300
|
||||
|
||||
MAX_VID_WIDTH = 10000
|
||||
MAX_VID_HEIGHT = 10000
|
||||
|
||||
MAX_VIDEO_SIZE = 1024 * 1024 * 1024 # 1 GB max for in-memory video processing
|
||||
|
||||
MOONVALLEY_MAREY_MAX_PROMPT_LENGTH = 5000
|
||||
|
||||
|
||||
def is_valid_task_creation_response(response: MoonvalleyPromptResponse) -> bool:
|
||||
"""Verifies that the initial response contains a task ID."""
|
||||
return bool(response.id)
|
||||
|
||||
|
||||
def validate_task_creation_response(response) -> None:
|
||||
if not is_valid_task_creation_response(response):
|
||||
error_msg = f"Moonvalley Marey API: Initial request failed. Code: {response.code}, Message: {response.message}, Data: {response}"
|
||||
logging.error(error_msg)
|
||||
raise RuntimeError(error_msg)
|
||||
|
||||
|
||||
def validate_video_to_video_input(video: Input.Video) -> Input.Video:
|
||||
"""
|
||||
Validates and processes video input for Moonvalley Video-to-Video generation.
|
||||
|
||||
Args:
|
||||
video: Input video to validate
|
||||
|
||||
Returns:
|
||||
Validated and potentially trimmed video
|
||||
|
||||
Raises:
|
||||
ValueError: If video doesn't meet requirements
|
||||
MoonvalleyApiError: If video duration is too short
|
||||
"""
|
||||
width, height = _get_video_dimensions(video)
|
||||
_validate_video_dimensions(width, height)
|
||||
validate_container_format_is_mp4(video)
|
||||
|
||||
return _validate_and_trim_duration(video)
|
||||
|
||||
|
||||
def _get_video_dimensions(video: Input.Video) -> tuple[int, int]:
|
||||
"""Extracts video dimensions with error handling."""
|
||||
try:
|
||||
return video.get_dimensions()
|
||||
except Exception as e:
|
||||
logging.error("Error getting dimensions of video: %s", e)
|
||||
raise ValueError(f"Cannot get video dimensions: {e}") from e
|
||||
|
||||
|
||||
def _validate_video_dimensions(width: int, height: int) -> None:
|
||||
"""Validates video dimensions meet Moonvalley V2V requirements."""
|
||||
supported_resolutions = {
|
||||
(1920, 1080),
|
||||
(1080, 1920),
|
||||
(1152, 1152),
|
||||
(1536, 1152),
|
||||
(1152, 1536),
|
||||
}
|
||||
|
||||
if (width, height) not in supported_resolutions:
|
||||
supported_list = ", ".join([f"{w}x{h}" for w, h in sorted(supported_resolutions)])
|
||||
raise ValueError(f"Resolution {width}x{height} not supported. Supported: {supported_list}")
|
||||
|
||||
|
||||
def _validate_and_trim_duration(video: Input.Video) -> Input.Video:
|
||||
"""Validates video duration and trims to 5 seconds if needed."""
|
||||
duration = video.get_duration()
|
||||
_validate_minimum_duration(duration)
|
||||
return _trim_if_too_long(video, duration)
|
||||
|
||||
|
||||
def _validate_minimum_duration(duration: float) -> None:
|
||||
"""Ensures video is at least 5 seconds long."""
|
||||
if duration < 5:
|
||||
raise ValueError("Input video must be at least 5 seconds long.")
|
||||
|
||||
|
||||
def _trim_if_too_long(video: Input.Video, duration: float) -> Input.Video:
|
||||
"""Trims video to 5 seconds if longer."""
|
||||
if duration > 5:
|
||||
return trim_video(video, 5)
|
||||
return video
|
||||
|
||||
|
||||
def parse_width_height_from_res(resolution: str):
|
||||
# Accepts a string like "16:9 (1920 x 1080)" and returns width, height as a dict
|
||||
res_map = {
|
||||
"16:9 (1920 x 1080)": {"width": 1920, "height": 1080},
|
||||
"9:16 (1080 x 1920)": {"width": 1080, "height": 1920},
|
||||
"1:1 (1152 x 1152)": {"width": 1152, "height": 1152},
|
||||
"4:3 (1536 x 1152)": {"width": 1536, "height": 1152},
|
||||
"3:4 (1152 x 1536)": {"width": 1152, "height": 1536},
|
||||
# "21:9 (2560 x 1080)": {"width": 2560, "height": 1080},
|
||||
}
|
||||
return res_map.get(resolution, {"width": 1920, "height": 1080})
|
||||
|
||||
|
||||
def parse_control_parameter(value):
|
||||
control_map = {
|
||||
"Motion Transfer": "motion_control",
|
||||
"Canny": "canny_control",
|
||||
"Pose Transfer": "pose_control",
|
||||
"Depth": "depth_control",
|
||||
}
|
||||
return control_map.get(value, control_map["Motion Transfer"])
|
||||
|
||||
|
||||
async def get_response(cls: type[IO.ComfyNode], task_id: str) -> MoonvalleyPromptResponse:
|
||||
return await poll_op(
|
||||
cls,
|
||||
ApiEndpoint(path=f"{API_PROMPTS_ENDPOINT}/{task_id}"),
|
||||
response_model=MoonvalleyPromptResponse,
|
||||
status_extractor=lambda r: (r.status if r and r.status else None),
|
||||
poll_interval=16.0,
|
||||
max_poll_attempts=240,
|
||||
)
|
||||
|
||||
|
||||
class MoonvalleyImg2VideoNode(IO.ComfyNode):
|
||||
|
||||
@classmethod
|
||||
def define_schema(cls) -> IO.Schema:
|
||||
return IO.Schema(
|
||||
node_id="MoonvalleyImg2VideoNode",
|
||||
display_name="Moonvalley Marey Image to Video",
|
||||
category="api node/video/Moonvalley Marey",
|
||||
description="Moonvalley Marey Image to Video Node",
|
||||
inputs=[
|
||||
IO.Image.Input(
|
||||
"image",
|
||||
tooltip="The reference image used to generate the video",
|
||||
),
|
||||
IO.String.Input(
|
||||
"prompt",
|
||||
multiline=True,
|
||||
),
|
||||
IO.String.Input(
|
||||
"negative_prompt",
|
||||
multiline=True,
|
||||
default="<synthetic> <scene cut> gopro, bright, contrast, static, overexposed, vignette, "
|
||||
"artifacts, still, noise, texture, scanlines, videogame, 360 camera, VR, transition, "
|
||||
"flare, saturation, distorted, warped, wide angle, saturated, vibrant, glowing, "
|
||||
"cross dissolve, cheesy, ugly hands, mutated hands, mutant, disfigured, extra fingers, "
|
||||
"blown out, horrible, blurry, worst quality, bad, dissolve, melt, fade in, fade out, "
|
||||
"wobbly, weird, low quality, plastic, stock footage, video camera, boring",
|
||||
tooltip="Negative prompt text",
|
||||
),
|
||||
IO.Combo.Input(
|
||||
"resolution",
|
||||
options=[
|
||||
"16:9 (1920 x 1080)",
|
||||
"9:16 (1080 x 1920)",
|
||||
"1:1 (1152 x 1152)",
|
||||
"4:3 (1536 x 1152)",
|
||||
"3:4 (1152 x 1536)",
|
||||
# "21:9 (2560 x 1080)",
|
||||
],
|
||||
default="16:9 (1920 x 1080)",
|
||||
tooltip="Resolution of the output video",
|
||||
),
|
||||
IO.Float.Input(
|
||||
"prompt_adherence",
|
||||
default=4.5,
|
||||
min=1.0,
|
||||
max=20.0,
|
||||
step=1.0,
|
||||
tooltip="Guidance scale for generation control",
|
||||
),
|
||||
IO.Int.Input(
|
||||
"seed",
|
||||
default=9,
|
||||
min=0,
|
||||
max=4294967295,
|
||||
step=1,
|
||||
display_mode=IO.NumberDisplay.number,
|
||||
tooltip="Random seed value",
|
||||
control_after_generate=True,
|
||||
),
|
||||
IO.Int.Input(
|
||||
"steps",
|
||||
default=80,
|
||||
min=75, # steps should be greater or equal to cooldown_steps(75) + warmup_steps(0)
|
||||
max=100,
|
||||
step=1,
|
||||
tooltip="Number of denoising steps",
|
||||
),
|
||||
],
|
||||
outputs=[IO.Video.Output()],
|
||||
hidden=[
|
||||
IO.Hidden.auth_token_comfy_org,
|
||||
IO.Hidden.api_key_comfy_org,
|
||||
IO.Hidden.unique_id,
|
||||
],
|
||||
is_api_node=True,
|
||||
price_badge=IO.PriceBadge(
|
||||
depends_on=IO.PriceBadgeDepends(),
|
||||
expr="""{"type":"usd","usd": 1.5}""",
|
||||
),
|
||||
)
|
||||
|
||||
@classmethod
|
||||
async def execute(
|
||||
cls,
|
||||
image: Input.Image,
|
||||
prompt: str,
|
||||
negative_prompt: str,
|
||||
resolution: str,
|
||||
prompt_adherence: float,
|
||||
seed: int,
|
||||
steps: int,
|
||||
) -> IO.NodeOutput:
|
||||
validate_image_dimensions(image, min_width=300, min_height=300, max_height=MAX_HEIGHT, max_width=MAX_WIDTH)
|
||||
validate_string(prompt, min_length=1, max_length=MOONVALLEY_MAREY_MAX_PROMPT_LENGTH)
|
||||
validate_string(negative_prompt, field_name="negative_prompt", max_length=MOONVALLEY_MAREY_MAX_PROMPT_LENGTH)
|
||||
width_height = parse_width_height_from_res(resolution)
|
||||
|
||||
inference_params = MoonvalleyTextToVideoInferenceParams(
|
||||
negative_prompt=negative_prompt,
|
||||
steps=steps,
|
||||
seed=seed,
|
||||
guidance_scale=prompt_adherence,
|
||||
width=width_height["width"],
|
||||
height=width_height["height"],
|
||||
use_negative_prompts=True,
|
||||
)
|
||||
|
||||
# Get MIME type from tensor - assuming PNG format for image tensors
|
||||
mime_type = "image/png"
|
||||
image_url = (await upload_images_to_comfyapi(cls, image, max_images=1, mime_type=mime_type))[0]
|
||||
task_creation_response = await sync_op(
|
||||
cls,
|
||||
endpoint=ApiEndpoint(path=API_IMG2VIDEO_ENDPOINT, method="POST"),
|
||||
response_model=MoonvalleyPromptResponse,
|
||||
data=MoonvalleyTextToVideoRequest(
|
||||
image_url=image_url, prompt_text=prompt, inference_params=inference_params
|
||||
),
|
||||
)
|
||||
validate_task_creation_response(task_creation_response)
|
||||
final_response = await get_response(cls, task_creation_response.id)
|
||||
video = await download_url_to_video_output(final_response.output_url)
|
||||
return IO.NodeOutput(video)
|
||||
|
||||
|
||||
class MoonvalleyVideo2VideoNode(IO.ComfyNode):
|
||||
|
||||
@classmethod
|
||||
def define_schema(cls) -> IO.Schema:
|
||||
return IO.Schema(
|
||||
node_id="MoonvalleyVideo2VideoNode",
|
||||
display_name="Moonvalley Marey Video to Video",
|
||||
category="api node/video/Moonvalley Marey",
|
||||
description="",
|
||||
inputs=[
|
||||
IO.String.Input(
|
||||
"prompt",
|
||||
multiline=True,
|
||||
tooltip="Describes the video to generate",
|
||||
),
|
||||
IO.String.Input(
|
||||
"negative_prompt",
|
||||
multiline=True,
|
||||
default="<synthetic> <scene cut> gopro, bright, contrast, static, overexposed, vignette, "
|
||||
"artifacts, still, noise, texture, scanlines, videogame, 360 camera, VR, transition, "
|
||||
"flare, saturation, distorted, warped, wide angle, saturated, vibrant, glowing, "
|
||||
"cross dissolve, cheesy, ugly hands, mutated hands, mutant, disfigured, extra fingers, "
|
||||
"blown out, horrible, blurry, worst quality, bad, dissolve, melt, fade in, fade out, "
|
||||
"wobbly, weird, low quality, plastic, stock footage, video camera, boring",
|
||||
tooltip="Negative prompt text",
|
||||
),
|
||||
IO.Int.Input(
|
||||
"seed",
|
||||
default=9,
|
||||
min=0,
|
||||
max=4294967295,
|
||||
step=1,
|
||||
display_mode=IO.NumberDisplay.number,
|
||||
tooltip="Random seed value",
|
||||
control_after_generate=False,
|
||||
),
|
||||
IO.Video.Input(
|
||||
"video",
|
||||
tooltip="The reference video used to generate the output video. Must be at least 5 seconds long. "
|
||||
"Videos longer than 5s will be automatically trimmed. Only MP4 format supported.",
|
||||
),
|
||||
IO.Combo.Input(
|
||||
"control_type",
|
||||
options=["Motion Transfer", "Pose Transfer"],
|
||||
default="Motion Transfer",
|
||||
optional=True,
|
||||
),
|
||||
IO.Int.Input(
|
||||
"motion_intensity",
|
||||
default=100,
|
||||
min=0,
|
||||
max=100,
|
||||
step=1,
|
||||
tooltip="Only used if control_type is 'Motion Transfer'",
|
||||
optional=True,
|
||||
),
|
||||
IO.Int.Input(
|
||||
"steps",
|
||||
default=60,
|
||||
min=60, # steps should be greater or equal to cooldown_steps(36) + warmup_steps(24)
|
||||
max=100,
|
||||
step=1,
|
||||
display_mode=IO.NumberDisplay.number,
|
||||
tooltip="Number of inference steps",
|
||||
),
|
||||
],
|
||||
outputs=[IO.Video.Output()],
|
||||
hidden=[
|
||||
IO.Hidden.auth_token_comfy_org,
|
||||
IO.Hidden.api_key_comfy_org,
|
||||
IO.Hidden.unique_id,
|
||||
],
|
||||
is_api_node=True,
|
||||
price_badge=IO.PriceBadge(
|
||||
depends_on=IO.PriceBadgeDepends(),
|
||||
expr="""{"type":"usd","usd": 2.25}""",
|
||||
),
|
||||
)
|
||||
|
||||
@classmethod
|
||||
async def execute(
|
||||
cls,
|
||||
prompt: str,
|
||||
negative_prompt: str,
|
||||
seed: int,
|
||||
video: Input.Video | None = None,
|
||||
control_type: str = "Motion Transfer",
|
||||
motion_intensity: int | None = 100,
|
||||
steps=60,
|
||||
prompt_adherence=4.5,
|
||||
) -> IO.NodeOutput:
|
||||
validated_video = validate_video_to_video_input(video)
|
||||
video_url = await upload_video_to_comfyapi(cls, validated_video)
|
||||
validate_string(prompt, min_length=1, max_length=MOONVALLEY_MAREY_MAX_PROMPT_LENGTH)
|
||||
validate_string(negative_prompt, field_name="negative_prompt", max_length=MOONVALLEY_MAREY_MAX_PROMPT_LENGTH)
|
||||
|
||||
# Only include motion_intensity for Motion Transfer
|
||||
control_params = {}
|
||||
if control_type == "Motion Transfer" and motion_intensity is not None:
|
||||
control_params["motion_intensity"] = motion_intensity
|
||||
|
||||
inference_params = MoonvalleyVideoToVideoInferenceParams(
|
||||
negative_prompt=negative_prompt,
|
||||
seed=seed,
|
||||
control_params=control_params,
|
||||
steps=steps,
|
||||
guidance_scale=prompt_adherence,
|
||||
)
|
||||
|
||||
task_creation_response = await sync_op(
|
||||
cls,
|
||||
endpoint=ApiEndpoint(path=API_VIDEO2VIDEO_ENDPOINT, method="POST"),
|
||||
response_model=MoonvalleyPromptResponse,
|
||||
data=MoonvalleyVideoToVideoRequest(
|
||||
control_type=parse_control_parameter(control_type),
|
||||
video_url=video_url,
|
||||
prompt_text=prompt,
|
||||
inference_params=inference_params,
|
||||
),
|
||||
)
|
||||
validate_task_creation_response(task_creation_response)
|
||||
final_response = await get_response(cls, task_creation_response.id)
|
||||
return IO.NodeOutput(await download_url_to_video_output(final_response.output_url))
|
||||
|
||||
|
||||
class MoonvalleyTxt2VideoNode(IO.ComfyNode):
|
||||
|
||||
@classmethod
|
||||
def define_schema(cls) -> IO.Schema:
|
||||
return IO.Schema(
|
||||
node_id="MoonvalleyTxt2VideoNode",
|
||||
display_name="Moonvalley Marey Text to Video",
|
||||
category="api node/video/Moonvalley Marey",
|
||||
description="",
|
||||
inputs=[
|
||||
IO.String.Input(
|
||||
"prompt",
|
||||
multiline=True,
|
||||
),
|
||||
IO.String.Input(
|
||||
"negative_prompt",
|
||||
multiline=True,
|
||||
default="<synthetic> <scene cut> gopro, bright, contrast, static, overexposed, vignette, "
|
||||
"artifacts, still, noise, texture, scanlines, videogame, 360 camera, VR, transition, "
|
||||
"flare, saturation, distorted, warped, wide angle, saturated, vibrant, glowing, "
|
||||
"cross dissolve, cheesy, ugly hands, mutated hands, mutant, disfigured, extra fingers, "
|
||||
"blown out, horrible, blurry, worst quality, bad, dissolve, melt, fade in, fade out, "
|
||||
"wobbly, weird, low quality, plastic, stock footage, video camera, boring",
|
||||
tooltip="Negative prompt text",
|
||||
),
|
||||
IO.Combo.Input(
|
||||
"resolution",
|
||||
options=[
|
||||
"16:9 (1920 x 1080)",
|
||||
"9:16 (1080 x 1920)",
|
||||
"1:1 (1152 x 1152)",
|
||||
"4:3 (1536 x 1152)",
|
||||
"3:4 (1152 x 1536)",
|
||||
"21:9 (2560 x 1080)",
|
||||
],
|
||||
default="16:9 (1920 x 1080)",
|
||||
tooltip="Resolution of the output video",
|
||||
),
|
||||
IO.Float.Input(
|
||||
"prompt_adherence",
|
||||
default=4.0,
|
||||
min=1.0,
|
||||
max=20.0,
|
||||
step=1.0,
|
||||
tooltip="Guidance scale for generation control",
|
||||
),
|
||||
IO.Int.Input(
|
||||
"seed",
|
||||
default=9,
|
||||
min=0,
|
||||
max=4294967295,
|
||||
step=1,
|
||||
display_mode=IO.NumberDisplay.number,
|
||||
control_after_generate=True,
|
||||
tooltip="Random seed value",
|
||||
),
|
||||
IO.Int.Input(
|
||||
"steps",
|
||||
default=80,
|
||||
min=75, # steps should be greater or equal to cooldown_steps(75) + warmup_steps(0)
|
||||
max=100,
|
||||
step=1,
|
||||
tooltip="Inference steps",
|
||||
),
|
||||
],
|
||||
outputs=[IO.Video.Output()],
|
||||
hidden=[
|
||||
IO.Hidden.auth_token_comfy_org,
|
||||
IO.Hidden.api_key_comfy_org,
|
||||
IO.Hidden.unique_id,
|
||||
],
|
||||
is_api_node=True,
|
||||
price_badge=IO.PriceBadge(
|
||||
depends_on=IO.PriceBadgeDepends(),
|
||||
expr="""{"type":"usd","usd": 1.5}""",
|
||||
),
|
||||
)
|
||||
|
||||
@classmethod
|
||||
async def execute(
|
||||
cls,
|
||||
prompt: str,
|
||||
negative_prompt: str,
|
||||
resolution: str,
|
||||
prompt_adherence: float,
|
||||
seed: int,
|
||||
steps: int,
|
||||
) -> IO.NodeOutput:
|
||||
validate_string(prompt, min_length=1, max_length=MOONVALLEY_MAREY_MAX_PROMPT_LENGTH)
|
||||
validate_string(negative_prompt, field_name="negative_prompt", max_length=MOONVALLEY_MAREY_MAX_PROMPT_LENGTH)
|
||||
width_height = parse_width_height_from_res(resolution)
|
||||
|
||||
inference_params = MoonvalleyTextToVideoInferenceParams(
|
||||
negative_prompt=negative_prompt,
|
||||
steps=steps,
|
||||
seed=seed,
|
||||
guidance_scale=prompt_adherence,
|
||||
num_frames=128,
|
||||
width=width_height["width"],
|
||||
height=width_height["height"],
|
||||
)
|
||||
|
||||
task_creation_response = await sync_op(
|
||||
cls,
|
||||
endpoint=ApiEndpoint(path=API_TXT2VIDEO_ENDPOINT, method="POST"),
|
||||
response_model=MoonvalleyPromptResponse,
|
||||
data=MoonvalleyTextToVideoRequest(prompt_text=prompt, inference_params=inference_params),
|
||||
)
|
||||
validate_task_creation_response(task_creation_response)
|
||||
final_response = await get_response(cls, task_creation_response.id)
|
||||
return IO.NodeOutput(await download_url_to_video_output(final_response.output_url))
|
||||
|
||||
|
||||
class MoonvalleyExtension(ComfyExtension):
|
||||
@override
|
||||
async def get_node_list(self) -> list[type[IO.ComfyNode]]:
|
||||
return [
|
||||
MoonvalleyImg2VideoNode,
|
||||
MoonvalleyTxt2VideoNode,
|
||||
MoonvalleyVideo2VideoNode,
|
||||
]
|
||||
|
||||
|
||||
async def comfy_entrypoint() -> MoonvalleyExtension:
|
||||
return MoonvalleyExtension()
|
||||
@@ -454,7 +454,6 @@ class OpenAIGPTImage1(IO.ComfyNode):
|
||||
step=16,
|
||||
tooltip="Used only when `size` is 'Custom'. Must be a multiple of 16 (GPT Image 2 only).",
|
||||
optional=True,
|
||||
advanced=True,
|
||||
),
|
||||
IO.Int.Input(
|
||||
"custom_height",
|
||||
@@ -464,7 +463,6 @@ class OpenAIGPTImage1(IO.ComfyNode):
|
||||
step=16,
|
||||
tooltip="Used only when `size` is 'Custom'. Must be a multiple of 16 (GPT Image 2 only).",
|
||||
optional=True,
|
||||
advanced=True,
|
||||
),
|
||||
],
|
||||
outputs=[
|
||||
|
||||
@@ -36,11 +36,15 @@ from comfy_api_nodes.util import (
|
||||
)
|
||||
|
||||
UPSCALER_MODELS_MAP = {
|
||||
"Astra 2": "ast-2",
|
||||
"Starlight (Astra) Fast": "slf-1",
|
||||
"Starlight (Astra) Creative": "slc-1",
|
||||
"Starlight Precise 2.5": "slp-2.5",
|
||||
}
|
||||
|
||||
AST2_MAX_FRAMES = 9000
|
||||
AST2_MAX_FRAMES_WITH_PROMPT = 450
|
||||
|
||||
|
||||
class TopazImageEnhance(IO.ComfyNode):
|
||||
@classmethod
|
||||
@@ -230,13 +234,20 @@ class TopazVideoEnhance(IO.ComfyNode):
|
||||
def define_schema(cls):
|
||||
return IO.Schema(
|
||||
node_id="TopazVideoEnhance",
|
||||
display_name="Topaz Video Enhance",
|
||||
display_name="Topaz Video Enhance (Legacy)",
|
||||
category="api node/video/Topaz",
|
||||
description="Breathe new life into video with powerful upscaling and recovery technology.",
|
||||
inputs=[
|
||||
IO.Video.Input("video"),
|
||||
IO.Boolean.Input("upscaler_enabled", default=True),
|
||||
IO.Combo.Input("upscaler_model", options=list(UPSCALER_MODELS_MAP.keys())),
|
||||
IO.Combo.Input(
|
||||
"upscaler_model",
|
||||
options=[
|
||||
"Starlight (Astra) Fast",
|
||||
"Starlight (Astra) Creative",
|
||||
"Starlight Precise 2.5",
|
||||
],
|
||||
),
|
||||
IO.Combo.Input("upscaler_resolution", options=["FullHD (1080p)", "4K (2160p)"]),
|
||||
IO.Combo.Input(
|
||||
"upscaler_creativity",
|
||||
@@ -304,6 +315,7 @@ class TopazVideoEnhance(IO.ComfyNode):
|
||||
IO.Hidden.unique_id,
|
||||
],
|
||||
is_api_node=True,
|
||||
is_deprecated=True,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
@@ -453,7 +465,350 @@ class TopazVideoEnhance(IO.ComfyNode):
|
||||
progress_extractor=lambda x: getattr(x, "progress", 0),
|
||||
price_extractor=lambda x: (x.estimates.cost[0] * 0.08 if x.estimates and x.estimates.cost[0] else None),
|
||||
poll_interval=10.0,
|
||||
max_poll_attempts=320,
|
||||
)
|
||||
return IO.NodeOutput(await download_url_to_video_output(final_response.download.url))
|
||||
|
||||
|
||||
class TopazVideoEnhanceV2(IO.ComfyNode):
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return IO.Schema(
|
||||
node_id="TopazVideoEnhanceV2",
|
||||
display_name="Topaz Video Enhance",
|
||||
category="api node/video/Topaz",
|
||||
description="Breathe new life into video with powerful upscaling and recovery technology.",
|
||||
inputs=[
|
||||
IO.Video.Input("video"),
|
||||
IO.DynamicCombo.Input(
|
||||
"upscaler_model",
|
||||
options=[
|
||||
IO.DynamicCombo.Option(
|
||||
"Astra 2",
|
||||
[
|
||||
IO.Combo.Input("upscaler_resolution", options=["FullHD (1080p)", "4K (2160p)"]),
|
||||
IO.Float.Input(
|
||||
"creativity",
|
||||
default=0.5,
|
||||
min=0.0,
|
||||
max=1.0,
|
||||
step=0.1,
|
||||
display_mode=IO.NumberDisplay.slider,
|
||||
tooltip="Creative strength of the upscale.",
|
||||
),
|
||||
IO.String.Input(
|
||||
"prompt",
|
||||
multiline=True,
|
||||
default="",
|
||||
tooltip="Optional descriptive (not instructive) scene prompt."
|
||||
f"Capping input at {AST2_MAX_FRAMES_WITH_PROMPT} frames (~15s @ 30fps) when set.",
|
||||
),
|
||||
IO.Float.Input(
|
||||
"sharp",
|
||||
default=0.5,
|
||||
min=0.0,
|
||||
max=1.0,
|
||||
step=0.01,
|
||||
display_mode=IO.NumberDisplay.slider,
|
||||
tooltip="Pre-enhance sharpness: "
|
||||
"0.0=Gaussian blur, 0.5=passthrough (default), 1.0=USM sharpening.",
|
||||
advanced=True,
|
||||
),
|
||||
IO.Float.Input(
|
||||
"realism",
|
||||
default=0.0,
|
||||
min=0.0,
|
||||
max=1.0,
|
||||
step=0.01,
|
||||
display_mode=IO.NumberDisplay.slider,
|
||||
tooltip="Pulls output toward photographic realism."
|
||||
"Leave at 0 for the model default.",
|
||||
advanced=True,
|
||||
),
|
||||
],
|
||||
),
|
||||
IO.DynamicCombo.Option(
|
||||
"Starlight (Astra) Fast",
|
||||
[IO.Combo.Input("upscaler_resolution", options=["FullHD (1080p)", "4K (2160p)"]),],
|
||||
),
|
||||
IO.DynamicCombo.Option(
|
||||
"Starlight (Astra) Creative",
|
||||
[
|
||||
IO.Combo.Input("upscaler_resolution", options=["FullHD (1080p)", "4K (2160p)"]),
|
||||
IO.Combo.Input(
|
||||
"creativity",
|
||||
options=["low", "middle", "high"],
|
||||
default="low",
|
||||
tooltip="Creative strength of the upscale.",
|
||||
),
|
||||
],
|
||||
),
|
||||
IO.DynamicCombo.Option(
|
||||
"Starlight Precise 2.5",
|
||||
[IO.Combo.Input("upscaler_resolution", options=["FullHD (1080p)", "4K (2160p)"])],
|
||||
),
|
||||
IO.DynamicCombo.Option("Disabled", []),
|
||||
],
|
||||
),
|
||||
IO.DynamicCombo.Input(
|
||||
"interpolation_model",
|
||||
options=[
|
||||
IO.DynamicCombo.Option("Disabled", []),
|
||||
IO.DynamicCombo.Option(
|
||||
"apo-8",
|
||||
[
|
||||
IO.Int.Input(
|
||||
"interpolation_frame_rate",
|
||||
default=60,
|
||||
min=15,
|
||||
max=240,
|
||||
display_mode=IO.NumberDisplay.number,
|
||||
tooltip="Output frame rate.",
|
||||
),
|
||||
IO.Int.Input(
|
||||
"interpolation_slowmo",
|
||||
default=1,
|
||||
min=1,
|
||||
max=16,
|
||||
display_mode=IO.NumberDisplay.number,
|
||||
tooltip="Slow-motion factor applied to the input video. "
|
||||
"For example, 2 makes the output twice as slow and doubles the duration.",
|
||||
advanced=True,
|
||||
),
|
||||
IO.Boolean.Input(
|
||||
"interpolation_duplicate",
|
||||
default=False,
|
||||
tooltip="Analyze the input for duplicate frames and remove them.",
|
||||
advanced=True,
|
||||
),
|
||||
IO.Float.Input(
|
||||
"interpolation_duplicate_threshold",
|
||||
default=0.01,
|
||||
min=0.001,
|
||||
max=0.1,
|
||||
step=0.001,
|
||||
display_mode=IO.NumberDisplay.number,
|
||||
tooltip="Detection sensitivity for duplicate frames.",
|
||||
advanced=True,
|
||||
),
|
||||
],
|
||||
),
|
||||
],
|
||||
),
|
||||
IO.Combo.Input(
|
||||
"dynamic_compression_level",
|
||||
options=["Low", "Mid", "High"],
|
||||
default="Low",
|
||||
tooltip="CQP level.",
|
||||
optional=True,
|
||||
),
|
||||
],
|
||||
outputs=[
|
||||
IO.Video.Output(),
|
||||
],
|
||||
hidden=[
|
||||
IO.Hidden.auth_token_comfy_org,
|
||||
IO.Hidden.api_key_comfy_org,
|
||||
IO.Hidden.unique_id,
|
||||
],
|
||||
is_api_node=True,
|
||||
price_badge=IO.PriceBadge(
|
||||
depends_on=IO.PriceBadgeDepends(widgets=[
|
||||
"upscaler_model",
|
||||
"upscaler_model.upscaler_resolution",
|
||||
"interpolation_model",
|
||||
]),
|
||||
expr="""
|
||||
(
|
||||
$model := $lookup(widgets, "upscaler_model");
|
||||
$res := $lookup(widgets, "upscaler_model.upscaler_resolution");
|
||||
$interp := $lookup(widgets, "interpolation_model");
|
||||
$is4k := $contains($res, "4k");
|
||||
$hasInterp := $interp != "disabled";
|
||||
$rates := {
|
||||
"starlight (astra) fast": {"hd": 0.43, "uhd": 0.85},
|
||||
"starlight precise 2.5": {"hd": 0.70, "uhd": 1.54},
|
||||
"astra 2": {"hd": 1.72, "uhd": 2.85},
|
||||
"starlight (astra) creative": {"hd": 2.25, "uhd": 3.99}
|
||||
};
|
||||
$surcharge := $is4k ? 0.28 : 0.14;
|
||||
$entry := $lookup($rates, $model);
|
||||
$base := $is4k ? $entry.uhd : $entry.hd;
|
||||
$hi := $base + ($hasInterp ? $surcharge : 0);
|
||||
$model = "disabled"
|
||||
? {"type":"text","text":"Interpolation only"}
|
||||
: ($hasInterp
|
||||
? {"type":"text","text":"~" & $string($base) & "–" & $string($hi) & " credits/src frame"}
|
||||
: {"type":"text","text":"~" & $string($base) & " credits/src frame"})
|
||||
)
|
||||
""",
|
||||
),
|
||||
)
|
||||
|
||||
@classmethod
|
||||
async def execute(
|
||||
cls,
|
||||
video: Input.Video,
|
||||
upscaler_model: dict,
|
||||
interpolation_model: dict,
|
||||
dynamic_compression_level: str = "Low",
|
||||
) -> IO.NodeOutput:
|
||||
upscaler_choice = upscaler_model["upscaler_model"]
|
||||
interpolation_choice = interpolation_model["interpolation_model"]
|
||||
if upscaler_choice == "Disabled" and interpolation_choice == "Disabled":
|
||||
raise ValueError("There is nothing to do: both upscaling and interpolation are disabled.")
|
||||
validate_container_format_is_mp4(video)
|
||||
src_width, src_height = video.get_dimensions()
|
||||
src_frame_rate = int(video.get_frame_rate())
|
||||
duration_sec = video.get_duration()
|
||||
src_video_stream = video.get_stream_source()
|
||||
target_width = src_width
|
||||
target_height = src_height
|
||||
target_frame_rate = src_frame_rate
|
||||
filters = []
|
||||
if upscaler_choice != "Disabled":
|
||||
if "1080p" in upscaler_model["upscaler_resolution"]:
|
||||
target_pixel_p = 1080
|
||||
max_long_side = 1920
|
||||
else:
|
||||
target_pixel_p = 2160
|
||||
max_long_side = 3840
|
||||
ar = src_width / src_height
|
||||
if src_width >= src_height:
|
||||
# Landscape or Square; Attempt to set height to target (e.g., 2160), calculate width
|
||||
target_height = target_pixel_p
|
||||
target_width = int(target_height * ar)
|
||||
# Check if width exceeds standard bounds (for ultra-wide e.g., 21:9 ARs)
|
||||
if target_width > max_long_side:
|
||||
target_width = max_long_side
|
||||
target_height = int(target_width / ar)
|
||||
else:
|
||||
# Portrait; Attempt to set width to target (e.g., 2160), calculate height
|
||||
target_width = target_pixel_p
|
||||
target_height = int(target_width / ar)
|
||||
# Check if height exceeds standard bounds
|
||||
if target_height > max_long_side:
|
||||
target_height = max_long_side
|
||||
target_width = int(target_height * ar)
|
||||
if target_width % 2 != 0:
|
||||
target_width += 1
|
||||
if target_height % 2 != 0:
|
||||
target_height += 1
|
||||
model_id = UPSCALER_MODELS_MAP[upscaler_choice]
|
||||
if model_id == "slc-1":
|
||||
filters.append(
|
||||
VideoEnhancementFilter(
|
||||
model=model_id,
|
||||
creativity=upscaler_model["creativity"],
|
||||
isOptimizedMode=True,
|
||||
)
|
||||
)
|
||||
elif model_id == "ast-2":
|
||||
n_frames = video.get_frame_count()
|
||||
ast2_prompt = (upscaler_model["prompt"] or "").strip()
|
||||
if ast2_prompt and n_frames > AST2_MAX_FRAMES_WITH_PROMPT:
|
||||
raise ValueError(
|
||||
f"Astra 2 with a prompt is limited to {AST2_MAX_FRAMES_WITH_PROMPT} input frames "
|
||||
f"(~15s @ 30fps); video has {n_frames}. Clear the prompt or shorten the clip."
|
||||
)
|
||||
if n_frames > AST2_MAX_FRAMES:
|
||||
raise ValueError(f"Astra 2 is limited to {AST2_MAX_FRAMES} input frames; video has {n_frames}.")
|
||||
realism = upscaler_model["realism"]
|
||||
filters.append(
|
||||
VideoEnhancementFilter(
|
||||
model=model_id,
|
||||
creativity=upscaler_model["creativity"],
|
||||
prompt=(ast2_prompt or None),
|
||||
sharp=upscaler_model["sharp"],
|
||||
realism=(realism if realism > 0 else None),
|
||||
)
|
||||
)
|
||||
else:
|
||||
filters.append(VideoEnhancementFilter(model=model_id))
|
||||
if interpolation_choice != "Disabled":
|
||||
target_frame_rate = interpolation_model["interpolation_frame_rate"]
|
||||
filters.append(
|
||||
VideoFrameInterpolationFilter(
|
||||
model=interpolation_choice,
|
||||
slowmo=interpolation_model["interpolation_slowmo"],
|
||||
fps=interpolation_model["interpolation_frame_rate"],
|
||||
duplicate=interpolation_model["interpolation_duplicate"],
|
||||
duplicate_threshold=interpolation_model["interpolation_duplicate_threshold"],
|
||||
),
|
||||
)
|
||||
initial_res = await sync_op(
|
||||
cls,
|
||||
ApiEndpoint(path="/proxy/topaz/video/", method="POST"),
|
||||
response_model=CreateVideoResponse,
|
||||
data=CreateVideoRequest(
|
||||
source=CreateVideoRequestSource(
|
||||
container="mp4",
|
||||
size=get_fs_object_size(src_video_stream),
|
||||
duration=int(duration_sec),
|
||||
frameCount=video.get_frame_count(),
|
||||
frameRate=src_frame_rate,
|
||||
resolution=Resolution(width=src_width, height=src_height),
|
||||
),
|
||||
filters=filters,
|
||||
output=OutputInformationVideo(
|
||||
resolution=Resolution(width=target_width, height=target_height),
|
||||
frameRate=target_frame_rate,
|
||||
audioCodec="AAC",
|
||||
audioTransfer="Copy",
|
||||
dynamicCompressionLevel=dynamic_compression_level,
|
||||
),
|
||||
),
|
||||
wait_label="Creating task",
|
||||
final_label_on_success="Task created",
|
||||
)
|
||||
upload_res = await sync_op(
|
||||
cls,
|
||||
ApiEndpoint(
|
||||
path=f"/proxy/topaz/video/{initial_res.requestId}/accept",
|
||||
method="PATCH",
|
||||
),
|
||||
response_model=VideoAcceptResponse,
|
||||
wait_label="Preparing upload",
|
||||
final_label_on_success="Upload started",
|
||||
)
|
||||
if len(upload_res.urls) > 1:
|
||||
raise NotImplementedError(
|
||||
"Large files are not currently supported. Please open an issue in the ComfyUI repository."
|
||||
)
|
||||
async with aiohttp.ClientSession(headers={"Content-Type": "video/mp4"}) as session:
|
||||
if isinstance(src_video_stream, BytesIO):
|
||||
src_video_stream.seek(0)
|
||||
async with session.put(upload_res.urls[0], data=src_video_stream, raise_for_status=True) as res:
|
||||
upload_etag = res.headers["Etag"]
|
||||
else:
|
||||
with builtins.open(src_video_stream, "rb") as video_file:
|
||||
async with session.put(upload_res.urls[0], data=video_file, raise_for_status=True) as res:
|
||||
upload_etag = res.headers["Etag"]
|
||||
await sync_op(
|
||||
cls,
|
||||
ApiEndpoint(
|
||||
path=f"/proxy/topaz/video/{initial_res.requestId}/complete-upload",
|
||||
method="PATCH",
|
||||
),
|
||||
response_model=VideoCompleteUploadResponse,
|
||||
data=VideoCompleteUploadRequest(
|
||||
uploadResults=[
|
||||
VideoCompleteUploadRequestPart(
|
||||
partNum=1,
|
||||
eTag=upload_etag,
|
||||
),
|
||||
],
|
||||
),
|
||||
wait_label="Finalizing upload",
|
||||
final_label_on_success="Upload completed",
|
||||
)
|
||||
final_response = await poll_op(
|
||||
cls,
|
||||
ApiEndpoint(path=f"/proxy/topaz/video/{initial_res.requestId}/status"),
|
||||
response_model=VideoStatusResponse,
|
||||
status_extractor=lambda x: x.status,
|
||||
progress_extractor=lambda x: getattr(x, "progress", 0),
|
||||
price_extractor=lambda x: (x.estimates.cost[0] * 0.08 if x.estimates and x.estimates.cost[0] else None),
|
||||
poll_interval=10.0,
|
||||
)
|
||||
return IO.NodeOutput(await download_url_to_video_output(final_response.download.url))
|
||||
|
||||
@@ -464,6 +819,7 @@ class TopazExtension(ComfyExtension):
|
||||
return [
|
||||
TopazImageEnhance,
|
||||
TopazVideoEnhance,
|
||||
TopazVideoEnhanceV2,
|
||||
]
|
||||
|
||||
|
||||
|
||||
@@ -38,7 +38,7 @@ async def execute_task(
|
||||
cls: type[IO.ComfyNode],
|
||||
vidu_endpoint: str,
|
||||
payload: TaskCreationRequest | TaskExtendCreationRequest | TaskMultiFrameCreationRequest,
|
||||
max_poll_attempts: int = 320,
|
||||
max_poll_attempts: int = 480,
|
||||
) -> list[TaskResult]:
|
||||
task_creation_response = await sync_op(
|
||||
cls,
|
||||
@@ -1097,7 +1097,6 @@ class ViduExtendVideoNode(IO.ComfyNode):
|
||||
video_url=await upload_video_to_comfyapi(cls, video, wait_label="Uploading video"),
|
||||
images=[image_url] if image_url else None,
|
||||
),
|
||||
max_poll_attempts=480,
|
||||
)
|
||||
return IO.NodeOutput(await download_url_to_video_output(results[0].url))
|
||||
|
||||
|
||||
@@ -818,7 +818,6 @@ class WanReferenceVideoApi(IO.ComfyNode):
|
||||
response_model=VideoTaskStatusResponse,
|
||||
status_extractor=lambda x: x.output.task_status,
|
||||
poll_interval=6,
|
||||
max_poll_attempts=280,
|
||||
)
|
||||
return IO.NodeOutput(await download_url_to_video_output(response.output.video_url))
|
||||
|
||||
|
||||
@@ -84,7 +84,6 @@ class WavespeedFlashVSRNode(IO.ComfyNode):
|
||||
response_model=TaskResultResponse,
|
||||
status_extractor=lambda x: "failed" if x.data is None else x.data.status,
|
||||
poll_interval=10.0,
|
||||
max_poll_attempts=480,
|
||||
)
|
||||
if final_response.code != 200:
|
||||
raise ValueError(
|
||||
@@ -156,7 +155,6 @@ class WavespeedImageUpscaleNode(IO.ComfyNode):
|
||||
response_model=TaskResultResponse,
|
||||
status_extractor=lambda x: "failed" if x.data is None else x.data.status,
|
||||
poll_interval=10.0,
|
||||
max_poll_attempts=480,
|
||||
)
|
||||
if final_response.code != 200:
|
||||
raise ValueError(
|
||||
|
||||
@@ -148,7 +148,7 @@ async def poll_op(
|
||||
queued_statuses: list[str | int] | None = None,
|
||||
data: BaseModel | None = None,
|
||||
poll_interval: float = 5.0,
|
||||
max_poll_attempts: int = 160,
|
||||
max_poll_attempts: int = 480,
|
||||
timeout_per_poll: float = 120.0,
|
||||
max_retries_per_poll: int = 10,
|
||||
retry_delay_per_poll: float = 1.0,
|
||||
@@ -254,7 +254,7 @@ async def poll_op_raw(
|
||||
queued_statuses: list[str | int] | None = None,
|
||||
data: dict[str, Any] | BaseModel | None = None,
|
||||
poll_interval: float = 5.0,
|
||||
max_poll_attempts: int = 160,
|
||||
max_poll_attempts: int = 480,
|
||||
timeout_per_poll: float = 120.0,
|
||||
max_retries_per_poll: int = 10,
|
||||
retry_delay_per_poll: float = 1.0,
|
||||
|
||||
@@ -1,16 +0,0 @@
|
||||
"""comfy_api_sealed_worker — torch-free type definitions for sealed worker children.
|
||||
|
||||
Drop-in replacement for comfy_api.latest._util type imports in sealed workers
|
||||
that do not have torch installed. Contains only data type definitions (TrimeshData,
|
||||
etc.) with numpy-only dependencies.
|
||||
|
||||
Usage in serializers:
|
||||
if _IMPORT_TORCH:
|
||||
from comfy_api.latest._util.trimesh_types import TrimeshData
|
||||
else:
|
||||
from comfy_api_sealed_worker.trimesh_types import TrimeshData
|
||||
"""
|
||||
|
||||
from .trimesh_types import TrimeshData
|
||||
|
||||
__all__ = ["TrimeshData"]
|
||||
@@ -1,259 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import numpy as np
|
||||
|
||||
|
||||
class TrimeshData:
|
||||
"""Triangular mesh payload for cross-process transfer.
|
||||
|
||||
Lightweight carrier for mesh geometry that does not depend on the
|
||||
``trimesh`` library. Serializers create this on the host side;
|
||||
isolated child processes convert to/from ``trimesh.Trimesh`` as needed.
|
||||
|
||||
Supports both ColorVisuals (vertex_colors) and TextureVisuals
|
||||
(uv + material with textures).
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
vertices: np.ndarray,
|
||||
faces: np.ndarray,
|
||||
vertex_normals: np.ndarray | None = None,
|
||||
face_normals: np.ndarray | None = None,
|
||||
vertex_colors: np.ndarray | None = None,
|
||||
uv: np.ndarray | None = None,
|
||||
material: dict | None = None,
|
||||
vertex_attributes: dict | None = None,
|
||||
face_attributes: dict | None = None,
|
||||
metadata: dict | None = None,
|
||||
) -> None:
|
||||
self.vertices = np.ascontiguousarray(vertices, dtype=np.float64)
|
||||
self.faces = np.ascontiguousarray(faces, dtype=np.int64)
|
||||
self.vertex_normals = (
|
||||
np.ascontiguousarray(vertex_normals, dtype=np.float64)
|
||||
if vertex_normals is not None
|
||||
else None
|
||||
)
|
||||
self.face_normals = (
|
||||
np.ascontiguousarray(face_normals, dtype=np.float64)
|
||||
if face_normals is not None
|
||||
else None
|
||||
)
|
||||
self.vertex_colors = (
|
||||
np.ascontiguousarray(vertex_colors, dtype=np.uint8)
|
||||
if vertex_colors is not None
|
||||
else None
|
||||
)
|
||||
self.uv = (
|
||||
np.ascontiguousarray(uv, dtype=np.float64)
|
||||
if uv is not None
|
||||
else None
|
||||
)
|
||||
self.material = material
|
||||
self.vertex_attributes = vertex_attributes or {}
|
||||
self.face_attributes = face_attributes or {}
|
||||
self.metadata = self._detensorize_dict(metadata) if metadata else {}
|
||||
|
||||
@staticmethod
|
||||
def _detensorize_dict(d):
|
||||
"""Recursively convert any tensors in a dict back to numpy arrays."""
|
||||
if not isinstance(d, dict):
|
||||
return d
|
||||
result = {}
|
||||
for k, v in d.items():
|
||||
if hasattr(v, "numpy"):
|
||||
result[k] = v.cpu().numpy() if hasattr(v, "cpu") else v.numpy()
|
||||
elif isinstance(v, dict):
|
||||
result[k] = TrimeshData._detensorize_dict(v)
|
||||
elif isinstance(v, list):
|
||||
result[k] = [
|
||||
item.cpu().numpy() if hasattr(item, "numpy") and hasattr(item, "cpu")
|
||||
else item.numpy() if hasattr(item, "numpy")
|
||||
else item
|
||||
for item in v
|
||||
]
|
||||
else:
|
||||
result[k] = v
|
||||
return result
|
||||
|
||||
@staticmethod
|
||||
def _to_numpy(arr, dtype):
|
||||
if arr is None:
|
||||
return None
|
||||
if hasattr(arr, "numpy"):
|
||||
arr = arr.cpu().numpy() if hasattr(arr, "cpu") else arr.numpy()
|
||||
return np.ascontiguousarray(arr, dtype=dtype)
|
||||
|
||||
@property
|
||||
def num_vertices(self) -> int:
|
||||
return self.vertices.shape[0]
|
||||
|
||||
@property
|
||||
def num_faces(self) -> int:
|
||||
return self.faces.shape[0]
|
||||
|
||||
@property
|
||||
def has_texture(self) -> bool:
|
||||
return self.uv is not None and self.material is not None
|
||||
|
||||
def to_trimesh(self):
|
||||
"""Convert to trimesh.Trimesh (requires trimesh in the environment)."""
|
||||
import trimesh
|
||||
from trimesh.visual import TextureVisuals
|
||||
|
||||
kwargs = {}
|
||||
if self.vertex_normals is not None:
|
||||
kwargs["vertex_normals"] = self.vertex_normals
|
||||
if self.face_normals is not None:
|
||||
kwargs["face_normals"] = self.face_normals
|
||||
if self.metadata:
|
||||
kwargs["metadata"] = self.metadata
|
||||
|
||||
mesh = trimesh.Trimesh(
|
||||
vertices=self.vertices, faces=self.faces, process=False, **kwargs
|
||||
)
|
||||
|
||||
# Reconstruct visual
|
||||
if self.has_texture:
|
||||
material = self._dict_to_material(self.material)
|
||||
mesh.visual = TextureVisuals(uv=self.uv, material=material)
|
||||
elif self.vertex_colors is not None:
|
||||
mesh.visual.vertex_colors = self.vertex_colors
|
||||
|
||||
for k, v in self.vertex_attributes.items():
|
||||
mesh.vertex_attributes[k] = v
|
||||
|
||||
for k, v in self.face_attributes.items():
|
||||
mesh.face_attributes[k] = v
|
||||
|
||||
return mesh
|
||||
|
||||
@staticmethod
|
||||
def _material_to_dict(material) -> dict:
|
||||
"""Serialize a trimesh material to a plain dict."""
|
||||
import base64
|
||||
from io import BytesIO
|
||||
from trimesh.visual.material import PBRMaterial, SimpleMaterial
|
||||
|
||||
result = {"type": type(material).__name__, "name": getattr(material, "name", None)}
|
||||
|
||||
if isinstance(material, PBRMaterial):
|
||||
result["baseColorFactor"] = material.baseColorFactor
|
||||
result["metallicFactor"] = material.metallicFactor
|
||||
result["roughnessFactor"] = material.roughnessFactor
|
||||
result["emissiveFactor"] = material.emissiveFactor
|
||||
result["alphaMode"] = material.alphaMode
|
||||
result["alphaCutoff"] = material.alphaCutoff
|
||||
result["doubleSided"] = material.doubleSided
|
||||
|
||||
for tex_name in ("baseColorTexture", "normalTexture", "emissiveTexture",
|
||||
"metallicRoughnessTexture", "occlusionTexture"):
|
||||
tex = getattr(material, tex_name, None)
|
||||
if tex is not None:
|
||||
buf = BytesIO()
|
||||
tex.save(buf, format="PNG")
|
||||
result[tex_name] = base64.b64encode(buf.getvalue()).decode("ascii")
|
||||
|
||||
elif isinstance(material, SimpleMaterial):
|
||||
result["main_color"] = list(material.main_color) if material.main_color is not None else None
|
||||
result["glossiness"] = material.glossiness
|
||||
if hasattr(material, "image") and material.image is not None:
|
||||
buf = BytesIO()
|
||||
material.image.save(buf, format="PNG")
|
||||
result["image"] = base64.b64encode(buf.getvalue()).decode("ascii")
|
||||
|
||||
return result
|
||||
|
||||
@staticmethod
|
||||
def _dict_to_material(d: dict):
|
||||
"""Reconstruct a trimesh material from a plain dict."""
|
||||
import base64
|
||||
from io import BytesIO
|
||||
from PIL import Image
|
||||
from trimesh.visual.material import PBRMaterial, SimpleMaterial
|
||||
|
||||
mat_type = d.get("type", "PBRMaterial")
|
||||
|
||||
if mat_type == "PBRMaterial":
|
||||
kwargs = {
|
||||
"name": d.get("name"),
|
||||
"baseColorFactor": d.get("baseColorFactor"),
|
||||
"metallicFactor": d.get("metallicFactor"),
|
||||
"roughnessFactor": d.get("roughnessFactor"),
|
||||
"emissiveFactor": d.get("emissiveFactor"),
|
||||
"alphaMode": d.get("alphaMode"),
|
||||
"alphaCutoff": d.get("alphaCutoff"),
|
||||
"doubleSided": d.get("doubleSided"),
|
||||
}
|
||||
for tex_name in ("baseColorTexture", "normalTexture", "emissiveTexture",
|
||||
"metallicRoughnessTexture", "occlusionTexture"):
|
||||
if tex_name in d and d[tex_name] is not None:
|
||||
img = Image.open(BytesIO(base64.b64decode(d[tex_name])))
|
||||
kwargs[tex_name] = img
|
||||
return PBRMaterial(**{k: v for k, v in kwargs.items() if v is not None})
|
||||
|
||||
elif mat_type == "SimpleMaterial":
|
||||
kwargs = {
|
||||
"name": d.get("name"),
|
||||
"glossiness": d.get("glossiness"),
|
||||
}
|
||||
if d.get("main_color") is not None:
|
||||
kwargs["diffuse"] = d["main_color"]
|
||||
if d.get("image") is not None:
|
||||
kwargs["image"] = Image.open(BytesIO(base64.b64decode(d["image"])))
|
||||
return SimpleMaterial(**kwargs)
|
||||
|
||||
raise ValueError(f"Unknown material type: {mat_type}")
|
||||
|
||||
@classmethod
|
||||
def from_trimesh(cls, mesh) -> TrimeshData:
|
||||
"""Create from a trimesh.Trimesh object."""
|
||||
from trimesh.visual.texture import TextureVisuals
|
||||
|
||||
vertex_normals = None
|
||||
if mesh._cache.cache.get("vertex_normals") is not None:
|
||||
vertex_normals = np.asarray(mesh.vertex_normals)
|
||||
|
||||
face_normals = None
|
||||
if mesh._cache.cache.get("face_normals") is not None:
|
||||
face_normals = np.asarray(mesh.face_normals)
|
||||
|
||||
vertex_colors = None
|
||||
uv = None
|
||||
material = None
|
||||
|
||||
if isinstance(mesh.visual, TextureVisuals):
|
||||
if mesh.visual.uv is not None:
|
||||
uv = np.asarray(mesh.visual.uv, dtype=np.float64)
|
||||
if mesh.visual.material is not None:
|
||||
material = cls._material_to_dict(mesh.visual.material)
|
||||
else:
|
||||
try:
|
||||
vc = mesh.visual.vertex_colors
|
||||
if vc is not None and len(vc) > 0:
|
||||
vertex_colors = np.asarray(vc, dtype=np.uint8)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
va = {}
|
||||
if hasattr(mesh, "vertex_attributes") and mesh.vertex_attributes:
|
||||
for k, v in mesh.vertex_attributes.items():
|
||||
va[k] = np.asarray(v) if hasattr(v, "__array__") else v
|
||||
|
||||
fa = {}
|
||||
if hasattr(mesh, "face_attributes") and mesh.face_attributes:
|
||||
for k, v in mesh.face_attributes.items():
|
||||
fa[k] = np.asarray(v) if hasattr(v, "__array__") else v
|
||||
|
||||
return cls(
|
||||
vertices=np.asarray(mesh.vertices),
|
||||
faces=np.asarray(mesh.faces),
|
||||
vertex_normals=vertex_normals,
|
||||
face_normals=face_normals,
|
||||
vertex_colors=vertex_colors,
|
||||
uv=uv,
|
||||
material=material,
|
||||
vertex_attributes=va if va else None,
|
||||
face_attributes=fa if fa else None,
|
||||
metadata=mesh.metadata if mesh.metadata else None,
|
||||
)
|
||||
@@ -459,27 +459,23 @@ class SDPoseKeypointExtractor(io.ComfyNode):
|
||||
total_images = image.shape[0]
|
||||
captured_feat = None
|
||||
|
||||
model_h = int(head.heatmap_size[0]) * 4 # e.g. 192 * 4 = 768
|
||||
model_w = int(head.heatmap_size[1]) * 4 # e.g. 256 * 4 = 1024
|
||||
model_w = int(head.heatmap_size[0]) * 4 # 192 * 4 = 768
|
||||
model_h = int(head.heatmap_size[1]) * 4 # 256 * 4 = 1024
|
||||
|
||||
def _resize_to_model(imgs):
|
||||
"""Aspect-preserving resize + zero-pad BHWC images to (model_h, model_w). Returns (resized_bhwc, scale, pad_top, pad_left)."""
|
||||
"""Stretch BHWC images to (model_h, model_w), model expects no aspect preservation."""
|
||||
h, w = imgs.shape[-3], imgs.shape[-2]
|
||||
scale = min(model_h / h, model_w / w)
|
||||
sh, sw = int(round(h * scale)), int(round(w * scale))
|
||||
pt, pl = (model_h - sh) // 2, (model_w - sw) // 2
|
||||
method = "area" if (model_h <= h and model_w <= w) else "bilinear"
|
||||
chw = imgs.permute(0, 3, 1, 2).float()
|
||||
scaled = comfy.utils.common_upscale(chw, sw, sh, upscale_method="bilinear", crop="disabled")
|
||||
padded = torch.zeros(scaled.shape[0], scaled.shape[1], model_h, model_w, dtype=scaled.dtype, device=scaled.device)
|
||||
padded[:, :, pt:pt + sh, pl:pl + sw] = scaled
|
||||
return padded.permute(0, 2, 3, 1), scale, pt, pl
|
||||
scaled = comfy.utils.common_upscale(chw, model_w, model_h, upscale_method=method, crop="disabled")
|
||||
return scaled.permute(0, 2, 3, 1), model_w / w, model_h / h
|
||||
|
||||
def _remap_keypoints(kp, scale, pad_top, pad_left, offset_x=0, offset_y=0):
|
||||
def _remap_keypoints(kp, scale_x, scale_y, offset_x=0, offset_y=0):
|
||||
"""Remap keypoints from model space back to original image space."""
|
||||
kp = kp.copy() if isinstance(kp, np.ndarray) else np.array(kp, dtype=np.float32)
|
||||
invalid = kp[..., 0] < 0
|
||||
kp[..., 0] = (kp[..., 0] - pad_left) / scale + offset_x
|
||||
kp[..., 1] = (kp[..., 1] - pad_top) / scale + offset_y
|
||||
kp[..., 0] = kp[..., 0] / scale_x + offset_x
|
||||
kp[..., 1] = kp[..., 1] / scale_y + offset_y
|
||||
kp[invalid] = -1
|
||||
return kp
|
||||
|
||||
@@ -529,18 +525,18 @@ class SDPoseKeypointExtractor(io.ComfyNode):
|
||||
continue
|
||||
|
||||
crop = img[:, y1:y2, x1:x2, :] # (1, crop_h, crop_w, C)
|
||||
crop_resized, scale, pad_top, pad_left = _resize_to_model(crop)
|
||||
crop_resized, sx, sy = _resize_to_model(crop)
|
||||
|
||||
latent_crop = vae.encode(crop_resized)
|
||||
kp_batch, sc_batch = _run_on_latent(latent_crop)
|
||||
kp = _remap_keypoints(kp_batch[0], scale, pad_top, pad_left, x1, y1)
|
||||
kp = _remap_keypoints(kp_batch[0], sx, sy, x1, y1)
|
||||
img_keypoints.append(kp)
|
||||
img_scores.append(sc_batch[0])
|
||||
else:
|
||||
img_resized, scale, pad_top, pad_left = _resize_to_model(img)
|
||||
img_resized, sx, sy = _resize_to_model(img)
|
||||
latent_img = vae.encode(img_resized)
|
||||
kp_batch, sc_batch = _run_on_latent(latent_img)
|
||||
img_keypoints.append(_remap_keypoints(kp_batch[0], scale, pad_top, pad_left))
|
||||
img_keypoints.append(_remap_keypoints(kp_batch[0], sx, sy))
|
||||
img_scores.append(sc_batch[0])
|
||||
|
||||
all_keypoints.append(img_keypoints)
|
||||
@@ -549,12 +545,12 @@ class SDPoseKeypointExtractor(io.ComfyNode):
|
||||
|
||||
else: # full-image mode, batched
|
||||
for batch_start in tqdm(range(0, total_images, batch_size), desc="Extracting keypoints"):
|
||||
batch_resized, scale, pad_top, pad_left = _resize_to_model(image[batch_start:batch_start + batch_size])
|
||||
batch_resized, sx, sy = _resize_to_model(image[batch_start:batch_start + batch_size])
|
||||
latent_batch = vae.encode(batch_resized)
|
||||
kp_batch, sc_batch = _run_on_latent(latent_batch)
|
||||
|
||||
for kp, sc in zip(kp_batch, sc_batch):
|
||||
all_keypoints.append([_remap_keypoints(kp, scale, pad_top, pad_left)])
|
||||
all_keypoints.append([_remap_keypoints(kp, sx, sy)])
|
||||
all_scores.append([sc])
|
||||
|
||||
pbar.update(len(kp_batch))
|
||||
@@ -727,13 +723,13 @@ class CropByBBoxes(io.ComfyNode):
|
||||
scale = min(output_width / crop_w, output_height / crop_h)
|
||||
scaled_w = int(round(crop_w * scale))
|
||||
scaled_h = int(round(crop_h * scale))
|
||||
scaled = comfy.utils.common_upscale(crop_chw, scaled_w, scaled_h, upscale_method="bilinear", crop="disabled")
|
||||
scaled = comfy.utils.common_upscale(crop_chw, scaled_w, scaled_h, upscale_method="area", crop="disabled")
|
||||
pad_left = (output_width - scaled_w) // 2
|
||||
pad_top = (output_height - scaled_h) // 2
|
||||
resized = torch.zeros(1, num_ch, output_height, output_width, dtype=image.dtype, device=image.device)
|
||||
resized[:, :, pad_top:pad_top + scaled_h, pad_left:pad_left + scaled_w] = scaled
|
||||
else: # "stretch"
|
||||
resized = comfy.utils.common_upscale(crop_chw, output_width, output_height, upscale_method="bilinear", crop="disabled")
|
||||
resized = comfy.utils.common_upscale(crop_chw, output_width, output_height, upscale_method="area", crop="disabled")
|
||||
crops.append(resized)
|
||||
|
||||
if not crops:
|
||||
|
||||
@@ -32,6 +32,8 @@ class TextGenerate(io.ComfyNode):
|
||||
io.Clip.Input("clip"),
|
||||
io.String.Input("prompt", multiline=True, dynamic_prompts=True, default=""),
|
||||
io.Image.Input("image", optional=True),
|
||||
io.Image.Input("video", optional=True, tooltip="Video frames as image batch. Assumed to be 24 FPS; subsampled to 1 FPS internally."),
|
||||
io.Audio.Input("audio", optional=True),
|
||||
io.Int.Input("max_length", default=256, min=1, max=2048),
|
||||
io.DynamicCombo.Input("sampling_mode", options=sampling_options, display_name="Sampling Mode"),
|
||||
io.Boolean.Input("thinking", optional=True, default=False, tooltip="Operate in thinking mode if the model supports it."),
|
||||
@@ -43,9 +45,9 @@ class TextGenerate(io.ComfyNode):
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, clip, prompt, max_length, sampling_mode, image=None, thinking=False, use_default_template=True) -> io.NodeOutput:
|
||||
def execute(cls, clip, prompt, max_length, sampling_mode, image=None, thinking=False, use_default_template=True, video=None, audio=None) -> io.NodeOutput:
|
||||
|
||||
tokens = clip.tokenize(prompt, image=image, skip_template=not use_default_template, min_length=1, thinking=thinking)
|
||||
tokens = clip.tokenize(prompt, image=image, skip_template=not use_default_template, min_length=1, thinking=thinking, video=video, audio=audio)
|
||||
|
||||
# Get sampling parameters from dynamic combo
|
||||
do_sample = sampling_mode.get("sampling_mode") == "on"
|
||||
@@ -70,7 +72,8 @@ class TextGenerate(io.ComfyNode):
|
||||
seed=seed
|
||||
)
|
||||
|
||||
generated_text = clip.decode(generated_ids, skip_special_tokens=True)
|
||||
generated_text = clip.decode(generated_ids)
|
||||
|
||||
return io.NodeOutput(generated_text)
|
||||
|
||||
|
||||
@@ -161,12 +164,12 @@ class TextGenerateLTX2Prompt(TextGenerate):
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, clip, prompt, max_length, sampling_mode, image=None, thinking=False, use_default_template=True) -> io.NodeOutput:
|
||||
def execute(cls, clip, prompt, max_length, sampling_mode, image=None, thinking=False, use_default_template=True, video=None, audio=None) -> io.NodeOutput:
|
||||
if image is None:
|
||||
formatted_prompt = f"<start_of_turn>system\n{LTX2_T2V_SYSTEM_PROMPT.strip()}<end_of_turn>\n<start_of_turn>user\nUser Raw Input Prompt: {prompt}.<end_of_turn>\n<start_of_turn>model\n"
|
||||
else:
|
||||
formatted_prompt = f"<start_of_turn>system\n{LTX2_I2V_SYSTEM_PROMPT.strip()}<end_of_turn>\n<start_of_turn>user\n\n<image_soft_token>\n\nUser Raw Input Prompt: {prompt}.<end_of_turn>\n<start_of_turn>model\n"
|
||||
return super().execute(clip, formatted_prompt, max_length, sampling_mode, image, thinking, use_default_template)
|
||||
return super().execute(clip, formatted_prompt, max_length, sampling_mode, image=image, thinking=thinking, use_default_template=use_default_template, video=video, audio=audio)
|
||||
|
||||
|
||||
class TextgenExtension(ComfyExtension):
|
||||
|
||||
+1
-1
@@ -92,7 +92,7 @@ if args.cuda_malloc:
|
||||
env_var = os.environ.get('PYTORCH_CUDA_ALLOC_CONF', None)
|
||||
if env_var is None:
|
||||
env_var = "backend:cudaMallocAsync"
|
||||
elif not args.use_process_isolation:
|
||||
else:
|
||||
env_var += ",backend:cudaMallocAsync"
|
||||
|
||||
os.environ['PYTORCH_CUDA_ALLOC_CONF'] = env_var
|
||||
|
||||
+22
-145
@@ -1,9 +1,7 @@
|
||||
import copy
|
||||
import gc
|
||||
import heapq
|
||||
import inspect
|
||||
import logging
|
||||
import os
|
||||
import sys
|
||||
import threading
|
||||
import time
|
||||
@@ -17,6 +15,7 @@ import torch
|
||||
from comfy.cli_args import args
|
||||
import comfy.memory_management
|
||||
import comfy.model_management
|
||||
import comfy.model_prefetch
|
||||
import comfy_aimdo.model_vbar
|
||||
|
||||
from latent_preview import set_preview_method
|
||||
@@ -44,8 +43,6 @@ from comfy_api.internal import _ComfyNodeInternal, _NodeOutputInternal, first_re
|
||||
from comfy_api.latest import io, _io
|
||||
from comfy_execution.cache_provider import _has_cache_providers, _get_cache_providers, _logger as _cache_logger
|
||||
|
||||
_AIMDO_VBAR_RESET_UNAVAILABLE_LOGGED = False
|
||||
|
||||
|
||||
class ExecutionResult(Enum):
|
||||
SUCCESS = 0
|
||||
@@ -266,31 +263,20 @@ async def _async_map_node_over_list(prompt_id, unique_id, obj, input_data_all, f
|
||||
pre_execute_cb(index)
|
||||
# V3
|
||||
if isinstance(obj, _ComfyNodeInternal) or (is_class(obj) and issubclass(obj, _ComfyNodeInternal)):
|
||||
# Check for isolated node - skip validation and class cloning
|
||||
if hasattr(obj, "_pyisolate_extension"):
|
||||
# Isolated Node: The stub is just a proxy; real validation happens in child process
|
||||
if v3_data is not None:
|
||||
inputs = _io.build_nested_inputs(inputs, v3_data)
|
||||
# Inject hidden inputs so they're available in the isolated child process
|
||||
inputs.update(v3_data.get("hidden_inputs", {}))
|
||||
f = getattr(obj, func)
|
||||
# Standard V3 Node (Existing Logic)
|
||||
|
||||
# if is just a class, then assign no state, just create clone
|
||||
if is_class(obj):
|
||||
type_obj = obj
|
||||
obj.VALIDATE_CLASS()
|
||||
class_clone = obj.PREPARE_CLASS_CLONE(v3_data)
|
||||
# otherwise, use class instance to populate/reuse some fields
|
||||
else:
|
||||
# if is just a class, then assign no resources or state, just create clone
|
||||
if is_class(obj):
|
||||
type_obj = obj
|
||||
obj.VALIDATE_CLASS()
|
||||
class_clone = obj.PREPARE_CLASS_CLONE(v3_data)
|
||||
# otherwise, use class instance to populate/reuse some fields
|
||||
else:
|
||||
type_obj = type(obj)
|
||||
type_obj.VALIDATE_CLASS()
|
||||
class_clone = type_obj.PREPARE_CLASS_CLONE(v3_data)
|
||||
f = make_locked_method_func(type_obj, func, class_clone)
|
||||
# in case of dynamic inputs, restructure inputs to expected nested dict
|
||||
if v3_data is not None:
|
||||
inputs = _io.build_nested_inputs(inputs, v3_data)
|
||||
type_obj = type(obj)
|
||||
type_obj.VALIDATE_CLASS()
|
||||
class_clone = type_obj.PREPARE_CLASS_CLONE(v3_data)
|
||||
f = make_locked_method_func(type_obj, func, class_clone)
|
||||
# in case of dynamic inputs, restructure inputs to expected nested dict
|
||||
if v3_data is not None:
|
||||
inputs = _io.build_nested_inputs(inputs, v3_data)
|
||||
# V1
|
||||
else:
|
||||
f = getattr(obj, func)
|
||||
@@ -552,17 +538,8 @@ async def execute(server, dynprompt, caches, current_item, extra_data, executed,
|
||||
if args.verbose == "DEBUG":
|
||||
comfy_aimdo.control.analyze()
|
||||
comfy.model_management.reset_cast_buffers()
|
||||
vbar_lib = getattr(comfy_aimdo.model_vbar, "lib", None)
|
||||
if vbar_lib is not None:
|
||||
comfy_aimdo.model_vbar.vbars_reset_watermark_limits()
|
||||
else:
|
||||
global _AIMDO_VBAR_RESET_UNAVAILABLE_LOGGED
|
||||
if not _AIMDO_VBAR_RESET_UNAVAILABLE_LOGGED:
|
||||
logging.warning(
|
||||
"DynamicVRAM backend unavailable for watermark reset; "
|
||||
"skipping vbar reset for this process."
|
||||
)
|
||||
_AIMDO_VBAR_RESET_UNAVAILABLE_LOGGED = True
|
||||
comfy.model_prefetch.cleanup_prefetch_queues()
|
||||
comfy_aimdo.model_vbar.vbars_reset_watermark_limits()
|
||||
|
||||
if has_pending_tasks:
|
||||
pending_async_nodes[unique_id] = output_data
|
||||
@@ -571,29 +548,8 @@ async def execute(server, dynprompt, caches, current_item, extra_data, executed,
|
||||
tasks = [x for x in output_data if isinstance(x, asyncio.Task)]
|
||||
await asyncio.gather(*tasks, return_exceptions=True)
|
||||
unblock()
|
||||
|
||||
# Keep isolation node execution deterministic by default, but allow
|
||||
# opt-out for diagnostics.
|
||||
isolation_sequential = os.environ.get("COMFY_ISOLATE_SEQUENTIAL", "1").lower() in ("1", "true", "yes")
|
||||
if args.use_process_isolation and isolation_sequential:
|
||||
await await_completion()
|
||||
results = []
|
||||
for r in pending_async_nodes[unique_id]:
|
||||
if isinstance(r, asyncio.Task):
|
||||
try:
|
||||
results.append(r.result())
|
||||
except Exception as ex:
|
||||
del pending_async_nodes[unique_id]
|
||||
raise ex
|
||||
else:
|
||||
results.append(r)
|
||||
del pending_async_nodes[unique_id]
|
||||
output_data, output_ui, has_subgraph = get_output_from_returns(results, class_def)
|
||||
has_pending_tasks = False
|
||||
|
||||
else:
|
||||
asyncio.create_task(await_completion())
|
||||
return (ExecutionResult.PENDING, None, None)
|
||||
asyncio.create_task(await_completion())
|
||||
return (ExecutionResult.PENDING, None, None)
|
||||
if len(output_ui) > 0:
|
||||
ui_outputs[unique_id] = {
|
||||
"meta": {
|
||||
@@ -703,46 +659,6 @@ class PromptExecutor:
|
||||
self.status_messages = []
|
||||
self.success = True
|
||||
|
||||
async def _notify_execution_graph_safe(self, class_types: set[str], *, fail_loud: bool = False) -> None:
|
||||
if not args.use_process_isolation:
|
||||
return
|
||||
try:
|
||||
from comfy.isolation import notify_execution_graph
|
||||
await notify_execution_graph(class_types, caches=self.caches.all)
|
||||
except Exception:
|
||||
if fail_loud:
|
||||
raise
|
||||
logging.debug("][ EX:notify_execution_graph failed", exc_info=True)
|
||||
|
||||
async def _flush_running_extensions_transport_state_safe(self) -> None:
|
||||
if not args.use_process_isolation:
|
||||
return
|
||||
try:
|
||||
from comfy.isolation import flush_running_extensions_transport_state
|
||||
await flush_running_extensions_transport_state()
|
||||
except Exception:
|
||||
logging.debug("][ EX:flush_running_extensions_transport_state failed", exc_info=True)
|
||||
|
||||
async def _wait_model_patcher_quiescence_safe(
|
||||
self,
|
||||
*,
|
||||
fail_loud: bool = False,
|
||||
timeout_ms: int = 120000,
|
||||
marker: str = "EX:wait_model_patcher_idle",
|
||||
) -> None:
|
||||
if not args.use_process_isolation:
|
||||
return
|
||||
try:
|
||||
from comfy.isolation import wait_for_model_patcher_quiescence
|
||||
|
||||
await wait_for_model_patcher_quiescence(
|
||||
timeout_ms=timeout_ms, fail_loud=fail_loud, marker=marker
|
||||
)
|
||||
except Exception:
|
||||
if fail_loud:
|
||||
raise
|
||||
logging.debug("][ EX:wait_model_patcher_quiescence failed", exc_info=True)
|
||||
|
||||
def add_message(self, event, data: dict, broadcast: bool):
|
||||
data = {
|
||||
**data,
|
||||
@@ -797,18 +713,6 @@ class PromptExecutor:
|
||||
asyncio.run(self.execute_async(prompt, prompt_id, extra_data, execute_outputs))
|
||||
|
||||
async def execute_async(self, prompt, prompt_id, extra_data={}, execute_outputs=[]):
|
||||
if args.use_process_isolation:
|
||||
# Update RPC event loops for all isolated extensions.
|
||||
# This is critical for serial workflow execution - each asyncio.run() creates
|
||||
# a new event loop, and RPC instances must be updated to use it.
|
||||
try:
|
||||
from comfy.isolation import update_rpc_event_loops
|
||||
update_rpc_event_loops()
|
||||
except ImportError:
|
||||
pass # Isolation not available
|
||||
except Exception as e:
|
||||
logging.getLogger(__name__).warning(f"Failed to update RPC event loops: {e}")
|
||||
|
||||
set_preview_method(extra_data.get("preview_method"))
|
||||
|
||||
nodes.interrupt_processing(False)
|
||||
@@ -821,25 +725,6 @@ class PromptExecutor:
|
||||
self.status_messages = []
|
||||
self.add_message("execution_start", { "prompt_id": prompt_id}, broadcast=False)
|
||||
|
||||
if args.use_process_isolation:
|
||||
try:
|
||||
# Boundary cleanup runs at the start of the next workflow in
|
||||
# isolation mode, matching non-isolated "next prompt" timing.
|
||||
self.caches = CacheSet(cache_type=self.cache_type, cache_args=self.cache_args)
|
||||
await self._wait_model_patcher_quiescence_safe(
|
||||
fail_loud=False,
|
||||
timeout_ms=120000,
|
||||
marker="EX:boundary_cleanup_wait_idle",
|
||||
)
|
||||
await self._flush_running_extensions_transport_state_safe()
|
||||
comfy.model_management.unload_all_models()
|
||||
comfy.model_management.cleanup_models_gc()
|
||||
comfy.model_management.cleanup_models()
|
||||
gc.collect()
|
||||
comfy.model_management.soft_empty_cache()
|
||||
except Exception:
|
||||
logging.debug("][ EX:isolation_boundary_cleanup_start failed", exc_info=True)
|
||||
|
||||
self._notify_prompt_lifecycle("start", prompt_id)
|
||||
ram_headroom = int(self.cache_args["ram"] * (1024 ** 3))
|
||||
ram_release_callback = self.caches.outputs.ram_release if self.cache_type == CacheType.RAM_PRESSURE else None
|
||||
@@ -877,18 +762,6 @@ class PromptExecutor:
|
||||
for node_id in list(execute_outputs):
|
||||
execution_list.add_node(node_id)
|
||||
|
||||
if args.use_process_isolation:
|
||||
pending_class_types = set()
|
||||
for node_id in execution_list.pendingNodes.keys():
|
||||
class_type = dynamic_prompt.get_node(node_id)["class_type"]
|
||||
pending_class_types.add(class_type)
|
||||
await self._wait_model_patcher_quiescence_safe(
|
||||
fail_loud=True,
|
||||
timeout_ms=120000,
|
||||
marker="EX:notify_graph_wait_idle",
|
||||
)
|
||||
await self._notify_execution_graph_safe(pending_class_types, fail_loud=True)
|
||||
|
||||
while not execution_list.is_empty():
|
||||
node_id, error, ex = await execution_list.stage_node_execution()
|
||||
if error is not None:
|
||||
@@ -1143,6 +1016,10 @@ async def validate_inputs(prompt_id, prompt, item, validated, visiting=None):
|
||||
|
||||
if isinstance(input_type, list) or input_type == io.Combo.io_type:
|
||||
if input_type == io.Combo.io_type:
|
||||
# Skip validation for combos with remote options — options
|
||||
# are fetched client-side and not available on the server.
|
||||
if extra_info.get("remote_combo"):
|
||||
continue
|
||||
combo_options = extra_info.get("options", [])
|
||||
else:
|
||||
combo_options = input_type
|
||||
|
||||
@@ -1,27 +1,7 @@
|
||||
import os
|
||||
import sys
|
||||
|
||||
IS_PYISOLATE_CHILD = os.environ.get("PYISOLATE_CHILD") == "1"
|
||||
|
||||
if __name__ == "__main__" and IS_PYISOLATE_CHILD:
|
||||
del os.environ["PYISOLATE_CHILD"]
|
||||
IS_PYISOLATE_CHILD = False
|
||||
|
||||
CURRENT_DIR = os.path.dirname(os.path.realpath(__file__))
|
||||
if CURRENT_DIR not in sys.path:
|
||||
sys.path.insert(0, CURRENT_DIR)
|
||||
|
||||
if not IS_PYISOLATE_CHILD:
|
||||
python_scripts_dir = os.path.dirname(os.path.realpath(sys.executable))
|
||||
path_entries = os.environ.get("PATH", "").split(os.pathsep)
|
||||
if python_scripts_dir and python_scripts_dir not in path_entries:
|
||||
os.environ["PATH"] = os.pathsep.join([python_scripts_dir, *path_entries])
|
||||
|
||||
IS_PRIMARY_PROCESS = (not IS_PYISOLATE_CHILD) and __name__ == "__main__"
|
||||
|
||||
import comfy.options
|
||||
comfy.options.enable_args_parsing()
|
||||
|
||||
import os
|
||||
import importlib.util
|
||||
import shutil
|
||||
import importlib.metadata
|
||||
@@ -29,56 +9,26 @@ import folder_paths
|
||||
import time
|
||||
from comfy.cli_args import args, enables_dynamic_vram
|
||||
from app.logger import setup_logger
|
||||
setup_logger(log_level=args.verbose, use_stdout=args.log_stdout)
|
||||
|
||||
from app.assets.seeder import asset_seeder
|
||||
from app.assets.services import register_output_files
|
||||
import itertools
|
||||
import utils.extra_config # noqa: F401
|
||||
import utils.extra_config
|
||||
from utils.mime_types import init_mime_types
|
||||
import faulthandler
|
||||
import logging
|
||||
import sys
|
||||
from comfy_execution.progress import get_progress_state
|
||||
from comfy_execution.utils import get_executing_context
|
||||
from comfy_api import feature_flags
|
||||
from app.database.db import init_db, dependencies_available
|
||||
|
||||
import comfy_aimdo.control
|
||||
|
||||
if enables_dynamic_vram():
|
||||
if not comfy_aimdo.control.init():
|
||||
logging.warning(
|
||||
"DynamicVRAM requested, but comfy-aimdo failed to initialize early. "
|
||||
"Will fall back to legacy model loading if device init fails."
|
||||
)
|
||||
|
||||
if args.use_process_isolation:
|
||||
from comfy.isolation import initialize_proxies
|
||||
initialize_proxies()
|
||||
|
||||
# Explicitly register the ComfyUI adapter for pyisolate (v1.0 architecture)
|
||||
try:
|
||||
import pyisolate
|
||||
from comfy.isolation.adapter import ComfyUIAdapter
|
||||
pyisolate.register_adapter(ComfyUIAdapter())
|
||||
logging.info("PyIsolate adapter registered: comfyui")
|
||||
except ImportError:
|
||||
logging.warning("PyIsolate not installed or version too old for explicit registration")
|
||||
except Exception as e:
|
||||
logging.error(f"Failed to register PyIsolate adapter: {e}")
|
||||
|
||||
if not IS_PYISOLATE_CHILD:
|
||||
if 'PYTORCH_CUDA_ALLOC_CONF' not in os.environ:
|
||||
os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'backend:native'
|
||||
|
||||
if not IS_PYISOLATE_CHILD:
|
||||
from comfy_execution.progress import get_progress_state
|
||||
from comfy_execution.utils import get_executing_context
|
||||
from comfy_api import feature_flags
|
||||
|
||||
if IS_PRIMARY_PROCESS:
|
||||
if __name__ == "__main__":
|
||||
#NOTE: These do not do anything on core ComfyUI, they are for custom nodes.
|
||||
os.environ['HF_HUB_DISABLE_TELEMETRY'] = '1'
|
||||
os.environ['DO_NOT_TRACK'] = '1'
|
||||
|
||||
if not IS_PYISOLATE_CHILD:
|
||||
setup_logger(log_level=args.verbose, use_stdout=args.log_stdout)
|
||||
|
||||
faulthandler.enable(file=sys.stderr, all_threads=False)
|
||||
|
||||
import comfy_aimdo.control
|
||||
@@ -143,15 +93,14 @@ if args.enable_manager:
|
||||
|
||||
|
||||
def apply_custom_paths():
|
||||
from utils import extra_config # Deferred import - spawn re-runs main.py
|
||||
# extra model paths
|
||||
extra_model_paths_config_path = os.path.join(os.path.dirname(os.path.realpath(__file__)), "extra_model_paths.yaml")
|
||||
if os.path.isfile(extra_model_paths_config_path):
|
||||
extra_config.load_extra_path_config(extra_model_paths_config_path)
|
||||
utils.extra_config.load_extra_path_config(extra_model_paths_config_path)
|
||||
|
||||
if args.extra_model_paths_config:
|
||||
for config_path in itertools.chain(*args.extra_model_paths_config):
|
||||
extra_config.load_extra_path_config(config_path)
|
||||
utils.extra_config.load_extra_path_config(config_path)
|
||||
|
||||
# --output-directory, --input-directory, --user-directory
|
||||
if args.output_directory:
|
||||
@@ -224,17 +173,15 @@ def execute_prestartup_script():
|
||||
else:
|
||||
import_message = " (PRESTARTUP FAILED)"
|
||||
logging.info("{:6.1f} seconds{}: {}".format(n[0], import_message, n[1]))
|
||||
logging.info("")
|
||||
logging.info("")
|
||||
|
||||
if not IS_PYISOLATE_CHILD:
|
||||
apply_custom_paths()
|
||||
init_mime_types()
|
||||
apply_custom_paths()
|
||||
init_mime_types()
|
||||
|
||||
if args.enable_manager and not IS_PYISOLATE_CHILD:
|
||||
if args.enable_manager:
|
||||
comfyui_manager.prestartup()
|
||||
|
||||
if not IS_PYISOLATE_CHILD:
|
||||
execute_prestartup_script()
|
||||
execute_prestartup_script()
|
||||
|
||||
|
||||
# Main code
|
||||
@@ -245,17 +192,17 @@ import gc
|
||||
if 'torch' in sys.modules:
|
||||
logging.warning("WARNING: Potential Error in code: Torch already imported, torch should never be imported before this point.")
|
||||
|
||||
|
||||
import comfy.utils
|
||||
|
||||
if not IS_PYISOLATE_CHILD:
|
||||
import execution
|
||||
import server
|
||||
from protocol import BinaryEventTypes
|
||||
import nodes
|
||||
import comfy.model_management
|
||||
import comfyui_version
|
||||
import app.logger
|
||||
import hook_breaker_ac10a0
|
||||
import execution
|
||||
import server
|
||||
from protocol import BinaryEventTypes
|
||||
import nodes
|
||||
import comfy.model_management
|
||||
import comfyui_version
|
||||
import app.logger
|
||||
import hook_breaker_ac10a0
|
||||
|
||||
import comfy.memory_management
|
||||
import comfy.model_patcher
|
||||
@@ -515,10 +462,6 @@ def start_comfyui(asyncio_loop=None):
|
||||
asyncio.set_event_loop(asyncio_loop)
|
||||
prompt_server = server.PromptServer(asyncio_loop)
|
||||
|
||||
if args.use_process_isolation:
|
||||
from comfy.isolation import start_isolation_loading_early
|
||||
start_isolation_loading_early(asyncio_loop)
|
||||
|
||||
if args.enable_manager and not args.disable_manager_ui:
|
||||
comfyui_manager.start()
|
||||
|
||||
@@ -563,13 +506,12 @@ def start_comfyui(asyncio_loop=None):
|
||||
if __name__ == "__main__":
|
||||
# Running directly, just start ComfyUI.
|
||||
logging.info("Python version: {}".format(sys.version))
|
||||
if not IS_PYISOLATE_CHILD:
|
||||
logging.info("ComfyUI version: {}".format(comfyui_version.__version__))
|
||||
for package in ("comfy-aimdo", "comfy-kitchen"):
|
||||
try:
|
||||
logging.info("{} version: {}".format(package, importlib.metadata.version(package)))
|
||||
except:
|
||||
pass
|
||||
logging.info("ComfyUI version: {}".format(comfyui_version.__version__))
|
||||
for package in ("comfy-aimdo", "comfy-kitchen"):
|
||||
try:
|
||||
logging.info("{} version: {}".format(package, importlib.metadata.version(package)))
|
||||
except:
|
||||
pass
|
||||
|
||||
if sys.version_info.major == 3 and sys.version_info.minor < 10:
|
||||
logging.warning("WARNING: You are using a python version older than 3.10, please upgrade to a newer one. 3.12 and above is recommended.")
|
||||
|
||||
+2
-2
@@ -86,6 +86,6 @@ def image_alpha_fix(destination, source):
|
||||
if destination.shape[-1] < source.shape[-1]:
|
||||
source = source[...,:destination.shape[-1]]
|
||||
elif destination.shape[-1] > source.shape[-1]:
|
||||
destination = torch.nn.functional.pad(destination, (0, 1))
|
||||
destination[..., -1] = 1.0
|
||||
source = torch.nn.functional.pad(source, (0, 1))
|
||||
source[..., -1] = 1.0
|
||||
return destination, source
|
||||
|
||||
@@ -1694,26 +1694,27 @@ class LoadImage:
|
||||
|
||||
RETURN_TYPES = ("IMAGE", "MASK")
|
||||
FUNCTION = "load_image"
|
||||
|
||||
def load_image(self, image):
|
||||
image_path = folder_paths.get_annotated_filepath(image)
|
||||
|
||||
dtype = comfy.model_management.intermediate_dtype()
|
||||
device = comfy.model_management.intermediate_device()
|
||||
|
||||
components = InputImpl.VideoFromFile(image_path).get_components()
|
||||
if components.images.shape[0] > 0:
|
||||
return (components.images, 1.0 - components.alpha[..., -1] if components.alpha is not None else torch.zeros((components.images.shape[0], 64, 64), dtype=torch.float32, device="cpu"))
|
||||
return (components.images.to(device=device, dtype=dtype), (1.0 - components.alpha[..., -1]).to(device=device, dtype=dtype) if components.alpha is not None else torch.zeros((components.images.shape[0], 64, 64), dtype=dtype, device=device))
|
||||
|
||||
# This code is left here to handle animated webp which pyav does not support loading
|
||||
img = node_helpers.pillow(Image.open, image_path)
|
||||
|
||||
output_images = []
|
||||
output_masks = []
|
||||
w, h = None, None
|
||||
|
||||
dtype = comfy.model_management.intermediate_dtype()
|
||||
|
||||
for i in ImageSequence.Iterator(img):
|
||||
i = node_helpers.pillow(ImageOps.exif_transpose, i)
|
||||
|
||||
if i.mode == 'I':
|
||||
i = i.point(lambda i: i * (1 / 255))
|
||||
image = i.convert("RGB")
|
||||
|
||||
if len(output_images) == 0:
|
||||
@@ -1728,25 +1729,15 @@ class LoadImage:
|
||||
if 'A' in i.getbands():
|
||||
mask = np.array(i.getchannel('A')).astype(np.float32) / 255.0
|
||||
mask = 1. - torch.from_numpy(mask)
|
||||
elif i.mode == 'P' and 'transparency' in i.info:
|
||||
mask = np.array(i.convert('RGBA').getchannel('A')).astype(np.float32) / 255.0
|
||||
mask = 1. - torch.from_numpy(mask)
|
||||
else:
|
||||
mask = torch.zeros((64,64), dtype=torch.float32, device="cpu")
|
||||
mask = torch.zeros((64, 64), dtype=torch.float32, device="cpu")
|
||||
output_images.append(image.to(dtype=dtype))
|
||||
output_masks.append(mask.unsqueeze(0).to(dtype=dtype))
|
||||
|
||||
if img.format == "MPO":
|
||||
break # ignore all frames except the first one for MPO format
|
||||
output_image = torch.cat(output_images, dim=0)
|
||||
output_mask = torch.cat(output_masks, dim=0)
|
||||
|
||||
if len(output_images) > 1:
|
||||
output_image = torch.cat(output_images, dim=0)
|
||||
output_mask = torch.cat(output_masks, dim=0)
|
||||
else:
|
||||
output_image = output_images[0]
|
||||
output_mask = output_masks[0]
|
||||
|
||||
return (output_image, output_mask)
|
||||
return (output_image.to(device=device, dtype=dtype), output_mask.to(device=device, dtype=dtype))
|
||||
|
||||
@classmethod
|
||||
def IS_CHANGED(s, image):
|
||||
@@ -1912,7 +1903,6 @@ class ImageInvert:
|
||||
|
||||
class ImageBatch:
|
||||
SEARCH_ALIASES = ["combine images", "merge images", "stack images"]
|
||||
ESSENTIALS_CATEGORY = "Image Tools"
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
@@ -2296,27 +2286,6 @@ async def init_external_custom_nodes():
|
||||
Returns:
|
||||
None
|
||||
"""
|
||||
whitelist = set()
|
||||
isolated_module_paths = set()
|
||||
if args.use_process_isolation:
|
||||
from pathlib import Path
|
||||
from comfy.isolation import await_isolation_loading, get_claimed_paths
|
||||
from comfy.isolation.host_policy import load_host_policy
|
||||
|
||||
# Load Global Host Policy
|
||||
host_policy = load_host_policy(Path(folder_paths.base_path))
|
||||
whitelist_dict = host_policy.get("whitelist", {})
|
||||
# Normalize whitelist keys to lowercase for case-insensitive matching
|
||||
# (matches ComfyUI-Manager's normalization: project.name.strip().lower())
|
||||
whitelist = set(k.strip().lower() for k in whitelist_dict.keys())
|
||||
logging.info(f"][ Loaded Whitelist: {len(whitelist)} nodes allowed.")
|
||||
|
||||
isolated_specs = await await_isolation_loading()
|
||||
for spec in isolated_specs:
|
||||
NODE_CLASS_MAPPINGS.setdefault(spec.node_name, spec.stub_class)
|
||||
NODE_DISPLAY_NAME_MAPPINGS.setdefault(spec.node_name, spec.display_name)
|
||||
isolated_module_paths = get_claimed_paths()
|
||||
|
||||
base_node_names = set(NODE_CLASS_MAPPINGS.keys())
|
||||
node_paths = folder_paths.get_folder_paths("custom_nodes")
|
||||
node_import_times = []
|
||||
@@ -2340,16 +2309,6 @@ async def init_external_custom_nodes():
|
||||
logging.info(f"Blocked by policy: {module_path}")
|
||||
continue
|
||||
|
||||
if args.use_process_isolation:
|
||||
if Path(module_path).resolve() in isolated_module_paths:
|
||||
continue
|
||||
|
||||
# Tri-State Enforcement: If not Isolated (checked above), MUST be Whitelisted.
|
||||
# Normalize to lowercase for case-insensitive matching (matches ComfyUI-Manager)
|
||||
if possible_module.strip().lower() not in whitelist:
|
||||
logging.warning(f"][ REJECTED: Node '{possible_module}' is blocked by security policy (not whitelisted/isolated).")
|
||||
continue
|
||||
|
||||
time_before = time.perf_counter()
|
||||
success = await load_custom_node(module_path, base_node_names, module_parent="custom_nodes")
|
||||
node_import_times.append((time.perf_counter() - time_before, module_path, success))
|
||||
@@ -2364,14 +2323,6 @@ async def init_external_custom_nodes():
|
||||
logging.info("{:6.1f} seconds{}: {}".format(n[0], import_message, n[1]))
|
||||
logging.info("")
|
||||
|
||||
if args.use_process_isolation:
|
||||
from comfy.isolation import isolated_node_timings
|
||||
if isolated_node_timings:
|
||||
logging.info("\nImport times for isolated custom nodes:")
|
||||
for timing, path, count in sorted(isolated_node_timings):
|
||||
logging.info("{:6.1f} seconds: {} ({})".format(timing, path, count))
|
||||
logging.info("")
|
||||
|
||||
async def init_builtin_extra_nodes():
|
||||
"""
|
||||
Initializes the built-in extra nodes in ComfyUI.
|
||||
|
||||
@@ -10,17 +10,6 @@ homepage = "https://www.comfy.org/"
|
||||
repository = "https://github.com/comfyanonymous/ComfyUI"
|
||||
documentation = "https://docs.comfy.org/"
|
||||
|
||||
[tool.comfy.host]
|
||||
sandbox_mode = "required"
|
||||
allow_network = false
|
||||
writable_paths = ["/dev/shm"]
|
||||
|
||||
[tool.comfy.host.whitelist]
|
||||
"ComfyUI-GGUF" = "*"
|
||||
"ComfyUI-KJNodes" = "*"
|
||||
"ComfyUI-Manager" = "*"
|
||||
"websocket_image_save.py" = "*"
|
||||
|
||||
[tool.ruff]
|
||||
lint.select = [
|
||||
"N805", # invalid-first-argument-name-for-method
|
||||
|
||||
+1
-3
@@ -1,5 +1,5 @@
|
||||
comfyui-frontend-package==1.42.15
|
||||
comfyui-workflow-templates==0.9.63
|
||||
comfyui-workflow-templates==0.9.68
|
||||
comfyui-embedded-docs==0.4.4
|
||||
torch
|
||||
torchsde
|
||||
@@ -35,5 +35,3 @@ pydantic~=2.0
|
||||
pydantic-settings~=2.0
|
||||
PyOpenGL
|
||||
glfw
|
||||
uv
|
||||
pyisolate==0.10.2
|
||||
|
||||
@@ -3,6 +3,7 @@ import sys
|
||||
import asyncio
|
||||
import traceback
|
||||
import time
|
||||
|
||||
import nodes
|
||||
import folder_paths
|
||||
import execution
|
||||
@@ -201,8 +202,6 @@ def create_block_external_middleware():
|
||||
class PromptServer():
|
||||
def __init__(self, loop):
|
||||
PromptServer.instance = self
|
||||
if loop is None:
|
||||
loop = asyncio.get_event_loop()
|
||||
|
||||
self.user_manager = UserManager()
|
||||
self.model_file_manager = ModelFileManager()
|
||||
@@ -353,17 +352,6 @@ class PromptServer():
|
||||
extensions.extend(list(map(lambda f: "/extensions/" + urllib.parse.quote(
|
||||
name) + "/" + os.path.relpath(f, dir).replace("\\", "/"), files)))
|
||||
|
||||
# Include JS files from proxied web directories (isolated nodes)
|
||||
if args.use_process_isolation:
|
||||
from comfy.isolation.proxies.web_directory_proxy import get_web_directory_cache
|
||||
cache = get_web_directory_cache()
|
||||
for ext_name in cache.extension_names:
|
||||
for entry in cache.list_files(ext_name):
|
||||
if entry["relative_path"].endswith(".js"):
|
||||
extensions.append(
|
||||
"/extensions/" + urllib.parse.quote(ext_name) + "/" + entry["relative_path"]
|
||||
)
|
||||
|
||||
return web.json_response(extensions)
|
||||
|
||||
def get_dir_by_type(dir_type):
|
||||
@@ -1079,36 +1067,6 @@ class PromptServer():
|
||||
for name, dir in nodes.EXTENSION_WEB_DIRS.items():
|
||||
self.app.add_routes([web.static('/extensions/' + name, dir)])
|
||||
|
||||
# Add dynamic handler for proxied web directories (isolated nodes)
|
||||
if args.use_process_isolation:
|
||||
from comfy.isolation.proxies.web_directory_proxy import (
|
||||
get_web_directory_cache,
|
||||
ALLOWED_EXTENSIONS,
|
||||
)
|
||||
|
||||
async def proxied_web_handler(request):
|
||||
ext_name = request.match_info["ext_name"]
|
||||
file_path = request.match_info["file_path"]
|
||||
|
||||
suffix = os.path.splitext(file_path)[1].lower()
|
||||
if suffix not in ALLOWED_EXTENSIONS:
|
||||
return web.Response(status=403, text="Forbidden file type")
|
||||
|
||||
cache = get_web_directory_cache()
|
||||
result = cache.get_file(ext_name, file_path)
|
||||
if result is None:
|
||||
return web.Response(status=404, text="Not found")
|
||||
|
||||
return web.Response(
|
||||
body=result["content"],
|
||||
content_type=result["content_type"],
|
||||
)
|
||||
|
||||
self.app.router.add_get(
|
||||
"/extensions/{ext_name}/{file_path:.*}",
|
||||
proxied_web_handler,
|
||||
)
|
||||
|
||||
installed_templates_version = FrontendManager.get_installed_templates_version()
|
||||
use_legacy_templates = True
|
||||
if installed_templates_version:
|
||||
|
||||
@@ -0,0 +1,139 @@
|
||||
import pytest
|
||||
|
||||
from comfy_api.latest._io import (
|
||||
Combo,
|
||||
RemoteComboOptions,
|
||||
RemoteItemSchema,
|
||||
RemoteOptions,
|
||||
)
|
||||
|
||||
|
||||
def _schema(**overrides):
|
||||
defaults = dict(value_field="id", label_field="name")
|
||||
return RemoteItemSchema(**{**defaults, **overrides})
|
||||
|
||||
|
||||
def _combo(**overrides):
|
||||
defaults = dict(route="/proxy/foo", item_schema=_schema())
|
||||
return RemoteComboOptions(**{**defaults, **overrides})
|
||||
|
||||
|
||||
def test_item_schema_defaults_accepted():
|
||||
d = _schema().as_dict()
|
||||
assert d == {"value_field": "id", "label_field": "name", "preview_type": "image"}
|
||||
|
||||
|
||||
def test_item_schema_full_config_accepted():
|
||||
d = _schema(
|
||||
preview_url_field="preview",
|
||||
preview_type="audio",
|
||||
description_field="desc",
|
||||
search_fields=["first", "last", "profile.email"],
|
||||
).as_dict()
|
||||
assert d["preview_type"] == "audio"
|
||||
assert d["search_fields"] == ["first", "last", "profile.email"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"bad_fields",
|
||||
[
|
||||
["{first} {last}"],
|
||||
["name", "{age}"],
|
||||
["leading{"],
|
||||
["trailing}"],
|
||||
],
|
||||
)
|
||||
def test_item_schema_rejects_template_strings_in_search_fields(bad_fields):
|
||||
with pytest.raises(ValueError, match="search_fields"):
|
||||
_schema(search_fields=bad_fields)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("bad_preview_type", ["middle", "IMAGE", "", "gif"])
|
||||
def test_item_schema_rejects_unknown_preview_type(bad_preview_type):
|
||||
with pytest.raises(ValueError, match="preview_type"):
|
||||
_schema(preview_type=bad_preview_type)
|
||||
|
||||
|
||||
def test_combo_options_minimal_accepted():
|
||||
d = _combo().as_dict()
|
||||
assert d["route"] == "/proxy/foo"
|
||||
assert d["refresh_button"] is True
|
||||
assert "item_schema" in d
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"route",
|
||||
[
|
||||
"/proxy/foo",
|
||||
"/voices",
|
||||
],
|
||||
)
|
||||
def test_combo_options_accepts_valid_routes(route):
|
||||
_combo(route=route)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"route",
|
||||
[
|
||||
"",
|
||||
"api.example.com/voices",
|
||||
"voices",
|
||||
"ftp-no-scheme",
|
||||
"http://localhost:9000/voices",
|
||||
"https://api.example.com/v1/voices",
|
||||
],
|
||||
)
|
||||
def test_combo_options_rejects_non_relative_routes(route):
|
||||
with pytest.raises(ValueError, match="'route'"):
|
||||
_combo(route=route)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("bad_auto_select", ["middle", "FIRST", "", "firstlast"])
|
||||
def test_combo_options_rejects_unknown_auto_select(bad_auto_select):
|
||||
with pytest.raises(ValueError, match="auto_select"):
|
||||
_combo(auto_select=bad_auto_select)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("bad_refresh", [1, 127])
|
||||
def test_combo_options_refresh_in_forbidden_range_rejected(bad_refresh):
|
||||
with pytest.raises(ValueError, match="refresh"):
|
||||
_combo(refresh=bad_refresh)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("ok_refresh", [0, -1, 128])
|
||||
def test_combo_options_refresh_valid_values_accepted(ok_refresh):
|
||||
_combo(refresh=ok_refresh)
|
||||
|
||||
|
||||
def test_combo_options_timeout_negative_rejected():
|
||||
with pytest.raises(ValueError, match="timeout"):
|
||||
_combo(timeout=-1)
|
||||
|
||||
|
||||
def test_combo_options_max_retries_negative_rejected():
|
||||
with pytest.raises(ValueError, match="max_retries"):
|
||||
_combo(max_retries=-1)
|
||||
|
||||
|
||||
def test_combo_options_as_dict_prunes_none_fields():
|
||||
d = _combo().as_dict()
|
||||
for pruned in ("response_key", "refresh", "timeout", "max_retries", "auto_select"):
|
||||
assert pruned not in d
|
||||
|
||||
|
||||
def test_combo_input_accepts_remote_combo_alone():
|
||||
Combo.Input("voice", remote_combo=_combo())
|
||||
|
||||
|
||||
def test_combo_input_rejects_remote_plus_remote_combo():
|
||||
with pytest.raises(ValueError, match="remote.*remote_combo"):
|
||||
Combo.Input(
|
||||
"voice",
|
||||
remote=RemoteOptions(route="/r", refresh_button=True),
|
||||
remote_combo=_combo(),
|
||||
)
|
||||
|
||||
|
||||
def test_combo_input_rejects_options_plus_remote_combo():
|
||||
with pytest.raises(ValueError, match="options.*remote_combo"):
|
||||
Combo.Input("voice", options=["a", "b"], remote_combo=_combo())
|
||||
@@ -1,211 +0,0 @@
|
||||
# pylint: disable=import-outside-toplevel,import-error
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import os
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _artifact_dir() -> Path | None:
|
||||
raw = os.environ.get("PYISOLATE_ARTIFACT_DIR")
|
||||
if not raw:
|
||||
return None
|
||||
path = Path(raw)
|
||||
path.mkdir(parents=True, exist_ok=True)
|
||||
return path
|
||||
|
||||
|
||||
def _write_artifact(name: str, content: str) -> None:
|
||||
artifact_dir = _artifact_dir()
|
||||
if artifact_dir is None:
|
||||
return
|
||||
(artifact_dir / name).write_text(content, encoding="utf-8")
|
||||
|
||||
|
||||
def _contains_tensor_marker(value: Any) -> bool:
|
||||
if isinstance(value, dict):
|
||||
if value.get("__type__") == "TensorValue":
|
||||
return True
|
||||
return any(_contains_tensor_marker(v) for v in value.values())
|
||||
if isinstance(value, (list, tuple)):
|
||||
return any(_contains_tensor_marker(v) for v in value)
|
||||
return False
|
||||
|
||||
|
||||
class InspectRuntimeNode:
|
||||
RETURN_TYPES = (
|
||||
"STRING",
|
||||
"STRING",
|
||||
"BOOLEAN",
|
||||
"BOOLEAN",
|
||||
"STRING",
|
||||
"STRING",
|
||||
"BOOLEAN",
|
||||
)
|
||||
RETURN_NAMES = (
|
||||
"path_dump",
|
||||
"runtime_report",
|
||||
"saw_comfy_root",
|
||||
"imported_comfy_wrapper",
|
||||
"comfy_module_dump",
|
||||
"python_exe",
|
||||
"saw_user_site",
|
||||
)
|
||||
FUNCTION = "inspect"
|
||||
CATEGORY = "PyIsolated/SealedWorker"
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls) -> dict[str, Any]: # noqa: N802
|
||||
return {"required": {}}
|
||||
|
||||
def inspect(self) -> tuple[str, str, bool, bool, str, str, bool]:
|
||||
import cfgrib
|
||||
import eccodes
|
||||
import xarray as xr
|
||||
|
||||
path_dump = "\n".join(sys.path)
|
||||
comfy_root = os.environ.get("COMFYUI_ROOT")
|
||||
if comfy_root is None:
|
||||
comfy_root = str(Path(__file__).resolve().parents[3])
|
||||
saw_comfy_root = any(
|
||||
entry == comfy_root
|
||||
or entry.startswith(f"{comfy_root}/comfy")
|
||||
or entry.startswith(f"{comfy_root}/.venv")
|
||||
for entry in sys.path
|
||||
)
|
||||
imported_comfy_wrapper = "comfy.isolation.extension_wrapper" in sys.modules
|
||||
comfy_module_dump = "\n".join(
|
||||
sorted(name for name in sys.modules if name.startswith("comfy"))
|
||||
)
|
||||
saw_user_site = any("/.local/lib/" in entry for entry in sys.path)
|
||||
python_exe = sys.executable
|
||||
|
||||
runtime_lines = [
|
||||
"Conda sealed worker runtime probe",
|
||||
f"python_exe={python_exe}",
|
||||
f"xarray_origin={getattr(xr, '__file__', '<missing>')}",
|
||||
f"cfgrib_origin={getattr(cfgrib, '__file__', '<missing>')}",
|
||||
f"eccodes_origin={getattr(eccodes, '__file__', '<missing>')}",
|
||||
f"saw_comfy_root={saw_comfy_root}",
|
||||
f"imported_comfy_wrapper={imported_comfy_wrapper}",
|
||||
f"saw_user_site={saw_user_site}",
|
||||
]
|
||||
runtime_report = "\n".join(runtime_lines)
|
||||
|
||||
_write_artifact("child_bootstrap_paths.txt", path_dump)
|
||||
_write_artifact("child_import_trace.txt", comfy_module_dump)
|
||||
_write_artifact("child_dependency_dump.txt", runtime_report)
|
||||
logger.warning("][ Conda sealed runtime probe executed")
|
||||
logger.warning("][ conda python executable: %s", python_exe)
|
||||
logger.warning(
|
||||
"][ conda dependency origins: xarray=%s cfgrib=%s eccodes=%s",
|
||||
getattr(xr, "__file__", "<missing>"),
|
||||
getattr(cfgrib, "__file__", "<missing>"),
|
||||
getattr(eccodes, "__file__", "<missing>"),
|
||||
)
|
||||
|
||||
return (
|
||||
path_dump,
|
||||
runtime_report,
|
||||
saw_comfy_root,
|
||||
imported_comfy_wrapper,
|
||||
comfy_module_dump,
|
||||
python_exe,
|
||||
saw_user_site,
|
||||
)
|
||||
|
||||
|
||||
class OpenWeatherDatasetNode:
|
||||
RETURN_TYPES = ("FLOAT", "STRING", "STRING")
|
||||
RETURN_NAMES = ("sum_value", "grib_path", "dependency_report")
|
||||
FUNCTION = "open_dataset"
|
||||
CATEGORY = "PyIsolated/SealedWorker"
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls) -> dict[str, Any]: # noqa: N802
|
||||
return {"required": {}}
|
||||
|
||||
def open_dataset(self) -> tuple[float, str, str]:
|
||||
import eccodes
|
||||
import xarray as xr
|
||||
|
||||
artifact_dir = _artifact_dir()
|
||||
if artifact_dir is None:
|
||||
artifact_dir = Path(os.environ.get("HOME", ".")) / "pyisolate_artifacts"
|
||||
artifact_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
grib_path = artifact_dir / "toolkit_weather_fixture.grib2"
|
||||
|
||||
gid = eccodes.codes_grib_new_from_samples("GRIB2")
|
||||
for key, value in [
|
||||
("gridType", "regular_ll"),
|
||||
("Nx", 2),
|
||||
("Ny", 2),
|
||||
("latitudeOfFirstGridPointInDegrees", 1.0),
|
||||
("longitudeOfFirstGridPointInDegrees", 0.0),
|
||||
("latitudeOfLastGridPointInDegrees", 0.0),
|
||||
("longitudeOfLastGridPointInDegrees", 1.0),
|
||||
("iDirectionIncrementInDegrees", 1.0),
|
||||
("jDirectionIncrementInDegrees", 1.0),
|
||||
("jScansPositively", 0),
|
||||
("shortName", "t"),
|
||||
("typeOfLevel", "surface"),
|
||||
("level", 0),
|
||||
("date", 20260315),
|
||||
("time", 0),
|
||||
("step", 0),
|
||||
]:
|
||||
eccodes.codes_set(gid, key, value)
|
||||
|
||||
eccodes.codes_set_values(gid, [1.0, 2.0, 3.0, 4.0])
|
||||
with grib_path.open("wb") as handle:
|
||||
eccodes.codes_write(gid, handle)
|
||||
eccodes.codes_release(gid)
|
||||
|
||||
dataset = xr.open_dataset(grib_path, engine="cfgrib")
|
||||
sum_value = float(dataset["t"].sum().item())
|
||||
dependency_report = "\n".join(
|
||||
[
|
||||
f"dataset_sum={sum_value}",
|
||||
f"grib_path={grib_path}",
|
||||
"xarray_engine=cfgrib",
|
||||
]
|
||||
)
|
||||
_write_artifact("weather_dependency_report.txt", dependency_report)
|
||||
logger.warning("][ cfgrib import ok")
|
||||
logger.warning("][ xarray open_dataset engine=cfgrib path=%s", grib_path)
|
||||
logger.warning("][ conda weather dataset sum=%s", sum_value)
|
||||
return sum_value, str(grib_path), dependency_report
|
||||
|
||||
|
||||
class EchoLatentNode:
|
||||
RETURN_TYPES = ("LATENT", "BOOLEAN")
|
||||
RETURN_NAMES = ("latent", "saw_json_tensor")
|
||||
FUNCTION = "echo_latent"
|
||||
CATEGORY = "PyIsolated/SealedWorker"
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls) -> dict[str, Any]: # noqa: N802
|
||||
return {"required": {"latent": ("LATENT",)}}
|
||||
|
||||
def echo_latent(self, latent: Any) -> tuple[Any, bool]:
|
||||
saw_json_tensor = _contains_tensor_marker(latent)
|
||||
logger.warning("][ conda latent echo json_marker=%s", saw_json_tensor)
|
||||
return latent, saw_json_tensor
|
||||
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"CondaSealedRuntimeProbe": InspectRuntimeNode,
|
||||
"CondaSealedOpenWeatherDataset": OpenWeatherDatasetNode,
|
||||
"CondaSealedLatentEcho": EchoLatentNode,
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"CondaSealedRuntimeProbe": "Conda Sealed Runtime Probe",
|
||||
"CondaSealedOpenWeatherDataset": "Conda Sealed Open Weather Dataset",
|
||||
"CondaSealedLatentEcho": "Conda Sealed Latent Echo",
|
||||
}
|
||||
@@ -1,13 +0,0 @@
|
||||
[project]
|
||||
name = "comfyui-toolkit-conda-sealed-worker"
|
||||
version = "0.1.0"
|
||||
dependencies = ["xarray", "cfgrib"]
|
||||
|
||||
[tool.comfy.isolation]
|
||||
can_isolate = true
|
||||
share_torch = false
|
||||
package_manager = "conda"
|
||||
execution_model = "sealed_worker"
|
||||
standalone = true
|
||||
conda_channels = ["conda-forge"]
|
||||
conda_dependencies = ["eccodes", "cfgrib"]
|
||||
@@ -1,7 +0,0 @@
|
||||
[tool.comfy.host]
|
||||
sandbox_mode = "required"
|
||||
allow_network = false
|
||||
writable_paths = [
|
||||
"/dev/shm",
|
||||
"/var/lib/comfyui/output",
|
||||
]
|
||||
@@ -1,6 +0,0 @@
|
||||
from .probe_nodes import (
|
||||
NODE_CLASS_MAPPINGS as NODE_CLASS_MAPPINGS,
|
||||
NODE_DISPLAY_NAME_MAPPINGS as NODE_DISPLAY_NAME_MAPPINGS,
|
||||
)
|
||||
|
||||
__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]
|
||||
@@ -1,75 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
|
||||
class InternalIsolationProbeImage:
|
||||
CATEGORY = "tests/isolation"
|
||||
RETURN_TYPES = ()
|
||||
FUNCTION = "run"
|
||||
OUTPUT_NODE = True
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {"required": {}}
|
||||
|
||||
def run(self):
|
||||
from comfy_api.latest import UI
|
||||
import torch
|
||||
|
||||
image = torch.zeros((1, 2, 2, 3), dtype=torch.float32)
|
||||
image[:, :, :, 0] = 1.0
|
||||
ui = UI.PreviewImage(image)
|
||||
return {"ui": ui.as_dict(), "result": ()}
|
||||
|
||||
|
||||
class InternalIsolationProbeAudio:
|
||||
CATEGORY = "tests/isolation"
|
||||
RETURN_TYPES = ()
|
||||
FUNCTION = "run"
|
||||
OUTPUT_NODE = True
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {"required": {}}
|
||||
|
||||
def run(self):
|
||||
from comfy_api.latest import UI
|
||||
import torch
|
||||
|
||||
waveform = torch.zeros((1, 1, 32), dtype=torch.float32)
|
||||
audio = {"waveform": waveform, "sample_rate": 44100}
|
||||
ui = UI.PreviewAudio(audio)
|
||||
return {"ui": ui.as_dict(), "result": ()}
|
||||
|
||||
|
||||
class InternalIsolationProbeUI3D:
|
||||
CATEGORY = "tests/isolation"
|
||||
RETURN_TYPES = ()
|
||||
FUNCTION = "run"
|
||||
OUTPUT_NODE = True
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {"required": {}}
|
||||
|
||||
def run(self):
|
||||
from comfy_api.latest import UI
|
||||
import torch
|
||||
|
||||
bg_image = torch.zeros((1, 2, 2, 3), dtype=torch.float32)
|
||||
bg_image[:, :, :, 1] = 1.0
|
||||
camera_info = {"distance": 1.0}
|
||||
ui = UI.PreviewUI3D("internal_probe_preview.obj", camera_info, bg_image=bg_image)
|
||||
return {"ui": ui.as_dict(), "result": ()}
|
||||
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"InternalIsolationProbeImage": InternalIsolationProbeImage,
|
||||
"InternalIsolationProbeAudio": InternalIsolationProbeAudio,
|
||||
"InternalIsolationProbeUI3D": InternalIsolationProbeUI3D,
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"InternalIsolationProbeImage": "Internal Isolation Probe Image",
|
||||
"InternalIsolationProbeAudio": "Internal Isolation Probe Audio",
|
||||
"InternalIsolationProbeUI3D": "Internal Isolation Probe UI3D",
|
||||
}
|
||||
@@ -1,957 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import importlib.util
|
||||
import os
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
|
||||
COMFYUI_ROOT = Path(__file__).resolve().parents[2]
|
||||
UV_SEALED_WORKER_MODULE = COMFYUI_ROOT / "tests" / "isolation" / "uv_sealed_worker" / "__init__.py"
|
||||
FORBIDDEN_MINIMAL_SEALED_MODULES = (
|
||||
"torch",
|
||||
"folder_paths",
|
||||
"comfy.utils",
|
||||
"comfy.model_management",
|
||||
"main",
|
||||
"comfy.isolation.extension_wrapper",
|
||||
)
|
||||
FORBIDDEN_SEALED_SINGLETON_MODULES = (
|
||||
"torch",
|
||||
"folder_paths",
|
||||
"comfy.utils",
|
||||
"comfy_execution.progress",
|
||||
)
|
||||
FORBIDDEN_EXACT_SMALL_PROXY_MODULES = FORBIDDEN_SEALED_SINGLETON_MODULES
|
||||
FORBIDDEN_MODEL_MANAGEMENT_MODULES = (
|
||||
"comfy.model_management",
|
||||
)
|
||||
|
||||
|
||||
def _load_module_from_path(module_name: str, module_path: Path):
|
||||
spec = importlib.util.spec_from_file_location(module_name, module_path)
|
||||
if spec is None or spec.loader is None:
|
||||
raise RuntimeError(f"unable to build import spec for {module_path}")
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
sys.modules[module_name] = module
|
||||
try:
|
||||
spec.loader.exec_module(module)
|
||||
except Exception:
|
||||
sys.modules.pop(module_name, None)
|
||||
raise
|
||||
return module
|
||||
|
||||
|
||||
def matching_modules(prefixes: tuple[str, ...], modules: set[str]) -> list[str]:
|
||||
return sorted(
|
||||
module_name
|
||||
for module_name in modules
|
||||
if any(
|
||||
module_name == prefix or module_name.startswith(f"{prefix}.")
|
||||
for prefix in prefixes
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def _load_helper_proxy_service() -> Any | None:
|
||||
try:
|
||||
from comfy.isolation.proxies.helper_proxies import HelperProxiesService
|
||||
except (ImportError, AttributeError):
|
||||
return None
|
||||
return HelperProxiesService
|
||||
|
||||
|
||||
def _load_model_management_proxy() -> Any | None:
|
||||
try:
|
||||
from comfy.isolation.proxies.model_management_proxy import ModelManagementProxy
|
||||
except (ImportError, AttributeError):
|
||||
return None
|
||||
return ModelManagementProxy
|
||||
|
||||
|
||||
async def _capture_minimal_sealed_worker_imports() -> dict[str, object]:
|
||||
from pyisolate.sealed import SealedNodeExtension
|
||||
|
||||
module_name = "tests.isolation.uv_sealed_worker_boundary_probe"
|
||||
before = set(sys.modules)
|
||||
extension = SealedNodeExtension()
|
||||
module = _load_module_from_path(module_name, UV_SEALED_WORKER_MODULE)
|
||||
try:
|
||||
await extension.on_module_loaded(module)
|
||||
node_list = await extension.list_nodes()
|
||||
node_details = await extension.get_node_details("UVSealedRuntimeProbe")
|
||||
imported = set(sys.modules) - before
|
||||
return {
|
||||
"mode": "minimal_sealed_worker",
|
||||
"node_names": sorted(node_list),
|
||||
"runtime_probe_function": node_details["function"],
|
||||
"modules": sorted(imported),
|
||||
"forbidden_matches": matching_modules(FORBIDDEN_MINIMAL_SEALED_MODULES, imported),
|
||||
}
|
||||
finally:
|
||||
sys.modules.pop(module_name, None)
|
||||
|
||||
|
||||
def capture_minimal_sealed_worker_imports() -> dict[str, object]:
|
||||
return asyncio.run(_capture_minimal_sealed_worker_imports())
|
||||
|
||||
|
||||
class FakeSingletonCaller:
|
||||
def __init__(self, methods: dict[str, Any], calls: list[dict[str, Any]], object_id: str):
|
||||
self._methods = methods
|
||||
self._calls = calls
|
||||
self._object_id = object_id
|
||||
|
||||
def __getattr__(self, name: str):
|
||||
if name not in self._methods:
|
||||
raise AttributeError(name)
|
||||
|
||||
async def method(*args: Any, **kwargs: Any) -> Any:
|
||||
self._calls.append(
|
||||
{
|
||||
"object_id": self._object_id,
|
||||
"method": name,
|
||||
"args": list(args),
|
||||
"kwargs": dict(kwargs),
|
||||
}
|
||||
)
|
||||
result = self._methods[name]
|
||||
return result(*args, **kwargs) if callable(result) else result
|
||||
|
||||
return method
|
||||
|
||||
|
||||
class FakeSingletonRPC:
|
||||
def __init__(self) -> None:
|
||||
self.calls: list[dict[str, Any]] = []
|
||||
self._device = {"__pyisolate_torch_device__": "cpu"}
|
||||
self._services: dict[str, dict[str, Any]] = {
|
||||
"FolderPathsProxy": {
|
||||
"rpc_get_models_dir": lambda: "/sandbox/models",
|
||||
"rpc_get_folder_names_and_paths": lambda: {
|
||||
"checkpoints": {
|
||||
"paths": ["/sandbox/models/checkpoints"],
|
||||
"extensions": [".ckpt", ".safetensors"],
|
||||
}
|
||||
},
|
||||
"rpc_get_extension_mimetypes_cache": lambda: {"webp": "image"},
|
||||
"rpc_get_filename_list_cache": lambda: {},
|
||||
"rpc_get_temp_directory": lambda: "/sandbox/temp",
|
||||
"rpc_get_input_directory": lambda: "/sandbox/input",
|
||||
"rpc_get_output_directory": lambda: "/sandbox/output",
|
||||
"rpc_get_user_directory": lambda: "/sandbox/user",
|
||||
"rpc_get_annotated_filepath": self._get_annotated_filepath,
|
||||
"rpc_exists_annotated_filepath": lambda _name: False,
|
||||
"rpc_add_model_folder_path": lambda *_args, **_kwargs: None,
|
||||
"rpc_get_folder_paths": lambda folder_name: [f"/sandbox/models/{folder_name}"],
|
||||
"rpc_get_filename_list": lambda folder_name: [f"{folder_name}_fixture.safetensors"],
|
||||
"rpc_get_full_path": lambda folder_name, filename: f"/sandbox/models/{folder_name}/{filename}",
|
||||
},
|
||||
"UtilsProxy": {
|
||||
"progress_bar_hook": lambda value, total, preview=None, node_id=None: {
|
||||
"value": value,
|
||||
"total": total,
|
||||
"preview": preview,
|
||||
"node_id": node_id,
|
||||
}
|
||||
},
|
||||
"ProgressProxy": {
|
||||
"rpc_set_progress": lambda value, max_value, node_id=None, image=None: {
|
||||
"value": value,
|
||||
"max_value": max_value,
|
||||
"node_id": node_id,
|
||||
"image": image,
|
||||
}
|
||||
},
|
||||
"HelperProxiesService": {
|
||||
"rpc_restore_input_types": lambda raw: raw,
|
||||
},
|
||||
"ModelManagementProxy": {
|
||||
"rpc_call": self._model_management_rpc_call,
|
||||
},
|
||||
}
|
||||
|
||||
def _model_management_rpc_call(self, method_name: str, args: Any = None, kwargs: Any = None) -> Any:
|
||||
if method_name == "get_torch_device":
|
||||
return self._device
|
||||
elif method_name == "get_torch_device_name":
|
||||
return "cpu"
|
||||
elif method_name == "get_free_memory":
|
||||
return 34359738368
|
||||
raise AssertionError(f"unexpected model_management method {method_name}")
|
||||
|
||||
@staticmethod
|
||||
def _get_annotated_filepath(name: str, default_dir: str | None = None) -> str:
|
||||
if name.endswith("[output]"):
|
||||
return f"/sandbox/output/{name[:-8]}"
|
||||
if name.endswith("[input]"):
|
||||
return f"/sandbox/input/{name[:-7]}"
|
||||
if name.endswith("[temp]"):
|
||||
return f"/sandbox/temp/{name[:-6]}"
|
||||
base_dir = default_dir or "/sandbox/input"
|
||||
return f"{base_dir}/{name}"
|
||||
|
||||
def create_caller(self, cls: Any, object_id: str):
|
||||
methods = self._services.get(object_id) or self._services.get(getattr(cls, "__name__", object_id))
|
||||
if methods is None:
|
||||
raise KeyError(object_id)
|
||||
return FakeSingletonCaller(methods, self.calls, object_id)
|
||||
|
||||
|
||||
def _clear_proxy_rpcs() -> None:
|
||||
from comfy.isolation.proxies.folder_paths_proxy import FolderPathsProxy
|
||||
from comfy.isolation.proxies.progress_proxy import ProgressProxy
|
||||
from comfy.isolation.proxies.utils_proxy import UtilsProxy
|
||||
|
||||
FolderPathsProxy.clear_rpc()
|
||||
ProgressProxy.clear_rpc()
|
||||
UtilsProxy.clear_rpc()
|
||||
helper_proxy_service = _load_helper_proxy_service()
|
||||
if helper_proxy_service is not None:
|
||||
helper_proxy_service.clear_rpc()
|
||||
model_management_proxy = _load_model_management_proxy()
|
||||
if model_management_proxy is not None and hasattr(model_management_proxy, "clear_rpc"):
|
||||
model_management_proxy.clear_rpc()
|
||||
|
||||
|
||||
def prepare_sealed_singleton_proxies(fake_rpc: FakeSingletonRPC) -> None:
|
||||
os.environ["PYISOLATE_CHILD"] = "1"
|
||||
os.environ["PYISOLATE_IMPORT_TORCH"] = "0"
|
||||
_clear_proxy_rpcs()
|
||||
|
||||
from comfy.isolation.proxies.folder_paths_proxy import FolderPathsProxy
|
||||
from comfy.isolation.proxies.progress_proxy import ProgressProxy
|
||||
from comfy.isolation.proxies.utils_proxy import UtilsProxy
|
||||
|
||||
FolderPathsProxy.set_rpc(fake_rpc)
|
||||
ProgressProxy.set_rpc(fake_rpc)
|
||||
UtilsProxy.set_rpc(fake_rpc)
|
||||
helper_proxy_service = _load_helper_proxy_service()
|
||||
if helper_proxy_service is not None:
|
||||
helper_proxy_service.set_rpc(fake_rpc)
|
||||
model_management_proxy = _load_model_management_proxy()
|
||||
if model_management_proxy is not None and hasattr(model_management_proxy, "set_rpc"):
|
||||
model_management_proxy.set_rpc(fake_rpc)
|
||||
|
||||
|
||||
def reset_forbidden_singleton_modules() -> None:
|
||||
for module_name in (
|
||||
"folder_paths",
|
||||
"comfy.utils",
|
||||
"comfy_execution.progress",
|
||||
):
|
||||
sys.modules.pop(module_name, None)
|
||||
|
||||
|
||||
class FakeExactRelayCaller:
|
||||
def __init__(self, methods: dict[str, Any], transcripts: list[dict[str, Any]], object_id: str):
|
||||
self._methods = methods
|
||||
self._transcripts = transcripts
|
||||
self._object_id = object_id
|
||||
|
||||
def __getattr__(self, name: str):
|
||||
if name not in self._methods:
|
||||
raise AttributeError(name)
|
||||
|
||||
async def method(*args: Any, **kwargs: Any) -> Any:
|
||||
self._transcripts.append(
|
||||
{
|
||||
"phase": "child_call",
|
||||
"object_id": self._object_id,
|
||||
"method": name,
|
||||
"args": list(args),
|
||||
"kwargs": dict(kwargs),
|
||||
}
|
||||
)
|
||||
impl = self._methods[name]
|
||||
self._transcripts.append(
|
||||
{
|
||||
"phase": "host_invocation",
|
||||
"object_id": self._object_id,
|
||||
"method": name,
|
||||
"target": impl["target"],
|
||||
"args": list(args),
|
||||
"kwargs": dict(kwargs),
|
||||
}
|
||||
)
|
||||
result = impl["result"](*args, **kwargs) if callable(impl["result"]) else impl["result"]
|
||||
self._transcripts.append(
|
||||
{
|
||||
"phase": "result",
|
||||
"object_id": self._object_id,
|
||||
"method": name,
|
||||
"result": result,
|
||||
}
|
||||
)
|
||||
return result
|
||||
|
||||
return method
|
||||
|
||||
|
||||
class FakeExactRelayRPC:
|
||||
def __init__(self) -> None:
|
||||
self.transcripts: list[dict[str, Any]] = []
|
||||
self._device = {"__pyisolate_torch_device__": "cpu"}
|
||||
self._services: dict[str, dict[str, Any]] = {
|
||||
"FolderPathsProxy": {
|
||||
"rpc_get_models_dir": {
|
||||
"target": "folder_paths.models_dir",
|
||||
"result": "/sandbox/models",
|
||||
},
|
||||
"rpc_get_temp_directory": {
|
||||
"target": "folder_paths.get_temp_directory",
|
||||
"result": "/sandbox/temp",
|
||||
},
|
||||
"rpc_get_input_directory": {
|
||||
"target": "folder_paths.get_input_directory",
|
||||
"result": "/sandbox/input",
|
||||
},
|
||||
"rpc_get_output_directory": {
|
||||
"target": "folder_paths.get_output_directory",
|
||||
"result": "/sandbox/output",
|
||||
},
|
||||
"rpc_get_user_directory": {
|
||||
"target": "folder_paths.get_user_directory",
|
||||
"result": "/sandbox/user",
|
||||
},
|
||||
"rpc_get_folder_names_and_paths": {
|
||||
"target": "folder_paths.folder_names_and_paths",
|
||||
"result": {
|
||||
"checkpoints": {
|
||||
"paths": ["/sandbox/models/checkpoints"],
|
||||
"extensions": [".ckpt", ".safetensors"],
|
||||
}
|
||||
},
|
||||
},
|
||||
"rpc_get_extension_mimetypes_cache": {
|
||||
"target": "folder_paths.extension_mimetypes_cache",
|
||||
"result": {"webp": "image"},
|
||||
},
|
||||
"rpc_get_filename_list_cache": {
|
||||
"target": "folder_paths.filename_list_cache",
|
||||
"result": {},
|
||||
},
|
||||
"rpc_get_annotated_filepath": {
|
||||
"target": "folder_paths.get_annotated_filepath",
|
||||
"result": lambda name, default_dir=None: FakeSingletonRPC._get_annotated_filepath(name, default_dir),
|
||||
},
|
||||
"rpc_exists_annotated_filepath": {
|
||||
"target": "folder_paths.exists_annotated_filepath",
|
||||
"result": False,
|
||||
},
|
||||
"rpc_add_model_folder_path": {
|
||||
"target": "folder_paths.add_model_folder_path",
|
||||
"result": None,
|
||||
},
|
||||
"rpc_get_folder_paths": {
|
||||
"target": "folder_paths.get_folder_paths",
|
||||
"result": lambda folder_name: [f"/sandbox/models/{folder_name}"],
|
||||
},
|
||||
"rpc_get_filename_list": {
|
||||
"target": "folder_paths.get_filename_list",
|
||||
"result": lambda folder_name: [f"{folder_name}_fixture.safetensors"],
|
||||
},
|
||||
"rpc_get_full_path": {
|
||||
"target": "folder_paths.get_full_path",
|
||||
"result": lambda folder_name, filename: f"/sandbox/models/{folder_name}/{filename}",
|
||||
},
|
||||
},
|
||||
"UtilsProxy": {
|
||||
"progress_bar_hook": {
|
||||
"target": "comfy.utils.PROGRESS_BAR_HOOK",
|
||||
"result": lambda value, total, preview=None, node_id=None: {
|
||||
"value": value,
|
||||
"total": total,
|
||||
"preview": preview,
|
||||
"node_id": node_id,
|
||||
},
|
||||
},
|
||||
},
|
||||
"ProgressProxy": {
|
||||
"rpc_set_progress": {
|
||||
"target": "comfy_execution.progress.get_progress_state().update_progress",
|
||||
"result": None,
|
||||
},
|
||||
},
|
||||
"HelperProxiesService": {
|
||||
"rpc_restore_input_types": {
|
||||
"target": "comfy.isolation.proxies.helper_proxies.restore_input_types",
|
||||
"result": lambda raw: raw,
|
||||
}
|
||||
},
|
||||
"ModelManagementProxy": {
|
||||
"rpc_call": {
|
||||
"target": "comfy.model_management.*",
|
||||
"result": self._model_management_rpc_call,
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
def _model_management_rpc_call(self, method_name: str, args: Any = None, kwargs: Any = None) -> Any:
|
||||
device = {"__pyisolate_torch_device__": "cpu"}
|
||||
if method_name == "get_torch_device":
|
||||
return device
|
||||
elif method_name == "get_torch_device_name":
|
||||
return "cpu"
|
||||
elif method_name == "get_free_memory":
|
||||
return 34359738368
|
||||
raise AssertionError(f"unexpected exact-relay method {method_name}")
|
||||
|
||||
def create_caller(self, cls: Any, object_id: str):
|
||||
methods = self._services.get(object_id) or self._services.get(getattr(cls, "__name__", object_id))
|
||||
if methods is None:
|
||||
raise KeyError(object_id)
|
||||
return FakeExactRelayCaller(methods, self.transcripts, object_id)
|
||||
|
||||
|
||||
def capture_exact_small_proxy_relay() -> dict[str, object]:
|
||||
reset_forbidden_singleton_modules()
|
||||
fake_rpc = FakeExactRelayRPC()
|
||||
previous_child = os.environ.get("PYISOLATE_CHILD")
|
||||
previous_import_torch = os.environ.get("PYISOLATE_IMPORT_TORCH")
|
||||
try:
|
||||
prepare_sealed_singleton_proxies(fake_rpc)
|
||||
|
||||
from comfy.isolation.proxies.folder_paths_proxy import FolderPathsProxy
|
||||
from comfy.isolation.proxies.helper_proxies import restore_input_types
|
||||
from comfy.isolation.proxies.progress_proxy import ProgressProxy
|
||||
from comfy.isolation.proxies.utils_proxy import UtilsProxy
|
||||
|
||||
folder_proxy = FolderPathsProxy()
|
||||
utils_proxy = UtilsProxy()
|
||||
progress_proxy = ProgressProxy()
|
||||
before = set(sys.modules)
|
||||
|
||||
restored = restore_input_types(
|
||||
{
|
||||
"required": {
|
||||
"image": {"__pyisolate_any_type__": True, "value": "*"},
|
||||
}
|
||||
}
|
||||
)
|
||||
folder_path = folder_proxy.get_annotated_filepath("demo.png[input]")
|
||||
models_dir = folder_proxy.models_dir
|
||||
folder_names_and_paths = folder_proxy.folder_names_and_paths
|
||||
asyncio.run(utils_proxy.progress_bar_hook(2, 5, node_id="node-17"))
|
||||
progress_proxy.set_progress(1.5, 5.0, node_id="node-17")
|
||||
|
||||
imported = set(sys.modules) - before
|
||||
return {
|
||||
"mode": "exact_small_proxy_relay",
|
||||
"folder_path": folder_path,
|
||||
"models_dir": models_dir,
|
||||
"folder_names_and_paths": folder_names_and_paths,
|
||||
"restored_any_type": str(restored["required"]["image"]),
|
||||
"transcripts": fake_rpc.transcripts,
|
||||
"modules": sorted(imported),
|
||||
"forbidden_matches": matching_modules(FORBIDDEN_EXACT_SMALL_PROXY_MODULES, imported),
|
||||
}
|
||||
finally:
|
||||
_clear_proxy_rpcs()
|
||||
if previous_child is None:
|
||||
os.environ.pop("PYISOLATE_CHILD", None)
|
||||
else:
|
||||
os.environ["PYISOLATE_CHILD"] = previous_child
|
||||
if previous_import_torch is None:
|
||||
os.environ.pop("PYISOLATE_IMPORT_TORCH", None)
|
||||
else:
|
||||
os.environ["PYISOLATE_IMPORT_TORCH"] = previous_import_torch
|
||||
|
||||
|
||||
class FakeModelManagementExactRelayRPC:
|
||||
def __init__(self) -> None:
|
||||
self.transcripts: list[dict[str, object]] = []
|
||||
self._device = {"__pyisolate_torch_device__": "cpu"}
|
||||
self._services: dict[str, dict[str, Any]] = {
|
||||
"ModelManagementProxy": {
|
||||
"rpc_call": self._rpc_call,
|
||||
}
|
||||
}
|
||||
|
||||
def create_caller(self, cls: Any, object_id: str):
|
||||
methods = self._services.get(object_id) or self._services.get(getattr(cls, "__name__", object_id))
|
||||
if methods is None:
|
||||
raise KeyError(object_id)
|
||||
return _ModelManagementExactRelayCaller(methods)
|
||||
|
||||
def _rpc_call(self, method_name: str, args: Any, kwargs: Any) -> Any:
|
||||
self.transcripts.append(
|
||||
{
|
||||
"phase": "child_call",
|
||||
"object_id": "ModelManagementProxy",
|
||||
"method": method_name,
|
||||
"args": _json_safe(args),
|
||||
"kwargs": _json_safe(kwargs),
|
||||
}
|
||||
)
|
||||
target = f"comfy.model_management.{method_name}"
|
||||
self.transcripts.append(
|
||||
{
|
||||
"phase": "host_invocation",
|
||||
"object_id": "ModelManagementProxy",
|
||||
"method": method_name,
|
||||
"target": target,
|
||||
"args": _json_safe(args),
|
||||
"kwargs": _json_safe(kwargs),
|
||||
}
|
||||
)
|
||||
if method_name == "get_torch_device":
|
||||
result = self._device
|
||||
elif method_name == "get_torch_device_name":
|
||||
result = "cpu"
|
||||
elif method_name == "get_free_memory":
|
||||
result = 34359738368
|
||||
else:
|
||||
raise AssertionError(f"unexpected exact-relay method {method_name}")
|
||||
self.transcripts.append(
|
||||
{
|
||||
"phase": "result",
|
||||
"object_id": "ModelManagementProxy",
|
||||
"method": method_name,
|
||||
"result": _json_safe(result),
|
||||
}
|
||||
)
|
||||
return result
|
||||
|
||||
|
||||
class _ModelManagementExactRelayCaller:
|
||||
def __init__(self, methods: dict[str, Any]):
|
||||
self._methods = methods
|
||||
|
||||
def __getattr__(self, name: str):
|
||||
if name not in self._methods:
|
||||
raise AttributeError(name)
|
||||
|
||||
async def method(*args: Any, **kwargs: Any) -> Any:
|
||||
impl = self._methods[name]
|
||||
return impl(*args, **kwargs) if callable(impl) else impl
|
||||
|
||||
return method
|
||||
|
||||
|
||||
def _json_safe(value: Any) -> Any:
|
||||
if callable(value):
|
||||
return f"<callable {getattr(value, '__name__', 'anonymous')}>"
|
||||
if isinstance(value, tuple):
|
||||
return [_json_safe(item) for item in value]
|
||||
if isinstance(value, list):
|
||||
return [_json_safe(item) for item in value]
|
||||
if isinstance(value, dict):
|
||||
return {key: _json_safe(inner) for key, inner in value.items()}
|
||||
return value
|
||||
|
||||
|
||||
def capture_model_management_exact_relay() -> dict[str, object]:
|
||||
for module_name in FORBIDDEN_MODEL_MANAGEMENT_MODULES:
|
||||
sys.modules.pop(module_name, None)
|
||||
|
||||
fake_rpc = FakeModelManagementExactRelayRPC()
|
||||
previous_child = os.environ.get("PYISOLATE_CHILD")
|
||||
previous_import_torch = os.environ.get("PYISOLATE_IMPORT_TORCH")
|
||||
try:
|
||||
os.environ["PYISOLATE_CHILD"] = "1"
|
||||
os.environ["PYISOLATE_IMPORT_TORCH"] = "0"
|
||||
|
||||
from comfy.isolation.proxies.model_management_proxy import ModelManagementProxy
|
||||
|
||||
if hasattr(ModelManagementProxy, "clear_rpc"):
|
||||
ModelManagementProxy.clear_rpc()
|
||||
if hasattr(ModelManagementProxy, "set_rpc"):
|
||||
ModelManagementProxy.set_rpc(fake_rpc)
|
||||
|
||||
proxy = ModelManagementProxy()
|
||||
before = set(sys.modules)
|
||||
device = proxy.get_torch_device()
|
||||
device_name = proxy.get_torch_device_name(device)
|
||||
free_memory = proxy.get_free_memory(device)
|
||||
imported = set(sys.modules) - before
|
||||
return {
|
||||
"mode": "model_management_exact_relay",
|
||||
"device": str(device),
|
||||
"device_type": getattr(device, "type", None),
|
||||
"device_name": device_name,
|
||||
"free_memory": free_memory,
|
||||
"transcripts": fake_rpc.transcripts,
|
||||
"modules": sorted(imported),
|
||||
"forbidden_matches": matching_modules(FORBIDDEN_MODEL_MANAGEMENT_MODULES, imported),
|
||||
}
|
||||
finally:
|
||||
model_management_proxy = _load_model_management_proxy()
|
||||
if model_management_proxy is not None and hasattr(model_management_proxy, "clear_rpc"):
|
||||
model_management_proxy.clear_rpc()
|
||||
if previous_child is None:
|
||||
os.environ.pop("PYISOLATE_CHILD", None)
|
||||
else:
|
||||
os.environ["PYISOLATE_CHILD"] = previous_child
|
||||
if previous_import_torch is None:
|
||||
os.environ.pop("PYISOLATE_IMPORT_TORCH", None)
|
||||
else:
|
||||
os.environ["PYISOLATE_IMPORT_TORCH"] = previous_import_torch
|
||||
|
||||
|
||||
FORBIDDEN_PROMPT_WEB_MODULES = (
|
||||
"server",
|
||||
"aiohttp",
|
||||
"comfy.isolation.extension_wrapper",
|
||||
)
|
||||
FORBIDDEN_EXACT_BOOTSTRAP_MODULES = (
|
||||
"comfy.isolation.adapter",
|
||||
"folder_paths",
|
||||
"comfy.utils",
|
||||
"comfy.model_management",
|
||||
"server",
|
||||
"main",
|
||||
"comfy.isolation.extension_wrapper",
|
||||
)
|
||||
|
||||
|
||||
class _PromptServiceExactRelayCaller:
|
||||
def __init__(self, methods: dict[str, Any], transcripts: list[dict[str, Any]], object_id: str):
|
||||
self._methods = methods
|
||||
self._transcripts = transcripts
|
||||
self._object_id = object_id
|
||||
|
||||
def __getattr__(self, name: str):
|
||||
if name not in self._methods:
|
||||
raise AttributeError(name)
|
||||
|
||||
async def method(*args: Any, **kwargs: Any) -> Any:
|
||||
self._transcripts.append(
|
||||
{
|
||||
"phase": "child_call",
|
||||
"object_id": self._object_id,
|
||||
"method": name,
|
||||
"args": _json_safe(args),
|
||||
"kwargs": _json_safe(kwargs),
|
||||
}
|
||||
)
|
||||
impl = self._methods[name]
|
||||
self._transcripts.append(
|
||||
{
|
||||
"phase": "host_invocation",
|
||||
"object_id": self._object_id,
|
||||
"method": name,
|
||||
"target": impl["target"],
|
||||
"args": _json_safe(args),
|
||||
"kwargs": _json_safe(kwargs),
|
||||
}
|
||||
)
|
||||
result = impl["result"](*args, **kwargs) if callable(impl["result"]) else impl["result"]
|
||||
self._transcripts.append(
|
||||
{
|
||||
"phase": "result",
|
||||
"object_id": self._object_id,
|
||||
"method": name,
|
||||
"result": _json_safe(result),
|
||||
}
|
||||
)
|
||||
return result
|
||||
|
||||
return method
|
||||
|
||||
|
||||
class FakePromptWebRPC:
|
||||
def __init__(self) -> None:
|
||||
self.transcripts: list[dict[str, Any]] = []
|
||||
self._services = {
|
||||
"PromptServerService": {
|
||||
"ui_send_progress_text": {
|
||||
"target": "server.PromptServer.instance.send_progress_text",
|
||||
"result": None,
|
||||
},
|
||||
"register_route_rpc": {
|
||||
"target": "server.PromptServer.instance.routes.add_route",
|
||||
"result": None,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
def create_caller(self, cls: Any, object_id: str):
|
||||
methods = self._services.get(object_id) or self._services.get(getattr(cls, "__name__", object_id))
|
||||
if methods is None:
|
||||
raise KeyError(object_id)
|
||||
return _PromptServiceExactRelayCaller(methods, self.transcripts, object_id)
|
||||
|
||||
|
||||
class FakeWebDirectoryProxy:
|
||||
def __init__(self, transcripts: list[dict[str, Any]]):
|
||||
self._transcripts = transcripts
|
||||
|
||||
def get_web_file(self, extension_name: str, relative_path: str) -> dict[str, Any]:
|
||||
self._transcripts.append(
|
||||
{
|
||||
"phase": "child_call",
|
||||
"object_id": "WebDirectoryProxy",
|
||||
"method": "get_web_file",
|
||||
"args": [extension_name, relative_path],
|
||||
"kwargs": {},
|
||||
}
|
||||
)
|
||||
self._transcripts.append(
|
||||
{
|
||||
"phase": "host_invocation",
|
||||
"object_id": "WebDirectoryProxy",
|
||||
"method": "get_web_file",
|
||||
"target": "comfy.isolation.proxies.web_directory_proxy.WebDirectoryProxy.get_web_file",
|
||||
"args": [extension_name, relative_path],
|
||||
"kwargs": {},
|
||||
}
|
||||
)
|
||||
result = {
|
||||
"content": "Y29uc29sZS5sb2coJ2RlbycpOw==",
|
||||
"content_type": "application/javascript",
|
||||
}
|
||||
self._transcripts.append(
|
||||
{
|
||||
"phase": "result",
|
||||
"object_id": "WebDirectoryProxy",
|
||||
"method": "get_web_file",
|
||||
"result": result,
|
||||
}
|
||||
)
|
||||
return result
|
||||
|
||||
|
||||
def capture_prompt_web_exact_relay() -> dict[str, object]:
|
||||
for module_name in FORBIDDEN_PROMPT_WEB_MODULES:
|
||||
sys.modules.pop(module_name, None)
|
||||
|
||||
fake_rpc = FakePromptWebRPC()
|
||||
|
||||
from comfy.isolation.proxies.prompt_server_impl import PromptServerStub
|
||||
from comfy.isolation.proxies.web_directory_proxy import WebDirectoryCache
|
||||
|
||||
PromptServerStub.set_rpc(fake_rpc)
|
||||
PromptServerStub._pending_child_routes = []
|
||||
stub = PromptServerStub()
|
||||
cache = WebDirectoryCache()
|
||||
cache.register_proxy("demo_ext", FakeWebDirectoryProxy(fake_rpc.transcripts))
|
||||
|
||||
before = set(sys.modules)
|
||||
|
||||
def demo_handler(_request):
|
||||
return {"ok": True}
|
||||
|
||||
stub.send_progress_text("hello", "node-17")
|
||||
stub.routes.get("/demo")(demo_handler)
|
||||
asyncio.run(PromptServerStub.flush_child_routes())
|
||||
web_file = cache.get_file("demo_ext", "js/app.js")
|
||||
imported = set(sys.modules) - before
|
||||
return {
|
||||
"mode": "prompt_web_exact_relay",
|
||||
"web_file": {
|
||||
"content_type": web_file["content_type"] if web_file else None,
|
||||
"content": web_file["content"].decode("utf-8") if web_file else None,
|
||||
},
|
||||
"transcripts": fake_rpc.transcripts,
|
||||
"modules": sorted(imported),
|
||||
"forbidden_matches": matching_modules(FORBIDDEN_PROMPT_WEB_MODULES, imported),
|
||||
}
|
||||
|
||||
|
||||
class FakeExactBootstrapRPC:
|
||||
def __init__(self) -> None:
|
||||
self.transcripts: list[dict[str, Any]] = []
|
||||
self._device = {"__pyisolate_torch_device__": "cpu"}
|
||||
self._services: dict[str, dict[str, Any]] = {
|
||||
"FolderPathsProxy": FakeExactRelayRPC()._services["FolderPathsProxy"],
|
||||
"HelperProxiesService": FakeExactRelayRPC()._services["HelperProxiesService"],
|
||||
"ProgressProxy": FakeExactRelayRPC()._services["ProgressProxy"],
|
||||
"UtilsProxy": FakeExactRelayRPC()._services["UtilsProxy"],
|
||||
"PromptServerService": {
|
||||
"ui_send_sync": {
|
||||
"target": "server.PromptServer.instance.send_sync",
|
||||
"result": None,
|
||||
},
|
||||
"ui_send": {
|
||||
"target": "server.PromptServer.instance.send",
|
||||
"result": None,
|
||||
},
|
||||
"ui_send_progress_text": {
|
||||
"target": "server.PromptServer.instance.send_progress_text",
|
||||
"result": None,
|
||||
},
|
||||
"register_route_rpc": {
|
||||
"target": "server.PromptServer.instance.routes.add_route",
|
||||
"result": None,
|
||||
},
|
||||
},
|
||||
"ModelManagementProxy": {
|
||||
"rpc_call": self._rpc_call,
|
||||
},
|
||||
}
|
||||
|
||||
def create_caller(self, cls: Any, object_id: str):
|
||||
methods = self._services.get(object_id) or self._services.get(getattr(cls, "__name__", object_id))
|
||||
if methods is None:
|
||||
raise KeyError(object_id)
|
||||
if object_id == "ModelManagementProxy":
|
||||
return _ModelManagementExactRelayCaller(methods)
|
||||
return _PromptServiceExactRelayCaller(methods, self.transcripts, object_id)
|
||||
|
||||
def _rpc_call(self, method_name: str, args: Any, kwargs: Any) -> Any:
|
||||
self.transcripts.append(
|
||||
{
|
||||
"phase": "child_call",
|
||||
"object_id": "ModelManagementProxy",
|
||||
"method": method_name,
|
||||
"args": _json_safe(args),
|
||||
"kwargs": _json_safe(kwargs),
|
||||
}
|
||||
)
|
||||
self.transcripts.append(
|
||||
{
|
||||
"phase": "host_invocation",
|
||||
"object_id": "ModelManagementProxy",
|
||||
"method": method_name,
|
||||
"target": f"comfy.model_management.{method_name}",
|
||||
"args": _json_safe(args),
|
||||
"kwargs": _json_safe(kwargs),
|
||||
}
|
||||
)
|
||||
result = self._device if method_name == "get_torch_device" else None
|
||||
self.transcripts.append(
|
||||
{
|
||||
"phase": "result",
|
||||
"object_id": "ModelManagementProxy",
|
||||
"method": method_name,
|
||||
"result": _json_safe(result),
|
||||
}
|
||||
)
|
||||
return result
|
||||
|
||||
|
||||
def capture_exact_proxy_bootstrap_contract() -> dict[str, object]:
|
||||
from pyisolate._internal.rpc_protocol import get_child_rpc_instance, set_child_rpc_instance
|
||||
|
||||
from comfy.isolation.adapter import ComfyUIAdapter
|
||||
from comfy.isolation.child_hooks import initialize_child_process
|
||||
from comfy.isolation.proxies.folder_paths_proxy import FolderPathsProxy
|
||||
from comfy.isolation.proxies.helper_proxies import HelperProxiesService
|
||||
from comfy.isolation.proxies.model_management_proxy import ModelManagementProxy
|
||||
from comfy.isolation.proxies.progress_proxy import ProgressProxy
|
||||
from comfy.isolation.proxies.prompt_server_impl import PromptServerStub
|
||||
from comfy.isolation.proxies.utils_proxy import UtilsProxy
|
||||
|
||||
host_services = sorted(cls.__name__ for cls in ComfyUIAdapter().provide_rpc_services())
|
||||
|
||||
for module_name in FORBIDDEN_EXACT_BOOTSTRAP_MODULES:
|
||||
sys.modules.pop(module_name, None)
|
||||
|
||||
previous_child = os.environ.get("PYISOLATE_CHILD")
|
||||
previous_import_torch = os.environ.get("PYISOLATE_IMPORT_TORCH")
|
||||
os.environ["PYISOLATE_CHILD"] = "1"
|
||||
os.environ["PYISOLATE_IMPORT_TORCH"] = "0"
|
||||
|
||||
_clear_proxy_rpcs()
|
||||
if hasattr(PromptServerStub, "clear_rpc"):
|
||||
PromptServerStub.clear_rpc()
|
||||
else:
|
||||
PromptServerStub._rpc = None # type: ignore[attr-defined]
|
||||
fake_rpc = FakeExactBootstrapRPC()
|
||||
set_child_rpc_instance(fake_rpc)
|
||||
|
||||
before = set(sys.modules)
|
||||
try:
|
||||
initialize_child_process()
|
||||
imported = set(sys.modules) - before
|
||||
matrix = {
|
||||
"base.py": {
|
||||
"bound": get_child_rpc_instance() is fake_rpc,
|
||||
"details": {"child_rpc_instance": get_child_rpc_instance() is fake_rpc},
|
||||
},
|
||||
"folder_paths_proxy.py": {
|
||||
"bound": "FolderPathsProxy" in host_services and FolderPathsProxy._rpc is not None,
|
||||
"details": {"host_service": "FolderPathsProxy" in host_services, "child_rpc": FolderPathsProxy._rpc is not None},
|
||||
},
|
||||
"helper_proxies.py": {
|
||||
"bound": "HelperProxiesService" in host_services and HelperProxiesService._rpc is not None,
|
||||
"details": {"host_service": "HelperProxiesService" in host_services, "child_rpc": HelperProxiesService._rpc is not None},
|
||||
},
|
||||
"model_management_proxy.py": {
|
||||
"bound": "ModelManagementProxy" in host_services and ModelManagementProxy._rpc is not None,
|
||||
"details": {"host_service": "ModelManagementProxy" in host_services, "child_rpc": ModelManagementProxy._rpc is not None},
|
||||
},
|
||||
"progress_proxy.py": {
|
||||
"bound": "ProgressProxy" in host_services and ProgressProxy._rpc is not None,
|
||||
"details": {"host_service": "ProgressProxy" in host_services, "child_rpc": ProgressProxy._rpc is not None},
|
||||
},
|
||||
"prompt_server_impl.py": {
|
||||
"bound": "PromptServerService" in host_services and PromptServerStub._rpc is not None,
|
||||
"details": {"host_service": "PromptServerService" in host_services, "child_rpc": PromptServerStub._rpc is not None},
|
||||
},
|
||||
"utils_proxy.py": {
|
||||
"bound": "UtilsProxy" in host_services and UtilsProxy._rpc is not None,
|
||||
"details": {"host_service": "UtilsProxy" in host_services, "child_rpc": UtilsProxy._rpc is not None},
|
||||
},
|
||||
"web_directory_proxy.py": {
|
||||
"bound": "WebDirectoryProxy" in host_services,
|
||||
"details": {"host_service": "WebDirectoryProxy" in host_services},
|
||||
},
|
||||
}
|
||||
finally:
|
||||
set_child_rpc_instance(None)
|
||||
if previous_child is None:
|
||||
os.environ.pop("PYISOLATE_CHILD", None)
|
||||
else:
|
||||
os.environ["PYISOLATE_CHILD"] = previous_child
|
||||
if previous_import_torch is None:
|
||||
os.environ.pop("PYISOLATE_IMPORT_TORCH", None)
|
||||
else:
|
||||
os.environ["PYISOLATE_IMPORT_TORCH"] = previous_import_torch
|
||||
|
||||
omitted = sorted(name for name, status in matrix.items() if not status["bound"])
|
||||
return {
|
||||
"mode": "exact_proxy_bootstrap_contract",
|
||||
"host_services": host_services,
|
||||
"matrix": matrix,
|
||||
"omitted_proxies": omitted,
|
||||
"modules": sorted(imported),
|
||||
"forbidden_matches": matching_modules(FORBIDDEN_EXACT_BOOTSTRAP_MODULES, imported),
|
||||
}
|
||||
|
||||
def capture_sealed_singleton_imports() -> dict[str, object]:
|
||||
reset_forbidden_singleton_modules()
|
||||
fake_rpc = FakeSingletonRPC()
|
||||
previous_child = os.environ.get("PYISOLATE_CHILD")
|
||||
previous_import_torch = os.environ.get("PYISOLATE_IMPORT_TORCH")
|
||||
before = set(sys.modules)
|
||||
try:
|
||||
prepare_sealed_singleton_proxies(fake_rpc)
|
||||
|
||||
from comfy.isolation.proxies.folder_paths_proxy import FolderPathsProxy
|
||||
from comfy.isolation.proxies.progress_proxy import ProgressProxy
|
||||
from comfy.isolation.proxies.utils_proxy import UtilsProxy
|
||||
|
||||
folder_proxy = FolderPathsProxy()
|
||||
progress_proxy = ProgressProxy()
|
||||
utils_proxy = UtilsProxy()
|
||||
|
||||
folder_path = folder_proxy.get_annotated_filepath("demo.png[input]")
|
||||
temp_dir = folder_proxy.get_temp_directory()
|
||||
models_dir = folder_proxy.models_dir
|
||||
asyncio.run(utils_proxy.progress_bar_hook(2, 5, node_id="node-17"))
|
||||
progress_proxy.set_progress(1.5, 5.0, node_id="node-17")
|
||||
|
||||
imported = set(sys.modules) - before
|
||||
return {
|
||||
"mode": "sealed_singletons",
|
||||
"folder_path": folder_path,
|
||||
"temp_dir": temp_dir,
|
||||
"models_dir": models_dir,
|
||||
"rpc_calls": fake_rpc.calls,
|
||||
"modules": sorted(imported),
|
||||
"forbidden_matches": matching_modules(FORBIDDEN_SEALED_SINGLETON_MODULES, imported),
|
||||
}
|
||||
finally:
|
||||
_clear_proxy_rpcs()
|
||||
if previous_child is None:
|
||||
os.environ.pop("PYISOLATE_CHILD", None)
|
||||
else:
|
||||
os.environ["PYISOLATE_CHILD"] = previous_child
|
||||
if previous_import_torch is None:
|
||||
os.environ.pop("PYISOLATE_IMPORT_TORCH", None)
|
||||
else:
|
||||
os.environ["PYISOLATE_IMPORT_TORCH"] = previous_import_torch
|
||||
@@ -1,69 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import shutil
|
||||
import sys
|
||||
import tempfile
|
||||
from contextlib import contextmanager
|
||||
from pathlib import Path
|
||||
from typing import Iterator
|
||||
|
||||
|
||||
COMFYUI_ROOT = Path(__file__).resolve().parents[2]
|
||||
PROBE_SOURCE_ROOT = COMFYUI_ROOT / "tests" / "isolation" / "internal_probe_node"
|
||||
PROBE_NODE_NAME = "InternalIsolationProbeNode"
|
||||
|
||||
PYPROJECT_CONTENT = """[project]
|
||||
name = "InternalIsolationProbeNode"
|
||||
version = "0.0.1"
|
||||
|
||||
[tool.comfy.isolation]
|
||||
can_isolate = true
|
||||
share_torch = true
|
||||
"""
|
||||
|
||||
|
||||
def _probe_target_root(comfy_root: Path) -> Path:
|
||||
return Path(comfy_root) / "custom_nodes" / PROBE_NODE_NAME
|
||||
|
||||
|
||||
def stage_probe_node(comfy_root: Path) -> Path:
|
||||
if not PROBE_SOURCE_ROOT.is_dir():
|
||||
raise RuntimeError(f"Missing probe source directory: {PROBE_SOURCE_ROOT}")
|
||||
|
||||
target_root = _probe_target_root(comfy_root)
|
||||
target_root.mkdir(parents=True, exist_ok=True)
|
||||
for source_path in PROBE_SOURCE_ROOT.iterdir():
|
||||
destination_path = target_root / source_path.name
|
||||
if source_path.is_dir():
|
||||
shutil.copytree(source_path, destination_path, dirs_exist_ok=True)
|
||||
else:
|
||||
shutil.copy2(source_path, destination_path)
|
||||
|
||||
(target_root / "pyproject.toml").write_text(PYPROJECT_CONTENT, encoding="utf-8")
|
||||
return target_root
|
||||
|
||||
|
||||
@contextmanager
|
||||
def staged_probe_node() -> Iterator[Path]:
|
||||
staging_root = Path(tempfile.mkdtemp(prefix="comfyui_internal_probe_"))
|
||||
try:
|
||||
yield stage_probe_node(staging_root)
|
||||
finally:
|
||||
shutil.rmtree(staging_root, ignore_errors=True)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Stage the internal isolation probe node under an explicit ComfyUI root."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--target-root",
|
||||
type=Path,
|
||||
required=True,
|
||||
help="Explicit ComfyUI root to stage under. Caller owns cleanup.",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
staged = stage_probe_node(args.target_root)
|
||||
sys.stdout.write(f"{staged}\n")
|
||||
@@ -1,122 +0,0 @@
|
||||
"""Tests for pyisolate._internal.client import-time snapshot handling."""
|
||||
|
||||
import json
|
||||
import os
|
||||
import subprocess
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
# Paths needed for subprocess
|
||||
PYISOLATE_ROOT = str(Path(__file__).parent.parent)
|
||||
COMFYUI_ROOT = os.environ.get("COMFYUI_ROOT") or str(Path.home() / "ComfyUI")
|
||||
|
||||
SCRIPT = """
|
||||
import json, sys
|
||||
import pyisolate._internal.client # noqa: F401 # triggers snapshot logic
|
||||
print(json.dumps(sys.path[:6]))
|
||||
"""
|
||||
|
||||
|
||||
def _run_client_process(env):
|
||||
# Ensure subprocess can find pyisolate and ComfyUI
|
||||
pythonpath_parts = [PYISOLATE_ROOT, COMFYUI_ROOT]
|
||||
existing = env.get("PYTHONPATH", "")
|
||||
if existing:
|
||||
pythonpath_parts.append(existing)
|
||||
env["PYTHONPATH"] = os.pathsep.join(pythonpath_parts)
|
||||
|
||||
result = subprocess.run( # noqa: S603
|
||||
[sys.executable, "-c", SCRIPT],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
env=env,
|
||||
check=True,
|
||||
)
|
||||
stdout = result.stdout.strip().splitlines()[-1]
|
||||
return json.loads(stdout)
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def comfy_module_path(tmp_path):
|
||||
comfy_root = tmp_path / "ComfyUI"
|
||||
module_path = comfy_root / "custom_nodes" / "TestNode"
|
||||
module_path.mkdir(parents=True)
|
||||
return comfy_root, module_path
|
||||
|
||||
|
||||
def test_snapshot_applied_and_comfy_root_prepend(tmp_path, comfy_module_path):
|
||||
comfy_root, module_path = comfy_module_path
|
||||
# Must include real ComfyUI path for utils validation to pass
|
||||
host_paths = [COMFYUI_ROOT, "/host/lib1", "/host/lib2"]
|
||||
snapshot = {
|
||||
"sys_path": host_paths,
|
||||
"sys_executable": sys.executable,
|
||||
"sys_prefix": sys.prefix,
|
||||
"environment": {},
|
||||
}
|
||||
snapshot_path = tmp_path / "snapshot.json"
|
||||
snapshot_path.write_text(json.dumps(snapshot), encoding="utf-8")
|
||||
|
||||
env = os.environ.copy()
|
||||
env.update(
|
||||
{
|
||||
"PYISOLATE_CHILD": "1",
|
||||
"PYISOLATE_HOST_SNAPSHOT": str(snapshot_path),
|
||||
"PYISOLATE_MODULE_PATH": str(module_path),
|
||||
}
|
||||
)
|
||||
|
||||
path_prefix = _run_client_process(env)
|
||||
|
||||
# Current client behavior preserves the runtime bootstrap path order and
|
||||
# keeps the resolved ComfyUI root available for imports.
|
||||
assert COMFYUI_ROOT in path_prefix
|
||||
# Module path should not override runtime root selection.
|
||||
assert str(comfy_root) not in path_prefix
|
||||
|
||||
|
||||
def test_missing_snapshot_file_does_not_crash(tmp_path, comfy_module_path):
|
||||
_, module_path = comfy_module_path
|
||||
missing_snapshot = tmp_path / "missing.json"
|
||||
|
||||
env = os.environ.copy()
|
||||
env.update(
|
||||
{
|
||||
"PYISOLATE_CHILD": "1",
|
||||
"PYISOLATE_HOST_SNAPSHOT": str(missing_snapshot),
|
||||
"PYISOLATE_MODULE_PATH": str(module_path),
|
||||
}
|
||||
)
|
||||
|
||||
# Should not raise even though snapshot path is missing
|
||||
paths = _run_client_process(env)
|
||||
assert len(paths) > 0
|
||||
|
||||
|
||||
def test_no_comfy_root_when_module_path_absent(tmp_path):
|
||||
# Must include real ComfyUI path for utils validation to pass
|
||||
host_paths = [COMFYUI_ROOT, "/alpha", "/beta"]
|
||||
snapshot = {
|
||||
"sys_path": host_paths,
|
||||
"sys_executable": sys.executable,
|
||||
"sys_prefix": sys.prefix,
|
||||
"environment": {},
|
||||
}
|
||||
snapshot_path = tmp_path / "snapshot.json"
|
||||
snapshot_path.write_text(json.dumps(snapshot), encoding="utf-8")
|
||||
|
||||
env = os.environ.copy()
|
||||
env.update(
|
||||
{
|
||||
"PYISOLATE_CHILD": "1",
|
||||
"PYISOLATE_HOST_SNAPSHOT": str(snapshot_path),
|
||||
}
|
||||
)
|
||||
|
||||
paths = _run_client_process(env)
|
||||
# Runtime path bootstrap keeps ComfyUI importability regardless of host
|
||||
# snapshot extras.
|
||||
assert COMFYUI_ROOT in paths
|
||||
assert "/alpha" not in paths and "/beta" not in paths
|
||||
@@ -1,637 +0,0 @@
|
||||
"""Synthetic integration coverage for manifest plumbing and env flags.
|
||||
|
||||
These tests do not perform a real wheel install or a real ComfyUI E2E run.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
import os
|
||||
import sys
|
||||
from types import SimpleNamespace
|
||||
from typing import Any, cast
|
||||
|
||||
import pytest
|
||||
|
||||
import comfy.isolation as isolation_pkg
|
||||
from comfy.isolation import runtime_helpers
|
||||
from comfy.isolation import extension_loader as extension_loader_module
|
||||
from comfy.isolation import extension_wrapper as extension_wrapper_module
|
||||
from comfy.isolation import model_patcher_proxy_utils
|
||||
from comfy.isolation.extension_loader import ExtensionLoadError, load_isolated_node
|
||||
from comfy.isolation.extension_wrapper import ComfyNodeExtension
|
||||
from comfy.isolation.model_patcher_proxy_utils import maybe_wrap_model_for_isolation
|
||||
from pyisolate._internal.environment_conda import _generate_pixi_toml
|
||||
|
||||
|
||||
class _DummyExtension:
|
||||
def __init__(self) -> None:
|
||||
self.name = "demo-extension"
|
||||
|
||||
async def stop(self) -> None:
|
||||
return None
|
||||
|
||||
|
||||
def _write_manifest(node_dir, manifest_text: str) -> None:
|
||||
(node_dir / "pyproject.toml").write_text(manifest_text, encoding="utf-8")
|
||||
|
||||
|
||||
def test_load_isolated_node_passes_normalized_cuda_wheels_config(tmp_path, monkeypatch):
|
||||
node_dir = tmp_path / "node"
|
||||
node_dir.mkdir()
|
||||
manifest_path = node_dir / "pyproject.toml"
|
||||
_write_manifest(
|
||||
node_dir,
|
||||
"""
|
||||
[project]
|
||||
name = "demo-node"
|
||||
dependencies = ["flash-attn>=1.0", "sageattention==0.1"]
|
||||
|
||||
[tool.comfy.isolation]
|
||||
can_isolate = true
|
||||
share_torch = true
|
||||
|
||||
[tool.comfy.isolation.cuda_wheels]
|
||||
index_url = "https://example.invalid/cuda-wheels"
|
||||
packages = ["flash_attn", "sageattention"]
|
||||
|
||||
[tool.comfy.isolation.cuda_wheels.package_map]
|
||||
flash_attn = "flash-attn-special"
|
||||
""".strip(),
|
||||
)
|
||||
|
||||
captured: dict[str, object] = {}
|
||||
|
||||
class DummyManager:
|
||||
def __init__(self, *args, **kwargs) -> None:
|
||||
return None
|
||||
|
||||
def load_extension(self, config):
|
||||
captured.update(config)
|
||||
return _DummyExtension()
|
||||
|
||||
monkeypatch.setattr(extension_loader_module.pyisolate, "ExtensionManager", DummyManager)
|
||||
monkeypatch.setattr(
|
||||
extension_loader_module,
|
||||
"load_host_policy",
|
||||
lambda base_path: {
|
||||
"sandbox_mode": "required",
|
||||
"allow_network": False,
|
||||
"writable_paths": [],
|
||||
"readonly_paths": [],
|
||||
},
|
||||
)
|
||||
monkeypatch.setattr(extension_loader_module, "is_cache_valid", lambda *args, **kwargs: True)
|
||||
monkeypatch.setattr(
|
||||
extension_loader_module,
|
||||
"load_from_cache",
|
||||
lambda *args, **kwargs: {"Node": {"display_name": "Node", "schema_v1": {}}},
|
||||
)
|
||||
monkeypatch.setitem(sys.modules, "folder_paths", SimpleNamespace(base_path=str(tmp_path)))
|
||||
|
||||
specs = asyncio.run(
|
||||
load_isolated_node(
|
||||
node_dir,
|
||||
manifest_path,
|
||||
logging.getLogger("test"),
|
||||
lambda *args, **kwargs: object,
|
||||
tmp_path / "venvs",
|
||||
[],
|
||||
)
|
||||
)
|
||||
|
||||
assert len(specs) == 1
|
||||
assert captured["sandbox_mode"] == "required"
|
||||
assert captured["cuda_wheels"] == {
|
||||
"index_url": "https://example.invalid/cuda-wheels/",
|
||||
"packages": ["flash-attn", "sageattention"],
|
||||
"package_map": {"flash-attn": "flash-attn-special"},
|
||||
}
|
||||
|
||||
|
||||
def test_load_isolated_node_passes_share_torch_no_deps(tmp_path, monkeypatch):
|
||||
node_dir = tmp_path / "node"
|
||||
node_dir.mkdir()
|
||||
manifest_path = node_dir / "pyproject.toml"
|
||||
_write_manifest(
|
||||
node_dir,
|
||||
"""
|
||||
[project]
|
||||
name = "demo-node"
|
||||
dependencies = ["timm", "pyyaml"]
|
||||
|
||||
[tool.comfy.isolation]
|
||||
can_isolate = true
|
||||
share_torch = true
|
||||
share_torch_no_deps = ["timm"]
|
||||
""".strip(),
|
||||
)
|
||||
|
||||
captured: dict[str, object] = {}
|
||||
|
||||
class DummyManager:
|
||||
def __init__(self, *args, **kwargs) -> None:
|
||||
return None
|
||||
|
||||
def load_extension(self, config):
|
||||
captured.update(config)
|
||||
return _DummyExtension()
|
||||
|
||||
monkeypatch.setattr(extension_loader_module.pyisolate, "ExtensionManager", DummyManager)
|
||||
monkeypatch.setattr(
|
||||
extension_loader_module,
|
||||
"load_host_policy",
|
||||
lambda base_path: {
|
||||
"sandbox_mode": "disabled",
|
||||
"allow_network": False,
|
||||
"writable_paths": [],
|
||||
"readonly_paths": [],
|
||||
},
|
||||
)
|
||||
monkeypatch.setattr(extension_loader_module, "is_cache_valid", lambda *args, **kwargs: True)
|
||||
monkeypatch.setattr(
|
||||
extension_loader_module,
|
||||
"load_from_cache",
|
||||
lambda *args, **kwargs: {"Node": {"display_name": "Node", "schema_v1": {}}},
|
||||
)
|
||||
monkeypatch.setitem(sys.modules, "folder_paths", SimpleNamespace(base_path=str(tmp_path)))
|
||||
|
||||
specs = asyncio.run(
|
||||
load_isolated_node(
|
||||
node_dir,
|
||||
manifest_path,
|
||||
logging.getLogger("test"),
|
||||
lambda *args, **kwargs: object,
|
||||
tmp_path / "venvs",
|
||||
[],
|
||||
)
|
||||
)
|
||||
|
||||
assert len(specs) == 1
|
||||
assert captured["share_torch_no_deps"] == ["timm"]
|
||||
|
||||
|
||||
def test_on_module_loaded_registers_legacy_routes(monkeypatch):
|
||||
captured: list[tuple[str, str, Any]] = []
|
||||
|
||||
def demo_handler(body):
|
||||
return body
|
||||
|
||||
module = SimpleNamespace(
|
||||
__file__="/tmp/demo_node/__init__.py",
|
||||
__name__="demo_node",
|
||||
NODE_CLASS_MAPPINGS={},
|
||||
NODE_DISPLAY_NAME_MAPPINGS={},
|
||||
ROUTES=[
|
||||
{"method": "POST", "path": "/sam3/interactive_segment_one", "handler": "demo_handler"},
|
||||
],
|
||||
demo_handler=demo_handler,
|
||||
)
|
||||
|
||||
def fake_register_route(self, method, path, handler):
|
||||
captured.append((method, path, handler))
|
||||
|
||||
monkeypatch.setattr(
|
||||
"comfy.isolation.proxies.prompt_server_impl.PromptServerStub.register_route",
|
||||
fake_register_route,
|
||||
)
|
||||
|
||||
extension = ComfyNodeExtension()
|
||||
asyncio.run(extension.on_module_loaded(module))
|
||||
|
||||
assert captured == [("POST", "/sam3/interactive_segment_one", demo_handler)]
|
||||
|
||||
|
||||
def test_prompt_server_stub_buffers_routes_without_rpc():
|
||||
from comfy.isolation.proxies.prompt_server_impl import PromptServerStub
|
||||
|
||||
def demo_handler(body):
|
||||
return body
|
||||
|
||||
old_rpc = PromptServerStub._rpc
|
||||
old_pending = list(PromptServerStub._pending_child_routes)
|
||||
try:
|
||||
PromptServerStub._rpc = None
|
||||
PromptServerStub._pending_child_routes = []
|
||||
PromptServerStub().register_route("POST", "/sam3/interactive_segment_one", demo_handler)
|
||||
assert PromptServerStub._pending_child_routes == [
|
||||
("POST", "/sam3/interactive_segment_one", demo_handler)
|
||||
]
|
||||
finally:
|
||||
PromptServerStub._rpc = old_rpc
|
||||
PromptServerStub._pending_child_routes = old_pending
|
||||
|
||||
|
||||
def test_load_isolated_node_rejects_undeclared_cuda_wheel_dependency(
|
||||
tmp_path, monkeypatch
|
||||
):
|
||||
node_dir = tmp_path / "node"
|
||||
node_dir.mkdir()
|
||||
manifest_path = node_dir / "pyproject.toml"
|
||||
_write_manifest(
|
||||
node_dir,
|
||||
"""
|
||||
[project]
|
||||
name = "demo-node"
|
||||
dependencies = ["numpy>=1.0"]
|
||||
|
||||
[tool.comfy.isolation]
|
||||
can_isolate = true
|
||||
|
||||
[tool.comfy.isolation.cuda_wheels]
|
||||
index_url = "https://example.invalid/cuda-wheels"
|
||||
packages = ["flash-attn"]
|
||||
""".strip(),
|
||||
)
|
||||
|
||||
monkeypatch.setitem(sys.modules, "folder_paths", SimpleNamespace(base_path=str(tmp_path)))
|
||||
|
||||
with pytest.raises(ExtensionLoadError, match="undeclared dependencies"):
|
||||
asyncio.run(
|
||||
load_isolated_node(
|
||||
node_dir,
|
||||
manifest_path,
|
||||
logging.getLogger("test"),
|
||||
lambda *args, **kwargs: object,
|
||||
tmp_path / "venvs",
|
||||
[],
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def test_conda_cuda_wheels_declared_packages_do_not_force_pixi_solve(tmp_path, monkeypatch):
|
||||
node_dir = tmp_path / "node"
|
||||
node_dir.mkdir()
|
||||
manifest_path = node_dir / "pyproject.toml"
|
||||
_write_manifest(
|
||||
node_dir,
|
||||
"""
|
||||
[project]
|
||||
name = "demo-node"
|
||||
dependencies = ["numpy>=1.0", "spconv", "cumm", "flash-attn"]
|
||||
|
||||
[tool.comfy.isolation]
|
||||
can_isolate = true
|
||||
package_manager = "conda"
|
||||
conda_channels = ["conda-forge"]
|
||||
|
||||
[tool.comfy.isolation.cuda_wheels]
|
||||
index_url = "https://example.invalid/cuda-wheels"
|
||||
packages = ["spconv", "cumm", "flash-attn"]
|
||||
""".strip(),
|
||||
)
|
||||
|
||||
captured: dict[str, object] = {}
|
||||
|
||||
class DummyManager:
|
||||
def __init__(self, *args, **kwargs) -> None:
|
||||
return None
|
||||
|
||||
def load_extension(self, config):
|
||||
captured.update(config)
|
||||
return _DummyExtension()
|
||||
|
||||
monkeypatch.setattr(extension_loader_module.pyisolate, "ExtensionManager", DummyManager)
|
||||
monkeypatch.setattr(
|
||||
extension_loader_module,
|
||||
"load_host_policy",
|
||||
lambda base_path: {
|
||||
"sandbox_mode": "disabled",
|
||||
"allow_network": False,
|
||||
"writable_paths": [],
|
||||
"readonly_paths": [],
|
||||
},
|
||||
)
|
||||
monkeypatch.setattr(extension_loader_module, "is_cache_valid", lambda *args, **kwargs: True)
|
||||
monkeypatch.setattr(
|
||||
extension_loader_module,
|
||||
"load_from_cache",
|
||||
lambda *args, **kwargs: {"Node": {"display_name": "Node", "schema_v1": {}}},
|
||||
)
|
||||
monkeypatch.setitem(sys.modules, "folder_paths", SimpleNamespace(base_path=str(tmp_path)))
|
||||
|
||||
asyncio.run(
|
||||
load_isolated_node(
|
||||
node_dir,
|
||||
manifest_path,
|
||||
logging.getLogger("test"),
|
||||
lambda *args, **kwargs: object,
|
||||
tmp_path / "venvs",
|
||||
[],
|
||||
)
|
||||
)
|
||||
|
||||
generated = _generate_pixi_toml(captured)
|
||||
assert 'numpy = ">=1.0"' in generated
|
||||
assert "spconv =" not in generated
|
||||
assert "cumm =" not in generated
|
||||
assert "flash-attn =" not in generated
|
||||
|
||||
|
||||
def test_conda_cuda_wheels_loader_accepts_sam3d_contract(tmp_path, monkeypatch):
|
||||
node_dir = tmp_path / "node"
|
||||
node_dir.mkdir()
|
||||
manifest_path = node_dir / "pyproject.toml"
|
||||
_write_manifest(
|
||||
node_dir,
|
||||
"""
|
||||
[project]
|
||||
name = "demo-node"
|
||||
dependencies = [
|
||||
"torch",
|
||||
"torchvision",
|
||||
"pytorch3d",
|
||||
"gsplat",
|
||||
"nvdiffrast",
|
||||
"flash-attn",
|
||||
"sageattention",
|
||||
"spconv",
|
||||
"cumm",
|
||||
]
|
||||
|
||||
[tool.comfy.isolation]
|
||||
can_isolate = true
|
||||
package_manager = "conda"
|
||||
conda_channels = ["conda-forge"]
|
||||
|
||||
[tool.comfy.isolation.cuda_wheels]
|
||||
index_url = "https://example.invalid/cuda-wheels"
|
||||
packages = ["pytorch3d", "gsplat", "nvdiffrast", "flash-attn", "sageattention", "spconv", "cumm"]
|
||||
""".strip(),
|
||||
)
|
||||
|
||||
captured: dict[str, object] = {}
|
||||
|
||||
class DummyManager:
|
||||
def __init__(self, *args, **kwargs) -> None:
|
||||
return None
|
||||
|
||||
def load_extension(self, config):
|
||||
captured.update(config)
|
||||
return _DummyExtension()
|
||||
|
||||
monkeypatch.setattr(extension_loader_module.pyisolate, "ExtensionManager", DummyManager)
|
||||
monkeypatch.setattr(
|
||||
extension_loader_module,
|
||||
"load_host_policy",
|
||||
lambda base_path: {
|
||||
"sandbox_mode": "disabled",
|
||||
"allow_network": False,
|
||||
"writable_paths": [],
|
||||
"readonly_paths": [],
|
||||
},
|
||||
)
|
||||
monkeypatch.setattr(extension_loader_module, "is_cache_valid", lambda *args, **kwargs: True)
|
||||
monkeypatch.setattr(
|
||||
extension_loader_module,
|
||||
"load_from_cache",
|
||||
lambda *args, **kwargs: {"Node": {"display_name": "Node", "schema_v1": {}}},
|
||||
)
|
||||
monkeypatch.setitem(sys.modules, "folder_paths", SimpleNamespace(base_path=str(tmp_path)))
|
||||
|
||||
asyncio.run(
|
||||
load_isolated_node(
|
||||
node_dir,
|
||||
manifest_path,
|
||||
logging.getLogger("test"),
|
||||
lambda *args, **kwargs: object,
|
||||
tmp_path / "venvs",
|
||||
[],
|
||||
)
|
||||
)
|
||||
|
||||
assert captured["package_manager"] == "conda"
|
||||
assert captured["cuda_wheels"] == {
|
||||
"index_url": "https://example.invalid/cuda-wheels/",
|
||||
"packages": [
|
||||
"pytorch3d",
|
||||
"gsplat",
|
||||
"nvdiffrast",
|
||||
"flash-attn",
|
||||
"sageattention",
|
||||
"spconv",
|
||||
"cumm",
|
||||
],
|
||||
"package_map": {},
|
||||
}
|
||||
|
||||
|
||||
def test_load_isolated_node_omits_cuda_wheels_when_not_configured(tmp_path, monkeypatch):
|
||||
node_dir = tmp_path / "node"
|
||||
node_dir.mkdir()
|
||||
manifest_path = node_dir / "pyproject.toml"
|
||||
_write_manifest(
|
||||
node_dir,
|
||||
"""
|
||||
[project]
|
||||
name = "demo-node"
|
||||
dependencies = ["numpy>=1.0"]
|
||||
|
||||
[tool.comfy.isolation]
|
||||
can_isolate = true
|
||||
""".strip(),
|
||||
)
|
||||
|
||||
captured: dict[str, object] = {}
|
||||
|
||||
class DummyManager:
|
||||
def __init__(self, *args, **kwargs) -> None:
|
||||
return None
|
||||
|
||||
def load_extension(self, config):
|
||||
captured.update(config)
|
||||
return _DummyExtension()
|
||||
|
||||
monkeypatch.setattr(extension_loader_module.pyisolate, "ExtensionManager", DummyManager)
|
||||
monkeypatch.setattr(
|
||||
extension_loader_module,
|
||||
"load_host_policy",
|
||||
lambda base_path: {
|
||||
"sandbox_mode": "disabled",
|
||||
"allow_network": False,
|
||||
"writable_paths": [],
|
||||
"readonly_paths": [],
|
||||
},
|
||||
)
|
||||
monkeypatch.setattr(extension_loader_module, "is_cache_valid", lambda *args, **kwargs: True)
|
||||
monkeypatch.setattr(
|
||||
extension_loader_module,
|
||||
"load_from_cache",
|
||||
lambda *args, **kwargs: {"Node": {"display_name": "Node", "schema_v1": {}}},
|
||||
)
|
||||
monkeypatch.setitem(sys.modules, "folder_paths", SimpleNamespace(base_path=str(tmp_path)))
|
||||
|
||||
asyncio.run(
|
||||
load_isolated_node(
|
||||
node_dir,
|
||||
manifest_path,
|
||||
logging.getLogger("test"),
|
||||
lambda *args, **kwargs: object,
|
||||
tmp_path / "venvs",
|
||||
[],
|
||||
)
|
||||
)
|
||||
|
||||
assert captured["sandbox_mode"] == "disabled"
|
||||
assert "cuda_wheels" not in captured
|
||||
|
||||
|
||||
def test_load_isolated_node_passes_extra_index_urls(tmp_path, monkeypatch):
|
||||
node_dir = tmp_path / "node"
|
||||
node_dir.mkdir()
|
||||
manifest_path = node_dir / "pyproject.toml"
|
||||
_write_manifest(
|
||||
node_dir,
|
||||
"""
|
||||
[project]
|
||||
name = "demo-node"
|
||||
dependencies = ["fbxsdkpy==2020.1.post2", "numpy>=1.0"]
|
||||
|
||||
[tool.comfy.isolation]
|
||||
can_isolate = true
|
||||
share_torch = true
|
||||
extra_index_urls = ["https://gitlab.inria.fr/api/v4/projects/18692/packages/pypi/simple"]
|
||||
""".strip(),
|
||||
)
|
||||
|
||||
captured: dict[str, object] = {}
|
||||
|
||||
class DummyManager:
|
||||
def __init__(self, *args, **kwargs) -> None:
|
||||
return None
|
||||
|
||||
def load_extension(self, config):
|
||||
captured.update(config)
|
||||
return _DummyExtension()
|
||||
|
||||
monkeypatch.setattr(extension_loader_module.pyisolate, "ExtensionManager", DummyManager)
|
||||
monkeypatch.setattr(
|
||||
extension_loader_module,
|
||||
"load_host_policy",
|
||||
lambda base_path: {
|
||||
"sandbox_mode": "disabled",
|
||||
"allow_network": False,
|
||||
"writable_paths": [],
|
||||
"readonly_paths": [],
|
||||
},
|
||||
)
|
||||
monkeypatch.setattr(extension_loader_module, "is_cache_valid", lambda *args, **kwargs: True)
|
||||
monkeypatch.setattr(
|
||||
extension_loader_module,
|
||||
"load_from_cache",
|
||||
lambda *args, **kwargs: {"Node": {"display_name": "Node", "schema_v1": {}}},
|
||||
)
|
||||
monkeypatch.setitem(sys.modules, "folder_paths", SimpleNamespace(base_path=str(tmp_path)))
|
||||
|
||||
specs = asyncio.run(
|
||||
load_isolated_node(
|
||||
node_dir,
|
||||
manifest_path,
|
||||
logging.getLogger("test"),
|
||||
lambda *args, **kwargs: object,
|
||||
tmp_path / "venvs",
|
||||
[],
|
||||
)
|
||||
)
|
||||
|
||||
assert len(specs) == 1
|
||||
assert captured["extra_index_urls"] == [
|
||||
"https://gitlab.inria.fr/api/v4/projects/18692/packages/pypi/simple"
|
||||
]
|
||||
|
||||
|
||||
def test_maybe_wrap_model_for_isolation_uses_runtime_flag(monkeypatch):
|
||||
class DummyRegistry:
|
||||
def register(self, model):
|
||||
return "model-123"
|
||||
|
||||
class DummyProxy:
|
||||
def __init__(self, model_id, registry, manage_lifecycle):
|
||||
self.model_id = model_id
|
||||
self.registry = registry
|
||||
self.manage_lifecycle = manage_lifecycle
|
||||
|
||||
monkeypatch.setattr(model_patcher_proxy_utils.args, "use_process_isolation", True)
|
||||
monkeypatch.delenv("PYISOLATE_ISOLATION_ACTIVE", raising=False)
|
||||
monkeypatch.delenv("PYISOLATE_CHILD", raising=False)
|
||||
monkeypatch.setitem(
|
||||
sys.modules,
|
||||
"comfy.isolation.model_patcher_proxy_registry",
|
||||
SimpleNamespace(ModelPatcherRegistry=DummyRegistry),
|
||||
)
|
||||
monkeypatch.setitem(
|
||||
sys.modules,
|
||||
"comfy.isolation.model_patcher_proxy",
|
||||
SimpleNamespace(ModelPatcherProxy=DummyProxy),
|
||||
)
|
||||
|
||||
wrapped = cast(Any, maybe_wrap_model_for_isolation(object()))
|
||||
|
||||
assert isinstance(wrapped, DummyProxy)
|
||||
assert getattr(wrapped, "model_id") == "model-123"
|
||||
assert getattr(wrapped, "manage_lifecycle") is True
|
||||
|
||||
|
||||
def test_flush_transport_state_uses_child_env_without_legacy_flag(monkeypatch):
|
||||
monkeypatch.setenv("PYISOLATE_CHILD", "1")
|
||||
monkeypatch.delenv("PYISOLATE_ISOLATION_ACTIVE", raising=False)
|
||||
monkeypatch.setattr(extension_wrapper_module, "_flush_tensor_transport_state", lambda marker: 3)
|
||||
monkeypatch.setitem(
|
||||
sys.modules,
|
||||
"comfy.isolation.model_patcher_proxy_registry",
|
||||
SimpleNamespace(
|
||||
ModelPatcherRegistry=lambda: SimpleNamespace(
|
||||
sweep_pending_cleanup=lambda: 0
|
||||
)
|
||||
),
|
||||
)
|
||||
|
||||
flushed = asyncio.run(
|
||||
ComfyNodeExtension.flush_transport_state(SimpleNamespace(name="demo"))
|
||||
)
|
||||
|
||||
assert flushed == 3
|
||||
|
||||
|
||||
def test_build_stub_class_relieves_host_vram_without_legacy_flag(monkeypatch):
|
||||
relieve_calls: list[str] = []
|
||||
|
||||
async def deserialize_from_isolation(result, extension):
|
||||
return result
|
||||
|
||||
monkeypatch.delenv("PYISOLATE_CHILD", raising=False)
|
||||
monkeypatch.delenv("PYISOLATE_ISOLATION_ACTIVE", raising=False)
|
||||
monkeypatch.setattr(
|
||||
runtime_helpers, "_relieve_host_vram_pressure", lambda marker, logger: relieve_calls.append(marker)
|
||||
)
|
||||
monkeypatch.setattr(runtime_helpers, "scan_shm_forensics", lambda *args, **kwargs: None)
|
||||
monkeypatch.setattr(isolation_pkg, "_RUNNING_EXTENSIONS", {}, raising=False)
|
||||
monkeypatch.setitem(
|
||||
sys.modules,
|
||||
"pyisolate._internal.model_serialization",
|
||||
SimpleNamespace(
|
||||
serialize_for_isolation=lambda payload: payload,
|
||||
deserialize_from_isolation=deserialize_from_isolation,
|
||||
),
|
||||
)
|
||||
|
||||
class DummyExtension:
|
||||
name = "demo-extension"
|
||||
module_path = os.getcwd()
|
||||
|
||||
async def execute_node(self, node_name, **inputs):
|
||||
return inputs
|
||||
|
||||
stub_cls = runtime_helpers.build_stub_class(
|
||||
"DemoNode",
|
||||
{"input_types": {}},
|
||||
DummyExtension(),
|
||||
{},
|
||||
logging.getLogger("test"),
|
||||
)
|
||||
|
||||
result = asyncio.run(
|
||||
getattr(stub_cls, "_pyisolate_execute")(SimpleNamespace(), value=1)
|
||||
)
|
||||
|
||||
assert relieve_calls == ["RUNTIME:pre_execute"]
|
||||
assert result == {"value": 1}
|
||||
@@ -1,22 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from tests.isolation.singleton_boundary_helpers import (
|
||||
capture_exact_proxy_bootstrap_contract,
|
||||
)
|
||||
|
||||
|
||||
def test_no_proxy_omission_allowed() -> None:
|
||||
payload = capture_exact_proxy_bootstrap_contract()
|
||||
|
||||
assert payload["omitted_proxies"] == []
|
||||
assert payload["forbidden_matches"] == []
|
||||
|
||||
matrix = payload["matrix"]
|
||||
assert matrix["base.py"]["bound"] is True
|
||||
assert matrix["folder_paths_proxy.py"]["bound"] is True
|
||||
assert matrix["helper_proxies.py"]["bound"] is True
|
||||
assert matrix["model_management_proxy.py"]["bound"] is True
|
||||
assert matrix["progress_proxy.py"]["bound"] is True
|
||||
assert matrix["prompt_server_impl.py"]["bound"] is True
|
||||
assert matrix["utils_proxy.py"]["bound"] is True
|
||||
assert matrix["web_directory_proxy.py"]["bound"] is True
|
||||
@@ -1,128 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from tests.isolation.singleton_boundary_helpers import (
|
||||
capture_exact_small_proxy_relay,
|
||||
capture_model_management_exact_relay,
|
||||
capture_prompt_web_exact_relay,
|
||||
)
|
||||
|
||||
|
||||
def _transcripts_for(payload: dict[str, object], object_id: str, method: str) -> list[dict[str, object]]:
|
||||
return [
|
||||
entry
|
||||
for entry in payload["transcripts"]
|
||||
if entry["object_id"] == object_id and entry["method"] == method
|
||||
]
|
||||
|
||||
|
||||
def test_folder_paths_exact_relay() -> None:
|
||||
payload = capture_exact_small_proxy_relay()
|
||||
|
||||
assert payload["forbidden_matches"] == []
|
||||
assert payload["models_dir"] == "/sandbox/models"
|
||||
assert payload["folder_path"] == "/sandbox/input/demo.png"
|
||||
|
||||
models_dir_calls = _transcripts_for(payload, "FolderPathsProxy", "rpc_get_models_dir")
|
||||
annotated_calls = _transcripts_for(payload, "FolderPathsProxy", "rpc_get_annotated_filepath")
|
||||
|
||||
assert models_dir_calls
|
||||
assert annotated_calls
|
||||
assert all(entry["phase"] != "child_call" or entry["method"] != "rpc_snapshot" for entry in payload["transcripts"])
|
||||
|
||||
|
||||
def test_progress_exact_relay() -> None:
|
||||
payload = capture_exact_small_proxy_relay()
|
||||
|
||||
progress_calls = _transcripts_for(payload, "ProgressProxy", "rpc_set_progress")
|
||||
|
||||
assert progress_calls
|
||||
host_targets = [entry["target"] for entry in progress_calls if entry["phase"] == "host_invocation"]
|
||||
assert host_targets == ["comfy_execution.progress.get_progress_state().update_progress"]
|
||||
result_entries = [entry for entry in progress_calls if entry["phase"] == "result"]
|
||||
assert result_entries == [{"phase": "result", "object_id": "ProgressProxy", "method": "rpc_set_progress", "result": None}]
|
||||
|
||||
|
||||
def test_utils_exact_relay() -> None:
|
||||
payload = capture_exact_small_proxy_relay()
|
||||
|
||||
utils_calls = _transcripts_for(payload, "UtilsProxy", "progress_bar_hook")
|
||||
|
||||
assert utils_calls
|
||||
host_targets = [entry["target"] for entry in utils_calls if entry["phase"] == "host_invocation"]
|
||||
assert host_targets == ["comfy.utils.PROGRESS_BAR_HOOK"]
|
||||
result_entries = [entry for entry in utils_calls if entry["phase"] == "result"]
|
||||
assert result_entries
|
||||
assert result_entries[0]["result"]["value"] == 2
|
||||
assert result_entries[0]["result"]["total"] == 5
|
||||
|
||||
|
||||
def test_helper_proxy_exact_relay() -> None:
|
||||
payload = capture_exact_small_proxy_relay()
|
||||
|
||||
helper_calls = _transcripts_for(payload, "HelperProxiesService", "rpc_restore_input_types")
|
||||
|
||||
assert helper_calls
|
||||
host_targets = [entry["target"] for entry in helper_calls if entry["phase"] == "host_invocation"]
|
||||
assert host_targets == ["comfy.isolation.proxies.helper_proxies.restore_input_types"]
|
||||
assert payload["restored_any_type"] == "*"
|
||||
|
||||
|
||||
def test_model_management_exact_relay() -> None:
|
||||
payload = capture_model_management_exact_relay()
|
||||
|
||||
model_calls = _transcripts_for(payload, "ModelManagementProxy", "get_torch_device")
|
||||
model_calls += _transcripts_for(payload, "ModelManagementProxy", "get_torch_device_name")
|
||||
model_calls += _transcripts_for(payload, "ModelManagementProxy", "get_free_memory")
|
||||
|
||||
assert payload["forbidden_matches"] == []
|
||||
assert model_calls
|
||||
host_targets = [
|
||||
entry["target"]
|
||||
for entry in payload["transcripts"]
|
||||
if entry["phase"] == "host_invocation"
|
||||
]
|
||||
assert host_targets == [
|
||||
"comfy.model_management.get_torch_device",
|
||||
"comfy.model_management.get_torch_device_name",
|
||||
"comfy.model_management.get_free_memory",
|
||||
]
|
||||
|
||||
|
||||
def test_model_management_capability_preserved() -> None:
|
||||
payload = capture_model_management_exact_relay()
|
||||
|
||||
assert payload["device"] == "cpu"
|
||||
assert payload["device_type"] == "cpu"
|
||||
assert payload["device_name"] == "cpu"
|
||||
assert payload["free_memory"] == 34359738368
|
||||
|
||||
|
||||
def test_prompt_server_exact_relay() -> None:
|
||||
payload = capture_prompt_web_exact_relay()
|
||||
|
||||
prompt_calls = _transcripts_for(payload, "PromptServerService", "ui_send_progress_text")
|
||||
prompt_calls += _transcripts_for(payload, "PromptServerService", "register_route_rpc")
|
||||
|
||||
assert payload["forbidden_matches"] == []
|
||||
assert prompt_calls
|
||||
host_targets = [
|
||||
entry["target"]
|
||||
for entry in payload["transcripts"]
|
||||
if entry["object_id"] == "PromptServerService" and entry["phase"] == "host_invocation"
|
||||
]
|
||||
assert host_targets == [
|
||||
"server.PromptServer.instance.send_progress_text",
|
||||
"server.PromptServer.instance.routes.add_route",
|
||||
]
|
||||
|
||||
|
||||
def test_web_directory_exact_relay() -> None:
|
||||
payload = capture_prompt_web_exact_relay()
|
||||
|
||||
web_calls = _transcripts_for(payload, "WebDirectoryProxy", "get_web_file")
|
||||
|
||||
assert web_calls
|
||||
host_targets = [entry["target"] for entry in web_calls if entry["phase"] == "host_invocation"]
|
||||
assert host_targets == ["comfy.isolation.proxies.web_directory_proxy.WebDirectoryProxy.get_web_file"]
|
||||
assert payload["web_file"]["content_type"] == "application/javascript"
|
||||
assert payload["web_file"]["content"] == "console.log('deo');"
|
||||
@@ -1,428 +0,0 @@
|
||||
"""Tests for conda config parsing in extension_loader.py (Slice 5).
|
||||
|
||||
These tests verify that extension_loader.py correctly parses conda-related
|
||||
fields from pyproject.toml manifests and passes them into the extension config
|
||||
dict given to pyisolate. The torch import chain is broken by pre-mocking
|
||||
extension_wrapper before importing extension_loader.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import importlib
|
||||
import sys
|
||||
import types
|
||||
from pathlib import Path
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
def _make_manifest(
|
||||
*,
|
||||
package_manager: str = "uv",
|
||||
conda_channels: list[str] | None = None,
|
||||
conda_dependencies: list[str] | None = None,
|
||||
conda_platforms: list[str] | None = None,
|
||||
share_torch: bool = False,
|
||||
can_isolate: bool = True,
|
||||
dependencies: list[str] | None = None,
|
||||
cuda_wheels: list[str] | None = None,
|
||||
) -> dict:
|
||||
"""Build a manifest dict matching tomllib.load() output."""
|
||||
isolation: dict = {"can_isolate": can_isolate}
|
||||
if package_manager != "uv":
|
||||
isolation["package_manager"] = package_manager
|
||||
if conda_channels is not None:
|
||||
isolation["conda_channels"] = conda_channels
|
||||
if conda_dependencies is not None:
|
||||
isolation["conda_dependencies"] = conda_dependencies
|
||||
if conda_platforms is not None:
|
||||
isolation["conda_platforms"] = conda_platforms
|
||||
if share_torch:
|
||||
isolation["share_torch"] = True
|
||||
if cuda_wheels is not None:
|
||||
isolation["cuda_wheels"] = cuda_wheels
|
||||
|
||||
return {
|
||||
"project": {
|
||||
"name": "test-extension",
|
||||
"dependencies": dependencies or ["numpy"],
|
||||
},
|
||||
"tool": {"comfy": {"isolation": isolation}},
|
||||
}
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def manifest_file(tmp_path):
|
||||
"""Create a dummy pyproject.toml so manifest_path.open('rb') succeeds."""
|
||||
path = tmp_path / "pyproject.toml"
|
||||
path.write_bytes(b"") # content is overridden by tomllib mock
|
||||
return path
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def loader_module(monkeypatch):
|
||||
"""Import extension_loader under a mocked isolation package for this test only."""
|
||||
mock_wrapper = MagicMock()
|
||||
mock_wrapper.ComfyNodeExtension = type("ComfyNodeExtension", (), {})
|
||||
|
||||
iso_mod = types.ModuleType("comfy.isolation")
|
||||
iso_mod.__path__ = [ # type: ignore[attr-defined]
|
||||
str(Path(__file__).resolve().parent.parent.parent / "comfy" / "isolation")
|
||||
]
|
||||
iso_mod.__package__ = "comfy.isolation"
|
||||
|
||||
manifest_loader = types.SimpleNamespace(
|
||||
is_cache_valid=lambda *args, **kwargs: False,
|
||||
load_from_cache=lambda *args, **kwargs: None,
|
||||
save_to_cache=lambda *args, **kwargs: None,
|
||||
)
|
||||
host_policy = types.SimpleNamespace(
|
||||
load_host_policy=lambda base_path: {
|
||||
"sandbox_mode": "required",
|
||||
"allow_network": False,
|
||||
"writable_paths": [],
|
||||
"readonly_paths": [],
|
||||
}
|
||||
)
|
||||
folder_paths = types.SimpleNamespace(base_path="/fake/comfyui")
|
||||
|
||||
monkeypatch.setitem(sys.modules, "comfy.isolation", iso_mod)
|
||||
monkeypatch.setitem(sys.modules, "comfy.isolation.extension_wrapper", mock_wrapper)
|
||||
monkeypatch.setitem(sys.modules, "comfy.isolation.runtime_helpers", MagicMock())
|
||||
monkeypatch.setitem(sys.modules, "comfy.isolation.manifest_loader", manifest_loader)
|
||||
monkeypatch.setitem(sys.modules, "comfy.isolation.host_policy", host_policy)
|
||||
monkeypatch.setitem(sys.modules, "folder_paths", folder_paths)
|
||||
sys.modules.pop("comfy.isolation.extension_loader", None)
|
||||
|
||||
module = importlib.import_module("comfy.isolation.extension_loader")
|
||||
try:
|
||||
yield module, mock_wrapper
|
||||
finally:
|
||||
sys.modules.pop("comfy.isolation.extension_loader", None)
|
||||
comfy_pkg = sys.modules.get("comfy")
|
||||
if comfy_pkg is not None and hasattr(comfy_pkg, "isolation"):
|
||||
delattr(comfy_pkg, "isolation")
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_pyisolate(loader_module):
|
||||
"""Mock pyisolate to avoid real venv creation."""
|
||||
module, mock_wrapper = loader_module
|
||||
mock_ext = AsyncMock()
|
||||
mock_ext.list_nodes = AsyncMock(return_value={})
|
||||
|
||||
mock_manager = MagicMock()
|
||||
mock_manager.load_extension = MagicMock(return_value=mock_ext)
|
||||
sealed_type = type("SealedNodeExtension", (), {})
|
||||
|
||||
with patch.object(module, "pyisolate") as mock_pi:
|
||||
mock_pi.ExtensionManager = MagicMock(return_value=mock_manager)
|
||||
mock_pi.SealedNodeExtension = sealed_type
|
||||
yield module, mock_pi, mock_manager, mock_ext, mock_wrapper
|
||||
|
||||
|
||||
def load_isolated_node(*args, **kwargs):
|
||||
return sys.modules["comfy.isolation.extension_loader"].load_isolated_node(
|
||||
*args, **kwargs
|
||||
)
|
||||
|
||||
|
||||
class TestCondaPackageManagerParsing:
|
||||
"""Verify extension_loader.py parses conda config from pyproject.toml."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_conda_package_manager_in_config(
|
||||
self, mock_pyisolate, manifest_file, tmp_path
|
||||
):
|
||||
"""package_manager='conda' must appear in extension_config."""
|
||||
|
||||
manifest = _make_manifest(
|
||||
package_manager="conda",
|
||||
conda_channels=["conda-forge"],
|
||||
conda_dependencies=["eccodes"],
|
||||
)
|
||||
|
||||
_, _, mock_manager, _, _ = mock_pyisolate
|
||||
|
||||
with patch("comfy.isolation.extension_loader.tomllib") as mock_tomllib:
|
||||
mock_tomllib.load.return_value = manifest
|
||||
await load_isolated_node(
|
||||
node_dir=tmp_path,
|
||||
manifest_path=manifest_file,
|
||||
logger=MagicMock(),
|
||||
build_stub_class=MagicMock(),
|
||||
venv_root=tmp_path / "venvs",
|
||||
extension_managers=[],
|
||||
)
|
||||
|
||||
config = mock_manager.load_extension.call_args[0][0]
|
||||
assert config["package_manager"] == "conda"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_conda_channels_in_config(
|
||||
self, mock_pyisolate, manifest_file, tmp_path
|
||||
):
|
||||
"""conda_channels must be passed through to extension_config."""
|
||||
|
||||
manifest = _make_manifest(
|
||||
package_manager="conda",
|
||||
conda_channels=["conda-forge", "nvidia"],
|
||||
conda_dependencies=["eccodes"],
|
||||
)
|
||||
|
||||
_, _, mock_manager, _, _ = mock_pyisolate
|
||||
|
||||
with patch("comfy.isolation.extension_loader.tomllib") as mock_tomllib:
|
||||
mock_tomllib.load.return_value = manifest
|
||||
await load_isolated_node(
|
||||
node_dir=tmp_path,
|
||||
manifest_path=manifest_file,
|
||||
logger=MagicMock(),
|
||||
build_stub_class=MagicMock(),
|
||||
venv_root=tmp_path / "venvs",
|
||||
extension_managers=[],
|
||||
)
|
||||
|
||||
config = mock_manager.load_extension.call_args[0][0]
|
||||
assert config["conda_channels"] == ["conda-forge", "nvidia"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_conda_dependencies_in_config(
|
||||
self, mock_pyisolate, manifest_file, tmp_path
|
||||
):
|
||||
"""conda_dependencies must be passed through to extension_config."""
|
||||
|
||||
manifest = _make_manifest(
|
||||
package_manager="conda",
|
||||
conda_channels=["conda-forge"],
|
||||
conda_dependencies=["eccodes", "cfgrib"],
|
||||
)
|
||||
|
||||
_, _, mock_manager, _, _ = mock_pyisolate
|
||||
|
||||
with patch("comfy.isolation.extension_loader.tomllib") as mock_tomllib:
|
||||
mock_tomllib.load.return_value = manifest
|
||||
await load_isolated_node(
|
||||
node_dir=tmp_path,
|
||||
manifest_path=manifest_file,
|
||||
logger=MagicMock(),
|
||||
build_stub_class=MagicMock(),
|
||||
venv_root=tmp_path / "venvs",
|
||||
extension_managers=[],
|
||||
)
|
||||
|
||||
config = mock_manager.load_extension.call_args[0][0]
|
||||
assert config["conda_dependencies"] == ["eccodes", "cfgrib"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_conda_platforms_in_config(
|
||||
self, mock_pyisolate, manifest_file, tmp_path
|
||||
):
|
||||
"""conda_platforms must be passed through to extension_config."""
|
||||
|
||||
manifest = _make_manifest(
|
||||
package_manager="conda",
|
||||
conda_channels=["conda-forge"],
|
||||
conda_dependencies=["eccodes"],
|
||||
conda_platforms=["linux-64"],
|
||||
)
|
||||
|
||||
_, _, mock_manager, _, _ = mock_pyisolate
|
||||
|
||||
with patch("comfy.isolation.extension_loader.tomllib") as mock_tomllib:
|
||||
mock_tomllib.load.return_value = manifest
|
||||
await load_isolated_node(
|
||||
node_dir=tmp_path,
|
||||
manifest_path=manifest_file,
|
||||
logger=MagicMock(),
|
||||
build_stub_class=MagicMock(),
|
||||
venv_root=tmp_path / "venvs",
|
||||
extension_managers=[],
|
||||
)
|
||||
|
||||
config = mock_manager.load_extension.call_args[0][0]
|
||||
assert config["conda_platforms"] == ["linux-64"]
|
||||
|
||||
|
||||
class TestCondaForcedOverrides:
|
||||
"""Verify conda forces share_torch=False, share_cuda_ipc=False."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_conda_forces_share_torch_false(
|
||||
self, mock_pyisolate, manifest_file, tmp_path
|
||||
):
|
||||
"""share_torch must be forced False for conda, even if manifest says True."""
|
||||
|
||||
manifest = _make_manifest(
|
||||
package_manager="conda",
|
||||
conda_channels=["conda-forge"],
|
||||
conda_dependencies=["eccodes"],
|
||||
share_torch=True, # manifest requests True — must be overridden
|
||||
)
|
||||
|
||||
_, _, mock_manager, _, _ = mock_pyisolate
|
||||
|
||||
with patch("comfy.isolation.extension_loader.tomllib") as mock_tomllib:
|
||||
mock_tomllib.load.return_value = manifest
|
||||
await load_isolated_node(
|
||||
node_dir=tmp_path,
|
||||
manifest_path=manifest_file,
|
||||
logger=MagicMock(),
|
||||
build_stub_class=MagicMock(),
|
||||
venv_root=tmp_path / "venvs",
|
||||
extension_managers=[],
|
||||
)
|
||||
|
||||
config = mock_manager.load_extension.call_args[0][0]
|
||||
assert config["share_torch"] is False
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_conda_forces_share_cuda_ipc_false(
|
||||
self, mock_pyisolate, manifest_file, tmp_path
|
||||
):
|
||||
"""share_cuda_ipc must be forced False for conda."""
|
||||
|
||||
manifest = _make_manifest(
|
||||
package_manager="conda",
|
||||
conda_channels=["conda-forge"],
|
||||
conda_dependencies=["eccodes"],
|
||||
share_torch=True,
|
||||
)
|
||||
|
||||
_, _, mock_manager, _, _ = mock_pyisolate
|
||||
|
||||
with patch("comfy.isolation.extension_loader.tomllib") as mock_tomllib:
|
||||
mock_tomllib.load.return_value = manifest
|
||||
await load_isolated_node(
|
||||
node_dir=tmp_path,
|
||||
manifest_path=manifest_file,
|
||||
logger=MagicMock(),
|
||||
build_stub_class=MagicMock(),
|
||||
venv_root=tmp_path / "venvs",
|
||||
extension_managers=[],
|
||||
)
|
||||
|
||||
config = mock_manager.load_extension.call_args[0][0]
|
||||
assert config["share_cuda_ipc"] is False
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_conda_sealed_worker_uses_host_policy_sandbox_config(
|
||||
self, mock_pyisolate, manifest_file, tmp_path
|
||||
):
|
||||
"""Conda sealed_worker must carry the host-policy sandbox config on Linux."""
|
||||
|
||||
manifest = _make_manifest(
|
||||
package_manager="conda",
|
||||
conda_channels=["conda-forge"],
|
||||
conda_dependencies=["eccodes"],
|
||||
)
|
||||
|
||||
_, _, mock_manager, _, _ = mock_pyisolate
|
||||
|
||||
with (
|
||||
patch("comfy.isolation.extension_loader.tomllib") as mock_tomllib,
|
||||
patch(
|
||||
"comfy.isolation.extension_loader.platform.system",
|
||||
return_value="Linux",
|
||||
),
|
||||
):
|
||||
mock_tomllib.load.return_value = manifest
|
||||
await load_isolated_node(
|
||||
node_dir=tmp_path,
|
||||
manifest_path=manifest_file,
|
||||
logger=MagicMock(),
|
||||
build_stub_class=MagicMock(),
|
||||
venv_root=tmp_path / "venvs",
|
||||
extension_managers=[],
|
||||
)
|
||||
|
||||
config = mock_manager.load_extension.call_args[0][0]
|
||||
assert config["sandbox"] == {
|
||||
"network": False,
|
||||
"writable_paths": [],
|
||||
"readonly_paths": [],
|
||||
}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_conda_uses_sealed_extension_type(
|
||||
self, mock_pyisolate, manifest_file, tmp_path
|
||||
):
|
||||
"""Conda must not launch through ComfyNodeExtension."""
|
||||
|
||||
_, mock_pi, _, _, mock_wrapper = mock_pyisolate
|
||||
manifest = _make_manifest(
|
||||
package_manager="conda",
|
||||
conda_channels=["conda-forge"],
|
||||
conda_dependencies=["eccodes"],
|
||||
)
|
||||
|
||||
with patch("comfy.isolation.extension_loader.tomllib") as mock_tomllib:
|
||||
mock_tomllib.load.return_value = manifest
|
||||
await load_isolated_node(
|
||||
node_dir=tmp_path,
|
||||
manifest_path=manifest_file,
|
||||
logger=MagicMock(),
|
||||
build_stub_class=MagicMock(),
|
||||
venv_root=tmp_path / "venvs",
|
||||
extension_managers=[],
|
||||
)
|
||||
|
||||
extension_type = mock_pi.ExtensionManager.call_args[0][0]
|
||||
assert extension_type.__name__ == "SealedNodeExtension"
|
||||
assert extension_type is not mock_wrapper.ComfyNodeExtension
|
||||
|
||||
|
||||
class TestUvUnchanged:
|
||||
"""Verify uv extensions are NOT affected by conda changes."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_uv_default_no_conda_keys(
|
||||
self, mock_pyisolate, manifest_file, tmp_path
|
||||
):
|
||||
"""Default uv extension must NOT have package_manager or conda keys."""
|
||||
|
||||
manifest = _make_manifest() # defaults: uv, no conda fields
|
||||
|
||||
_, _, mock_manager, _, _ = mock_pyisolate
|
||||
|
||||
with patch("comfy.isolation.extension_loader.tomllib") as mock_tomllib:
|
||||
mock_tomllib.load.return_value = manifest
|
||||
await load_isolated_node(
|
||||
node_dir=tmp_path,
|
||||
manifest_path=manifest_file,
|
||||
logger=MagicMock(),
|
||||
build_stub_class=MagicMock(),
|
||||
venv_root=tmp_path / "venvs",
|
||||
extension_managers=[],
|
||||
)
|
||||
|
||||
config = mock_manager.load_extension.call_args[0][0]
|
||||
# uv extensions should not have conda-specific keys
|
||||
assert config.get("package_manager", "uv") == "uv"
|
||||
assert "conda_channels" not in config
|
||||
assert "conda_dependencies" not in config
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_uv_keeps_comfy_extension_type(
|
||||
self, mock_pyisolate, manifest_file, tmp_path
|
||||
):
|
||||
"""uv keeps the existing ComfyNodeExtension path."""
|
||||
|
||||
_, mock_pi, _, _, _ = mock_pyisolate
|
||||
manifest = _make_manifest()
|
||||
|
||||
with patch("comfy.isolation.extension_loader.tomllib") as mock_tomllib:
|
||||
mock_tomllib.load.return_value = manifest
|
||||
await load_isolated_node(
|
||||
node_dir=tmp_path,
|
||||
manifest_path=manifest_file,
|
||||
logger=MagicMock(),
|
||||
build_stub_class=MagicMock(),
|
||||
venv_root=tmp_path / "venvs",
|
||||
extension_managers=[],
|
||||
)
|
||||
|
||||
extension_type = mock_pi.ExtensionManager.call_args[0][0]
|
||||
assert extension_type.__name__ == "ComfyNodeExtension"
|
||||
assert extension_type is not mock_pi.SealedNodeExtension
|
||||
@@ -1,281 +0,0 @@
|
||||
"""Tests for execution_model parsing and sealed-worker loader selection."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import importlib
|
||||
import sys
|
||||
import types
|
||||
from pathlib import Path
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
def _make_manifest(
|
||||
*,
|
||||
package_manager: str = "uv",
|
||||
execution_model: str | None = None,
|
||||
can_isolate: bool = True,
|
||||
dependencies: list[str] | None = None,
|
||||
sealed_host_ro_paths: list[str] | None = None,
|
||||
) -> dict:
|
||||
isolation: dict = {"can_isolate": can_isolate}
|
||||
if package_manager != "uv":
|
||||
isolation["package_manager"] = package_manager
|
||||
if execution_model is not None:
|
||||
isolation["execution_model"] = execution_model
|
||||
if sealed_host_ro_paths is not None:
|
||||
isolation["sealed_host_ro_paths"] = sealed_host_ro_paths
|
||||
|
||||
return {
|
||||
"project": {
|
||||
"name": "test-extension",
|
||||
"dependencies": dependencies or ["numpy"],
|
||||
},
|
||||
"tool": {"comfy": {"isolation": isolation}},
|
||||
}
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def manifest_file(tmp_path):
|
||||
path = tmp_path / "pyproject.toml"
|
||||
path.write_bytes(b"")
|
||||
return path
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def loader_module(monkeypatch):
|
||||
mock_wrapper = MagicMock()
|
||||
mock_wrapper.ComfyNodeExtension = type("ComfyNodeExtension", (), {})
|
||||
|
||||
iso_mod = types.ModuleType("comfy.isolation")
|
||||
iso_mod.__path__ = [ # type: ignore[attr-defined]
|
||||
str(Path(__file__).resolve().parent.parent.parent / "comfy" / "isolation")
|
||||
]
|
||||
iso_mod.__package__ = "comfy.isolation"
|
||||
|
||||
manifest_loader = types.SimpleNamespace(
|
||||
is_cache_valid=lambda *args, **kwargs: False,
|
||||
load_from_cache=lambda *args, **kwargs: None,
|
||||
save_to_cache=lambda *args, **kwargs: None,
|
||||
)
|
||||
host_policy = types.SimpleNamespace(
|
||||
load_host_policy=lambda base_path: {
|
||||
"sandbox_mode": "required",
|
||||
"allow_network": False,
|
||||
"writable_paths": [],
|
||||
"readonly_paths": [],
|
||||
"sealed_worker_ro_import_paths": [],
|
||||
}
|
||||
)
|
||||
folder_paths = types.SimpleNamespace(base_path="/fake/comfyui")
|
||||
|
||||
monkeypatch.setitem(sys.modules, "comfy.isolation", iso_mod)
|
||||
monkeypatch.setitem(sys.modules, "comfy.isolation.extension_wrapper", mock_wrapper)
|
||||
monkeypatch.setitem(sys.modules, "comfy.isolation.runtime_helpers", MagicMock())
|
||||
monkeypatch.setitem(sys.modules, "comfy.isolation.manifest_loader", manifest_loader)
|
||||
monkeypatch.setitem(sys.modules, "comfy.isolation.host_policy", host_policy)
|
||||
monkeypatch.setitem(sys.modules, "folder_paths", folder_paths)
|
||||
sys.modules.pop("comfy.isolation.extension_loader", None)
|
||||
|
||||
module = importlib.import_module("comfy.isolation.extension_loader")
|
||||
try:
|
||||
yield module
|
||||
finally:
|
||||
sys.modules.pop("comfy.isolation.extension_loader", None)
|
||||
comfy_pkg = sys.modules.get("comfy")
|
||||
if comfy_pkg is not None and hasattr(comfy_pkg, "isolation"):
|
||||
delattr(comfy_pkg, "isolation")
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_pyisolate(loader_module):
|
||||
mock_ext = AsyncMock()
|
||||
mock_ext.list_nodes = AsyncMock(return_value={})
|
||||
|
||||
mock_manager = MagicMock()
|
||||
mock_manager.load_extension = MagicMock(return_value=mock_ext)
|
||||
sealed_type = type("SealedNodeExtension", (), {})
|
||||
|
||||
with patch.object(loader_module, "pyisolate") as mock_pi:
|
||||
mock_pi.ExtensionManager = MagicMock(return_value=mock_manager)
|
||||
mock_pi.SealedNodeExtension = sealed_type
|
||||
yield loader_module, mock_pi, mock_manager, mock_ext, sealed_type
|
||||
|
||||
|
||||
def load_isolated_node(*args, **kwargs):
|
||||
return sys.modules["comfy.isolation.extension_loader"].load_isolated_node(*args, **kwargs)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_uv_sealed_worker_selects_sealed_extension_type(
|
||||
mock_pyisolate, manifest_file, tmp_path
|
||||
):
|
||||
manifest = _make_manifest(execution_model="sealed_worker")
|
||||
|
||||
_, mock_pi, mock_manager, _, sealed_type = mock_pyisolate
|
||||
|
||||
with patch("comfy.isolation.extension_loader.tomllib") as mock_tomllib:
|
||||
mock_tomllib.load.return_value = manifest
|
||||
await load_isolated_node(
|
||||
node_dir=tmp_path,
|
||||
manifest_path=manifest_file,
|
||||
logger=MagicMock(),
|
||||
build_stub_class=MagicMock(),
|
||||
venv_root=tmp_path / "venvs",
|
||||
extension_managers=[],
|
||||
)
|
||||
|
||||
extension_type = mock_pi.ExtensionManager.call_args[0][0]
|
||||
config = mock_manager.load_extension.call_args[0][0]
|
||||
assert extension_type is sealed_type
|
||||
assert config["execution_model"] == "sealed_worker"
|
||||
assert "apis" not in config
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_default_uv_keeps_host_coupled_extension_type(
|
||||
mock_pyisolate, manifest_file, tmp_path
|
||||
):
|
||||
manifest = _make_manifest()
|
||||
|
||||
_, mock_pi, mock_manager, _, sealed_type = mock_pyisolate
|
||||
|
||||
with patch("comfy.isolation.extension_loader.tomllib") as mock_tomllib:
|
||||
mock_tomllib.load.return_value = manifest
|
||||
await load_isolated_node(
|
||||
node_dir=tmp_path,
|
||||
manifest_path=manifest_file,
|
||||
logger=MagicMock(),
|
||||
build_stub_class=MagicMock(),
|
||||
venv_root=tmp_path / "venvs",
|
||||
extension_managers=[],
|
||||
)
|
||||
|
||||
extension_type = mock_pi.ExtensionManager.call_args[0][0]
|
||||
config = mock_manager.load_extension.call_args[0][0]
|
||||
assert extension_type is not sealed_type
|
||||
assert "execution_model" not in config
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_conda_without_execution_model_remains_sealed_worker(
|
||||
mock_pyisolate, manifest_file, tmp_path
|
||||
):
|
||||
manifest = _make_manifest(package_manager="conda")
|
||||
manifest["tool"]["comfy"]["isolation"]["conda_channels"] = ["conda-forge"]
|
||||
manifest["tool"]["comfy"]["isolation"]["conda_dependencies"] = ["eccodes"]
|
||||
|
||||
_, mock_pi, mock_manager, _, sealed_type = mock_pyisolate
|
||||
|
||||
with patch("comfy.isolation.extension_loader.tomllib") as mock_tomllib:
|
||||
mock_tomllib.load.return_value = manifest
|
||||
await load_isolated_node(
|
||||
node_dir=tmp_path,
|
||||
manifest_path=manifest_file,
|
||||
logger=MagicMock(),
|
||||
build_stub_class=MagicMock(),
|
||||
venv_root=tmp_path / "venvs",
|
||||
extension_managers=[],
|
||||
)
|
||||
|
||||
extension_type = mock_pi.ExtensionManager.call_args[0][0]
|
||||
config = mock_manager.load_extension.call_args[0][0]
|
||||
assert extension_type is sealed_type
|
||||
assert config["execution_model"] == "sealed_worker"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_sealed_worker_uses_host_policy_ro_import_paths(
|
||||
mock_pyisolate, manifest_file, tmp_path
|
||||
):
|
||||
manifest = _make_manifest(execution_model="sealed_worker")
|
||||
|
||||
module, _, mock_manager, _, _ = mock_pyisolate
|
||||
|
||||
with (
|
||||
patch("comfy.isolation.extension_loader.tomllib") as mock_tomllib,
|
||||
patch.object(
|
||||
module,
|
||||
"load_host_policy",
|
||||
return_value={
|
||||
"sandbox_mode": "required",
|
||||
"allow_network": False,
|
||||
"writable_paths": [],
|
||||
"readonly_paths": [],
|
||||
"sealed_worker_ro_import_paths": ["/home/johnj/ComfyUI"],
|
||||
},
|
||||
),
|
||||
):
|
||||
mock_tomllib.load.return_value = manifest
|
||||
await load_isolated_node(
|
||||
node_dir=tmp_path,
|
||||
manifest_path=manifest_file,
|
||||
logger=MagicMock(),
|
||||
build_stub_class=MagicMock(),
|
||||
venv_root=tmp_path / "venvs",
|
||||
extension_managers=[],
|
||||
)
|
||||
|
||||
config = mock_manager.load_extension.call_args[0][0]
|
||||
assert config["sealed_host_ro_paths"] == ["/home/johnj/ComfyUI"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_host_coupled_does_not_emit_sealed_host_ro_paths(
|
||||
mock_pyisolate, manifest_file, tmp_path
|
||||
):
|
||||
manifest = _make_manifest(execution_model="host-coupled")
|
||||
|
||||
module, _, mock_manager, _, _ = mock_pyisolate
|
||||
|
||||
with (
|
||||
patch("comfy.isolation.extension_loader.tomllib") as mock_tomllib,
|
||||
patch.object(
|
||||
module,
|
||||
"load_host_policy",
|
||||
return_value={
|
||||
"sandbox_mode": "required",
|
||||
"allow_network": False,
|
||||
"writable_paths": [],
|
||||
"readonly_paths": [],
|
||||
"sealed_worker_ro_import_paths": ["/home/johnj/ComfyUI"],
|
||||
},
|
||||
),
|
||||
):
|
||||
mock_tomllib.load.return_value = manifest
|
||||
await load_isolated_node(
|
||||
node_dir=tmp_path,
|
||||
manifest_path=manifest_file,
|
||||
logger=MagicMock(),
|
||||
build_stub_class=MagicMock(),
|
||||
venv_root=tmp_path / "venvs",
|
||||
extension_managers=[],
|
||||
)
|
||||
|
||||
config = mock_manager.load_extension.call_args[0][0]
|
||||
assert "sealed_host_ro_paths" not in config
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_sealed_worker_manifest_ro_import_paths_blocked(
|
||||
mock_pyisolate, manifest_file, tmp_path
|
||||
):
|
||||
manifest = _make_manifest(
|
||||
execution_model="sealed_worker",
|
||||
sealed_host_ro_paths=["/home/johnj/ComfyUI"],
|
||||
)
|
||||
|
||||
_, _, _mock_manager, _, _ = mock_pyisolate
|
||||
|
||||
with patch("comfy.isolation.extension_loader.tomllib") as mock_tomllib:
|
||||
mock_tomllib.load.return_value = manifest
|
||||
with pytest.raises(ValueError, match="Manifest field 'sealed_host_ro_paths' is not allowed"):
|
||||
await load_isolated_node(
|
||||
node_dir=tmp_path,
|
||||
manifest_path=manifest_file,
|
||||
logger=MagicMock(),
|
||||
build_stub_class=MagicMock(),
|
||||
venv_root=tmp_path / "venvs",
|
||||
extension_managers=[],
|
||||
)
|
||||
@@ -1,122 +0,0 @@
|
||||
"""Unit tests for FolderPathsProxy."""
|
||||
|
||||
import pytest
|
||||
from pathlib import Path
|
||||
|
||||
from comfy.isolation.proxies.folder_paths_proxy import FolderPathsProxy
|
||||
from tests.isolation.singleton_boundary_helpers import capture_sealed_singleton_imports
|
||||
|
||||
|
||||
class TestFolderPathsProxy:
|
||||
"""Test FolderPathsProxy methods."""
|
||||
|
||||
@pytest.fixture
|
||||
def proxy(self):
|
||||
"""Create a FolderPathsProxy instance for testing."""
|
||||
return FolderPathsProxy()
|
||||
|
||||
def test_get_temp_directory_returns_string(self, proxy):
|
||||
"""Verify get_temp_directory returns a non-empty string."""
|
||||
result = proxy.get_temp_directory()
|
||||
assert isinstance(result, str), f"Expected str, got {type(result)}"
|
||||
assert len(result) > 0, "Temp directory path is empty"
|
||||
|
||||
def test_get_temp_directory_returns_absolute_path(self, proxy):
|
||||
"""Verify get_temp_directory returns an absolute path."""
|
||||
result = proxy.get_temp_directory()
|
||||
path = Path(result)
|
||||
assert path.is_absolute(), f"Path is not absolute: {result}"
|
||||
|
||||
def test_get_input_directory_returns_string(self, proxy):
|
||||
"""Verify get_input_directory returns a non-empty string."""
|
||||
result = proxy.get_input_directory()
|
||||
assert isinstance(result, str), f"Expected str, got {type(result)}"
|
||||
assert len(result) > 0, "Input directory path is empty"
|
||||
|
||||
def test_get_input_directory_returns_absolute_path(self, proxy):
|
||||
"""Verify get_input_directory returns an absolute path."""
|
||||
result = proxy.get_input_directory()
|
||||
path = Path(result)
|
||||
assert path.is_absolute(), f"Path is not absolute: {result}"
|
||||
|
||||
def test_get_annotated_filepath_plain_name(self, proxy):
|
||||
"""Verify get_annotated_filepath works with plain filename."""
|
||||
result = proxy.get_annotated_filepath("test.png")
|
||||
assert isinstance(result, str), f"Expected str, got {type(result)}"
|
||||
assert "test.png" in result, f"Filename not in result: {result}"
|
||||
|
||||
def test_get_annotated_filepath_with_output_annotation(self, proxy):
|
||||
"""Verify get_annotated_filepath handles [output] annotation."""
|
||||
result = proxy.get_annotated_filepath("test.png[output]")
|
||||
assert isinstance(result, str), f"Expected str, got {type(result)}"
|
||||
assert "test.pn" in result, f"Filename base not in result: {result}"
|
||||
# Should resolve to output directory
|
||||
assert "output" in result.lower() or Path(result).parent.name == "output"
|
||||
|
||||
def test_get_annotated_filepath_with_input_annotation(self, proxy):
|
||||
"""Verify get_annotated_filepath handles [input] annotation."""
|
||||
result = proxy.get_annotated_filepath("test.png[input]")
|
||||
assert isinstance(result, str), f"Expected str, got {type(result)}"
|
||||
assert "test.pn" in result, f"Filename base not in result: {result}"
|
||||
|
||||
def test_get_annotated_filepath_with_temp_annotation(self, proxy):
|
||||
"""Verify get_annotated_filepath handles [temp] annotation."""
|
||||
result = proxy.get_annotated_filepath("test.png[temp]")
|
||||
assert isinstance(result, str), f"Expected str, got {type(result)}"
|
||||
assert "test.pn" in result, f"Filename base not in result: {result}"
|
||||
|
||||
def test_exists_annotated_filepath_returns_bool(self, proxy):
|
||||
"""Verify exists_annotated_filepath returns a boolean."""
|
||||
result = proxy.exists_annotated_filepath("nonexistent.png")
|
||||
assert isinstance(result, bool), f"Expected bool, got {type(result)}"
|
||||
|
||||
def test_exists_annotated_filepath_nonexistent_file(self, proxy):
|
||||
"""Verify exists_annotated_filepath returns False for nonexistent file."""
|
||||
result = proxy.exists_annotated_filepath("definitely_does_not_exist_12345.png")
|
||||
assert result is False, "Expected False for nonexistent file"
|
||||
|
||||
def test_exists_annotated_filepath_with_annotation(self, proxy):
|
||||
"""Verify exists_annotated_filepath works with annotation suffix."""
|
||||
# Even for nonexistent files, should return bool without error
|
||||
result = proxy.exists_annotated_filepath("test.png[output]")
|
||||
assert isinstance(result, bool), f"Expected bool, got {type(result)}"
|
||||
|
||||
def test_models_dir_property_returns_string(self, proxy):
|
||||
"""Verify models_dir property returns valid path string."""
|
||||
result = proxy.models_dir
|
||||
assert isinstance(result, str), f"Expected str, got {type(result)}"
|
||||
assert len(result) > 0, "Models directory path is empty"
|
||||
|
||||
def test_models_dir_is_absolute_path(self, proxy):
|
||||
"""Verify models_dir returns an absolute path."""
|
||||
result = proxy.models_dir
|
||||
path = Path(result)
|
||||
assert path.is_absolute(), f"Path is not absolute: {result}"
|
||||
|
||||
def test_add_model_folder_path_runs_without_error(self, proxy):
|
||||
"""Verify add_model_folder_path executes without raising."""
|
||||
test_path = "/tmp/test_models_florence2"
|
||||
# Should not raise
|
||||
proxy.add_model_folder_path("TEST_FLORENCE2", test_path)
|
||||
|
||||
def test_get_folder_paths_returns_list(self, proxy):
|
||||
"""Verify get_folder_paths returns a list."""
|
||||
# Use known folder type that should exist
|
||||
result = proxy.get_folder_paths("checkpoints")
|
||||
assert isinstance(result, list), f"Expected list, got {type(result)}"
|
||||
|
||||
def test_get_folder_paths_checkpoints_not_empty(self, proxy):
|
||||
"""Verify checkpoints folder paths list is not empty."""
|
||||
result = proxy.get_folder_paths("checkpoints")
|
||||
# Should have at least one checkpoint path registered
|
||||
assert len(result) > 0, "Checkpoints folder paths is empty"
|
||||
|
||||
def test_sealed_child_safe_uses_rpc_without_importing_folder_paths(self, monkeypatch):
|
||||
monkeypatch.setenv("PYISOLATE_CHILD", "1")
|
||||
monkeypatch.setenv("PYISOLATE_IMPORT_TORCH", "0")
|
||||
|
||||
payload = capture_sealed_singleton_imports()
|
||||
|
||||
assert payload["temp_dir"] == "/sandbox/temp"
|
||||
assert payload["models_dir"] == "/sandbox/models"
|
||||
assert "folder_paths" not in payload["modules"]
|
||||
@@ -1,209 +0,0 @@
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
def _write_pyproject(path: Path, content: str) -> None:
|
||||
path.write_text(content, encoding="utf-8")
|
||||
|
||||
|
||||
def test_load_host_policy_defaults_when_pyproject_missing(tmp_path):
|
||||
from comfy.isolation.host_policy import DEFAULT_POLICY, load_host_policy
|
||||
|
||||
policy = load_host_policy(tmp_path)
|
||||
|
||||
assert policy["sandbox_mode"] == DEFAULT_POLICY["sandbox_mode"]
|
||||
assert policy["allow_network"] == DEFAULT_POLICY["allow_network"]
|
||||
assert policy["writable_paths"] == DEFAULT_POLICY["writable_paths"]
|
||||
assert policy["readonly_paths"] == DEFAULT_POLICY["readonly_paths"]
|
||||
assert policy["whitelist"] == DEFAULT_POLICY["whitelist"]
|
||||
|
||||
|
||||
def test_load_host_policy_defaults_when_section_missing(tmp_path):
|
||||
from comfy.isolation.host_policy import DEFAULT_POLICY, load_host_policy
|
||||
|
||||
_write_pyproject(
|
||||
tmp_path / "pyproject.toml",
|
||||
"""
|
||||
[project]
|
||||
name = "ComfyUI"
|
||||
""".strip(),
|
||||
)
|
||||
|
||||
policy = load_host_policy(tmp_path)
|
||||
assert policy["sandbox_mode"] == DEFAULT_POLICY["sandbox_mode"]
|
||||
assert policy["allow_network"] == DEFAULT_POLICY["allow_network"]
|
||||
assert policy["whitelist"] == {}
|
||||
|
||||
|
||||
def test_load_host_policy_reads_values(tmp_path):
|
||||
from comfy.isolation.host_policy import load_host_policy
|
||||
|
||||
_write_pyproject(
|
||||
tmp_path / "pyproject.toml",
|
||||
"""
|
||||
[tool.comfy.host]
|
||||
sandbox_mode = "disabled"
|
||||
allow_network = true
|
||||
writable_paths = ["/tmp/a", "/tmp/b"]
|
||||
readonly_paths = ["/opt/readonly"]
|
||||
|
||||
[tool.comfy.host.whitelist]
|
||||
ExampleNode = "*"
|
||||
""".strip(),
|
||||
)
|
||||
|
||||
policy = load_host_policy(tmp_path)
|
||||
assert policy["sandbox_mode"] == "disabled"
|
||||
assert policy["allow_network"] is True
|
||||
assert policy["writable_paths"] == ["/tmp/a", "/tmp/b"]
|
||||
assert policy["readonly_paths"] == ["/opt/readonly"]
|
||||
assert policy["whitelist"] == {"ExampleNode": "*"}
|
||||
|
||||
|
||||
def test_load_host_policy_ignores_invalid_whitelist_type(tmp_path):
|
||||
from comfy.isolation.host_policy import DEFAULT_POLICY, load_host_policy
|
||||
|
||||
_write_pyproject(
|
||||
tmp_path / "pyproject.toml",
|
||||
"""
|
||||
[tool.comfy.host]
|
||||
allow_network = true
|
||||
whitelist = ["bad"]
|
||||
""".strip(),
|
||||
)
|
||||
|
||||
policy = load_host_policy(tmp_path)
|
||||
assert policy["allow_network"] is True
|
||||
assert policy["whitelist"] == DEFAULT_POLICY["whitelist"]
|
||||
|
||||
|
||||
def test_load_host_policy_ignores_invalid_sandbox_mode(tmp_path):
|
||||
from comfy.isolation.host_policy import DEFAULT_POLICY, load_host_policy
|
||||
|
||||
_write_pyproject(
|
||||
tmp_path / "pyproject.toml",
|
||||
"""
|
||||
[tool.comfy.host]
|
||||
sandbox_mode = "surprise"
|
||||
""".strip(),
|
||||
)
|
||||
|
||||
policy = load_host_policy(tmp_path)
|
||||
|
||||
assert policy["sandbox_mode"] == DEFAULT_POLICY["sandbox_mode"]
|
||||
|
||||
|
||||
def test_load_host_policy_uses_env_override_path(tmp_path, monkeypatch):
|
||||
from comfy.isolation.host_policy import load_host_policy
|
||||
|
||||
override_path = tmp_path / "host_policy_override.toml"
|
||||
_write_pyproject(
|
||||
override_path,
|
||||
"""
|
||||
[tool.comfy.host]
|
||||
sandbox_mode = "disabled"
|
||||
allow_network = true
|
||||
""".strip(),
|
||||
)
|
||||
|
||||
monkeypatch.setenv("COMFY_HOST_POLICY_PATH", str(override_path))
|
||||
|
||||
policy = load_host_policy(tmp_path / "missing-root")
|
||||
|
||||
assert policy["sandbox_mode"] == "disabled"
|
||||
assert policy["allow_network"] is True
|
||||
|
||||
|
||||
def test_disallows_host_tmp_default_or_override_defaults(tmp_path):
|
||||
from comfy.isolation.host_policy import DEFAULT_POLICY, load_host_policy
|
||||
|
||||
policy = load_host_policy(tmp_path)
|
||||
|
||||
assert "/tmp" not in DEFAULT_POLICY["writable_paths"]
|
||||
assert "/tmp" not in policy["writable_paths"]
|
||||
|
||||
|
||||
def test_disallows_host_tmp_default_or_override_config(tmp_path):
|
||||
from comfy.isolation.host_policy import load_host_policy
|
||||
|
||||
_write_pyproject(
|
||||
tmp_path / "pyproject.toml",
|
||||
"""
|
||||
[tool.comfy.host]
|
||||
writable_paths = ["/dev/shm", "/tmp", "/tmp/", "/work/cache"]
|
||||
""".strip(),
|
||||
)
|
||||
|
||||
policy = load_host_policy(tmp_path)
|
||||
|
||||
assert policy["writable_paths"] == ["/dev/shm", "/work/cache"]
|
||||
|
||||
|
||||
def test_sealed_worker_ro_import_paths_defaults_off_and_parse(tmp_path):
|
||||
from comfy.isolation.host_policy import load_host_policy
|
||||
|
||||
policy = load_host_policy(tmp_path)
|
||||
assert policy["sealed_worker_ro_import_paths"] == []
|
||||
|
||||
_write_pyproject(
|
||||
tmp_path / "pyproject.toml",
|
||||
"""
|
||||
[tool.comfy.host]
|
||||
sealed_worker_ro_import_paths = ["/home/johnj/ComfyUI", "/opt/comfy-shared"]
|
||||
""".strip(),
|
||||
)
|
||||
|
||||
policy = load_host_policy(tmp_path)
|
||||
assert policy["sealed_worker_ro_import_paths"] == [
|
||||
"/home/johnj/ComfyUI",
|
||||
"/opt/comfy-shared",
|
||||
]
|
||||
|
||||
|
||||
def test_sealed_worker_ro_import_paths_rejects_non_list_or_relative(tmp_path):
|
||||
from comfy.isolation.host_policy import load_host_policy
|
||||
|
||||
_write_pyproject(
|
||||
tmp_path / "pyproject.toml",
|
||||
"""
|
||||
[tool.comfy.host]
|
||||
sealed_worker_ro_import_paths = "/home/johnj/ComfyUI"
|
||||
""".strip(),
|
||||
)
|
||||
with pytest.raises(ValueError, match="must be a list of absolute paths"):
|
||||
load_host_policy(tmp_path)
|
||||
|
||||
_write_pyproject(
|
||||
tmp_path / "pyproject.toml",
|
||||
"""
|
||||
[tool.comfy.host]
|
||||
sealed_worker_ro_import_paths = ["relative/path"]
|
||||
""".strip(),
|
||||
)
|
||||
with pytest.raises(ValueError, match="entries must be absolute paths"):
|
||||
load_host_policy(tmp_path)
|
||||
|
||||
|
||||
def test_host_policy_path_override_controls_ro_import_paths(tmp_path, monkeypatch):
|
||||
from comfy.isolation.host_policy import load_host_policy
|
||||
|
||||
_write_pyproject(
|
||||
tmp_path / "pyproject.toml",
|
||||
"""
|
||||
[tool.comfy.host]
|
||||
sealed_worker_ro_import_paths = ["/ignored/base/path"]
|
||||
""".strip(),
|
||||
)
|
||||
override_path = tmp_path / "host_policy_override.toml"
|
||||
_write_pyproject(
|
||||
override_path,
|
||||
"""
|
||||
[tool.comfy.host]
|
||||
sealed_worker_ro_import_paths = ["/override/ro/path"]
|
||||
""".strip(),
|
||||
)
|
||||
monkeypatch.setenv("COMFY_HOST_POLICY_PATH", str(override_path))
|
||||
|
||||
policy = load_host_policy(tmp_path)
|
||||
assert policy["sealed_worker_ro_import_paths"] == ["/override/ro/path"]
|
||||
@@ -1,80 +0,0 @@
|
||||
"""Unit tests for PyIsolate isolation system initialization."""
|
||||
|
||||
import importlib
|
||||
import sys
|
||||
|
||||
from tests.isolation.singleton_boundary_helpers import (
|
||||
FakeSingletonRPC,
|
||||
reset_forbidden_singleton_modules,
|
||||
)
|
||||
|
||||
|
||||
def test_log_prefix():
|
||||
"""Verify LOG_PREFIX constant is correctly defined."""
|
||||
from comfy.isolation import LOG_PREFIX
|
||||
assert LOG_PREFIX == "]["
|
||||
assert isinstance(LOG_PREFIX, str)
|
||||
|
||||
|
||||
def test_module_initialization():
|
||||
"""Verify module initializes without errors."""
|
||||
isolation_pkg = importlib.import_module("comfy.isolation")
|
||||
assert hasattr(isolation_pkg, "LOG_PREFIX")
|
||||
assert hasattr(isolation_pkg, "initialize_proxies")
|
||||
|
||||
|
||||
class TestInitializeProxies:
|
||||
def test_initialize_proxies_runs_without_error(self):
|
||||
from comfy.isolation import initialize_proxies
|
||||
initialize_proxies()
|
||||
|
||||
def test_initialize_proxies_registers_folder_paths_proxy(self):
|
||||
from comfy.isolation import initialize_proxies
|
||||
from comfy.isolation.proxies.folder_paths_proxy import FolderPathsProxy
|
||||
initialize_proxies()
|
||||
proxy = FolderPathsProxy()
|
||||
assert proxy is not None
|
||||
assert hasattr(proxy, "get_temp_directory")
|
||||
|
||||
def test_initialize_proxies_registers_model_management_proxy(self):
|
||||
from comfy.isolation import initialize_proxies
|
||||
from comfy.isolation.proxies.model_management_proxy import ModelManagementProxy
|
||||
initialize_proxies()
|
||||
proxy = ModelManagementProxy()
|
||||
assert proxy is not None
|
||||
assert hasattr(proxy, "get_torch_device")
|
||||
|
||||
def test_initialize_proxies_can_be_called_multiple_times(self):
|
||||
from comfy.isolation import initialize_proxies
|
||||
initialize_proxies()
|
||||
initialize_proxies()
|
||||
initialize_proxies()
|
||||
|
||||
def test_dev_proxies_accessible_when_dev_mode(self, monkeypatch):
|
||||
"""Verify dev mode does not break core proxy initialization."""
|
||||
monkeypatch.setenv("PYISOLATE_DEV", "1")
|
||||
from comfy.isolation import initialize_proxies
|
||||
from comfy.isolation.proxies.folder_paths_proxy import FolderPathsProxy
|
||||
from comfy.isolation.proxies.utils_proxy import UtilsProxy
|
||||
initialize_proxies()
|
||||
folder_proxy = FolderPathsProxy()
|
||||
utils_proxy = UtilsProxy()
|
||||
assert folder_proxy is not None
|
||||
assert utils_proxy is not None
|
||||
|
||||
def test_sealed_child_safe_initialize_proxies_avoids_real_utils_import(self, monkeypatch):
|
||||
monkeypatch.setenv("PYISOLATE_CHILD", "1")
|
||||
monkeypatch.setenv("PYISOLATE_IMPORT_TORCH", "0")
|
||||
reset_forbidden_singleton_modules()
|
||||
|
||||
from pyisolate._internal import rpc_protocol
|
||||
from comfy.isolation import initialize_proxies
|
||||
|
||||
fake_rpc = FakeSingletonRPC()
|
||||
monkeypatch.setattr(rpc_protocol, "get_child_rpc_instance", lambda: fake_rpc)
|
||||
|
||||
initialize_proxies()
|
||||
|
||||
assert "comfy.utils" not in sys.modules
|
||||
assert "folder_paths" not in sys.modules
|
||||
assert "comfy_execution.progress" not in sys.modules
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user