youtube.py

"""Download a single YouTube video's audio for the separation UI."""

from html import escape
import os
from pathlib import Path
import tempfile
from urllib.parse import urlparse

import gradio as gr
import yt_dlp


FORMATS = {
    "Best available": "bestaudio/best",
    "M4A (if available)": "bestaudio[ext=m4a]/bestaudio/best",
    "Opus/WebM (if available)": "bestaudio[ext=webm]/bestaudio/best",
}


def _validate_url(url: str) -> str:
    url = (url or "").strip()
    parsed = urlparse(url)
    host = (parsed.hostname or "").lower()
    if parsed.scheme not in ("http", "https") or host not in {
        "youtube.com", "www.youtube.com", "m.youtube.com",
        "music.youtube.com", "youtu.be", "www.youtu.be",
    }:
        raise gr.Error("Enter a valid YouTube video URL.")
    return url


def _options(audio_format: str = "Best available") -> dict:
    if audio_format not in FORMATS:
        raise gr.Error("Select a valid audio format.")
    options = {
        "format": FORMATS[audio_format],
        "noplaylist": True,
        "max_filesize": 300 * 1024 * 1024,
        "source_address": "0.0.0.0",
        "quiet": True,
        "no_warnings": True,
    }
    proxy = os.environ.get("YTDLP_PROXY", "").strip()
    if proxy:
        parsed_proxy = urlparse(proxy)
        if parsed_proxy.scheme not in ("http", "https", "socks5", "socks5h") or not parsed_proxy.hostname:
            raise gr.Error("YTDLP_PROXY must be an HTTP(S) or SOCKS5 proxy URL.")
        options["proxy"] = proxy
    return options


def _download_error(exc: Exception, action: str) -> gr.Error:
    if "Sign in to confirm you’re not a bot" in str(exc) or "Sign in to confirm you're not a bot" in str(exc):
        return gr.Error(
            "YouTube is requiring verification from this Space's network. "
            "Use a trusted proxy configured as the YTDLP_PROXY Space Secret, "
            "or upload an audio file instead."
        )
    return gr.Error(f"Could not {action} the YouTube video. Check the URL and try again.")


def _extract(url: str, options: dict, download: bool):
    """Retry with a YouTube client that does not require a PO Token."""
    last_error = None
    for client in (None, "android_vr"):
        attempt = options.copy()
        if client:
            attempt["extractor_args"] = {"youtube": {"player_client": [client]}}
        try:
            with yt_dlp.YoutubeDL(attempt) as ydl:
                info = ydl.extract_info(url, download=download)
                return info, ydl
        except yt_dlp.utils.DownloadError as exc:
            last_error = exc
    raise last_error


def inspect_youtube(url: str) -> str:
    """Show video details without downloading the media."""
    url = _validate_url(url)
    try:
        info, _ = _extract(url, _options(), download=False)
    except yt_dlp.utils.DownloadError as exc:
        raise _download_error(exc, "inspect") from exc
    if not info or info.get("_type") == "playlist":
        raise gr.Error("Select a single video from the playlist.")
    title = escape(str(info.get("title") or "Untitled"))
    channel = escape(str(info.get("channel") or info.get("uploader") or "Unknown"))
    duration = info.get("duration")
    length = f"{int(duration) // 60}:{int(duration) % 60:02d}" if duration else "Unknown"
    return f"**{title}** · {channel} · {length}"


def youtube(url: str, audio_format: str = "Best available") -> str:
    """Keep the selected source audio without an unnecessary MP3 conversion."""
    url = _validate_url(url)
    directory = Path(tempfile.mkdtemp(prefix="vocal-separation-yt-"))
    options = _options(audio_format)
    options["outtmpl"] = str(directory / f"{directory.name}.%(ext)s")
    try:
        info, ydl = _extract(url, options, download=True)
        if not info or info.get("_type") == "playlist":
            raise gr.Error("Select a single video from the playlist.")
        downloads = info.get("requested_downloads") or []
        candidate = downloads[0].get("filepath") if downloads else None
        path = Path(candidate or ydl.prepare_filename(info))
        if not path.is_file():
            raise gr.Error("The downloaded audio file was not found.")
        return str(path)
    except yt_dlp.utils.DownloadError as exc:
        raise _download_error(exc, "download") from exc