Spaces:
Runtime error
Runtime error
File size: 4,403 Bytes
c19ca42 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 |
import os
import tempfile
from collections import namedtuple
from pathlib import Path
from PIL import Image, PngImagePlugin
from modules import shared, errors, paths
Savedfile = namedtuple("Savedfile", ["name"])
debug = errors.log.trace if os.environ.get('SD_PATH_DEBUG', None) is not None else lambda *args, **kwargs: None
def register_tmp_file(gradio, filename):
if hasattr(gradio, 'temp_file_sets'):
gradio.temp_file_sets[0] = gradio.temp_file_sets[0] | {os.path.abspath(filename)}
def check_tmp_file(gradio, filename):
ok = False
if hasattr(gradio, 'temp_file_sets'):
ok = ok or any(filename in fileset for fileset in gradio.temp_file_sets)
if shared.opts.outdir_samples != '':
ok = ok or Path(shared.opts.outdir_samples).resolve() in Path(filename).resolve().parents
else:
ok = ok or Path(shared.opts.outdir_txt2img_samples).resolve() in Path(filename).resolve().parents
ok = ok or Path(shared.opts.outdir_img2img_samples).resolve() in Path(filename).resolve().parents
ok = ok or Path(shared.opts.outdir_extras_samples).resolve() in Path(filename).resolve().parents
if shared.opts.outdir_grids != '':
ok = ok or Path(shared.opts.outdir_grids).resolve() in Path(filename).resolve().parents
else:
ok = ok or Path(shared.opts.outdir_txt2img_grids).resolve() in Path(filename).resolve().parents
ok = ok or Path(shared.opts.outdir_img2img_grids).resolve() in Path(filename).resolve().parents
ok = ok or Path(shared.opts.outdir_save).resolve() in Path(filename).resolve().parents
ok = ok or Path(shared.opts.outdir_init_images).resolve() in Path(filename).resolve().parents
return ok
def pil_to_temp_file(self, img: Image, dir: str, format="png") -> str: # pylint: disable=redefined-builtin,unused-argument
"""
# original gradio implementation
bytes_data = gr.processing_utils.encode_pil_to_bytes(img, format)
temp_dir = Path(dir) / self.hash_bytes(bytes_data)
temp_dir.mkdir(exist_ok=True, parents=True)
filename = str(temp_dir / f"image.{format}")
img.save(filename, pnginfo=gr.processing_utils.get_pil_metadata(img))
"""
folder = dir
already_saved_as = getattr(img, 'already_saved_as', None)
exists = os.path.isfile(already_saved_as) if already_saved_as is not None else False
debug(f'Image lookup: {already_saved_as} exists={exists}')
if already_saved_as and exists:
register_tmp_file(shared.demo, already_saved_as)
file_obj = Savedfile(already_saved_as)
name = file_obj.name
debug(f'Image registered: {name}')
return name
if shared.opts.temp_dir != "":
folder = shared.opts.temp_dir
use_metadata = False
metadata = PngImagePlugin.PngInfo()
for key, value in img.info.items():
if isinstance(key, str) and isinstance(value, str):
metadata.add_text(key, value)
use_metadata = True
if not os.path.exists(folder):
os.makedirs(folder, exist_ok=True)
shared.log.debug(f'Created temp folder: path="{folder}"')
with tempfile.NamedTemporaryFile(delete=False, suffix=".png", dir=folder) as tmp:
name = tmp.name
img.save(name, pnginfo=(metadata if use_metadata else None))
img.already_saved_as = name
size = os.path.getsize(name)
shared.log.debug(f'Saving temp: image="{name}" resolution={img.width}x{img.height} size={size}')
params = ', '.join([f'{k}: {v}' for k, v in img.info.items()])
params = params[12:] if params.startswith('parameters: ') else params
with open(os.path.join(paths.data_path, "params.txt"), "w", encoding="utf8") as file:
file.write(params)
return name
# override save to file function so that it also writes PNG info
def on_tmpdir_changed():
if shared.opts.temp_dir == "":
return
register_tmp_file(shared.demo, os.path.join(shared.opts.temp_dir, "x"))
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" and extension != ".jpg" and extension != ".webp":
continue
filename = os.path.join(root, name)
os.remove(filename)
|