diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..9ffd20c --- /dev/null +++ b/__init__.py @@ -0,0 +1,254 @@ +import os +import sys +import json +import time +import subprocess +import threading +import asyncio +from server import PromptServer + +# Global variable to store GPU stats +gpu_stats = { + "gpu_utilization": 0, + "vram_used": 0, + "vram_total": 0, + "vram_used_percent": 0, + "gpu_temperature": 0, + "last_update": 0 +} + +# Monitor thread control +monitor_thread = None +thread_control = threading.Event() +monitor_update_interval = 1 # seconds + +def find_rocm_smi(): + """Find the rocm-smi or amd-smi executable""" + rocm_paths = [ + "/opt/rocm/bin/rocm-smi", + "/usr/bin/rocm-smi", + "/usr/local/bin/rocm-smi", + "/opt/amdgpu-pro/bin/amd-smi", + "/usr/bin/amd-smi", + ] + + for path in rocm_paths: + if os.path.isfile(path) and os.access(path, os.X_OK): + return path + + # Try to find it in PATH + try: + result = subprocess.run(["which", "rocm-smi"], capture_output=True, text=True) + if result.returncode == 0: + return result.stdout.strip() + except: + pass + + try: + result = subprocess.run(["which", "amd-smi"], capture_output=True, text=True) + if result.returncode == 0: + return result.stdout.strip() + except: + pass + + return None + +def run_rocm_smi_command(rocm_smi_path, *args): + """Run a rocm-smi command and return the parsed JSON output""" + if not rocm_smi_path: + return {} + + cmd = [rocm_smi_path] + list(args) + try: + result = subprocess.run(cmd, capture_output=True, text=True, timeout=5) + if result.returncode != 0: + return {} + + # Check if output is JSON + if '--json' in args: + try: + return json.loads(result.stdout) + except json.JSONDecodeError: + return {} + else: + return result.stdout + except subprocess.TimeoutExpired: + return {} + except Exception as e: + return {} + +def get_gpu_info(rocm_smi_path): + """Get current GPU information""" + global gpu_stats + + # Get GPU utilization + try: + info = run_rocm_smi_command(rocm_smi_path, '--showuse', '--json') + if isinstance(info, dict) and 'card0' in info: + card_info = info['card0'] # Use first GPU + if 'GPU use (%)' in card_info: + gpu_use = card_info['GPU use (%)'] + if isinstance(gpu_use, str): + gpu_use = gpu_use.replace('%', '') + gpu_stats["gpu_utilization"] = int(float(gpu_use)) + except: + pass + + # Get VRAM information + try: + info = run_rocm_smi_command(rocm_smi_path, '--showmeminfo', 'vram', '--json') + if isinstance(info, dict) and 'card0' in info: + card_info = info['card0'] # Use first GPU + + # Parse the B (bytes) format ROCm 5.x/6.x uses + if 'VRAM Total Memory (B)' in card_info and 'VRAM Total Used Memory (B)' in card_info: + vram_total_bytes = int(card_info['VRAM Total Memory (B)']) + vram_used_bytes = int(card_info['VRAM Total Used Memory (B)']) + + # Convert to MB for display + vram_total = vram_total_bytes / (1024 * 1024) + vram_used = vram_used_bytes / (1024 * 1024) + + gpu_stats["vram_total"] = int(vram_total) + gpu_stats["vram_used"] = int(vram_used) + gpu_stats["vram_used_percent"] = int((vram_used / vram_total) * 100) + except: + pass + + # Get temperature + try: + info = run_rocm_smi_command(rocm_smi_path, '--showtemp', '--json') + if isinstance(info, dict) and 'card0' in info: + card_info = info['card0'] # Use first GPU + + # Try different temperature sensors, starting with edge + if 'Temperature (Sensor edge) (C)' in card_info: + temp_str = card_info['Temperature (Sensor edge) (C)'] + if isinstance(temp_str, str): + temp_str = temp_str.replace('°C', '').strip() + gpu_stats["gpu_temperature"] = int(float(temp_str)) + elif 'Temperature (Sensor junction) (C)' in card_info: + temp_str = card_info['Temperature (Sensor junction) (C)'] + if isinstance(temp_str, str): + temp_str = temp_str.replace('°C', '').strip() + gpu_stats["gpu_temperature"] = int(float(temp_str)) + except: + pass + + gpu_stats["last_update"] = time.time() + return gpu_stats + +def send_monitor_update(): + """Format and send monitor update data""" + data = { + 'device_type': 'rocm', + 'gpus': [{ + 'gpu_utilization': gpu_stats['gpu_utilization'], + 'gpu_temperature': gpu_stats['gpu_temperature'], + 'vram_total': gpu_stats['vram_total'], + 'vram_used': gpu_stats['vram_used'], + 'vram_used_percent': gpu_stats['vram_used_percent'] + }] + } + + # Send the data + try: + PromptServer.instance.send_sync('amd_gpu_monitor', data) + except: + pass + +def monitor_thread_function(): + """Thread function to continuously monitor GPU stats""" + global monitor_update_interval + + rocm_smi_path = find_rocm_smi() + if not rocm_smi_path: + print("ERROR: Could not find rocm-smi or amd-smi executable") + return + + print(f"Using AMD SMI tool: {rocm_smi_path}") + + while not thread_control.is_set(): + try: + get_gpu_info(rocm_smi_path) + # Send update to UI + send_monitor_update() + except: + pass + + # Sleep for the update interval + time.sleep(monitor_update_interval) + +def start_monitor_thread(): + """Start the GPU monitoring thread""" + global monitor_thread, thread_control + + if monitor_thread is not None and monitor_thread.is_alive(): + # Already running + return + + # Clear the control event and start new thread + thread_control.clear() + monitor_thread = threading.Thread(target=monitor_thread_function) + monitor_thread.daemon = True + monitor_thread.start() + print("AMD GPU Monitor thread started") + +def stop_monitor_thread(): + """Stop the GPU monitoring thread""" + global monitor_thread, thread_control + + if monitor_thread is None or not monitor_thread.is_alive(): + # Not running + return + + # Set control event to stop thread + thread_control.set() + monitor_thread.join(timeout=5) + print("AMD GPU Monitor thread stopped") + +# Start the monitor thread when this module is loaded +start_monitor_thread() + +# Define our nodes (not really used, but required for ComfyUI to load the extension) +class AMDGPUMonitor: + """ + A placeholder node for ComfyUI. The actual monitoring is done via a separate thread. + """ + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "update_interval": ("FLOAT", {"default": 1.0, "min": 0.1, "max": 10.0, "step": 0.1}), + }, + } + + RETURN_TYPES = ("STRING",) + FUNCTION = "monitor_gpu" + CATEGORY = "AMD GPU" + + def monitor_gpu(self, update_interval): + """Update interval can be changed via input""" + global monitor_update_interval + monitor_update_interval = update_interval + + # Return current stats as a string for debugging + stats = f"GPU: {gpu_stats['gpu_utilization']}% | VRAM: {gpu_stats['vram_used']}MB/{gpu_stats['vram_total']}MB ({gpu_stats['vram_used_percent']}%) | Temp: {gpu_stats['gpu_temperature']}°C" + return (stats,) + +# Register our node when this script is imported +NODE_CLASS_MAPPINGS = { + "AMDGPUMonitor": AMDGPUMonitor, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "AMDGPUMonitor": "AMD GPU Monitor", +} + +# This cleanup will be called when ComfyUI is shutting down +def cleanup(): + stop_monitor_thread() + +# Web directory setup for ComfyUI to find our JS files +WEB_DIRECTORY = os.path.join(os.path.dirname(os.path.realpath(__file__)), "web") +print(f"AMD GPU Monitor: Web directory set to {WEB_DIRECTORY}") diff --git a/screenshot.png b/screenshot.png new file mode 100644 index 0000000..65e9992 Binary files /dev/null and b/screenshot.png differ diff --git a/web/index.js b/web/index.js new file mode 100644 index 0000000..cb13995 --- /dev/null +++ b/web/index.js @@ -0,0 +1 @@ +import "./monitor.js"; diff --git a/web/monitor.js b/web/monitor.js new file mode 100644 index 0000000..afd0c5c --- /dev/null +++ b/web/monitor.js @@ -0,0 +1,403 @@ +import { app } from "../../scripts/app.js"; +import { api } from "../../scripts/api.js"; + +// Create the monitor UI element +const createMonitorElement = () => { + // Create main container + const container = document.createElement("div"); + container.className = "amd-gpu-monitor"; + container.style.position = "absolute"; + container.style.top = "40px"; // Moved down to avoid toolbar + container.style.right = "5px"; + container.style.zIndex = "1000"; + container.style.backgroundColor = "#1a1a1a"; + container.style.color = "#fff"; + container.style.padding = "10px"; + container.style.borderRadius = "5px"; + container.style.fontFamily = "monospace"; + container.style.fontSize = "12px"; + container.style.boxShadow = "0 0 10px rgba(0, 0, 0, 0.5)"; + container.style.width = "220px"; + container.style.userSelect = "none"; + + // Add title + const title = document.createElement("div"); + title.style.fontWeight = "bold"; + title.style.marginBottom = "5px"; + title.style.display = "flex"; + title.style.alignItems = "center"; + title.style.justifyContent = "space-between"; + title.innerHTML = 'AMD GPU Monitor'; + + // Add collapse button + const collapseButton = document.createElement("button"); + collapseButton.innerHTML = "−"; // Unicode minus sign + collapseButton.style.background = "none"; + collapseButton.style.border = "none"; + collapseButton.style.color = "#888"; + collapseButton.style.cursor = "pointer"; + collapseButton.style.fontSize = "14px"; + collapseButton.style.padding = "0 5px"; + collapseButton.title = "Collapse/Expand"; + + // Add close button + const closeButton = document.createElement("button"); + closeButton.innerHTML = "×"; // Unicode times sign + closeButton.style.background = "none"; + closeButton.style.border = "none"; + closeButton.style.color = "#888"; + closeButton.style.cursor = "pointer"; + closeButton.style.fontSize = "14px"; + closeButton.style.padding = "0 5px"; + closeButton.title = "Close"; + + const buttonContainer = document.createElement("div"); + buttonContainer.appendChild(collapseButton); + buttonContainer.appendChild(closeButton); + + title.appendChild(buttonContainer); + container.appendChild(title); + + // Content container that can be collapsed + const content = document.createElement("div"); + content.className = "amd-gpu-monitor-content"; + container.appendChild(content); + + // GPU Utilization section + const gpuSection = document.createElement("div"); + gpuSection.style.marginBottom = "8px"; + + const gpuLabel = document.createElement("div"); + gpuLabel.textContent = "GPU Utilization:"; + gpuLabel.style.marginBottom = "2px"; + + const gpuBarContainer = document.createElement("div"); + gpuBarContainer.style.height = "15px"; + gpuBarContainer.style.backgroundColor = "#333"; + gpuBarContainer.style.borderRadius = "3px"; + gpuBarContainer.style.position = "relative"; + + const gpuBar = document.createElement("div"); + gpuBar.className = "amd-gpu-utilization-bar"; + gpuBar.style.height = "100%"; + gpuBar.style.width = "0%"; + gpuBar.style.backgroundColor = "#47a0ff"; + gpuBar.style.borderRadius = "3px"; + gpuBar.style.transition = "width 0.5s ease-out, background-color 0.3s"; + + const gpuText = document.createElement("div"); + gpuText.className = "amd-gpu-utilization-text"; + gpuText.textContent = "0%"; + gpuText.style.position = "absolute"; + gpuText.style.top = "0"; + gpuText.style.left = "5px"; + gpuText.style.lineHeight = "15px"; + gpuText.style.textShadow = "1px 1px 1px #000"; + + gpuBarContainer.appendChild(gpuBar); + gpuBarContainer.appendChild(gpuText); + gpuSection.appendChild(gpuLabel); + gpuSection.appendChild(gpuBarContainer); + content.appendChild(gpuSection); + + // VRAM Usage section + const vramSection = document.createElement("div"); + vramSection.style.marginBottom = "8px"; + + const vramLabel = document.createElement("div"); + vramLabel.textContent = "VRAM Usage:"; + vramLabel.style.marginBottom = "2px"; + + const vramBarContainer = document.createElement("div"); + vramBarContainer.style.height = "15px"; + vramBarContainer.style.backgroundColor = "#333"; + vramBarContainer.style.borderRadius = "3px"; + vramBarContainer.style.position = "relative"; + + const vramBar = document.createElement("div"); + vramBar.className = "amd-vram-bar"; + vramBar.style.height = "100%"; + vramBar.style.width = "0%"; + vramBar.style.backgroundColor = "#47a0ff"; + vramBar.style.borderRadius = "3px"; + vramBar.style.transition = "width 0.5s ease-out, background-color 0.3s"; + + const vramText = document.createElement("div"); + vramText.className = "amd-vram-text"; + vramText.textContent = "0MB / 0MB (0%)"; + vramText.style.position = "absolute"; + vramText.style.top = "0"; + vramText.style.left = "5px"; + vramText.style.lineHeight = "15px"; + vramText.style.textShadow = "1px 1px 1px #000"; + + vramBarContainer.appendChild(vramBar); + vramBarContainer.appendChild(vramText); + vramSection.appendChild(vramLabel); + vramSection.appendChild(vramBarContainer); + content.appendChild(vramSection); + + // Temperature section + const tempSection = document.createElement("div"); + + const tempLabel = document.createElement("div"); + tempLabel.textContent = "GPU Temperature:"; + tempLabel.style.marginBottom = "2px"; + + const tempBarContainer = document.createElement("div"); + tempBarContainer.style.height = "15px"; + tempBarContainer.style.backgroundColor = "#333"; + tempBarContainer.style.borderRadius = "3px"; + tempBarContainer.style.position = "relative"; + + const tempBar = document.createElement("div"); + tempBar.className = "amd-temp-bar"; + tempBar.style.height = "100%"; + tempBar.style.width = "0%"; + tempBar.style.backgroundColor = "#47a0ff"; + tempBar.style.borderRadius = "3px"; + tempBar.style.transition = "width 0.5s ease-out, background-color 0.3s"; + + const tempText = document.createElement("div"); + tempText.className = "amd-temp-text"; + tempText.textContent = "0°C"; + tempText.style.position = "absolute"; + tempText.style.top = "0"; + tempText.style.left = "5px"; + tempText.style.lineHeight = "15px"; + tempText.style.textShadow = "1px 1px 1px #000"; + + tempBarContainer.appendChild(tempBar); + tempBarContainer.appendChild(tempText); + tempSection.appendChild(tempLabel); + tempSection.appendChild(tempBarContainer); + content.appendChild(tempSection); + + // Add event listener for collapsing + let isCollapsed = false; + collapseButton.addEventListener("click", () => { + if (isCollapsed) { + content.style.display = "block"; + collapseButton.innerHTML = "−"; + isCollapsed = false; + } else { + content.style.display = "none"; + collapseButton.innerHTML = "+"; + isCollapsed = true; + } + }); + + // Add event listener for closing + closeButton.addEventListener("click", () => { + container.style.display = "none"; + // Store the closed state in localStorage + localStorage.setItem("amd-gpu-monitor-closed", "true"); + }); + + // Make the monitor draggable + let isDragging = false; + let dragOffsetX, dragOffsetY; + + title.addEventListener("mousedown", (e) => { + // Only handle main button (left button) + if (e.button !== 0) return; + + isDragging = true; + dragOffsetX = e.clientX - container.offsetLeft; + dragOffsetY = e.clientY - container.offsetTop; + + // Prevent text selection during drag + e.preventDefault(); + }); + + document.addEventListener("mousemove", (e) => { + if (!isDragging) return; + + const x = e.clientX - dragOffsetX; + const y = e.clientY - dragOffsetY; + + // Keep monitor within window bounds + const maxX = window.innerWidth - container.offsetWidth; + const maxY = window.innerHeight - container.offsetHeight; + + container.style.left = Math.max(0, Math.min(x, maxX)) + "px"; + container.style.top = Math.max(0, Math.min(y, maxY)) + "px"; + + // We're now positioning with left instead of right + container.style.right = "auto"; + }); + + document.addEventListener("mouseup", () => { + isDragging = false; + + // Save position to localStorage + if (container.style.left && container.style.top) { + localStorage.setItem("amd-gpu-monitor-position", JSON.stringify({ + left: container.style.left, + top: container.style.top + })); + } + }); + + // Load saved position if available + const savedPosition = localStorage.getItem("amd-gpu-monitor-position"); + if (savedPosition) { + try { + const { left, top } = JSON.parse(savedPosition); + container.style.left = left; + container.style.top = top; + container.style.right = "auto"; + } catch (e) { + // Silently fail and use default position + } + } + + // Check if monitor was closed previously + if (localStorage.getItem("amd-gpu-monitor-closed") === "true") { + container.style.display = "none"; + } + + // Add a button to show the monitor again + const showButton = document.createElement("button"); + showButton.textContent = "Show AMD GPU Monitor"; + showButton.style.position = "fixed"; + showButton.style.top = "5px"; + showButton.style.right = "5px"; + showButton.style.zIndex = "999"; + showButton.style.padding = "5px"; + showButton.style.borderRadius = "3px"; + showButton.style.backgroundColor = "#333"; + showButton.style.color = "#fff"; + showButton.style.border = "none"; + showButton.style.fontSize = "12px"; + showButton.style.cursor = "pointer"; + showButton.style.display = "none"; + + showButton.addEventListener("click", () => { + container.style.display = "block"; + showButton.style.display = "none"; + localStorage.removeItem("amd-gpu-monitor-closed"); + }); + + document.body.appendChild(showButton); + + // Toggle showButton visibility based on monitor visibility + const updateShowButtonVisibility = () => { + if (container.style.display === "none") { + showButton.style.display = "block"; + } else { + showButton.style.display = "none"; + } + }; + + // Create a MutationObserver to watch for changes to container's display style + const observer = new MutationObserver((mutations) => { + mutations.forEach((mutation) => { + if (mutation.attributeName === "style") { + updateShowButtonVisibility(); + } + }); + }); + + observer.observe(container, { attributes: true }); + + // Initial visibility check + updateShowButtonVisibility(); + + return { container, gpuBar, gpuText, vramBar, vramText, tempBar, tempText }; +}; + +// Update the monitor UI with new data +const updateMonitorUI = (monitor, data) => { + // Check if we have GPU data + if (!data || !data.gpus || data.gpus.length === 0) return; + + const gpu = data.gpus[0]; // Use the first GPU + + // Update GPU utilization + if (monitor.gpuBar && monitor.gpuText) { + const utilization = gpu.gpu_utilization || 0; + monitor.gpuBar.style.width = `${utilization}%`; + monitor.gpuText.textContent = `${utilization}%`; + + // Change color based on utilization + if (utilization > 80) { + monitor.gpuBar.style.backgroundColor = '#ff4d4d'; // Red for high + } else if (utilization > 50) { + monitor.gpuBar.style.backgroundColor = '#ffad33'; // Orange for medium + } else { + monitor.gpuBar.style.backgroundColor = '#47a0ff'; // Blue for low + } + } + + // Update VRAM usage + if (monitor.vramBar && monitor.vramText) { + const vramPercent = gpu.vram_used_percent || 0; + const vramUsed = gpu.vram_used || 0; + const vramTotal = gpu.vram_total || 1; + + monitor.vramBar.style.width = `${vramPercent}%`; + + // Format the text to show MB or GB + let vramUsedText = vramUsed; + let vramTotalText = vramTotal; + let unit = 'MB'; + + if (vramTotal >= 1024) { + vramUsedText = (vramUsed / 1024).toFixed(1); + vramTotalText = (vramTotal / 1024).toFixed(1); + unit = 'GB'; + } + + monitor.vramText.textContent = `${vramUsedText}${unit} / ${vramTotalText}${unit} (${vramPercent}%)`; + + // Change color based on VRAM usage + if (vramPercent > 85) { + monitor.vramBar.style.backgroundColor = '#ff4d4d'; // Red for high + } else if (vramPercent > 70) { + monitor.vramBar.style.backgroundColor = '#ffad33'; // Orange for medium + } else { + monitor.vramBar.style.backgroundColor = '#47a0ff'; // Blue for low + } + } + + // Update temperature + if (monitor.tempBar && monitor.tempText) { + const temp = gpu.gpu_temperature || 0; + + // Assume max reasonable temp is 100°C for the progress bar + const tempPercent = Math.min(temp, 100); + monitor.tempBar.style.width = `${tempPercent}%`; + monitor.tempText.textContent = `${temp}°C`; + + // Change color based on temperature + if (temp > 80) { + monitor.tempBar.style.backgroundColor = '#ff4d4d'; // Red for high + } else if (temp > 60) { + monitor.tempBar.style.backgroundColor = '#ffad33'; // Orange for medium + } else { + monitor.tempBar.style.backgroundColor = '#47a0ff'; // Blue for low + } + } +}; + +// Main app function +const main = () => { + // Create the monitor UI + const monitor = createMonitorElement(); + document.body.appendChild(monitor.container); + + // Set up WebSocket listener for GPU updates + api.addEventListener("amd_gpu_monitor", (event) => { + updateMonitorUI(monitor, event.detail); + }); +}; + +// Wait for DOM to be loaded +app.registerExtension({ + name: "amd.gpu.monitor", + async setup() { + // Wait a bit for the UI to be fully loaded + setTimeout(main, 1000); + }, +});