'commit'

BaiHu/Tools/TestComfyUIApi.py

@@ -6,30 +6,6 @@ config = ConfigUtil.getConfig()
 server_address = config.get('comfyui', 'server_address')
 
 
-# 获取webui使用情况
-def get_webui_used():
-    # webui 服务器地址
-    webui_server = config['webui']['webui_server']
-    vram_url = webui_server + config['webui']['get_vram_url']
-    res = submit_post(vram_url, None)
-    if res.status_code == 200:
-        print(res.text)
-    else:
-        print("调用集成的获取显卡显存代码失败！")
-
-
-# 清理webui的缓存
-def clear_webui_cache():
-    # webui 服务器地址
-    webui_server = config['webui']['webui_server']
-    empty_cache_url = webui_server + config['webui']['empty_cache_url']
-    res = submit_post(empty_cache_url, None)
-    if res.status_code == 200:
-        print(res.text)
-    else:
-        print("调用集成的清理显卡显存代码失败！")
-
-
 # 清理GPU缓存
 info = get_comfyui_used(server_address)
 print(info)
@@ -38,8 +14,8 @@ info = get_comfyui_used(server_address)
 print(info)
 
 # 获取使显存数量
-get_webui_used()
+get_webui_used(config)
 # 清理缓存
-clear_webui_cache()
+clear_webui_cache(config)
 # 获取使显存数量
-get_webui_used()
+get_webui_used(config)

BaiHu/Util/ComfyUIUtil.py

@@ -1,3 +1,4 @@
+import json
 import os.path
 import time
 import urllib.parse
@@ -165,3 +166,36 @@ def upload_file(server_address, file, subfolder="", overwrite=False):
     except Exception as error:
         print(error)
     return path
+
+
+# 获取webui使用情况
+def get_webui_used(config):
+    # webui 服务器地址
+    webui_server = config['webui']['webui_server']
+    vram_url = webui_server + config['webui']['get_vram_url']
+    res = submit_post(vram_url, None)
+    if res.status_code == 200:
+        jo = json.loads(res.text)
+        total = jo[0].replace('total', '').replace(' ', '').replace(':', '').replace('GB', '')
+        # free = jo[1].replace('free', '').replace(' ', '').replace(':', '').replace('GB', '')
+        used = jo[2].replace('used', '').replace(' ', '').replace(':', '').replace('GB', '')
+        used_lv = round(1.0 * (float(used)) / float(total) * 100, 2)
+        print("显存:" + total + "GB,已使用:" + used + "GB,使用率：" + str(used_lv) + '%')
+    else:
+        print("调用集成的获取显卡显存代码失败！")
+
+
+# 清理webui的缓存
+def clear_webui_cache(config):
+    # webui 服务器地址
+    webui_server = config['webui']['webui_server']
+    empty_cache_url = webui_server + config['webui']['empty_cache_url']
+    res = submit_post(empty_cache_url, None)
+    if res.status_code == 200:
+        jo = json.loads(res.text)
+        if jo['success']:
+            print("显存清理成功！")
+        else:
+            print("显存清理失败！")
+    else:
+        print("调用集成的清理显卡显存代码失败！")