Repository: gradio-app/gradio

Issue:
Improve /upload experience while file is still loading (?)
From the 4.0 tracking issue.




Discussion:

Resolving patch:
diff --git a/gradio/route_utils.py b/gradio/route_utils.py
--- a/gradio/route_utils.py
+++ b/gradio/route_utils.py
@@ -2,6 +2,8 @@
 
 import hashlib
 import json
+from collections import deque
+from dataclasses import dataclass as python_dataclass
 from tempfile import NamedTemporaryFile, _TemporaryFileWrapper
 from typing import TYPE_CHECKING, AsyncGenerator, BinaryIO, List, Optional, Tuple, Union
 
@@ -294,6 +296,44 @@ def __init__(
         self.sha = hashlib.sha1()
 
 
+@python_dataclass(frozen=True)
+class FileUploadProgressUnit:
+    filename: str
+    chunk_size: int
+    is_done: bool
+
+
+class FileUploadProgress:
+    def __init__(self) -> None:
+        self._statuses: dict[str, deque[FileUploadProgressUnit]] = {}
+
+    def track(self, upload_id: str):
+        if upload_id not in self._statuses:
+            self._statuses[upload_id] = deque()
+
+    def update(self, upload_id: str, filename: str, message_bytes: bytes):
+        if upload_id not in self._statuses:
+            self._statuses[upload_id] = deque()
+        self._statuses[upload_id].append(
+            FileUploadProgressUnit(filename, len(message_bytes), is_done=False)
+        )
+
+    def set_done(self, upload_id: str):
+        self._statuses[upload_id].append(FileUploadProgressUnit("", 0, is_done=True))
+
+    def stop_tracking(self, upload_id: str):
+        if upload_id in self._statuses:
+            del self._statuses[upload_id]
+
+    def status(self, upload_id: str) -> deque[FileUploadProgressUnit]:
+        if upload_id not in self._statuses:
+            return deque()
+        return self._statuses[upload_id]
+
+    def is_tracked(self, upload_id: str):
+        return upload_id in self._statuses
+
+
 class GradioMultiPartParser:
     """Vendored from starlette.MultipartParser.
 
@@ -315,6 +355,8 @@ def __init__(
         *,
         max_files: Union[int, float] = 1000,
         max_fields: Union[int, float] = 1000,
+        upload_id: str | None = None,
+        upload_progress: FileUploadProgress | None = None,
     ) -> None:
         assert (
             multipart is not None
@@ -324,6 +366,8 @@ def __init__(
         self.max_files = max_files
         self.max_fields = max_fields
         self.items: List[Tuple[str, Union[str, UploadFile]]] = []
+        self.upload_id = upload_id
+        self.upload_progress = upload_progress
         self._current_files = 0
         self._current_fields = 0
         self._current_partial_header_name: bytes = b""
@@ -339,6 +383,10 @@ def on_part_begin(self) -> None:
 
     def on_part_data(self, data: bytes, start: int, end: int) -> None:
         message_bytes = data[start:end]
+        if self.upload_progress is not None:
+            self.upload_progress.update(
+                self.upload_id, self._current_part.file.filename, message_bytes  # type: ignore
+            )
         if self._current_part.file is None:
             self._current_part.data += message_bytes
         else:
@@ -464,4 +512,6 @@ async def parse(self) -> FormData:
             raise exc
 
         parser.finalize()
+        if self.upload_progress is not None:
+            self.upload_progress.set_done(self.upload_id)  # type: ignore
         return FormData(self.items)
diff --git a/gradio/routes.py b/gradio/routes.py
--- a/gradio/routes.py
+++ b/gradio/routes.py
@@ -56,6 +56,7 @@
 from gradio.oauth import attach_oauth
 from gradio.queueing import Estimation, Event
 from gradio.route_utils import (  # noqa: F401
+    FileUploadProgress,
     GradioMultiPartParser,
     GradioUploadFile,
     MultiPartException,
@@ -121,6 +122,9 @@ def move_uploaded_files_to_cache(files: list[str], destinations: list[str]) -> N
         shutil.move(file, dest)
 
 
+file_upload_statuses = FileUploadProgress()
+
+
 class App(FastAPI):
     """
     FastAPI App Wrapper
@@ -681,8 +685,57 @@ def component_server(body: ComponentServerBody):
         async def get_queue_status():
             return app.get_blocks()._queue.get_estimation()
 
+        @app.get("/upload_progress")
+        def get_upload_progress(upload_id: str, request: fastapi.Request):
+            async def sse_stream(request: fastapi.Request):
+                last_heartbeat = time.perf_counter()
+                is_done = False
+                while True:
+                    if await request.is_disconnected():
+                        file_upload_statuses.stop_tracking(upload_id)
+                        return
+                    if is_done:
+                        file_upload_statuses.stop_tracking(upload_id)
+                        return
+
+                    heartbeat_rate = 15
+                    check_rate = 0.05
+                    message = None
+                    try:
+                        if update := file_upload_statuses.status(upload_id).popleft():
+                            if update.is_done:
+                                message = {"msg": "done"}
+                                is_done = True
+                            else:
+                                message = {
+                                    "msg": "update",
+                                    "orig_name": update.filename,
+                                    "chunk_size": update.chunk_size,
+                                }
+                        else:
+                            await asyncio.sleep(check_rate)
+                            if time.perf_counter() - last_heartbeat > heartbeat_rate:
+                                message = {"msg": "heartbeat"}
+                                last_heartbeat = time.perf_counter()
+                        if message:
+                            yield f"data: {json.dumps(message)}\n\n"
+                    except IndexError:
+                        if not file_upload_statuses.is_tracked(upload_id):
+                            return
+                        # pop from empty queue
+                        continue
+
+            return StreamingResponse(
+                sse_stream(request),
+                media_type="text/event-stream",
+            )
+
         @app.post("/upload", dependencies=[Depends(login_check)])
-        async def upload_file(request: fastapi.Request, bg_tasks: BackgroundTasks):
+        async def upload_file(
+            request: fastapi.Request,
+            bg_tasks: BackgroundTasks,
+            upload_id: Optional[str] = None,
+        ):
             content_type_header = request.headers.get("Content-Type")
             content_type: bytes
             content_type, _ = parse_options_header(content_type_header)
@@ -690,11 +743,15 @@ async def upload_file(request: fastapi.Request, bg_tasks: BackgroundTasks):
                 raise HTTPException(status_code=400, detail="Invalid content type.")
 
             try:
+                if upload_id:
+                    file_upload_statuses.track(upload_id)
                 multipart_parser = GradioMultiPartParser(
                     request.headers,
                     request.stream(),
                     max_files=1000,
                     max_fields=1000,
+                    upload_id=upload_id if upload_id else None,
+                    upload_progress=file_upload_statuses if upload_id else None,
                 )
                 form = await multipart_parser.parse()
             except MultiPartException as exc:


Test patch:
diff --git a/js/app/test/audio_component_events.spec.ts b/js/app/test/audio_component_events.spec.ts
--- a/js/app/test/audio_component_events.spec.ts
+++ b/js/app/test/audio_component_events.spec.ts
@@ -7,7 +7,7 @@ test("Audio click-to-upload uploads audio successfuly.", async ({ page }) => {
 	const uploader = await page.locator("input[type=file]");
 	await Promise.all([
 		uploader.setInputFiles(["../../test/test_files/audio_sample.wav"]),
-		page.waitForResponse("**/upload")
+		page.waitForResponse("**/upload?*")
 	]);
 
 	await expect(page.getByLabel("# Change Events")).toHaveValue("1");
@@ -21,7 +21,7 @@ test("Audio click-to-upload uploads audio successfuly.", async ({ page }) => {
 
 	await Promise.all([
 		uploader.setInputFiles(["../../test/test_files/audio_sample.wav"]),
-		page.waitForResponse("**/upload")
+		page.waitForResponse("**/upload?*")
 	]);
 
 	await expect(page.getByLabel("# Change Events")).toHaveValue("3");
@@ -39,7 +39,7 @@ test("Audio drag-and-drop uploads a file to the server correctly.", async ({
 			"audio_sample.wav",
 			"audio/wav"
 		),
-		page.waitForResponse("**/upload")
+		page.waitForResponse("**/upload?*")
 	]);
 	await expect(page.getByLabel("# Change Events")).toHaveValue("1");
 	await expect(page.getByLabel("# Upload Events")).toHaveValue("1");
diff --git a/js/app/test/video_component_events.spec.ts b/js/app/test/video_component_events.spec.ts
--- a/js/app/test/video_component_events.spec.ts
+++ b/js/app/test/video_component_events.spec.ts
@@ -9,7 +9,7 @@ test("Video click-to-upload uploads video successfuly. Clear, play, and pause bu
 	const uploader = await page.locator("input[type=file]");
 	await Promise.all([
 		uploader.setInputFiles(["./test/files/file_test.ogg"]),
-		page.waitForResponse("**/upload")
+		page.waitForResponse("**/upload?*?*")
 	]);
 
 	await expect(page.getByLabel("# Change Events")).toHaveValue("1");
@@ -28,7 +28,7 @@ test("Video click-to-upload uploads video successfuly. Clear, play, and pause bu
 
 	await Promise.all([
 		uploader.setInputFiles(["./test/files/file_test.ogg"]),
-		page.waitForResponse("**/upload")
+		page.waitForResponse("**/upload?*")
 	]);
 
 	await expect(page.getByLabel("# Change Events")).toHaveValue("3");
@@ -50,7 +50,7 @@ test("Video drag-and-drop uploads a file to the server correctly.", async ({
 		"file_test.ogg",
 		"video/*"
 	);
-	await page.waitForResponse("**/upload");
+	await page.waitForResponse("**/upload?*");
 	await expect(page.getByLabel("# Change Events")).toHaveValue("1");
 	await expect(page.getByLabel("# Upload Events")).toHaveValue("1");
 });