ADD 增加midjourney文生图和描述图片节点

__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, \
+from .nodes.llm_nodes import LLMChat, LLMChatMultiModalImageUpload, LLMChatMultiModalImageTensor, Jinja2RenderTemplate
-    ModalClothesMask, ModalEditCustom
 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": "描述图片内容"
 }

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

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)