145 lines
6.0 KiB
Python
145 lines
6.0 KiB
Python
import os
|
|
import yt_dlp
|
|
import subprocess
|
|
import re
|
|
|
|
class YouTubeVideoDownloader:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"youtube_url": ("STRING", {"default": ""}),
|
|
"download_path": ("STRING", {"default": "D:\\ComfyUI_windows_portable\\ComfyUI\\downloads"}),
|
|
"resolution": (["360p", "480p", "720p", "1080p"], {"default": "720p"}),
|
|
"save_as_mp4": ("BOOLEAN", {"default": False}),
|
|
"audio_only": ("BOOLEAN", {"default": False})
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING", "STRING") # Remove the embed video output
|
|
RETURN_NAMES = ("video_path", "audio_path")
|
|
FUNCTION = "download_youtube_video"
|
|
CATEGORY = "🎬🔀MediaMixer"
|
|
|
|
def sanitize_filename(self, title):
|
|
"""Convert title to safe filename"""
|
|
safe_title = re.sub(r'[<>:"/\\|?*]', '', title)
|
|
safe_title = safe_title.replace(' ', '_')
|
|
safe_title = re.sub(r'[^\w\-_.]', '', safe_title)
|
|
return safe_title[:100] # Limit to 100 characters
|
|
|
|
def convert_mkv_to_mp4(self, mkv_path):
|
|
"""Convert MKV to MP4 using FFmpeg"""
|
|
mp4_path = mkv_path.replace('.mkv', '.mp4')
|
|
try:
|
|
cmd = ['ffmpeg', '-i', mkv_path, '-codec', 'copy', mp4_path, '-y']
|
|
subprocess.run(cmd, check=True, capture_output=True)
|
|
if os.path.exists(mp4_path) and os.path.getsize(mp4_path) > 0:
|
|
os.remove(mkv_path)
|
|
return mp4_path
|
|
except subprocess.CalledProcessError as e:
|
|
print(f"[ERROR] FFmpeg conversion failed: {e}")
|
|
return mkv_path
|
|
return mkv_path
|
|
|
|
def download_youtube_video(self, youtube_url, download_path, resolution, save_as_mp4, audio_only):
|
|
try:
|
|
os.makedirs(download_path, exist_ok=True)
|
|
print(f"[INFO] Download path verified or created at: {download_path}")
|
|
|
|
print("[INFO] Getting video information...")
|
|
with yt_dlp.YoutubeDL({'quiet': True}) as ydl:
|
|
info = ydl.extract_info(youtube_url, download=False)
|
|
video_title = self.sanitize_filename(info['title'])
|
|
print(f"[INFO] Video title: {info['title']}")
|
|
|
|
base_filename = video_title
|
|
counter = 1
|
|
while os.path.exists(os.path.join(download_path, f"{base_filename}.mkv")) or \
|
|
os.path.exists(os.path.join(download_path, f"{base_filename}.mp4")) or \
|
|
os.path.exists(os.path.join(download_path, f"{base_filename}.wav")):
|
|
base_filename = f"{video_title}_{counter}"
|
|
counter += 1
|
|
|
|
video_path = os.path.join(download_path, f"{base_filename}.mkv")
|
|
audio_path = os.path.join(download_path, f"{base_filename}.wav")
|
|
final_video_path = video_path
|
|
|
|
print("[INFO] Downloading and extracting audio...")
|
|
ydl_opts_audio = {
|
|
'format': 'bestaudio',
|
|
'outtmpl': audio_path[:-4],
|
|
'postprocessors': [{
|
|
'key': 'FFmpegExtractAudio',
|
|
'preferredcodec': 'wav',
|
|
'preferredquality': '192',
|
|
}],
|
|
'quiet': False,
|
|
'no_warnings': False,
|
|
}
|
|
|
|
with yt_dlp.YoutubeDL(ydl_opts_audio) as ydl:
|
|
ydl.extract_info(youtube_url, download=True)
|
|
|
|
if not os.path.exists(audio_path):
|
|
raise FileNotFoundError(f"Audio file not found at {audio_path}")
|
|
if os.path.getsize(audio_path) == 0:
|
|
raise ValueError(f"Downloaded audio file is empty: {audio_path}")
|
|
|
|
if audio_only:
|
|
print("[INFO] Audio-only mode: Skipping video download")
|
|
return ("", audio_path)
|
|
|
|
print("[INFO] Downloading video...")
|
|
height = resolution[:-1]
|
|
|
|
ydl_opts_video = {
|
|
'format': f'bestvideo[height<={height}][ext=mp4]+bestaudio[ext=m4a]/best[ext=mp4]/best',
|
|
'outtmpl': video_path,
|
|
'merge_output_format': 'mkv',
|
|
'postprocessors': [{
|
|
'key': 'FFmpegVideoConvertor',
|
|
'preferedformat': 'mkv',
|
|
}],
|
|
'quiet': False,
|
|
'no_warnings': False,
|
|
}
|
|
|
|
with yt_dlp.YoutubeDL(ydl_opts_video) as ydl:
|
|
ydl.extract_info(youtube_url, download=True)
|
|
print("[INFO] Video download completed")
|
|
|
|
if save_as_mp4:
|
|
print("[INFO] Converting to MP4...")
|
|
final_video_path = self.convert_mkv_to_mp4(video_path)
|
|
if final_video_path.endswith('.mp4'):
|
|
print("[INFO] Successfully converted to MP4")
|
|
else:
|
|
print("[WARNING] MP4 conversion failed, keeping MKV format")
|
|
|
|
if not os.path.exists(final_video_path):
|
|
raise FileNotFoundError(f"Video file not found at {final_video_path}")
|
|
if os.path.getsize(final_video_path) == 0:
|
|
raise ValueError(f"Downloaded video file is empty: {final_video_path}")
|
|
|
|
print(f"[INFO] Files saved to:\nVideo: {final_video_path}\nAudio: {audio_path}")
|
|
|
|
return (final_video_path, audio_path) # Only return video and audio paths
|
|
|
|
except Exception as e:
|
|
error_message = f"[ERROR] Error during download: {str(e)}"
|
|
print(error_message)
|
|
return (error_message, error_message)
|
|
|
|
# Register node class mappings for ComfyUI
|
|
NODE_CLASS_MAPPINGS = {
|
|
"YouTubeVideoDownloader": YouTubeVideoDownloader
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"YouTubeVideoDownloader": "YouTube Video Downloader"
|
|
}
|