Release v53.2.0
This commit is contained in:
chenjunhong
+10
-4
@@ -3,12 +3,18 @@
|
||||
__all__ = [
|
||||
"NODE_CLASS_MAPPINGS",
|
||||
"NODE_DISPLAY_NAME_MAPPINGS",
|
||||
|
||||
|
||||
]
|
||||
|
||||
__author__ = """xuyucai"""
|
||||
__email__ = "986327386@qq.com"
|
||||
__version__ = "53.1.2"
|
||||
__version__ = "53.2.0"
|
||||
|
||||
from .src.Comfyui_Fd_Nodes.nodes import NODE_CLASS_MAPPINGS
|
||||
from .src.Comfyui_Fd_Nodes.nodes import NODE_DISPLAY_NAME_MAPPINGS
|
||||
try:
|
||||
from .src.Comfyui_Fd_Nodes.nodes import NODE_CLASS_MAPPINGS
|
||||
from .src.Comfyui_Fd_Nodes.nodes import NODE_DISPLAY_NAME_MAPPINGS
|
||||
except ModuleNotFoundError as exc:
|
||||
if not exc.name.startswith("comfy"):
|
||||
raise
|
||||
NODE_CLASS_MAPPINGS = {}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {}
|
||||
|
||||
+1
-1
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
|
||||
|
||||
[project]
|
||||
name = "Comfyui_Fd_Nodes"
|
||||
version = "53.1.2" # 999开头的版本是debug版
|
||||
version = "53.2.0" # 999开头的版本是debug版
|
||||
description = "Custom nodes for fd project"
|
||||
authors = [
|
||||
{name = "xuyucai", email = "986327386@qq.com"}
|
||||
|
||||
@@ -0,0 +1,198 @@
|
||||
# 服装退底接口 remove-bg-by-meitu 重构文档
|
||||
|
||||
## 一、接口概述
|
||||
|
||||
**业务含义**: 接收一张服装图片 URL，调用外部 AI 服务（美图）进行背景移除（抠图），返回处理后的图片、mask 和红边图片 URL。
|
||||
|
||||
**当前实现方式**: Spring Cloud OpenFeign 直连外部 AI 服务，Controller 层透传，无额外业务逻辑。
|
||||
|
||||
**重构目标**: 重新实现请求逻辑（具体目标待确认，以下文档供 Codex 理解现有代码结构）。
|
||||
|
||||
---
|
||||
|
||||
## 二、调用链路
|
||||
|
||||
```
|
||||
HTTP 请求 (POST /image-combine/mt/remove-bg-by-meitu)
|
||||
→ ImageCombineMTController.removeBgByMeitu()
|
||||
→ ImageCombineMTServingClient.removeBgByMeitu() [Feign 调用]
|
||||
→ POST {image.detail.url}/image/remove_bg_by_meitu [外部 AI 服务]
|
||||
← ClothRemoveBGDTO
|
||||
← ApiResult<ClothRemoveBGDTO>
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 三、涉及文件清单（绝对路径）
|
||||
|
||||
### 3.1 Controller 层
|
||||
|
||||
**文件**: `/Users/chenjunhong/Documents/javaProject/fashion-parent/app-business-clothingAI/src/main/java/com/zhiyi/fashion/clothingai/server/facade/http/controller/ImageCombineMTController.java`
|
||||
|
||||
- 第 74-78 行: `remove-bg-by-meitu` 接口入口
|
||||
- 依赖注入: `ImageCombineMTServingClient`, `XinChenManager`, `ImageProcessStatisticsManager`
|
||||
- 注意: 当前该接口**没有** XinChen 团队判断和统计埋点，但同文件中 `merge-forge-back` 接口（第 108-118 行）有这些逻辑，可作为参考
|
||||
|
||||
```java
|
||||
// 第 74-78 行
|
||||
@ApiOperation(value = "服装退底", httpMethod = "POST")
|
||||
@PostMapping("/remove-bg-by-meitu")
|
||||
public ApiResult<ClothRemoveBGDTO> removeBgByMeitu(HttpServletRequest request, @RequestBody ClothRemoveBGRequest clothRemoveBGRequest) {
|
||||
ClothRemoveBGDTO clothRemoveBGDTO = imageCombineMTServingClient.removeBgByMeitu(clothRemoveBGRequest);
|
||||
return ApiResult.success(clothRemoveBGDTO);
|
||||
}
|
||||
```
|
||||
|
||||
### 3.2 Feign Client 层（HTTP 客户端）
|
||||
|
||||
**文件**: `/Users/chenjunhong/Documents/javaProject/fashion-parent/app-business-clothingAI/src/main/java/com/zhiyi/fashion/clothingai/server/infra/api/ai/ImageCombineMTServingClient.java`
|
||||
|
||||
- 第 51-52 行: `removeBgByMeitu` 方法定义
|
||||
- 第 15 行: `@FeignClient` 配置，基础 URL 为 `${image.detail.url}`
|
||||
- 第 60-61 行: 还有一个 API 版本 `remove_bg_by_api`，接收相同请求/响应
|
||||
|
||||
```java
|
||||
// 第 51-52 行
|
||||
@PostMapping("/image/remove_bg_by_meitu")
|
||||
ClothRemoveBGDTO removeBgByMeitu(@RequestBody ClothRemoveBGRequest request);
|
||||
```
|
||||
|
||||
### 3.3 Feign 序列化配置
|
||||
|
||||
**文件**: `/Users/chenjunhong/Documents/javaProject/fashion-parent/app-business-clothingAI/src/main/java/com/zhiyi/fashion/clothingai/server/infra/api/ai/config/FeignSnakeCaseConfig.java`
|
||||
|
||||
- 功能: 将 Java camelCase 字段自动转为 snake_case 发送给下游服务
|
||||
- 影响范围: 所有使用此 configuration 的 Feign Client
|
||||
|
||||
### 3.4 请求对象
|
||||
|
||||
**文件**: `/Users/chenjunhong/Documents/javaProject/fashion-parent/app-business-clothingAI/src/main/java/com/zhiyi/fashion/clothingai/server/infra/api/ai/request/ClothRemoveBGRequest.java`
|
||||
|
||||
```java
|
||||
@Data
|
||||
public class ClothRemoveBGRequest {
|
||||
@ApiModelProperty("图片链接")
|
||||
private String imageUrl;
|
||||
|
||||
private Integer repaintEdge;
|
||||
|
||||
@ApiModelProperty("红边框粗细，默认40")
|
||||
private Integer edgeThickness = 40;
|
||||
}
|
||||
```
|
||||
|
||||
经 FeignSnakeCaseConfig 转换后，实际发送给下游的 JSON 格式:
|
||||
```json
|
||||
{
|
||||
"image_url": "https://xxx/xxx.jpg",
|
||||
"repaint_edge": 1,
|
||||
"edge_thickness": 40
|
||||
}
|
||||
```
|
||||
|
||||
### 3.5 响应对象
|
||||
|
||||
**基础 DTO 文件**: `/Users/chenjunhong/Documents/javaProject/fashion-parent/app-business-clothingAI/src/main/java/com/zhiyi/fashion/clothingai/server/infra/api/ai/dto/ImageCombineBaseDTO.java`
|
||||
|
||||
```java
|
||||
@Data
|
||||
public class ImageCombineBaseDTO implements Serializable {
|
||||
private Boolean status;
|
||||
@JsonProperty("cost_time")
|
||||
private String costTime;
|
||||
@JsonProperty("server_name")
|
||||
private String serverName;
|
||||
private ImageCombineError error;
|
||||
|
||||
@Data
|
||||
public static class ImageCombineError {
|
||||
private String code;
|
||||
private String message;
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
**业务 DTO 文件**: `/Users/chenjunhong/Documents/javaProject/fashion-parent/app-business-clothingAI/src/main/java/com/zhiyi/fashion/clothingai/server/infra/api/ai/dto/ClothRemoveBGDTO.java`
|
||||
|
||||
```java
|
||||
@Data
|
||||
public class ClothRemoveBGDTO extends ImageCombineBaseDTO implements Serializable {
|
||||
private String resultUrl;
|
||||
private String mainMaskUrl;
|
||||
private String RedEdgeImageUrl; // 注意: 这里是大写 R 开头，snake_case 转换后为 red_edge_image_url
|
||||
}
|
||||
```
|
||||
|
||||
下游返回的 JSON 格式:
|
||||
```json
|
||||
{
|
||||
"status": true,
|
||||
"cost_time": "1.23s",
|
||||
"server_name": "xxx",
|
||||
"error": null,
|
||||
"result_url": "https://xxx/result.jpg",
|
||||
"main_mask_url": "https://xxx/mask.png",
|
||||
"red_edge_image_url": "https://xxx/edge.png"
|
||||
}
|
||||
```
|
||||
|
||||
### 3.6 可参考的辅助类（同模块已有逻辑）
|
||||
|
||||
**XinChen 团队处理**: `/Users/chenjunhong/Documents/javaProject/fashion-parent/app-business-clothingAI/src/main/java/com/zhiyi/fashion/clothingai/server/domain/manager/XinChenManager.java`
|
||||
|
||||
**图片处理统计**: `/Users/chenjunhong/Documents/javaProject/fashion-parent/app-business-clothingAI/src/main/java/com/zhiyi/fashion/clothingai/server/domain/manager/ImageProcessStatisticsManager.java`
|
||||
|
||||
---
|
||||
|
||||
## 四、同文件中可参考的完整业务逻辑（merge-forge-back）
|
||||
|
||||
`ImageCombineMTController.java` 第 106-118 行的 `merge-forge-back` 接口包含了完整的业务处理模式，包含:
|
||||
1. Feign 调用
|
||||
2. XinChen 团队图片转存判断
|
||||
3. 统计埋点
|
||||
|
||||
```java
|
||||
@ApiOperation(value = "背景融合接口", httpMethod = "POST")
|
||||
@PostMapping("/merge-forge-back")
|
||||
public ApiResult<MergeForgeBackDTO> mergeForgeBack(HttpServletRequest request, @RequestBody MergeForgeBackRequest mergeForgeBackRequest) {
|
||||
MergeForgeBackDTO mainMaskDTO = imageCombineMTServingClient.mergeForgeBack(mergeForgeBackRequest);
|
||||
RequestContext requestContext = RequestHelper.formContext(request);
|
||||
if (xinChenManager.isXinChenTeam(requestContext.getTeamId())) {
|
||||
mainMaskDTO.setImageUrl(xinChenManager.copy2XinChenImage(mainMaskDTO.getImageUrl()));
|
||||
mainMaskDTO.setMainMaskUrl(xinChenManager.copy2XinChenImage(mainMaskDTO.getMainMaskUrl()));
|
||||
mainMaskDTO.setResultUrl(xinChenManager.copy2XinChenImage(mainMaskDTO.getResultUrl()));
|
||||
}
|
||||
imageProcessStatisticsManager.saveImageProcessStatistics(mergeForgeBackRequest.getImageUrl(), 5);
|
||||
return ApiResult.success(mainMaskDTO);
|
||||
}
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 五、项目包结构
|
||||
|
||||
```
|
||||
com.zhiyi.fashion.clothingai.server
|
||||
├── facade.http.controller → Controller 层
|
||||
│ └── ImageCombineMTController.java
|
||||
├── infra.api.ai → 外部 API 调用层
|
||||
│ ├── ImageCombineMTServingClient.java (Feign Client 接口)
|
||||
│ ├── config/
|
||||
│ │ └── FeignSnakeCaseConfig.java (snake_case 序列化配置)
|
||||
│ ├── request/
|
||||
│ │ └── ClothRemoveBGRequest.java (请求体)
|
||||
│ └── dto/
|
||||
│ ├── ImageCombineBaseDTO.java (基础响应)
|
||||
│ └── ClothRemoveBGDTO.java (业务响应)
|
||||
└── domain.manager → 业务逻辑层
|
||||
├── XinChenManager.java
|
||||
└── ImageProcessStatisticsManager.java
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 六、关键配置
|
||||
|
||||
- **外部服务基础 URL**: `${image.detail.url}` — 在 Apollo 或环境变量中配置
|
||||
- **Feign 超时**: 在 `bootstrap.yml` 中配置
|
||||
- **命名策略**: camelCase → snake_case（双向，请求和响应都自动转换）
|
||||
@@ -1,6 +1,3 @@
|
||||
import json
|
||||
import os
|
||||
|
||||
from .config import (
|
||||
FD_LITELLM_API_KEY,
|
||||
FD_LITELLM_BASE_URL,
|
||||
|
||||
@@ -0,0 +1,5 @@
|
||||
from .gpt_image_node import FD_GTPImage
|
||||
|
||||
GPTImageEditNode = FD_GTPImage
|
||||
|
||||
__all__ = ["GPTImageEditNode"]
|
||||
@@ -2,8 +2,6 @@ import io
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import random
|
||||
import time
|
||||
import traceback
|
||||
from datetime import datetime
|
||||
from inspect import cleandoc
|
||||
@@ -50,6 +48,7 @@ from .gpt_multi_image_node import FD_GPTMultiImage
|
||||
from .prompt_nodes import EcommercePromptGenerator, PromptListSelector
|
||||
from .zhiyi_text_node import ZhiYiTextGenNode
|
||||
from .zhiyi_image_text_node import ZhiYiImageTextNode
|
||||
from .zhiyi_image_text_combo_node import ZhiYiImageTextComboNode
|
||||
from .zhiyi_image_to_image_node import ZhiYiImageToImageNode
|
||||
from .toggle import NodeToggleByID
|
||||
from .utils.logging_utils import configure_default_logging
|
||||
@@ -57,6 +56,7 @@ from .zhiyi_image_combo_node import ZhiYiImageComboNode
|
||||
from .zhiyi_image_to_image_combo_node import ZhiYiImageToImageComboNode
|
||||
from .aistudio_image_combo_node import ZhiYiAiStudioImageComboNode
|
||||
from .remove_bg_by_meitu_node import ZhiYiRemoveBgByMeituNode
|
||||
from .utils.gpt_image_size import resolution_to_edit_size
|
||||
|
||||
configure_default_logging()
|
||||
logger = logging.getLogger(__name__)
|
||||
@@ -65,6 +65,10 @@ logger = logging.getLogger(__name__)
|
||||
FD_REMOVE_WATERMARK_SERVICE_URL = os.getenv("FD_REMOVE_WATERMARK_SERVICE_URL", "http://localhost:8000/v1/process")
|
||||
|
||||
|
||||
def _resolution_to_edit_size(resolution: str, aspect_ratio: str) -> str:
|
||||
return resolution_to_edit_size(resolution, aspect_ratio)
|
||||
|
||||
|
||||
class FD_RemoveWatermark:
|
||||
"""
|
||||
Remove watermark from image using AI inpainting service
|
||||
@@ -883,6 +887,7 @@ NODE_CLASS_MAPPINGS = {
|
||||
"MaoziEcommercePromptGenerator": EcommercePromptGenerator,
|
||||
"MaoziPromptListSelector": PromptListSelector,
|
||||
"ZhiYiImageTextNode": ZhiYiImageTextNode,
|
||||
"ZhiYiImageTextComboNode": ZhiYiImageTextComboNode,
|
||||
"ZhiYiImageToImageNode": ZhiYiImageToImageNode,
|
||||
"ZhiYiTextGenNode": ZhiYiTextGenNode,
|
||||
"NodeToggleByID": NodeToggleByID,
|
||||
@@ -907,6 +912,7 @@ NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"MaoziEcommercePromptGenerator": "猫子提示词节点-详情页生成器",
|
||||
"MaoziPromptListSelector": "猫子提示词节点-列表选择器",
|
||||
"ZhiYiImageTextNode": "知衣-图生文",
|
||||
"ZhiYiImageTextComboNode": "知衣-图生文-combo",
|
||||
"ZhiYiImageToImageNode": "知衣-图生图",
|
||||
"ZhiYiTextGenNode": "知衣-文生文",
|
||||
"NodeToggleByID": "节点开关 (按ID)",
|
||||
|
||||
@@ -1,5 +1,4 @@
|
||||
import json
|
||||
import os
|
||||
import time
|
||||
|
||||
import oss2
|
||||
@@ -107,4 +106,4 @@ class FD_imgToText_Doubao:
|
||||
return (data["response"]["Result"],)
|
||||
except Exception as e:
|
||||
print(e)
|
||||
return (defaultPrompt,)
|
||||
return (defaultPrompt,)
|
||||
|
||||
@@ -26,6 +26,8 @@ logger = logging.getLogger(__name__)
|
||||
MODEL_NAME_MAP = {
|
||||
"gemini-2.5-flash-image-preview": "google/gemini-2.5-flash-image-preview",
|
||||
"gemini-3-pro-image-preview": "google/gemini-3-pro-image-preview",
|
||||
"gemini-3-pro-image-preview-aistudio": "google/gemini-3-pro-image-preview-official",
|
||||
"gemini-3-pro-image-preview-official": "google/gemini-3-pro-image-preview-official",
|
||||
"gemini-3.1-flash-image-preview": "google/gemini-3.1-flash-image-preview",
|
||||
}
|
||||
|
||||
|
||||
@@ -16,7 +16,6 @@ logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
LITELLM_GEMINI_MODELS = {
|
||||
"gemini-3-pro-image-preview-aistudio",
|
||||
"gemini-3-pro-image-preview-siphonlab",
|
||||
}
|
||||
|
||||
|
||||
@@ -0,0 +1,399 @@
|
||||
import json
|
||||
import logging
|
||||
import time
|
||||
import traceback
|
||||
from concurrent.futures import ThreadPoolExecutor, as_completed
|
||||
|
||||
import requests
|
||||
|
||||
from .config_manager import load_config
|
||||
from .utils.error_utils import normalize_error_message
|
||||
from .utils.logging_utils import configure_default_logging
|
||||
from .zhiyi_image_text_node import ZhiYiImageTextNode
|
||||
|
||||
configure_default_logging()
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class ZhiYiImageTextComboNode(ZhiYiImageTextNode):
|
||||
"""知衣多图图生文节点 - 接收最多8个图片组合并发调用豆包视觉理解 API。"""
|
||||
|
||||
DEFAULT_PROMPT = "介绍这幅图片"
|
||||
MODEL = "doubao-seed-2.0-mini"
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"max_concurrency": ("INT", {
|
||||
"default": 8,
|
||||
"min": 1,
|
||||
"max": 64,
|
||||
"step": 1,
|
||||
"tooltip": "同时发送请求的最大数量，避免触发 API 限流",
|
||||
}),
|
||||
"temperature": ("FLOAT", {
|
||||
"default": 0.7,
|
||||
"min": 0.0,
|
||||
"max": 2.0,
|
||||
"step": 0.05,
|
||||
}),
|
||||
"max_tokens": ("INT", {
|
||||
"default": 2048,
|
||||
"min": 1,
|
||||
"max": 8192,
|
||||
"step": 1,
|
||||
}),
|
||||
"retry_count": ("INT", {
|
||||
"default": 1,
|
||||
"min": 0,
|
||||
"max": 3,
|
||||
"step": 1,
|
||||
"tooltip": "请求失败后的重试次数。仅重试网络错误、超时、429 和 5xx",
|
||||
}),
|
||||
},
|
||||
"optional": {
|
||||
"node_switch": ("INT", {
|
||||
"default": 0,
|
||||
"min": 0,
|
||||
"max": 1,
|
||||
"step": 1,
|
||||
"display": "number",
|
||||
}),
|
||||
"combo_1": ("ZHIYI_COMBO",),
|
||||
"combo_2": ("ZHIYI_COMBO",),
|
||||
"combo_3": ("ZHIYI_COMBO",),
|
||||
"combo_4": ("ZHIYI_COMBO",),
|
||||
"combo_5": ("ZHIYI_COMBO",),
|
||||
"combo_6": ("ZHIYI_COMBO",),
|
||||
"combo_7": ("ZHIYI_COMBO",),
|
||||
"combo_8": ("ZHIYI_COMBO",),
|
||||
"system_prompt": ("STRING", {
|
||||
"default": "",
|
||||
"multiline": True,
|
||||
}),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("STRING", "STRING")
|
||||
RETURN_NAMES = ("text", "log")
|
||||
FUNCTION = "generate"
|
||||
CATEGORY = "知衣/图生文"
|
||||
OUTPUT_NODE = False
|
||||
|
||||
def _expand_images(self, combo_images):
|
||||
result = []
|
||||
for image in combo_images:
|
||||
if image is None:
|
||||
continue
|
||||
if image.ndim == 4:
|
||||
for i in range(image.shape[0]):
|
||||
result.append(image[i:i + 1])
|
||||
elif image.ndim == 3:
|
||||
result.append(image.unsqueeze(0))
|
||||
else:
|
||||
raise RuntimeError(f"图片 tensor 维度错误: {image.ndim}")
|
||||
return result
|
||||
|
||||
def _select_prompt(self, combo):
|
||||
combo_prompts = combo.get("prompts")
|
||||
if isinstance(combo_prompts, str):
|
||||
combo_prompts = [combo_prompts]
|
||||
if isinstance(combo_prompts, list):
|
||||
for prompt in combo_prompts:
|
||||
if isinstance(prompt, str) and prompt.strip():
|
||||
return prompt.strip()
|
||||
|
||||
prompt = combo.get("prompt", "")
|
||||
if isinstance(prompt, str) and prompt.strip():
|
||||
return prompt.strip()
|
||||
return self.DEFAULT_PROMPT
|
||||
|
||||
def _build_messages(self, prompt, data_urls, system_prompt=""):
|
||||
messages = []
|
||||
system_text = (system_prompt or "").strip()
|
||||
if system_text:
|
||||
messages.append({
|
||||
"role": "system",
|
||||
"content": [{"type": "text", "text": system_text}],
|
||||
})
|
||||
|
||||
content = [{"type": "text", "text": prompt}]
|
||||
for data_url in data_urls:
|
||||
content.append({"type": "image_url", "image_url": {"url": data_url}})
|
||||
|
||||
messages.append({"role": "user", "content": content})
|
||||
return messages
|
||||
|
||||
def _extract_text(self, result):
|
||||
text = result["choices"][0]["message"]["content"]
|
||||
if isinstance(text, list):
|
||||
return "".join(
|
||||
part.get("text", "") for part in text if isinstance(part, dict)
|
||||
)
|
||||
if text is None:
|
||||
return ""
|
||||
return str(text)
|
||||
|
||||
def _is_retryable_status(self, status_code):
|
||||
return status_code == 429 or 500 <= status_code < 600
|
||||
|
||||
def _is_retryable_exception(self, exc):
|
||||
if isinstance(exc, (requests.exceptions.Timeout, requests.exceptions.ConnectionError)):
|
||||
return True
|
||||
|
||||
response = getattr(exc, "response", None)
|
||||
status_code = getattr(response, "status_code", None)
|
||||
return status_code is not None and self._is_retryable_status(status_code)
|
||||
|
||||
def _sleep_before_retry(self, attempt_index):
|
||||
time.sleep(min(2 ** attempt_index, 2))
|
||||
|
||||
def _single_request(self, url, api_key, messages, temperature, max_tokens, retry_count):
|
||||
payload = {
|
||||
"stream": False,
|
||||
"model": self.MODEL,
|
||||
"messages": messages,
|
||||
"temperature": temperature,
|
||||
"max_tokens": max_tokens,
|
||||
}
|
||||
max_attempts = max(1, int(retry_count) + 1)
|
||||
last_error = None
|
||||
|
||||
for attempt_index in range(max_attempts):
|
||||
response = None
|
||||
try:
|
||||
logger.info(
|
||||
"Calling ZhiYi image-to-text combo API with payload=%s",
|
||||
{
|
||||
"url": url,
|
||||
"stream": payload["stream"],
|
||||
"model": payload["model"],
|
||||
"temperature": payload["temperature"],
|
||||
"max_tokens": payload["max_tokens"],
|
||||
"attempt": attempt_index + 1,
|
||||
"max_attempts": max_attempts,
|
||||
"messages": self._summarize_messages_for_log(messages),
|
||||
},
|
||||
)
|
||||
response = requests.post(
|
||||
url=url,
|
||||
headers={
|
||||
"Authorization": f"Bearer {api_key}",
|
||||
"Content-Type": "application/json",
|
||||
},
|
||||
data=json.dumps(payload),
|
||||
timeout=600,
|
||||
)
|
||||
|
||||
status_code = getattr(response, "status_code", None)
|
||||
if status_code is not None and status_code >= 400:
|
||||
response_text = getattr(response, "text", "")[:1000]
|
||||
last_error = RuntimeError(f"API 请求失败: {status_code}\n{response_text}")
|
||||
if self._is_retryable_status(status_code) and attempt_index < max_attempts - 1:
|
||||
logger.warning(
|
||||
"ZhiYi image-to-text combo API attempt %s/%s failed with HTTP %s, retrying",
|
||||
attempt_index + 1,
|
||||
max_attempts,
|
||||
status_code,
|
||||
)
|
||||
self._sleep_before_retry(attempt_index)
|
||||
continue
|
||||
raise last_error
|
||||
|
||||
response.raise_for_status()
|
||||
result = response.json()
|
||||
logger.info(
|
||||
"ZhiYi image-to-text combo API response summary: %s",
|
||||
{
|
||||
"status_code": response.status_code,
|
||||
**self._summarize_response_for_log(result),
|
||||
},
|
||||
)
|
||||
return self._extract_text(result), attempt_index
|
||||
except (KeyError, IndexError, TypeError, ValueError) as exc:
|
||||
response_text = response.text[:500] if response is not None else ""
|
||||
raise RuntimeError(f"解析响应失败: {exc}\n原始响应: {response_text}") from exc
|
||||
except requests.exceptions.RequestException as exc:
|
||||
last_error = exc
|
||||
if self._is_retryable_exception(exc) and attempt_index < max_attempts - 1:
|
||||
logger.warning(
|
||||
"ZhiYi image-to-text combo API attempt %s/%s failed, retrying: %s",
|
||||
attempt_index + 1,
|
||||
max_attempts,
|
||||
exc,
|
||||
)
|
||||
self._sleep_before_retry(attempt_index)
|
||||
continue
|
||||
raise RuntimeError(f"API 请求失败: {exc}") from exc
|
||||
|
||||
raise RuntimeError(last_error or "API 请求失败")
|
||||
|
||||
def _run_concurrent(self, tasks, max_workers, label="请求"):
|
||||
results = {}
|
||||
log_lines = []
|
||||
last_error_message = None
|
||||
workers = min(len(tasks), max(1, int(max_workers)))
|
||||
with ThreadPoolExecutor(max_workers=workers) as executor:
|
||||
futures = {
|
||||
executor.submit(fn, *args): (result_index, combo_slot)
|
||||
for result_index, combo_slot, fn, args in tasks
|
||||
}
|
||||
for future in as_completed(futures):
|
||||
result_index, combo_slot = futures[future]
|
||||
try:
|
||||
text, retry_used = future.result()
|
||||
results[result_index] = {"ok": True, "text": text}
|
||||
suffix = f"，重试 {retry_used} 次" if retry_used else ""
|
||||
log_lines.append(f"[{label} combo_{combo_slot}] 成功{suffix}")
|
||||
except Exception as exc:
|
||||
normalized_error = normalize_error_message(exc)
|
||||
last_error_message = normalized_error
|
||||
results[result_index] = {"ok": False, "text": f"ERROR: {normalized_error}"}
|
||||
log_lines.append(f"[{label} combo_{combo_slot}] 失败: {normalized_error}")
|
||||
print(f"[知衣图生文-combo] combo_{combo_slot} 失败: {normalized_error}")
|
||||
traceback.print_exc()
|
||||
return results, log_lines, last_error_message
|
||||
|
||||
def _format_text_output(self, combo_inputs, result_items):
|
||||
successful = [item for item in result_items if item and item.get("ok")]
|
||||
if len(combo_inputs) == 1 and len(successful) == 1:
|
||||
return successful[0]["text"]
|
||||
|
||||
sections = []
|
||||
for index, (combo_slot, _combo) in enumerate(combo_inputs):
|
||||
item = result_items[index] or {
|
||||
"ok": False,
|
||||
"text": "ERROR: UNKNOWN: 未执行",
|
||||
}
|
||||
sections.append(f"[combo_{combo_slot}]\n{item['text']}")
|
||||
return "\n\n".join(sections)
|
||||
|
||||
def generate(
|
||||
self,
|
||||
max_concurrency=8,
|
||||
temperature=0.7,
|
||||
max_tokens=2048,
|
||||
retry_count=1,
|
||||
node_switch=0,
|
||||
combo_1=None,
|
||||
combo_2=None,
|
||||
combo_3=None,
|
||||
combo_4=None,
|
||||
combo_5=None,
|
||||
combo_6=None,
|
||||
combo_7=None,
|
||||
combo_8=None,
|
||||
system_prompt="",
|
||||
):
|
||||
if node_switch == 1:
|
||||
return ("", "node skipped")
|
||||
|
||||
cfg = load_config()
|
||||
base_url = cfg["base_url"].rstrip("/")
|
||||
api_key = cfg["api_key"]
|
||||
url = f"{base_url}/v1/chat/completions"
|
||||
|
||||
combo_inputs = [
|
||||
(slot, combo)
|
||||
for slot, combo in enumerate(
|
||||
[combo_1, combo_2, combo_3, combo_4, combo_5, combo_6, combo_7, combo_8],
|
||||
start=1,
|
||||
)
|
||||
if combo is not None
|
||||
]
|
||||
if not combo_inputs:
|
||||
raise RuntimeError("未提供任何组合输入，请连接至少一个 combo")
|
||||
|
||||
result_items = [None] * len(combo_inputs)
|
||||
tasks = []
|
||||
pre_errors = []
|
||||
last_error_message = None
|
||||
|
||||
for result_index, (combo_slot, combo) in enumerate(combo_inputs):
|
||||
try:
|
||||
if not isinstance(combo, dict):
|
||||
raise RuntimeError(f"数据格式错误: 期望 dict，实际 {type(combo).__name__}")
|
||||
|
||||
combo_images = combo.get("images", [])
|
||||
if not combo_images:
|
||||
raise RuntimeError("无图片")
|
||||
|
||||
expanded_images = self._expand_images(combo_images)
|
||||
if not expanded_images:
|
||||
raise RuntimeError("无有效图片")
|
||||
|
||||
data_urls = []
|
||||
for image_index, image in enumerate(expanded_images, start=1):
|
||||
data_url, image_info = self._image_tensor_to_data_url(image)
|
||||
logger.info(
|
||||
"ZhiYi image-to-text combo encoded image: %s",
|
||||
{
|
||||
"combo": combo_slot,
|
||||
"image": image_index,
|
||||
**image_info,
|
||||
},
|
||||
)
|
||||
data_urls.append(data_url)
|
||||
|
||||
prompt = self._select_prompt(combo)
|
||||
messages = self._build_messages(prompt, data_urls, system_prompt)
|
||||
tasks.append(
|
||||
(
|
||||
result_index,
|
||||
combo_slot,
|
||||
self._single_request,
|
||||
(
|
||||
url,
|
||||
api_key,
|
||||
messages,
|
||||
temperature,
|
||||
max_tokens,
|
||||
retry_count,
|
||||
),
|
||||
)
|
||||
)
|
||||
except Exception as exc:
|
||||
normalized_error = normalize_error_message(exc)
|
||||
last_error_message = normalized_error
|
||||
result_items[result_index] = {"ok": False, "text": f"ERROR: {normalized_error}"}
|
||||
msg = f"[combo_{combo_slot}] 预处理失败: {normalized_error}"
|
||||
pre_errors.append(msg)
|
||||
print(f"[知衣图生文-combo] {msg}")
|
||||
traceback.print_exc()
|
||||
|
||||
run_results = {}
|
||||
run_log_lines = []
|
||||
if tasks:
|
||||
print(
|
||||
f"[知衣图生文-combo] 并发发送 {len(tasks)} 个请求（{len(combo_inputs)} 个组合），"
|
||||
f"并发上限 {max_concurrency}"
|
||||
)
|
||||
run_results, run_log_lines, run_last_error = self._run_concurrent(
|
||||
tasks,
|
||||
max_concurrency,
|
||||
label="请求",
|
||||
)
|
||||
if run_last_error:
|
||||
last_error_message = run_last_error
|
||||
|
||||
for result_index, item in run_results.items():
|
||||
result_items[result_index] = item
|
||||
|
||||
successful_count = sum(1 for item in result_items if item and item.get("ok"))
|
||||
header = (
|
||||
f"总计: {successful_count}/{len(combo_inputs)} 成功, "
|
||||
f"url={url}, 并发上限={max_concurrency}, retry_count={retry_count}"
|
||||
)
|
||||
log_lines = [header]
|
||||
if pre_errors:
|
||||
log_lines.extend(pre_errors)
|
||||
log_lines.extend(run_log_lines)
|
||||
log_text = "\n".join(log_lines)
|
||||
|
||||
if successful_count == 0:
|
||||
raise RuntimeError(
|
||||
last_error_message or normalize_error_message(f"所有请求均失败，无文本返回\n{log_text}")
|
||||
)
|
||||
|
||||
return (self._format_text_output(combo_inputs, result_items), log_text)
|
||||
@@ -11,6 +11,11 @@ from .utils.logging_utils import configure_default_logging
|
||||
configure_default_logging()
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
MAX_IMAGE_DATA_URL_BYTES = 10_000_000
|
||||
JPEG_QUALITY_STEPS = (85, 80, 75, 70, 65, 60)
|
||||
MIN_IMAGE_LONG_EDGE = 1024
|
||||
IMAGE_RESIZE_FACTOR = 0.8
|
||||
|
||||
|
||||
class ZhiYiImageTextNode:
|
||||
"""知衣图生文节点 - 传入图片，调用 Gemini 模型生成描述文本"""
|
||||
@@ -58,12 +63,66 @@ class ZhiYiImageTextNode:
|
||||
CATEGORY = "知衣/图生文"
|
||||
OUTPUT_NODE = False
|
||||
|
||||
def _image_tensor_to_base64(self, image_tensor):
|
||||
arr = (image_tensor[0].numpy() * 255).clip(0, 255).astype(np.uint8)
|
||||
pil_img = Image.fromarray(arr)
|
||||
def _image_tensor_to_pil(self, image_tensor):
|
||||
if getattr(image_tensor, "ndim", None) == 4:
|
||||
image_tensor = image_tensor[0]
|
||||
arr = (image_tensor.detach().cpu().numpy() * 255).clip(0, 255).astype(np.uint8)
|
||||
return Image.fromarray(arr).convert("RGB")
|
||||
|
||||
def _encode_jpeg_data_url(self, pil_img, quality):
|
||||
buf = io.BytesIO()
|
||||
pil_img.save(buf, format="PNG")
|
||||
return base64.b64encode(buf.getvalue()).decode("utf-8")
|
||||
pil_img.save(buf, format="JPEG", quality=quality, optimize=True)
|
||||
img_b64 = base64.b64encode(buf.getvalue()).decode("utf-8")
|
||||
return f"data:image/jpeg;base64,{img_b64}", len(buf.getvalue())
|
||||
|
||||
def _resize_to_long_edge(self, pil_img, long_edge):
|
||||
width, height = pil_img.size
|
||||
current_long_edge = max(width, height)
|
||||
if current_long_edge <= long_edge:
|
||||
return pil_img
|
||||
|
||||
scale = long_edge / current_long_edge
|
||||
new_size = (
|
||||
max(1, int(round(width * scale))),
|
||||
max(1, int(round(height * scale))),
|
||||
)
|
||||
return pil_img.resize(new_size, Image.Resampling.LANCZOS)
|
||||
|
||||
def _image_tensor_to_data_url(self, image_tensor):
|
||||
original_img = self._image_tensor_to_pil(image_tensor)
|
||||
original_size = original_img.size
|
||||
current_long_edge = max(original_size)
|
||||
best_result = None
|
||||
|
||||
while True:
|
||||
candidate_img = self._resize_to_long_edge(original_img, current_long_edge)
|
||||
for quality in JPEG_QUALITY_STEPS:
|
||||
data_url, image_bytes = self._encode_jpeg_data_url(candidate_img, quality)
|
||||
data_url_bytes = len(data_url.encode("utf-8"))
|
||||
info = {
|
||||
"original_size": original_size,
|
||||
"final_size": candidate_img.size,
|
||||
"mime_type": "image/jpeg",
|
||||
"quality": quality,
|
||||
"image_bytes": image_bytes,
|
||||
"data_url_bytes": data_url_bytes,
|
||||
"max_data_url_bytes": MAX_IMAGE_DATA_URL_BYTES,
|
||||
"resized": candidate_img.size != original_size,
|
||||
}
|
||||
best_result = (data_url, info)
|
||||
if data_url_bytes <= MAX_IMAGE_DATA_URL_BYTES:
|
||||
return best_result
|
||||
|
||||
if current_long_edge <= MIN_IMAGE_LONG_EDGE:
|
||||
break
|
||||
current_long_edge = max(MIN_IMAGE_LONG_EDGE, int(current_long_edge * IMAGE_RESIZE_FACTOR))
|
||||
|
||||
_, info = best_result
|
||||
raise RuntimeError(
|
||||
"图生文输入图片压缩后仍超过 10MB 传输限制: "
|
||||
f"{info['data_url_bytes']} bytes > {MAX_IMAGE_DATA_URL_BYTES} bytes; "
|
||||
"请降低输入图片分辨率后重试"
|
||||
)
|
||||
|
||||
def _summarize_messages_for_log(self, messages):
|
||||
summarized_messages = []
|
||||
@@ -120,8 +179,8 @@ class ZhiYiImageTextNode:
|
||||
base_url = base_url.rstrip("/")
|
||||
url = f"{base_url}/v1/chat/completions"
|
||||
|
||||
img_b64 = self._image_tensor_to_base64(image)
|
||||
data_url = f"data:image/png;base64,{img_b64}"
|
||||
data_url, image_info = self._image_tensor_to_data_url(image)
|
||||
logger.info("ZhiYi image-to-text encoded image: %s", image_info)
|
||||
|
||||
messages = []
|
||||
if system_prompt.strip():
|
||||
|
||||
@@ -22,9 +22,14 @@ class ZhiYiImageToImageComboNode:
|
||||
"""知衣图生图节点 - 接收最多8个图片组合并发调用 API"""
|
||||
|
||||
MODELS = [
|
||||
"google/gemini-3-pro-image-preview",
|
||||
"google/gemini-3.1-flash-image-preview",
|
||||
"google/gemini-2.5-flash-image-preview",
|
||||
"google/gemini-3-pro-image-preview-official",
|
||||
"gemini-3-pro-image-preview",
|
||||
"gemini-3.1-flash-image-preview",
|
||||
"gemini-2.5-flash-image-preview",
|
||||
"gemini-3-pro-image-preview-official",
|
||||
"gemini-3-pro-image-preview-aistudio",
|
||||
"gemini-3-pro-image-preview-siphonlab",
|
||||
]
|
||||
|
||||
@@ -23,10 +23,16 @@ class ZhiYiImageToImageNode:
|
||||
"""知衣图生图节点 - 最多6张输入图片，batch_size 控制重复请求次数，支持提示词列表并发"""
|
||||
|
||||
MODELS = [
|
||||
"google/gemini-3-pro-image-preview",
|
||||
"google/gemini-3.1-flash-image-preview",
|
||||
"google/gemini-2.5-flash-image-preview",
|
||||
"google/gemini-3-pro-image-preview-official",
|
||||
"gemini-3-pro-image-preview",
|
||||
"gemini-3.1-flash-image-preview",
|
||||
"gemini-2.5-flash-image-preview",
|
||||
"gemini-3-pro-image-preview-official",
|
||||
"gemini-3-pro-image-preview-aistudio",
|
||||
"gemini-3-pro-image-preview-siphonlab",
|
||||
]
|
||||
|
||||
ASPECT_RATIOS = ["", "1:1", "2:3", "3:2", "3:4", "4:3", "4:5", "5:4", "16:9", "9:16", "21:9"]
|
||||
|
||||
@@ -1,6 +1,52 @@
|
||||
import os
|
||||
import sys
|
||||
import types
|
||||
|
||||
os.environ.setdefault("FD_LITELLM_BASE_URL", "https://example.com")
|
||||
os.environ.setdefault("FD_LITELLM_API_KEY", "test-key")
|
||||
|
||||
# Add the project root directory to Python path
|
||||
# This allows the tests to import the project
|
||||
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
|
||||
|
||||
|
||||
def _install_comfy_test_stubs():
|
||||
if "comfy.comfy_types.node_typing" in sys.modules:
|
||||
return
|
||||
|
||||
comfy_module = sys.modules.get("comfy") or types.ModuleType("comfy")
|
||||
comfy_types_module = types.ModuleType("comfy.comfy_types")
|
||||
node_typing_module = types.ModuleType("comfy.comfy_types.node_typing")
|
||||
utils_module = types.ModuleType("comfy.utils")
|
||||
|
||||
class IO:
|
||||
STRING = "STRING"
|
||||
COMBO = "COMBO"
|
||||
INT = "INT"
|
||||
FLOAT = "FLOAT"
|
||||
BOOLEAN = "BOOLEAN"
|
||||
IMAGE = "IMAGE"
|
||||
|
||||
class ComfyNodeABC:
|
||||
pass
|
||||
|
||||
def common_upscale(samples, width, height, _upscale_method, _crop):
|
||||
import torch.nn.functional as F
|
||||
|
||||
return F.interpolate(samples, size=(height, width), mode="bilinear", align_corners=False)
|
||||
|
||||
node_typing_module.IO = IO
|
||||
node_typing_module.ComfyNodeABC = ComfyNodeABC
|
||||
node_typing_module.InputTypeDict = dict
|
||||
utils_module.common_upscale = common_upscale
|
||||
comfy_module.comfy_types = comfy_types_module
|
||||
comfy_module.utils = utils_module
|
||||
comfy_types_module.node_typing = node_typing_module
|
||||
|
||||
sys.modules["comfy"] = comfy_module
|
||||
sys.modules["comfy.comfy_types"] = comfy_types_module
|
||||
sys.modules["comfy.comfy_types.node_typing"] = node_typing_module
|
||||
sys.modules["comfy.utils"] = utils_module
|
||||
|
||||
|
||||
_install_comfy_test_stubs()
|
||||
|
||||
+397
-24
@@ -17,8 +17,10 @@ from src.Comfyui_Fd_Nodes.nodes import (
|
||||
_resolution_to_edit_size,
|
||||
)
|
||||
from src.Comfyui_Fd_Nodes.prompt_nodes import EcommercePromptGenerator, PromptListSelector
|
||||
from src.Comfyui_Fd_Nodes import zhiyi_image_text_combo_node as zhiyi_image_text_combo_module
|
||||
from src.Comfyui_Fd_Nodes import zhiyi_image_text_node as zhiyi_image_text_module
|
||||
from src.Comfyui_Fd_Nodes import zhiyi_text_node as zhiyi_text_module
|
||||
from src.Comfyui_Fd_Nodes.zhiyi_image_text_combo_node import ZhiYiImageTextComboNode
|
||||
from src.Comfyui_Fd_Nodes.zhiyi_image_text_node import ZhiYiImageTextNode
|
||||
from src.Comfyui_Fd_Nodes.zhiyi_image_to_image_combo_node import ZhiYiImageToImageComboNode
|
||||
from src.Comfyui_Fd_Nodes.zhiyi_image_to_image_node import ZhiYiImageToImageNode
|
||||
@@ -97,6 +99,7 @@ def test_new_nodes_hide_base_url_and_api_key_inputs():
|
||||
assert {"base_url", "api_key"}.isdisjoint(GPTImageEditNode.INPUT_TYPES()["required"])
|
||||
assert {"api_url", "api_key"}.isdisjoint(EcommercePromptGenerator.INPUT_TYPES()["required"])
|
||||
assert {"base_url", "api_key"}.isdisjoint(ZhiYiImageTextNode.INPUT_TYPES()["required"])
|
||||
assert {"base_url", "api_key"}.isdisjoint(ZhiYiImageTextComboNode.INPUT_TYPES()["required"])
|
||||
assert {"base_url", "api_key"}.isdisjoint(ZhiYiImageToImageNode.INPUT_TYPES()["required"])
|
||||
assert {"base_url", "api_key"}.isdisjoint(ZhiYiTextGenNode.INPUT_TYPES()["required"])
|
||||
|
||||
@@ -114,6 +117,22 @@ def test_zhiyi_image_to_image_exposes_out_request_id():
|
||||
assert optional_inputs["out_request_id"][1]["default"] == "default"
|
||||
|
||||
|
||||
def test_zhiyi_image_to_image_nodes_expose_legacy_and_channel_models():
|
||||
expected_models = {
|
||||
"google/gemini-2.5-flash-image-preview",
|
||||
"google/gemini-3-pro-image-preview",
|
||||
"google/gemini-3-pro-image-preview-official",
|
||||
"google/gemini-3.1-flash-image-preview",
|
||||
"gemini-3-pro-image-preview-aistudio",
|
||||
"gemini-3-pro-image-preview-siphonlab",
|
||||
}
|
||||
|
||||
assert expected_models.issubset(set(ZhiYiImageToImageNode.MODELS))
|
||||
assert expected_models.issubset(set(ZhiYiImageToImageComboNode.MODELS))
|
||||
assert ZhiYiImageToImageNode.INPUT_TYPES()["required"]["model"][1]["default"] == "google/gemini-3-pro-image-preview"
|
||||
assert ZhiYiImageToImageComboNode.INPUT_TYPES()["required"]["model"][1]["default"] == "google/gemini-3-pro-image-preview"
|
||||
|
||||
|
||||
def test_gemini_service_builds_internal_request_body():
|
||||
client = GeminiImageServiceClient(service_url="https://gemini.internal", oss_url_prefix="https://oss/")
|
||||
|
||||
@@ -136,11 +155,21 @@ def test_gemini_service_builds_internal_request_body():
|
||||
}
|
||||
assert client.summarize_request_body(body)["image_count"] == 1
|
||||
|
||||
aistudio_body = client.build_request_body(
|
||||
prompt="draw product",
|
||||
model="gemini-3-pro-image-preview-aistudio",
|
||||
image_url_list=["https://oss/input.png"],
|
||||
)
|
||||
assert aistudio_body["model"] == "google/gemini-3-pro-image-preview-official"
|
||||
|
||||
|
||||
def test_gemini_service_model_and_prompt_helpers():
|
||||
assert normalize_gemini_model_name("gemini-2.5-flash-image-preview") == "google/gemini-2.5-flash-image-preview"
|
||||
assert normalize_gemini_model_name("google/gemini-3-pro-image-preview") == "google/gemini-3-pro-image-preview"
|
||||
assert should_use_litellm_gemini("gemini-3-pro-image-preview-aistudio") is True
|
||||
assert normalize_gemini_model_name("gemini-3-pro-image-preview-official") == "google/gemini-3-pro-image-preview-official"
|
||||
assert normalize_gemini_model_name("google/gemini-3-pro-image-preview-official") == "google/gemini-3-pro-image-preview-official"
|
||||
assert normalize_gemini_model_name("gemini-3-pro-image-preview-aistudio") == "google/gemini-3-pro-image-preview-official"
|
||||
assert should_use_litellm_gemini("gemini-3-pro-image-preview-aistudio") is False
|
||||
assert should_use_litellm_gemini("gemini-3-pro-image-preview-siphonlab") is True
|
||||
assert should_use_litellm_gemini("gemini-3-pro-image-preview") is False
|
||||
assert compose_prompt("user prompt", "system prompt") == "system prompt\n\nuser prompt"
|
||||
@@ -187,25 +216,21 @@ def test_zhiyi_image_to_image_uses_internal_gemini_service(monkeypatch):
|
||||
assert request["out_request_id"] == "req-123"
|
||||
|
||||
|
||||
def test_zhiyi_image_to_image_aistudio_uses_litellm_without_model_rename(monkeypatch):
|
||||
def test_zhiyi_image_to_image_aistudio_uses_internal_official_model(monkeypatch):
|
||||
node = ZhiYiImageToImageNode()
|
||||
calls = []
|
||||
|
||||
monkeypatch.setattr("src.Comfyui_Fd_Nodes.zhiyi_image_to_image_node.tensor_to_base64", lambda _tensor: "AAA")
|
||||
class FakeClient:
|
||||
def upload_images(self, image_tensors):
|
||||
calls.append(("upload_images", len(image_tensors)))
|
||||
return ["https://oss/input.png"]
|
||||
|
||||
def fake_litellm(messages, model, aspect_ratio, image_size, seed, out_request_id):
|
||||
calls.append({
|
||||
"messages": messages,
|
||||
"model": model,
|
||||
"aspect_ratio": aspect_ratio,
|
||||
"image_size": image_size,
|
||||
"seed": seed,
|
||||
"out_request_id": out_request_id,
|
||||
})
|
||||
return torch.ones((1, 2, 2, 3), dtype=torch.float32)
|
||||
def call_with_image_urls(self, **kwargs):
|
||||
calls.append(("call", kwargs))
|
||||
return torch.ones((1, 2, 2, 3), dtype=torch.float32), "https://oss/result.png", "ok"
|
||||
|
||||
monkeypatch.setattr(node, "_single_litellm_request", fake_litellm)
|
||||
monkeypatch.setattr(node.gemini_client, "upload_images", lambda _image_tensors: pytest.fail("aistudio should not upload images to Gemini service"))
|
||||
node.gemini_client = FakeClient()
|
||||
monkeypatch.setattr(node, "_single_litellm_request", lambda *args: pytest.fail("aistudio should use internal Gemini service"))
|
||||
|
||||
result, actual_seed = node.generate(
|
||||
image_1=torch.zeros((1, 2, 2, 3), dtype=torch.float32),
|
||||
@@ -222,10 +247,12 @@ def test_zhiyi_image_to_image_aistudio_uses_litellm_without_model_rename(monkeyp
|
||||
|
||||
assert result.shape == (1, 2, 2, 3)
|
||||
assert actual_seed == 42
|
||||
assert calls[0]["model"] == "gemini-3-pro-image-preview-aistudio"
|
||||
assert calls[0]["messages"][0]["role"] == "system"
|
||||
assert calls[0]["messages"][1]["content"][0]["text"] == "user prompt"
|
||||
assert calls[0]["messages"][1]["content"][1]["image_url"]["url"] == "data:image/png;base64,AAA"
|
||||
assert calls[0] == ("upload_images", 1)
|
||||
request = calls[1][1]
|
||||
assert request["image_url_list"] == ["https://oss/input.png"]
|
||||
assert request["prompt"] == "system prompt\n\nuser prompt"
|
||||
assert request["model"] == "gemini-3-pro-image-preview-aistudio"
|
||||
assert normalize_gemini_model_name(request["model"]) == "google/gemini-3-pro-image-preview-official"
|
||||
|
||||
|
||||
def test_normalize_error_message_classifies_timeout_and_nsfw():
|
||||
@@ -345,6 +372,55 @@ def test_zhiyi_image_to_image_combo_uses_internal_gemini_service(monkeypatch):
|
||||
assert calls[1][1] == "system prompt\n\nprompt two"
|
||||
|
||||
|
||||
def test_zhiyi_image_to_image_combo_aistudio_uses_internal_official_model(monkeypatch):
|
||||
node = ZhiYiImageToImageComboNode()
|
||||
node.gemini_client.service_url = "https://gemini.internal/generate"
|
||||
combo = {
|
||||
"images": [torch.zeros((1, 2, 2, 3), dtype=torch.float32)],
|
||||
"prompts": ["prompt one"],
|
||||
}
|
||||
calls = []
|
||||
|
||||
monkeypatch.setattr(node.gemini_client, "upload_images", lambda image_tensors: calls.append(("upload_images", len(image_tensors))) or ["https://oss/input.png"])
|
||||
monkeypatch.setattr(node, "_single_litellm_request", lambda *args: pytest.fail("aistudio should use internal Gemini service"))
|
||||
|
||||
def fake_single_request(*args):
|
||||
calls.append(("single_request", args))
|
||||
return torch.ones((1, 2, 2, 3), dtype=torch.float32)
|
||||
|
||||
monkeypatch.setattr(node, "_single_request", fake_single_request)
|
||||
|
||||
images, actual_seed, log_text = node.generate(
|
||||
model="gemini-3-pro-image-preview-aistudio",
|
||||
aspect_ratio="1:1",
|
||||
image_size="4K",
|
||||
batch_size=1,
|
||||
max_concurrency=1,
|
||||
seed_mode="固定种子",
|
||||
seed=7,
|
||||
out_request_id="req-456",
|
||||
combo_1=combo,
|
||||
system_prompt="system prompt",
|
||||
)
|
||||
|
||||
assert len(images) == 1
|
||||
assert actual_seed == 7
|
||||
assert "url=https://gemini.internal/generate" in log_text
|
||||
assert calls[0] == ("upload_images", 1)
|
||||
assert calls[1][0] == "single_request"
|
||||
request_args = calls[1][1]
|
||||
assert request_args == (
|
||||
["https://oss/input.png"],
|
||||
"system prompt\n\nprompt one",
|
||||
"gemini-3-pro-image-preview-aistudio",
|
||||
"1:1",
|
||||
"4K",
|
||||
7,
|
||||
"req-456",
|
||||
)
|
||||
assert normalize_gemini_model_name(request_args[2]) == "google/gemini-3-pro-image-preview-official"
|
||||
|
||||
|
||||
def test_zhiyi_image_to_image_combo_siphonlab_uses_litellm_without_model_rename(monkeypatch):
|
||||
node = ZhiYiImageToImageComboNode()
|
||||
combo = {
|
||||
@@ -403,15 +479,27 @@ def test_zhiyi_image_text_request_logs_request_and_response_without_images(monke
|
||||
captured_logs.append(("post_payload", json.loads(data)))
|
||||
return DummyResponse()
|
||||
|
||||
def fake_image_tensor_to_base64(_image):
|
||||
return "AAA"
|
||||
def fake_image_tensor_to_data_url(_image):
|
||||
return (
|
||||
"data:image/jpeg;base64,AAA",
|
||||
{
|
||||
"original_size": (1, 1),
|
||||
"final_size": (1, 1),
|
||||
"mime_type": "image/jpeg",
|
||||
"quality": 85,
|
||||
"image_bytes": 3,
|
||||
"data_url_bytes": 26,
|
||||
"max_data_url_bytes": zhiyi_image_text_module.MAX_IMAGE_DATA_URL_BYTES,
|
||||
"resized": False,
|
||||
},
|
||||
)
|
||||
|
||||
def fake_logger_info(message, *args):
|
||||
captured_logs.append((message, args))
|
||||
|
||||
monkeypatch.setattr(zhiyi_image_text_module, "load_config", lambda: {"base_url": "https://example.com", "api_key": "secret"})
|
||||
monkeypatch.setattr(zhiyi_image_text_module.requests, "post", fake_post)
|
||||
monkeypatch.setattr(node, "_image_tensor_to_base64", fake_image_tensor_to_base64)
|
||||
monkeypatch.setattr(node, "_image_tensor_to_data_url", fake_image_tensor_to_data_url)
|
||||
monkeypatch.setattr(zhiyi_image_text_module.logger, "info", fake_logger_info)
|
||||
|
||||
result = node.generate(
|
||||
@@ -425,18 +513,283 @@ def test_zhiyi_image_text_request_logs_request_and_response_without_images(monke
|
||||
|
||||
assert result == ("hello",)
|
||||
request_payload = next(value for key, value in captured_logs if key == "post_payload")
|
||||
assert request_payload["model"] == "gemini-3-pro-preview"
|
||||
assert request_payload["model"] == "doubao-seed-2.0-mini"
|
||||
assert request_payload["messages"][1]["content"][1]["image_url"]["url"] == "data:image/jpeg;base64,AAA"
|
||||
request_log = next(args[0] for message, args in captured_logs if message == "Calling ZhiYi image-to-text API with payload=%s")
|
||||
assert request_log["messages"] == [
|
||||
{"role": "system", "text": "system prompt"},
|
||||
{"role": "user", "text": "describe this image", "image_count": 1},
|
||||
]
|
||||
assert "data:image/png;base64,AAA" not in json.dumps(request_log, ensure_ascii=False)
|
||||
encoded_image_log = next(args[0] for message, args in captured_logs if message == "ZhiYi image-to-text encoded image: %s")
|
||||
assert encoded_image_log["mime_type"] == "image/jpeg"
|
||||
assert "data:image/jpeg;base64,AAA" not in json.dumps(request_log, ensure_ascii=False)
|
||||
response_log = next(args[0] for message, args in captured_logs if message == "ZhiYi image-to-text API response summary: %s")
|
||||
assert response_log["status_code"] == 200
|
||||
assert response_log["choice_count"] == 1
|
||||
|
||||
|
||||
def test_zhiyi_image_text_encodes_small_image_as_jpeg_data_url():
|
||||
"""Small images should be encoded as JPEG data URLs without resizing."""
|
||||
node = ZhiYiImageTextNode()
|
||||
image = torch.zeros((1, 16, 16, 3), dtype=torch.float32)
|
||||
|
||||
data_url, info = node._image_tensor_to_data_url(image)
|
||||
|
||||
assert data_url.startswith("data:image/jpeg;base64,")
|
||||
assert len(data_url.encode("utf-8")) <= zhiyi_image_text_module.MAX_IMAGE_DATA_URL_BYTES
|
||||
assert info["original_size"] == (16, 16)
|
||||
assert info["final_size"] == (16, 16)
|
||||
assert info["quality"] == 85
|
||||
assert info["resized"] is False
|
||||
|
||||
|
||||
def test_zhiyi_image_text_recompresses_until_data_url_is_under_limit(monkeypatch):
|
||||
"""Large/noisy inputs should be retried with lower quality and smaller dimensions until under the transport limit."""
|
||||
node = ZhiYiImageTextNode()
|
||||
image = torch.rand((1, 512, 512, 3), dtype=torch.float32)
|
||||
|
||||
monkeypatch.setattr(zhiyi_image_text_module, "MAX_IMAGE_DATA_URL_BYTES", 50_000)
|
||||
monkeypatch.setattr(zhiyi_image_text_module, "MIN_IMAGE_LONG_EDGE", 128)
|
||||
monkeypatch.setattr(zhiyi_image_text_module, "IMAGE_RESIZE_FACTOR", 0.5)
|
||||
|
||||
data_url, info = node._image_tensor_to_data_url(image)
|
||||
|
||||
assert data_url.startswith("data:image/jpeg;base64,")
|
||||
assert len(data_url.encode("utf-8")) <= 50_000
|
||||
assert info["data_url_bytes"] <= 50_000
|
||||
assert info["final_size"][0] <= 512
|
||||
assert info["final_size"][1] <= 512
|
||||
assert info["resized"] is True
|
||||
|
||||
|
||||
def test_zhiyi_image_text_raises_clear_error_when_image_cannot_fit(monkeypatch):
|
||||
"""If every compression attempt remains too large, the node should fail before making the API request."""
|
||||
node = ZhiYiImageTextNode()
|
||||
image = torch.zeros((1, 4, 4, 3), dtype=torch.float32)
|
||||
|
||||
monkeypatch.setattr(zhiyi_image_text_module, "MAX_IMAGE_DATA_URL_BYTES", 1)
|
||||
monkeypatch.setattr(zhiyi_image_text_module, "MIN_IMAGE_LONG_EDGE", 1)
|
||||
|
||||
with pytest.raises(RuntimeError, match="10MB 传输限制"):
|
||||
node._image_tensor_to_data_url(image)
|
||||
|
||||
|
||||
def test_zhiyi_image_text_combo_registered_and_accepts_combo_inputs():
|
||||
input_types = ZhiYiImageTextComboNode.INPUT_TYPES()
|
||||
|
||||
assert "ZhiYiImageTextComboNode" in NODE_CLASS_MAPPINGS
|
||||
assert NODE_CLASS_MAPPINGS["ZhiYiImageTextComboNode"] is ZhiYiImageTextComboNode
|
||||
assert NODE_DISPLAY_NAME_MAPPINGS["ZhiYiImageTextComboNode"] == "知衣-图生文-combo"
|
||||
assert set(input_types["required"]) == {"max_concurrency", "temperature", "max_tokens", "retry_count"}
|
||||
assert "combo_1" in input_types["optional"]
|
||||
assert "combo_8" in input_types["optional"]
|
||||
assert ZhiYiImageTextComboNode.RETURN_TYPES == ("STRING", "STRING")
|
||||
assert ZhiYiImageTextComboNode.RETURN_NAMES == ("text", "log")
|
||||
|
||||
|
||||
def test_zhiyi_image_text_combo_sends_multiple_images_in_one_request(monkeypatch):
|
||||
node = ZhiYiImageTextComboNode()
|
||||
captured_payloads = []
|
||||
captured_logs = []
|
||||
image = torch.zeros((2, 2, 2, 3), dtype=torch.float32)
|
||||
combo = {"images": [image], "prompts": ["describe both images"]}
|
||||
|
||||
class DummyResponse:
|
||||
status_code = 200
|
||||
text = '{"choices":[{"message":{"content":"both images"}}]}'
|
||||
|
||||
def raise_for_status(self):
|
||||
return None
|
||||
|
||||
def json(self):
|
||||
return {"choices": [{"message": {"content": "both images"}}]}
|
||||
|
||||
def fake_post(url, headers, data, timeout):
|
||||
captured_payloads.append(json.loads(data))
|
||||
return DummyResponse()
|
||||
|
||||
def fake_image_tensor_to_data_url(image_tensor):
|
||||
index = len(captured_logs) + 1
|
||||
return (
|
||||
f"data:image/jpeg;base64,IMG{index}",
|
||||
{
|
||||
"original_size": (2, 2),
|
||||
"final_size": (2, 2),
|
||||
"mime_type": "image/jpeg",
|
||||
"quality": 85,
|
||||
"image_bytes": index,
|
||||
"data_url_bytes": 30,
|
||||
"max_data_url_bytes": zhiyi_image_text_module.MAX_IMAGE_DATA_URL_BYTES,
|
||||
"resized": False,
|
||||
},
|
||||
)
|
||||
|
||||
def fake_logger_info(message, *args):
|
||||
captured_logs.append((message, args))
|
||||
|
||||
monkeypatch.setattr(zhiyi_image_text_combo_module, "load_config", lambda: {"base_url": "https://example.com", "api_key": "secret"})
|
||||
monkeypatch.setattr(zhiyi_image_text_combo_module.requests, "post", fake_post)
|
||||
monkeypatch.setattr(node, "_image_tensor_to_data_url", fake_image_tensor_to_data_url)
|
||||
monkeypatch.setattr(zhiyi_image_text_combo_module.logger, "info", fake_logger_info)
|
||||
|
||||
text, log_text = node.generate(
|
||||
max_concurrency=1,
|
||||
temperature=0.2,
|
||||
max_tokens=256,
|
||||
retry_count=0,
|
||||
combo_1=combo,
|
||||
system_prompt="system prompt",
|
||||
)
|
||||
|
||||
assert text == "both images"
|
||||
assert "总计: 1/1 成功" in log_text
|
||||
payload = captured_payloads[0]
|
||||
assert payload["model"] == "doubao-seed-2.0-mini"
|
||||
assert payload["messages"][0]["role"] == "system"
|
||||
content = payload["messages"][1]["content"]
|
||||
assert content[0] == {"type": "text", "text": "describe both images"}
|
||||
assert [part["image_url"]["url"] for part in content[1:]] == [
|
||||
"data:image/jpeg;base64,IMG1",
|
||||
"data:image/jpeg;base64,IMG2",
|
||||
]
|
||||
request_log = next(args[0] for message, args in captured_logs if message == "Calling ZhiYi image-to-text combo API with payload=%s")
|
||||
assert request_log["messages"] == [
|
||||
{"role": "system", "text": "system prompt"},
|
||||
{"role": "user", "text": "describe both images", "image_count": 2},
|
||||
]
|
||||
assert "data:image/jpeg;base64,IMG1" not in json.dumps(request_log, ensure_ascii=False)
|
||||
|
||||
|
||||
def test_zhiyi_image_text_combo_keeps_output_order_when_requests_finish_out_of_order(monkeypatch):
|
||||
node = ZhiYiImageTextComboNode()
|
||||
image = torch.zeros((1, 2, 2, 3), dtype=torch.float32)
|
||||
|
||||
monkeypatch.setattr(zhiyi_image_text_combo_module, "load_config", lambda: {"base_url": "https://example.com", "api_key": "secret"})
|
||||
monkeypatch.setattr(node, "_image_tensor_to_data_url", lambda _image: ("data:image/jpeg;base64,AAA", {}))
|
||||
|
||||
def fake_single_request(url, api_key, messages, temperature, max_tokens, retry_count):
|
||||
prompt = messages[-1]["content"][0]["text"]
|
||||
if prompt == "first":
|
||||
return "text first", 0
|
||||
return "text second", 0
|
||||
|
||||
def fake_run_concurrent(tasks, max_workers, label="请求"):
|
||||
second = tasks[1]
|
||||
first = tasks[0]
|
||||
results = {
|
||||
second[0]: {"ok": True, "text": second[2](*second[3])[0]},
|
||||
first[0]: {"ok": True, "text": first[2](*first[3])[0]},
|
||||
}
|
||||
return results, ["[请求 combo_2] 成功", "[请求 combo_1] 成功"], None
|
||||
|
||||
monkeypatch.setattr(node, "_single_request", fake_single_request)
|
||||
monkeypatch.setattr(node, "_run_concurrent", fake_run_concurrent)
|
||||
|
||||
text, log_text = node.generate(
|
||||
max_concurrency=2,
|
||||
retry_count=0,
|
||||
combo_1={"images": [image], "prompts": ["first"]},
|
||||
combo_2={"images": [image], "prompts": ["second"]},
|
||||
)
|
||||
|
||||
assert text == "[combo_1]\ntext first\n\n[combo_2]\ntext second"
|
||||
assert "总计: 2/2 成功" in log_text
|
||||
|
||||
|
||||
def test_zhiyi_image_text_combo_retries_once_after_retryable_failure(monkeypatch):
|
||||
node = ZhiYiImageTextComboNode()
|
||||
calls = []
|
||||
|
||||
class DummyResponse:
|
||||
def __init__(self, status_code, text, payload=None):
|
||||
self.status_code = status_code
|
||||
self.text = text
|
||||
self.payload = payload or {"choices": [{"message": {"content": "retry success"}}]}
|
||||
|
||||
def raise_for_status(self):
|
||||
return None
|
||||
|
||||
def json(self):
|
||||
return self.payload
|
||||
|
||||
def fake_post(url, headers, data, timeout):
|
||||
calls.append(json.loads(data))
|
||||
if len(calls) == 1:
|
||||
return DummyResponse(500, "internal error")
|
||||
return DummyResponse(200, '{"choices":[{"message":{"content":"retry success"}}]}')
|
||||
|
||||
monkeypatch.setattr(zhiyi_image_text_combo_module.requests, "post", fake_post)
|
||||
monkeypatch.setattr(node, "_sleep_before_retry", lambda _attempt_index: None)
|
||||
|
||||
text, retry_used = node._single_request(
|
||||
url="https://example.com/v1/chat/completions",
|
||||
api_key="secret",
|
||||
messages=[{"role": "user", "content": [{"type": "text", "text": "prompt"}]}],
|
||||
temperature=0.7,
|
||||
max_tokens=128,
|
||||
retry_count=1,
|
||||
)
|
||||
|
||||
assert text == "retry success"
|
||||
assert retry_used == 1
|
||||
assert len(calls) == 2
|
||||
|
||||
|
||||
def test_zhiyi_image_text_combo_keeps_failure_placeholder_for_partial_failure(monkeypatch):
|
||||
node = ZhiYiImageTextComboNode()
|
||||
image = torch.zeros((1, 2, 2, 3), dtype=torch.float32)
|
||||
|
||||
monkeypatch.setattr(zhiyi_image_text_combo_module, "load_config", lambda: {"base_url": "https://example.com", "api_key": "secret"})
|
||||
monkeypatch.setattr(node, "_image_tensor_to_data_url", lambda _image: ("data:image/jpeg;base64,AAA", {}))
|
||||
|
||||
def fake_run_concurrent(tasks, max_workers, label="请求"):
|
||||
return (
|
||||
{
|
||||
0: {"ok": True, "text": "ok text"},
|
||||
1: {"ok": False, "text": "ERROR: TIMEOUT: request timed out"},
|
||||
},
|
||||
["[请求 combo_1] 成功", "[请求 combo_2] 失败: TIMEOUT: request timed out"],
|
||||
"TIMEOUT: request timed out",
|
||||
)
|
||||
|
||||
monkeypatch.setattr(node, "_run_concurrent", fake_run_concurrent)
|
||||
|
||||
text, log_text = node.generate(
|
||||
max_concurrency=2,
|
||||
retry_count=0,
|
||||
combo_1={"images": [image], "prompts": ["first"]},
|
||||
combo_2={"images": [image], "prompts": ["second"]},
|
||||
)
|
||||
|
||||
assert text == "[combo_1]\nok text\n\n[combo_2]\nERROR: TIMEOUT: request timed out"
|
||||
assert "总计: 1/2 成功" in log_text
|
||||
assert "[请求 combo_2] 失败: TIMEOUT: request timed out" in log_text
|
||||
|
||||
|
||||
def test_zhiyi_image_text_combo_raises_when_all_requests_fail(monkeypatch):
|
||||
node = ZhiYiImageTextComboNode()
|
||||
image = torch.zeros((1, 2, 2, 3), dtype=torch.float32)
|
||||
|
||||
monkeypatch.setattr(zhiyi_image_text_combo_module, "load_config", lambda: {"base_url": "https://example.com", "api_key": "secret"})
|
||||
monkeypatch.setattr(node, "_image_tensor_to_data_url", lambda _image: ("data:image/jpeg;base64,AAA", {}))
|
||||
monkeypatch.setattr(
|
||||
node,
|
||||
"_run_concurrent",
|
||||
lambda tasks, max_workers, label="请求": (
|
||||
{0: {"ok": False, "text": "ERROR: UNKNOWN: API 请求失败: 500"}},
|
||||
["[请求 combo_1] 失败: UNKNOWN: API 请求失败: 500"],
|
||||
"UNKNOWN: API 请求失败: 500",
|
||||
),
|
||||
)
|
||||
|
||||
with pytest.raises(RuntimeError, match=r"^UNKNOWN: API 请求失败: 500$"):
|
||||
node.generate(
|
||||
max_concurrency=1,
|
||||
retry_count=0,
|
||||
combo_1={"images": [image], "prompts": ["first"]},
|
||||
)
|
||||
|
||||
|
||||
def test_zhiyi_text_gen_request_logs_request_and_response(monkeypatch):
|
||||
"""The text API should log request/response summaries."""
|
||||
node = ZhiYiTextGenNode()
|
||||
@@ -491,6 +844,16 @@ def test_configure_default_logging_installs_timestamp_format(monkeypatch):
|
||||
class FakeRootLogger:
|
||||
handlers = []
|
||||
|
||||
def addHandler(self, handler):
|
||||
self.handlers.append(handler)
|
||||
|
||||
def removeHandler(self, handler):
|
||||
if handler in self.handlers:
|
||||
self.handlers.remove(handler)
|
||||
|
||||
def setLevel(self, _level):
|
||||
return None
|
||||
|
||||
def fake_get_logger(name=None):
|
||||
if name is None:
|
||||
return FakeRootLogger()
|
||||
@@ -523,6 +886,16 @@ def test_configure_default_logging_upgrades_existing_handler_without_timestamp(m
|
||||
def __init__(self):
|
||||
self.handlers = [FakeHandler()]
|
||||
|
||||
def addHandler(self, handler):
|
||||
self.handlers.append(handler)
|
||||
|
||||
def removeHandler(self, handler):
|
||||
if handler in self.handlers:
|
||||
self.handlers.remove(handler)
|
||||
|
||||
def setLevel(self, _level):
|
||||
return None
|
||||
|
||||
fake_root_logger = FakeRootLogger()
|
||||
|
||||
def fake_get_logger(name=None):
|
||||
|
||||
Reference in New Issue
Block a user