Spaces:
Sleeping
Sleeping
| import os | |
| import math | |
| import gradio as gr | |
| import torch | |
| from transformers import AutoModelForSpeechSeq2Seq, AutoProcessor, pipeline | |
| from moviepy.editor import VideoFileClip, concatenate_audioclips | |
| def transcribe(video_file, transcribe_to_text, transcribe_to_srt, language): | |
| device = "cuda:0" if torch.cuda.is_available() else "cpu" | |
| torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32 | |
| model_id = "openai/whisper-large-v3" | |
| model = AutoModelForSpeechSeq2Seq.from_pretrained( | |
| model_id, torch_dtype=torch_dtype, low_cpu_mem_usage=True, use_safetensors=True | |
| ) | |
| model.to(device) | |
| processor = AutoProcessor.from_pretrained(model_id) | |
| pipe = pipeline( | |
| "automatic-speech-recognition", | |
| model=model, | |
| tokenizer=processor.tokenizer, | |
| feature_extractor=processor.feature_extractor, | |
| max_new_tokens=128, | |
| chunk_length_s=10, # Shorter chunk length to prevent overly long captions | |
| batch_size=2, | |
| return_timestamps=True, | |
| torch_dtype=torch_dtype, | |
| device=device, | |
| ) | |
| if video_file is None: | |
| yield "Error: No video file provided.", None | |
| return | |
| video_path = video_file.name if hasattr(video_file, 'name') else video_file | |
| try: | |
| video = VideoFileClip(video_path) | |
| except Exception as e: | |
| yield f"Error processing video file: {str(e)}", None | |
| return | |
| audio = video.audio | |
| duration = video.duration | |
| n_chunks = math.ceil(duration / 10) # Split into 10-second chunks | |
| transcription_txt = "" | |
| transcription_srt = [] | |
| for i in range(n_chunks): | |
| start = i * 10 | |
| end = min((i + 1) * 10, duration) | |
| audio_chunk = audio.subclip(start, end) | |
| temp_file_path = f"temp_audio_{i}.wav" | |
| audio_chunk.write_audiofile(temp_file_path, codec='pcm_s16le') | |
| with open(temp_file_path, "rb") as temp_file: | |
| result = pipe(temp_file_path, generate_kwargs={"language": language}) | |
| transcription_txt += result["text"] | |
| if transcribe_to_srt: | |
| for chunk in result["chunks"]: | |
| start_time, end_time = chunk["timestamp"] | |
| if start_time is not None and end_time is not None: | |
| transcription_srt.append({ | |
| "start": start_time + i * 10, | |
| "end": end_time + i * 10, | |
| "text": chunk["text"] | |
| }) | |
| else: | |
| print(f"Warning: Invalid timestamp for chunk: {chunk}") | |
| os.remove(temp_file_path) | |
| yield f"Progress: {int(((i + 1) / n_chunks) * 100)}%", None | |
| output = "" | |
| srt_file_path = None | |
| if transcribe_to_text: | |
| output += "Text Transcription:\n" + transcription_txt + "\n\n" | |
| if transcribe_to_srt: | |
| output += "SRT Transcription:\n" | |
| srt_content = "" | |
| for i, sub in enumerate(transcription_srt, 1): | |
| srt_entry = f"{i}\n{format_time(sub['start'])} --> {format_time(sub['end'])}\n{sub['text']}\n\n" | |
| output += srt_entry | |
| srt_content += srt_entry | |
| srt_file_path = "transcription.srt" | |
| with open(srt_file_path, "w", encoding="utf-8") as srt_file: | |
| srt_file.write(srt_content) | |
| output += f"\nSRT file saved as: {srt_file_path}" | |
| yield output, srt_file_path | |
| def format_time(seconds): | |
| m, s = divmod(seconds, 60) | |
| h, m = divmod(m, 60) | |
| return f"{int(h):02d}:{int(m):02d}:{s:06.3f}".replace('.', ',') | |
| iface = gr.Interface( | |
| fn=transcribe, | |
| inputs=[ | |
| gr.Video(), | |
| gr.Checkbox(label="Transcribe to Text"), | |
| gr.Checkbox(label="Transcribe to SRT"), | |
| gr.Dropdown(choices=['en', 'he', 'it', 'es', 'fr', 'de', 'zh', 'ar'], label="Language") | |
| ], | |
| outputs=[ | |
| gr.Textbox(label="Transcription Output"), | |
| gr.File(label="Download SRT") | |
| ], | |
| title="WhisperCap Video Transcription", | |
| description="Upload a video file to transcribe its audio using Whisper. You can download the SRT file if generated.", | |
| ) | |
| iface.launch(share=True) | |