Merge pull request #1214 from WDevelopsWebApps/saving

Save Button file and directory naming

modules/images.py

@@ -311,7 +311,12 @@ def apply_filename_pattern(x, p, seed, prompt):
         x = x.replace("[cfg]", str(p.cfg_scale))
         x = x.replace("[width]", str(p.width))
         x = x.replace("[height]", str(p.height))
-        x = x.replace("[styles]", sanitize_filename_part(", ".join([x for x in p.styles if not x == "None"]), replace_spaces=False))
+        
+        #currently disabled if using the save button, will work otherwise 
+        # if enabled it will cause a bug because styles is not included in the save_files data dictionary
+        if hasattr(p, "styles"):
+            x = x.replace("[styles]", sanitize_filename_part(", ".join([x for x in p.styles if not x == "None"]), replace_spaces=False))
+
         x = x.replace("[sampler]", sanitize_filename_part(sd_samplers.samplers[p.sampler_index].name, replace_spaces=False))
 
     x = x.replace("[model_hash]", shared.sd_model.sd_model_hash)

modules/ui.py

@@ -33,6 +33,7 @@ import modules.gfpgan_model
 import modules.codeformer_model
 import modules.styles
 import modules.generation_parameters_copypaste
+from modules.images import apply_filename_pattern, get_next_sequence_number
 import modules.textual_inversion.ui
 
 # this is a fix for Windows users. Without it, javascript files will be served with text/html content-type and the bowser will not show any UI
@@ -96,14 +97,31 @@ def send_gradio_gallery_to_image(x):
 
 
 def save_files(js_data, images, index):
-    import csv
+    import csv    
-
-    os.makedirs(opts.outdir_save, exist_ok=True)
-
     filenames = []
 
+    #quick dictionary to class object conversion. Its neccesary due apply_filename_pattern requiring it
+    class MyObject:
+        def __init__(self, d=None):
+            if d is not None:
+                for key, value in d.items():
+                    setattr(self, key, value)
+
     data = json.loads(js_data)
+
+    p = MyObject(data)
+    path = opts.outdir_save
+    save_to_dirs = opts.save_to_dirs
+
+    if save_to_dirs:
+        dirname = apply_filename_pattern(opts.directories_filename_pattern or "[prompt_words]", p, p.seed, p.prompt)
+        path = os.path.join(opts.outdir_save, dirname)
+
+    os.makedirs(path, exist_ok=True)
+  
+
     if index > -1 and opts.save_selected_only and (index >= data["index_of_first_image"]):  # ensures we are looking at a specific non-grid picture, and we have save_selected_only
+
         images = [images[index]]
         infotexts = [data["infotexts"][index]]
     else:
@@ -115,11 +133,20 @@ def save_files(js_data, images, index):
         if at_start:
             writer.writerow(["prompt", "seed", "width", "height", "sampler", "cfgs", "steps", "filename", "negative_prompt"])
 
-        filename_base = str(int(time.time() * 1000))
+        file_decoration = opts.samples_filename_pattern or "[seed]-[prompt_spaces]"
+        if file_decoration != "":
+            file_decoration = "-" + file_decoration.lower()
+        file_decoration = apply_filename_pattern(file_decoration, p, p.seed, p.prompt)
+        truncated = (file_decoration[:240] + '..') if len(file_decoration) > 240 else file_decoration
+        filename_base = truncated
         extension = opts.samples_format.lower()
+
+        basecount = get_next_sequence_number(path, "")
         for i, filedata in enumerate(images):
-            filename = filename_base + ("" if len(images) == 1 else "-" + str(i + 1)) + f".{extension}"
+            file_number = f"{basecount+i:05}"
-            filepath = os.path.join(opts.outdir_save, filename)
+            filename = file_number + filename_base + f".{extension}"
+            filepath = os.path.join(path, filename)
+
 
             if filedata.startswith("data:image/png;base64,"):
                 filedata = filedata[len("data:image/png;base64,"):]