Update app.py

app.py

@@ -6,10 +6,23 @@ import requests
 from concurrent.futures import ThreadPoolExecutor
 import time
 from yt_dlp import YoutubeDL
+from yt_dlp.utils import DownloadError
 import subprocess
 import shutil
 from typing import List, Tuple
 import pandas as pd
+import asyncio
+import aiohttp
+import aiofiles
+import json
+from functools import lru_cache
+import tenacity
+from tqdm import tqdm
+from pathlib import Path
+import logging
+
+logging.basicConfig(level=logging.DEBUG)
+logger = logging.getLogger(__name__)
 
 def sanitize_title(title):
     return re.sub(r'[\\/*?:"<>|]', "", title)
@@ -17,6 +30,13 @@ def sanitize_title(title):
 def format_time(seconds):
     return time.strftime('%H:%M:%S', time.gmtime(seconds))
 
+@tenacity.retry(
+    stop=tenacity.stop_after_attempt(3),
+    wait=tenacity.wait_exponential(multiplier=1, min=4, max=10),
+    retry=tenacity.retry_if_exception_type((requests.RequestException, aiohttp.ClientError)),
+    reraise=True
+)
+@lru_cache(maxsize=100)
 def get_video_info(video_url):
     with YoutubeDL({'quiet': True, 'no_warnings': True}) as ydl:
         try:
@@ -43,62 +63,125 @@ def get_video_info(video_url):
                     return info['title'], best_format['url'], None
                 else:
                     raise Exception("No suitable video formats found")
-        except Exception as e:
+        except DownloadError as e:
             raise Exception(f"Error extracting video info: {str(e)}")
+        except Exception as e:
+            raise Exception(f"Unexpected error: {str(e)}")
 
-def download_segment(url, start_time, end_time, output_path):
+@tenacity.retry(
+    stop=tenacity.stop_after_attempt(3),
+    wait=tenacity.wait_exponential(multiplier=1, min=4, max=10),
+    retry=tenacity.retry_if_exception_type((subprocess.CalledProcessError, asyncio.TimeoutError)),
+    reraise=True
+)
+async def download_segment_async(url, start_time, end_time, output_path):
+    output_path = Path(output_path)
     command = [
         'ffmpeg',
         '-ss', format_time(start_time),
         '-i', url,
         '-t', format_time(end_time - start_time),
         '-c', 'copy',
+        '-avoid_negative_ts', 'make_zero',
         '-y',
-        output_path
+        str(output_path)
     ]
-    process = subprocess.Popen(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, universal_newlines=True)
-    
-    while True:
-        output = process.stderr.readline()
-        if output == '' and process.poll() is not None:
-            break
-        if output:
-            yield output.strip()
+    logger.debug(f"Executing command: {' '.join(command)}")
+    process = await asyncio.create_subprocess_exec(
+        *command,
+        stdout=asyncio.subprocess.PIPE,
+        stderr=asyncio.subprocess.PIPE
+    )
     
-    rc = process.poll()
-    return rc == 0
-
-def combine_segments(video_segments, audio_segments, output_path):
-    temp_video = 'temp_video.mp4'
-    temp_audio = 'temp_audio.m4a'
+    stdout, stderr = await process.communicate()
     
-    # Concatenate video segments
-    with open('video_list.txt', 'w') as f:
-        for segment in video_segments:
-            f.write(f"file '{segment}'\n")
+    if process.returncode != 0:
+        error_message = f"FFmpeg error: {stderr.decode()}"
+        logger.error(error_message)
+        yield error_message
+        raise Exception(error_message)
     
-    subprocess.run(['ffmpeg', '-f', 'concat', '-safe', '0', '-i', 'video_list.txt', '-c', 'copy', temp_video])
+    if not output_path.exists():
+        error_message = f"Output file not created: {output_path}"
+        logger.error(error_message)
+        yield error_message
+        raise FileNotFoundError(error_message)
     
-    # Concatenate audio segments if they exist
-    if audio_segments:
-        with open('audio_list.txt', 'w') as f:
-            for segment in audio_segments:
-                f.write(f"file '{segment}'\n")
-        subprocess.run(['ffmpeg', '-f', 'concat', '-safe', '0', '-i', 'audio_list.txt', '-c', 'copy', temp_audio])
+    logger.info(f"Successfully downloaded segment to {output_path}")
+    yield f"Successfully downloaded segment to {output_path}"
+
+async def combine_segments_async(video_segments, audio_segments, output_path):
+    if not audio_segments:
+        raise Exception("No audio segments to combine")
+
+    video_list = None
+    audio_list = None
+    temp_video = None
+    temp_audio = None
+
+    try:
+        if video_segments:
+            temp_video = 'temp_video.mp4'
+            temp_audio = 'temp_audio.m4a'
+            
+            # Concatenate video segments
+            video_list = 'video_list.txt'
+            async with aiofiles.open(video_list, 'w') as f:
+                await f.write('\n'.join(f"file '{segment}'" for segment in video_segments))
+            
+            process = await asyncio.create_subprocess_exec(
+                'ffmpeg', '-f', 'concat', '-safe', '0', '-i', video_list, '-c', 'copy', temp_video,
+                stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE
+            )
+            stdout, stderr = await process.communicate()
+            if process.returncode != 0:
+                raise Exception(f"Error concatenating video segments: {stderr.decode()}")
+            
+            # Concatenate audio segments
+            audio_list = 'audio_list.txt'
+            async with aiofiles.open(audio_list, 'w') as f:
+                await f.write('\n'.join(f"file '{segment}'" for segment in audio_segments))
+            
+            process = await asyncio.create_subprocess_exec(
+                'ffmpeg', '-f', 'concat', '-safe', '0', '-i', audio_list, '-c', 'copy', temp_audio,
+                stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE
+            )
+            stdout, stderr = await process.communicate()
+            if process.returncode != 0:
+                raise Exception(f"Error concatenating audio segments: {stderr.decode()}")
+            
+            # Combine video and audio
+            process = await asyncio.create_subprocess_exec(
+                'ffmpeg', '-i', temp_video, '-i', temp_audio, '-c', 'copy', output_path,
+                stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE
+            )
+            stdout, stderr = await process.communicate()
+            if process.returncode != 0:
+                raise Exception(f"Error combining video and audio: {stderr.decode()}")
+        else:
+            # Audio only
+            audio_list = 'audio_list.txt'
+            async with aiofiles.open(audio_list, 'w') as f:
+                await f.write('\n'.join(f"file '{segment}'" for segment in audio_segments))
+            
+            process = await asyncio.create_subprocess_exec(
+                'ffmpeg', '-f', 'concat', '-safe', '0', '-i', audio_list, '-c', 'copy', output_path,
+                stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE
+            )
+            stdout, stderr = await process.communicate()
+            if process.returncode != 0:
+                raise Exception(f"Error combining audio segments: {stderr.decode()}")
         
-        # Combine video and audio
-        subprocess.run(['ffmpeg', '-i', temp_video, '-i', temp_audio, '-c', 'copy', output_path])
-    else:
-        shutil.move(temp_video, output_path)
-    
-    # Clean up temporary files
-    os.remove('video_list.txt')
-    if os.path.exists('audio_list.txt'):
-        os.remove('audio_list.txt')
-    if os.path.exists(temp_video):
-        os.remove(temp_video)
-    if os.path.exists(temp_audio):
-        os.remove(temp_audio)
+        return output_path
+    except Exception as e:
+        raise Exception(f"Error in combine_segments_async: {str(e)}")
+    finally:
+        # Clean up temporary files
+        for file in [f for f in [video_list, audio_list, temp_video, temp_audio] if f]:
+            try:
+                Path(file).unlink(missing_ok=True)
+            except Exception as e:
+                print(f"Error removing temporary file {file}: {str(e)}")
 
 def add_segment(start_hours, start_minutes, start_seconds, end_hours, end_minutes, end_seconds, segments):
     start_time = f"{start_hours:02d}:{start_minutes:02d}:{start_seconds:02d}"
@@ -132,116 +215,225 @@ def parse_segments(segments: pd.DataFrame) -> List[Tuple[int, int]]:
             continue  # Skip invalid segments
     return parsed_segments
 
-def process_video(video_url, segments, combine, progress=gr.Progress()):
+async def process_video(video_url, segments, combine, audio_only, progress=gr.Progress()):
     if not video_url.strip():
-        return 0, "Error: Please provide a valid YouTube URL", None
+        yield 0, "Error: Please provide a valid YouTube URL", None
+        return
 
-    # Extract segments from the Dataframe
     parsed_segments = parse_segments(segments)
     if not parsed_segments:
-        return 0, "Error: No valid segments provided", None
+        yield 0, "Error: No valid segments provided", None
+        return
 
-    output_dir = 'output'
-    os.makedirs(output_dir, exist_ok=True)
+    output_dir = Path('output')
+    output_dir.mkdir(exist_ok=True)
 
     try:
+        progress(0, "Extracting video info...")
         video_title, video_url, audio_url = get_video_info(video_url)
     except Exception as e:
-        return 0, f"Error: {str(e)}", None
+        yield 0, f"Error: {str(e)}", None
+        return
 
     video_segments = []
     audio_segments = []
     total_segments = len(parsed_segments)
 
     for i, (start_time, end_time) in enumerate(parsed_segments):
-        video_output = os.path.join(output_dir, f"{sanitize_title(video_title)}_video_segment_{i+1}.mp4")
-        for output in download_segment(video_url, start_time, end_time, video_output):
-            progress((i / total_segments) + (1 / total_segments) * 0.5)
-            yield i * 100 // total_segments, f"Downloading video segment {i+1}/{total_segments}: {output}", None
-        video_segments.append(video_output)
-
-        if audio_url:
-            audio_output = os.path.join(output_dir, f"{sanitize_title(video_title)}_audio_segment_{i+1}.m4a")
-            for output in download_segment(audio_url, start_time, end_time, audio_output):
-                progress((i / total_segments) + (1 / total_segments) * 0.75)
-                yield i * 100 // total_segments + 50, f"Downloading audio segment {i+1}/{total_segments}: {output}", None
-            audio_segments.append(audio_output)
-
-    if combine:
-        output_path = os.path.join(output_dir, f"{sanitize_title(video_title)}_combined.mp4")
-        combine_segments(video_segments, audio_segments, output_path)
-        yield 100, f"Segments combined and saved as {output_path}", output_path
-    else:
-        # If not combining, return the first video segment (you might want to modify this behavior)
-        output_path = video_segments[0] if video_segments else None
-        yield 100, "All segments downloaded successfully", output_path
-
-    # Clean up individual segments if combined
-    if combine:
-        for segment in video_segments + audio_segments:
-            os.remove(segment)
+        progress((i / total_segments) * 0.8, f"Downloading segment {i+1}/{total_segments}")
+        
+        try:
+            if not audio_only:
+                video_output = output_dir / f"{sanitize_title(video_title)}_video_segment_{i+1}.mp4"
+                async for output in download_segment_async(video_url, start_time, end_time, str(video_output)):
+                    progress((i / total_segments) * 0.8 + (1 / total_segments) * 0.4, f"Downloading video segment {i+1}/{total_segments}: {output}")
+                if not video_output.exists():
+                    raise FileNotFoundError(f"Video segment file not found: {video_output}")
+                video_segments.append(str(video_output))
+
+            audio_output = output_dir / f"{sanitize_title(video_title)}_audio_segment_{i+1}.m4a"
+            async for output in download_segment_async(audio_url or video_url, start_time, end_time, str(audio_output)):
+                progress((i / total_segments) * 0.8 + (1 / total_segments) * 0.8, f"Downloading audio segment {i+1}/{total_segments}: {output}")
+            if not audio_output.exists():
+                raise FileNotFoundError(f"Audio segment file not found: {audio_output}")
+            audio_segments.append(str(audio_output))
+        
+        except Exception as e:
+            yield (i / total_segments) * 100, f"Error downloading segment {i+1}: {str(e)}", None
+            return
+
+    try:
+        if combine:
+            progress(90, "Combining segments...")
+            if audio_only:
+                output_path = output_dir / f"{sanitize_title(video_title)}_combined.m4a"
+            else:
+                output_path = output_dir / f"{sanitize_title(video_title)}_combined.mp4"
+            
+            output_path = await combine_segments_async(video_segments if not audio_only else [], audio_segments, str(output_path))
+            if Path(output_path).exists():
+                yield 100, f"Segments combined and saved as {output_path}", output_path
+            else:
+                raise FileNotFoundError(f"Combined output file not found: {output_path}")
+        else:
+            output_path = audio_segments[0] if audio_only else video_segments[0]
+            if Path(output_path).exists():
+                yield 100, "All segments downloaded successfully", output_path
+            else:
+                raise FileNotFoundError(f"Output file not found: {output_path}")
+
+    except Exception as e:
+        yield 100, f"Error: {str(e)}", None
+    finally:
+        # Clean up individual segment files if they were combined
+        if combine:
+            for segment in video_segments + audio_segments:
+                try:
+                    Path(segment).unlink(missing_ok=True)
+                except Exception as e:
+                    print(f"Error removing segment file {segment}: {str(e)}")
+
+def get_video_qualities(video_url):
+    if not video_url.strip():
+        return []
+    
+    with YoutubeDL({'quiet': True, 'no_warnings': True}) as ydl:
+        try:
+            info = ydl.extract_info(video_url, download=False)
+            formats = info.get('formats', [])
+            
+            qualities = set()
+            for f in formats:
+                if f.get('vcodec') != 'none' and f.get('acodec') != 'none':
+                    height = f.get('height')
+                    if height:
+                        qualities.add(f"{height}p")
+            
+            return sorted(list(qualities), key=lambda x: int(x[:-1]), reverse=True)
+        except DownloadError:
+            return []
+        except Exception as e:
+            print(f"Error in get_video_qualities: {str(e)}")
+            return []
 
 # Disable Gradio analytics
 utils.colab_check = lambda: True
 
-with gr.Blocks(title="YouTube Segment Downloader", theme=gr.themes.Soft(primary_hue="indigo", secondary_hue="blue")) as iface:
-    gr.Markdown(
-        """
-        # YouTube Segment Downloader
-        Download specific segments of YouTube videos with ease.
-        """
-    )
-    
-    with gr.Row():
-        video_url = gr.Textbox(label="YouTube URL", placeholder="Enter YouTube URL here")
-    
-    with gr.Row():
-        with gr.Column(scale=1):
-            with gr.Group():
-                gr.Markdown("### Add Segment")
-                with gr.Row():
-                    start_hours = gr.Number(label="Start", minimum=0, maximum=23, step=1, value=0)
-                    start_minutes = gr.Number(label="", minimum=0, maximum=59, step=1, value=0)
-                    start_seconds = gr.Number(label="", minimum=0, maximum=59, step=1, value=0)
+custom_css = """
+    :root {
+        --primary-hue: 210;
+        --secondary-hue: 160;
+        --neutral-hue: 0;
+        --primary-color: hsl(var(--primary-hue), 70%, 55%);
+        --secondary-color: hsl(var(--secondary-hue), 70%, 55%);
+        --neutral-color: hsl(var(--neutral-hue), 0%, 95%);
+        --text-color: hsl(var(--neutral-hue), 0%, 20%);
+        --border-radius: 8px;
+    }
+    body { font-family: 'Inter', sans-serif; margin: 0; padding: 0; height: 100vh; overflow: hidden; }
+    #component-0 { height: 100vh; overflow-y: auto; padding: 1rem; box-sizing: border-box; }
+    .container { max-width: 100%; margin: 0 auto; display: flex; flex-direction: column; height: 100%; }
+    .header { text-align: center; margin-bottom: 1rem; }
+    .main-content { display: flex; flex-wrap: wrap; gap: 1rem; flex-grow: 1; overflow-y: auto; }
+    .input-section, .segment-controls, .options-section, .output-section {
+        background-color: var(--neutral-color);
+        border-radius: var(--border-radius);
+        padding: 1rem;
+        box-shadow: 0 2px 4px rgba(0,0,0,0.1);
+        flex: 1 1 300px;
+    }
+    .segment-input, .time-input, .button-row {
+        display: flex;
+        flex-wrap: wrap;
+        gap: 0.5rem;
+        align-items: center;
+    }
+    .time-input { width: 60px !important; text-align: center; }
+    .time-label, .separator { font-weight: bold; color: var(--primary-color); margin-bottom: 0.5rem; }
+    .button-row { justify-content: space-between; }
+    @media (max-width: 768px) {
+        .main-content { flex-direction: column; }
+        .input-section, .segment-controls, .options-section, .output-section { flex: 1 1 100%; }
+    }
+    .url-input { flex-grow: 1; }
+    .quality-dropdown { width: 100px !important; }
+"""
+
+with gr.Blocks(
+    title="YouTube Segment Downloader",
+    theme=gr.themes.Soft(
+        primary_hue="blue",
+        secondary_hue="green",
+        neutral_hue="gray"
+    ).set(
+        body_text_color="*neutral_800",
+        background_fill_primary="*neutral_50",
+        button_primary_background_fill="*primary_500",
+        button_primary_background_fill_hover="*primary_600",
+        button_secondary_background_fill="*secondary_500",
+        button_secondary_background_fill_hover="*secondary_600",
+    ),
+    css=custom_css
+) as iface:
+    with gr.Column(elem_classes="container"):
+        gr.Markdown("# 🎬 YouTube Segment Downloader", elem_classes="header")
+        
+        with gr.Row(elem_classes="main-content"):
+            with gr.Column(elem_classes="input-section"):
                 with gr.Row():
-                    end_hours = gr.Number(label="End", minimum=0, maximum=23, step=1, value=0)
-                    end_minutes = gr.Number(label="", minimum=0, maximum=59, step=1, value=0)
-                    end_seconds = gr.Number(label="", minimum=0, maximum=59, step=1, value=0)
-                add_btn = gr.Button("Add Segment", variant="primary")
+                    video_url = gr.Textbox(
+                        label="YouTube URL",
+                        placeholder="Paste URL here",
+                        elem_classes="url-input"
+                    )
+                    quality = gr.Dropdown(
+                        label="Quality",
+                        choices=[],
+                        interactive=True,
+                        elem_classes="quality-dropdown",
+                        visible=False
+                    )
+                url_status = gr.Markdown(visible=False)
+                
+                with gr.Column(elem_classes="segment-controls"):
+                    gr.Markdown("### Add Segments", elem_classes="time-label")
+                    with gr.Row(elem_classes="segment-input"):
+                        start_hours = gr.Number(label="Start HH", minimum=0, maximum=23, step=1, value=0, elem_classes="time-input")
+                        start_minutes = gr.Number(label="MM", minimum=0, maximum=59, step=1, value=0, elem_classes="time-input")
+                        start_seconds = gr.Number(label="SS", minimum=0, maximum=59, step=1, value=0, elem_classes="time-input")
+                        
+                        gr.Markdown("to", elem_classes="separator")
+                        
+                        end_hours = gr.Number(label="End HH", minimum=0, maximum=23, step=1, value=0, elem_classes="time-input")
+                        end_minutes = gr.Number(label="MM", minimum=0, maximum=59, step=1, value=0, elem_classes="time-input")
+                        end_seconds = gr.Number(label="SS", minimum=0, maximum=59, step=1, value=0, elem_classes="time_input")
+                    
+                    add_btn = gr.Button("Add Segment", variant="primary")
+                
+                segments = gr.Dataframe(
+                    headers=["Segment"],
+                    row_count=3,
+                    col_count=1,
+                    datatype=["str"],
+                    interactive=True,
+                    label="Segments"
+                )
+            
+            with gr.Column(elem_classes="options-section"):
+                combine = gr.Checkbox(label="Combine segments", value=True)
+                audio_only = gr.Checkbox(label="Audio only", value=False)
+                remove_index = gr.Number(label="Remove segment", minimum=0, step=1, value=0)
+                remove_btn = gr.Button("Remove", variant="secondary")
+                old_index = gr.Number(label="Move from", minimum=0, step=1, value=0)
+                new_index = gr.Number(label="to", minimum=0, step=1, value=0)
+                move_btn = gr.Button("Move", variant="secondary")
+                
+                submit_btn = gr.Button("🚀 Download", variant="primary")
         
-        with gr.Column(scale=2):
-            segments = gr.Dataframe(
-                headers=["Segment"],
-                row_count=5,
-                col_count=1,
-                datatype=["str"],
-                interactive=True,
-                label="Segments to Download"
-            )
-    
-    with gr.Row():
-        with gr.Column(scale=1):
-            combine = gr.Checkbox(label="Combine Segments", value=True)
-        with gr.Column(scale=2):
-            submit_btn = gr.Button("Download Segments", variant="primary", size="lg")
-    
-    with gr.Row():
-        progress = gr.Slider(label="Progress", minimum=0, maximum=100, step=1, interactive=False)
-    
-    with gr.Row():
-        status = gr.Textbox(label="Status", lines=3)
-    
-    with gr.Row():
-        output_file = gr.File(label="Downloaded Video")
-
-    with gr.Accordion("Advanced Options", open=False):
-        with gr.Row():
-            remove_index = gr.Number(label="Remove Segment Index", minimum=0, step=1, value=0)
-            remove_btn = gr.Button("Remove Segment", variant="secondary")
-        with gr.Row():
-            old_index = gr.Number(label="Move From Index", minimum=0, step=1, value=0)
-            new_index = gr.Number(label="Move To Index", minimum=0, step=1, value=0)
-            move_btn = gr.Button("Move Segment", variant="secondary")
+        with gr.Column(elem_classes="output-section"):
+            progress = gr.Slider(label="Progress", minimum=0, maximum=100, step=1, interactive=False)
+            status = gr.Textbox(label="Status", lines=1)
+            output_file = gr.File(label="Downloaded File")
 
     add_btn.click(
         add_segment,
@@ -251,7 +443,7 @@ with gr.Blocks(title="YouTube Segment Downloader", theme=gr.themes.Soft(primary_
 
     submit_btn.click(
         process_video,
-        inputs=[video_url, segments, combine],
+        inputs=[video_url, segments, combine, audio_only],
         outputs=[progress, status, output_file]
     )
 
@@ -267,4 +459,28 @@ with gr.Blocks(title="YouTube Segment Downloader", theme=gr.themes.Soft(primary_
         outputs=[segments]
     )
 
-iface.launch()
+    def update_qualities(url):
+        qualities = get_video_qualities(url)
+        if qualities:
+            return (
+                gr.Dropdown(choices=qualities, value=qualities[0], visible=True),
+                gr.Markdown(visible=False)
+            )
+        elif url.strip():  # Only show error if URL is not empty
+            return (
+                gr.Dropdown(choices=[], visible=False),
+                gr.Markdown("Unable to fetch video qualities. The URL might be invalid or the video might be unavailable.", visible=True)
+            )
+        else:
+            return (
+                gr.Dropdown(choices=[], visible=False),
+                gr.Markdown(visible=False)
+            )
+
+    video_url.change(
+        update_qualities,
+        inputs=[video_url],
+        outputs=[quality, url_status]
+    )
+
+iface.launch(server_name="0.0.0.0", server_port=7860)