Update sd_save_intermediate_images.py

Saving intermediate and final image to Scripts folder temporary and adding button to open image viewer HTML file
pull/41/head
heloess 2024-06-19 21:21:52 +03:00 committed by GitHub
parent 24c3ce206a
commit 99ca140365
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 50 additions and 28 deletions

View File

@ -7,7 +7,8 @@ import platform
import re
import shutil
import sys
import base64
import webbrowser
from modules import paths
from modules import scripts
from modules import script_callbacks
@ -501,9 +502,12 @@ class Script(scripts.Script):
inputs=[ssii_is_active, ssii_final_save, ssii_intermediate_type, ssii_every_n, ssii_start_at_n, ssii_stop_at_n, ssii_mode, ssii_video_format, ssii_mp4_parms, ssii_video_fps, ssii_add_first_frames, ssii_add_last_frames, ssii_smooth, ssii_seconds, ssii_lores, ssii_hires, ssii_ffmpeg_bat, ssii_bat_only, ssii_debug],
outputs=[ssii_message],
)
open_image_viewer_button = gr.Button("Open Image Viewer")
open_image_viewer_button.click(fn=self.open_image_viewer)
return [ssii_is_active, ssii_final_save, ssii_intermediate_type, ssii_every_n, ssii_start_at_n, ssii_stop_at_n, ssii_mode, ssii_video_format, ssii_mp4_parms, ssii_video_fps, ssii_add_first_frames, ssii_add_last_frames, ssii_smooth, ssii_seconds, ssii_lores, ssii_hires, ssii_ffmpeg_bat, ssii_bat_only, ssii_debug]
def open_image_viewer(self):
filepath = os.path.join(os.path.dirname(__file__), 'Image Viewer.html')
webbrowser.open_new_tab('file://' + os.path.abspath(filepath))
def save_image_only_get_name(image, path, basename, seed=None, prompt=None, extension='png', info=None, short_filename=False, no_prompt=False, grid=False, pnginfo_section_name='parameters', p=None, existing_info=None, forced_filename=None, suffix="", save_to_dirs=None):
# for description see modules.images.save_image, same code up saving of files
logger.debug(f"func: {sys._getframe(0).f_code.co_name}")
@ -773,7 +777,13 @@ class Script(scripts.Script):
logger.debug(f"filename: {filename_clean(filename)}")
p.intermed_files.append((index, filename + ".png", None))
p.intermed_last[index] = (filename + ".png", p.intermed_outpath, False)
current_dir = os.path.dirname(os.path.abspath(__file__))
file_path_temp = os.path.join(current_dir,"saved.png")
image.save(file_path_temp)
logger.debug(f"index, p.intermed_last[index]: {index}, {filename_clean(p.intermed_last[index][0])}, {filename_clean(p.intermed_last[index][1])}, {p.intermed_last[index][2]}")
return orig_callback_state(self, d)
setattr(KDiffusionSampler, "callback_state", callback_state)
@ -787,34 +797,46 @@ class Script(scripts.Script):
# Make video for last batch_count
make_video(p, ssii_is_active, ssii_final_save, ssii_intermediate_type, ssii_every_n, ssii_start_at_n, ssii_stop_at_n, ssii_mode, ssii_video_format, ssii_mp4_parms, ssii_video_fps, ssii_add_first_frames, ssii_add_last_frames, ssii_smooth, ssii_seconds, ssii_lores, ssii_hires, ssii_ffmpeg_bat, ssii_bat_only, ssii_debug)
def handle_image_saved(params : script_callbacks.ImageSaveParams):
logger.debug(f"func: {sys._getframe(0).f_code.co_name}")
if hasattr(params.p, "intermed_is_active"):
# Copy final image to intermediates folder
if params.p.intermed_is_active and params.p.intermed_final_save:
directories = os.path.normpath(params.filename).split(os.sep)
if "intermediates" not in directories:
# Get last file name of current index
last_found = False
for index, (last_filename, last_path, last_done) in enumerate(params.p.intermed_last.values()):
if not last_done:
last_found = True
params.p.intermed_last[index] = (last_filename, last_path, True)
break
if last_found:
# Convert final file name to intermediates filename
match1 = re.search(r'^(.*?)-(\d+)(.*)$', last_filename)
def handle_image_saved(params):
logger.debug(f"Function: {sys._getframe(0).f_code.co_name}")
if hasattr(params.p, "intermed_is_active") and params.p.intermed_is_active and params.p.intermed_final_save:
filename = params.filename
directories = os.path.normpath(filename).split(os.sep)
if "intermediates" not in directories:
# Find the last unprocessed file
last_found = False
for index, (last_filename, last_path, last_done) in enumerate(params.p.intermed_last.values()):
if not last_done:
last_found = True
params.p.intermed_last[index] = (last_filename, last_path, True)
break
if last_found:
# Construct new filename for intermediate copy
match1 = re.search(r'^(.*?)-(\d+)(.*)$', last_filename)
if match1:
new_number = str(int(match1.group(2)) + 1).zfill(3)
file_new = match1.group(1) + '-' + new_number + match1.group(3)
#file_new_path = os.path.join(os.path.dirname(params.p.intermed_lastfile), file_new)
file_new_path = os.path.join(last_path, file_new)
logger.debug(f"last_filename: {filename_clean(last_filename)}")
logger.debug(f"file_new_path: {filename_clean(file_new_path)}")
logger.debug(f"Original Filename: {filename}")
logger.debug(f"New Intermediate Filename: {file_new_path}")
shutil.copy(params.filename, file_new_path)
try:
shutil.copy(filename, file_new_path)
# Update information for video processing or further steps
params.p.intermed_files.append((index, file_new, None))
# Save final image to another location (assuming filename is the final image path)
current_dir = current_dir = os.path.dirname(os.path.abspath(__file__))
file_path_temp = os.path.join(current_dir, "saved.png")
shutil.copy(filename, file_path_temp)
logger.debug(f"Final image saved to: {file_path_temp}")
except Exception as e:
logger.error(f"Error copying file: {str(e)}")
else:
logger.error(f"Failed to parse filename: {filename}")
# Add info for make video
params.p.intermed_files.append((index, file_new, None))
script_callbacks.on_image_saved(handle_image_saved)