Spaces:
Running
Running
| """ | |
| YouTube downloader for Video Model Studio. | |
| Handles downloading videos from YouTube URLs. | |
| """ | |
| import logging | |
| import gradio as gr | |
| from pathlib import Path | |
| from typing import Optional, Any, Union, Callable | |
| from pytubefix import YouTube | |
| from vms.config import VIDEOS_TO_SPLIT_PATH | |
| logger = logging.getLogger(__name__) | |
| class YouTubeDownloader: | |
| """Handles downloading videos from YouTube""" | |
| def download_video(self, url: str, enable_splitting: bool, progress: Optional[Callable] = None) -> str: | |
| """Download a video from YouTube | |
| Args: | |
| url: YouTube video URL | |
| enable_splitting: Whether to enable automatic video splitting | |
| progress: Optional Gradio progress indicator | |
| Returns: | |
| Status message string | |
| """ | |
| if not url or not url.strip(): | |
| logger.warning("No YouTube URL provided") | |
| return "Please enter a YouTube URL" | |
| try: | |
| logger.info(f"Downloading YouTube video: {url}") | |
| # Extract video ID and create YouTube object | |
| yt = YouTube(url, on_progress_callback=lambda stream, chunk, bytes_remaining: | |
| progress((1 - bytes_remaining / stream.filesize), desc="Downloading...") | |
| if progress else None) | |
| video_id = yt.video_id | |
| # Choose target directory based on auto-splitting setting | |
| target_dir = VIDEOS_TO_SPLIT_PATH if enable_splitting else STAGING_PATH | |
| output_path = target_dir / f"{video_id}.mp4" | |
| # Download highest quality progressive MP4 | |
| if progress: | |
| logger.debug("Getting video streams...") | |
| progress(0, desc="Getting video streams...") | |
| video = yt.streams.filter(progressive=True, file_extension='mp4').order_by('resolution').desc().first() | |
| if not video: | |
| logger.error("Could not find a compatible video format") | |
| gr.Error("Could not find a compatible video format") | |
| return "Could not find a compatible video format" | |
| # Download the video | |
| if progress: | |
| logger.info("Starting YouTube video download...") | |
| progress(0, desc="Starting download...") | |
| video.download(output_path=str(target_dir), filename=f"{video_id}.mp4") | |
| # Update UI | |
| if progress: | |
| logger.info("YouTube video download complete!") | |
| gr.Info("YouTube video download complete!") | |
| progress(1, desc="Download complete!") | |
| return f"Successfully downloaded video: {yt.title}" | |
| except Exception as e: | |
| logger.error(f"Error downloading YouTube video: {str(e)}", exc_info=True) | |
| gr.Error(f"Error downloading video: {str(e)}") | |
| return f"Error downloading video: {str(e)}" |