diff --git a/BaiHu/Tools/TestComfyUIApi.py b/BaiHu/Tools/TestComfyUIApi.py
index 191f513f..68daba75 100644
--- a/BaiHu/Tools/TestComfyUIApi.py
+++ b/BaiHu/Tools/TestComfyUIApi.py
@@ -13,17 +13,6 @@ clear_comfyui_cache(comfyui_address)
 info = get_comfyui_used(comfyui_address)
 print(info)
 
-# 获取使显存数量
-total, used = get_webui_used(webui_address)
+# 清理一下sd
+release_sd()
 
-# 显卡显存占用超过20GB
-if float(used) > 20:
-    print("发现显存占用高达" + used + "GB，准备重启SD...")
-    # 重新启动一下SD
-    restart_server(webui_address)
-    # 休息2分钟以保证SD重启完成
-    time.sleep(2 * 60)
-else:
-    # 清理一下缓存即可
-    clear_webui_cache(webui_address)
-    print("显存占用正常，无需重新启动SD，清理一下缓存...")
diff --git a/BaiHu/Util/SDUtil.py b/BaiHu/Util/SDUtil.py
index f7be27e2..a73d5695 100644
--- a/BaiHu/Util/SDUtil.py
+++ b/BaiHu/Util/SDUtil.py
@@ -202,3 +202,20 @@ def clear_webui_cache(webui_address):
 def restart_server(webui_address):
     server_restart_url = "http://{}/sdapi/v1/server-restart".format(webui_address)
     submit_post(server_restart_url, None)
+
+
+# 清理一下SD
+def release_sd(webui_address):
+    # 获取使显存数量
+    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:
+        # 清理一下缓存即可
+        clear_webui_cache(webui_address)
+        print("显存占用正常，无需重新启动SD，清理一下缓存...")