From eb63521ff98e76d08775ab5cb43fe485c4fe1b3b Mon Sep 17 00:00:00 2001 From: zenfyr Date: Mon, 10 Aug 2026 07:10:56 +0700 Subject: [PATCH] media: use disk instead of ram --- atproto/models.py | 5 +- atproto/xrpc.py | 8 ++- bluesky/client.py | 5 +- bluesky/input.py | 6 ++- bluesky/output.py | 22 ++++---- cross/media.py | 125 ++++++++++++++++++++++++++------------------- mastodon/input.py | 4 +- mastodon/output.py | 45 ++++++++-------- misskey/input.py | 4 +- 9 files changed, 127 insertions(+), 97 deletions(-) diff --git a/atproto/models.py b/atproto/models.py index 79b152d..ef99efa 100644 --- a/atproto/models.py +++ b/atproto/models.py @@ -3,6 +3,7 @@ from abc import ABC, abstractmethod from dataclasses import dataclass, field from typing import Any +from cross.media import Blob from util.util import Result @@ -126,7 +127,7 @@ class Facet: @dataclass(kw_only=True) class ImageEmbed: - image: bytes + image: Blob alt: str | None = None aspect_ratio: tuple[int, int] | None = None @@ -153,7 +154,7 @@ class GalleryImage(ImageEmbed): @dataclass(kw_only=True) class VideoEmbed: - video: bytes + video: Blob alt: str | None = None aspect_ratio: tuple[int, int] | None = None diff --git a/atproto/xrpc.py b/atproto/xrpc.py index d37bfec..ab71d45 100644 --- a/atproto/xrpc.py +++ b/atproto/xrpc.py @@ -4,6 +4,7 @@ from typing import Any, TypeVar import httpx from atproto.store import AtprotoStore, IdentityInfo, Session +from cross.media import Blob from util.util import LOGGER, normalize_service_url @@ -176,7 +177,7 @@ class XRPCClient: def upload_blob( self, - blob: bytes, + blob: Blob, content_type: str, did: str, ) -> dict[str, Any]: @@ -191,9 +192,12 @@ class XRPCClient: } try: - response = self.http.post(url, content=blob, headers=headers, timeout=60) + with open(blob.path, "rb") as f: + response = self.http.post(url, content=f, headers=headers, timeout=60) except httpx.RequestError as e: raise XRPCError(f"Blob upload request failed: {e}") from e + except OSError as e: + raise XRPCError(f"Could not read blob file: {e}") from e if response.status_code != 200: error_data = response.json() if response.content else {} diff --git a/bluesky/client.py b/bluesky/client.py index d441ec2..0fcddf4 100644 --- a/bluesky/client.py +++ b/bluesky/client.py @@ -22,6 +22,7 @@ from atproto.models import ( ) from atproto.store import AtprotoStore from atproto.xrpc import XRPCClient, XRPCError, resolve_identity +from cross.media import Blob from util.util import normalize_service_url @@ -52,7 +53,7 @@ class BlueskyClient: return time_iso return datetime.now(UTC).isoformat().replace("+00:00", "Z") - def _upload_blob(self, data: bytes, content_type: str) -> dict[str, Any]: + def _upload_blob(self, data: Blob, content_type: str) -> dict[str, Any]: return self.xrpc.upload_blob(data, content_type, self.did) def send_post( @@ -147,7 +148,7 @@ class BlueskyClient: def send_video( self, text: str, - video: bytes, + video: Blob, alt: str | None = None, aspect_ratio: tuple[int, int] | None = None, facets: list[Facet] | None = None, diff --git a/bluesky/input.py b/bluesky/input.py index 1506646..d87833e 100644 --- a/bluesky/input.py +++ b/bluesky/input.py @@ -20,7 +20,7 @@ from cross.attachments import ( QuoteAttachment, RemoteUrlAttachment, ) -from cross.media import Blob, download_blob +from cross.media import Blob, cleanup_blobs, download_blob from cross.post import Post, PostRef from cross.service import InputService from database.connection import DatabasePool @@ -152,8 +152,8 @@ class BlueskyBaseInputService(BlueskyService, InputService, ABC): for a in attachments: post.attachments.put(a) + blobs: list[Blob] = [] if blob_urls: - blobs: list[Blob] = [] for url, cid, alt in blob_urls: self.log.info("Downloading '%s'...", cid) blob = download_blob(url, alt, client=self.http) @@ -203,6 +203,8 @@ class BlueskyBaseInputService(BlueskyService, InputService, ABC): for out in self.outputs: self.submitter(lambda: out.accept_post(post)) + self.submitter(lambda: cleanup_blobs(blobs)) + def _on_repost(self, record: dict[str, Any]): post_uri = cast(str, record["$xpost.strongRef"]["uri"]) post_cid = cast(str, record["$xpost.strongRef"]["cid"]) diff --git a/bluesky/output.py b/bluesky/output.py index 7246f26..4d70da1 100644 --- a/bluesky/output.py +++ b/bluesky/output.py @@ -357,7 +357,7 @@ class BlueskyOutputService(BlueskyService, OutputService): return for blob in supported_media: - if blob.mime.startswith("image/") and len(blob.io) > 2_000_000: + if blob.mime.startswith("image/") and blob.size() > 2_000_000: self.log.error( "Skipping '%s': image too large", post.id, @@ -370,7 +370,7 @@ class BlueskyOutputService(BlueskyService, OutputService): post.id, ) return - if len(blob.io) > 100_000_000: + if blob.size() > 100_000_000: self.log.error( "Skipping '%s': video too large", post.id, @@ -434,21 +434,19 @@ class BlueskyOutputService(BlueskyService, OutputService): images: list[ImageEmbed] = [] images_len: int = len(pwa["attachments"]) for img_blob in pwa["attachments"]: - image_io = img_blob.io - if len(image_io) > 2_000_000: + if img_blob.size() > 2_000_000: self.log.info("Compressing %s...", img_blob.name or "image") - compressed = compress_image(img_blob) - image_io = compressed.io + img_blob = compress_image(img_blob) try: - meta = get_media_meta(image_io) + meta = get_media_meta(img_blob.path) aspect_ratio = (meta.width, meta.height) except Exception as e: self.log.error(e) aspect_ratio = None final_embed = ImageEmbed( - image=image_io, + image=img_blob, alt=img_blob.alt, aspect_ratio=aspect_ratio, ) @@ -467,15 +465,13 @@ class BlueskyOutputService(BlueskyService, OutputService): ) else: video_blob = pwa["attachments"][0] - video_io = video_blob.io if video_blob.mime != "video/mp4": self.log.info("Converting %s to mp4...", video_blob.name or "video") - converted = convert_to_mp4(video_blob) - video_io = converted.io + video_blob = convert_to_mp4(video_blob) try: - meta = get_media_meta(video_io) + meta = get_media_meta(video_blob.path) aspect_ratio = (meta.width, meta.height) duration = meta.duration except Exception as e: @@ -493,7 +489,7 @@ class BlueskyOutputService(BlueskyService, OutputService): response = self._client.send_video( text=text or "", - video=video_io, + video=video_blob, alt=video_blob.alt, aspect_ratio=aspect_ratio, embed=embed, diff --git a/cross/media.py b/cross/media.py index 487709c..6661ae6 100644 --- a/cross/media.py +++ b/cross/media.py @@ -2,8 +2,10 @@ import json import os import re import subprocess +import tempfile import urllib.parse -from dataclasses import dataclass, field +from dataclasses import dataclass +from pathlib import Path from typing import Any, cast import httpx @@ -19,11 +21,14 @@ MAGIC = magic.Magic(mime=True) @dataclass class Blob: - url: str - mime: str - io: bytes = field(repr=False) - name: str | None = None - alt: str | None = None + url: str # remote url + mime: str # mime type + path: Path # on disk path + name: str | None = None # filename + alt: str | None = None # alt text/description + + def size(self) -> int: + return self.path.stat().st_size @dataclass @@ -33,8 +38,8 @@ class MediaInfo: duration: float | None = None -def mime_from_bytes(io: bytes) -> str: - mime = MAGIC.from_buffer(io) +def mime_from_file(io: Path) -> str: + mime = MAGIC.from_file(io) if not mime: mime = "application/octet-stream" return str(mime) @@ -47,35 +52,51 @@ def download_blob( client: httpx.Client | None = None, ) -> Result[Blob, str]: name = get_filename_from_url(url, client) - io = download_chuncked(url, max_bytes, client) - if not io.is_ok(): - return Result.err(io.error()) - return Result.ok(Blob(url, mime_from_bytes(io.value()), io.value(), name, alt)) + dest = download_chuncked(url, max_bytes, client) + if not dest.is_ok(): + return Result.err(dest.error()) + return Result.ok(Blob(url, mime_from_file(dest.value()), dest.value(), name, alt)) + + +def cleanup_blobs(blobs: list[Blob]): + if not blobs: + return + + for blob in blobs: + blob.path.unlink(missing_ok=True) def download_chuncked( url: str, max_bytes: int = 100_000_000, client: httpx.Client | None = None -) -> Result[bytes, str]: +) -> Result[Path, str]: if client is None: client = httpx.Client() + + with ( + tempfile.NamedTemporaryFile( + prefix="xpost-", suffix=".bin", delete=False + ) as download_tmp, + ): + download_path = Path(download_tmp.name) + with client.stream("GET", url, timeout=20) as response: if response.status_code != 200: return Result.err(f"HTTP {response.status_code}: {response.text}") - downloaded_bytes = b"" current_size = 0 - for chunk in response.iter_bytes(chunk_size=8192): - if not chunk: - continue + with open(download_path, "wb") as f: + for chunk in response.iter_bytes(chunk_size=8192): + if not chunk: + continue - current_size += len(chunk) - if current_size > max_bytes: - return Result.err(f"'{url}' larger than max_bytes ({max_bytes})") + current_size += len(chunk) + if current_size > max_bytes: + return Result.err(f"'{url}' larger than max_bytes ({max_bytes})") - downloaded_bytes += chunk + f.write(chunk) - return Result.ok(downloaded_bytes) + return Result.ok(download_path) def get_filename_from_url(url: str, client: httpx.Client | None = None) -> str: @@ -104,69 +125,69 @@ def get_filename_from_url(url: str, client: httpx.Client | None = None) -> str: def convert_to_mp4(video: Blob) -> Blob: + with ( + tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as out_tmp, + ): + out_path = Path(out_tmp.name) + cmd = [ env.FFMPEG_PATH, - "-i", "pipe:0", + "-i", str(video.path), "-c:v", "copy", "-c:a", "aac", "-b:a", "128k", "-movflags", "frag_keyframe+empty_moov+default_base_moof", "-f", "mp4", - "pipe:1", + str(out_path), ] # fmt: skip - proc = subprocess.Popen( - cmd, stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE - ) - out_bytes, err = proc.communicate(input=video.io) - - if proc.returncode != 0: - raise RuntimeError(f"ffmpeg compress failed: {err.decode()}") + result = subprocess.run(cmd, capture_output=True, text=True) + if result.returncode != 0: + raise RuntimeError(f"ffmpeg compress failed: {result.stderr}") - return Blob(video.url, mime_from_bytes(out_bytes), out_bytes, video.name, video.alt) + return Blob(video.url, mime_from_file(out_path), out_path, video.name, video.alt) def compress_image(image: Blob, quality: int = 95) -> Blob: + with ( + tempfile.NamedTemporaryFile(suffix=".webp", delete=False) as out_tmp, + ): + out_path = Path(out_tmp.name) + cmd = [ env.FFMPEG_PATH, - "-f", "image2pipe", - "-i", "pipe:0", + "-i", str(image.path), "-c:v", "webp", "-q:v", str(quality), - "-f", "image2pipe", - "pipe:1", + str(out_path), ] # fmt: skip - proc = subprocess.Popen( - cmd, stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE - ) - out_bytes, err = proc.communicate(input=image.io) - - if proc.returncode != 0: - raise RuntimeError(f"ffmpeg compress failed: {err.decode()}") + result = subprocess.run(cmd, capture_output=True, text=True) + if result.returncode != 0: + raise RuntimeError(f"ffmpeg compress failed: {result.stderr}") - return Blob(image.url, "image/webp", out_bytes, image.name, image.alt) + return Blob(image.url, "image/webp", out_path, image.name, image.alt) -def probe_bytes(bytes: bytes) -> dict[str, Any]: +def probe_file(path: Path) -> dict[str, Any]: cmd = [ env.FFPROBE_PATH, "-v", "error", "-show_format", "-show_streams", "-print_format", "json", - "pipe:0", + str(path), ] # fmt: skip - proc = subprocess.run(cmd, input=bytes, capture_output=True) - if proc.returncode != 0: - raise RuntimeError(f"ffprobe failed: {proc.stderr.decode()}") + result = subprocess.run(cmd, capture_output=True, text=True) + if result.returncode != 0: + raise RuntimeError(f"ffprobe failed: {result.stderr}") - return json.loads(proc.stdout) # type: ignore[no-any-return] + return json.loads(result.stdout) # type: ignore[no-any-return] -def get_media_meta(bytes: bytes) -> MediaInfo: - probe = probe_bytes(bytes) +def get_media_meta(path: Path) -> MediaInfo: + probe = probe_file(path) streams = [s for s in probe["streams"] if s["codec_type"] == "video"] if not streams: raise ValueError("No video stream found") diff --git a/mastodon/input.py b/mastodon/input.py index c607287..53fcc3c 100644 --- a/mastodon/input.py +++ b/mastodon/input.py @@ -14,7 +14,7 @@ from cross.attachments import ( RemoteUrlAttachment, SensitiveAttachment, ) -from cross.media import Blob, download_blob +from cross.media import Blob, cleanup_blobs, download_blob from cross.post import Post, PostRef from cross.service import InputService from database.connection import DatabasePool @@ -199,6 +199,8 @@ class MastodonInputService(MastodonService, InputService): for out in self.outputs: self.submitter(lambda: out.accept_post(post)) + self.submitter(lambda: cleanup_blobs(blobs)) + def _on_reblog(self, status: dict[str, Any], reblog: dict[str, Any]): reposted = self._get_post(self.url, self.user_id, reblog["id"]) if not reposted: diff --git a/mastodon/output.py b/mastodon/output.py index 2e95363..fde5088 100644 --- a/mastodon/output.py +++ b/mastodon/output.py @@ -168,46 +168,47 @@ class MastodonOutputService(MastodonService, OutputService): for blob in attachments: if ( blob.mime.startswith("image/") - and len(blob.io) > self.instance_info.image_size_limit + and blob.size() > self.instance_info.image_size_limit ): return Result.err( - f"image too large: {len(blob.io)} bytes (limit: {self.instance_info.image_size_limit})" + f"image too large: {blob.size()} bytes (limit: {self.instance_info.image_size_limit})" ) if ( blob.mime.startswith("video/") - and len(blob.io) > self.instance_info.video_size_limit + and blob.size() > self.instance_info.video_size_limit ): return Result.err( - f"video too large: {len(blob.io)} bytes (limit: {self.instance_info.video_size_limit})" + f"video too large: {blob.size()} bytes (limit: {self.instance_info.video_size_limit})" ) if ( not blob.mime.startswith(("image/", "video/")) - and len(blob.io) > 7_000_000 + and blob.size() > 7_000_000 ): return Result.err( - f"file too large: {len(blob.io)} bytes (limit: 7000000)" + f"file too large: {blob.size()} bytes (limit: 7000000)" ) uploads: list[MediaUploadResult] = [] for blob in attachments: - files = { - "file": ( - blob.name or "file", - blob.io, - blob.mime, - ) - } - data = {} - if blob.alt: - data["description"] = blob.alt + with open(blob.path, "rb") as f: + files = { + "file": ( + blob.name or "file", + f, + blob.mime, + ) + } + data = {} + if blob.alt: + data["description"] = blob.alt - response = self.http.post( - f"{self.url}/api/v2/media", - headers={"Authorization": f"Bearer {self._get_token()}"}, - files=files, - data=data, - ) + response = self.http.post( + f"{self.url}/api/v2/media", + headers={"Authorization": f"Bearer {self._get_token()}"}, + files=files, + data=data, + ) if response.status_code == 200: self.log.info( diff --git a/misskey/input.py b/misskey/input.py index 8651f3c..cc4815b 100644 --- a/misskey/input.py +++ b/misskey/input.py @@ -14,7 +14,7 @@ from cross.attachments import ( RemoteUrlAttachment, SensitiveAttachment, ) -from cross.media import Blob, download_blob +from cross.media import Blob, cleanup_blobs, download_blob from cross.post import Post, PostRef from cross.service import InputService from database.connection import DatabasePool @@ -191,6 +191,8 @@ class MisskeyInputService(MisskeyService, InputService): for out in self.outputs: self.submitter(lambda: out.accept_post(post)) + self.submitter(lambda: cleanup_blobs(blobs)) + def _on_renote(self, note: dict[str, Any], renote: dict[str, Any]): reposted = self._get_post(self.url, self.user_id, renote["id"]) if not reposted: -- 2.51.2