From 751287479a1d0ee8aa68f9c09ea60ec5d7cb8d63 Mon Sep 17 00:00:00 2001 From: nexy7574 Date: Mon, 15 Apr 2024 22:28:38 +0100 Subject: [PATCH] Properly take advantage of HW accel --- src/cogs/auto_responder.py | 155 +++++++++++++++++++++++-------------- 1 file changed, 99 insertions(+), 56 deletions(-) diff --git a/src/cogs/auto_responder.py b/src/cogs/auto_responder.py index 0026126..37740e3 100644 --- a/src/cogs/auto_responder.py +++ b/src/cogs/auto_responder.py @@ -2,6 +2,7 @@ import asyncio import pathlib import subprocess +import aiohttp import discord import re import tempfile @@ -66,37 +67,67 @@ class AutoResponder(commands.Cog): else: self.log.info("No HEVC streams found in %s", uri) return - with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as tmp: - tmp_path = pathlib.Path(tmp.name) - self.log.info("Transcoding %r to %r", uri, tmp_path) - args = [ - "-i", str(uri), - "-c:v", "libx264", - "-crf", "25", - "-c:a", "libopus", - "-b:a", "64k", - "-preset", "slower", - "-y" - ] - process = await asyncio.create_subprocess_exec( - "ffmpeg", - *args, - str(tmp_path), - stdout=asyncio.subprocess.PIPE, - stderr=asyncio.subprocess.PIPE, - ) - stdout, stderr = await process.communicate() - self.log.info("finished transcode with return code %d", process.returncode) - self.log.debug("stdout: %r", stdout.decode) - self.log.debug("stderr: %r", stderr.decode) - if process.returncode != 0: - raise subprocess.SubprocessError( - process.returncode, - " ".join(args), - stdout, - stderr, + extension = pathlib.Path(uri).suffix + with tempfile.NamedTemporaryFile(suffix=extension) as tmp_dl: + self.log.info("Downloading %r to %r", uri, tmp_dl.name) + async with aiohttp.ClientSession( + headers={ + "Accept": "text/html,application/xhtml+xml,application/xml;q=0.9,image/avif,image/webp," + "*/*;q=0.8", + "Accept-Language": "en-US,en;q=0.5", + "Sec-Fetch-Dest": "document", + "Sec-Fetch-Mode": "navigate", + "Sec-Fetch-Site": "none", + "Sec-Fetch-User": "?1", + "Sec-Gpc": "1", + "Upgrade-Insecure-Requests": "1", + "User-Agent": "Mozilla/5.0 (X11; Linux x86_64; rv:124.0) Gecko/20100101 Firefox/124.0", + } + ) as session: + async with session.get(uri) as resp: + resp.raise_for_status() + async for chunk in resp.content.iter_any(): + # noinspection PyTypeChecker + tmp_dl.write(chunk) + self.log.info("Finished downloading %r to %r", uri, tmp_dl.name) + + with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as tmp: + tmp_path = pathlib.Path(tmp.name) + self.log.info("Transcoding %r to %r", uri, tmp_path) + args = [ + "-hide_banner", + "-hwaccel", "auto", + "-i", tmp_dl.name, + "-c:v", "libx264", + "-crf", "25", + "-c:a", "libopus", + "-b:a", "64k", + "-preset", "faster", + "-vsync", "2", + "-pix_fmt", "yuv420p", + "-movflags", "faststart", + "-profile:v", "main", + "-y" + ] + process = await asyncio.create_subprocess_exec( + "ffmpeg", + *args, + str(tmp_path), + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, ) - return discord.File(tmp_path), tmp_path + stdout, stderr = await process.communicate() + self.log.info("finished transcode with return code %d", process.returncode) + self.log.debug("stdout: %r", stdout.decode) + self.log.debug("stderr: %r", stderr.decode) + if process.returncode != 0: + raise subprocess.SubprocessError( + process.returncode, + " ".join(args), + stdout, + stderr, + ) + return discord.File(tmp_path), tmp_path @commands.Cog.listener("on_message") async def auto_responder(self, message: discord.Message): @@ -104,33 +135,45 @@ class AutoResponder(commands.Cog): return # Check for HEVC truth social links and convert into h264 - links = self.extract_links(message.content, "static-assets-1.truthsocial.com") - if links: - for link in links: - if link.lower().endswith( - (".mp4", ".mov", ".qtff", ".mkv", ".asf", ".avi", ".mxf", ".ps", ".ts", ".3gp", ".3g2") - ): - # All containers allowed to contain HEVC + if message.channel.name == "spam" and message.author.id in {1101439218334576742, 1229496078726860921}: + links = self.extract_links(message.content, "static-assets-1.truthsocial.com") + if links: + for link in links: + hevc_containers = { + ".mp4", + ".mov", + ".qtff", + ".mkv", + ".asf", + ".avi", + ".mxf", + ".ps", + ".ts", + ".3gp", + ".3g2" + } + # ^ All containers allowed to contain HEVC # per https://en.wikipedia.org/wiki/Comparison_of_video_container_formats - self.log.info("Found link to transcode: %r", link) - try: - async with message.channel.typing(): - _r = await self._transcode_hevc_to_h264(link) - if not _r: - continue - file, _p = _r - if file: - if _p.stat().st_size <= 24.5 * 1024 * 1024: - await message.reply(file=file) - else: - self.log.warning( - "Transcoded file too large: %r (%.2f)MB", - _p, - _p.stat().st_size / 1024 / 1024 - ) - _p.unlink() - except Exception as e: - self.log.error("Failed to transcode %r: %r", link, e) + if link.lower().endswith(tuple(hevc_containers)): + self.log.info("Found link to transcode: %r", link) + try: + async with message.channel.typing(): + _r = await self._transcode_hevc_to_h264(link) + if not _r: + continue + file, _p = _r + if file: + if _p.stat().st_size <= 24.5 * 1024 * 1024: + await message.reply(file=file) + else: + self.log.warning( + "Transcoded file too large: %r (%.2f)MB", + _p, + _p.stat().st_size / 1024 / 1024 + ) + _p.unlink() + except Exception as e: + self.log.error("Failed to transcode %r: %r", link, e) def setup(bot):