300 lines
12 KiB
Python
300 lines
12 KiB
Python
import os
|
||
from qiniu import Auth, put_data
|
||
import io
|
||
from PIL import Image
|
||
import numpy as np
|
||
import uuid
|
||
import os
|
||
import uuid
|
||
import tempfile
|
||
import subprocess
|
||
import imageio_ffmpeg
|
||
from ..cloud_utils import load_cloud_config, CloudUploader
|
||
|
||
|
||
class QiniuUploader:
|
||
def __init__(self, access_key: str, secret_key: str, bucket_name: str, domain: str,):
|
||
self.access_key = access_key
|
||
self.secret_key = secret_key
|
||
self.bucket_name = bucket_name
|
||
self.domain = domain
|
||
self.q = Auth(self.access_key, self.secret_key)
|
||
|
||
def upload_binary(self, data: bytes, key: str = None) -> str:
|
||
"""
|
||
上传二进制数据到七牛云
|
||
:param data: 二进制内容
|
||
:param key: 文件名（可选），不传则由七牛自动生成
|
||
:return: 文件外链URL
|
||
"""
|
||
token = self.q.upload_token(self.bucket_name, key, 3600)
|
||
ret, info = put_data(token, key, data)
|
||
print(f"七牛 put_data 返回 ret: {ret}, info: {info}") # 打印上传结果
|
||
# 七牛 put_data 返回 ret 可能为 None，info.key 才是真实 key
|
||
real_key = None
|
||
if ret is not None and 'key' in ret:
|
||
real_key = ret['key']
|
||
elif hasattr(info, 'key') and info.key:
|
||
real_key = info.key
|
||
if real_key:
|
||
url = f"{self.domain}/{real_key}"
|
||
if not url.startswith("http://") and not url.startswith("https://"):
|
||
url = "http://" + url.lstrip("/")
|
||
return url
|
||
else:
|
||
raise Exception(f"上传失败: {info}")
|
||
|
||
# class JDCloudUploader(CloudUploader):
|
||
# """
|
||
# 京东云对象存储上传实现（需安装 jdcloud-sdk-python）
|
||
# """
|
||
# def __init__(self, access_key: str, secret_key: str, bucket_name: str, domain: str, region: str = "cn-north-1"):
|
||
# from jdcloud_sdk.services.oss.client.OssClient import OssClient
|
||
# from jdcloud_sdk.core.credential import Credential
|
||
# self.access_key = access_key
|
||
# self.secret_key = secret_key
|
||
# self.bucket_name = bucket_name
|
||
# self.domain = domain
|
||
# self.region = region
|
||
# self.credential = Credential(self.access_key, self.secret_key)
|
||
# self.client = OssClient(self.credential)
|
||
|
||
# def upload_binary(self, data: bytes, key: str = None) -> str:
|
||
# if not key:
|
||
# import uuid
|
||
# key = uuid.uuid4().hex
|
||
# resp = self.client.put_object(self.bucket_name, key, body=data)
|
||
# if hasattr(resp, "error") and resp.error:
|
||
# raise Exception(f"京东云上传失败: {resp.error}")
|
||
# url = f"{self.domain}/{key}"
|
||
# if not url.startswith("http://") and not url.startswith("https://"):
|
||
# url = "http://" + url.lstrip("/")
|
||
# return url
|
||
|
||
class CloudImageUploadNode:
|
||
"""
|
||
ComfyUI 图片张量直接上传云节点
|
||
"""
|
||
@classmethod
|
||
def INPUT_TYPES(cls):
|
||
cloud_type, config = load_cloud_config()
|
||
return {
|
||
"required": {
|
||
"access_key": ("STRING", {"default": config.get("access_key", "")}),
|
||
"secret_key": ("STRING", {"default": config.get("secret_key", "")}),
|
||
"bucket_name": ("STRING", {"default": config.get("bucket_name", "")}),
|
||
"domain": ("STRING", {"default": config.get("domain", "")}),
|
||
"images": ("IMAGE", ),
|
||
"folder": ("STRING", {"default": "output"}),
|
||
"key_prefix": ("STRING", {"default": "comfyui_"}),
|
||
"format": (["PNG", "JPEG", "GIF"], {"default": "PNG"}),
|
||
}
|
||
}
|
||
|
||
RETURN_TYPES = ("STRING",)
|
||
RETURN_NAMES = ("urls",)
|
||
FUNCTION = "upload_images"
|
||
CATEGORY = "云服务"
|
||
OUTPUT_NODE = True
|
||
|
||
def upload_images(self, access_key, secret_key, bucket_name, domain, images, folder, key_prefix, format):
|
||
cloud_type, _ = load_cloud_config()
|
||
if cloud_type == "jdcloud":
|
||
# uploader = JDCloudUploader(access_key, secret_key, bucket_name, domain)
|
||
raise NotImplementedError(f"暂不支持的云类型: {cloud_type}")
|
||
else:
|
||
uploader = QiniuUploader(access_key, secret_key, bucket_name, domain)
|
||
urls = []
|
||
arr = images.cpu().numpy() if hasattr(images, 'cpu') else images
|
||
for idx, image in enumerate(arr):
|
||
img = Image.fromarray(np.clip(255. * image, 0, 255).astype(np.uint8))
|
||
buf = io.BytesIO()
|
||
if format == "GIF":
|
||
img = img.convert("P", palette=Image.ADAPTIVE)
|
||
img.save(buf, format="GIF")
|
||
else:
|
||
img.save(buf, format=format)
|
||
buf.seek(0)
|
||
random_name = uuid.uuid4().hex
|
||
# 拼接文件夹路径
|
||
folder_path = folder.strip().strip('/')
|
||
key = f"{folder_path}/{key_prefix}{random_name}.{format.lower()}"
|
||
url = uploader.upload_binary(buf.getvalue(), key)
|
||
print(f"上传图片返回 url: {url}, 类型: {type(url)}")
|
||
urls.append(str(url))
|
||
return (urls,)
|
||
|
||
class CloudVideoUploadNode:
|
||
"""
|
||
ComfyUI 视频文件上传云节点
|
||
"""
|
||
@classmethod
|
||
def INPUT_TYPES(cls):
|
||
cloud_type, config = load_cloud_config()
|
||
return {
|
||
"required": {
|
||
"access_key": ("STRING", {"default": config.get("access_key", "")}),
|
||
"secret_key": ("STRING", {"default": config.get("secret_key", "")}),
|
||
"bucket_name": ("STRING", {"default": config.get("bucket_name", "")}),
|
||
"domain": ("STRING", {"default": config.get("domain", "")}),
|
||
"video_path": ("STRING", {"default": ""}),
|
||
"folder": ("STRING", {"default": "video"}),
|
||
"key_prefix": ("STRING", {"default": "comfyui_"}),
|
||
"ext": (["mp4", "mov", "avi", "mkv"], {"default": "mp4"}),
|
||
}
|
||
}
|
||
|
||
RETURN_TYPES = ("STRING",)
|
||
RETURN_NAMES = ("url",)
|
||
FUNCTION = "upload_video"
|
||
CATEGORY = "云服务"
|
||
OUTPUT_NODE = True
|
||
|
||
def upload_video(self, access_key, secret_key, bucket_name, domain, video_path, folder, key_prefix, ext):
|
||
cloud_type, _ = load_cloud_config()
|
||
if cloud_type == "jdcloud":
|
||
# uploader = JDCloudUploader(access_key, secret_key, bucket_name, domain)
|
||
raise NotImplementedError(f"暂不支持的云类型: {cloud_type}")
|
||
else:
|
||
uploader = QiniuUploader(access_key, secret_key, bucket_name, domain)
|
||
if not os.path.isfile(video_path):
|
||
raise FileNotFoundError(f"视频文件不存在: {video_path}")
|
||
with open(video_path, "rb") as f:
|
||
data = f.read()
|
||
random_name = uuid.uuid4().hex
|
||
folder_path = folder.strip().strip('/')
|
||
key = f"{folder_path}/{key_prefix}{random_name}.{ext}"
|
||
url = uploader.upload_binary(data, key)
|
||
print(f"上传视频返回 url: {url}, 类型: {type(url)}")
|
||
return (url,)
|
||
|
||
class CloudImagesToVideoAndUpload:
|
||
"""
|
||
将图片张量序列合成为视频并上传到云存储（支持多云扩展）
|
||
输入: 图片张量 (N, H, W, C)
|
||
输出: 云存储视频URL
|
||
"""
|
||
@classmethod
|
||
def INPUT_TYPES(cls):
|
||
cloud_type, config = load_cloud_config()
|
||
return {
|
||
"required": {
|
||
"images": ("IMAGE",),
|
||
"fps": ("INT", {"default": 12, "min": 1, "max": 60}),
|
||
"cloud_type": ("STRING", {"default": cloud_type}),
|
||
"access_key": ("STRING", {"default": config["access_key"]}),
|
||
"secret_key": ("STRING", {"default": config["secret_key"]}),
|
||
"bucket_name": ("STRING", {"default": config["bucket_name"]}),
|
||
"domain": ("STRING", {"default": config["domain"]}),
|
||
"folder": ("STRING", {"default": "video"}),
|
||
"key_prefix": ("STRING", {"default": "comfyui_"}),
|
||
"ext": (["mp4", "mov", "avi", "mkv"], {"default": "mp4"}),
|
||
"audio": ("AUDIO", {"default": None}), # 改为AUDIO类型
|
||
}
|
||
}
|
||
|
||
RETURN_TYPES = ("STRING",)
|
||
RETURN_NAMES = ("url",)
|
||
FUNCTION = "images_to_video_and_upload"
|
||
CATEGORY = "云服务"
|
||
OUTPUT_NODE = True
|
||
|
||
def images_to_video_and_upload(self, images, fps, cloud_type, access_key, secret_key, bucket_name, domain, folder, key_prefix, ext, audio=None):
|
||
config = load_cloud_config()[1]
|
||
access_key = access_key or config.get("access_key", "")
|
||
secret_key = secret_key or config.get("secret_key", "")
|
||
bucket_name = bucket_name or config.get("bucket_name", "")
|
||
domain = domain or config.get("domain", "")
|
||
# 1. 张量转图片序列
|
||
arr = images.cpu().numpy() if hasattr(images, 'cpu') else images
|
||
img_list = [Image.fromarray(np.clip(255. * img, 0, 255).astype(np.uint8)) for img in arr]
|
||
# 2. 合成视频
|
||
height, width = img_list[0].height, img_list[0].width
|
||
pix_fmt = 'rgb24'
|
||
with tempfile.NamedTemporaryFile(suffix=f'.{ext}', delete=False) as tmpfile:
|
||
tmp_path = tmpfile.name
|
||
# 先生成无音频视频，临时文件名用 _noaudio 结尾但扩展名标准
|
||
tmp_video_path = tmp_path.replace(f'.{ext}', f'_noaudio.{ext}')
|
||
cmd = [
|
||
imageio_ffmpeg.get_ffmpeg_exe(),
|
||
'-y',
|
||
'-f', 'rawvideo',
|
||
'-vcodec', 'rawvideo',
|
||
'-s', f'{width}x{height}',
|
||
'-pix_fmt', pix_fmt,
|
||
'-r', str(fps),
|
||
'-i', '-',
|
||
'-an',
|
||
'-vcodec', 'libx264',
|
||
'-pix_fmt', 'yuv420p',
|
||
tmp_video_path
|
||
]
|
||
proc = subprocess.Popen(cmd, stdin=subprocess.PIPE)
|
||
for im in img_list:
|
||
proc.stdin.write(im.convert('RGB').tobytes())
|
||
proc.stdin.close()
|
||
proc.wait()
|
||
# 如果有AUDIO，保存为wav临时文件再合成
|
||
def is_valid_audio(audio):
|
||
if not (audio and isinstance(audio, dict) and "waveform" in audio and "sample_rate" in audio):
|
||
return False
|
||
waveform = audio["waveform"]
|
||
if not hasattr(waveform, 'numel') or waveform.numel() == 0:
|
||
return False
|
||
if waveform.abs().sum().item() == 0:
|
||
return False
|
||
return True
|
||
if is_valid_audio(audio):
|
||
with tempfile.NamedTemporaryFile(suffix='.wav', delete=False) as tmp_audio:
|
||
audio_path = tmp_audio.name
|
||
waveform = audio["waveform"]
|
||
sample_rate = audio["sample_rate"]
|
||
# waveform shape: (1, channels, samples) or (channels, samples)
|
||
if waveform.dim() == 3:
|
||
waveform = waveform.squeeze(0)
|
||
import torchaudio
|
||
torchaudio.save(audio_path, waveform, sample_rate)
|
||
# 合成音视频到tmp_path
|
||
merge_cmd = [
|
||
imageio_ffmpeg.get_ffmpeg_exe(),
|
||
'-y',
|
||
'-i', tmp_video_path,
|
||
'-i', audio_path,
|
||
'-c:v', 'copy',
|
||
'-c:a', 'aac',
|
||
'-shortest',
|
||
tmp_path
|
||
]
|
||
subprocess.run(merge_cmd, check=True)
|
||
os.remove(tmp_video_path)
|
||
os.remove(audio_path)
|
||
else:
|
||
# 无音频直接重命名
|
||
os.rename(tmp_video_path, tmp_path)
|
||
# 3. 上传到云存储
|
||
if cloud_type == "qiniu":
|
||
uploader = QiniuUploader(access_key, secret_key, bucket_name, domain)
|
||
elif cloud_type == "jdcloud":
|
||
# uploader = JDCloudUploader(access_key, secret_key, bucket_name, domain)
|
||
raise NotImplementedError(f"暂不支持的云类型: {cloud_type}")
|
||
random_name = uuid.uuid4().hex
|
||
folder_path = folder.strip().strip('/')
|
||
key = f"{folder_path}/{key_prefix}{random_name}.{ext}"
|
||
with open(tmp_path, "rb") as f:
|
||
data = f.read()
|
||
url = uploader.upload_binary(data, key)
|
||
os.remove(tmp_path)
|
||
return (url,)
|
||
|
||
# 注册到节点映射
|
||
NODE_CLASS_MAPPINGS = {
|
||
"CloudImageUploadNode": CloudImageUploadNode,
|
||
"CloudVideoUploadNode": CloudVideoUploadNode,
|
||
"CloudImagesToVideoAndUpload": CloudImagesToVideoAndUpload
|
||
}
|
||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||
"CloudImageUploadNode": "☁️ 图片上传到云 (IMAGE)",
|
||
"CloudVideoUploadNode": "☁️ 视频上传到云 (VIDEO)",
|
||
"CloudImagesToVideoAndUpload": "🖼️图片合成视频并上传到云"
|
||
}
|