VidyaMadugula/video-agent
0
1"""
2audio_pipeline.py
3
4Handles:
5- Any URL (YouTube, Vimeo, Twitter/X, TikTok, SoundCloud, etc. via yt-dlp,
6 1000+ sites supported) OR a direct link to an audio/video file
7- Local file uploads
8- Converts everything to mono 16kHz WAV
9- Chunks into fixed-length segments
10
11Each run gets its own subfolder: downloads/{job_id}/
12This makes end-to-end cleanup trivial — see cleanup_job_dir() in
13transcriber.py, which just shutil.rmtree()s the whole folder once
14a job is done.
15"""
16
17import os
18import glob
19import shutil
20import uuid
21import subprocess
22import requests
23import yt_dlp
24
25DOWNLOAD_DIR = "downloads"
26os.makedirs(DOWNLOAD_DIR, exist_ok=True)
27
28
29def make_job_dir() -> tuple[str, str]:
30 """Create a fresh per-job folder. Returns (job_id, job_dir)."""
31 job_id = uuid.uuid4().hex[:8]
32 job_dir = os.path.join(DOWNLOAD_DIR, job_id)
33 os.makedirs(job_dir, exist_ok=True)
34 return job_id, job_dir
35
36
37# --------------------------------------------------------------------------
38# 1. URL -> audio (YouTube + any yt-dlp-supported site)
39# --------------------------------------------------------------------------
40def download_via_ytdlp(url: str, job_dir: str) -> str:
41 """
42 Download audio from any yt-dlp-supported URL and convert to WAV.
43 Returns the path to the resulting .wav file, inside job_dir.
44 """
45 output_template = os.path.join(job_dir, "%(title).100B.%(ext)s")
46
47 ydl_opts = {
48 "format": "bestaudio/best",
49 "outtmpl": output_template,
50 "noplaylist": True,
51 "retries": 10,
52 "fragment_retries": 10,
53 "extractor_retries": 10,
54 "socket_timeout": 30,
55 "quiet": False, # keep this on while debugging
56 "no_warnings": False,
57 "extractor_args": {
58 "youtube": {
59 "player_client": ["web", "android"]
60 }
61 }
62 "postprocessors": [
63 {
64 "key": "FFmpegExtractAudio",
65 "preferredcodec": "wav",
66 "preferredquality": "192",
67 }
68 ],
69 # If YouTube throws "Sign in to confirm you're not a bot",
70 # uncomment ONE of the following (requires being logged in
71 # in that browser on this machine, or a cookies.txt export):
72 # "cookiesfrombrowser": ("chrome",),
73 # "cookiefile": "cookies.txt",
74 }
75
76 try:
77 with yt_dlp.YoutubeDL(ydl_opts) as ydl:
78 ydl.extract_info(url, download=True)
79 except yt_dlp.utils.DownloadError as e:
80 raise RuntimeError(f"yt-dlp could not download this URL: {e}") from e
81
82 # job_dir is unique per run, so we can just glob for the .wav
83 # inside it — no need for a filename prefix anymore.
84 matches = glob.glob(os.path.join(job_dir, "*.wav"))
85 if not matches:
86 raise FileNotFoundError(
87 f"yt-dlp reported success but no WAV was found in {job_dir}. "
88 f"Check that ffmpeg is installed and on PATH."
89 )
90 return matches[0]
91
92
93def is_ytdlp_supported(url: str) -> bool:
94 """
95 Does yt-dlp recognize this URL as belonging to a known extractor
96 (YouTube, Vimeo, etc.)? False for plain direct file links
97 (e.g. straight-up .mp3/.mp4 links), which we handle separately.
98 """
99 try:
100 with yt_dlp.YoutubeDL({"quiet": True}) as ydl:
101 ie = ydl.extract_info(url, download=False, process=False)
102 return ie is not None
103 except Exception:
104 return False
105
106
107# --------------------------------------------------------------------------
108# 2. Direct file URL -> local file (for links yt-dlp doesn't recognize,
109# e.g. a raw .mp3/.mp4/.wav hosted somewhere)
110# --------------------------------------------------------------------------
111_MEDIA_CONTENT_TYPES = ("audio/", "video/", "application/octet-stream")
112
113
114def download_direct_file(url: str, job_dir: str) -> str:
115 ext = os.path.splitext(url.split("?")[0])[1] or ".mp4"
116 local_path = os.path.join(job_dir, f"direct{ext}")
117
118 with requests.get(url, stream=True, timeout=60) as r:
119 r.raise_for_status()
120
121 content_type = r.headers.get("Content-Type", "")
122 if not any(content_type.startswith(ct) for ct in _MEDIA_CONTENT_TYPES):
123 raise RuntimeError(
124 f"URL did not return a media file (Content-Type: '{content_type}'). "
125 f"This usually means the link points to a webpage, not a raw "
126 f"audio/video file."
127 )
128
129 with open(local_path, "wb") as f:
130 for chunk in r.iter_content(chunk_size=8192):
131 f.write(chunk)
132
133 return local_path
134
135
136# --------------------------------------------------------------------------
137# 3. Any local file -> mono 16kHz WAV
138# --------------------------------------------------------------------------
139def convert_to_wav(input_path: str) -> str:
140 output_path = os.path.splitext(input_path)[0] + "_converted.wav"
141
142 result = subprocess.run(
143 [
144 "ffmpeg", "-i", input_path,
145 "-vn", "-ac", "1", "-ar", "16000", "-sample_fmt", "s16",
146 output_path, "-y",
147 ],
148 check=False,
149 capture_output=True,
150 text=True,
151 )
152
153 if result.returncode != 0 or not os.path.exists(output_path):
154 raise RuntimeError(f"ffmpeg conversion failed:\n{result.stderr}")
155
156 return output_path
157
158
159# --------------------------------------------------------------------------
160# 4. WAV -> chunks
161# --------------------------------------------------------------------------
162def chunk_audio(wav_path: str, chunk_minutes: int = 5) -> list:
163 chunk_seconds = chunk_minutes * 60
164 output_pattern = wav_path.replace(".wav", "_chunk_%03d.wav")
165
166 result = subprocess.run(
167 [
168 "ffmpeg", "-i", wav_path,
169 "-f", "segment", "-segment_time", str(chunk_seconds),
170 "-ar", "16000", "-ac", "1", "-sample_fmt", "s16", # re-encode, don't stream-copy
171 output_pattern, "-y",
172 ],
173 check=False,
174 capture_output=True,
175 text=True,
176 )
177
178 if result.returncode != 0:
179 raise RuntimeError(f"ffmpeg chunking failed:\n{result.stderr}")
180
181 base = wav_path.replace(".wav", "")
182 chunks = []
183 i = 0
184 while True:
185 path = f"{base}_chunk_{i:03d}.wav"
186 if not os.path.exists(path):
187 break
188 chunks.append(path)
189 i += 1
190
191 if not chunks:
192 raise RuntimeError("Chunking produced no output files.")
193
194 return chunks
195
196
197# --------------------------------------------------------------------------
198# 5. Entry point — handles URL (any site) OR local file path
199# --------------------------------------------------------------------------
200def process_input(source: str) -> tuple[list, str]:
201 """
202 Runs the full download -> convert -> chunk pipeline for a single job.
203
204 Returns (chunks, job_dir):
205 chunks - list of chunk WAV file paths, ready for transcription
206 job_dir - the per-job folder (downloads/{job_id}/) holding all
207 files for this run. Pass this to cleanup_job_dir()
208 once transcription (and any later steps) are done,
209 so the whole folder gets removed in one shot.
210 """
211 job_id, job_dir = make_job_dir()
212
213 try:
214 if source.startswith(("http://", "https://")):
215 print(f"[{job_id}] Detected URL.")
216
217 if is_ytdlp_supported(source):
218 # yt-dlp recognizes this site (YouTube, Vimeo, etc.)
219 # If extraction fails here, it's a real error
220 # (bad ID, bot-check, private/deleted video) —
221 # don't mask it by falling back to a raw HTTP GET.
222 print(f"[{job_id}] Site recognized by yt-dlp. Downloading...")
223 wav_path = download_via_ytdlp(source, job_dir)
224 else:
225 # Not a known site — likely a direct link to a
226 # media file (e.g. https://example.com/audio.mp3)
227 print(f"[{job_id}] Site not recognized by yt-dlp. Trying direct download...")
228 raw_path = download_direct_file(source, job_dir)
229 wav_path = convert_to_wav(raw_path)
230 else:
231 print(f"[{job_id}] Detected local file. Converting to WAV...")
232 # Copy the uploaded file into the job dir first, so
233 # everything for this job lives in one place and
234 # cleanup_job_dir() can remove it all together.
235 local_copy = os.path.join(job_dir, os.path.basename(source))
236 shutil.copy2(source, local_copy)
237 wav_path = convert_to_wav(local_copy)
238
239 print(f"[{job_id}] Chunking audio...")
240 chunks = chunk_audio(wav_path, chunk_minutes=5)
241 print(f"[{job_id}] Created {len(chunks)} chunk(s).")
242 return chunks, job_dir
243
244 except Exception as e:
245 raise RuntimeError(f"Processing failed for '{source}': {e}") from e
246
247
248if __name__ == "__main__":
249 # quick manual test
250 import sys
251 if len(sys.argv) < 2:
252 print("Usage: python audio_pipeline.py <url_or_filepath>")
253 else:
254 result_chunks, result_job_dir = process_input(sys.argv[1])
255 print("Chunks:", result_chunks)
256 print("Job dir:", result_job_dir)
257
258
259# import os
260# import uuid
261# import subprocess
262# import yt_dlp
263
264# DOWNLOAD_DIR = "downloads"
265# os.makedirs(DOWNLOAD_DIR, exist_ok=True)
266
267
268# def download_youtube_audio(url: str) -> str:
269# """
270# Download YouTube audio and convert to WAV.
271# """
272
273# unique_id = uuid.uuid4().hex[:8]
274
275# output_template = os.path.join(
276# DOWNLOAD_DIR,
277# f"{unique_id}_%(title)s.%(ext)s"
278# )
279
280# ydl_opts = {
281# "format": "bestaudio[ext=m4a]/bestaudio/best",
282# "outtmpl": output_template,
283# "noplaylist": True,
284# "retries": 10,
285# "fragment_retries": 10,
286# "extractor_retries": 10,
287# "quiet": True,
288
289# "postprocessors": [
290# {
291# "key": "FFmpegExtractAudio",
292# "preferredcodec": "wav",
293# "preferredquality": "192",
294# }
295# ],
296# }
297
298# # with yt_dlp.YoutubeDL(ydl_opts) as ydl:
299# # info = ydl.extract_info(url, download=True)
300# try:
301# with yt_dlp.YoutubeDL(ydl_opts) as ydl:
302# info = ydl.extract_info(
303# url,
304# download=True
305# )
306# except Exception as e:
307# print(f"YT-DLP ERROR: {e}")
308# raise
309
310# wav_path = (
311# os.path.splitext(
312# ydl.prepare_filename(info)
313# )[0]
314# + ".wav"
315# )
316
317# if not os.path.exists(wav_path):
318# raise FileNotFoundError(
319# f"WAV file not found: {wav_path}"
320# )
321
322# return wav_path
323
324
325# def convert_to_wav(input_path: str) -> str:
326# """
327# Convert any audio/video file to
328# mono 16kHz WAV.
329# """
330
331# output_path = (
332# os.path.splitext(input_path)[0]
333# + "_converted.wav"
334# )
335
336# subprocess.run(
337# [
338# "ffmpeg",
339# "-i",
340# input_path,
341# "-vn",
342# "-ac",
343# "1",
344# "-ar",
345# "16000",
346# "-sample_fmt",
347# "s16",
348# output_path,
349# "-y",
350# ],
351# check=True,
352# )
353
354# return output_path
355
356
357# def chunk_audio(
358# wav_path: str,
359# chunk_minutes: int = 5
360# ):
361# """
362# Split WAV into chunks.
363# """
364
365# chunk_seconds = chunk_minutes * 60
366
367# output_pattern = wav_path.replace(
368# ".wav",
369# "_chunk_%03d.wav"
370# )
371
372# subprocess.run(
373# [
374# "ffmpeg",
375# "-i",
376# wav_path,
377# "-f",
378# "segment",
379# "-segment_time",
380# str(chunk_seconds),
381# output_pattern,
382# "-y",
383# ],
384# check=True,
385# )
386
387# base = wav_path.replace(".wav", "")
388
389# chunks = []
390
391# i = 0
392
393# while True:
394# path = f"{base}_chunk_{i:03d}.wav"
395
396# if not os.path.exists(path):
397# break
398
399# chunks.append(path)
400# i += 1
401
402# return chunks
403
404
405# def process_input(source: str):
406# try:
407# if source.startswith(
408# ("http://", "https://")
409# ):
410# print(
411# "Detected YouTube URL."
412# )
413
414# wav_path = download_youtube_audio(
415# source
416# )
417
418# else:
419# print(
420# "Detected local file."
421# )
422
423# wav_path = convert_to_wav(
424# source
425# )
426
427# print("Chunking audio...")
428
429# chunks = chunk_audio(
430# wav_path,
431# chunk_minutes=5
432# )
433
434# print(
435# f"Created {len(chunks)} chunks."
436# )
437
438# return chunks
439
440# except Exception as e:
441# raise RuntimeError(
442# f"Processing failed: {e}"
443# )