diff options
Diffstat (limited to 'modules/postprocessing.py')
| -rw-r--r-- | modules/postprocessing.py | 65 |
1 files changed, 32 insertions, 33 deletions
diff --git a/modules/postprocessing.py b/modules/postprocessing.py index 136e9c88..fd0c0cc9 100644 --- a/modules/postprocessing.py +++ b/modules/postprocessing.py @@ -11,37 +11,32 @@ def run_postprocessing(extras_mode, image, image_folder, input_dir, output_dir, shared.state.begin(job="extras")
- image_data = []
- image_names = []
outputs = []
- if extras_mode == 1:
- for img in image_folder:
- if isinstance(img, Image.Image):
- image = img
- fn = ''
- else:
- image = Image.open(os.path.abspath(img.name))
- fn = os.path.splitext(img.orig_name)[0]
- image_data.append(image)
- image_names.append(fn)
- elif extras_mode == 2:
- assert not shared.cmd_opts.hide_ui_dir_config, '--hide-ui-dir-config option must be disabled'
- assert input_dir, 'input directory not selected'
-
- image_list = shared.listfiles(input_dir)
- for filename in image_list:
- try:
- image = Image.open(filename)
- except Exception:
- continue
- image_data.append(image)
- image_names.append(filename)
- else:
- assert image, 'image not selected'
-
- image_data.append(image)
- image_names.append(None)
+ def get_images(extras_mode, image, image_folder, input_dir):
+ if extras_mode == 1:
+ for img in image_folder:
+ if isinstance(img, Image.Image):
+ image = img
+ fn = ''
+ else:
+ image = Image.open(os.path.abspath(img.name))
+ fn = os.path.splitext(img.orig_name)[0]
+ yield image, fn
+ elif extras_mode == 2:
+ assert not shared.cmd_opts.hide_ui_dir_config, '--hide-ui-dir-config option must be disabled'
+ assert input_dir, 'input directory not selected'
+
+ image_list = shared.listfiles(input_dir)
+ for filename in image_list:
+ try:
+ image = Image.open(filename)
+ except Exception:
+ continue
+ yield image, filename
+ else:
+ assert image, 'image not selected'
+ yield image, None
if extras_mode == 2 and output_dir != '':
outpath = output_dir
@@ -50,14 +45,16 @@ def run_postprocessing(extras_mode, image, image_folder, input_dir, output_dir, infotext = ''
- for image, name in zip(image_data, image_names):
+ for image_data, name in get_images(extras_mode, image, image_folder, input_dir):
+ image_data: Image.Image
+
shared.state.textinfo = name
- parameters, existing_pnginfo = images.read_info_from_image(image)
+ parameters, existing_pnginfo = images.read_info_from_image(image_data)
if parameters:
existing_pnginfo["parameters"] = parameters
- pp = scripts_postprocessing.PostprocessedImage(image.convert("RGB"))
+ pp = scripts_postprocessing.PostprocessedImage(image_data.convert("RGB"))
scripts.scripts_postproc.run(pp, args)
@@ -78,8 +75,10 @@ def run_postprocessing(extras_mode, image, image_folder, input_dir, output_dir, if extras_mode != 2 or show_extras_results:
outputs.append(pp.image)
- devices.torch_gc()
+ image_data.close()
+ devices.torch_gc()
+ shared.state.end()
return outputs, ui_common.plaintext_to_html(infotext), ''
|