diff --git a/BaiHu/Config.ini b/BaiHu/Config.ini index b34fc789..a15db0b0 100644 --- a/BaiHu/Config.ini +++ b/BaiHu/Config.ini @@ -17,7 +17,6 @@ img2img_url = /sdapi/v1/img2img ; WEB服务器地址 [webServer] web_url = https://www.hzkjai.com -;web_url = http://10.10.21.20:9000 [system] ; 处理机编号 diff --git a/BaiHu/Test/WD14.py b/BaiHu/Test/WD14.py index 5bbaa0e3..312e5956 100644 --- a/BaiHu/Test/WD14.py +++ b/BaiHu/Test/WD14.py @@ -1,49 +1,5 @@ -import requests -import base64 -from PIL import Image +from Util.SDUtil import get_wd_14 -url = 'http://192.168.1.21:7860/tagger/v1/interrogate' +server_address = "http://192.168.1.21:7860" image_path = r'D:\KeCheng\BaiHu\Backup\mote2.png' - -# 反推模型 -model = 'wd14-vit-v2-git' # 'wd14-convnext' -# 阀值 -threshold = 0.35 - -# 确认照片为上传照片 -image = Image.open(image_path) -# 将图片转换为Base64字符串 -with open(image_path, 'rb') as file: - image_data = file.read() - base64_image = base64.b64encode(image_data).decode('utf-8') - -image.close() - -# 构建请求体的JSON数据 -data = { - "image": base64_image, - "model": model, - "threshold": threshold -} - -# 发送POST请求 -response = requests.post(url, json=data) - -# 检查响应状态码 -if response.status_code == 200: - json_data = response.json() - # 处理返回的JSON数据 - caption_dict = json_data['caption'] - sorted_items = sorted(caption_dict.items(), key=lambda x: x[1], reverse=True) - # output = '\n'.join([f'{k}: {v}, {int(v * 100)}%' for k, v in sorted_items]) - # output = ','.join([f'{k.replace("_"," ")}' for k, v in sorted_items]) - - output = '' - for k, v in sorted_items: - if v > threshold: - output = output + "," + k.replace("_", " ") - output = output[1:] - print(output) -else: - print('Error:', response.status_code) - print('Response body:', response.text) +print(get_wd_14(server_address, image_path)) diff --git a/BaiHu/Util/SDUtil.py b/BaiHu/Util/SDUtil.py index 2215cf38..d1bac1ea 100644 --- a/BaiHu/Util/SDUtil.py +++ b/BaiHu/Util/SDUtil.py @@ -3,12 +3,43 @@ import time import urllib.parse import urllib.request - import websocket from Util.CommonUtil import * +# 获取反推词 +def get_wd_14(url, v_image_path): + # 反推模型 + model = 'wd14-vit-v2-git' + # 阀值 + threshold = 0.35 + + # 将图片转换为Base64字符串 + with open(v_image_path, 'rb') as file: + image_data = file.read() + base64_image = base64.b64encode(image_data).decode('utf-8') + + # 构建请求体的JSON数据 + data = { + "image": base64_image, + "model": model, + "threshold": threshold + } + # 发送POST请求 + response = requests.post(url, json=data) + json_data = response.json() + # 处理返回的JSON数据 + caption_dict = json_data['caption'] + sorted_items = sorted(caption_dict.items(), key=lambda x: x[1], reverse=True) + output = '' + for k, v in sorted_items: + if v > threshold: + output = output + "," + k.replace("_", " ") + output = output[1:] + return output + + # 定义一个函数向服务器队列发送提示信息 def queue_prompt(server_address, client_id, prompt): p = {"prompt": prompt, "client_id": client_id} @@ -191,7 +222,7 @@ def restart_server(webui_address): # 清理一下SD -def release_sd(webui_address,comfyui_address): +def release_sd(webui_address, comfyui_address): # comfyui的显存先清理一下 clear_comfyui_cache(comfyui_address) diff --git a/BaiHu/Util/__pycache__/SDUtil.cpython-310.pyc b/BaiHu/Util/__pycache__/SDUtil.cpython-310.pyc index a7b70481..b0da8a47 100644 Binary files a/BaiHu/Util/__pycache__/SDUtil.cpython-310.pyc and b/BaiHu/Util/__pycache__/SDUtil.cpython-310.pyc differ diff --git a/BaiHu/doTask.py b/BaiHu/doTask.py index 13400388..aea54c42 100644 --- a/BaiHu/doTask.py +++ b/BaiHu/doTask.py @@ -316,6 +316,10 @@ if __name__ == '__main__': web_url = config['webServer']['web_url'] # COMFYUI服务器地址 comfyui_address = config.get('comfyui', 'server_address') + # 反推接口地址 + wd_url = 'http://' + webui_address + '/tagger/v1/interrogate' + #image_path = r'D:\KeCheng\BaiHu\Backup\mote2.png' + #print(get_wd_14(wd_url, image_path)) while True: try: