261 lines
10 KiB
Python
261 lines
10 KiB
Python
import requests
|
|
import time
|
|
import json
|
|
from PIL import Image, ImageOps
|
|
from io import BytesIO
|
|
import numpy as np
|
|
import torch
|
|
import os # 确保导入 os 库
|
|
|
|
class ExecuteNode:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"apiConfig": ("STRUCT",), # 设置节点的输入
|
|
"nodeInfoList": ("ARRAY", {"default": []}), # NodeInfoList节点的输出
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",) # 支持单张和多张图片输出
|
|
RETURN_NAMES = ("images",) # 定义返回名称，与 ComfyUI 的预期匹配
|
|
|
|
CATEGORY = "RunningHub"
|
|
FUNCTION = "process" # 指向 process 方法
|
|
|
|
def process(self, apiConfig, nodeInfoList):
|
|
"""
|
|
该节点通过调用 RunningHub API 创建任务并返回生成的图片链接。
|
|
"""
|
|
# 打印请求数据，方便调试
|
|
print(f"API request data: {apiConfig}")
|
|
print(f"Node info list: {nodeInfoList}")
|
|
|
|
# 1. 查询账户状态，检查是否可以提交任务
|
|
account_status = self.check_account_status(apiConfig["apiKey"])
|
|
if int(account_status["currentTaskCounts"]) > 0:
|
|
print("There are tasks running, waiting for them to finish.")
|
|
# 等待最多 10 分钟，如果任务未完成，则超时
|
|
start_time = time.time()
|
|
while account_status["currentTaskCounts"] > 0 and time.time() - start_time < 600:
|
|
time.sleep(2) # 每 2 秒查询一次
|
|
account_status = self.check_account_status(apiConfig["apiKey"])
|
|
if int(account_status["currentTaskCounts"]) > 0:
|
|
raise Exception("Timeout: There are still running tasks after 10 minutes.")
|
|
|
|
# 2. 创建任务
|
|
task_creation_result = self.create_task(apiConfig, nodeInfoList)
|
|
if task_creation_result["code"] != 0:
|
|
raise Exception(f"Task creation failed: {task_creation_result['msg']}")
|
|
|
|
task_id = task_creation_result["data"]["taskId"]
|
|
task_status = task_creation_result["data"]["taskStatus"]
|
|
print(f"Task created successfully, taskId: {task_id}, status: {task_status}")
|
|
|
|
# 3. 查询任务状态直到任务完成
|
|
while task_status != "success":
|
|
print(f"Task still running, checking again in 2 seconds...")
|
|
time.sleep(2) # 每 2 秒检查一次任务状态
|
|
task_status_result = self.check_task_status(task_id, apiConfig["apiKey"])
|
|
print(f"Task info, taskId: {task_id}, status: {task_status_result}")
|
|
|
|
if isinstance(task_status_result, dict):
|
|
task_status = task_status_result.get("taskStatus", "unknown") # 从结果中获取任务状态
|
|
elif isinstance(task_status_result, list):
|
|
# 假设任务完成后返回的数据列表表示任务成功
|
|
task_status = "success"
|
|
else:
|
|
task_status = "unknown"
|
|
|
|
if task_status not in ["RUNNING", "success"]:
|
|
print(f"Task failed or completed with status: {task_status}")
|
|
break
|
|
|
|
# 4. 任务完成，处理输出
|
|
return self.process_task_output(task_id, apiConfig["apiKey"])
|
|
|
|
def check_account_status(self, api_key):
|
|
"""
|
|
查询账户状态，检查是否可以提交新任务
|
|
"""
|
|
url = "https://www.runninghub.cn/uc/openapi/accountStatus"
|
|
headers = {
|
|
"User-Agent": "Apifox/1.0.0 (https://apifox.com)",
|
|
"Content-Type": "application/json",
|
|
}
|
|
data = {
|
|
"apikey": api_key
|
|
}
|
|
|
|
response = requests.post(url, json=data, headers=headers)
|
|
result = response.json()
|
|
if result["code"] != 0:
|
|
raise Exception(f"Failed to get account status: {result['msg']}")
|
|
# 检查并确保 currentTaskCounts 是整数
|
|
try:
|
|
current_task_counts = int(result["data"]["currentTaskCounts"])
|
|
except ValueError:
|
|
raise Exception("Invalid value for currentTaskCounts. It should be an integer.")
|
|
|
|
result["data"]["currentTaskCounts"] = current_task_counts
|
|
return result["data"]
|
|
|
|
def create_task(self, apiConfig, nodeInfoList):
|
|
"""
|
|
创建任务
|
|
"""
|
|
url = "https://www.runninghub.cn/task/openapi/create"
|
|
headers = {
|
|
"Content-Type": "application/json",
|
|
"User-Agent": "Apifox/1.0.0 (https://apifox.com)",
|
|
}
|
|
data = {
|
|
"workflowId": apiConfig["workflowId"],
|
|
"apiKey": apiConfig["apiKey"],
|
|
"nodeInfoList": [
|
|
{
|
|
"nodeId": int(nodeInfo["nodeId"]), # 确保 nodeId 为整数类型
|
|
"fieldName": nodeInfo["fieldName"],
|
|
"fieldValue": nodeInfo["fieldValue"],
|
|
}
|
|
for nodeInfo in nodeInfoList
|
|
],
|
|
}
|
|
|
|
response = requests.post(url, json=data, headers=headers)
|
|
return response.json()
|
|
|
|
def check_task_status(self, task_id, api_key):
|
|
"""
|
|
查询任务状态
|
|
"""
|
|
url = "https://www.runninghub.cn/task/openapi/outputs"
|
|
headers = {
|
|
"User-Agent": "Apifox/1.0.0 (https://apifox.com)",
|
|
"Content-Type": "application/json",
|
|
}
|
|
data = {
|
|
"taskId": task_id,
|
|
"apiKey": api_key
|
|
}
|
|
|
|
response = requests.post(url, json=data, headers=headers)
|
|
|
|
# 打印响应以便调试
|
|
print("Response Status Code:", response.status_code)
|
|
try:
|
|
response_json = response.json()
|
|
print("Response JSON:", json.dumps(response_json, indent=4, ensure_ascii=False))
|
|
except ValueError:
|
|
print("Response Text:", response.text)
|
|
|
|
if response.status_code != 200:
|
|
raise Exception(f"HTTP request failed with status code: {response.status_code}")
|
|
|
|
result = response.json()
|
|
|
|
# 检查 'code' 和 'msg' 字段
|
|
if result.get("code") != 0:
|
|
# 如果任务正在运行，返回一个特定的状态
|
|
if result.get("msg") == "APIKEY_TASK_IS_RUNNING":
|
|
return {"taskStatus": "RUNNING"}
|
|
# 检查 'data' 是否存在并且是列表类型
|
|
if result.get("data") and isinstance(result["data"], list):
|
|
if len(result["data"]) > 0:
|
|
return result["data"] # 返回整个列表
|
|
else:
|
|
return {"taskStatus": "RUNNING"} # 如果 data 是空列表，任务仍在运行
|
|
else:
|
|
return {"taskStatus": "RUNNING"} # 如果 data 是 None，任务仍在运行
|
|
|
|
def process_task_output(self, task_id, api_key):
|
|
"""
|
|
处理任务输出，返回文件链接。
|
|
"""
|
|
task_status_result = self.check_task_status(task_id, api_key)
|
|
|
|
# 记录任务状态结果以了解其结构
|
|
print("Task Status Result:", json.dumps(task_status_result, indent=4, ensure_ascii=False))
|
|
|
|
image_urls = []
|
|
|
|
# 确保 task_status_result 是字典或列表
|
|
if isinstance(task_status_result, dict):
|
|
# 检查 fileUrl 和 fileType
|
|
file_url = task_status_result.get("fileUrl")
|
|
file_type = task_status_result.get("fileType")
|
|
if file_url and file_type.lower() in ["png", "jpg", "jpeg"]:
|
|
image_urls.append(file_url)
|
|
elif isinstance(task_status_result, list):
|
|
for output in task_status_result:
|
|
if isinstance(output, dict):
|
|
file_url = output.get("fileUrl")
|
|
file_type = output.get("fileType")
|
|
if file_url and file_type.lower() in ["png", "jpg", "jpeg"]:
|
|
image_urls.append(file_url)
|
|
|
|
if not image_urls:
|
|
raise Exception("No valid image output found.")
|
|
|
|
# 下载并处理所有图像
|
|
image_data_list = []
|
|
for url in image_urls:
|
|
print("Downloading image from URL:", url) # 记录图像 URL
|
|
image_tensor = self.download_image(url) # 下载并处理图像
|
|
print("Image downloaded and processed successfully.")
|
|
|
|
# 打印张量信息，避免在 uint8 上调用 mean()
|
|
print(f"Image tensor shape: {image_tensor.shape}")
|
|
mean_val = image_tensor.mean()
|
|
print(f"Image tensor min: {image_tensor.min()}, max: {image_tensor.max()}, mean: {mean_val}")
|
|
|
|
# 确保返回的是 torch.Tensor
|
|
if not isinstance(image_tensor, torch.Tensor):
|
|
raise TypeError(f"Expected torch.Tensor, got {type(image_tensor)}")
|
|
|
|
image_data_list.append(image_tensor)
|
|
|
|
print(f"Returning {len(image_data_list)} images.")
|
|
return (image_data_list,) # 返回一个包含图像列表的元组，与 RETURN_TYPES 和 RETURN_NAMES 匹配
|
|
|
|
def download_image(self, image_url):
|
|
"""
|
|
从 URL 下载图像并转换为适合预览或保存的 torch.Tensor 格式。
|
|
"""
|
|
response = requests.get(image_url)
|
|
|
|
if response.status_code == 200:
|
|
img = Image.open(BytesIO(response.content)).convert("RGB")
|
|
|
|
# 转换为 numpy 数组并调整数据类型和范围
|
|
img_array = np.array(img).astype(np.float32) / 255.0 # 归一化到 [0, 1]
|
|
|
|
# **保持形状为 [H, W, C]，不进行 permute 操作**
|
|
img_tensor = torch.from_numpy(img_array) # 形状 [H, W, C]
|
|
|
|
# 打印图像尺寸
|
|
print(f"Downloaded image dimensions: {img_tensor.shape}") # 例如 (高度, 宽度, 3)
|
|
|
|
return img_tensor
|
|
|
|
else:
|
|
raise Exception(f"Failed to download image: {image_url}")
|
|
|
|
def download_video(self, video_url):
|
|
"""
|
|
从 URL 下载视频。
|
|
根据 ComfyUI 的要求实现此方法。
|
|
"""
|
|
response = requests.get(video_url, stream=True)
|
|
if response.status_code == 200:
|
|
# 示例：将视频保存到临时位置并返回路径或数据
|
|
video_content = response.content
|
|
# 您可能需要根据 ComfyUI 的要求处理视频数据
|
|
# 目前，返回原始字节
|
|
return video_content
|
|
else:
|
|
raise Exception(f"Failed to download video: {video_url}")
|