perf(video): stream uploads into a single on-disk copy

invokeai/app/api/routers/videos.py

@@ -6,12 +6,13 @@
 from pathlib import Path
 from typing import Annotated, BinaryIO, Optional
 
-from fastapi import Body, HTTPException, Query, Request, Response, UploadFile
+from fastapi import Body, HTTPException, Query, Request, Response
 from fastapi import Path as PathParam
 from fastapi.responses import StreamingResponse
 from fastapi.routing import APIRouter
 from PIL import Image as PILImage
 from pydantic import BaseModel, Field, StringConstraints, ValidationError
+from python_multipart.multipart import MultipartParser, parse_options_header
 from starlette.concurrency import run_in_threadpool
 
 from invokeai.app.api.auth_dependencies import CurrentMediaUserOrDefault, CurrentUserOrDefault
@@ -48,15 +49,17 @@
 # Per-chunk size for HTTP Range responses (1 MB)
 RANGE_CHUNK_SIZE = 1024 * 1024
 
-# Upload streaming chunk size (1 MB) and a coarse per-upload size cap. The cap is generous
+# Coarse per-upload size cap, enforced against the file part as it streams in. Generous
 # because Wan-generated MP4s for long sequences can run into the hundreds of megabytes;
-# the goal is to prevent a single client from exhausting RAM, not to be a content policy.
-UPLOAD_CHUNK_SIZE = 1024 * 1024
+# the goal is to prevent a single client from exhausting RAM/disk, not to be a content policy.
 MAX_UPLOAD_SIZE = 1024 * 1024 * 1024  # 1 GB
 # Pre-parse ingress cap enforced by VideoUploadLimitASGIMiddleware, applied to the whole
-# request body *before* the multipart parser spools it to temp storage. Slightly larger
-# than MAX_UPLOAD_SIZE to allow for multipart framing and the metadata form field.
+# request body *before* the upload route parses it. Slightly larger than MAX_UPLOAD_SIZE
+# to allow for multipart framing and the metadata form field.
 MAX_UPLOAD_REQUEST_SIZE = MAX_UPLOAD_SIZE + 10 * 1024 * 1024
+# The `metadata` form field is a stringified JSON dict; it is buffered in memory while the
+# body streams, so it gets its own (generous) cap.
+MAX_UPLOAD_METADATA_SIZE = 1024 * 1024
 # Global bound on concurrent video uploads — each in-flight upload can hold up to two
 # copies of the file in temp storage (the multipart spool + the route's own tmp file).
 MAX_CONCURRENT_VIDEO_UPLOADS = 2
@@ -163,14 +166,147 @@ def _assert_video_read_access(video_name: str, current_user: CurrentUserOrDefaul
     raise HTTPException(status_code=403, detail="Not authorized to access this video")
 
 
-def _is_accepted_video_upload(file: UploadFile) -> bool:
-    if file.content_type and file.content_type.startswith(ACCEPTED_VIDEO_MIME_PREFIXES):
+def _is_accepted_video_upload(filename: Optional[str], content_type: Optional[str]) -> bool:
+    if content_type and content_type.startswith(ACCEPTED_VIDEO_MIME_PREFIXES):
         return True
-    if file.filename:
-        return file.filename.lower().endswith(ACCEPTED_VIDEO_EXTENSIONS)
+    if filename:
+        return filename.lower().endswith(ACCEPTED_VIDEO_EXTENSIONS)
     return False
 
 
+class _VideoUploadStreamParser:
+    """Parses the multipart upload body, writing the file part straight to `destination`.
+
+    Declaring `file: UploadFile` on the route makes Starlette parse the body into its own
+    spooled temp file first, so the route's copy to a named temp file was a SECOND
+    full-size copy: every in-flight upload occupied up to 2 x MAX_UPLOAD_SIZE of temp
+    storage (x MAX_CONCURRENT_VIDEO_UPLOADS) for the whole probe/thumbnail/create phase.
+    The spool's path cannot be reused instead — once rolled over it is an unlinked
+    anonymous file, and both ffmpeg and videos.create need a real path.
+
+    Parsing the stream ourselves keeps exactly one copy on disk. It also lets the
+    file-type and size checks fire while the body is still arriving, rather than after the
+    whole thing has been written somewhere.
+
+    Callbacks run inside `MultipartParser.write`, which the route calls in a worker thread
+    — the disk writes must not happen on the event loop.
+    """
+
+    def __init__(self, destination: BinaryIO) -> None:
+        self._destination = destination
+        self._header_field = bytearray()
+        self._header_value = bytearray()
+        self._headers: dict[bytes, bytes] = {}
+        self._part_name: Optional[bytes] = None
+        self._metadata_chunks: list[bytes] = []
+        self._metadata_size = 0
+        self.filename: Optional[str] = None
+        self.content_type: Optional[str] = None
+        self.metadata: Optional[str] = None
+        self.file_size = 0
+        self.saw_file_part = False
+
+    @property
+    def callbacks(self) -> dict[str, object]:
+        return {
+            "on_part_begin": self._on_part_begin,
+            "on_part_data": self._on_part_data,
+            "on_part_end": self._on_part_end,
+            "on_header_field": self._on_header_field,
+            "on_header_value": self._on_header_value,
+            "on_header_end": self._on_header_end,
+            "on_headers_finished": self._on_headers_finished,
+        }
+
+    def _on_part_begin(self) -> None:
+        self._headers = {}
+        self._header_field = bytearray()
+        self._header_value = bytearray()
+        self._part_name = None
+        self._metadata_chunks = []
+        self._metadata_size = 0
+
+    def _on_header_field(self, data: bytes, start: int, end: int) -> None:
+        self._header_field.extend(data[start:end])
+
+    def _on_header_value(self, data: bytes, start: int, end: int) -> None:
+        self._header_value.extend(data[start:end])
+
+    def _on_header_end(self) -> None:
+        self._headers[bytes(self._header_field).lower()] = bytes(self._header_value)
+        self._header_field = bytearray()
+        self._header_value = bytearray()
+
+    def _on_headers_finished(self) -> None:
+        _, options = parse_options_header(self._headers.get(b"content-disposition", b""))
+        self._part_name = options.get(b"name")
+        if self._part_name != b"file":
+            return
+        if self.saw_file_part:
+            raise HTTPException(status_code=422, detail="Expected exactly one video file")
+        self.saw_file_part = True
+        filename = options.get(b"filename")
+        self.filename = filename.decode("utf-8", errors="replace") if filename is not None else None
+        content_type = self._headers.get(b"content-type")
+        self.content_type = content_type.decode("latin-1") if content_type is not None else None
+        # Reject the wrong kind of file before any of its bytes reach the disk.
+        if not _is_accepted_video_upload(self.filename, self.content_type):
+            raise HTTPException(status_code=415, detail="Not a supported video file")
+
+    def _on_part_data(self, data: bytes, start: int, end: int) -> None:
+        chunk = data[start:end]
+        if self._part_name == b"file":
+            self.file_size += len(chunk)
+            if self.file_size > MAX_UPLOAD_SIZE:
+                raise HTTPException(
+                    status_code=413,
+                    detail=f"Video upload exceeds maximum size ({MAX_UPLOAD_SIZE} bytes)",
+                )
+            self._destination.write(chunk)
+        elif self._part_name == b"metadata":
+            self._metadata_size += len(chunk)
+            if self._metadata_size > MAX_UPLOAD_METADATA_SIZE:
+                raise HTTPException(
+                    status_code=413,
+                    detail=f"Video metadata exceeds maximum size ({MAX_UPLOAD_METADATA_SIZE} bytes)",
+                )
+            self._metadata_chunks.append(chunk)
+        # Any other field is dropped rather than buffered: an unknown part must not be a
+        # way to make the server hold arbitrary bytes in memory.
+
+    def _on_part_end(self) -> None:
+        if self._part_name == b"metadata":
+            try:
+                self.metadata = b"".join(self._metadata_chunks).decode("utf-8")
+            except UnicodeDecodeError as error:
+                raise HTTPException(status_code=422, detail="Metadata must be UTF-8 encoded") from error
+        self._part_name = None
+        self._metadata_chunks = []
+        self._metadata_size = 0
+
+
+async def _stream_video_upload(request: Request, destination: BinaryIO) -> _VideoUploadStreamParser:
+    """Streams the request body through the multipart parser into `destination`."""
+    media_type, options = parse_options_header(request.headers.get("content-type", ""))
+    boundary = options.get(b"boundary")
+    if media_type != b"multipart/form-data" or boundary is None:
+        raise HTTPException(status_code=422, detail="Expected a multipart/form-data video upload")
+
+    parser_state = _VideoUploadStreamParser(destination)
+    # max_size is the pre-parse ingress cap the middleware already enforces; repeating it
+    # here bounds the parser itself for any path that reaches it directly.
+    parser = MultipartParser(boundary, parser_state.callbacks, max_size=MAX_UPLOAD_REQUEST_SIZE)
+    async for chunk in request.stream():
+        # Parsing writes to disk, so it belongs in the thread pool alongside the rest of
+        # the blocking upload work.
+        await run_in_threadpool(parser.write, chunk)
+    await run_in_threadpool(parser.finalize)
+
+    if not parser_state.saw_file_part:
+        raise HTTPException(status_code=422, detail="Expected a video file in the upload")
+    return parser_state
+
+
 def _is_mp4_file(path: Path) -> bool:
     try:
         with open(path, "rb") as video_file:
@@ -230,29 +366,43 @@ def _probe_decodable_video(path: Path) -> tuple[tuple[int, int, float, Optional[
     },
     status_code=201,
     response_model=VideoDTO,
+    # The body is parsed by hand (see _stream_video_upload) so the file lands in exactly
+    # one temp file, which means FastAPI cannot infer the request schema from the
+    # signature. This spells out the same multipart body the `file` + `metadata`
+    # parameters used to generate, so the documented contract is unchanged.
+    openapi_extra={
+        "requestBody": {
+            "required": True,
+            "content": {
+                "multipart/form-data": {
+                    "schema": {
+                        "title": "Body_upload_video",
+                        "type": "object",
+                        "properties": {
+                            "file": {"title": "File", "type": "string", "format": "binary"},
+                            "metadata": {
+                                "title": "Metadata",
+                                "description": "The metadata to associate with the video, must be a stringified JSON dict",
+                                "anyOf": [{"type": "string"}, {"type": "null"}],
+                            },
+                        },
+                        "required": ["file"],
+                    }
+                }
+            },
+        }
+    },
 )
 async def upload_video(
     current_user: CurrentUserOrDefault,
-    file: UploadFile,
     request: Request,
     response: Response,
     video_category: ImageCategory = Query(description="The category of the video"),
     is_intermediate: bool = Query(description="Whether this is an intermediate video"),
     board_id: Optional[str] = Query(default=None, description="The board to add this video to, if any"),
     session_id: Optional[str] = Query(default=None, description="The session ID associated with this upload, if any"),
-    metadata: Optional[str] = Body(
-        default=None,
-        description="The metadata to associate with the video, must be a stringified JSON dict",
-        embed=True,
-    ),
 ) -> VideoDTO:
     """Uploads a video for the current user."""
-    if metadata is not None:
-        try:
-            MetadataFieldValidator.validate_json(metadata)
-        except ValidationError as e:
-            raise HTTPException(status_code=422, detail="Metadata must be a JSON object") from e
-
     # Check board access for uploads to a specific board.
     if board_id is not None:
         from invokeai.app.services.board_records.board_records_common import BoardVisibility
@@ -268,35 +418,24 @@ async def upload_video(
         ):
             raise HTTPException(status_code=403, detail="Not authorized to upload to this board")
 
-    if not _is_accepted_video_upload(file):
-        raise HTTPException(status_code=415, detail="Not a supported video file")
-
-    # Stream the upload to a tmp file so we can probe and then hand its path to the service.
-    # Reading the full body into memory first risked exhausting RAM on multi-GB uploads;
-    # chunk-stream instead and enforce a hard size cap. Filesystem writes, container
-    # validation, ffmpeg probing, and thumbnail extraction are all blocking — run them in
-    # the thread pool so a slow (or hostile) file can't stall the event loop and every
-    # other API request with it.
+    # Stream the upload straight into a tmp file so we can probe it and then hand its path
+    # to the service. Reading the full body into memory first risked exhausting RAM on
+    # multi-GB uploads; the parser streams it instead and enforces a hard size cap as the
+    # bytes arrive. Filesystem writes, container validation, ffmpeg probing, and thumbnail
+    # extraction are all blocking — run them in the thread pool so a slow (or hostile) file
+    # can't stall the event loop and every other API request with it.
     tmp = tempfile.NamedTemporaryFile(prefix="invokeai_upload_", suffix=".mp4", delete=False)
     tmp_path = Path(tmp.name)
     try:
-        total = 0
-        while chunk := await file.read(UPLOAD_CHUNK_SIZE):
-            total += len(chunk)
-            if total > MAX_UPLOAD_SIZE:
-                tmp.close()
-                raise HTTPException(
-                    status_code=413,
-                    detail=f"Video upload exceeds maximum size ({MAX_UPLOAD_SIZE} bytes)",
-                )
-            await run_in_threadpool(tmp.write, chunk)
+        upload = await _stream_video_upload(request, tmp)
         tmp.close()
-        # Release the multipart spool now that the body is copied: each in-flight
-        # upload otherwise holds TWO on-disk copies (Starlette's spool + our tmp file)
-        # through the probe/thumbnail/create phase — up to 2 x MAX_UPLOAD_SIZE x
-        # MAX_CONCURRENT_VIDEO_UPLOADS of temp disk. Closing shrinks the double-copy
-        # window to the copy loop itself.
-        await file.close()
+
+        metadata = upload.metadata
+        if metadata is not None:
+            try:
+                MetadataFieldValidator.validate_json(metadata)
+            except ValidationError as e:
+                raise HTTPException(status_code=422, detail="Metadata must be a JSON object") from e
 
         if not await run_in_threadpool(_is_mp4_file, tmp_path):
             raise HTTPException(status_code=415, detail="Not an MP4 video file")

invokeai/frontend/web/openapi.json

@@ -6191,16 +6191,6 @@
             "description": "The session ID associated with this upload, if any"
           }
         ],
-        "requestBody": {
-          "required": true,
-          "content": {
-            "multipart/form-data": {
-              "schema": {
-                "$ref": "#/components/schemas/Body_upload_video"
-              }
-            }
-          }
-        },
         "responses": {
           "201": {
             "description": "The video was uploaded successfully",
@@ -6225,6 +6215,37 @@
               }
             }
           }
+        },
+        "requestBody": {
+          "required": true,
+          "content": {
+            "multipart/form-data": {
+              "schema": {
+                "title": "Body_upload_video",
+                "type": "object",
+                "properties": {
+                  "file": {
+                    "title": "File",
+                    "type": "string",
+                    "format": "binary"
+                  },
+                  "metadata": {
+                    "title": "Metadata",
+                    "description": "The metadata to associate with the video, must be a stringified JSON dict",
+                    "anyOf": [
+                      {
+                        "type": "string"
+                      },
+                      {
+                        "type": "null"
+                      }
+                    ]
+                  }
+                },
+                "required": ["file"]
+              }
+            }
+          }
         }
       }
     },
@@ -16129,30 +16150,6 @@
         "required": ["file"],
         "title": "Body_upload_image"
       },
-      "Body_upload_video": {
-        "properties": {
-          "file": {
-            "type": "string",
-            "format": "binary",
-            "title": "File"
-          },
-          "metadata": {
-            "anyOf": [
-              {
-                "type": "string"
-              },
-              {
-                "type": "null"
-              }
-            ],
-            "title": "Metadata",
-            "description": "The metadata to associate with the video, must be a stringified JSON dict"
-          }
-        },
-        "type": "object",
-        "required": ["file"],
-        "title": "Body_upload_video"
-      },
       "BooleanCollectionInvocation": {
         "category": "primitives",
         "class": "invocation",