From f7603952d589f949bae33bb65f7f65c3d9832aa1 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E9=BB=84=E6=B5=B7?= <10402852@qq.com>
Date: Mon, 1 Jul 2024 08:04:01 +0800
Subject: [PATCH] 'commit'

---
 BaiHu/Tools/TestComfyUIApi.py | 33 +++++++++++++++++++--------------
 1 file changed, 19 insertions(+), 14 deletions(-)

diff --git a/BaiHu/Tools/TestComfyUIApi.py b/BaiHu/Tools/TestComfyUIApi.py
index 2f69149e..53071f3f 100644
--- a/BaiHu/Tools/TestComfyUIApi.py
+++ b/BaiHu/Tools/TestComfyUIApi.py
@@ -20,6 +20,21 @@ def queue_prompt(server_address, client_id, prompt):
         print(err)
     return json.loads(urllib.request.urlopen(req).read())
 
+def getUse(server_address):
+    req = urllib.request.Request("http://{}/system_stats".format(server_address))
+    res = json.loads(urllib.request.urlopen(req).read())
+
+    vram_total = res['devices'][0]['vram_total']
+    vram_total_str = str(int(vram_total / 1024 / 1024 / 1024 + 0.5))
+    vram_free = res['devices'][0]['vram_free']
+    used_vram = vram_total - vram_free
+    used_vram_str = str(int((vram_total - vram_free) / 1024 / 1024 / 1024 + 0.5))
+    used_lv = round(1.0 * (used_vram) / vram_total * 100, 2)
+    print("显存共：" + vram_total_str + "GB,已使用：" + used_vram_str + "GB,使用率：" + str(used_lv) + "% ")
+    # vram_total # 显存容量：25756696576 即24GB
+    # vram_free  # 显存剩余：25756696576
+    # torch_vram_total 16005464064
+    # torch_vram_free  331041996
 
 # GET /system_stats
 # 系统统计信息接口
@@ -28,20 +43,7 @@ def queue_prompt(server_address, client_id, prompt):
 config = ConfigUtil.getConfig()
 
 server_address = config.get('comfyui', 'server_address')
-req = urllib.request.Request("http://{}/system_stats".format(server_address))
-res = json.loads(urllib.request.urlopen(req).read())
-
-vram_total = res['devices'][0]['vram_total']
-vram_total_str = str(int(vram_total / 1024 / 1024 / 1024 + 0.5))
-vram_free = res['devices'][0]['vram_free']
-used_vram = vram_total - vram_free
-used_vram_str = str(int((vram_total - vram_free) / 1024 / 1024 / 1024 + 0.5))
-used_lv = round(1.0 * (used_vram) / vram_total * 100, 2)
-print("显存共：" + vram_total_str + "GB,已使用：" + used_vram_str + "GB,使用率：" + str(used_lv) + "% ")
-# vram_total # 显存容量：25756696576 即24GB
-# vram_free  # 显存剩余：25756696576
-# torch_vram_total 16005464064
-# torch_vram_free  331041996
+getUse(server_address)
 
 # 打开文件并读取内容
 file_path = r'../JSON/27.json'
@@ -68,3 +70,6 @@ print(urllib.request.urlopen(req).read())
 # https://www.runcomfy.com/comfyui-nodes/ComfyUI-Easy-Use/easy-cleanGpuUsed
 
 # https://www.yunrobot.cn/showdoc/web/#/641840309/231516860
+
+# https://www.reddit.com/r/comfyui/comments/1cdhz5v/best_way_to_clear_vram_after_each_generation/
+