Update app.py

app.py

@@ -1,126 +1,83 @@
-import streamlit as st
-import yt_dlp
 import os
-
-# Streamlit 제목
-st.title('유튜브 영상 다운로드')
-st.write("<br>", unsafe_allow_html=True)  # 1줄 띄어쓰기
-st.write("사용방법")
-st.write("<br>", unsafe_allow_html=True)  # 1줄 띄어쓰기
-st.write("1. 다운받으려는 유튜브 영상의 링크를 입력해주세요")
-st.write("2. 링크를 입력하고 확인버튼을 누르면 해당 영상의 썸네일이 등장합니다")
-st.write("3. 썸네일 밑의 선택란을 통해 해상도와 영상 용량을 확인한 후 원하는 해상도를 선택하고 추출 버튼을 눌러주세요")
-st.write("4. 잠시 후 다운받으려는 영상이 밑에 등장하고 재생하거나 다운로드 받을 수 있습니다 (영상의 용량이 클수록 추출하는데 시간이 더 걸립니다. 20분 영상 추출하는데 1분내로 소요)")
-
-# 유저로부터 유튜브 링크 입력받기
-col1, col2 = st.columns([4, 1])
-with col1:
-    youtube_url = st.text_input('Enter YouTube video URL')
-
-def fetch_video_info(youtube_url):
-    ydl_opts = {'nocheckcertificate': True}
-    with yt_dlp.YoutubeDL(ydl_opts) as ydl:
-        info_dict = ydl.extract_info(youtube_url, download=False)
-        return info_dict
-
-def get_available_formats(info_dict):
-    formats = {}
-    for fmt in info_dict['formats']:
-        if fmt.get('height') and fmt.get('ext') == 'mp4':
-            filesize = fmt.get('filesize', 0)
-            if filesize:
-                filesize_mb = round(filesize / (1024 * 1024), 2)
-                resolution = f"{fmt['height']}p"
-                formats[resolution] = {
-                    'id': fmt['format_id'],
-                    'info': f"{resolution} - {filesize_mb} MB"
-                }
-    return formats
-
-def reset_session_state():
-    st.session_state.pop('info_dict', None)
-    st.session_state.pop('thumbnail_url', None)
-    st.session_state.pop('formats', None)
-    st.session_state.pop('youtube_url', None)
-    st.session_state.pop('confirmed', None)
-    st.session_state.pop('selected_format_id', None)
-
-def delete_downloaded_video():
-    if os.path.exists('downloaded_video.mp4'):
-        os.remove('downloaded_video.mp4')
-
-# "확인" 버튼을 누르면 처리할 로직
-if st.button('확인'):
-    if youtube_url:
-        try:
-            # 기존 다운로드된 영상 삭제
-            delete_downloaded_video()
-
-            # 세션 상태 초기화 및 URL 변경 확인
-            if 'youtube_url' in st.session_state and st.session_state.youtube_url != youtube_url:
-                reset_session_state()
-
-            st.session_state.youtube_url = youtube_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.formats = get_available_formats(info_dict)
-            st.session_state.confirmed = True
-        except Exception as e:
-            st.error(f"An error occurred: {e}")
-    else:
-        st.warning('Please enter a valid YouTube video URL')
-
-if 'confirmed' in st.session_state and st.session_state.confirmed:
-    st.image(st.session_state.thumbnail_url, caption='Video Thumbnail', use_column_width=True)
-
-    # 해상도 선택
-    format_options = [value['info'] for value in st.session_state.formats.values()]
-    format_selection = st.selectbox('Select resolution', format_options)
-
-    # 선택된 해상도의 format_id 저장
-    st.session_state.selected_format_id = next(
-        (value['id'] for key, value in st.session_state.formats.items() if value['info'] == format_selection), None
-    )
-
-    def download_video(youtube_url, format_id):
-        ydl_opts = {
-            'format': f"{format_id}+bestaudio/best",
-            'merge_output_format': 'mp4',
-            'outtmpl': 'downloaded_video.%(ext)s',
-            'nocache': True,  # 캐시 비활성화
-            'force_generic_extractor': True  # 항상 새로운 URL 사용
-        }
-        with yt_dlp.YoutubeDL(ydl_opts) as ydl:
-            ydl.download([youtube_url])
-
-    # "추출" 버튼 추가
-    if st.button('추출'):
-        format_id = st.session_state.selected_format_id
-        if format_id is None:
-            st.error('Selected format not found.')
-        else:
-            try:
-                # 세션 상태의 youtube_url 사용
-                download_video(st.session_state.youtube_url, format_id)
-                st.success('Video downloaded successfully!')
-                st.video('downloaded_video.mp4')
-
-                # 다운로드된 파일 제공 (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'
-                    )
-
-                # 추출된 영상 정보 출력
-                st.write(f"**제목:** {st.session_state.info_dict['title']}")
-                st.write(f"**채널명:** {st.session_state.info_dict['uploader']}")
-                st.write(f"**업로드 날짜:** {st.session_state.info_dict['upload_date']}")
-
-                # 다운로드 후 세션 상태 리셋
-                reset_session_state()
-            except Exception as e:
-                st.error(f"An error occurred during download: {e}")
+import streamlit as st
+import torch
+from transformers import WhisperForConditionalGeneration, WhisperProcessor
+import librosa
+import moviepy.editor as mp
+import srt
+from datetime import timedelta
+from tempfile import NamedTemporaryFile
+
+# 모델 및 프로세서 로드
+@st.cache_resource
+def load_model():
+    model = WhisperForConditionalGeneration.from_pretrained("lcjln/AIME_Project_The_Final")
+    processor = WhisperProcessor.from_pretrained("lcjln/AIME_The_Final")
+    return model, processor
+
+model, processor = load_model()
+
+# 웹 애플리케이션 인터페이스
+st.title("Whisper 자막 생성기")
+st.write("영상 파일을 업로드하여 자막을 생성하세요.")
+
+# 영상 파일 업로드
+uploaded_file = st.file_uploader("여기에 영상 파일을 드래그 앤 드롭 하세요", type=["mp4", "mkv", "mov"])
+
+if uploaded_file is not None:
+    # 진행바 초기화
+    progress_bar = st.progress(0)
+    
+    # 업로드된 영상 파일에서 오디오 추출
+    st.write("영상을 처리하는 중입니다...")
+    progress_bar.progress(10)
+
+    with NamedTemporaryFile(suffix=".mp4") as temp_video_file:
+        temp_video_file.write(uploaded_file.read())
+        video = mp.VideoFileClip(temp_video_file.name)
+        audio_path = temp_video_file.name.replace(".mp4", ".wav")
+        video.audio.write_audiofile(audio_path, codec='pcm_s16le')
+
+    progress_bar.progress(30)
+
+    # 오디오 파일 로드 및 처리
+    st.write("오디오 파일을 처리하는 중입니다...")
+    audio, sr = librosa.load(audio_path, sr=16000)
+
+    progress_bar.progress(50)
+
+    # Whisper 모델로 변환
+    st.write("모델을 통해 자막을 생성하는 중입니다...")
+    inputs = processor(audio, return_tensors="pt", sampling_rate=16000)
+    with torch.no_grad():
+        predicted_ids = model.generate(inputs["input_features"], max_length=2048)
+
+    transcription = processor.batch_decode(predicted_ids, skip_special_tokens=True)[0].strip()
+
+    progress_bar.progress(80)
+
+    # 자막을 SRT 파일로 저장
+    st.write("SRT 파일을 생성하는 중입니다...")
+    lines = transcription.split(". ")
+    subs = []
+    step = len(audio) / sr / len(lines)
+    start_time = 0.0
+
+    for i, line in enumerate(lines):
+        end_time = start_time + step
+        subs.append(srt.Subtitle(index=i+1, start=timedelta(seconds=start_time), end=timedelta(seconds=end_time), content=line))
+        start_time = end_time
+
+    srt_content = srt.compose(subs)
+
+    srt_file_path = audio_path.replace(".wav", ".srt")
+    with open(srt_file_path, "w", encoding="utf-8") as f:
+        f.write(srt_content)
+
+    progress_bar.progress(100)
+
+    st.success("SRT 파일이 성공적으로 생성되었습니다!")
+
+    # SRT 파일 다운로드 버튼
+    with open(srt_file_path, "rb") as srt_file:
+        st.download_button(label="SRT 파일 다운로드", data=srt_file, file_name=os.path.basename(srt_file_path), mime="text/srt")