From 5c66b15b5115c4ed90c13e8c62c651af829d0353 Mon Sep 17 00:00:00 2001 From: BennyThink Date: Sun, 27 Feb 2022 10:35:04 +0800 Subject: [PATCH] progress bar for ffmpeg --- requirements.txt | 1 + ytdlbot/downloader.py | 30 +++++++++++++++++++++++++----- ytdlbot/tasks.py | 12 ++++++++---- ytdlbot/ytdl_bot.py | 5 ++--- 4 files changed, 36 insertions(+), 12 deletions(-) diff --git a/requirements.txt b/requirements.txt index 527a556..a5a000f 100644 --- a/requirements.txt +++ b/requirements.txt @@ -18,3 +18,4 @@ redis==4.1.4 requests==2.27.1 tqdm==4.62.3 requests-toolbelt==0.9.1 +ffpb==0.4.1 \ No newline at end of file diff --git a/ytdlbot/downloader.py b/ytdlbot/downloader.py index b400981..6715f50 100644 --- a/ytdlbot/downloader.py +++ b/ytdlbot/downloader.py @@ -19,6 +19,7 @@ from unittest.mock import MagicMock import fakeredis import ffmpeg +import ffpb import filetype import yt_dlp as ytdl from tqdm import tqdm @@ -65,7 +66,7 @@ def tqdm_progress(desc, total, finished, speed="", eta=""): raw_output = f.getvalue() tqdm_output = raw_output.split("|") progress = f"`[{tqdm_output[1]}]`" - detail = tqdm_output[2] + detail = tqdm_output[2].replace("[A", "") text = f""" {desc} @@ -142,13 +143,32 @@ def convert_to_mp4(resp: dict, bot_msg): edit_text(bot_msg, f"{current_time()}: Converting {path.name} to mp4. Please wait.") new_file_path = path.with_suffix(".mp4") logging.info("Detected %s, converting to mp4...", mime) - subprocess.check_output(["ffmpeg", "-y", "-i", path, new_file_path]) + run_ffmpeg(["ffmpeg", "-y", "-i", path, new_file_path], bot_msg) index = resp["filepath"].index(path) resp["filepath"][index] = new_file_path return resp +class ProgressBar(tqdm): + b = None + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.bot_msg = self.b + + def update(self, n=1): + super().update(n) + t = tqdm_progress("Converting...", self.total, self.n) + edit_text(self.bot_msg, t) + + +def run_ffmpeg(cmd_list, bm): + cmd_list = cmd_list.copy()[1:] + ProgressBar.b = bm + ffpb.main(cmd_list, tqdm=ProgressBar) + + def can_convert_mp4(video_path, uid): if not ENABLE_VIP: return True @@ -225,20 +245,20 @@ def ytdl_download(url, tempdir, bm) -> dict: # only convert if send type is video convert_to_mp4(response, bm) if settings[2] == "audio": - check_audio_format(response) + convert_audio_format(response, bm) # disable it for now # split_large_video(response) return response -def check_audio_format(resp: "dict"): +def convert_audio_format(resp: "dict", bm): if resp["status"]: # all_converted = [] path: pathlib.PosixPath for path in resp["filepath"]: if path.suffix != f".{AUDIO_FORMAT}": new_path = path.with_suffix(f".{AUDIO_FORMAT}") - subprocess.check_output(["ffmpeg", "-y", "-i", path, new_path]) + run_ffmpeg(["ffmpeg", "-y", "-i", path, new_path], bm) path.unlink() index = resp["filepath"].index(path) resp["filepath"][index] = new_path diff --git a/ytdlbot/tasks.py b/ytdlbot/tasks.py index ad7766d..813f91e 100644 --- a/ytdlbot/tasks.py +++ b/ytdlbot/tasks.py @@ -34,8 +34,8 @@ from config import (ARCHIVE_ID, AUDIO_FORMAT, BROKER, ENABLE_CELERY, ENABLE_VIP, TG_MAX_SIZE, WORKERS) from constant import BotText from db import Redis -from downloader import (edit_text, sizeof_fmt, tqdm_progress, upload_hook, - ytdl_download) +from downloader import (edit_text, run_ffmpeg, sizeof_fmt, tqdm_progress, + upload_hook, ytdl_download) from limit import VIP from utils import (apply_log_formatter, auto_restart, customize_logger, get_metadata, get_revision, get_user_settings) @@ -211,24 +211,28 @@ def direct_normal_download(bot_msg, client, url): def normal_audio(bot_msg, client): chat_id = bot_msg.chat.id fn = getattr(bot_msg.video, "file_name", None) or getattr(bot_msg.document, "file_name", None) + status_msg = bot_msg.reply_text("Converting to audio...please wait patiently", quote=True) with tempfile.TemporaryDirectory(prefix="ytdl-") as tmp: logging.info("downloading to %s", tmp) base_path = pathlib.Path(tmp) video_path = base_path.joinpath(fn) audio = base_path.joinpath(fn).with_suffix(f".{AUDIO_FORMAT}") client.send_chat_action(chat_id, 'record_video_note') + status_msg.edit_text("Preparing your conversion....") client.download_media(bot_msg, video_path) logging.info("downloading complete %s", video_path) # execute ffmpeg client.send_chat_action(chat_id, 'record_audio') try: - subprocess.check_output(["ffmpeg", "-y", "-i", video_path, "-vn", "-acodec", "copy", audio]) + run_ffmpeg(["ffmpeg", "-y", "-i", video_path, "-vn", "-acodec", "copy", audio], status_msg) except subprocess.CalledProcessError: # CPU consuming if re-encoding. - subprocess.check_output(["ffmpeg", "-y", "-i", video_path, audio]) + run_ffmpeg(["ffmpeg", "-y", "-i", video_path, audio], status_msg) + status_msg.edit_text("Sending audio now...") client.send_chat_action(chat_id, 'upload_audio') client.send_audio(chat_id, audio) + status_msg.edit_text("✅ Conversion complete.") Redis().update_metrics("audio_success") diff --git a/ytdlbot/ytdl_bot.py b/ytdlbot/ytdl_bot.py index c1e7b0a..1a34f21 100644 --- a/ytdlbot/ytdl_bot.py +++ b/ytdlbot/ytdl_bot.py @@ -266,7 +266,6 @@ def download_handler(client: "Client", message: "types.Message"): red.update_metrics("video_request") text = bot_text.get_receive_link_text() - time.sleep(random.random() / 2) try: # raise pyrogram.errors.exceptions.FloodWait(10) bot_msg: typing.Union["types.Message", "typing.Any"] = message.reply_text(text, quote=True) @@ -308,8 +307,8 @@ def audio_callback(client: "Client", callback_query: types.CallbackQuery): callback_query.answer(f"Converting to audio...please wait patiently") Redis().update_metrics("audio_request") - msg = callback_query.message - audio_entrance(msg, client) + vmsg = callback_query.message + audio_entrance(vmsg, client) @app.on_callback_query(filters.regex(r"Local|Celery"))