99 lines
3.1 KiB
Python
99 lines
3.1 KiB
Python
# -*- coding: utf-8 -*-
|
||
"""
|
||
File video_preview.py
|
||
Author charon
|
||
Date 2025/9/6 07:01
|
||
"""
|
||
import os
|
||
import requests
|
||
import urllib.parse
|
||
from uuid import uuid4
|
||
import folder_paths
|
||
|
||
|
||
class VideoDownloaderNode:
|
||
OUTPUT_DIR = folder_paths.get_input_directory()
|
||
|
||
def __init__(self):
|
||
pass
|
||
|
||
@classmethod
|
||
def INPUT_TYPES(s):
|
||
"""
|
||
定义节点的输入参数。
|
||
"""
|
||
return {
|
||
"required": {
|
||
"url": ("STRING", {
|
||
"multiline": False,
|
||
"default": "视频链接"
|
||
}),
|
||
"filename": ("STRING", {
|
||
"multiline": False,
|
||
"default": ""
|
||
}),
|
||
}
|
||
}
|
||
|
||
RETURN_TYPES = ("STRING",)
|
||
RETURN_NAMES = ("file_path",)
|
||
FUNCTION = "download_video"
|
||
CATEGORY = "不忘科技-自定义节点🚩/utils/下载视频"
|
||
|
||
def download_video(self, url, filename=""):
|
||
if not url or not url.strip().startswith('http'):
|
||
print("[VideoDownloader] 无效的URL，跳过下载。")
|
||
return ("",)
|
||
|
||
try:
|
||
print(f"[VideoDownloader] 开始从 {url} 下载...")
|
||
response = requests.get(url, stream=True, timeout=10)
|
||
response.raise_for_status()
|
||
|
||
if not filename.strip():
|
||
try:
|
||
parsed_url = urllib.parse.urlparse(url)
|
||
filename = os.path.basename(parsed_url.path)
|
||
if not filename:
|
||
raise ValueError
|
||
except (ValueError, AttributeError):
|
||
content_type = response.headers.get('content-type')
|
||
ext = '.mp4'
|
||
if content_type and '/' in content_type:
|
||
mime_type = content_type.split('/')[1]
|
||
if len(mime_type) < 5: # 简单的扩展名检查
|
||
ext = '.' + mime_type
|
||
filename = f"downloaded_video_{uuid4().hex[:8]}{ext}"
|
||
|
||
# 清理文件名，防止路径问题
|
||
safe_filename = "".join(c for c in filename if c.isalnum() or c in ('.', '_', '-')).strip()
|
||
if not safe_filename: safe_filename = f"safe_video_{uuid4().hex[:8]}.mp4"
|
||
|
||
file_path = os.path.join(self.OUTPUT_DIR, safe_filename)
|
||
|
||
with open(file_path, 'wb') as f:
|
||
for chunk in response.iter_content(chunk_size=8192):
|
||
f.write(chunk)
|
||
|
||
print(f"[VideoDownloader] 视频已成功下载到: {file_path}")
|
||
|
||
ui_preview = {
|
||
"videos": [{
|
||
"filename": safe_filename,
|
||
"subfolder": "",
|
||
"type": "input"
|
||
}]
|
||
}
|
||
return {"ui": ui_preview, "result": (file_path,)}
|
||
|
||
except requests.exceptions.RequestException as e:
|
||
print(f"[VideoDownloader] 下载视频时出错: {e}")
|
||
return ("",)
|
||
|
||
# NODE_CLASS_MAPPINGS = {
|
||
# "VideoDownloaderNode": VideoDownloaderNode
|
||
# }
|
||
#
|
||
# NODE_DISPLAY_NAME_MAPPINGS = {
|
||
# "VideoDownloaderNode": "视频下载器 (带预览)"
|
||
# } |