Avoid loading model_management.py so early

Because model_management.py has a top-level `logging.info`, we have to
be careful not to import that file before we call `setup_logging`. If we
do, we end up having the default logging handler registered in addition
to our custom one.
This commit is contained in:
Jacob Segal 2025-07-09 20:46:42 -07:00
parent 5a36b1cb4c
commit 85dd205616

View File

@ -4,7 +4,9 @@ from PIL import Image
from enum import Enum from enum import Enum
from abc import ABC from abc import ABC
from tqdm import tqdm from tqdm import tqdm
from comfy_execution.graph import DynamicPrompt from typing import TYPE_CHECKING
if TYPE_CHECKING:
from comfy_execution.graph import DynamicPrompt
from protocol import BinaryEventTypes from protocol import BinaryEventTypes
from comfy_api import feature_flags from comfy_api import feature_flags
@ -235,7 +237,7 @@ class ProgressRegistry:
Registry that maintains node progress state and notifies registered handlers. Registry that maintains node progress state and notifies registered handlers.
""" """
def __init__(self, prompt_id: str, dynprompt: DynamicPrompt): def __init__(self, prompt_id: str, dynprompt: "DynamicPrompt"):
self.prompt_id = prompt_id self.prompt_id = prompt_id
self.dynprompt = dynprompt self.dynprompt = dynprompt
self.nodes: Dict[str, NodeProgressState] = {} self.nodes: Dict[str, NodeProgressState] = {}
@ -314,14 +316,10 @@ class ProgressRegistry:
for handler in self.handlers.values(): for handler in self.handlers.values():
handler.reset() handler.reset()
# Global registry instance # Global registry instance
global_progress_registry: ProgressRegistry = ProgressRegistry( global_progress_registry: ProgressRegistry | None = None
prompt_id="", dynprompt=DynamicPrompt({})
)
def reset_progress_state(prompt_id: str, dynprompt: "DynamicPrompt") -> None:
def reset_progress_state(prompt_id: str, dynprompt: DynamicPrompt) -> None:
global global_progress_registry global global_progress_registry
# Reset existing handlers if registry exists # Reset existing handlers if registry exists
@ -333,9 +331,17 @@ def reset_progress_state(prompt_id: str, dynprompt: DynamicPrompt) -> None:
def add_progress_handler(handler: ProgressHandler) -> None: def add_progress_handler(handler: ProgressHandler) -> None:
handler.set_registry(global_progress_registry) registry = get_progress_state()
global_progress_registry.register_handler(handler) handler.set_registry(registry)
registry.register_handler(handler)
def get_progress_state() -> ProgressRegistry: def get_progress_state() -> ProgressRegistry:
global global_progress_registry
if global_progress_registry is None:
from comfy_execution.graph import DynamicPrompt
global_progress_registry = ProgressRegistry(
prompt_id="", dynprompt=DynamicPrompt({})
)
return global_progress_registry return global_progress_registry