test commit

xpu_server.py

@@ -0,0 +1,155 @@
+"""
+xpu_server.py — WebSocket / HTTP endpoint for GPU metrics.
+
+Registers a /xpusys/stats route on ComfyUI's PromptServer that returns
+the latest GPUSnapshot as JSON.  A background task also broadcasts data
+to all connected WebSocket clients at the configured interval.
+
+This module is provider-agnostic: it works with any BaseGPUProvider
+implementation (Intel, NVIDIA, AMD, ...).
+"""
+
+import asyncio
+import json
+import logging
+import os
+from aiohttp import web
+
+try:
+    import folder_paths as _fp
+except ImportError:
+    _fp = None
+
+logger = logging.getLogger("XPUSYSMonitor")
+
+# Will be set by __init__.py after creating the provider
+_provider = None
+
+
+def set_provider(provider):
+    global _provider
+    _provider = provider
+
+
+def _snapshot_to_dict(snap) -> dict:
+    return {
+        # VRAM breakdown
+        "vram_total_gb":       round(snap.vram_total_gb,       2),
+        "vram_free_gb":        round(snap.vram_free_gb,        2),
+        "vram_driver_used_gb": round(snap.vram_driver_used_gb, 2),
+        "vram_allocated_gb":   round(snap.vram_allocated_gb,   2),
+        "vram_reserved_gb":    round(snap.vram_reserved_gb,    2),
+        # GPU
+        "gpu_load_pct":        round(snap.gpu_load_pct, 1),
+        "gpu_freq_mhz":        round(snap.gpu_freq_mhz, 0),
+        "gpu_temp_c":          round(snap.gpu_temp_c,   1),
+        # Power
+        "power_w":             round(snap.power_w, 1),
+        "power_available":     snap.power_available,
+        "tgp_w":               round(snap.tgp_w, 1),
+        "device_name":         snap.device_name,
+        # CPU
+        "cpu_pct":             round(snap.cpu_pct,      1),
+        "cpu_freq_ghz":        round(snap.cpu_freq_ghz, 2),
+        "cpu_model":           snap.cpu_model,
+        "cpu_threads":         snap.cpu_threads,
+        # RAM
+        "ram_pct":             round(snap.ram_pct,          1),
+        "ram_total_gb":        round(snap.ram_total_gb,     2),
+        "ram_used_gb":         round(snap.ram_used_gb,      2),
+        "ram_free_gb":         round(snap.ram_free_gb,      2),
+        "commit_used_gb":      round(snap.commit_used_gb,   2),
+        "commit_limit_gb":     round(snap.commit_limit_gb,  2),
+        # Meta
+        "is_admin":            snap.is_admin,
+        "gpu_vendor":          snap.gpu_vendor,
+        "error":               snap.error,
+    }
+
+
+def register_routes(server):
+    """Call this with the PromptServer instance to attach our HTTP route."""
+
+    @server.routes.get("/xpusys/stats")
+    async def get_stats(request):
+        if _provider is None:
+            return web.json_response({"error": "provider not ready"}, status=503)
+        snap = _provider.get_snapshot()
+        return web.json_response(_snapshot_to_dict(snap))
+
+    @server.routes.post("/xpusys/model_sizes")
+    async def get_model_sizes(request):
+        """
+        Accept a list of {name, type} model descriptors from the frontend,
+        resolve each to a physical file, and return {name, size} in GB.
+        Route name is intentionally distinct from /vram_predict/* used by
+        the standalone ComfyUI-Vram-Predictor plugin.
+        """
+        if _fp is None:
+            return web.json_response({"models": []})
+        try:
+            data  = await request.json()
+            items = data.get("models", [])
+            _ALLOWED = {".safetensors", ".gguf", ".ckpt", ".pt", ".pth", ".bin", ".onnx", ".pkl"}
+            _SEARCH  = ["checkpoints", "vae", "loras", "controlnet", "clip",
+                        "upscale_models", "unet", "diffusion_models",
+                        "ultralytics", "annotator", "bbox", "onnx",
+                        "mmaudio", "audio", "rife", "vfi"]
+            results = []
+            for m in items:
+                name = m.get("name", "")
+                model_path = m.get("path", "")
+                if not name:
+                    continue
+                if os.path.splitext(name)[1].lower() not in _ALLOWED:
+                    continue
+
+                # 在所有搜索目录中按优先级查找
+                path = None
+                for folder in _SEARCH:
+                    # 优先使用完整路径查找（支持子文件夹）
+                    if model_path:
+                        p = _fp.get_full_path(folder, model_path)
+                        if p and os.path.isfile(p):
+                            path = p
+                            break
+
+                # 兜底：用文件名在所有目录下递归搜索
+                if not path:
+                    for folder in _SEARCH:
+                        base_folder = _fp.get_full_path(folder, "")
+                        if base_folder and os.path.isdir(base_folder):
+                            for root, _, files in os.walk(base_folder):
+                                if name in files:
+                                    path = os.path.join(root, name)
+                                    break
+                            if path:
+                                break
+
+                if path and os.path.isfile(path):
+                    size_gb = os.path.getsize(path) / (1024 ** 3)
+                    if size_gb > 0.001:
+                        results.append({"name": name, "size": round(size_gb, 2)})
+            return web.json_response({"models": results})
+        except Exception as exc:
+            logger.debug(f"XPUSYSMonitor: model_sizes error — {exc}")
+            return web.json_response({"models": []})
+
+    logger.info("XPUSYSMonitor: HTTP routes /xpusys/stats and /xpusys/model_sizes registered.")
+
+
+async def broadcast_loop(server, interval_s: float = 1.0):
+    """
+    Continuously broadcast XPU stats to all WebSocket clients via
+    ComfyUI's built-in send_json helper.
+    """
+    while True:
+        await asyncio.sleep(interval_s)
+        if _provider is None:
+            continue
+        try:
+            snap = _provider.get_snapshot()
+            data = _snapshot_to_dict(snap)
+            await server.send_json("xpusys_stats", data)
+        except Exception as exc:
+            logger.debug(f"XPUSYSMonitor: broadcast error — {exc}")