Will Miao
260b21766b
ComfyUI runs a full gc.collect() after every prompt, and its cost scales with the number of GC-tracked objects. With a large library the hydrated model/recipe caches account for ~1M long-lived objects, so each post-prompt collection held the GIL for seconds and froze the whole server. After all scanner init tasks finish, run gc.collect() + gc.freeze() once to move the long-lived objects into the permanent generation, which the collector no longer scans. Matches the workaround the reporter verified manually.
483 lines
19 KiB
Python
483 lines
19 KiB
Python
import asyncio
|
|
import sys
|
|
import os
|
|
import logging
|
|
from .utils.logging_config import setup_logging
|
|
|
|
# Check if we're in standalone mode
|
|
standalone_mode = (
|
|
os.environ.get("LORA_MANAGER_STANDALONE", "0") == "1"
|
|
or os.environ.get("HF_HUB_DISABLE_TELEMETRY", "0") == "0"
|
|
)
|
|
|
|
# Only setup logging prefix if not in standalone mode
|
|
if not standalone_mode:
|
|
setup_logging()
|
|
|
|
from server import PromptServer # pyright: ignore[reportMissingImports]
|
|
|
|
from .config import config
|
|
from .services.model_service_factory import (
|
|
ModelServiceFactory,
|
|
register_default_model_types,
|
|
)
|
|
from .routes.recipe_routes import RecipeRoutes
|
|
from .routes.stats_routes import StatsRoutes
|
|
from .routes.update_routes import UpdateRoutes
|
|
from .routes.misc_routes import MiscRoutes
|
|
from .routes.pending_delete_routes import PendingDeleteRoutes
|
|
from .routes.preview_routes import PreviewRoutes
|
|
from .routes.example_images_routes import ExampleImagesRoutes
|
|
from .services.service_registry import ServiceRegistry
|
|
from .services.settings_manager import get_settings_manager
|
|
from .services.pending_delete_service import get_pending_delete_service
|
|
from .utils.example_images_migration import ExampleImagesMigration
|
|
from .services.websocket_manager import ws_manager
|
|
from .services.example_images_cleanup_service import ExampleImagesCleanupService
|
|
from .middleware.csp_middleware import relax_csp_for_remote_media
|
|
from .middleware.error_middleware import api_json_error
|
|
from .utils.gc_utils import freeze_long_lived_objects
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
HEADER_SIZE_LIMIT = 16384
|
|
|
|
|
|
def _sanitize_size_limit(value):
|
|
"""Return a non-negative integer size for ``handler_args`` comparisons."""
|
|
|
|
try:
|
|
coerced = int(value)
|
|
except (TypeError, ValueError):
|
|
return 0
|
|
return coerced if coerced >= 0 else 0
|
|
|
|
|
|
class _SettingsProxy:
|
|
def __init__(self):
|
|
self._manager = None
|
|
|
|
def _resolve(self):
|
|
if self._manager is None:
|
|
self._manager = get_settings_manager()
|
|
return self._manager
|
|
|
|
def get(self, *args, **kwargs):
|
|
return self._resolve().get(*args, **kwargs)
|
|
|
|
def __getattr__(self, item):
|
|
return getattr(self._resolve(), item)
|
|
|
|
|
|
settings = _SettingsProxy()
|
|
|
|
|
|
class LoraManager:
|
|
"""Main entry point for LoRA Manager plugin"""
|
|
|
|
@classmethod
|
|
def add_routes(cls):
|
|
"""Initialize and register all routes using the new refactored architecture"""
|
|
app = PromptServer.instance.app
|
|
|
|
# Register JSON error middleware for /api/* routes as the outermost
|
|
# middleware so it catches errors from all other middlewares.
|
|
if api_json_error not in app.middlewares:
|
|
app.middlewares.insert(0, api_json_error)
|
|
|
|
if relax_csp_for_remote_media not in app.middlewares:
|
|
# Ensure CSP relaxer executes after ComfyUI's block_external_middleware so it can
|
|
# see and extend the restrictive header instead of being overwritten by it.
|
|
block_middleware_index = next(
|
|
(
|
|
idx
|
|
for idx, middleware in enumerate(app.middlewares)
|
|
if getattr(middleware, "__name__", "")
|
|
== "block_external_middleware"
|
|
),
|
|
None,
|
|
)
|
|
|
|
if block_middleware_index is None:
|
|
app.middlewares.append(relax_csp_for_remote_media)
|
|
else:
|
|
app.middlewares.insert(
|
|
block_middleware_index, relax_csp_for_remote_media
|
|
)
|
|
|
|
# Increase allowed header sizes so browsers with large localhost cookie
|
|
# jars (multiple UIs on 127.0.0.1) don't trip aiohttp's 8KB default
|
|
# limits. Cookies for unrelated apps are still sent to the plugin and
|
|
# may otherwise raise LineTooLong errors when the request parser reads
|
|
# them. Preserve any previously configured handler arguments while
|
|
# ensuring our minimum sizes are applied.
|
|
handler_args = getattr(app, "_handler_args", {}) or {}
|
|
updated_handler_args = dict(handler_args)
|
|
updated_handler_args["max_field_size"] = max(
|
|
_sanitize_size_limit(handler_args.get("max_field_size", 0)),
|
|
HEADER_SIZE_LIMIT,
|
|
)
|
|
updated_handler_args["max_line_size"] = max(
|
|
_sanitize_size_limit(handler_args.get("max_line_size", 0)),
|
|
HEADER_SIZE_LIMIT,
|
|
)
|
|
app._handler_args = updated_handler_args
|
|
|
|
# Configure aiohttp access logger to be less verbose
|
|
logging.getLogger("aiohttp.access").setLevel(logging.WARNING)
|
|
|
|
# Add specific suppression for connection reset errors
|
|
class ConnectionResetFilter(logging.Filter):
|
|
def filter(self, record):
|
|
# Filter out connection reset errors that are not critical
|
|
if "ConnectionResetError" in str(record.getMessage()):
|
|
return False
|
|
if "_call_connection_lost" in str(record.getMessage()):
|
|
return False
|
|
if "WinError 10054" in str(record.getMessage()):
|
|
return False
|
|
return True
|
|
|
|
# Apply the filter to asyncio logger
|
|
asyncio_logger = logging.getLogger("asyncio")
|
|
asyncio_logger.addFilter(ConnectionResetFilter())
|
|
|
|
# Add static route for example images if the path exists in settings
|
|
example_images_path = settings.get("example_images_path")
|
|
logger.info(f"Example images path: {example_images_path}")
|
|
if example_images_path and os.path.exists(example_images_path):
|
|
app.router.add_static("/example_images_static", example_images_path)
|
|
logger.info(
|
|
f"Added static route for example images: /example_images_static -> {example_images_path}"
|
|
)
|
|
|
|
# Add static route for locales JSON files
|
|
if os.path.exists(config.i18n_path):
|
|
app.router.add_static("/locales", config.i18n_path)
|
|
logger.info(
|
|
f"Added static route for locales: /locales -> {config.i18n_path}"
|
|
)
|
|
|
|
# Add static route for plugin assets
|
|
app.router.add_static("/loras_static", config.static_path)
|
|
|
|
# Register default model types with the factory
|
|
register_default_model_types()
|
|
|
|
# Setup all model routes using the factory
|
|
ModelServiceFactory.setup_all_routes(app)
|
|
|
|
# Setup non-model-specific routes
|
|
stats_routes = StatsRoutes()
|
|
stats_routes.setup_routes(app)
|
|
RecipeRoutes.setup_routes(app)
|
|
UpdateRoutes.setup_routes(app)
|
|
MiscRoutes.setup_routes(app)
|
|
PendingDeleteRoutes.setup_routes(app)
|
|
ExampleImagesRoutes.setup_routes(app, ws_manager=ws_manager)
|
|
PreviewRoutes.setup_routes(app)
|
|
|
|
# Setup WebSocket routes that are shared across all model types
|
|
app.router.add_get("/ws/fetch-progress", ws_manager.handle_connection)
|
|
app.router.add_get(
|
|
"/ws/download-progress", ws_manager.handle_download_connection
|
|
)
|
|
app.router.add_get("/ws/init-progress", ws_manager.handle_init_connection)
|
|
|
|
# Schedule service initialization
|
|
app.on_startup.append(lambda app: cls._initialize_services())
|
|
|
|
# Add cleanup
|
|
app.on_shutdown.append(cls._cleanup)
|
|
|
|
@classmethod
|
|
async def _initialize_services(cls):
|
|
"""Initialize all services using the ServiceRegistry"""
|
|
try:
|
|
# Initialize CivitaiClient first to ensure it's ready for other services
|
|
await ServiceRegistry.get_civitai_client()
|
|
|
|
# Register DownloadManager with ServiceRegistry
|
|
await ServiceRegistry.get_download_manager()
|
|
|
|
# Initialize DownloadQueueService for persistent queue/history
|
|
await ServiceRegistry.get_download_queue_service()
|
|
|
|
await ServiceRegistry.get_backup_service()
|
|
|
|
from .services.metadata_service import initialize_metadata_providers
|
|
|
|
await initialize_metadata_providers()
|
|
|
|
# Initialize WebSocket manager
|
|
await ServiceRegistry.get_websocket_manager()
|
|
|
|
# Preload LLM model catalog (background task, non-blocking)
|
|
from .services.llm_service import LLMService
|
|
await LLMService.get_instance()
|
|
|
|
# Initialize scanners in background
|
|
lora_scanner = await ServiceRegistry.get_lora_scanner()
|
|
checkpoint_scanner = await ServiceRegistry.get_checkpoint_scanner()
|
|
embedding_scanner = await ServiceRegistry.get_embedding_scanner()
|
|
other_scanner = await ServiceRegistry.get_other_scanner()
|
|
|
|
# Initialize recipe scanner if needed
|
|
recipe_scanner = await ServiceRegistry.get_recipe_scanner()
|
|
|
|
# Create low-priority initialization tasks
|
|
init_tasks = [
|
|
asyncio.create_task(
|
|
lora_scanner.initialize_in_background(), name="lora_cache_init"
|
|
),
|
|
asyncio.create_task(
|
|
checkpoint_scanner.initialize_in_background(),
|
|
name="checkpoint_cache_init",
|
|
),
|
|
asyncio.create_task(
|
|
embedding_scanner.initialize_in_background(),
|
|
name="embedding_cache_init",
|
|
),
|
|
asyncio.create_task(
|
|
other_scanner.initialize_in_background(),
|
|
name="other_cache_init",
|
|
),
|
|
asyncio.create_task(
|
|
recipe_scanner.initialize_in_background(), name="recipe_cache_init"
|
|
),
|
|
]
|
|
|
|
await ExampleImagesMigration.check_and_run_migrations()
|
|
|
|
# Schedule post-initialization tasks to run after scanners complete
|
|
asyncio.create_task(
|
|
cls._run_post_initialization_tasks(init_tasks), name="post_init_tasks"
|
|
)
|
|
|
|
# Startup sweep: purge pending-delete batches that expired during a
|
|
# previous run. Non-blocking (fire-and-forget); purge_expired only
|
|
# removes already-expired batches, so a staged undo that survived a
|
|
# restart stays restorable. scan_roots=True runs the reconciliation
|
|
# pass first so leftover batches (the in-process registry is empty
|
|
# after a restart) are re-discovered on disk. Covers both plugin
|
|
# and standalone modes (StandaloneLoraManager reuses this
|
|
# classmethod).
|
|
pending_delete_service = await get_pending_delete_service()
|
|
asyncio.create_task(
|
|
pending_delete_service.purge_expired(scan_roots=True),
|
|
name="pending_delete_startup_sweep",
|
|
)
|
|
|
|
logger.debug(
|
|
"LoRA Manager: All services initialized and background tasks scheduled"
|
|
)
|
|
|
|
except Exception as e:
|
|
logger.error(
|
|
f"LoRA Manager: Error initializing services: {e}", exc_info=True
|
|
)
|
|
|
|
@classmethod
|
|
async def _run_post_initialization_tasks(cls, init_tasks):
|
|
"""Run post-initialization tasks after all scanners complete"""
|
|
try:
|
|
logger.debug(
|
|
"LoRA Manager: Waiting for scanner initialization to complete..."
|
|
)
|
|
|
|
# Wait for all scanner initialization tasks to complete
|
|
await asyncio.gather(*init_tasks, return_exceptions=True)
|
|
|
|
# The caches are hydrated now: move their long-lived objects out of
|
|
# GC tracking so ComfyUI's post-prompt gc.collect() no longer walks
|
|
# them on every prompt (issue #1147).
|
|
freeze_long_lived_objects("cache initialization")
|
|
|
|
logger.debug(
|
|
"LoRA Manager: Scanner initialization completed, starting post-initialization tasks..."
|
|
)
|
|
|
|
# Run post-initialization tasks
|
|
post_tasks = [
|
|
asyncio.create_task(
|
|
cls._cleanup_backup_files(), name="cleanup_bak_files"
|
|
),
|
|
# Add more post-initialization tasks here as needed
|
|
# asyncio.create_task(cls._another_post_task(), name='another_task'),
|
|
]
|
|
|
|
# Run all post-initialization tasks
|
|
results = await asyncio.gather(*post_tasks, return_exceptions=True)
|
|
|
|
# Log results
|
|
for i, result in enumerate(results):
|
|
task_name = post_tasks[i].get_name()
|
|
if isinstance(result, Exception):
|
|
logger.error(
|
|
f"Post-initialization task '{task_name}' failed: {result}"
|
|
)
|
|
else:
|
|
logger.debug(
|
|
f"Post-initialization task '{task_name}' completed successfully"
|
|
)
|
|
|
|
logger.debug("LoRA Manager: All post-initialization tasks completed")
|
|
|
|
except Exception as e:
|
|
logger.error(
|
|
f"LoRA Manager: Error in post-initialization tasks: {e}", exc_info=True
|
|
)
|
|
|
|
@classmethod
|
|
async def _cleanup_backup_files(cls):
|
|
"""Clean up .bak files in all model roots"""
|
|
try:
|
|
logger.debug("Starting cleanup of .bak files in model directories...")
|
|
|
|
# Collect all model roots
|
|
all_roots = set()
|
|
all_roots.update(config.loras_roots)
|
|
all_roots.update(config.base_models_roots or [])
|
|
all_roots.update(config.embeddings_roots or [])
|
|
all_roots.update(config.other_roots or [])
|
|
|
|
total_deleted = 0
|
|
total_size_freed = 0
|
|
|
|
for root_path in all_roots:
|
|
if not os.path.exists(root_path):
|
|
continue
|
|
|
|
try:
|
|
(
|
|
deleted_count,
|
|
size_freed,
|
|
) = await cls._cleanup_backup_files_in_directory(root_path)
|
|
total_deleted += deleted_count
|
|
total_size_freed += size_freed
|
|
|
|
if deleted_count > 0:
|
|
logger.debug(
|
|
f"Cleaned up {deleted_count} .bak files in {root_path} (freed {size_freed / (1024 * 1024):.2f} MB)"
|
|
)
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error cleaning up .bak files in {root_path}: {e}")
|
|
|
|
# Yield control periodically
|
|
await asyncio.sleep(0.01)
|
|
|
|
if total_deleted > 0:
|
|
logger.debug(
|
|
f"Backup cleanup completed: removed {total_deleted} .bak files, freed {total_size_freed / (1024 * 1024):.2f} MB total"
|
|
)
|
|
else:
|
|
logger.debug("Backup cleanup completed: no .bak files found")
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error during backup file cleanup: {e}", exc_info=True)
|
|
|
|
@classmethod
|
|
async def _cleanup_backup_files_in_directory(cls, directory_path: str):
|
|
"""Clean up .bak files in a specific directory recursively
|
|
|
|
Args:
|
|
directory_path: Path to the directory to clean
|
|
|
|
Returns:
|
|
Tuple[int, int]: (number of files deleted, total size freed in bytes)
|
|
"""
|
|
deleted_count = 0
|
|
size_freed = 0
|
|
visited_paths = set()
|
|
|
|
def cleanup_recursive(path):
|
|
nonlocal deleted_count, size_freed
|
|
|
|
try:
|
|
real_path = os.path.realpath(path)
|
|
if real_path in visited_paths:
|
|
return
|
|
visited_paths.add(real_path)
|
|
|
|
with os.scandir(path) as it:
|
|
for entry in it:
|
|
try:
|
|
if entry.is_file(
|
|
follow_symlinks=True
|
|
) and entry.name.endswith(".bak"):
|
|
file_size = entry.stat().st_size
|
|
os.remove(entry.path)
|
|
deleted_count += 1
|
|
size_freed += file_size
|
|
logger.debug(f"Deleted .bak file: {entry.path}")
|
|
|
|
elif entry.is_dir(follow_symlinks=True):
|
|
cleanup_recursive(entry.path)
|
|
|
|
except Exception as e:
|
|
logger.warning(
|
|
f"Could not delete .bak file {entry.path}: {e}"
|
|
)
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error scanning directory {path} for .bak files: {e}")
|
|
|
|
# Run the recursive cleanup in a thread pool to avoid blocking
|
|
loop = asyncio.get_event_loop()
|
|
await loop.run_in_executor(None, cleanup_recursive, directory_path)
|
|
|
|
return deleted_count, size_freed
|
|
|
|
@classmethod
|
|
async def _cleanup_example_images_folders(cls):
|
|
"""Invoke the example images cleanup service for manual execution."""
|
|
try:
|
|
service = ExampleImagesCleanupService()
|
|
result = await service.cleanup_example_image_folders()
|
|
|
|
if result.get("success"):
|
|
logger.debug(
|
|
"Manual example images cleanup completed: moved=%s",
|
|
result.get("moved_total"),
|
|
)
|
|
elif result.get("partial_success"):
|
|
logger.warning(
|
|
"Manual example images cleanup partially succeeded: moved=%s failures=%s",
|
|
result.get("moved_total"),
|
|
result.get("move_failures"),
|
|
)
|
|
else:
|
|
logger.debug(
|
|
"Manual example images cleanup skipped or failed: %s",
|
|
result.get("error", "no changes"),
|
|
)
|
|
|
|
return result
|
|
|
|
except Exception as e: # pragma: no cover - defensive guard
|
|
logger.error(f"Error during example images cleanup: {e}", exc_info=True)
|
|
return {
|
|
"success": False,
|
|
"error": str(e),
|
|
"error_code": "unexpected_error",
|
|
}
|
|
|
|
@classmethod
|
|
async def _cleanup(cls, app):
|
|
"""Cleanup resources using ServiceRegistry"""
|
|
try:
|
|
logger.info("LoRA Manager: Cleaning up services")
|
|
|
|
# Cancel any in-flight scanner initialization tasks so thread-pool
|
|
# workers (e.g. _initialize_cache_sync) can break out of their loops
|
|
# when the server shuts down (e.g. Ctrl+C on WSL).
|
|
for name in ("lora_scanner", "checkpoint_scanner", "embedding_scanner", "other_scanner"):
|
|
scanner = ServiceRegistry.get_service_sync(name)
|
|
if scanner is not None and hasattr(scanner, "cancel_task"):
|
|
scanner.cancel_task()
|
|
logger.debug("LoRA Manager: Cancelled %s", name)
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error during cleanup: {e}", exc_info=True)
|