Run Pre/Post processing for components in a separate thread (#13168)

Author: freddyaboulton

.changeset/wild-hounds-laugh.md

@@ -0,0 +1,5 @@
+---
+"gradio": patch
+---
+
+feat:Run Pre/Post processing for components in a separate thread

gradio/blocks.py

@@ -1771,11 +1771,14 @@ async def preprocess_data(
                     inputs[i].get("value", None) if is_prop_input else inputs[i]
                 )
 
-                inputs_cached = await processing_utils.async_move_files_to_cache(
-                    value_to_process,
-                    block,
-                    check_in_upload_folder=not explicit_call,
-                )
+                from gradio.profiling import trace_phase
+
+                async with trace_phase("preprocess_move_to_cache"):
+                    inputs_cached = await processing_utils.async_move_files_to_cache(
+                        value_to_process,
+                        block,
+                        check_in_upload_folder=not explicit_call,
+                    )
                 if getattr(block, "data_model", None) and inputs_cached is not None:
                     data_model = cast(
                         Union[GradioModel, GradioRootModel], block.data_model
@@ -1793,7 +1796,9 @@ async def preprocess_data(
                 state._update_value_in_config(block._id, inputs_serialized)
 
                 if block_fn.preprocess:
-                    processed_value = block.preprocess(inputs_cached)
+                    processed_value = await anyio.to_thread.run_sync(
+                        block.preprocess, inputs_cached, limiter=self.limiter
+                    )
                 else:
                     processed_value = inputs_serialized
 
@@ -1855,6 +1860,8 @@ async def postprocess_data(
         predictions: list | dict,
         state: SessionState | None,
     ) -> list[Any]:
+        from gradio.profiling import trace_phase
+
         state = state or SessionState(self)
         if (
             isinstance(predictions, dict)
@@ -1943,33 +1950,36 @@ async def postprocess_data(
                         )
                     if block._id in state:
                         block = state[block._id]
-                    prediction_value = block.postprocess(prediction_value)
+                    prediction_value = await anyio.to_thread.run_sync(
+                        block.postprocess, prediction_value, limiter=self.limiter
+                    )
                     if isinstance(prediction_value, (GradioModel, GradioRootModel)):
                         prediction_value_serialized = prediction_value.model_dump()
                     else:
                         prediction_value_serialized = prediction_value
-                    prediction_value_serialized = (
-                        await processing_utils.async_move_files_to_cache(
-                            prediction_value_serialized,
-                            block,
-                            postprocess=True,
+                    async with trace_phase("postprocess_update_state_in_config"):
+                        prediction_value_serialized = (
+                            await processing_utils.async_move_files_to_cache(
+                                prediction_value_serialized,
+                                block,
+                                postprocess=True,
+                            )
+                        )
+                        if block._id not in state:
+                            state[block._id] = block
+                        state._update_value_in_config(
+                            block._id, prediction_value_serialized
                         )
-                    )
-                    if block._id not in state:
-                        state[block._id] = block
-                    state._update_value_in_config(
-                        block._id, prediction_value_serialized
-                    )
                 elif not block_fn.postprocess:
                     if block._id not in state:
                         state[block._id] = block
                     state._update_value_in_config(block._id, prediction_value)
-
-                outputs_cached = await processing_utils.async_move_files_to_cache(
-                    prediction_value,
-                    block,
-                    postprocess=True,
-                )
+                async with trace_phase("postprocess_move_to_cache"):
+                    outputs_cached = await processing_utils.async_move_files_to_cache(
+                        prediction_value,
+                        block,
+                        postprocess=True,
+                    )
                 output.append(outputs_cached)
 
         return output

gradio/components/audio.py

@@ -4,6 +4,7 @@
 
 import dataclasses
 import io
+import json
 from collections.abc import Callable, Sequence
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, Literal
@@ -414,9 +415,6 @@ def _process_json_subtitles(
     def _process_subtitle_file(
         self, subtitle_file: str | Path
     ) -> FileData | list[dict[str, Any]]:
-        import json
-        from pathlib import Path
-
         file_path = Path(subtitle_file)
         if file_path.suffix.lower() == ".json":
             try:

gradio/components/video.py

@@ -22,6 +22,7 @@
 from gradio.data_classes import FileData, MediaStreamChunk
 from gradio.events import Events
 from gradio.i18n import I18nData
+from gradio.profiling import trace_phase_sync, traced_sync
 from gradio.utils import get_upload_folder, set_default_buttons
 
 if TYPE_CHECKING:
@@ -189,6 +190,7 @@ def __init__(
         )
         self._value_description = "a string filepath to a video"
 
+    @traced_sync("preprocess_video")
     def preprocess(self, payload: FileData | None) -> str | None:
         """
         Parameters:
@@ -250,6 +252,7 @@ def preprocess(self, payload: FileData | None) -> str | None:
         else:
             return str(file_name)
 
+    @traced_sync("postprocess_video")
     def postprocess(self, value: str | Path | None) -> FileData | None:
         """
         Parameters:
@@ -272,6 +275,7 @@ def _format_video(self, video: str | Path | None) -> FileData | None:
         Processes a video to ensure that it is in the correct format
         and adds a watermark if requested.
         """
+
         if video is None:
             return None
         video = str(video)
@@ -301,7 +305,8 @@ def _format_video(self, video: str | Path | None) -> FileData | None:
             warnings.warn(
                 "Video does not have browser-compatible container or codec. Converting to mp4."
             )
-            video = processing_utils.convert_video_to_playable_mp4(video)
+            with trace_phase_sync("postprocess_video_convert_video_to_playable_mp4"):
+                video = processing_utils.convert_video_to_playable_mp4(video)
         # Recalculate the format in case convert_video_to_playable_mp4 already made it the selected format
         returned_format = utils.get_extension_from_file_path_or_url(video).lower()
         if (
@@ -404,8 +409,6 @@ def _format_subtitles(self, subtitle: str | Path | None) -> FileData | None:
         """
         Convert subtitle format to VTT and process the video to ensure it meets the HTML5 requirements.
         """
-        import json
-        from pathlib import Path
 
         def srt_to_vtt(srt_file_path, vtt_file_path):
             """Convert an SRT subtitle file to a VTT subtitle file"""

gradio/image_utils.py

@@ -17,6 +17,7 @@
 from gradio.components.image_editor import WatermarkOptions
 from gradio.data_classes import ImageData
 from gradio.exceptions import Error
+from gradio.profiling import traced_sync
 
 PIL.Image.init()  # fixes https://github.com/gradio-app/gradio/issues/2843 (remove when requiring Pillow 9.4+)
 
@@ -258,6 +259,7 @@ def extract_svg_content(image_file: str | Path) -> str:
         return svg_content
 
 
+@traced_sync("preprocess_format_image")
 def preprocess_image(
     payload: ImageData | None,
     cache_dir: str,
@@ -311,6 +313,7 @@ def preprocess_image(
             warnings.simplefilter("ignore")
             if image_mode is not None:
                 im = im.convert(image_mode)
+
     return format_image(
         im,
         type=type,

gradio/processing_utils.py

@@ -30,6 +30,7 @@
 from gradio.context import LocalContext
 from gradio.data_classes import FileData, GradioModel, GradioRootModel, JsonData
 from gradio.exceptions import Error, InvalidPathError
+from gradio.profiling import traced_sync
 from gradio.route_utils import API_PREFIX
 from gradio.utils import abspath, get_hash_seed, get_upload_folder, is_in_or_equal
 
@@ -155,6 +156,7 @@ def hash_base64(base64_encoding: str, chunk_num_blocks: int = 128) -> str:
     return sha.hexdigest()
 
 
+@traced_sync("postprocess_save_pil_to_cache")
 def save_pil_to_cache(
     img: Image.Image,
     cache_dir: str,
@@ -169,13 +171,15 @@ def save_pil_to_cache(
     return filename
 
 
+@traced_sync("postprocess_save_img_array_to_cache")
 def save_img_array_to_cache(
     arr: np.ndarray, cache_dir: str, format: str = "webp"
 ) -> str:
     pil_image = Image.fromarray(_convert(arr, np.uint8, force_copy=False))
     return save_pil_to_cache(pil_image, cache_dir, format=format)
 
 
+@traced_sync("postprocess_save_audio_to_cache")
 def save_audio_to_cache(
     data: np.ndarray, sample_rate: int, format: str, cache_dir: str
 ) -> str:
@@ -207,6 +211,7 @@ def detect_audio_format(data: bytes) -> str:
     return ""
 
 
+@traced_sync("postprocess_save_bytes_to_cache")
 def save_bytes_to_cache(data: bytes, file_name: str, cache_dir: str) -> str:
     path = Path(cache_dir) / hash_bytes(data)
     path.mkdir(exist_ok=True, parents=True)
@@ -218,6 +223,7 @@ def save_bytes_to_cache(data: bytes, file_name: str, cache_dir: str) -> str:
     return str(path.resolve())
 
 
+@traced_sync("save_file_to_cache")
 def save_file_to_cache(file_path: str | Path, cache_dir: str) -> str:
     """Returns a temporary file path for a copy of the given file path if it does
     not already exist. Otherwise returns the path to the existing temp file."""
@@ -644,28 +650,31 @@ def resize_and_crop(img, size, crop_type="center"):
 def audio_from_file(
     filename: str, crop_min: float = 0, crop_max: float = 100
 ) -> tuple[int, np.ndarray]:
-    try:
-        audio = AudioSegment.from_file(filename)
-    except FileNotFoundError as e:
-        isfile = Path(filename).is_file()
-        msg = (
-            f"Cannot load audio from file: `{'ffprobe' if isfile else filename}` not found."
-            + " Please install `ffmpeg` in your system to use non-WAV audio file formats"
-            " and make sure `ffprobe` is in your PATH."
-            if isfile
-            else ""
-        )
-        raise RuntimeError(msg) from e
-    except OSError as e:
-        raise e
-    if crop_min != 0 or crop_max != 100:
-        audio_start = len(audio) * crop_min / 100
-        audio_end = len(audio) * crop_max / 100
-        audio = audio[audio_start:audio_end]
-    data = np.array(audio.get_array_of_samples())
-    if audio.channels > 1:
-        data = data.reshape(-1, audio.channels)
-    return audio.frame_rate, data
+    from gradio.profiling import trace_phase_sync
+
+    with trace_phase_sync("preprocess_audio_from_file"):
+        try:
+            audio = AudioSegment.from_file(filename)
+        except FileNotFoundError as e:
+            isfile = Path(filename).is_file()
+            msg = (
+                f"Cannot load audio from file: `{'ffprobe' if isfile else filename}` not found."
+                + " Please install `ffmpeg` in your system to use non-WAV audio file formats"
+                " and make sure `ffprobe` is in your PATH."
+                if isfile
+                else ""
+            )
+            raise RuntimeError(msg) from e
+        except OSError as e:
+            raise e
+        if crop_min != 0 or crop_max != 100:
+            audio_start = len(audio) * crop_min / 100
+            audio_end = len(audio) * crop_max / 100
+            audio = audio[audio_start:audio_end]
+        data = np.array(audio.get_array_of_samples())
+        if audio.channels > 1:
+            data = data.reshape(-1, audio.channels)
+        return audio.frame_rate, data
 
 
 def audio_to_file(sample_rate, data, filename, format="wav"):