CoolFace
Apppublic

VidyaMadugula/video-agent

sourceHugging Faceupdated 3mo agoView on Hugging Face
0likes
audio_processor.py443 linesDownload Raw Back to utils
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#         )