import os.path 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} data = json.dumps(p).encode('utf-8') req = urllib.request.Request("http://{}/prompt".format(server_address), data=data) return json.loads(urllib.request.urlopen(req).read()) # 定义一个函数来获取图片 def get_image(server_address, filename, subfolder, folder_type): data = {"filename": filename, "subfolder": subfolder, "type": folder_type} url_values = urllib.parse.urlencode(data) with urllib.request.urlopen("http://{}/view?{}".format(server_address, url_values)) as response: return response.read() # 定义一个函数来获取历史记录 def get_history(server_address, prompt_id): with urllib.request.urlopen("http://{}/history/{}".format(server_address, prompt_id)) as response: return json.loads(response.read()) # 定义一个函数来获取图片，这涉及到监听WebSocket消息 def get_images(ws, server_address, client_id, prompt): prompt_id = queue_prompt(server_address, client_id, prompt)['prompt_id'] output_images = {} while True: out = ws.recv() if isinstance(out, str): message = json.loads(out) print(message) if message['type'] == 'executing': data = message['data'] if data['node'] is None and data['prompt_id'] == prompt_id: break # 执行完成 else: # printf("我在这里...") continue # 预览为二进制数据 history = get_history(server_address, prompt_id)[prompt_id] for o in history['outputs']: for node_id in history['outputs']: node_output = history['outputs'][node_id] # 图片分支 if 'images' in node_output: images_output = [] for image in node_output['images']: image_data = get_image(server_address, image['filename'], image['subfolder'], image['type']) images_output.append(image_data) output_images[node_id] = images_output # 视频分支 if 'videos' in node_output: videos_output = [] for video in node_output['videos']: video_data = get_image(server_address, video['filename'], video['subfolder'], video['type']) videos_output.append(video_data) output_images[node_id] = videos_output return output_images # 获取显卡使用率 def get_comfyui_used(server_address): req = urllib.request.Request("http://{}/system_stats".format(server_address)) res = json.loads(urllib.request.urlopen(req).read()) vram_total = res['devices'][0]['vram_total'] vram_total_str = str(int(vram_total / 1024 / 1024 / 1024 + 0.5)) vram_free = res['devices'][0]['vram_free'] used_vram = vram_total - vram_free used_vram_str = str(int((vram_total - vram_free) / 1024 / 1024 / 1024 + 0.5)) used_lv = round(1.0 * (used_vram) / vram_total * 100, 2) return "显存共" + vram_total_str + "GB,已使用" + used_vram_str + "GB,使用率：" + str(used_lv) + "% " # 参考文档：如何清空Comfyui的gpu缓存 # https://wailikeji.blog.csdn.net/article/details/140035515 # 清理GPU显存 def clear_comfyui_cache(server_address): # 打开文件并读取内容 file_path = r'../JSON/clearGPU.json' if not os.path.exists(file_path): file_path = r'./JSON/clearGPU.json' with open(file_path, 'r', encoding='utf-8') as file: prompt_data = json.load(file) # 清理 queue_prompt(server_address, "cleanGpuRam", prompt_data) # 生成图像 def generate_clip(server_address, prompt_data, client_id, output_path, myfilter): ws = websocket.WebSocket() ws.connect("ws://{}/ws?clientId={}".format(server_address, client_id)) images = get_images(ws, server_address, client_id, prompt_data) files = [] for node_id in images: # 过滤节点 if (myfilter is not None) and (node_id not in myfilter): continue for image_data in images[node_id]: # 使用格式化的时间戳在文件名中 file_id = str(str(uuid.uuid4())) # 创建文件的完整路径 files.append(file_id) GIF_LOCATION = "{}/{}.png".format(output_path, file_id) with open(GIF_LOCATION, "wb") as binary_file: # 写入二进制文件 binary_file.write(image_data) ws.close() return files def upload_file(server_address, file, subfolder="", overwrite=False): path = '' try: body = {"image": file} data = {} if overwrite: data["overwrite"] = "true" if subfolder: data["subfolder"] = subfolder resp = requests.post(f"http://{server_address}/upload/image", files=body, data=data) if resp.status_code == 200: data = resp.json() path = data["name"] if "subfolder" in data: if data["subfolder"] != "": path = data["subfolder"] + "/" + path else: print(f"{resp.status_code} - {resp.reason}") except Exception as error: print(error) return path # 获取webui使用情况 def get_webui_used(webui_server): # webui 服务器地址 vram_url = "http://{}/sdapi/v1/get_vram".format(webui_server) 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', '') used = jo[2].replace('used', '').replace(' ', '').replace(':', '').replace('GB', '') else: pass return total, used # 清理webui的缓存 def clear_webui_cache(webui_address): # webui 服务器地址 empty_cache_url = "http://{}/sdapi/v1/empty_cache".format(webui_address) res = submit_post(empty_cache_url, None) if res.status_code == 200: jo = json.loads(res.text) if jo['success']: pass else: print("显存清理失败！") else: print("调用集成的清理显卡显存代码失败！") # 重新启动SD 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, comfyui_address): pass # comfyui的显存先清理一下 # clear_comfyui_cache(comfyui_address) # 文生图一下，据测试在img_2_img报错，暴显存的时候，需要先清一下显存 # txt_2_img_blank(webui_address) # 清理一下缓存 # clear_webui_cache(webui_address)