mirror of
https://github.com/comfyanonymous/ComfyUI.git
synced 2026-04-01 06:23:37 +08:00
88 lines
1.9 KiB
Python
88 lines
1.9 KiB
Python
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"]
|