diff options
| author | AUTOMATIC1111 <16777216c@gmail.com> | 2022-11-27 09:54:45 +0000 |
|---|---|---|
| committer | GitHub <noreply@github.com> | 2022-11-27 09:54:45 +0000 |
| commit | 6df49457183fedec635d30dd0c611844b0932c85 (patch) | |
| tree | 7cd692e17d8ff273916c52d30c77370f1871de54 /modules/ui_tempdir.py | |
| parent | 45fd785436068f3b1c09fb7bc575118b6059fc7b (diff) | |
| parent | b48b7999c86fd6d7f006f76adf5a484175782c37 (diff) | |
| download | stable-diffusion-webui-gfx803-6df49457183fedec635d30dd0c611844b0932c85.tar.gz stable-diffusion-webui-gfx803-6df49457183fedec635d30dd0c611844b0932c85.tar.bz2 stable-diffusion-webui-gfx803-6df49457183fedec635d30dd0c611844b0932c85.zip | |
Merge branch 'master' into DPM++SDE
Diffstat (limited to 'modules/ui_tempdir.py')
| -rw-r--r-- | modules/ui_tempdir.py | 62 |
1 files changed, 62 insertions, 0 deletions
diff --git a/modules/ui_tempdir.py b/modules/ui_tempdir.py new file mode 100644 index 00000000..9c6d3a9d --- /dev/null +++ b/modules/ui_tempdir.py @@ -0,0 +1,62 @@ +import os
+import tempfile
+from collections import namedtuple
+
+import gradio as gr
+
+from PIL import PngImagePlugin
+
+from modules import shared
+
+
+Savedfile = namedtuple("Savedfile", ["name"])
+
+
+def save_pil_to_file(pil_image, dir=None):
+ already_saved_as = getattr(pil_image, 'already_saved_as', None)
+ if already_saved_as:
+ shared.demo.temp_dirs = shared.demo.temp_dirs | {os.path.abspath(os.path.dirname(already_saved_as))}
+ file_obj = Savedfile(already_saved_as)
+ return file_obj
+
+ if shared.opts.temp_dir != "":
+ dir = shared.opts.temp_dir
+
+ use_metadata = False
+ metadata = PngImagePlugin.PngInfo()
+ for key, value in pil_image.info.items():
+ if isinstance(key, str) and isinstance(value, str):
+ metadata.add_text(key, value)
+ use_metadata = True
+
+ file_obj = tempfile.NamedTemporaryFile(delete=False, suffix=".png", dir=dir)
+ pil_image.save(file_obj, pnginfo=(metadata if use_metadata else None))
+ return file_obj
+
+
+# override save to file function so that it also writes PNG info
+gr.processing_utils.save_pil_to_file = save_pil_to_file
+
+
+def on_tmpdir_changed():
+ if shared.opts.temp_dir == "" or shared.demo is None:
+ return
+
+ os.makedirs(shared.opts.temp_dir, exist_ok=True)
+
+ shared.demo.temp_dirs = shared.demo.temp_dirs | {os.path.abspath(shared.opts.temp_dir)}
+
+
+def cleanup_tmpdr():
+ temp_dir = shared.opts.temp_dir
+ if temp_dir == "" or not os.path.isdir(temp_dir):
+ return
+
+ for root, dirs, files in os.walk(temp_dir, topdown=False):
+ for name in files:
+ _, extension = os.path.splitext(name)
+ if extension != ".png":
+ continue
+
+ filename = os.path.join(root, name)
+ os.remove(filename)
|