StevenChen16 commited on
Commit
1e923d6
1 Parent(s): 4036c8e

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +36 -54
app.py CHANGED
@@ -6,63 +6,48 @@ from transformers.pipelines.audio_utils import ffmpeg_read
6
  import tempfile
7
  import gc
8
  import os
9
- import time
10
 
11
  # Constants
12
  DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
13
- BATCH_SIZE = 2 # Reduced batch size
14
- COMPUTE_TYPE = "int8" # Changed to int8 for lower memory usage
15
- FILE_LIMIT_MB = 25 # Reduced file size limit
16
 
17
- def clean_gpu_memory():
18
- """Helper function to clean GPU memory"""
19
- if torch.cuda.is_available():
20
- torch.cuda.empty_cache()
21
- gc.collect()
22
-
23
- @spaces.GPU
24
  def transcribe_audio(inputs, task):
25
  if inputs is None:
26
  raise gr.Error("No audio file submitted! Please upload or record an audio file before submitting your request.")
27
 
28
  try:
29
- # Check file size
30
- file_size = os.path.getsize(inputs) / (1024 * 1024) # Convert to MB
31
- if file_size > FILE_LIMIT_MB:
32
- raise gr.Error(f"File size ({file_size:.1f}MB) exceeds limit of {FILE_LIMIT_MB}MB")
33
-
34
- # Load audio with error handling
35
- try:
36
  audio = whisperx.load_audio(inputs)
37
- except Exception as e:
38
- raise gr.Error(f"Error loading audio file: {str(e)}")
39
 
40
  # 1. Transcribe with base Whisper model
41
- try:
42
- model = whisperx.load_model("large-v3", device=DEVICE, compute_type=COMPUTE_TYPE)
43
- result = model.transcribe(audio, batch_size=BATCH_SIZE)
44
- finally:
45
- clean_gpu_memory()
46
- if 'model' in locals():
47
- del model
48
 
49
  # 2. Align whisper output
50
- try:
51
- model_a, metadata = whisperx.load_align_model(language_code=result["language"], device=DEVICE)
52
- result = whisperx.align(result["segments"], model_a, metadata, audio, DEVICE, return_char_alignments=False)
53
- finally:
54
- clean_gpu_memory()
55
- if 'model_a' in locals():
56
- del model_a
57
 
58
  # 3. Diarize audio
59
- try:
60
- diarize_model = whisperx.DiarizationPipeline(use_auth_token=os.environ.get("HF_TOKEN"), device=DEVICE)
61
- diarize_segments = diarize_model(audio)
62
- finally:
63
- if 'diarize_model' in locals():
64
- del diarize_model
65
- clean_gpu_memory()
66
 
67
  # 4. Assign speaker labels
68
  result = whisperx.assign_word_speakers(diarize_segments, result)
@@ -77,8 +62,12 @@ def transcribe_audio(inputs, task):
77
  return output_text
78
 
79
  except Exception as e:
80
- clean_gpu_memory()
81
  raise gr.Error(f"Error processing audio: {str(e)}")
 
 
 
 
 
82
 
83
  # Create Gradio interface
84
  demo = gr.Blocks(theme=gr.themes.Ocean())
@@ -91,7 +80,7 @@ with demo:
91
  audio_input = gr.Audio(
92
  sources=["microphone", "upload"],
93
  type="filepath",
94
- label=f"Audio Input (Max {FILE_LIMIT_MB}MB)",
95
  )
96
  task = gr.Radio(
97
  ["transcribe", "translate"],
@@ -107,17 +96,15 @@ with demo:
107
  placeholder="Transcribed text will appear here..."
108
  )
109
 
110
- gr.Markdown(f"""
111
  ### Features:
112
  - High-accuracy transcription using WhisperX
113
  - Automatic speaker diarization
114
  - Support for both microphone recording and file upload
115
- - File size limit: {FILE_LIMIT_MB}MB
116
 
117
  ### Note:
118
- - Processing may take a few moments
119
- - For optimal results, use clear audio with minimal background noise
120
- - If you encounter errors, try with a shorter audio clip
121
  """)
122
 
123
  submit_button.click(
@@ -126,9 +113,4 @@ with demo:
126
  outputs=output_text
127
  )
128
 
129
- demo.queue(max_size=1).launch(
130
- share=False,
131
- debug=True,
132
- show_error=True,
133
- ssr_mode=False
134
- )
 
6
  import tempfile
7
  import gc
8
  import os
 
9
 
10
  # Constants
11
  DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
12
+ BATCH_SIZE = 4 # reduce if low on GPU mem
13
+ COMPUTE_TYPE = "float32" # change to "int8" if low on GPU mem
14
+ FILE_LIMIT_MB = 1000
15
 
16
+ @spaces.GPU(duration=200)
 
 
 
 
 
 
17
  def transcribe_audio(inputs, task):
18
  if inputs is None:
19
  raise gr.Error("No audio file submitted! Please upload or record an audio file before submitting your request.")
20
 
21
  try:
22
+ # Load audio
23
+ if isinstance(inputs, str):
24
+ # For file path input
25
+ audio = whisperx.load_audio(inputs)
26
+ else:
27
+ # For microphone input (needs conversion)
 
28
  audio = whisperx.load_audio(inputs)
 
 
29
 
30
  # 1. Transcribe with base Whisper model
31
+ model = whisperx.load_model("large-v3", device=DEVICE, compute_type=COMPUTE_TYPE)
32
+ result = model.transcribe(audio, batch_size=BATCH_SIZE)
33
+
34
+ # Clear GPU memory
35
+ del model
36
+ gc.collect()
37
+ torch.cuda.empty_cache()
38
 
39
  # 2. Align whisper output
40
+ model_a, metadata = whisperx.load_align_model(language_code=result["language"], device=DEVICE)
41
+ result = whisperx.align(result["segments"], model_a, metadata, audio, DEVICE, return_char_alignments=False)
42
+
43
+ # Clear GPU memory again
44
+ del model_a
45
+ gc.collect()
46
+ torch.cuda.empty_cache()
47
 
48
  # 3. Diarize audio
49
+ diarize_model = whisperx.DiarizationPipeline(use_auth_token=os.environ["HF_TOKEN"], device=DEVICE)
50
+ diarize_segments = diarize_model(audio)
 
 
 
 
 
51
 
52
  # 4. Assign speaker labels
53
  result = whisperx.assign_word_speakers(diarize_segments, result)
 
62
  return output_text
63
 
64
  except Exception as e:
 
65
  raise gr.Error(f"Error processing audio: {str(e)}")
66
+
67
+ finally:
68
+ # Final cleanup
69
+ gc.collect()
70
+ torch.cuda.empty_cache()
71
 
72
  # Create Gradio interface
73
  demo = gr.Blocks(theme=gr.themes.Ocean())
 
80
  audio_input = gr.Audio(
81
  sources=["microphone", "upload"],
82
  type="filepath",
83
+ label="Audio Input (Microphone or File Upload)"
84
  )
85
  task = gr.Radio(
86
  ["transcribe", "translate"],
 
96
  placeholder="Transcribed text will appear here..."
97
  )
98
 
99
+ gr.Markdown("""
100
  ### Features:
101
  - High-accuracy transcription using WhisperX
102
  - Automatic speaker diarization
103
  - Support for both microphone recording and file upload
104
+ - GPU-accelerated processing
105
 
106
  ### Note:
107
+ Processing may take a few moments depending on the audio length and system resources.
 
 
108
  """)
109
 
110
  submit_button.click(
 
113
  outputs=output_text
114
  )
115
 
116
+ demo.queue().launch(ssr_mode=False)