diff --git a/image2prompt_node.py b/image2prompt_node.py index f75eeaa..0f8402b 100644 --- a/image2prompt_node.py +++ b/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", }