if-ai-ComfyUI-yt_dl/__init__.py

from .YouTubeDL import NODE_CLASS_MAPPINGS, _log, _YtDlpLogger
import os
import nodes
from aiohttp import web
from pathlib import Path

js_dir = os.path.join(os.path.dirname(os.path.realpath(__file__)), "js")

nodes.EXTENSION_WEB_DIRS["ComfyUI-yt_dl"] = js_dir

__all__ = ['NODE_CLASS_MAPPINGS']

from server import PromptServer

routes = PromptServer.instance.routes


def _get_output_dir():
    try:
        import folder_paths
        base_output = None
        if hasattr(folder_paths, 'get_output_directory'):
            base_output = folder_paths.get_output_directory()
        elif hasattr(folder_paths, 'get_temp_directory'):
            base_output = folder_paths.get_temp_directory()
        else:
            base_output = os.path.abspath(os.path.join(os.getcwd(), 'output'))
        yt_dir = os.path.join(base_output, 'yt_dl')
        os.makedirs(yt_dir, exist_ok=True)
        return yt_dir
    except Exception:
        return os.path.abspath(os.path.join(os.getcwd(), 'output', 'yt_dl'))


@routes.post('/yt_dl/api/ytdl')
async def ytdl_endpoint(request):
    try:
        import yt_dlp

        data = await request.json()
        url = data.get('url', '')
        mode = data.get('mode', 'video')
        quality = data.get('quality', 'highest')
        if not url:
            return web.json_response({'error': 'Missing url'}, status=400)

        output_dir = _get_output_dir()
        os.makedirs(output_dir, exist_ok=True)

        if mode == 'video':
            if quality == 'highest':
                format_string = 'bestvideo[ext=mp4]+bestaudio[ext=m4a]/best[ext=mp4]/best'
            else:
                # quality arrives as "360p"; the height filter needs a bare number
                # and must constrain every alternative, not just the last one.
                h = quality.rstrip('p')
                format_string = (
                    f"bestvideo[ext=mp4][height<={h}]+bestaudio[ext=m4a]"
                    f"/best[ext=mp4][height<={h}]"
                    f"/best[height<={h}]"
                )
        elif mode == 'audio':
            format_string = 'bestaudio[ext=m4a]/best[ext=m4a]'
        else:
            format_string = None

        outtmpl = os.path.join(output_dir, '%(title)s.%(ext)s')
        ydl_opts = {
            'format': format_string or 'best',
            'outtmpl': outtmpl,
            'writethumbnail': True,
            'quiet': False,
            # Never write to stdout: ComfyUI's stdout can be a closed pipe.
            'logger': _YtDlpLogger(),
            'noprogress': True,
            # No http_headers / extractor_args overrides: yt-dlp sends a
            # User-Agent matching whichever player client it impersonates, and
            # pinning clients or skipping the webpage hides adaptive formats
            # (caps downloads at 360p) or breaks extraction entirely.
            # Add retry options
            'retries': 10,
            'fragment_retries': 10,
            'skip_unavailable_fragments': False,
            # Add sleep requests to avoid rate limiting
            'sleep_interval_requests': 1,
            'sleep_interval': 1,
            # Add cookie support
            'cookiefile': None,  # Will be set if cookies file exists
        }

        # Check if cookies file exists in the output directory
        cookies_file = os.path.join(output_dir, 'cookies.txt')
        if os.path.exists(cookies_file):
            ydl_opts['cookiefile'] = cookies_file

        paths = []

        class CaptureHook(yt_dlp.postprocessor.common.PostProcessor):
            def run(self, information):
                filepath = information.get('filepath') or information.get('requested_downloads', [{}])[0].get('filepath')
                if filepath:
                    paths.append(filepath)
                return [], information

        with yt_dlp.YoutubeDL(ydl_opts) as ydl:
            try:
                ydl.add_post_processor(CaptureHook(ydl))
            except Exception:
                pass
            
            try:
                ydl.download([url])
            except yt_dlp.utils.DownloadError as e:
                # Handle specific download errors
                error_msg = str(e)
                if 'HTTP Error 403' in error_msg:
                    _log('[yt_dl] HTTP 403 error: YouTube is blocking the request. Trying alternative methods...')
                    # Fall back to alternative player clients
                    retry_opts = dict(ydl_opts)
                    retry_opts['extractor_args'] = {'youtube': {'player_client': ['android', 'web']}}
                    with yt_dlp.YoutubeDL(retry_opts) as ydl2:
                        ydl2.add_post_processor(CaptureHook(ydl2))
                        ydl2.download([url])
                else:
                    _log(f'[yt_dl] Download error: {error_msg}')
                    raise
            except Exception as e:
                _log(f'[yt_dl] Unexpected error during download: {str(e)}')
                raise

        file_path = paths[-1] if paths else ''
        return web.json_response({'file_path': file_path})
    except yt_dlp.utils.DownloadError as e:
        return web.json_response({'error': f'Download failed: {str(e)}'}, status=400)
    except Exception as e:
        return web.json_response({'error': f'Unexpected error: {str(e)}'}, status=500)