[SD][web] Add support for saving generated images

Signed-Off-by: Gaurav Shukla <[email protected]>

web/models/stable_diffusion/main.py

@@ -9,6 +9,10 @@
 import numpy as np
 import time
 import sys
+from datetime import datetime as dt
+from csv import DictWriter
+import re
+from pathlib import Path
 
 
 if args.clear_all:
@@ -65,6 +69,39 @@ def set_ui_params(
     args.variant = variant
 
 
+# save output images and the inputs correspoding to it.
+def save_output_img(output_img):
+    output_path = args.output_dir if args.output_dir else Path.cwd()
+    generated_imgs_path = Path(output_path, "generated_imgs")
+    generated_imgs_path.mkdir(parents=True, exist_ok=True)
+    csv_path = Path(output_path, "imgs_details.csv")
+
+    prompt_slice = re.sub("[^a-zA-Z0-9]", "_", args.prompts[0][:15])
+    out_img_name = (
+        f"{prompt_slice}_{args.seed}_{dt.now().strftime('%y%m%d_%H%M%S')}"
+    )
+    out_img_path = Path(generated_imgs_path, f"{out_img_name}.jpg")
+    output_img.save(out_img_path)
+
+    new_entry = {
+        "VARIANT": args.variant,
+        "VERSION": args.version,
+        "SCHEDULER": args.scheduler,
+        "PROMPT": args.prompts[0],
+        "NEG_PROMPT": args.negative_prompts[0],
+        "SEED": args.seed,
+        "CFG_SCALE": args.guidance_scale,
+        "PRECISION": args.precision,
+        "STEPS": args.steps,
+        "OUTPUT": out_img_path,
+    }
+
+    with open(csv_path, "a") as csv_obj:
+        dictwriter_obj = DictWriter(csv_obj, fieldnames=list(new_entry.keys()))
+        dictwriter_obj.writerow(new_entry)
+        csv_obj.close()
+
+
 def stable_diff_inf(
     prompt: str,
     negative_prompt: str,
@@ -226,4 +263,6 @@ def stable_diff_inf(
     text_output += f"\nAverage step time: {avg_ms:.4f}ms/it"
     text_output += f"\nTotal image generation time: {total_time:.4f}sec"
 
+    save_output_img(pil_images[0])
+
     return pil_images[0], text_output

web/models/stable_diffusion/stable_args.py

@@ -117,6 +117,13 @@
     help="other supported schedulers are [PNDM, DDIM, LMSDiscrete, EulerDiscrete, DPMSolverMultistep]",
 )
 
+p.add_argument(
+    "--output_dir",
+    type=str,
+    default=None,
+    help="Directory path to save the output images and json",
+)
+
 ##############################################################################
 ### IREE - Vulkan supported flags
 ##############################################################################