From 21dc2ece1bd8386a2f77617bfbd676196814eea9 Mon Sep 17 00:00:00 2001
From: cmdr2 <secondary.cmdr2@gmail.com>
Date: Sun, 4 Sep 2022 00:12:48 +0530
Subject: [PATCH] Store images during a session in the same folder; Store the
 metadata for each image as a txt file next to it

---
 ui/sd_internal/runtime.py | 31 ++++++++++++++++++++++++-------
 1 file changed, 24 insertions(+), 7 deletions(-)

diff --git a/ui/sd_internal/runtime.py b/ui/sd_internal/runtime.py
index 6bfef32e..b30910bb 100644
--- a/ui/sd_internal/runtime.py
+++ b/ui/sd_internal/runtime.py
@@ -15,6 +15,8 @@ from ldm.util import instantiate_from_config
 from optimizedSD.optimUtils import split_weighted_subprompts
 from transformers import logging
 
+import uuid
+
 logging.set_verbosity_error()
 
 # consts
@@ -26,6 +28,8 @@ import base64
 from io import BytesIO
 
 # local
+session_id = str(uuid.uuid4())
+
 ckpt = None
 model = None
 modelCS = None
@@ -140,7 +144,7 @@ def mk_img(req: Request):
     opt_init_img = req.init_image
     opt_format = 'png'
 
-    print(req.to_string(), 'device', device)
+    print(req.to_string(), '\n    device', device)
 
     seed_everything(opt_seed)
 
@@ -183,15 +187,19 @@ def mk_img(req: Request):
         t_enc = int(opt_strength * opt_ddim_steps)
         print(f"target t_enc is {t_enc} steps")
 
+    if opt_save_to_disk_path is not None:
+        session_out_path = os.path.join(opt_save_to_disk_path, 'session-' + session_id)
+        os.makedirs(session_out_path, exist_ok=True)
+    else:
+        session_out_path = None
+
     seeds = ""
     with torch.no_grad():
         for n in trange(opt_n_iter, desc="Sampling"):
             for prompts in tqdm(data, desc="data"):
 
                 if opt_save_to_disk_path is not None:
-                    sample_path = os.path.join(opt_save_to_disk_path, "_".join(re.split(":| ", prompts[0])))[:150]
-                    os.makedirs(sample_path, exist_ok=True)
-                    base_count = len(os.listdir(sample_path))
+                    base_count = len(os.listdir(session_out_path))
 
                 with precision_scope("cuda"):
                     modelCS.to(device)
@@ -234,9 +242,18 @@ def mk_img(req: Request):
                         res.images.append(ResponseImage(data=img_data, seed=opt_seed))
 
                         if opt_save_to_disk_path is not None:
-                            img.save(
-                                os.path.join(sample_path, "seed_" + str(opt_seed) + "_" + f"{base_count:05}.{opt_format}")
-                            )
+                            prompt_flattened = "_".join(re.split(":| ", prompts[0]))
+                            prompt_flattened = prompt_flattened[:150]
+
+                            file_path = f"sd_{prompt_flattened}_Seed-{opt_seed}_Steps-{opt_ddim_steps}_Guidance-{opt_scale}_{base_count:05}"
+                            img_out_path = os.path.join(session_out_path, f"{file_path}.{opt_format}")
+                            meta_out_path = os.path.join(session_out_path, f"{file_path}.txt")
+
+                            metadata = f"{prompts[0]}\nSeed: {opt_seed}\nSteps: {opt_ddim_steps}\nGuidance Scale: {opt_scale}"
+                            img.save(img_out_path)
+                            with open(meta_out_path, 'w') as f:
+                                f.write(metadata)
+
                             base_count += 1
 
                         seeds += str(opt_seed) + ","