TogetherAI commited on
Commit
ec615a5
1 Parent(s): f08ecef

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +74 -20
app.py CHANGED
@@ -1,16 +1,18 @@
1
  import torch
 
2
  import gradio as gr
3
  import yt_dlp as youtube_dl
4
  from transformers import pipeline
5
  from transformers.pipelines.audio_utils import ffmpeg_read
 
6
  import tempfile
7
  import os
8
- import time
9
 
10
- # Model and setup
11
  MODEL_NAME = "openai/whisper-large-v3"
12
  BATCH_SIZE = 8
13
- YT_LENGTH_LIMIT_S = 3600 # 1-hour limit for YouTube files
 
 
14
  device = 0 if torch.cuda.is_available() else "cpu"
15
 
16
  pipe = pipeline(
@@ -20,34 +22,74 @@ pipe = pipeline(
20
  device=device,
21
  )
22
 
23
- # Function to transcribe audio
24
  def transcribe(inputs, task):
25
  if inputs is None:
26
- raise gr.Error("No audio file submitted! Please upload or record an audio file.")
 
27
  text = pipe(inputs, batch_size=BATCH_SIZE, generate_kwargs={"task": task}, return_timestamps=True)["text"]
28
- return text
 
29
 
30
- # YouTube video processing functions
31
  def _return_yt_html_embed(yt_url):
32
  video_id = yt_url.split("?v=")[-1]
33
- return f'<center><iframe width="500" height="320" src="https://www.youtube.com/embed/{video_id}"></iframe></center>'
 
 
 
 
34
 
35
  def download_yt_audio(yt_url, filename):
36
- # [ ... existing code for download_yt_audio ... ]
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
37
 
38
- def yt_transcribe(yt_url, task):
 
39
  html_embed_str = _return_yt_html_embed(yt_url)
 
40
  with tempfile.TemporaryDirectory() as tmpdirname:
41
  filepath = os.path.join(tmpdirname, "video.mp4")
42
  download_yt_audio(yt_url, filepath)
43
  with open(filepath, "rb") as f:
44
  inputs = f.read()
 
45
  inputs = ffmpeg_read(inputs, pipe.feature_extractor.sampling_rate)
46
  inputs = {"array": inputs, "sampling_rate": pipe.feature_extractor.sampling_rate}
 
47
  text = pipe(inputs, batch_size=BATCH_SIZE, generate_kwargs={"task": task}, return_timestamps=True)["text"]
 
48
  return html_embed_str, text
49
 
50
- # Gradio interfaces
 
 
51
  mf_transcribe = gr.Interface(
52
  fn=transcribe,
53
  inputs=[
@@ -58,7 +100,12 @@ mf_transcribe = gr.Interface(
58
  layout="horizontal",
59
  theme="huggingface",
60
  title="Whisper Large V3: Transcribe Audio",
61
- description="Transcribe long-form microphone or audio inputs with the click of a button!"
 
 
 
 
 
62
  )
63
 
64
  file_transcribe = gr.Interface(
@@ -69,9 +116,14 @@ file_transcribe = gr.Interface(
69
  ],
70
  outputs="text",
71
  layout="horizontal",
72
- theme="NoCrypt/miku@1.2.1",
73
  title="Whisper Large V3: Transcribe Audio",
74
- description="Transcribe long-form microphone or audio inputs with the click of a button!"
 
 
 
 
 
75
  )
76
 
77
  yt_transcribe = gr.Interface(
@@ -82,16 +134,18 @@ yt_transcribe = gr.Interface(
82
  ],
83
  outputs=["html", "text"],
84
  layout="horizontal",
85
- theme="NoCrypt/miku@1.2.1",
86
  title="Whisper Large V3: Transcribe YouTube",
87
- description="Transcribe long-form YouTube videos with the click of a button!"
 
 
 
 
 
88
  )
89
 
90
- # Main Gradio application
91
- with gr.Blocks(theme="NoCrypt/miku@1.2.1") as demo:
92
- gr.HTML("<h1><center>AI Assistant<h1><center>")
93
  gr.TabbedInterface([mf_transcribe, file_transcribe, yt_transcribe], ["Microphone", "Audio file", "YouTube"])
94
 
95
  demo.launch(enable_queue=True)
96
 
97
-
 
1
  import torch
2
+
3
  import gradio as gr
4
  import yt_dlp as youtube_dl
5
  from transformers import pipeline
6
  from transformers.pipelines.audio_utils import ffmpeg_read
7
+
8
  import tempfile
9
  import os
 
10
 
 
11
  MODEL_NAME = "openai/whisper-large-v3"
12
  BATCH_SIZE = 8
13
+ FILE_LIMIT_MB = 1000
14
+ YT_LENGTH_LIMIT_S = 3600 # limit to 1 hour YouTube files
15
+
16
  device = 0 if torch.cuda.is_available() else "cpu"
17
 
18
  pipe = pipeline(
 
22
  device=device,
23
  )
24
 
25
+
26
  def transcribe(inputs, task):
27
  if inputs is None:
28
+ raise gr.Error("No audio file submitted! Please upload or record an audio file before submitting your request.")
29
+
30
  text = pipe(inputs, batch_size=BATCH_SIZE, generate_kwargs={"task": task}, return_timestamps=True)["text"]
31
+ return text
32
+
33
 
 
34
  def _return_yt_html_embed(yt_url):
35
  video_id = yt_url.split("?v=")[-1]
36
+ HTML_str = (
37
+ f'<center> <iframe width="500" height="320" src="https://www.youtube.com/embed/{video_id}"> </iframe>'
38
+ " </center>"
39
+ )
40
+ return HTML_str
41
 
42
  def download_yt_audio(yt_url, filename):
43
+ info_loader = youtube_dl.YoutubeDL()
44
+
45
+ try:
46
+ info = info_loader.extract_info(yt_url, download=False)
47
+ except youtube_dl.utils.DownloadError as err:
48
+ raise gr.Error(str(err))
49
+
50
+ file_length = info["duration_string"]
51
+ file_h_m_s = file_length.split(":")
52
+ file_h_m_s = [int(sub_length) for sub_length in file_h_m_s]
53
+
54
+ if len(file_h_m_s) == 1:
55
+ file_h_m_s.insert(0, 0)
56
+ if len(file_h_m_s) == 2:
57
+ file_h_m_s.insert(0, 0)
58
+ file_length_s = file_h_m_s[0] * 3600 + file_h_m_s[1] * 60 + file_h_m_s[2]
59
+
60
+ if file_length_s > YT_LENGTH_LIMIT_S:
61
+ yt_length_limit_hms = time.strftime("%HH:%MM:%SS", time.gmtime(YT_LENGTH_LIMIT_S))
62
+ file_length_hms = time.strftime("%HH:%MM:%SS", time.gmtime(file_length_s))
63
+ raise gr.Error(f"Maximum YouTube length is {yt_length_limit_hms}, got {file_length_hms} YouTube video.")
64
+
65
+ ydl_opts = {"outtmpl": filename, "format": "worstvideo[ext=mp4]+bestaudio[ext=m4a]/best[ext=mp4]/best"}
66
+
67
+ with youtube_dl.YoutubeDL(ydl_opts) as ydl:
68
+ try:
69
+ ydl.download([yt_url])
70
+ except youtube_dl.utils.ExtractorError as err:
71
+ raise gr.Error(str(err))
72
 
73
+
74
+ def yt_transcribe(yt_url, task, max_filesize=75.0):
75
  html_embed_str = _return_yt_html_embed(yt_url)
76
+
77
  with tempfile.TemporaryDirectory() as tmpdirname:
78
  filepath = os.path.join(tmpdirname, "video.mp4")
79
  download_yt_audio(yt_url, filepath)
80
  with open(filepath, "rb") as f:
81
  inputs = f.read()
82
+
83
  inputs = ffmpeg_read(inputs, pipe.feature_extractor.sampling_rate)
84
  inputs = {"array": inputs, "sampling_rate": pipe.feature_extractor.sampling_rate}
85
+
86
  text = pipe(inputs, batch_size=BATCH_SIZE, generate_kwargs={"task": task}, return_timestamps=True)["text"]
87
+
88
  return html_embed_str, text
89
 
90
+
91
+ demo = gr.Blocks()
92
+
93
  mf_transcribe = gr.Interface(
94
  fn=transcribe,
95
  inputs=[
 
100
  layout="horizontal",
101
  theme="huggingface",
102
  title="Whisper Large V3: Transcribe Audio",
103
+ description=(
104
+ "Transcribe long-form microphone or audio inputs with the click of a button! Demo uses the"
105
+ f" checkpoint [{MODEL_NAME}](https://huggingface.co/{MODEL_NAME}) and 🤗 Transformers to transcribe audio files"
106
+ " of arbitrary length."
107
+ ),
108
+ allow_flagging="never",
109
  )
110
 
111
  file_transcribe = gr.Interface(
 
116
  ],
117
  outputs="text",
118
  layout="horizontal",
119
+ theme="huggingface",
120
  title="Whisper Large V3: Transcribe Audio",
121
+ description=(
122
+ "Transcribe long-form microphone or audio inputs with the click of a button! Demo uses the"
123
+ f" checkpoint [{MODEL_NAME}](https://huggingface.co/{MODEL_NAME}) and 🤗 Transformers to transcribe audio files"
124
+ " of arbitrary length."
125
+ ),
126
+ allow_flagging="never",
127
  )
128
 
129
  yt_transcribe = gr.Interface(
 
134
  ],
135
  outputs=["html", "text"],
136
  layout="horizontal",
137
+ theme="huggingface",
138
  title="Whisper Large V3: Transcribe YouTube",
139
+ description=(
140
+ "Transcribe long-form YouTube videos with the click of a button! Demo uses the checkpoint"
141
+ f" [{MODEL_NAME}](https://huggingface.co/{MODEL_NAME}) and 🤗 Transformers to transcribe video files of"
142
+ " arbitrary length."
143
+ ),
144
+ allow_flagging="never",
145
  )
146
 
147
+ with demo:
 
 
148
  gr.TabbedInterface([mf_transcribe, file_transcribe, yt_transcribe], ["Microphone", "Audio file", "YouTube"])
149
 
150
  demo.launch(enable_queue=True)
151