From f4228a850c439f981bbcaf0c7b19fda4d666c07c Mon Sep 17 00:00:00 2001
From: Vito Sansevero <vito.sansevero@gmail.com>
Date: Mon, 21 Jul 2025 07:56:13 -0700
Subject: [PATCH] refactor(logic): improve path handling in image saving

---
 kikotools/tools/kiko_save_image/logic.py | 45 ++++++++++++++++--------
 1 file changed, 30 insertions(+), 15 deletions(-)

diff --git a/kikotools/tools/kiko_save_image/logic.py b/kikotools/tools/kiko_save_image/logic.py
index ddfc743..68faab5 100644
--- a/kikotools/tools/kiko_save_image/logic.py
+++ b/kikotools/tools/kiko_save_image/logic.py
@@ -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