Gradio 4 + WebUI 1.10

modules/ui_tempdir.py

@@ -4,6 +4,7 @@ from collections import namedtuple
 from pathlib import Path
 
 import gradio.components
+import gradio as gr
 
 from PIL import PngImagePlugin
 
@@ -13,25 +14,35 @@ from modules import shared
 Savedfile = namedtuple("Savedfile", ["name"])
 
 
-def register_tmp_file(gradio, filename):
-    if hasattr(gradio, 'temp_file_sets'):  # gradio 3.15
-        gradio.temp_file_sets[0] = gradio.temp_file_sets[0] | {os.path.abspath(filename)}
+def register_tmp_file(gradio_app, filename):
+    if hasattr(gradio_app, 'temp_file_sets'):  # gradio 3.15
+        if hasattr(gr.utils, 'abspath'):  # gradio 4.19
+            filename = gr.utils.abspath(filename)
+        else:
+            filename = os.path.abspath(filename)
 
-    if hasattr(gradio, 'temp_dirs'):  # gradio 3.9
-        gradio.temp_dirs = gradio.temp_dirs | {os.path.abspath(os.path.dirname(filename))}
+        gradio_app.temp_file_sets[0] = gradio_app.temp_file_sets[0] | {filename}
+
+    if hasattr(gradio_app, 'temp_dirs'):  # gradio 3.9
+        gradio_app.temp_dirs = gradio_app.temp_dirs | {os.path.abspath(os.path.dirname(filename))}
 
 
-def check_tmp_file(gradio, filename):
-    if hasattr(gradio, 'temp_file_sets'):
-        return any(filename in fileset for fileset in gradio.temp_file_sets)
+def check_tmp_file(gradio_app, filename):
+    if hasattr(gradio_app, 'temp_file_sets'):
+        if hasattr(gr.utils, 'abspath'):  # gradio 4.19
+            filename = gr.utils.abspath(filename)
+        else:
+            filename = os.path.abspath(filename)
 
-    if hasattr(gradio, 'temp_dirs'):
-        return any(Path(temp_dir).resolve() in Path(filename).resolve().parents for temp_dir in gradio.temp_dirs)
+        return any(filename in fileset for fileset in gradio_app.temp_file_sets)
+
+    if hasattr(gradio_app, 'temp_dirs'):
+        return any(Path(temp_dir).resolve() in Path(filename).resolve().parents for temp_dir in gradio_app.temp_dirs)
 
     return False
 
 
-def save_pil_to_file(self, pil_image, dir=None, format="png"):
+def save_pil_to_file(pil_image, cache_dir=None, format="png"):
     already_saved_as = getattr(pil_image, 'already_saved_as', None)
     if already_saved_as and os.path.isfile(already_saved_as):
         register_tmp_file(shared.demo, already_saved_as)
@@ -39,9 +50,10 @@ def save_pil_to_file(self, pil_image, dir=None, format="png"):
         register_tmp_file(shared.demo, filename_with_mtime)
         return filename_with_mtime
 
-    if shared.opts.temp_dir != "":
+    if shared.opts.temp_dir:
         dir = shared.opts.temp_dir
     else:
+        dir = cache_dir
         os.makedirs(dir, exist_ok=True)
 
     use_metadata = False
@@ -56,9 +68,96 @@ def save_pil_to_file(self, pil_image, dir=None, format="png"):
     return file_obj.name
 
 
+async def async_move_files_to_cache(data, block, postprocess=False, check_in_upload_folder=False, keep_in_cache=False):
+    """Move any files in `data` to cache and (optionally), adds URL prefixes (/file=...) needed to access the cached file.
+    Also handles the case where the file is on an external Gradio app (/proxy=...).
+
+    Runs after .postprocess() and before .preprocess().
+
+    Copied from gradio's processing_utils.py
+
+    Args:
+        data: The input or output data for a component. Can be a dictionary or a dataclass
+        block: The component whose data is being processed
+        postprocess: Whether its running from postprocessing
+        check_in_upload_folder: If True, instead of moving the file to cache, checks if the file is in already in cache (exception if not).
+        keep_in_cache: If True, the file will not be deleted from cache when the server is shut down.
+    """
+
+    from gradio import FileData
+    from gradio.data_classes import GradioRootModel
+    from gradio.data_classes import GradioModel
+    from gradio_client import utils as client_utils
+    from gradio.utils import get_upload_folder, is_in_or_equal, is_static_file
+
+    async def _move_to_cache(d: dict):
+        payload = FileData(**d)
+
+        # EDITED
+        payload.path = payload.path.rsplit('?', 1)[0]
+
+        # If the gradio app developer is returning a URL from
+        # postprocess, it means the component can display a URL
+        # without it being served from the gradio server
+        # This makes it so that the URL is not downloaded and speeds up event processing
+        if payload.url and postprocess and client_utils.is_http_url_like(payload.url):
+            payload.path = payload.url
+        elif is_static_file(payload):
+            pass
+        elif not block.proxy_url:
+            # EDITED
+            if check_tmp_file(shared.demo, payload.path):
+                temp_file_path = payload.path
+            else:
+                # If the file is on a remote server, do not move it to cache.
+                if check_in_upload_folder and not client_utils.is_http_url_like(
+                    payload.path
+                ):
+                    path = os.path.abspath(payload.path)
+                    if not is_in_or_equal(path, get_upload_folder()):
+                        raise ValueError(
+                            f"File {path} is not in the upload folder and cannot be accessed."
+                        )
+                if not payload.is_stream:
+                    temp_file_path = await block.async_move_resource_to_block_cache(
+                        payload.path
+                    )
+                    if temp_file_path is None:
+                        raise ValueError("Did not determine a file path for the resource.")
+                    payload.path = temp_file_path
+                    if keep_in_cache:
+                        block.keep_in_cache.add(payload.path)
+
+        url_prefix = "/stream/" if payload.is_stream else "/file="
+        if block.proxy_url:
+            proxy_url = block.proxy_url.rstrip("/")
+            url = f"/proxy={proxy_url}{url_prefix}{payload.path}"
+        elif client_utils.is_http_url_like(payload.path) or payload.path.startswith(
+            f"{url_prefix}"
+        ):
+            url = payload.path
+        else:
+            url = f"{url_prefix}{payload.path}"
+        payload.url = url
+
+        return payload.model_dump()
+
+    if isinstance(data, (GradioRootModel, GradioModel)):
+        data = data.model_dump()
+
+    return await client_utils.async_traverse(
+        data, _move_to_cache, client_utils.is_file_obj
+    )
+
+
 def install_ui_tempdir_override():
-    """override save to file function so that it also writes PNG info"""
-    gradio.components.IOComponent.pil_to_temp_file = save_pil_to_file
+    """
+    override save to file function so that it also writes PNG info.
+    override gradio4's move_files_to_cache function to prevent it from writing a copy into a temporary directory.
+    """
+
+    gradio.processing_utils.save_pil_to_cache = save_pil_to_file
+    gradio.processing_utils.async_move_files_to_cache = async_move_files_to_cache
 
 
 def on_tmpdir_changed():