Update app.py

app.py

@@ -10,9 +10,13 @@ import string
 import subprocess
 import glob
 import shutil
+from groq import Groq
+import tempfile
+from pydub import AudioSegment
 
-# Initialize the Hugging Face API with the token
-api = HfApi(token=os.getenv("HF_API_TOKEN"))
+# Initialize the Hugging Face and Groq APIs
+hf_api = HfApi(token=os.getenv("HF_API_TOKEN"))
+groq_client = Groq(api_key=os.getenv("GROQ_API_KEY"))
 
 def generate_random_string(length=4):
     return ''.join(random.choices(string.ascii_lowercase, k=length))
@@ -40,13 +44,13 @@ def save_custom_thumbnail(thumbnail_file, thumbnail_path):
 
 def get_video_length(video_path):
     video = cv2.VideoCapture(video_path)
-    fps = video.get(cv2.CAP_PROP_FPS)  # Frames per second
-    total_frames = int(video.get(cv2.CAP_PROP_FRAME_COUNT))  # Total frames in the video
-    duration = int(total_frames / fps) if fps > 0 else 0  # Duration in seconds, as an integer
+    fps = video.get(cv2.CAP_PROP_FPS)
+    total_frames = int(video.get(cv2.CAP_PROP_FRAME_COUNT))
+    duration = int(total_frames / fps) if fps > 0 else 0
     video.release()
     return duration
 
-def generate_metadata(video_name, title, description, uploader, file_location, thumbnail_location, duration):
+def generate_metadata(video_name, title, description, uploader, file_location, thumbnail_location, subtitle_location, duration):
     return {
         "fileName": video_name,
         "title": title,
@@ -55,46 +59,38 @@ def generate_metadata(video_name, title, description, uploader, file_location, t
         "uploadTimestamp": datetime.now().isoformat(),
         "fileLocation": file_location,
         "thumbnailLocation": thumbnail_location,
-        "duration": duration,  # Add duration here
+        "subtitleLocation": subtitle_location,
+        "duration": duration,
         "views": 0,
         "likes": 0
     }
 
 def update_index_file(new_metadata_path):
-    # Create a temporary directory for cloning
     temp_dir = "temp_repo"
     if os.path.exists(temp_dir):
         shutil.rmtree(temp_dir)
     
     try:
-        # Clone the repository
         subprocess.run(['git', 'clone', 'https://huggingface.co/spaces/vericudebuget/ok4231', temp_dir], check=True)
-        
-        # Find all JSON files in the metadata directory
         metadata_dir = os.path.join(temp_dir, 'metadata')
         json_files = glob.glob(os.path.join(metadata_dir, '*-index.json'))
         
-        # Create the paths string
         base_url = "huggingface.co/spaces/vericudebuget/ok4231/raw/main/metadata/"
         paths = [f"{base_url}{os.path.basename(f)}" for f in json_files]
         
-        # Add the new metadata path if it's not already there
         new_metadata_filename = os.path.basename(new_metadata_path)
         new_full_path = f"{base_url}{new_metadata_filename}"
         if new_full_path not in paths:
             paths.append(new_full_path)
         
-        # Create the index content
         index_content = "{ " + "; ".join(paths) + "; }"
         
-        # Write to a temporary file
         index_path = os.path.join(temp_dir, 'metadata', 'video-index.json')
         os.makedirs(os.path.dirname(index_path), exist_ok=True)
         with open(index_path, 'w') as f:
             f.write(index_content)
         
-        # Upload the index file
-        api.upload_file(
+        hf_api.upload_file(
             path_or_fileobj=index_path,
             path_in_repo="metadata/video-index.json",
             repo_id="vericudebuget/ok4231",
@@ -102,18 +98,45 @@ def update_index_file(new_metadata_path):
         )
     
     finally:
-        # Clean up
         if os.path.exists(temp_dir):
             shutil.rmtree(temp_dir)
 
-def upload_video_to_hf(video_file, original_video_name, title, description, uploader, custom_thumbnail=None):
-    # Create temp paths
+def generate_subtitles(video_path):
+    with tempfile.NamedTemporaryFile(suffix='.mp3', delete=False) as temp_audio:
+        # Convert video to mono 128kbps MP3
+        audio = AudioSegment.from_file(video_path)
+        audio = audio.set_channels(1).set_frame_rate(44100).set_sample_width(2)
+        audio.export(temp_audio.name, format='mp3', bitrate='128k')
+        
+        # Generate subtitles using Groq
+        with open(temp_audio.name, 'rb') as audio_file:
+            translation = groq_client.audio.translations.create(
+                file=(temp_audio.name, audio_file.read()),
+                model="whisper-large-v3",
+                response_format="verbose_json",
+                temperature=0.0
+            )
+    
+    # Generate VTT content
+    vtt_content = "WEBVTT\n\n"
+    for segment in translation.segments:
+        start_time = segment['start']
+        end_time = segment['end']
+        text = segment['text'].strip()
+        
+        start_time_vtt = f"{int(start_time // 3600):02}:{int((start_time % 3600) // 60):02}:{start_time % 60:06.3f}"
+        end_time_vtt = f"{int(end_time // 3600):02}:{int((end_time % 3600) // 60):02}:{end_time % 60:06.3f}"
+        
+        vtt_content += f"{start_time_vtt} --> {end_time_vtt}\n{text}\n\n"
+    
+    return vtt_content
+
+def upload_video_to_hf(video_file, original_video_name, title, description, uploader, generate_subs=False, custom_thumbnail=None):
     temp_dir = "temp"
     if not os.path.exists(temp_dir):
         os.makedirs(temp_dir)
     
     try:
-        # Generate randomized filenames
         video_name = add_random_to_filename(original_video_name)
         video_path = os.path.join(temp_dir, video_name)
         
@@ -124,11 +147,9 @@ def upload_video_to_hf(video_file, original_video_name, title, description, uplo
         json_name = f"{base_name}-index.json"
         json_path = os.path.join(temp_dir, json_name)
         
-        # Write the video content to a file
         with open(video_path, "wb") as f:
             f.write(video_file.read())
         
-        # Handle thumbnail
         if custom_thumbnail:
             thumbnail_extracted = save_custom_thumbnail(custom_thumbnail, thumbnail_path)
         else:
@@ -138,47 +159,61 @@ def upload_video_to_hf(video_file, original_video_name, title, description, uplo
             st.error("Failed to process thumbnail")
             return None
         
-        # Get video length
         video_length = get_video_length(video_path)
+        
+        # Generate and upload subtitles if requested and video is not too long
+        subtitle_location = ""
+        if generate_subs and video_length <= 3600:  # 1 hour in seconds
+            vtt_content = generate_subtitles(video_path)
+            subtitle_name = f"{base_name}.vtt"
+            subtitle_path = os.path.join(temp_dir, subtitle_name)
+            
+            with open(subtitle_path, 'w') as f:
+                f.write(vtt_content)
+            
+            subtitle_location = f"subtitles/{subtitle_name}"
+            hf_api.upload_file(
+                path_or_fileobj=subtitle_path,
+                path_in_repo=subtitle_location,
+                repo_id="vericudebuget/ok4231",
+                repo_type="space",
+            )
 
-        # Upload the video
+        # Upload video and thumbnail
         video_location = f"videos/{video_name}"
-        api.upload_file(
+        hf_api.upload_file(
             path_or_fileobj=video_path,
             path_in_repo=video_location,
             repo_id="vericudebuget/ok4231",
             repo_type="space",
         )
         
-        # Upload the thumbnail
         thumbnail_location = f"thumbnails/{thumbnail_name}"
-        api.upload_file(
+        hf_api.upload_file(
             path_or_fileobj=thumbnail_path,
             path_in_repo=thumbnail_location,
             repo_id="vericudebuget/ok4231",
             repo_type="space",
         )
         
-        # Generate and upload metadata JSON
-        metadata = generate_metadata(video_name, title, description, uploader, video_location, thumbnail_location, video_length)
+        # Generate and upload metadata
+        metadata = generate_metadata(video_name, title, description, uploader, video_location, thumbnail_location, subtitle_location, video_length)
         with open(json_path, "w") as f:
             json.dump(metadata, f, indent=2)
         
         metadata_location = f"metadata/{json_name}"
-        api.upload_file(
+        hf_api.upload_file(
             path_or_fileobj=json_path,
             path_in_repo=metadata_location,
             repo_id="vericudebuget/ok4231",
             repo_type="space",
         )
         
-        # Update the index file
         update_index_file(metadata_location)
         
         return metadata
     
     finally:
-        # Cleanup temp files
         if os.path.exists(temp_dir):
             shutil.rmtree(temp_dir)
 
@@ -186,38 +221,46 @@ def upload_video_to_hf(video_file, original_video_name, title, description, uplo
 st.title("Upload your video")
 st.markdown("---")
 
-# File uploader for video input
-uploaded_video = st.file_uploader("Choose video file", type=["mp4", "avi", "mov"])
+uploaded_video = st.file_uploader("Choose video file", type=["mp4", "avi", "mov", "webm", "mkv"])
 
 if uploaded_video:
-    # Show the video details form
     with st.form("video_details"):
         st.write("Video Details")
         title = st.text_input("Title", placeholder="Enter video title")
         description = st.text_area("Description", placeholder="Enter video description")
         uploader = st.text_input("Uploader Name", placeholder="Enter your name")
         
-        # Optional custom thumbnail uploader
+        # Create a temporary file to get video duration
+        with tempfile.NamedTemporaryFile(delete=False, suffix='.mp4') as temp_video:
+            temp_video.write(uploaded_video.getvalue())
+            video_duration = get_video_length(temp_video.name)
+        os.unlink(temp_video.name)  # Clean up temp file
+        
+        # Subtitle generation toggle, disabled if video is longer than 1 hour
+        generate_subtitles = st.toggle("Generate Subtitles", disabled=video_duration > 3600)
+        if video_duration > 3600 and generate_subtitles:
+            st.warning("Subtitle generation is disabled for videos longer than 1 hour.")
+        
         custom_thumbnail = st.file_uploader("Upload custom thumbnail (optional)", type=["jpg", "jpeg", "png"])
         
-        # Upload button within the form
         submit_button = st.form_submit_button("Upload Video")
         
         if submit_button:
             if not title or not uploader:
                 st.error("Please fill in the title and uploader name.")
             else:
-                with st.spinner("Uploading video, generating thumbnail and metadata..."):
+                with st.spinner("Uploading video, generating thumbnail and metadata... This may take some time. Please wait."):
                     metadata = upload_video_to_hf(
                         uploaded_video, 
                         uploaded_video.name, 
                         title, 
                         description, 
                         uploader,
+                        generate_subtitles,
                         custom_thumbnail
                     )
                     if metadata:
                         st.success("Upload completed successfully!")
                         st.json(metadata)
 else:
-    st.info("Please upload a video file to begin.")
+    st.info("Please upload a video file to begin.")