refactor(logic): improve path handling in image saving

This commit is contained in:
Vito Sansevero
2025-07-21 07:56:13 -07:00
parent 79042b78d2
commit f4228a850c
+30 -15
View File
@@ -42,20 +42,31 @@ def get_save_image_path(
Returns:
Tuple of (full_path, relative_filename)
"""
# Sanitize filename prefix to remove invalid characters
# Replace path separators and other problematic characters
safe_prefix = filename_prefix.replace("/", "_").replace("\\", "_").replace(":", "_")
# Split filename_prefix into directory path and actual filename prefix
# This allows for directory structures like "kittybear/anime/images/kittybear"
prefix_dir = os.path.dirname(filename_prefix)
prefix_name = os.path.basename(filename_prefix)
# Sanitize only the filename part (not the directory path)
safe_prefix = prefix_name.replace(":", "_") # Only sanitize problematic chars for filenames
safe_prefix = "".join(c for c in safe_prefix if c.isalnum() or c in "._-")
# Create unique filename with timestamp to avoid conflicts
timestamp = int(time.time())
filename = f"{safe_prefix}_{timestamp:010d}_{batch_number:05d}{format_ext}"
# Handle subfolder
# Handle subfolder and prefix directory (but not the filename part)
path_components = []
path_components.append(output_dir)
if subfolder:
full_output_folder = os.path.join(output_dir, subfolder)
else:
full_output_folder = output_dir
path_components.append(subfolder)
# Only add prefix_dir if it exists (the directory part, not the filename part)
if prefix_dir:
path_components.append(prefix_dir)
full_output_folder = os.path.join(*path_components)
# Ensure directory exists
os.makedirs(full_output_folder, exist_ok=True)
@@ -64,17 +75,21 @@ def get_save_image_path(
# For the preview, ComfyUI needs the filename and subfolder separately
# The subfolder needs to be relative to the output directory root
# `os.path.relpath` is not suitable here as it might resolve symbolic links
# and produce `../` in paths, which is not what we want.
# A simple string replacement or careful path joining is more robust.
# Build the relative subfolder path including prefix directory (but not filename part)
relative_path_components = []
if subfolder:
# Ensure subfolder is a relative path structure
relative_subfolder = subfolder.strip("/\\")
# The filename for the preview should not contain the subfolder path
preview_filename = filename
relative_path_components.append(subfolder.strip("/\\"))
if prefix_dir:
relative_path_components.append(prefix_dir.strip("/\\"))
if relative_path_components:
relative_subfolder = os.path.join(*relative_path_components)
else:
relative_subfolder = ""
preview_filename = filename
preview_filename = filename
return full_path, preview_filename, relative_subfolder