From 2a39183ffa2f0b993844a5c4a97f6e9667ed10b9 Mon Sep 17 00:00:00 2001 From: LuqP2 Date: Tue, 23 Dec 2025 14:22:59 -0300 Subject: [PATCH] feat: add GPU/performance metrics collection and tracking - Add collect_gpu_metrics() for VRAM/GPU auto-detection (CUDA/MPS/CPU) - Add collect_version_info() for ComfyUI/PyTorch/Python versions - Add optional override inputs (vram_peak_mb, gpu_device_override, generation_time_override) - Add auto-timing wrapper with time.time() in save_images() - Calculate derived metrics (generation_time_ms, steps_per_second) - Update build_imh_metadata() analytics section with Tier 1/2/3 metrics - Silent failures - never interrupts generation - Reset CUDA peak memory stats after collection --- metadata_utils.py | 138 ++++++++++++++++++++++++++++++++++++++++++- metahub_save_node.py | 57 ++++++++++++++++++ 2 files changed, 194 insertions(+), 1 deletion(-) diff --git a/metadata_utils.py b/metadata_utils.py index 1d5af3c..1b2276e 100644 --- a/metadata_utils.py +++ b/metadata_utils.py @@ -174,6 +174,128 @@ def calculate_lora_hashes(lora_list: List[dict]) -> Dict[str, str]: return hashes +# ============================================================================ +# PERFORMANCE METRICS COLLECTION +# ============================================================================ + +def collect_gpu_metrics() -> Dict[str, Any]: + """ + Auto-detects GPU metrics (VRAM peak, device name). + + Handles CUDA, MPS (Mac Metal), and CPU-only setups. + Silent failures - never interrupts generation. + + Returns: + dict with keys: + - vram_peak_mb: Peak VRAM usage in MB (None if unavailable) + - gpu_device: GPU device name string + - gpu_available: Boolean indicating if GPU is available + """ + metrics = { + "vram_peak_mb": None, + "gpu_device": None, + "gpu_available": False, + } + + try: + import torch + + # Check CUDA availability + if torch.cuda.is_available(): + metrics["gpu_available"] = True + + # Get GPU device name + try: + device_name = torch.cuda.get_device_name(0) + metrics["gpu_device"] = device_name + except Exception as e: + print(f"[MetaHub] Warning: Could not get GPU name: {e}") + + # Get peak VRAM usage (convert bytes to MB) + try: + vram_bytes = torch.cuda.max_memory_allocated(0) + vram_mb = vram_bytes / (1024 * 1024) + metrics["vram_peak_mb"] = round(vram_mb, 2) + + # Reset peak memory counter for next generation + torch.cuda.reset_peak_memory_stats(0) + except Exception as e: + print(f"[MetaHub] Warning: Could not get VRAM usage: {e}") + + # Check MPS (Mac Metal) + elif hasattr(torch.backends, 'mps') and torch.backends.mps.is_available(): + metrics["gpu_available"] = True + metrics["gpu_device"] = "Apple Metal Performance Shaders (MPS)" + # MPS doesn't support memory tracking yet + metrics["vram_peak_mb"] = None + + else: + # CPU-only mode + metrics["gpu_available"] = False + metrics["gpu_device"] = "CPU (CUDA not available)" + + except ImportError: + # torch not installed (should never happen in ComfyUI) + print("[MetaHub] Warning: PyTorch not found, GPU metrics unavailable") + except Exception as e: + print(f"[MetaHub] Warning: GPU metrics collection failed: {e}") + + return metrics + + +def collect_version_info() -> Dict[str, Optional[str]]: + """ + Auto-detects software versions (Python, PyTorch, ComfyUI). + + Silent failures - returns None for unavailable versions. + + Returns: + dict with keys: + - python_version: Python version string (e.g., "3.10.12") + - torch_version: PyTorch version string (e.g., "2.0.1+cu118") + - comfyui_version: ComfyUI version string (e.g., "0.1.0") + """ + versions = { + "comfyui_version": None, + "torch_version": None, + "python_version": None, + } + + # Python version (always available) + try: + import sys + versions["python_version"] = f"{sys.version_info.major}.{sys.version_info.minor}.{sys.version_info.micro}" + except Exception: + pass + + # PyTorch version + try: + import torch + versions["torch_version"] = torch.__version__ + except Exception: + pass + + # ComfyUI version (multiple detection methods) + try: + # Method 1: Try importing comfy module + try: + import comfy + if hasattr(comfy, '__version__'): + versions["comfyui_version"] = comfy.__version__ + except (ImportError, AttributeError): + pass + + # Method 2: Check for version file + if not versions["comfyui_version"]: + version_file = Path(__file__).parent.parent.parent / "comfy_version.txt" + if version_file.exists(): + versions["comfyui_version"] = version_file.read_text().strip() + except Exception: + pass + + return versions + + # ============================================================================ # METADATA FORMATTING # ============================================================================ @@ -280,8 +402,22 @@ def build_imh_metadata(params: dict, workflow_json: dict) -> dict: "project_name": params.get('project_name', ''), }, - # Analytics (custom extension) + # Analytics (custom extension) - Performance/Benchmark metrics "analytics": { + # Tier 1: CRITICAL metrics + "vram_peak_mb": params.get('vram_peak_mb'), + "gpu_device": params.get('gpu_device'), + "generation_time_ms": params.get('generation_time_ms'), + + # Tier 2: VERY USEFUL metrics + "steps_per_second": params.get('steps_per_second'), + "comfyui_version": params.get('comfyui_version'), + + # Tier 3: NICE-TO-HAVE metrics + "torch_version": params.get('torch_version'), + "python_version": params.get('python_version'), + + # Legacy field (kept for backward compatibility) "generation_time": params.get('generation_time', 0.0), }, diff --git a/metahub_save_node.py b/metahub_save_node.py index 4df0ad3..74d95c7 100644 --- a/metahub_save_node.py +++ b/metahub_save_node.py @@ -125,6 +125,22 @@ class MetaHubSaveNode: "default": "", "tooltip": "Upscale model name" }), + # Performance metrics overrides (advanced users) + "vram_peak_mb": ("FLOAT", { + "default": None, + "forceInput": True, + "tooltip": "Override VRAM peak (MB)" + }), + "gpu_device_override": ("STRING", { + "default": None, + "forceInput": True, + "tooltip": "Override GPU device name" + }), + "generation_time_override": ("FLOAT", { + "default": None, + "forceInput": True, + "tooltip": "Override generation time (seconds)" + }), }, "hidden": { "prompt": "PROMPT", @@ -162,12 +178,19 @@ class MetaHubSaveNode: generation_time=0.0, filename_prefix="ComfyUI", upscale_model="", + vram_peak_mb=None, + gpu_device_override=None, + generation_time_override=None, prompt=None, extra_pnginfo=None, unique_id=None, ): global _HINT_SHOWN + # Start timing for auto-measurement + import time + start_time = time.time() + try: workflow_json = utils.get_workflow_json(extra_pnginfo) prompt_data = prompt if isinstance(prompt, dict) else workflow_json.get("prompt", {}) @@ -230,6 +253,32 @@ class MetaHubSaveNode: height, width = images[0].shape[0], images[0].shape[1] + # Calculate elapsed time and collect performance metrics + elapsed_seconds = time.time() - start_time + + # Determine final generation time (use override if provided) + if generation_time_override is not None: + final_time = generation_time_override + elif elapsed_seconds > 0: + final_time = elapsed_seconds + else: + final_time = generation_time or 0.0 + + # Collect GPU metrics (auto-detect) + gpu_metrics = utils.collect_gpu_metrics() + + # Collect version info + version_info = utils.collect_version_info() + + # Calculate derived metrics + generation_time_ms = None + if final_time > 0: + generation_time_ms = int(final_time * 1000) + + steps_per_second = None + if generation_time_ms and generation_time_ms > 0 and steps_value > 0: + steps_per_second = round((steps_value / (generation_time_ms / 1000)), 2) + params = { "positive": positive_value, "negative": negative_value, @@ -251,6 +300,14 @@ class MetaHubSaveNode: "project_name": project_name, "lora_list": lora_list, "lora_hashes": lora_hashes, + # Performance metrics (Tier 1, 2, 3) + "vram_peak_mb": vram_peak_mb if vram_peak_mb is not None else gpu_metrics.get("vram_peak_mb"), + "gpu_device": gpu_device_override if gpu_device_override else gpu_metrics.get("gpu_device"), + "generation_time_ms": generation_time_ms, + "steps_per_second": steps_per_second, + "comfyui_version": version_info.get("comfyui_version"), + "torch_version": version_info.get("torch_version"), + "python_version": version_info.get("python_version"), } a1111_metadata = utils.build_a1111_metadata(params)