From d3df10ba7a293429c1446310ffcc90aa025e1c80 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E9=BB=84=E6=B5=B7?= <10402852@qq.com>
Date: Tue, 2 Jul 2024 14:51:10 +0800
Subject: [PATCH] 'commit'

---
 BaiHu/Config.ini              |  7 ++++---
 BaiHu/Doc/GPU监控.txt       |  5 ++---
 BaiHu/Tools/TestComfyUIApi.py | 19 ++++++++++++++++---
 BaiHu/doTask.py               |  7 +++++--
 4 files changed, 27 insertions(+), 11 deletions(-)

diff --git a/BaiHu/Config.ini b/BaiHu/Config.ini
index 02b084a8..106bdcb1 100644
--- a/BaiHu/Config.ini
+++ b/BaiHu/Config.ini
@@ -11,9 +11,10 @@ server_address = 192.168.1.21:8188
 
 [webui]
 ; 生图服务地址
-txt2img_url = http://192.168.1.21:7860/sdapi/v1/txt2img
-img2img_url = http://192.168.1.21:7860/sdapi/v1/img2img
-
+webui_server = http://192.168.1.21:7860
+txt2img_url = /sdapi/v1/txt2img
+img2img_url = /sdapi/v1/img2img
+get_vram_url = /sdapi/v1/get_vram
 ; WEB服务器地址
 [webServer]
 ; web_url = https://www.hzkjai.com
diff --git a/BaiHu/Doc/GPU监控.txt b/BaiHu/Doc/GPU监控.txt
index 645a9c85..2d147ef6 100644
--- a/BaiHu/Doc/GPU监控.txt
+++ b/BaiHu/Doc/GPU监控.txt
@@ -36,12 +36,11 @@ def empty_cache(self):
         gc.collect()
     return {"success":True,"message":"GPU is cleared!"}
 
-
-def getVRam(self):
+def get_vram(self):
     nvmlInit()
     h = nvmlDeviceGetHandleByIndex(0)
     info = nvmlDeviceGetMemoryInfo(h)
-    res=[]
+    res = []
     res.append(f'total: {round(info.total / 1024 / 1024 / 1024, 1)} GB')
     res.append(f'free : {round(info.free / 1024 / 1024 / 1024, 1)} GB')
     res.append(f'used : {round(info.used / 1024 / 1024 / 1024, 1)} GB')
diff --git a/BaiHu/Tools/TestComfyUIApi.py b/BaiHu/Tools/TestComfyUIApi.py
index b9fbec60..0e3b11c0 100644
--- a/BaiHu/Tools/TestComfyUIApi.py
+++ b/BaiHu/Tools/TestComfyUIApi.py
@@ -6,6 +6,19 @@ config = ConfigUtil.getConfig()
 server_address = config.get('comfyui', 'server_address')
 
 # 清理GPU缓存
-info = getUse(server_address)
-print(info)
-clear_cache(server_address)
+# info = getUse(server_address)
+# print(info)
+# clear_cache(server_address)
+
+def get_vram():
+    # 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("调用集成的获取显卡显存代码失败！")
+
+get_vram()
+
diff --git a/BaiHu/doTask.py b/BaiHu/doTask.py
index 8ab3183c..de84d42f 100644
--- a/BaiHu/doTask.py
+++ b/BaiHu/doTask.py
@@ -285,10 +285,13 @@ if __name__ == '__main__':
     config = ConfigUtil.getConfig()
     # 处理机编号
     machine_id = config['system']['machine_id']
+
+    # webui 服务器地址
+    webui_server = config['webui']['webui_server']
     # 文生图服务地址
-    txt2img_url = config['webui']['txt2img_url']
+    txt2img_url = webui_server + config['webui']['txt2img_url']
     # 图生图服务地址
-    img2img_url = config['webui']['img2img_url']
+    img2img_url = webui_server + config['webui']['img2img_url']
 
     # WEB服务器地址
     web_url = config['webServer']['web_url']