add png as default format

Author: PhaneeshB

shark/examples/shark_inference/stable_diffusion/main.py

@@ -283,8 +283,16 @@ def end_profiling(device):
             }
             prompt_slice = re.sub("[^a-zA-Z0-9]", "_", args.prompts[i][:15])
             img_name = f"{prompt_slice}_{args.seed}_{run}_{dt.now().strftime('%y%m%d_%H%M%S')}"
-            pil_images[i].save(
-                output_path / f"{img_name}.jpg", quality=95, subsampling=0
-            )
+            if args.output_img_format == "jpg":
+                pil_images[i].save(
+                    output_path / f"{img_name}.jpg", quality=95, subsampling=0
+                )
+            else:
+                pil_images[i].save(output_path / f"{img_name}.png", "PNG")
+                if args.output_img_format not in ["png", "jpg"]:
+                    print(
+                        f"[ERROR] Format {args.output_img_format} is not supported yet."
+                        "saving image as png. Supported formats png / jpg"
+                    )
             with open(output_path / f"{img_name}.json", "w") as f:
                 f.write(json.dumps(json_store, indent=4))

shark/examples/shark_inference/stable_diffusion/stable_args.py

@@ -123,6 +123,13 @@ def path_expand(s):
     help="other supported schedulers are [PNDM, DDIM, LMSDiscrete, EulerDiscrete, DPMSolverMultistep]",
 )
 
+p.add_argument(
+    "--output_img_format",
+    type=str,
+    default="png",
+    help="specify the format in which output image is save. Supported options: jpg / png",
+)
+
 p.add_argument(
     "--output_dir",
     type=str,

web/models/stable_diffusion/main.py

@@ -82,8 +82,17 @@ def save_output_img(output_img):
     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, quality=95, subsampling=0)
+    if args.output_img_format == "jpg":
+        out_img_path = Path(generated_imgs_path, f"{out_img_name}.jpg")
+        output_img.save(out_img_path, quality=95, subsampling=0)
+    else:
+        out_img_path = Path(generated_imgs_path, f"{out_img_name}.png")
+        output_img.save(out_img_path, "PNG")
+        if args.output_img_format not in ["png", "jpg"]:
+            print(
+                f"[ERROR] Format {args.output_img_format} is not supported yet."
+                "saving image as png. Supported formats png / jpg"
+            )
 
     new_entry = {
         "VARIANT": args.variant,

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_img_format",
+    type=str,
+    default="png",
+    help="specify the format in which output image is save. Supported options: jpg / png",
+)
+
 p.add_argument(
     "--output_dir",
     type=str,