From d3113f5c3f6aa266cf9cae498690cece6f3784c3 Mon Sep 17 00:00:00 2001
From: Manuel Schmid <9307310+mashb1t@users.noreply.github.com>
Date: Sun, 25 Feb 2024 22:56:38 +0100
Subject: [PATCH] feat: use consistent file name in gradio (#1932)

* feat: use consistent file name in gradio

returns and uses filepaths instead of numpy image by saving to temp dir
uses double the temp dir file storage on disk as it saves to temp dir and gradio temp dir when displaying the image, but reuses logged output image

* feat: delete temp images after yielding to gradio

* feat: use args temp path if given

* chore: code cleanup, remove redundant if statement
---
 args_manager.py                    |  6 ++++++
 ldm_patched/modules/args_parser.py |  1 -
 modules/async_worker.py            |  9 +++++----
 modules/private_logger.py          | 16 +++++++++-------
 webui.py                           |  5 +++++
 5 files changed, 25 insertions(+), 12 deletions(-)

diff --git a/args_manager.py b/args_manager.py
index eeb38e1..1675c31 100644
--- a/args_manager.py
+++ b/args_manager.py
@@ -1,5 +1,7 @@
 import ldm_patched.modules.args_parser as args_parser
+import os
 
+from tempfile import gettempdir
 
 args_parser.parser.add_argument("--share", action='store_true', help="Set whether to share on Gradio.")
 args_parser.parser.add_argument("--preset", type=str, default=None, help="Apply specified UI preset.")
@@ -40,7 +42,11 @@ args_parser.args.always_offload_from_vram = not args_parser.args.disable_offload
 if args_parser.args.disable_analytics:
     import os
     os.environ["GRADIO_ANALYTICS_ENABLED"] = "False"
+
 if args_parser.args.disable_in_browser:
     args_parser.args.in_browser = False
 
+if args_parser.args.temp_path is None:
+    args_parser.args.temp_path = os.path.join(gettempdir(), 'Fooocus')
+
 args = args_parser.args
diff --git a/ldm_patched/modules/args_parser.py b/ldm_patched/modules/args_parser.py
index 272deb8..0c6165a 100644
--- a/ldm_patched/modules/args_parser.py
+++ b/ldm_patched/modules/args_parser.py
@@ -102,7 +102,6 @@ vram_group.add_argument("--always-low-vram", action="store_true")
 vram_group.add_argument("--always-no-vram", action="store_true")
 vram_group.add_argument("--always-cpu", type=int, nargs="?", metavar="CPU_NUM_THREADS", const=-1)
 
-
 parser.add_argument("--always-offload-from-vram", action="store_true")
 parser.add_argument("--pytorch-deterministic", action="store_true")
 
diff --git a/modules/async_worker.py b/modules/async_worker.py
index 4fca096..2a31aae 100644
--- a/modules/async_worker.py
+++ b/modules/async_worker.py
@@ -563,8 +563,8 @@ def worker():
 
             if direct_return:
                 d = [('Upscale (Fast)', '2x')]
-                log(uov_input_image, d)
-                yield_result(async_task, uov_input_image, do_not_show_finished_images=True)
+                uov_input_image_path = log(uov_input_image, d)
+                yield_result(async_task, uov_input_image_path, do_not_show_finished_images=True)
                 return
 
             tiled = True
@@ -828,6 +828,7 @@ def worker():
                 if inpaint_worker.current_task is not None:
                     imgs = [inpaint_worker.current_task.post_process(x) for x in imgs]
 
+                img_paths = []
                 for x in imgs:
                     d = [
                         ('Prompt', task['log_positive_prompt']),
@@ -853,9 +854,9 @@ def worker():
                         if n != 'None':
                             d.append((f'LoRA {li + 1}', f'{n} : {w}'))
                     d.append(('Version', 'v' + fooocus_version.version))
-                    log(x, d)
+                    img_paths.append(log(x, d))
 
-                yield_result(async_task, imgs, do_not_show_finished_images=len(tasks) == 1 or disable_intermediate_results)
+                yield_result(async_task, img_paths, do_not_show_finished_images=len(tasks) == 1 or disable_intermediate_results)
             except ldm_patched.modules.model_management.InterruptProcessingException as e:
                 if async_task.last_stop == 'skip':
                     print('User skipped')
diff --git a/modules/private_logger.py b/modules/private_logger.py
index 49f17dc..506b105 100644
--- a/modules/private_logger.py
+++ b/modules/private_logger.py
@@ -6,7 +6,7 @@ import urllib.parse
 
 from PIL import Image
 from modules.util import generate_temp_filename
-
+from tempfile import gettempdir
 
 log_cache = {}
 
@@ -18,13 +18,15 @@ def get_current_html_path():
     return html_name
 
 
-def log(img, dic):
-    if args_manager.args.disable_image_log:
-        return
-
-    date_string, local_temp_filename, only_name = generate_temp_filename(folder=modules.config.path_outputs, extension='png')
+def log(img, dic) -> str:
+    path_outputs = args_manager.args.temp_path if args_manager.args.disable_image_log else modules.config.path_outputs
+    date_string, local_temp_filename, only_name = generate_temp_filename(folder=path_outputs, extension='png')
     os.makedirs(os.path.dirname(local_temp_filename), exist_ok=True)
     Image.fromarray(img).save(local_temp_filename)
+
+    if args_manager.args.disable_image_log:
+        return local_temp_filename
+
     html_name = os.path.join(os.path.dirname(local_temp_filename), 'log.html')
 
     css_styles = (
@@ -105,4 +107,4 @@ def log(img, dic):
 
     log_cache[html_name] = middle_part
 
-    return
+    return local_temp_filename
diff --git a/webui.py b/webui.py
index 14ba2a1..a3663df 100644
--- a/webui.py
+++ b/webui.py
@@ -72,6 +72,11 @@ def generate_clicked(task):
                     gr.update(visible=True, value=product)
                 finished = True
 
+                # delete Fooocus temp images, only keep gradio temp images
+                if args_manager.args.disable_image_log:
+                    for filepath in product:
+                        os.remove(filepath)
+
     execution_time = time.perf_counter() - execution_start_time
     print(f'Total time: {execution_time:.2f} seconds')
     return