Files
arcum42-ComfyUI_SageUtils/utils/performance_fix.py
T

219 lines
8.4 KiB
Python

#!/usr/bin/env python3
"""
Performance fix for INPUT_TYPES methods that are causing 32+ second delays during node registration.
The issue: Multiple LLM nodes call provider discovery methods during INPUT_TYPES evaluation.
These functions make network calls that can timeout, causing massive startup delays.
The fix: Create lightweight cached versions that don't block during node registration.
"""
from typing import List
from .logger import get_logger
from .llm.cache import get_llm_cache
logger = get_logger('performance.fix')
def _get_cached_models(provider: str, list_key: str) -> List[str] | None:
"""Return cached provider model list without triggering network fetch."""
cache = get_llm_cache()
return cache.peek_model_list(provider, list_key)
def get_cached_ollama_models_for_input_types() -> List[str]:
"""
Lightweight version for INPUT_TYPES - returns cached models without network calls.
Falls back to placeholder if no cache available.
"""
try:
cached = _get_cached_models('ollama', 'models')
if cached is not None:
return cached
# Return placeholder - models will be populated when cache is populated
return ["(Loading Ollama models...)"]
except Exception as e:
logger.debug(f"Failed to get cached Ollama models for INPUT_TYPES: {e}")
return ["(Ollama models unavailable)"]
def get_cached_lmstudio_models_for_input_types() -> List[str]:
"""
Lightweight version for INPUT_TYPES - returns cached models without network calls.
Falls back to placeholder if no cache available.
"""
try:
cached = _get_cached_models('lmstudio', 'models')
if cached is not None:
return cached
# Return placeholder - models will be populated when cache is populated
return ["(Loading LM Studio models...)"]
except Exception as e:
logger.debug(f"Failed to get cached LM Studio models for INPUT_TYPES: {e}")
return ["(LM Studio models unavailable)"]
def get_cached_ollama_vision_models_for_input_types() -> List[str]:
"""
Lightweight version for INPUT_TYPES - returns cached vision models without network calls.
Falls back to placeholder if no cache available.
"""
try:
cached = _get_cached_models('ollama', 'vision_models')
if cached is not None:
return cached
# Return placeholder - models will be populated when cache is populated
return ["(Loading Ollama vision models...)"]
except Exception as e:
logger.debug(f"Failed to get cached Ollama vision models for INPUT_TYPES: {e}")
return ["(Ollama vision models unavailable)"]
def get_cached_lmstudio_vision_models_for_input_types() -> List[str]:
"""
Lightweight version for INPUT_TYPES - returns cached vision models without network calls.
Falls back to placeholder if no cache available.
"""
try:
cached = _get_cached_models('lmstudio', 'vision_models')
if cached is not None:
return cached
# Return placeholder - models will be populated when cache is populated
return ["(Loading LM Studio vision models...)"]
except Exception as e:
logger.debug(f"Failed to get cached LM Studio vision models for INPUT_TYPES: {e}")
return ["(LM Studio vision models unavailable)"]
def get_cached_lmstudio_rest_models_for_input_types() -> List[str]:
"""
Lightweight version for INPUT_TYPES - returns cached LM Studio REST models without network calls.
Falls back to placeholder if no cache available.
"""
try:
cached = _get_cached_models('lmstudio_rest', 'models')
if cached is not None:
return cached
return ["(Loading LM Studio REST models...)"]
except Exception as e:
logger.debug(f"Failed to get cached LM Studio REST models for INPUT_TYPES: {e}")
return ["(LM Studio REST models unavailable)"]
def get_cached_lmstudio_rest_vision_models_for_input_types() -> List[str]:
"""
Lightweight version for INPUT_TYPES - returns cached LM Studio REST vision models without network calls.
Falls back to placeholder if no cache available.
"""
try:
cached = _get_cached_models('lmstudio_rest', 'vision_models')
if cached is not None:
return cached
return ["(Loading LM Studio REST vision models...)"]
except Exception as e:
logger.debug(f"Failed to get cached LM Studio REST vision models for INPUT_TYPES: {e}")
return ["(LM Studio REST vision models unavailable)"]
def get_cached_ollama_rest_models_for_input_types() -> List[str]:
"""
Lightweight version for INPUT_TYPES - returns cached Ollama REST models without network calls.
Falls back to placeholder if no cache available.
"""
try:
cached = _get_cached_models('ollama_rest', 'models')
if cached is not None:
return cached
return ["(Loading Ollama REST models...)"]
except Exception as e:
logger.debug(f"Failed to get cached Ollama REST models for INPUT_TYPES: {e}")
return ["(Ollama REST models unavailable)"]
def get_cached_ollama_rest_vision_models_for_input_types() -> List[str]:
"""
Lightweight version for INPUT_TYPES - returns cached Ollama REST vision models without network calls.
Falls back to placeholder if no cache available.
"""
try:
cached = _get_cached_models('ollama_rest', 'vision_models')
if cached is not None:
return cached
return ["(Loading Ollama REST vision models...)"]
except Exception as e:
logger.debug(f"Failed to get cached Ollama REST vision models for INPUT_TYPES: {e}")
return ["(Ollama REST vision models unavailable)"]
def get_cached_openai_models_for_input_types() -> List[str]:
"""
Lightweight version for INPUT_TYPES - returns cached OpenAI models without network calls.
Falls back to placeholder if no cache available.
"""
try:
cached = _get_cached_models('openai', 'models')
if cached is not None:
return cached
return ["(Loading OpenAI models...)"]
except Exception as e:
logger.debug(f"Failed to get cached OpenAI models for INPUT_TYPES: {e}")
return ["(OpenAI models unavailable)"]
def get_cached_openai_vision_models_for_input_types() -> List[str]:
"""
Lightweight version for INPUT_TYPES - returns cached OpenAI vision models without network calls.
Falls back to placeholder if no cache available.
"""
try:
cached = _get_cached_models('openai', 'vision_models')
if cached is not None:
return cached
return ["(Loading OpenAI vision models...)"]
except Exception as e:
logger.debug(f"Failed to get cached OpenAI vision models for INPUT_TYPES: {e}")
return ["(OpenAI vision models unavailable)"]
# Background task to populate cache without blocking startup
def populate_llm_cache_async():
"""
Populate LLM cache in background without blocking startup.
This should be called after node registration is complete.
"""
import threading
def _populate_cache():
try:
from .llm import service as llm
logger.info("Starting background LLM cache population...")
# These calls will populate the cache asynchronously
logger.info("Fetching LLM models to populate cache (LM Studio REST)...")
llm.get_lmstudio_rest_models()
logger.info("Fetching LLM models to populate cache (Ollama REST)...")
llm.get_ollama_rest_models()
logger.info("Fetching LLM models to populate cache (OpenAI)...")
llm.get_openai_models()
logger.info("Fetching LLM vision models to populate cache (LM Studio REST)...")
llm.get_lmstudio_rest_vision_models()
logger.info("Fetching LLM vision models to populate cache (Ollama REST)...")
llm.get_ollama_rest_vision_models()
logger.info("Fetching LLM vision models to populate cache (OpenAI)...")
llm.get_openai_vision_models()
logger.info("Background LLM cache population completed")
except Exception as e:
logger.error(f"Failed to populate LLM cache in background: {e}")
# Start background thread
thread = threading.Thread(target=_populate_cache, daemon=True)
thread.start()
logger.debug("Started background thread for LLM cache population")