Add files via upload

__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}")

web/index.js

@@ -0,0 +1 @@
+import "./monitor.js";

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 = '<span style="color: #ff5555;">AMD GPU Monitor</span>';
+    
+    // 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);
+    },
+});