Update app.py

app.py

@@ -1,7 +1,6 @@
 import streamlit as st
-from pytube import YouTube
+import yt_dlp
 import os
-import time
 
 # Streamlit 제목
 st.title('YouTube Video Downloader')
@@ -12,12 +11,18 @@ youtube_url = st.text_input('Enter YouTube video URL')
 if 'confirmed' not in st.session_state:
     st.session_state.confirmed = False
 
+def fetch_video_info(youtube_url):
+    ydl_opts = {}
+    with yt_dlp.YoutubeDL(ydl_opts) as ydl:
+        info_dict = ydl.extract_info(youtube_url, download=False)
+        return info_dict
+
 if st.button('확인'):
     if youtube_url:
         try:
-            yt = YouTube(youtube_url)
-            st.session_state.yt = yt
-            st.session_state.thumbnail_url = yt.thumbnail_url
+            info_dict = fetch_video_info(youtube_url)
+            st.session_state.info_dict = info_dict
+            st.session_state.thumbnail_url = info_dict.get('thumbnail')
             st.session_state.confirmed = True
         except Exception as e:
             st.error(f"An error occurred: {e}")
@@ -28,45 +33,33 @@ if st.session_state.confirmed:
     st.image(st.session_state.thumbnail_url, caption='Video Thumbnail', use_column_width=True)
 
     # 해상도 선택
-    resolution = st.selectbox('Select resolution', ['1080p', '1080p60', '1080pHD', '720p'])
-
-    def fetch_video_stream(youtube_url, resolution):
-        yt = st.session_state.yt
-
-        # 딜레이 추가
-        time.sleep(2)  # 2초 딜레이
-
-        # 선택된 해상도의 비디오 스트림 선택
-        stream = yt.streams.filter(res=resolution.replace('pHD', 'p'), progressive=True, file_extension='mp4').first()
+    resolutions = ['1080p', '1080p60', '1080pHD', '720p']
+    resolution = st.selectbox('Select resolution', resolutions)
 
-        # 1080p 변형들이 존재하지 않으면 720p 스트림 선택
-        if '1080p' in resolution and stream is None:
-            st.warning(f"{resolution}가 존재하지 않습니다. 720p를 다운로드합니다.")
-            stream = yt.streams.filter(res='720p', progressive=True, file_extension='mp4').first()
-
-        return yt, stream
+    def download_video(youtube_url, resolution):
+        ydl_opts = {
+            'format': f'bestvideo[height={resolution[:-1]}]+bestaudio/best[height={resolution[:-1]}]',
+            'outtmpl': 'downloaded_video.%(ext)s'
+        }
+        with yt_dlp.YoutubeDL(ydl_opts) as ydl:
+            ydl.download([youtube_url])
 
     # "추출" 버튼 추가
     if st.button('추출'):
-        yt, stream = fetch_video_stream(youtube_url, resolution)
-
-        if stream is None:
-            st.error(f"{resolution} 해상도의 스트림을 찾을 수 없습니다.")
-        else:
-            try:
-                # 다운로드
-                download_path = stream.download()
-                st.session_state.download_path = download_path
-                st.success('Video downloaded successfully!')
-                st.video(download_path)
+        try:
+            if resolution == '1080pHD':
+                resolution = '1080p'  # yt-dlp에서 '1080pHD'를 '1080p'로 매핑
+            download_video(youtube_url, resolution)
+            st.success('Video downloaded successfully!')
+            st.video('downloaded_video.mp4')
 
-                # 다운로드된 파일 제공 (optional)
-                with open(download_path, 'rb') as file:
-                    st.download_button(
-                        label="Download Video File",
-                        data=file,
-                        file_name=os.path.basename(download_path),
-                        mime='video/mp4'
-                    )
-            except Exception as e:
-                st.error(f"An error occurred during download: {e}")
+            # 다운로드된 파일 제공 (optional)
+            with open('downloaded_video.mp4', 'rb') as file:
+                st.download_button(
+                    label="Download Video File",
+                    data=file,
+                    file_name='downloaded_video.mp4',
+                    mime='video/mp4'
+                )
+        except Exception as e:
+            st.error(f"An error occurred during download: {e}")