bowong
/
ComfyUI-CustomNode
6
0
Fork 0
Code Pull Requests Releases Activity

REFINE 优化图像处理逻辑，统一代码风格并增强可读性

This commit is contained in:
iHeyTang 2025-08-12 15:46:46 +08:00
parent 667c546449
commit bbd0c5325e
1 changed files with 170 additions and 96 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

266
nodes/image_modal_nodes.py
View File

@ -1,6 +1,6 @@
import io import io
import json import json
from time import sleep from time import sleep, time
import folder_paths import folder_paths
import requests import requests
@ -29,7 +29,8 @@ def url_to_tensor(image_url: str, max_retries: int = 3):
ValueError: 无效图片格式 ValueError: 无效图片格式
""" """
headers = { headers = {
'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36'} "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36"
}
for attempt in range(max_retries): for attempt in range(max_retries):
try: try:
@ -38,8 +39,8 @@ def url_to_tensor(image_url: str, max_retries: int = 3):
response.raise_for_status() response.raise_for_status()
# 检查内容类型是否为图像 # 检查内容类型是否为图像
content_type = response.headers.get('Content-Type', '') content_type = response.headers.get("Content-Type", "")
if not content_type.startswith('image/'): if not content_type.startswith("image/"):
raise ValueError(f"URL返回非图像内容: {content_type}") raise ValueError(f"URL返回非图像内容: {content_type}")
# 验证图像完整性 # 验证图像完整性
@ -48,12 +49,10 @@ def url_to_tensor(image_url: str, max_retries: int = 3):
raise ValueError("下载的内容过小，可能不是完整图像") raise ValueError("下载的内容过小，可能不是完整图像")
# 尝试打开图像 # 尝试打开图像
img = Image.open(io.BytesIO(img_data)).convert('RGB') img = Image.open(io.BytesIO(img_data)).convert("RGB")
# 转换为张量 # 转换为张量
transform = transforms.Compose([ transform = transforms.Compose([transforms.ToTensor()])
transforms.ToTensor()
])
return transform(img).unsqueeze(0).permute(0, 2, 3, 1) return transform(img).unsqueeze(0).permute(0, 2, 3, 1)
except (requests.exceptions.RequestException, ValueError) as e: except (requests.exceptions.RequestException, ValueError) as e:
@ -70,7 +69,12 @@ class ModalClothesMask:
"image": ("IMAGE",), "image": ("IMAGE",),
"mask_color": ("STRING", {"default": "绿色"}), "mask_color": ("STRING", {"default": "绿色"}),
"clothes_type": ("STRING", {"default": "裤子"}), "clothes_type": ("STRING", {"default": "裤子"}),
"endpoint": ("STRING", {"default": "bowongai-dev--bowong-ai-video-gemini-fastapi-webapp.modal.run"}), "endpoint": (
"STRING",
{
"default": "bowongai-dev--bowong-ai-video-gemini-fastapi-webapp.modal.run"
},
),
}, },
} }
@ -80,25 +84,37 @@ class ModalClothesMask:
OUTPUT_NODE = False OUTPUT_NODE = False
CATEGORY = "不忘科技-自定义节点🚩/图片/Gemini" CATEGORY = "不忘科技-自定义节点🚩/图片/Gemini"
def process(self, image: torch.Tensor, mask_color: str, clothes_type: str, endpoint: str): def process(
self, image: torch.Tensor, mask_color: str, clothes_type: str, endpoint: str
):
try: try:
timeout = 60 timeout = 60
logger.info("获取token") logger.info("获取token")
api_key = send_request("get", f"https://{endpoint}/google/access-token", api_key = send_request(
headers={'Authorization': 'Bearer bowong7777'}, timeout=timeout).json()[ "get",
"access_token"] f"https://{endpoint}/google/access-token",
headers={"Authorization": "Bearer bowong7777"},
timeout=timeout,
).json()["access_token"]
format = "PNG" format = "PNG"
logger.info("请求图像编辑") logger.info("请求图像编辑")
job_resp = send_request("post", f"https://{endpoint}/google/image/clothes_mark", job_resp = send_request(
headers={'x-google-api-key': api_key}, "post",
data={ f"https://{endpoint}/google/image/clothes_mark",
"mark_clothes_type": clothes_type, headers={"x-google-api-key": api_key},
"mark_color": mask_color, data={
}, "mark_clothes_type": clothes_type,
files={"origin_image": ( "mark_color": mask_color,
'image.' + format.lower(), tensor_to_image_bytes(image, format), },
f'image/{format.lower()}')}, files={
timeout=timeout) "origin_image": (
"image." + format.lower(),
tensor_to_image_bytes(image, format),
f"image/{format.lower()}",
)
},
timeout=timeout,
)
job_resp.raise_for_status() job_resp.raise_for_status()
job_resp = job_resp.json() job_resp = job_resp.json()
if not job_resp["success"]: if not job_resp["success"]:
@ -111,8 +127,12 @@ class ModalClothesMask:
sleep(1) sleep(1)
for _ in range(0, wait_time, interval): for _ in range(0, wait_time, interval):
logger.info("查询任务状态") logger.info("查询任务状态")
result = send_request("get", f"https://{endpoint}/google/{job_id}", result = send_request(
headers={'Authorization': 'Bearer bowong7777'}, timeout=timeout) "get",
f"https://{endpoint}/google/{job_id}",
headers={"Authorization": "Bearer bowong7777"},
timeout=timeout,
)
if result.status_code == 200: if result.status_code == 200:
result = result.json() result = result.json()
if result["status"] == "success": if result["status"] == "success":
@ -133,14 +153,22 @@ class ModalEditCustom:
def INPUT_TYPES(cls): def INPUT_TYPES(cls):
return { return {
"required": { "required": {
"prompt": ("STRING", {"default": "将背景去除，输出原尺寸图片", "multiline": True}), "prompt": (
"STRING",
{"default": "将背景去除，输出原尺寸图片", "multiline": True},
),
"temperature": ("FLOAT", {"default": 0.1, "min": 0, "max": 2}), "temperature": ("FLOAT", {"default": 0.1, "min": 0, "max": 2}),
"topP": ("FLOAT", {"default": 0.7, "min": 0, "max": 1}), "topP": ("FLOAT", {"default": 0.7, "min": 0, "max": 1}),
"endpoint": ("STRING", {"default": "bowongai-dev--bowong-ai-video-gemini-fastapi-webapp.modal.run"}), "endpoint": (
"STRING",
{
"default": "bowongai-dev--bowong-ai-video-gemini-fastapi-webapp.modal.run"
},
),
}, },
"optional": { "optional": {
"image": ("IMAGE",), "image": ("IMAGE",),
} },
} }
RETURN_TYPES = ("IMAGE",) RETURN_TYPES = ("IMAGE",)
@ -149,31 +177,39 @@ class ModalEditCustom:
OUTPUT_NODE = False OUTPUT_NODE = False
CATEGORY = "不忘科技-自定义节点🚩/图片/Gemini" CATEGORY = "不忘科技-自定义节点🚩/图片/Gemini"
def process(self, prompt: str, temperature: float, topP: float, endpoint: str, **kwargs): def process(
self, prompt: str, temperature: float, topP: float, endpoint: str, **kwargs
):
try: try:
timeout = 60 timeout = 60
logger.info("获取token") logger.info("获取token")
api_key = send_request("get", f"https://{endpoint}/google/access-token", api_key = send_request(
headers={'Authorization': 'Bearer bowong7777'}, timeout=timeout).json()[ "get",
"access_token"] f"https://{endpoint}/google/access-token",
headers={"Authorization": "Bearer bowong7777"},
timeout=timeout,
).json()["access_token"]
format = "PNG" format = "PNG"
if "image" in kwargs and kwargs["image"] is not None: if "image" in kwargs and kwargs["image"] is not None:
image = kwargs["image"] image = kwargs["image"]
files = {"origin_image": ( files = {
'image.' + format.lower(), tensor_to_image_bytes(image, format), "origin_image": (
f'image/{format.lower()}')} "image." + format.lower(),
tensor_to_image_bytes(image, format),
f"image/{format.lower()}",
)
}
else: else:
files = None files = None
logger.info("请求图像编辑") logger.info("请求图像编辑")
job_resp = send_request("post", f"https://{endpoint}/google/image/edit_custom", job_resp = send_request(
headers={'x-google-api-key': api_key}, "post",
data={ f"https://{endpoint}/google/image/edit_custom",
"prompt": prompt, headers={"x-google-api-key": api_key},
"temperature": temperature, data={"prompt": prompt, "temperature": temperature, "topP": topP},
"topP": topP files=files,
}, timeout=timeout,
files=files, )
timeout=timeout)
job_resp.raise_for_status() job_resp.raise_for_status()
job_resp = job_resp.json() job_resp = job_resp.json()
if not job_resp["success"]: if not job_resp["success"]:
@ -186,8 +222,12 @@ class ModalEditCustom:
sleep(1) sleep(1)
for _ in range(0, wait_time, interval): for _ in range(0, wait_time, interval):
logger.info("查询任务状态") logger.info("查询任务状态")
result = send_request("get", f"https://{endpoint}/google/{job_id}", result = send_request(
headers={'Authorization': 'Bearer bowong7777'}, timeout=timeout) "get",
f"https://{endpoint}/google/{job_id}",
headers={"Authorization": "Bearer bowong7777"},
timeout=timeout,
)
if result.status_code == 200: if result.status_code == 200:
result = result.json() result = result.json()
if result["status"] == "success": if result["status"] == "success":
@ -208,14 +248,22 @@ class ModalMidJourneyGenerateImage:
def INPUT_TYPES(cls): def INPUT_TYPES(cls):
return { return {
"required": { "required": {
"prompt": ("STRING", {"default": "一幅宏大壮美的山川画卷", "multiline": True}), "prompt": (
"provider":(["ttapi","302ai"],), "STRING",
"endpoint": ("STRING", {"default": "bowongai-test--text-video-agent-fastapi-app.modal.run"}), {"default": "一幅宏大壮美的山川画卷", "multiline": True},
),
"provider": (["ttapi", "302ai"],),
"endpoint": (
"STRING",
{
"default": "bowongai-test--text-video-agent-fastapi-app.modal.run"
},
),
"timeout": ("INT", {"default": 300, "min": 10, "max": 1200}), "timeout": ("INT", {"default": 300, "min": 10, "max": 1200}),
}, },
"optional": { "optional": {
"image": ("IMAGE",), "image": ("IMAGE",),
} },
} }
RETURN_TYPES = ("IMAGE",) RETURN_TYPES = ("IMAGE",)
@ -224,53 +272,67 @@ class ModalMidJourneyGenerateImage:
OUTPUT_NODE = False OUTPUT_NODE = False
CATEGORY = "不忘科技-自定义节点🚩/图片/Midjourney" CATEGORY = "不忘科技-自定义节点🚩/图片/Midjourney"
def process(self, prompt: str, provider:str, endpoint: str, timeout: int, **kwargs): def process(
self, prompt: str, provider: str, endpoint: str, timeout: int, **kwargs
):
try: try:
logger.info("请求同步接口") logger.info("请求同步接口")
format = "PNG" format = "PNG"
if "image" in kwargs and kwargs["image"] is not None: if "image" in kwargs and kwargs["image"] is not None:
image = kwargs["image"] image = kwargs["image"]
files = {"img_file": ( files = {
'image.' + format.lower(), tensor_to_image_bytes(image, format), "img_file": (
f'image/{format.lower()}')} "image." + format.lower(),
tensor_to_image_bytes(image, format),
f"image/{format.lower()}",
)
}
else: else:
files = None files = None
if provider == "302ai": interval = 3
interval = 3 logger.info("提交任务")
logger.info("提交任务") logger.info(f"https://{endpoint}/api/custom/image/submit/task")
logger.info(f"https://{endpoint}/api/custom/image/submit/task") job_resp = send_request(
job_resp = send_request("post", f"https://{endpoint}/api/custom/image/submit/task", "post",
data={"model_name":"302/mj", "prompt": prompt, "mode": "turbo"}, f"https://{endpoint}/api/custom/image/submit/task",
files=files, data={"model_name": "302/mj", "prompt": prompt, "mode": "turbo"},
timeout=timeout) files=files,
job_resp.raise_for_status() timeout=timeout,
job_resp = job_resp.json() )
if not job_resp["status"]: job_resp.raise_for_status()
raise Exception("生成失败, 可能因为风控") job_resp = job_resp.json()
job_id = job_resp["data"] if not job_resp["status"]:
for _ in range(0, timeout // interval, interval): raise Exception("生成失败, 可能因为风控")
logger.info("等待" + str(interval) + "秒") job_id = job_resp["data"]
sleep(interval) start_time = time()
logger.info("查询结果") for _ in range(0, timeout // interval, interval):
resp = send_request("get", f"https://{endpoint}/api/custom/task/status?task_id={job_id}", timeout=30) logger.info(f"已等待 {time() - start_time} 秒，{interval} 秒后查询...")
resp.raise_for_status() sleep(interval)
if resp.json()["status"] == "running": logger.info(f"查询结果 {job_id}")
logger.info("任务正在运行") resp = send_request(
continue "get",
if resp.json()["status"] == "failed": f"https://{endpoint}/api/custom/task/status?task_id={job_id}",
raise Exception(f"生成失败: {resp.json()['msg']}") timeout=30,
if resp.json()["status"] == "success": )
result_url = resp.json()["data"] resp.raise_for_status()
if not isinstance(result_url, list): if resp.json()["status"] == "running":
raise Exception("生成失败，返回结果为空") logger.info(f"任务正在运行 {job_id}")
result_list = [] continue
for url in result_url: if resp.json()["status"] == "failed":
logger.success("img_url: " + url) raise Exception(f"生成失败: {resp.json()['msg']}")
result_list.append(url_to_tensor(url).squeeze(0)) if resp.json()["status"] == "success":
result_list = torch.stack(result_list, dim=0) result_url = resp.json()["data"]
return (result_list,) if not isinstance(result_url, list):
raise Exception("等待超时") raise Exception("生成失败，返回结果为空")
result_list = []
for url in result_url:
logger.success("img_url: " + url)
result_list.append(url_to_tensor(url).squeeze(0))
result_list = torch.stack(result_list, dim=0)
logger.success(f"生成成功 {job_id}")
return (result_list,)
raise Exception("等待超时")
except Exception as e: except Exception as e:
raise e raise e
@ -281,7 +343,12 @@ class ModalMidJourneyDescribeImage:
return { return {
"required": { "required": {
"image": ("IMAGE",), "image": ("IMAGE",),
"endpoint": ("STRING", {"default": "bowongai-test--text-video-agent-fastapi-app.modal.run"}), "endpoint": (
"STRING",
{
"default": "bowongai-test--text-video-agent-fastapi-app.modal.run"
},
),
}, },
} }
@ -295,12 +362,19 @@ class ModalMidJourneyDescribeImage:
try: try:
logger.info("请求同步接口") logger.info("请求同步接口")
format = "PNG" format = "PNG"
job_resp = send_request("post", f"https://{endpoint}/api/302/mj/sync/file/img/describe", job_resp = send_request(
headers={'Authorization': 'Bearer bowong7777'}, "post",
files={"img_file": ( f"https://{endpoint}/api/302/mj/sync/file/img/describe",
'image.' + format.lower(), tensor_to_image_bytes(image, format), headers={"Authorization": "Bearer bowong7777"},
f'image/{format.lower()}')}, files={
timeout=300) "img_file": (
"image." + format.lower(),
tensor_to_image_bytes(image, format),
f"image/{format.lower()}",
)
},
timeout=300,
)
job_resp.raise_for_status() job_resp.raise_for_status()
job_resp = job_resp.json() job_resp = job_resp.json()
if not job_resp["status"]: if not job_resp["status"]: