386 lines
14 KiB
Python
386 lines
14 KiB
Python
import io
|
||
import json
|
||
from time import sleep, time
|
||
|
||
import folder_paths
|
||
import requests
|
||
import torch
|
||
from PIL import Image
|
||
from loguru import logger
|
||
from torchvision import transforms
|
||
|
||
from ..utils.http_utils import send_request
|
||
from ..utils.image_utils import tensor_to_image_bytes, base64_to_tensor
|
||
|
||
|
||
def url_to_tensor(image_url: str, max_retries: int = 3):
|
||
"""
|
||
从URL下载图片并转换为PyTorch张量，增强错误处理能力
|
||
|
||
参数:
|
||
image_url (str): 图片URL
|
||
max_retries (int): 最大重试次数
|
||
|
||
返回:
|
||
torch.Tensor: 形状为[C, H, W]的张量
|
||
|
||
异常:
|
||
HTTPError: 网络请求失败
|
||
ValueError: 无效图片格式
|
||
"""
|
||
headers = {
|
||
"User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36"
|
||
}
|
||
|
||
for attempt in range(max_retries):
|
||
try:
|
||
# 发送带User-Agent的请求
|
||
response = requests.get(image_url, headers=headers, stream=True, timeout=15)
|
||
response.raise_for_status()
|
||
|
||
# 检查内容类型是否为图像
|
||
content_type = response.headers.get("Content-Type", "")
|
||
if not content_type.startswith("image/"):
|
||
raise ValueError(f"URL返回非图像内容: {content_type}")
|
||
|
||
# 验证图像完整性
|
||
img_data = response.content
|
||
if len(img_data) < 100: # 极小数据通常不是有效图像
|
||
raise ValueError("下载的内容过小，可能不是完整图像")
|
||
|
||
# 尝试打开图像
|
||
img = Image.open(io.BytesIO(img_data)).convert("RGB")
|
||
|
||
# 转换为张量
|
||
transform = transforms.Compose([transforms.ToTensor()])
|
||
return transform(img).unsqueeze(0).permute(0, 2, 3, 1)
|
||
|
||
except (requests.exceptions.RequestException, ValueError) as e:
|
||
logger.warning(f"尝试 {attempt + 1}/{max_retries} 失败: {e}")
|
||
if attempt == max_retries - 1:
|
||
raise e
|
||
|
||
|
||
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 = 2
|
||
logger.info("开始轮询任务状态")
|
||
sleep(1)
|
||
for _ in range(0, wait_time, interval):
|
||
logger.info("查询任务状态")
|
||
result = send_request(
|
||
"get",
|
||
f"https://{endpoint}/google/{job_id}",
|
||
headers={"Authorization": "Bearer bowong7777"},
|
||
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": {
|
||
"prompt": (
|
||
"STRING",
|
||
{"default": "将背景去除，输出原尺寸图片", "multiline": True},
|
||
),
|
||
"temperature": ("FLOAT", {"default": 0.1, "min": 0, "max": 2}),
|
||
"topP": ("FLOAT", {"default": 0.7, "min": 0, "max": 1}),
|
||
"endpoint": (
|
||
"STRING",
|
||
{
|
||
"default": "bowongai-dev--bowong-ai-video-gemini-fastapi-webapp.modal.run"
|
||
},
|
||
),
|
||
},
|
||
"optional": {
|
||
"image": ("IMAGE",),
|
||
},
|
||
}
|
||
|
||
RETURN_TYPES = ("IMAGE",)
|
||
RETURN_NAMES = ("image",)
|
||
FUNCTION = "process"
|
||
OUTPUT_NODE = False
|
||
CATEGORY = "不忘科技-自定义节点🚩/图片/Gemini"
|
||
|
||
def process(
|
||
self, prompt: str, temperature: float, topP: float, endpoint: str, **kwargs
|
||
):
|
||
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"
|
||
if "image" in kwargs and kwargs["image"] is not None:
|
||
image = kwargs["image"]
|
||
files = {
|
||
"origin_image": (
|
||
"image." + format.lower(),
|
||
tensor_to_image_bytes(image, format),
|
||
f"image/{format.lower()}",
|
||
)
|
||
}
|
||
else:
|
||
files = None
|
||
logger.info("请求图像编辑")
|
||
job_resp = send_request(
|
||
"post",
|
||
f"https://{endpoint}/google/image/edit_custom",
|
||
headers={"x-google-api-key": api_key},
|
||
data={"prompt": prompt, "temperature": temperature, "topP": topP},
|
||
files=files,
|
||
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 = 2
|
||
logger.info("开始轮询任务状态")
|
||
sleep(1)
|
||
for _ in range(0, wait_time, interval):
|
||
logger.info("查询任务状态")
|
||
result = send_request(
|
||
"get",
|
||
f"https://{endpoint}/google/{job_id}",
|
||
headers={"Authorization": "Bearer bowong7777"},
|
||
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 ModalMidJourneyGenerateImage:
|
||
@classmethod
|
||
def INPUT_TYPES(cls):
|
||
return {
|
||
"required": {
|
||
"prompt": (
|
||
"STRING",
|
||
{"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}),
|
||
},
|
||
"optional": {
|
||
"image": ("IMAGE",),
|
||
},
|
||
}
|
||
|
||
RETURN_TYPES = ("IMAGE",)
|
||
RETURN_NAMES = ("image",)
|
||
FUNCTION = "process"
|
||
OUTPUT_NODE = False
|
||
CATEGORY = "不忘科技-自定义节点🚩/图片/Midjourney"
|
||
|
||
def process(
|
||
self, prompt: str, provider: str, endpoint: str, timeout: int, **kwargs
|
||
):
|
||
try:
|
||
logger.info("请求同步接口")
|
||
format = "PNG"
|
||
if "image" in kwargs and kwargs["image"] is not None:
|
||
image = kwargs["image"]
|
||
files = {
|
||
"img_file": (
|
||
"image." + format.lower(),
|
||
tensor_to_image_bytes(image, format),
|
||
f"image/{format.lower()}",
|
||
)
|
||
}
|
||
else:
|
||
files = None
|
||
|
||
interval = 3
|
||
logger.info("提交任务")
|
||
logger.info(f"https://{endpoint}/api/custom/image/submit/task")
|
||
job_resp = send_request(
|
||
"post",
|
||
f"https://{endpoint}/api/custom/image/submit/task",
|
||
data={"model_name": "302/mj", "prompt": prompt, "mode": "turbo"},
|
||
files=files,
|
||
timeout=timeout,
|
||
)
|
||
job_resp.raise_for_status()
|
||
job_resp = job_resp.json()
|
||
if not job_resp["status"]:
|
||
raise Exception("生成失败, 可能因为风控")
|
||
job_id = job_resp["data"]
|
||
start_time = time()
|
||
for _ in range(0, timeout // interval, interval):
|
||
logger.info(f"已等待 {time() - start_time} 秒，{interval} 秒后查询...")
|
||
sleep(interval)
|
||
logger.info(f"查询结果 {job_id}")
|
||
resp = send_request(
|
||
"get",
|
||
f"https://{endpoint}/api/custom/task/status?task_id={job_id}",
|
||
timeout=30,
|
||
)
|
||
resp.raise_for_status()
|
||
if resp.json()["status"] == "running":
|
||
logger.info(f"任务正在运行 {job_id}")
|
||
continue
|
||
if resp.json()["status"] == "failed":
|
||
raise Exception(f"生成失败: {resp.json()['msg']}")
|
||
if resp.json()["status"] == "success":
|
||
result_url = resp.json()["data"]
|
||
if not isinstance(result_url, list):
|
||
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:
|
||
raise e
|
||
|
||
|
||
class ModalMidJourneyDescribeImage:
|
||
@classmethod
|
||
def INPUT_TYPES(cls):
|
||
return {
|
||
"required": {
|
||
"image": ("IMAGE",),
|
||
"endpoint": (
|
||
"STRING",
|
||
{
|
||
"default": "bowongai-test--text-video-agent-fastapi-app.modal.run"
|
||
},
|
||
),
|
||
},
|
||
}
|
||
|
||
RETURN_TYPES = ("STRING",)
|
||
RETURN_NAMES = ("描述内容",)
|
||
FUNCTION = "process"
|
||
OUTPUT_NODE = False
|
||
CATEGORY = "不忘科技-自定义节点🚩/图片/Midjourney"
|
||
|
||
def process(self, image: torch.Tensor, endpoint: str):
|
||
try:
|
||
logger.info("请求同步接口")
|
||
format = "PNG"
|
||
job_resp = send_request(
|
||
"post",
|
||
f"https://{endpoint}/api/302/mj/sync/file/img/describe",
|
||
headers={"Authorization": "Bearer bowong7777"},
|
||
files={
|
||
"img_file": (
|
||
"image." + format.lower(),
|
||
tensor_to_image_bytes(image, format),
|
||
f"image/{format.lower()}",
|
||
)
|
||
},
|
||
timeout=300,
|
||
)
|
||
job_resp.raise_for_status()
|
||
job_resp = job_resp.json()
|
||
if not job_resp["status"]:
|
||
raise Exception("描述失败")
|
||
result = job_resp["data"]
|
||
return (result,)
|
||
except Exception as e:
|
||
raise e
|