add text2prompt

image2prompt_node.py

@@ -201,7 +201,133 @@ class Image2Prompt:
         """
         url = f"{APIUrl}/generate"
         headers = {"authorization": APIKey}
-        data = {"base64": imgBase64}
+        data = { "type": "img2img", "base64": imgBase64 }
+
+        max_retries = max(0, retry_count)
+
+        for attempt in range(max_retries + 1):
+            try:
+                response = self.session.post(url, json=data, headers=headers, timeout=30)
+                response.raise_for_status()
+
+                result = response.json()
+
+                if result.get("code") == 200:
+                    if neutral:
+                        return result.get("neutralText")
+                    else:
+                        return result.get("promptText")
+                else:
+                    # API 返回错误码，记录日志
+                    print(f"API request failed with code: {result.get('code')}, message: {result.get('message', 'unknown error')}")
+
+            except requests.exceptions.RequestException as e:
+                print(f"Attempt {attempt + 1} failed: {str(e)}")
+
+            except Exception as e:
+                print(f"Unexpected error on attempt {attempt + 1}: {str(e)}")
+
+            # 如果不是最后一次尝试，等待一下再重试
+            if attempt < max_retries:
+                import time
+                time.sleep(1)
+
+        # 所有重试都失败了
+        print(f"Failed after {max_retries} retries")
+        return None
+
+
+class Text2Prompt:
+    """
+    ComfyUI 自定义节点：将文本描述转换为图片提示
+    """
+
+    @classmethod
+    def INPUT_TYPES(cls) -> Dict[str, Any]:
+        return {
+            "required": {
+                "imgDesc": ("STRING", {"default": "", "multiline": True}),
+                "APIUrl": ("STRING", {"default": "", "multiline": False}),
+                "APIKey": ("STRING", {"default": "", "multiline": False}),
+                "RetryCount": ("INT", {"default": 3, "min": 0, "max": 100}),
+                "Neutral": ("BOOLEAN", {"default": True}),
+            }
+        }
+
+    RETURN_TYPES = ("STRING",)
+    RETURN_NAMES = ("text",)
+    FUNCTION = "execute"
+    CATEGORY = "NudeJS"
+
+    def __init__(self):
+        self.session = requests.Session()
+        # 配置重试策略
+        retry_strategy = Retry(
+            total=3,
+            backoff_factor=1,
+            status_forcelist=[429, 500, 502, 503, 504],
+        )
+        adapter = HTTPAdapter(max_retries=retry_strategy)
+        self.session.mount("http://", adapter)
+        self.session.mount("https://", adapter)
+
+    def execute(
+        self,
+        imgDesc: str,
+        APIUrl: str,
+        APIKey: str,
+        RetryCount: int,
+        Neutral: bool,
+    ) -> Tuple[Union[str, None]]:
+        """
+        执行文本转提示词的主逻辑
+
+        Args:
+            imgDesc: 图片的文本描述
+            APIUrl: API 接口地址
+            APIKey: API 密钥
+            RetryCount: 重试次数
+            Neutral: 是否使用中性文本
+
+        Returns:
+            转换后的文本提示
+        """
+
+        # 发送请求获取提示词
+        prompt_text = self._call_api(
+            APIUrl=APIUrl,
+            APIKey=APIKey,
+            imgDesc=imgDesc,
+            retry_count=RetryCount,
+            neutral=Neutral,
+        )
+
+        return (prompt_text,)
+
+    def _call_api(
+        self,
+        APIUrl: str,
+        APIKey: str,
+        imgDesc: str,
+        retry_count: int,
+        neutral: bool,
+    ) -> Union[str, None]:
+        """
+        调用 API 接口获取提示词
+
+        Args:
+            APIUrl: API 地址
+            APIKey: API 密钥
+            imgDesc: 图片的文本描述
+            retry_count: 最大重试次数
+            neutral: 是否使用中性文本
+
+        Returns:
+            提示词文本，失败返回 None
+        """
+        url = f"{APIUrl}/generate"
+        headers = {"authorization": APIKey}
+        data = { "type": "text2img", "desc": imgDesc }
 
         max_retries = max(0, retry_count)
 
@@ -240,8 +366,10 @@ class Image2Prompt:
 # 注册节点
 NODE_CLASS_MAPPINGS = {
     "Image2Prompt": Image2Prompt,
+    "Text2Prompt": Text2Prompt,
 }
 
 NODE_DISPLAY_NAME_MAPPINGS = {
     "Image2Prompt": "Image2Prompt",
+    "Text2Prompt": "Text2Prompt",
 }