# LLM API 通过cloudflare gateway调用llm import base64 import io 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): if key in data: return data[key] # 递归检查所有其他键的值 for value in data.values(): result = find_value_recursive(key, value) if result is not None: return result elif isinstance(data, list): for item in data: result = find_value_recursive(key, item) if result is not None: return result def image_tensor_to_base64(image): pil_image = Image.fromarray(np.clip(255. * image.cpu().numpy().squeeze(), 0, 255).astype(np.uint8)) # 创建一个BytesIO对象，用于临时存储图像数据 image_data = io.BytesIO() # 将图像保存到BytesIO对象中，格式为PNG pil_image.save(image_data, format='PNG') # 将BytesIO对象的内容转换为字节串 image_data_bytes = image_data.getvalue() # 将图像数据编码为Base64字符串 encoded_image = "data:image/png;base64," + base64.b64encode(image_data_bytes).decode('utf-8') return encoded_image class LLMChat: """llm chat""" @classmethod def INPUT_TYPES(s): return { "required": { "llm_provider": (["claude-3-5-sonnet-20241022-v2", "claude-3-5-sonnet-20241022-v3", "claude-3-7-sonnet-20250219-v1", "claude-4-sonnet-20250514-v1", "gpt-4o-1120", "gpt-4.1", "deepseek-v3", "deepseek-r1"],), "prompt": ("STRING", {"multiline": True}), "temperature": ("FLOAT", {"default": 0.7, "min": 0.0, "max": 1.0}), "max_tokens": ("INT", {"default": 4096, "min": 1, "max": 65535}), "timeout": ("INT", {"default": 120, "min": 30, "max": 900}), } } RETURN_TYPES = ("STRING",) RETURN_NAMES = ("llm输出",) FUNCTION = "chat" CATEGORY = "不忘科技-自定义节点🚩/LLM" def chat(self, llm_provider: str, prompt: str, temperature: float, max_tokens: int, timeout: int): @retry(Exception, tries=3, delay=1) def _chat(): try: with httpx.Client(timeout=httpx.Timeout(timeout, connect=15)) as session: resp = session.post("https://gateway.bowong.cc/chat/completions", headers={ "Content-Type": "application/json", "Accept": "application/json", "Authorization": "Bearer auth-bowong7777" }, json={ "model": llm_provider, "messages": [ { "role": "user", "content": prompt } ], "temperature": temperature, "max_tokens": max_tokens }) resp.raise_for_status() resp = resp.json() content = find_value_recursive("content", resp) content = re.sub(r'\n{2,}', '\n', content) except Exception as e: raise Exception("llm调用失败 {}".format(e)) return (content,) return _chat() class LLMChatMultiModalImageUpload: """llm chat""" @classmethod def INPUT_TYPES(s): input_dir = folder_paths.get_input_directory() files = [f for f in os.listdir(input_dir) if os.path.isfile(os.path.join(input_dir, f))] files = folder_paths.filter_files_content_types(files, ["image"]) return { "required": { "llm_provider": (["gpt-4o-1120", "gpt-4.1"],), "prompt": ("STRING", {"multiline": True}), "image": (sorted(files), {"image_upload": True}), "temperature": ("FLOAT", {"default": 0.7, "min": 0.0, "max": 1.0}), "max_tokens": ("INT", {"default": 4096, "min": 1, "max": 65535}), "timeout": ("INT", {"default": 120, "min": 30, "max": 900}), } } RETURN_TYPES = ("STRING",) RETURN_NAMES = ("llm输出",) FUNCTION = "chat" CATEGORY = "不忘科技-自定义节点🚩/LLM" def chat(self, llm_provider: str, prompt: str, image, temperature: float, max_tokens: int, timeout: int): @retry(Exception, tries=3, delay=1) def _chat(): try: image_path = folder_paths.get_annotated_filepath(image) mime_type, _ = guess_type(image_path) with open(image_path, "rb") as image_file: base64_encoded_data = base64.b64encode(image_file.read()).decode('utf-8') with httpx.Client(timeout=httpx.Timeout(timeout, connect=15)) as session: resp = session.post("https://gateway.bowong.cc/chat/completions", headers={ "Content-Type": "application/json", "Accept": "application/json", "Authorization": "Bearer auth-bowong7777" }, json={ "model": llm_provider, "messages": [ { "role": "user", "content": [ {"type": "text", "text": prompt}, { "type": "image_url", "image_url": { "url": f"data:{mime_type};base64,{base64_encoded_data}"}, }, ] } ], "temperature": temperature, "max_tokens": max_tokens }) resp.raise_for_status() resp = resp.json() content = find_value_recursive("content", resp) content = re.sub(r'\n{2,}', '\n', content) except Exception as e: # logger.exception("llm调用失败 {}".format(e)) raise Exception("llm调用失败 {}".format(e)) return (content,) return _chat() class LLMChatMultiModalImageTensor: """llm chat""" @classmethod def INPUT_TYPES(s): return { "required": { "llm_provider": (["gpt-4o-1120", "gpt-4.1"],), "prompt": ("STRING", {"multiline": True}), "image": ("IMAGE",), "temperature": ("FLOAT", {"default": 0.7, "min": 0.0, "max": 1.0}), "max_tokens": ("INT", {"default": 4096, "min": 1, "max": 65535}), "timeout": ("INT", {"default": 120, "min": 30, "max": 900}), } } RETURN_TYPES = ("STRING",) RETURN_NAMES = ("llm输出",) FUNCTION = "chat" CATEGORY = "不忘科技-自定义节点🚩/LLM" def chat(self, llm_provider: str, prompt: str, image: torch.Tensor, temperature: float, max_tokens: int, timeout: int): @retry(Exception, tries=3, delay=1) def _chat(): try: with httpx.Client(timeout=httpx.Timeout(timeout, connect=15)) as session: resp = session.post("https://gateway.bowong.cc/chat/completions", headers={ "Content-Type": "application/json", "Accept": "application/json", "Authorization": "Bearer auth-bowong7777" }, json={ "model": llm_provider, "messages": [ { "role": "user", "content": [ {"type": "text", "text": prompt}, { "type": "image_url", "image_url": {"url": image_tensor_to_base64(image)}, }, ] } ], "temperature": temperature, "max_tokens": max_tokens }) resp.raise_for_status() resp = resp.json() content = find_value_recursive("content", resp) content = re.sub(r'\n{2,}', '\n', content) except Exception as e: # logger.exception("llm调用失败 {}".format(e)) raise Exception("llm调用失败 {}".format(e)) return (content,) return _chat() class Jinja2RenderTemplate: @classmethod def INPUT_TYPES(s): return { "required": { "template": ("STRING", {"multiline": True}), "kv_map": ("STRING", {"multiline": True}), } } RETURN_TYPES = ("STRING",) RETURN_NAMES = ("prompt",) FUNCTION = "render_prompt" CATEGORY = "不忘科技-自定义节点🚩/LLM" def render_prompt(self, template: str, kv_map: str) -> tuple: """ 使用Jinja2渲染prompt模板 参数: template: 包含Jinja2标记的模板字符串 kv_map: 键值映射字典，用于提供模板渲染所需的变量 返回: 渲染后的字符串 异常: 如果模板中有未定义的变量，抛出jinja2.exceptions.UndefinedError """ kv_map = json.loads(kv_map) # 创建模板对象，设置为严格模式，未定义变量会抛出异常 template = Template(template, undefined=StrictUndefined) # 渲染模板 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 = 3 logger.info("开始轮询任务状态") for _ in range(0, wait_time, interval): logger.info("查询任务状态") result = send_request("get", f"https://{endpoint}/google/{job_id}", 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 = 3 logger.info("开始轮询任务状态") for _ in range(0, wait_time, interval): logger.info("查询任务状态") result = send_request("get", f"https://{endpoint}/google/{job_id}", 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)