from Util import ConfigUtil
from Util.SDUtil import *

# 打开配置文件
config = ConfigUtil.getConfig()
comfyui_address = config.get('comfyui', 'server_address')
webui_address = config.get('webui', 'webui_address')

# 清理GPU缓存
info = get_comfyui_used(comfyui_address)
print(info)
clear_comfyui_cache(comfyui_address)
info = get_comfyui_used(comfyui_address)
print(info)

# 获取使显存数量
total, used = get_webui_used(webui_address)

# 显卡显存占用超过20GB
if float(used) > 20:
    print("发现显存占用高达" + used + "GB，准备重启SD...")
    # 重新启动一下SD
    restart_server(webui_address)
    # 休息2分钟以保证SD重启完成
    time.sleep(2 * 60)
else:
    print("显存占用正常，无法重新启动SD！")