diff --git a/BaiHu/Tools/TestComfyUIApi.py b/BaiHu/Tools/TestComfyUIApi.py index 6fb11b94..994d1216 100644 --- a/BaiHu/Tools/TestComfyUIApi.py +++ b/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) diff --git a/BaiHu/Util/ComfyUIUtil.py b/BaiHu/Util/ComfyUIUtil.py index c1759e8a..aac27baa 100644 --- a/BaiHu/Util/ComfyUIUtil.py +++ b/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("调用集成的清理显卡显存代码失败!") diff --git a/BaiHu/Util/__pycache__/ComfyUIUtil.cpython-310.pyc b/BaiHu/Util/__pycache__/ComfyUIUtil.cpython-310.pyc index 0646df14..572579df 100644 Binary files a/BaiHu/Util/__pycache__/ComfyUIUtil.cpython-310.pyc and b/BaiHu/Util/__pycache__/ComfyUIUtil.cpython-310.pyc differ