File size: 3,020 Bytes
b613c3c
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
246c64e
b613c3c
 
 
 
246c64e
b613c3c
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
246c64e
 
 
 
b613c3c
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
246c64e
b613c3c
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
"""
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)}"