diff --git a/comfy/app/colab.py b/comfy/app/colab.py index 4f560b9c2..d575e1205 100644 --- a/comfy/app/colab.py +++ b/comfy/app/colab.py @@ -4,15 +4,10 @@ import stat import subprocess import threading from asyncio import Task -from typing import NamedTuple +from typing import NamedTuple, Optional import requests -from ..cmd.folder_paths import init_default_paths, folder_names_and_paths # pylint: disable=import-error -# experimental workarounds for colab -from ..cmd.main import _start_comfyui -from ..execution_context import * - class _ColabTuple(NamedTuple): tunnel: "CloudflaredTunnel" @@ -52,23 +47,22 @@ class CloudflaredTunnel: for chunk in response.iter_content(chunk_size=8192): f.write(chunk) - # Make the file executable (add execute permission for the owner) - current_permissions = os.stat(self._executable_path).st_mode - os.chmod(self._executable_path, current_permissions | stat.S_IEXEC) + current_permissions = os.stat(self._executable_path).st_mode + os.chmod(self._executable_path, current_permissions | stat.S_IEXEC) def _start_tunnel(self) -> str: """Starts the tunnel and returns the public URL.""" command = [self._executable_path, "tunnel", "--url", f"http://localhost:{self._port}", "--no-autoupdate"] - # Using DEVNULL for stderr to keep the output clean, stdout is piped self._process = subprocess.Popen( command, + bufsize=1, stdout=subprocess.PIPE, - stderr=subprocess.DEVNULL, + stderr=subprocess.STDOUT, text=True ) - for line in iter(self._process.stdout.readline, ""): + for line in self._process.stdout: if ".trycloudflare.com" in line: # The line format is typically: "INFO | https://.trycloudflare.com |" try: @@ -78,7 +72,6 @@ class CloudflaredTunnel: except IndexError: continue - # If the loop finishes without finding a URL self.stop() raise RuntimeError("Failed to start cloudflared tunnel or find URL.") @@ -124,8 +117,17 @@ def start_server_in_colab() -> str: :return: """ if len(_colab_instances) == 0: + from ..execution_context import ExecutionContext, ServerStub, comfyui_execution_context + from ..component_model.folder_path_types import FolderNames + from ..nodes.package_typing import ExportedNodes + from ..progress_types import ProgressRegistryStub comfyui_execution_context.set(ExecutionContext(server=ServerStub(), folder_names_and_paths=FolderNames(is_root=True), custom_nodes=ExportedNodes(), progress_registry=ProgressRegistryStub())) + # now we're ready to import + from ..cmd.folder_paths import init_default_paths, folder_names_and_paths + # experimental workarounds for colab + from ..cmd.main import _start_comfyui + async def colab_server_loop(): init_default_paths(folder_names_and_paths) await _start_comfyui() diff --git a/comfy/client/embedded_comfy_client.py b/comfy/client/embedded_comfy_client.py index 8113fddef..3d7bfd5bc 100644 --- a/comfy/client/embedded_comfy_client.py +++ b/comfy/client/embedded_comfy_client.py @@ -4,9 +4,11 @@ import asyncio import copy import gc import json +import logging import threading import uuid from asyncio import get_event_loop +from dataclasses import dataclass from multiprocessing import RLock from typing import Optional @@ -28,6 +30,8 @@ from ..execution_context import current_execution_context _prompt_executor = threading.local() +logger = logging.getLogger(__name__) + def _execute_prompt( prompt: dict, @@ -216,7 +220,11 @@ class Comfy: prompt: PromptDict | dict, prompt_id: Optional[str] = None, client_id: Optional[str] = None, - partial_execution_targets: Optional[list[str]] = None) -> dict: + partial_execution_targets: Optional[list[str]] = None, + progress_handler: Optional[ExecutorToClientProgress] = None) -> dict: + if isinstance(self._executor, ProcessPoolExecutor) and progress_handler is not None: + logger.debug(f"a progress_handler={progress_handler} was passed, it must be pickleable to support ProcessPoolExecutor") + progress_handler = progress_handler or self._progress_handler with self._task_count_lock: self._task_count += 1 prompt_id = prompt_id or str(uuid.uuid4()) @@ -233,7 +241,7 @@ class Comfy: client_id, carrier, # todo: a proxy object or something more sophisticated will have to be done here to restore progress notifications for ProcessPoolExecutors - None if isinstance(self._executor, ProcessPoolExecutor) else self._progress_handler, + None if isinstance(self._executor, ProcessPoolExecutor) else progress_handler, self._configuration, partial_execution_targets, )