diff --git a/__init__.py b/__init__.py index 61fed63..557d9bd 100644 --- a/__init__.py +++ b/__init__.py @@ -1,8 +1,9 @@ +from .nodes.image_modal_nodes import ModalEditCustom, ModalClothesMask, ModalMidJourneyGenerateImage, \ + ModalMidJourneyDescribeImage from .nodes.image_face_nodes import FaceDetect, FaceExtract from .nodes.image_gesture_nodes import JMGestureCorrect from .nodes.image_nodes import SaveImagePath, LoadNetImg, SaveImageWithOutput -from .nodes.llm_nodes import LLMChat, LLMChatMultiModalImageUpload, LLMChatMultiModalImageTensor, Jinja2RenderTemplate, \ - ModalClothesMask, ModalEditCustom +from .nodes.llm_nodes import LLMChat, LLMChatMultiModalImageUpload, LLMChatMultiModalImageTensor, Jinja2RenderTemplate from .nodes.object_storage_nodes import COSUpload, COSDownload, S3Download, S3Upload, S3UploadURL from .nodes.text_nodes import StringEmptyJudgement, LoadTextLocal, LoadTextOnline, RandomLineSelector from .nodes.util_nodes import LogToDB, TaskIdGenerate, TraverseFolder, UnloadAllModels, VodToLocalNode, \ @@ -43,7 +44,9 @@ NODE_CLASS_MAPPINGS = { "Jinja2RenderTemplate": Jinja2RenderTemplate, "JMGestureCorrect": JMGestureCorrect, "ModalClothesMask": ModalClothesMask, - "ModalEditCustom": ModalEditCustom + "ModalEditCustom": ModalEditCustom, + "ModalMidJourneyGenerateImage": ModalMidJourneyGenerateImage, + "ModalMidJourneyDescribeImage": ModalMidJourneyDescribeImage } NODE_DISPLAY_NAME_MAPPINGS = { @@ -79,5 +82,7 @@ NODE_DISPLAY_NAME_MAPPINGS = { "Jinja2RenderTemplate": "Jinja2格式Prompt模板渲染", "JMGestureCorrect": "人物侧身图片转为正面图-即梦", "ModalClothesMask": "模特指定衣服替换为指定颜色", - "ModalEditCustom": "自定义Prompt修改图片" + "ModalEditCustom": "自定义Prompt修改图片", + "ModalMidJourneyGenerateImage": "Prompt生图", + "ModalMidJourneyDescribeImage": "描述图片内容" } diff --git a/nodes/image_modal_nodes.py b/nodes/image_modal_nodes.py new file mode 100644 index 0000000..d232fe2 --- /dev/null +++ b/nodes/image_modal_nodes.py @@ -0,0 +1,266 @@ +import io +import json +from time import sleep + +import folder_paths +import requests +import torch +from PIL import Image +from loguru import logger +from torchvision import transforms + +from ..utils.http_utils import send_request +from ..utils.image_utils import tensor_to_image_bytes, base64_to_tensor + + +def url_to_tensor(image_url: str, max_retries: int = 3): + """ + 从URL下载图片并转换为PyTorch张量，增强错误处理能力 + + 参数: + image_url (str): 图片URL + max_retries (int): 最大重试次数 + + 返回: + torch.Tensor: 形状为[C, H, W]的张量 + + 异常: + HTTPError: 网络请求失败 + ValueError: 无效图片格式 + """ + headers = { + 'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36'} + + for attempt in range(max_retries): + try: + # 发送带User-Agent的请求 + response = requests.get(image_url, headers=headers, stream=True, timeout=15) + response.raise_for_status() + + # 检查内容类型是否为图像 + content_type = response.headers.get('Content-Type', '') + if not content_type.startswith('image/'): + raise ValueError(f"URL返回非图像内容: {content_type}") + + # 验证图像完整性 + img_data = response.content + if len(img_data) < 100: # 极小数据通常不是有效图像 + raise ValueError("下载的内容过小，可能不是完整图像") + + # 尝试打开图像 + img = Image.open(io.BytesIO(img_data)).convert('RGB') + + # 转换为张量 + transform = transforms.Compose([ + transforms.ToTensor() + ]) + return transform(img).unsqueeze(0).permute(0, 2, 3, 1) + + except (requests.exceptions.RequestException, ValueError) as e: + logger.warning(f"尝试 {attempt + 1}/{max_retries} 失败: {e}") + if attempt == max_retries - 1: + raise e + + +class ModalClothesMask: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE",), + "mask_color": ("STRING", {"default": "绿色"}), + "clothes_type": ("STRING", {"default": "裤子"}), + "endpoint": ("STRING", {"default": "bowongai-dev--bowong-ai-video-gemini-fastapi-webapp.modal.run"}), + }, + } + + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("image",) + FUNCTION = "process" + OUTPUT_NODE = False + CATEGORY = "不忘科技-自定义节点🚩/图片/Gemini图像编辑" + + 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"] + 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.raise_for_status() + job_resp = job_resp.json() + if not job_resp["success"]: + raise Exception("请求Modal API失败") + job_id = job_resp["taskId"] + + wait_time = 240 + interval = 2 + logger.info("开始轮询任务状态") + 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) + if result.status_code == 200: + result = result.json() + if result["status"] == "success": + logger.success("任务成功") + image_b64 = json.loads(result["result"])[0]["image_b64"] + image_tensor = base64_to_tensor(image_b64) + return (image_tensor,) + elif "fail" in result["status"].lower(): + raise Exception("任务失败") + sleep(interval) + raise Exception("查询任务状态超时") + except Exception as e: + raise Exception(e) + + +class ModalEditCustom: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE",), + "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"}), + }, + } + + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("image",) + FUNCTION = "process" + OUTPUT_NODE = False + CATEGORY = "不忘科技-自定义节点🚩/图片/Gemini图像编辑" + + def process(self, image: torch.Tensor, prompt: str, temperature: float, topP: float, 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"] + format = "PNG" + 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={"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"]: + raise Exception("请求Modal API失败") + job_id = job_resp["taskId"] + + wait_time = 240 + interval = 2 + logger.info("开始轮询任务状态") + 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) + if result.status_code == 200: + result = result.json() + if result["status"] == "success": + logger.success("任务成功") + image_b64 = json.loads(result["result"])[0]["image_b64"] + image_tensor = base64_to_tensor(image_b64) + return (image_tensor,) + elif "fail" in result["status"].lower(): + raise Exception("任务失败") + sleep(interval) + raise Exception("查询任务状态超时") + except Exception as e: + raise Exception(e) + + +class ModalMidJourneyGenerateImage: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "prompt": ("STRING", {"default": "一幅宏大壮美的山川画卷", "multiline": True}), + "endpoint": ("STRING", {"default": "bowongai-dev--bowong-ai-video-gemini-fastapi-webapp.modal.run"}), + }, + } + + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("image",) + FUNCTION = "process" + OUTPUT_NODE = False + CATEGORY = "不忘科技-自定义节点🚩/图片/Midjourney" + + def process(self, prompt: str, endpoint: str): + try: + logger.info("请求同步接口") + job_resp = send_request("post", f"https://{endpoint}/mj_router/sync/generate/image", + headers={'Authorization': 'Bearer bowong7777'}, + data={ + "prompt": prompt, + }, + timeout=60) + job_resp.raise_for_status() + job_resp = job_resp.json() + if "失败" in job_resp["msg"] or "fail" in job_resp["msg"] or "error" in job_resp["msg"]: + raise Exception("生成失败") + result_url = job_resp["data"] + logger.success("img_url: "+result_url) + return (url_to_tensor(result_url),) + except Exception as e: + raise e + +class ModalMidJourneyDescribeImage: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "img_url": ("STRING", {"default": "https://vcg03.cfp.cn/creative/vcg/800/new/VCG41N948031096.jpg", "multiline": True}), + "endpoint": ("STRING", {"default": "bowongai-dev--bowong-ai-video-gemini-fastapi-webapp.modal.run"}), + }, + } + + RETURN_TYPES = ("TEXT",) + RETURN_NAMES = ("描述内容",) + FUNCTION = "process" + OUTPUT_NODE = False + CATEGORY = "不忘科技-自定义节点🚩/图片/Midjourney" + + def process(self, img_url: str, endpoint: str): + try: + logger.info("请求同步接口") + job_resp = send_request("post", f"https://{endpoint}/mj_router/sync/describe/image", + headers={'Authorization': 'Bearer bowong7777'}, + data={ + "image_url": img_url, + }, + timeout=60) + job_resp.raise_for_status() + job_resp = job_resp.json() + if "失败" in job_resp["msg"] or "fail" in job_resp["msg"] or "error" in job_resp["msg"]: + raise Exception("描述失败") + result = job_resp["data"] + return (result,) + except Exception as e: + raise e diff --git a/nodes/llm_nodes.py b/nodes/llm_nodes.py index 1a59dd5..e5a6258 100644 --- a/nodes/llm_nodes.py +++ b/nodes/llm_nodes.py @@ -5,22 +5,16 @@ import json import os import re from mimetypes import guess_type -from time import sleep from typing import Any, Union import folder_paths import httpx import numpy as np -import requests import torch from PIL import Image from jinja2 import Template, StrictUndefined -from loguru import logger from retry import retry -from ..utils.http_utils import send_request -from ..utils.image_utils import tensor_to_image_bytes, base64_to_tensor - def find_value_recursive(key: str, data: Union[dict, list]) -> str | None | Any: if isinstance(data, dict): @@ -283,133 +277,3 @@ class Jinja2RenderTemplate: # 渲染模板 return (template.render(kv_map),) - - -class ModalClothesMask: - @classmethod - def INPUT_TYPES(cls): - return { - "required": { - "image": ("IMAGE",), - "mask_color": ("STRING", {"default": "绿色"}), - "clothes_type": ("STRING", {"default": "裤子"}), - "endpoint": ("STRING", {"default": "bowongai-dev--bowong-ai-video-gemini-fastapi-webapp.modal.run"}), - }, - } - - RETURN_TYPES = ("IMAGE",) - RETURN_NAMES = ("image",) - FUNCTION = "process" - OUTPUT_NODE = False - CATEGORY = "不忘科技-自定义节点🚩/图片/Gemini图像编辑" - - 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"] - 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.raise_for_status() - job_resp = job_resp.json() - if not job_resp["success"]: - raise Exception("请求Modal API失败") - job_id = job_resp["taskId"] - - wait_time = 240 - interval = 2 - logger.info("开始轮询任务状态") - 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) - if result.status_code == 200: - result = result.json() - if result["status"] == "success": - logger.success("任务成功") - image_b64 = json.loads(result["result"])[0]["image_b64"] - image_tensor = base64_to_tensor(image_b64) - return (image_tensor,) - elif "fail" in result["status"].lower(): - raise Exception("任务失败") - sleep(interval) - raise Exception("查询任务状态超时") - except Exception as e: - raise Exception(e) - - -class ModalEditCustom: - @classmethod - def INPUT_TYPES(cls): - return { - "required": { - "image": ("IMAGE",), - "prompt": ("STRING", {"default": "将背景去除，输出原尺寸图片", "multiline": True}), - "endpoint": ("STRING", {"default": "bowongai-dev--bowong-ai-video-gemini-fastapi-webapp.modal.run"}), - }, - } - - RETURN_TYPES = ("IMAGE",) - RETURN_NAMES = ("image",) - FUNCTION = "process" - OUTPUT_NODE = False - CATEGORY = "不忘科技-自定义节点🚩/图片/Gemini图像编辑" - - def process(self, image: torch.Tensor, prompt: 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"] - format = "PNG" - logger.info("请求图像编辑") - job_resp = send_request("post", f"https://{endpoint}/google/image/edit_custom", - headers={'x-google-api-key': api_key}, - data={ - "prompt": prompt - }, - 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"]: - raise Exception("请求Modal API失败") - job_id = job_resp["taskId"] - - wait_time = 240 - interval = 2 - logger.info("开始轮询任务状态") - 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) - if result.status_code == 200: - result = result.json() - if result["status"] == "success": - logger.success("任务成功") - image_b64 = json.loads(result["result"])[0]["image_b64"] - image_tensor = base64_to_tensor(image_b64) - return (image_tensor,) - elif "fail" in result["status"].lower(): - raise Exception("任务失败") - sleep(interval) - raise Exception("查询任务状态超时") - except Exception as e: - raise Exception(e)