Antique3e
1c39aa6d18
Updated CPU detection to use subprocess and added placeholder for memory total.
137 lines
3.9 KiB
Python
137 lines
3.9 KiB
Python
"""
|
|
ComfyUI Credit Tracker
|
|
Simple extension to track Modal GPU credits
|
|
"""
|
|
|
|
import os
|
|
import json
|
|
import subprocess
|
|
import psutil # NEW: Added for CPU/Memory detection
|
|
from aiohttp import web
|
|
import server
|
|
|
|
# Export web directory
|
|
WEB_DIRECTORY = "./js"
|
|
|
|
# File paths
|
|
SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
|
|
CONFIG_FILE = os.path.join(SCRIPT_DIR, "config.json")
|
|
BALANCE_FILE = os.path.join(SCRIPT_DIR, "balance.json")
|
|
|
|
|
|
def load_config():
|
|
"""Load configuration file"""
|
|
if os.path.exists(CONFIG_FILE):
|
|
with open(CONFIG_FILE, 'r') as f:
|
|
return json.load(f)
|
|
return None
|
|
|
|
|
|
def load_balance():
|
|
"""Load balance file"""
|
|
if os.path.exists(BALANCE_FILE):
|
|
with open(BALANCE_FILE, 'r') as f:
|
|
return json.load(f)
|
|
return None
|
|
|
|
|
|
def save_balance(balance_data):
|
|
"""Save balance to file"""
|
|
with open(BALANCE_FILE, 'w') as f:
|
|
json.dump(balance_data, f, indent=2)
|
|
|
|
|
|
def get_gpu_info():
|
|
"""Detect GPU using nvidia-smi + CPU/Memory using psutil"""
|
|
result = {"gpu_name": "Unknown", "success": False}
|
|
|
|
# Detect GPU (existing code)
|
|
try:
|
|
gpu_result = subprocess.run(
|
|
['nvidia-smi', '--query-gpu=name', '--format=csv,noheader'],
|
|
capture_output=True,
|
|
text=True,
|
|
timeout=5
|
|
)
|
|
if gpu_result.returncode == 0:
|
|
result["gpu_name"] = gpu_result.stdout.strip()
|
|
result["success"] = True
|
|
except Exception as e:
|
|
print(f"GPU detection failed: {e}")
|
|
|
|
# NEW: Detect CPU and Memory (added to same function)
|
|
try:
|
|
cpu_result = subprocess.run(['nproc'], capture_output=True, text=True, timeout=5)
|
|
if cpu_result.returncode == 0:
|
|
result["cpu_cores"] = int(cpu_result.stdout.strip())
|
|
else:
|
|
result["cpu_cores"] = 0
|
|
result["memory_total_gb"] = 0 # Will use config value instead
|
|
except Exception as e:
|
|
print(f"CPU detection failed: {e}")
|
|
result["cpu_cores"] = 0
|
|
result["memory_total_gb"] = 0
|
|
|
|
return result
|
|
|
|
|
|
# API Routes
|
|
@server.PromptServer.instance.routes.get("/credit_tracker/config")
|
|
async def get_config(request):
|
|
"""Get configuration"""
|
|
config = load_config()
|
|
if config:
|
|
return web.json_response(config)
|
|
return web.json_response({"error": "Config not found"}, status=404)
|
|
|
|
|
|
@server.PromptServer.instance.routes.get("/credit_tracker/balance")
|
|
async def get_balance(request):
|
|
"""Get current balance"""
|
|
balance = load_balance()
|
|
if balance:
|
|
return web.json_response(balance)
|
|
return web.json_response({"error": "Balance not found"}, status=404)
|
|
|
|
|
|
@server.PromptServer.instance.routes.post("/credit_tracker/balance")
|
|
async def update_balance(request):
|
|
"""Save balance"""
|
|
try:
|
|
balance_data = await request.json()
|
|
save_balance(balance_data)
|
|
return web.json_response({"success": True})
|
|
except Exception as e:
|
|
return web.json_response({"error": str(e)}, status=500)
|
|
|
|
|
|
@server.PromptServer.instance.routes.get("/credit_tracker/gpu_info")
|
|
async def get_gpu(request):
|
|
"""Get GPU information"""
|
|
gpu_info = get_gpu_info()
|
|
return web.json_response(gpu_info)
|
|
|
|
|
|
@server.PromptServer.instance.routes.post("/credit_tracker/reset")
|
|
async def reset_balance(request):
|
|
"""Reset balance to starting amount"""
|
|
try:
|
|
config = load_config()
|
|
if not config:
|
|
return web.json_response({"error": "Config not found"}, status=404)
|
|
|
|
from datetime import datetime
|
|
balance_data = {
|
|
"last_updated": datetime.now().isoformat(),
|
|
"remaining_balance": config.get("starting_balance", 80.0)
|
|
}
|
|
save_balance(balance_data)
|
|
return web.json_response({"success": True, "balance": balance_data})
|
|
except Exception as e:
|
|
return web.json_response({"error": str(e)}, status=500)
|
|
|
|
|
|
print("✅ Credit Tracker: Extension loaded")
|
|
|
|
__all__ = ['WEB_DIRECTORY']
|