From bbd0c5325e5bcbfdcaf6b1d625114d591eb29aa0 Mon Sep 17 00:00:00 2001 From: iHeyTang Date: Tue, 12 Aug 2025 15:46:46 +0800 Subject: [PATCH] =?UTF-8?q?REFINE=20=E4=BC=98=E5=8C=96=E5=9B=BE=E5=83=8F?= =?UTF-8?q?=E5=A4=84=E7=90=86=E9=80=BB=E8=BE=91=EF=BC=8C=E7=BB=9F=E4=B8=80?= =?UTF-8?q?=E4=BB=A3=E7=A0=81=E9=A3=8E=E6=A0=BC=E5=B9=B6=E5=A2=9E=E5=BC=BA?= =?UTF-8?q?=E5=8F=AF=E8=AF=BB=E6=80=A7?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- nodes/image_modal_nodes.py | 266 ++++++++++++++++++++++++------------- 1 file changed, 170 insertions(+), 96 deletions(-) diff --git a/nodes/image_modal_nodes.py b/nodes/image_modal_nodes.py index f92c233..fa8af13 100644 --- a/nodes/image_modal_nodes.py +++ b/nodes/image_modal_nodes.py @@ -1,6 +1,6 @@ import io import json -from time import sleep +from time import sleep, time import folder_paths import requests @@ -29,7 +29,8 @@ def url_to_tensor(image_url: str, max_retries: int = 3): ValueError: 无效图片格式 """ headers = { - 'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36'} + "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36" + } for attempt in range(max_retries): try: @@ -38,8 +39,8 @@ def url_to_tensor(image_url: str, max_retries: int = 3): response.raise_for_status() # 检查内容类型是否为图像 - content_type = response.headers.get('Content-Type', '') - if not content_type.startswith('image/'): + content_type = response.headers.get("Content-Type", "") + if not content_type.startswith("image/"): raise ValueError(f"URL返回非图像内容: {content_type}") # 验证图像完整性 @@ -48,12 +49,10 @@ def url_to_tensor(image_url: str, max_retries: int = 3): raise ValueError("下载的内容过小，可能不是完整图像") # 尝试打开图像 - img = Image.open(io.BytesIO(img_data)).convert('RGB') + img = Image.open(io.BytesIO(img_data)).convert("RGB") # 转换为张量 - transform = transforms.Compose([ - transforms.ToTensor() - ]) + transform = transforms.Compose([transforms.ToTensor()]) return transform(img).unsqueeze(0).permute(0, 2, 3, 1) except (requests.exceptions.RequestException, ValueError) as e: @@ -70,7 +69,12 @@ class ModalClothesMask: "image": ("IMAGE",), "mask_color": ("STRING", {"default": "绿色"}), "clothes_type": ("STRING", {"default": "裤子"}), - "endpoint": ("STRING", {"default": "bowongai-dev--bowong-ai-video-gemini-fastapi-webapp.modal.run"}), + "endpoint": ( + "STRING", + { + "default": "bowongai-dev--bowong-ai-video-gemini-fastapi-webapp.modal.run" + }, + ), }, } @@ -80,25 +84,37 @@ class ModalClothesMask: OUTPUT_NODE = False CATEGORY = "不忘科技-自定义节点🚩/图片/Gemini" - def process(self, image: torch.Tensor, mask_color: str, clothes_type: str, endpoint: str): + def process( + self, image: torch.Tensor, mask_color: str, clothes_type: str, endpoint: str + ): try: timeout = 60 logger.info("获取token") - api_key = send_request("get", f"https://{endpoint}/google/access-token", - headers={'Authorization': 'Bearer bowong7777'}, timeout=timeout).json()[ - "access_token"] + api_key = send_request( + "get", + f"https://{endpoint}/google/access-token", + headers={"Authorization": "Bearer bowong7777"}, + timeout=timeout, + ).json()["access_token"] format = "PNG" logger.info("请求图像编辑") - job_resp = send_request("post", f"https://{endpoint}/google/image/clothes_mark", - headers={'x-google-api-key': api_key}, - data={ - "mark_clothes_type": clothes_type, - "mark_color": mask_color, - }, - files={"origin_image": ( - 'image.' + format.lower(), tensor_to_image_bytes(image, format), - f'image/{format.lower()}')}, - timeout=timeout) + job_resp = send_request( + "post", + f"https://{endpoint}/google/image/clothes_mark", + headers={"x-google-api-key": api_key}, + data={ + "mark_clothes_type": clothes_type, + "mark_color": mask_color, + }, + files={ + "origin_image": ( + "image." + format.lower(), + tensor_to_image_bytes(image, format), + f"image/{format.lower()}", + ) + }, + timeout=timeout, + ) job_resp.raise_for_status() job_resp = job_resp.json() if not job_resp["success"]: @@ -111,8 +127,12 @@ class ModalClothesMask: sleep(1) for _ in range(0, wait_time, interval): logger.info("查询任务状态") - result = send_request("get", f"https://{endpoint}/google/{job_id}", - headers={'Authorization': 'Bearer bowong7777'}, timeout=timeout) + result = send_request( + "get", + f"https://{endpoint}/google/{job_id}", + headers={"Authorization": "Bearer bowong7777"}, + timeout=timeout, + ) if result.status_code == 200: result = result.json() if result["status"] == "success": @@ -133,14 +153,22 @@ class ModalEditCustom: def INPUT_TYPES(cls): return { "required": { - "prompt": ("STRING", {"default": "将背景去除，输出原尺寸图片", "multiline": True}), + "prompt": ( + "STRING", + {"default": "将背景去除，输出原尺寸图片", "multiline": True}, + ), "temperature": ("FLOAT", {"default": 0.1, "min": 0, "max": 2}), "topP": ("FLOAT", {"default": 0.7, "min": 0, "max": 1}), - "endpoint": ("STRING", {"default": "bowongai-dev--bowong-ai-video-gemini-fastapi-webapp.modal.run"}), + "endpoint": ( + "STRING", + { + "default": "bowongai-dev--bowong-ai-video-gemini-fastapi-webapp.modal.run" + }, + ), }, "optional": { "image": ("IMAGE",), - } + }, } RETURN_TYPES = ("IMAGE",) @@ -149,31 +177,39 @@ class ModalEditCustom: OUTPUT_NODE = False CATEGORY = "不忘科技-自定义节点🚩/图片/Gemini" - def process(self, prompt: str, temperature: float, topP: float, endpoint: str, **kwargs): + def process( + self, prompt: str, temperature: float, topP: float, endpoint: str, **kwargs + ): try: timeout = 60 logger.info("获取token") - api_key = send_request("get", f"https://{endpoint}/google/access-token", - headers={'Authorization': 'Bearer bowong7777'}, timeout=timeout).json()[ - "access_token"] + api_key = send_request( + "get", + f"https://{endpoint}/google/access-token", + headers={"Authorization": "Bearer bowong7777"}, + timeout=timeout, + ).json()["access_token"] format = "PNG" if "image" in kwargs and kwargs["image"] is not None: image = kwargs["image"] - files = {"origin_image": ( - 'image.' + format.lower(), tensor_to_image_bytes(image, format), - f'image/{format.lower()}')} + files = { + "origin_image": ( + "image." + format.lower(), + tensor_to_image_bytes(image, format), + f"image/{format.lower()}", + ) + } else: files = None logger.info("请求图像编辑") - job_resp = send_request("post", f"https://{endpoint}/google/image/edit_custom", - headers={'x-google-api-key': api_key}, - data={ - "prompt": prompt, - "temperature": temperature, - "topP": topP - }, - files=files, - timeout=timeout) + job_resp = send_request( + "post", + f"https://{endpoint}/google/image/edit_custom", + headers={"x-google-api-key": api_key}, + data={"prompt": prompt, "temperature": temperature, "topP": topP}, + files=files, + timeout=timeout, + ) job_resp.raise_for_status() job_resp = job_resp.json() if not job_resp["success"]: @@ -186,8 +222,12 @@ class ModalEditCustom: sleep(1) for _ in range(0, wait_time, interval): logger.info("查询任务状态") - result = send_request("get", f"https://{endpoint}/google/{job_id}", - headers={'Authorization': 'Bearer bowong7777'}, timeout=timeout) + result = send_request( + "get", + f"https://{endpoint}/google/{job_id}", + headers={"Authorization": "Bearer bowong7777"}, + timeout=timeout, + ) if result.status_code == 200: result = result.json() if result["status"] == "success": @@ -208,14 +248,22 @@ class ModalMidJourneyGenerateImage: def INPUT_TYPES(cls): return { "required": { - "prompt": ("STRING", {"default": "一幅宏大壮美的山川画卷", "multiline": True}), - "provider":(["ttapi","302ai"],), - "endpoint": ("STRING", {"default": "bowongai-test--text-video-agent-fastapi-app.modal.run"}), + "prompt": ( + "STRING", + {"default": "一幅宏大壮美的山川画卷", "multiline": True}, + ), + "provider": (["ttapi", "302ai"],), + "endpoint": ( + "STRING", + { + "default": "bowongai-test--text-video-agent-fastapi-app.modal.run" + }, + ), "timeout": ("INT", {"default": 300, "min": 10, "max": 1200}), }, "optional": { "image": ("IMAGE",), - } + }, } RETURN_TYPES = ("IMAGE",) @@ -224,53 +272,67 @@ class ModalMidJourneyGenerateImage: OUTPUT_NODE = False CATEGORY = "不忘科技-自定义节点🚩/图片/Midjourney" - def process(self, prompt: str, provider:str, endpoint: str, timeout: int, **kwargs): + def process( + self, prompt: str, provider: str, endpoint: str, timeout: int, **kwargs + ): try: logger.info("请求同步接口") format = "PNG" if "image" in kwargs and kwargs["image"] is not None: image = kwargs["image"] - files = {"img_file": ( - 'image.' + format.lower(), tensor_to_image_bytes(image, format), - f'image/{format.lower()}')} + files = { + "img_file": ( + "image." + format.lower(), + tensor_to_image_bytes(image, format), + f"image/{format.lower()}", + ) + } else: files = None - if provider == "302ai": - interval = 3 - logger.info("提交任务") - logger.info(f"https://{endpoint}/api/custom/image/submit/task") - job_resp = send_request("post", f"https://{endpoint}/api/custom/image/submit/task", - data={"model_name":"302/mj", "prompt": prompt, "mode": "turbo"}, - files=files, - timeout=timeout) - job_resp.raise_for_status() - job_resp = job_resp.json() - if not job_resp["status"]: - raise Exception("生成失败, 可能因为风控") - job_id = job_resp["data"] - for _ in range(0, timeout // interval, interval): - logger.info("等待" + str(interval) + "秒") - sleep(interval) - logger.info("查询结果") - resp = send_request("get", f"https://{endpoint}/api/custom/task/status?task_id={job_id}", timeout=30) - resp.raise_for_status() - if resp.json()["status"] == "running": - logger.info("任务正在运行") - continue - if resp.json()["status"] == "failed": - raise Exception(f"生成失败: {resp.json()['msg']}") - if resp.json()["status"] == "success": - result_url = resp.json()["data"] - if not isinstance(result_url, list): - raise Exception("生成失败，返回结果为空") - result_list = [] - for url in result_url: - logger.success("img_url: " + url) - result_list.append(url_to_tensor(url).squeeze(0)) - result_list = torch.stack(result_list, dim=0) - return (result_list,) - raise Exception("等待超时") + interval = 3 + logger.info("提交任务") + logger.info(f"https://{endpoint}/api/custom/image/submit/task") + job_resp = send_request( + "post", + f"https://{endpoint}/api/custom/image/submit/task", + data={"model_name": "302/mj", "prompt": prompt, "mode": "turbo"}, + files=files, + timeout=timeout, + ) + job_resp.raise_for_status() + job_resp = job_resp.json() + if not job_resp["status"]: + raise Exception("生成失败, 可能因为风控") + job_id = job_resp["data"] + start_time = time() + for _ in range(0, timeout // interval, interval): + logger.info(f"已等待 {time() - start_time} 秒，{interval} 秒后查询...") + sleep(interval) + logger.info(f"查询结果 {job_id}") + resp = send_request( + "get", + f"https://{endpoint}/api/custom/task/status?task_id={job_id}", + timeout=30, + ) + resp.raise_for_status() + if resp.json()["status"] == "running": + logger.info(f"任务正在运行 {job_id}") + continue + if resp.json()["status"] == "failed": + raise Exception(f"生成失败: {resp.json()['msg']}") + if resp.json()["status"] == "success": + result_url = resp.json()["data"] + if not isinstance(result_url, list): + raise Exception("生成失败，返回结果为空") + result_list = [] + for url in result_url: + logger.success("img_url: " + url) + result_list.append(url_to_tensor(url).squeeze(0)) + result_list = torch.stack(result_list, dim=0) + logger.success(f"生成成功 {job_id}") + return (result_list,) + raise Exception("等待超时") except Exception as e: raise e @@ -281,7 +343,12 @@ class ModalMidJourneyDescribeImage: return { "required": { "image": ("IMAGE",), - "endpoint": ("STRING", {"default": "bowongai-test--text-video-agent-fastapi-app.modal.run"}), + "endpoint": ( + "STRING", + { + "default": "bowongai-test--text-video-agent-fastapi-app.modal.run" + }, + ), }, } @@ -295,12 +362,19 @@ class ModalMidJourneyDescribeImage: try: logger.info("请求同步接口") format = "PNG" - job_resp = send_request("post", f"https://{endpoint}/api/302/mj/sync/file/img/describe", - headers={'Authorization': 'Bearer bowong7777'}, - files={"img_file": ( - 'image.' + format.lower(), tensor_to_image_bytes(image, format), - f'image/{format.lower()}')}, - timeout=300) + job_resp = send_request( + "post", + f"https://{endpoint}/api/302/mj/sync/file/img/describe", + headers={"Authorization": "Bearer bowong7777"}, + files={ + "img_file": ( + "image." + format.lower(), + tensor_to_image_bytes(image, format), + f"image/{format.lower()}", + ) + }, + timeout=300, + ) job_resp.raise_for_status() job_resp = job_resp.json() if not job_resp["status"]: