From 6e4181a7bb5e2ef147aa8e1d0845098a709306a4 Mon Sep 17 00:00:00 2001 From: John Pollock Date: Fri, 15 Aug 2025 08:25:18 -0500 Subject: [PATCH] Refactor: Remove debugging and memory audit utilities This commit removes several utility modules used for debugging, memory inspection, and hardware information gathering. These tools are no longer required and their removal simplifies the codebase. The following files have been deleted: - `debug_utils.py` - `device_memory_audit.py` - `hardware_info.py` - `model_sig.py` Additionally, the call to log memory usage on startup has been removed from `__init__.py`. --- __init__.py | 5 - debug_utils.py | 27 ---- device_memory_audit.py | 351 ----------------------------------------- hardware_info.py | 83 ---------- model_sig.py | 23 --- pyproject.toml | 2 +- 6 files changed, 1 insertion(+), 490 deletions(-) delete mode 100644 debug_utils.py delete mode 100644 device_memory_audit.py delete mode 100644 hardware_info.py delete mode 100644 model_sig.py diff --git a/__init__.py b/__init__.py index d1e710a..721c769 100644 --- a/__init__.py +++ b/__init__.py @@ -368,8 +368,3 @@ logger.info(dash_line) logger.info(f"[MultiGPU] Registration complete. Final mappings: {', '.join(NODE_CLASS_MAPPINGS.keys())}") - -# --- Memory Logging Test --- -from .debug_utils import log_memory_usage -logger.debug("ComfyUI Startup Memory Log:") -log_memory_usage("ComfyUI Startup") diff --git a/debug_utils.py b/debug_utils.py deleted file mode 100644 index 97695d8..0000000 --- a/debug_utils.py +++ /dev/null @@ -1,27 +0,0 @@ -import logging -# Utility to get memory stats, using self-contained hardware_info module. - -def log_memory_usage(label=""): - try: - from .hardware_info import CHardwareInfo - - hardware_info = CHardwareInfo(switchRAM=True, switchVRAM=True) - status = hardware_info.getStatus() - - ram_used_gb = status.get('ram_used', 0) / (1024**3) - ram_total_gb = status.get('ram_total', 0) / (1024**3) - - log_message = f"[MultiGPU] {label} | RAM Used: {ram_used_gb:.2f}/{ram_total_gb:.2f} GB" - - if 'gpus' in status: - for i, gpu in enumerate(status['gpus']): - vram_used_gb = gpu.get('vram_used', 0) / (1024**3) - vram_total_gb = gpu.get('vram_total', 0) / (1024**3) - log_message += f" | VRAM cuda:{i}: {vram_used_gb:.2f}/{vram_total_gb:.2f} GB" - - logging.debug(log_message) - - except ImportError: - logging.warning("[MultiGPU] Could not import local CHardwareInfo. Cannot log memory usage.") - except Exception as e: - logging.error(f"[MultiGPU] Error getting memory usage: {e}") diff --git a/device_memory_audit.py b/device_memory_audit.py deleted file mode 100644 index 025991e..0000000 --- a/device_memory_audit.py +++ /dev/null @@ -1,351 +0,0 @@ -""" -Device Memory Audit Utility for ComfyUI MultiGPU -Provides tools to inspect and monitor GPU/CPU memory usage and model placement -""" - -import torch -import gc -import psutil -import logging -from collections import defaultdict -from typing import Dict, List, Tuple, Any -import comfy.model_management as mm - -def format_bytes(bytes_value: int) -> str: - """Convert bytes to human-readable format""" - for unit in ['B', 'KB', 'MB', 'GB', 'TB']: - if bytes_value < 1024.0: - return f"{bytes_value:.2f} {unit}" - bytes_value /= 1024.0 - return f"{bytes_value:.2f} PB" - -def get_device_memory_info() -> Dict[str, Dict[str, Any]]: - """ - Get current memory usage for all available devices - - Returns: - Dict with device names as keys and memory info as values - """ - memory_info = {} - - # Check CUDA devices - if torch.cuda.is_available(): - for i in range(torch.cuda.device_count()): - device_name = f"cuda:{i}" - torch.cuda.synchronize(i) # Ensure accurate memory reading - - allocated = torch.cuda.memory_allocated(i) - reserved = torch.cuda.memory_reserved(i) - total = torch.cuda.get_device_properties(i).total_memory - free = total - allocated - - memory_info[device_name] = { - "allocated": allocated, - "allocated_str": format_bytes(allocated), - "reserved": reserved, - "reserved_str": format_bytes(reserved), - "total": total, - "total_str": format_bytes(total), - "free": free, - "free_str": format_bytes(free), - "usage_percent": (allocated / total * 100) if total > 0 else 0 - } - - # Check CPU/System memory - vm = psutil.virtual_memory() - memory_info["cpu"] = { - "allocated": vm.used, - "allocated_str": format_bytes(vm.used), - "total": vm.total, - "total_str": format_bytes(vm.total), - "free": vm.available, - "free_str": format_bytes(vm.available), - "usage_percent": vm.percent - } - - return memory_info - -def audit_torch_tensors() -> Dict[str, List[Dict[str, Any]]]: - """ - Find all torch tensors in memory and group by device - - Returns: - Dict with device names as keys and list of tensor info as values - """ - tensors_by_device = defaultdict(list) - - # Iterate through all objects in memory - for obj in gc.get_objects(): - try: - if torch.is_tensor(obj): - device_str = str(obj.device) - size_bytes = obj.element_size() * obj.numel() - - tensor_info = { - "id": id(obj), - "shape": list(obj.shape), - "dtype": str(obj.dtype), - "size_bytes": size_bytes, - "size_str": format_bytes(size_bytes), - "requires_grad": obj.requires_grad, - "is_leaf": obj.is_leaf - } - - tensors_by_device[device_str].append(tensor_info) - except: - # Some objects might not be accessible - pass - - # Sort tensors by size for each device - for device in tensors_by_device: - tensors_by_device[device].sort(key=lambda x: x["size_bytes"], reverse=True) - - return dict(tensors_by_device) - -def audit_model_placement(model) -> Dict[str, List[Tuple[str, str, int]]]: - """ - Audit where each layer of a model is placed - - Args: - model: PyTorch model to audit - - Returns: - Dict with device names as keys and list of (layer_name, layer_type, size) tuples - """ - device_layers = defaultdict(list) - - for name, module in model.named_modules(): - # Skip container modules - if len(list(module.children())) > 0: - continue - - # Find device for this module - device = None - size_bytes = 0 - - # Check parameters - for param_name, param in module.named_parameters(recurse=False): - if param is not None: - device = str(param.device) - size_bytes += param.element_size() * param.numel() - - if device: - layer_type = type(module).__name__ - device_layers[device].append((name, layer_type, size_bytes)) - - return dict(device_layers) - -def audit_comfy_models() -> Dict[str, Any]: - """ - Audit ComfyUI's loaded models and their placement - - Returns: - Dict with model information - """ - model_info = {} - - # Try to get loaded models from ComfyUI's model management - try: - # Access currently loaded models if available - if hasattr(mm, 'current_loaded_models'): - loaded_models = mm.current_loaded_models - - for i, model_data in enumerate(loaded_models): - model_name = f"model_{i}" - if hasattr(model_data, 'model'): - model = model_data.model - placement = audit_model_placement(model) - - total_size = 0 - for device_layers in placement.values(): - for _, _, size in device_layers: - total_size += size - - model_info[model_name] = { - "type": type(model).__name__, - "device_placement": placement, - "total_size": total_size, - "total_size_str": format_bytes(total_size) - } - except Exception as e: - logging.warning(f"[MultiGPU] Could not audit ComfyUI models: {e}") - - return model_info - -def print_memory_audit(detailed=False): - """ - Print a formatted memory audit report - - Args: - detailed: If True, include tensor-level details - """ - logging.info("\n" + "=" * 60) - logging.info(" DEVICE MEMORY AUDIT REPORT") - logging.info("=" * 60) - - # Device memory overview - memory_info = get_device_memory_info() - logging.info("\nšŸ“Š MEMORY USAGE BY DEVICE:") - logging.info("-" * 60) - - for device, info in memory_info.items(): - logging.info(f"\n{device.upper()}:") - logging.info(f" Total: {info['total_str']:>12}") - logging.info(f" Allocated: {info['allocated_str']:>12} ({info['usage_percent']:.1f}%)") - logging.info(f" Free: {info['free_str']:>12}") - if 'reserved_str' in info: - logging.info(f" Reserved: {info['reserved_str']:>12}") - - # Tensor audit - if detailed: - tensors = audit_torch_tensors() - logging.info("\nšŸ“¦ TENSORS BY DEVICE:") - logging.info("-" * 60) - - for device, tensor_list in tensors.items(): - total_size = sum(t["size_bytes"] for t in tensor_list) - logging.info(f"\n{device}: {len(tensor_list)} tensors, {format_bytes(total_size)} total") - - # Show top 5 largest tensors - for i, tensor in enumerate(tensor_list[:5]): - logging.info(f" #{i+1}: Shape {tensor['shape']}, {tensor['size_str']}, {tensor['dtype']}") - - # ComfyUI model audit - models = audit_comfy_models() - if models: - logging.info("\nšŸ¤– COMFYUI LOADED MODELS:") - logging.info("-" * 60) - - for model_name, info in models.items(): - logging.info(f"\n{model_name} ({info['type']}):") - logging.info(f" Total size: {info['total_size_str']}") - - for device, layers in info['device_placement'].items(): - device_size = sum(size for _, _, size in layers) - logging.info(f" {device}: {len(layers)} layers, {format_bytes(device_size)}") - - logging.info("\n" + "=" * 60) - -def get_memory_snapshot() -> Dict[str, Any]: - """ - Get a complete memory snapshot for benchmarking - - Returns: - Dict containing all memory information - """ - return { - "device_memory": get_device_memory_info(), - "tensors": audit_torch_tensors(), - "models": audit_comfy_models() - } - -def compare_snapshots(before: Dict[str, Any], after: Dict[str, Any]) -> Dict[str, Any]: - """ - Compare two memory snapshots to see what changed - - Args: - before: Snapshot taken before an operation - after: Snapshot taken after an operation - - Returns: - Dict with differences between snapshots - """ - differences = {} - - # Compare device memory - memory_diff = {} - for device in after["device_memory"]: - if device in before["device_memory"]: - before_mem = before["device_memory"][device]["allocated"] - after_mem = after["device_memory"][device]["allocated"] - diff = after_mem - before_mem - - memory_diff[device] = { - "before": format_bytes(before_mem), - "after": format_bytes(after_mem), - "difference": format_bytes(abs(diff)), - "increased": diff > 0 - } - - differences["memory_changes"] = memory_diff - - # Compare tensor counts - tensor_diff = {} - for device in after["tensors"]: - after_count = len(after["tensors"][device]) - before_count = len(before["tensors"].get(device, [])) - - if after_count != before_count: - tensor_diff[device] = { - "before": before_count, - "after": after_count, - "difference": after_count - before_count - } - - differences["tensor_count_changes"] = tensor_diff - - return differences - -def print_snapshot_comparison(before: Dict[str, Any], after: Dict[str, Any], label: str = "Operation"): - """ - Print a formatted comparison of two snapshots - - Args: - before: Snapshot before operation - after: Snapshot after operation - label: Label for the operation being measured - """ - diff = compare_snapshots(before, after) - - logging.info(f"\nšŸ“ˆ MEMORY CHANGES AFTER {label}:") - logging.info("-" * 60) - - for device, changes in diff["memory_changes"].items(): - symbol = "ā†‘" if changes["increased"] else "ā†“" - logging.info(f"{device}: {changes['before']} ā†’ {changes['after']} ({symbol} {changes['difference']})") - - if diff["tensor_count_changes"]: - logging.info("\nTensor count changes:") - for device, changes in diff["tensor_count_changes"].items(): - logging.info(f"{device}: {changes['before']} ā†’ {changes['after']} ({changes['difference']:+d})") - -# Example usage functions for benchmarking -def benchmark_memory_points(): - """ - Example function showing how to capture the 4 key memory points - """ - logging.info("\nšŸ” Starting Memory Benchmark...") - - # 1. Pre-load memory - snapshot_pre = get_memory_snapshot() - print_memory_audit(detailed=False) - - # User would load model here - logging.info("\n[Load your model now]") - input("Press Enter when model is loaded...") - - # 2. Post-load memory - snapshot_post_load = get_memory_snapshot() - print_snapshot_comparison(snapshot_pre, snapshot_post_load, "MODEL LOAD") - - # User would run inference here - logging.info("\n[Run inference now]") - input("Press Enter when inference is complete...") - - # 3. Active inference memory (captured during) - snapshot_active = get_memory_snapshot() - print_snapshot_comparison(snapshot_post_load, snapshot_active, "INFERENCE") - - # 4. Post-generation memory - logging.info("\n[Waiting for cleanup...]") - torch.cuda.empty_cache() # Force cleanup - gc.collect() - - snapshot_post_gen = get_memory_snapshot() - print_snapshot_comparison(snapshot_active, snapshot_post_gen, "CLEANUP") - - logging.info("\nāœ… Benchmark complete!") - -if __name__ == "__main__": - # Run a basic audit when script is executed directly - print_memory_audit(detailed=True) diff --git a/hardware_info.py b/hardware_info.py deleted file mode 100644 index fb92e7e..0000000 --- a/hardware_info.py +++ /dev/null @@ -1,83 +0,0 @@ -import torch -import logging -import psutil -import comfy.model_management as mm - -# --- Surgically lifted from ComfyUI-Crystools by Crystian --- -# This is a self-contained version of the necessary hardware monitoring classes -# to avoid cross-node import issues related to load order. - -# --- From gpu.py --- -class CGPUInfo: - def __init__(self): - self.pynvmlLoaded = False - self.cudaDevicesFound = 0 - self.gpus = [] - - try: - import pynvml - self.pynvml = pynvml - self.pynvml.nvmlInit() - self.pynvmlLoaded = True - self.cudaDevicesFound = self.pynvml.nvmlDeviceGetCount() - - for i in range(self.cudaDevicesFound): - handle = self.pynvml.nvmlDeviceGetHandleByIndex(i) - gpu_name = self.pynvml.nvmlDeviceGetName(handle) - self.gpus.append({'index': i, 'name': gpu_name}) - - except ImportError: - logging.warning("[MultiGPU Hardware] pynvml not installed, VRAM monitoring disabled.") - except Exception as e: - logging.error(f"[MultiGPU Hardware] Could not initialize pynvml: {e}") - - def getStatus(self): - gpus_status = [] - if self.pynvmlLoaded: - for i in range(self.cudaDevicesFound): - handle = self.pynvml.nvmlDeviceGetHandleByIndex(i) - try: - mem = self.pynvml.nvmlDeviceGetMemoryInfo(handle) - gpus_status.append({ - 'vram_total': mem.total, - 'vram_used': mem.used, - 'vram_used_percent': mem.used / mem.total * 100 if mem.total > 0 else 0, - }) - except Exception as e: - logging.error(f"Could not get VRAM info for GPU {i}: {e}") - gpus_status.append({'vram_total': 0, 'vram_used': 0, 'vram_used_percent': 0}) - - return { - 'device_type': 'cuda' if self.pynvmlLoaded else 'cpu', - 'gpus': gpus_status, - } - -# --- From hardware.py --- -class CHardwareInfo: - def __init__(self, switchRAM=False, switchVRAM=False): - self.switchRAM = switchRAM - self.GPUInfo = CGPUInfo() - self.switchVRAM = switchVRAM - - def getStatus(self): - ramTotal = -1 - ramUsed = -1 - ramUsedPercent = -1 - - if self.switchRAM: - ram = psutil.virtual_memory() - ramTotal = ram.total - ramUsed = ram.used - ramUsedPercent = ram.percent - - gpu_status = {'device_type': 'cpu', 'gpus': []} - if self.switchVRAM: - gpu_status = self.GPUInfo.getStatus() - - return { - 'ram_total': ramTotal, - 'ram_used': ramUsed, - 'ram_used_percent': ramUsedPercent, - 'device_type': gpu_status['device_type'], - 'gpus': gpu_status['gpus'], - } diff --git a/model_sig.py b/model_sig.py deleted file mode 100644 index af5cb89..0000000 --- a/model_sig.py +++ /dev/null @@ -1,23 +0,0 @@ -def get_model_type(model_patcher): - """ - Identifies the model type using a multi-layered approach for robustness. - It first checks the diffusion model's class name, then falls back to the - model_type enum. - """ - if hasattr(model_patcher, 'model') and hasattr(model_patcher.model, 'diffusion_model'): - class_name = type(model_patcher.model.diffusion_model).__name__ - - # Prioritize class name for accuracy - if "Flux" in class_name: - return "FLUX" - if "Qwen" in class_name: - return "QWEN" - if "WanModel" in class_name: - return "WANVIDEO" - - # Fallback to the model_type enum for other cases - if hasattr(model_patcher, 'model') and hasattr(model_patcher.model, 'model_type'): - # model_type is an enum, so we return its name as a string - return model_patcher.model.model_type.name - - return "UNKNOWN" diff --git a/pyproject.toml b/pyproject.toml index f91f8f9..e90aeb7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,7 +1,7 @@ [project] name = "comfyui-multigpu" description = "Provides a suite of custom nodes to manage multiple GPUs for ComfyUI, including advanced model offloading for both GGUF and Safetensor formats with DisTorch, and bespoke MultiGPU support for WanVideoWrapper and other custom nodes." -version = "2.0.0" +version = "2.0.1" license = {file = "LICENSE"} [project.urls]