Added flags for metadata information. (huggingface#894)

EvanGuanSF · web-flow · commit 8cafe56eb4ca · 2023-02-01T05:16:11.000-08:00
diff --git a/apps/stable_diffusion/scripts/txt2img.py b/apps/stable_diffusion/scripts/txt2img.py
@@ -2,10 +2,12 @@
 
 os.environ["AMD_ENABLE_LLPC"] = "1"
 
+import json
 import torch
 import re
 import time
 from pathlib import Path
+from PIL import PngImagePlugin
 from datetime import datetime as dt
 from dataclasses import dataclass
 from csv import DictWriter
@@ -61,7 +63,29 @@ def save_output_img(output_img):
         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")
+        pngInfo = PngImagePlugin.PngInfo()
+
+        if args.write_metadata_to_png:
+            pngInfo.add_text(
+                "parameters",
+                f"{args.prompts[0]}\nNegative prompt: {args.negative_prompts[0]}\nSteps:{args.steps}, Sampler: {args.scheduler}, CFG scale: {args.guidance_scale}, Seed: {args.seed}, Size: {args.width}x{args.height}, Model: {args.hf_model_id}",
+            )
+
+        output_img.save(
+            output_path / f"{out_img_name}.png", "PNG", pnginfo=pngInfo
+        )
+
+        if args.output_img_format not in ["png", "jpg"]:
+            print(
+                f"[ERROR] Format {args.output_img_format} is not supported yet."
+                "Image saved as png instead. Supported formats: png / jpg"
+            )
 
     new_entry = {
         "VARIANT": args.hf_model_id,
@@ -83,6 +107,11 @@ def save_output_img(output_img):
         dictwriter_obj.writerow(new_entry)
         csv_obj.close()
 
+    if args.save_metadata_to_json:
+        del new_entry["OUTPUT"]
+        with open(f"{output_path}/{out_img_name}.json", "w") as f:
+            json.dump(new_entry, f, indent=4)
+
 
 txt2img_obj = None
 config_obj = None
@@ -106,6 +135,8 @@ def txt2img_inf(
     precision: str,
     device: str,
     max_length: int,
+    save_metadata_to_json: bool,
+    save_metadata_to_png: bool,
 ):
     global txt2img_obj
     global config_obj
@@ -119,6 +150,8 @@ def txt2img_inf(
     args.scheduler = scheduler
     args.hf_model_id = custom_model_id if custom_model_id else model_id
     args.ckpt_loc = ckpt_file_obj.name if ckpt_file_obj else ""
+    args.save_metadata_to_json = save_metadata_to_json
+    args.write_metadata_to_png = save_metadata_to_png
     dtype = torch.float32 if precision == "fp32" else torch.half
     cpu_scheduling = not scheduler.startswith("Shark")
     new_config_obj = Config(
diff --git a/apps/stable_diffusion/src/utils/stable_args.py b/apps/stable_diffusion/src/utils/stable_args.py
@@ -270,6 +270,20 @@ def path_expand(s):
     help="flag to clear all mlir and vmfb from common locations. Recompiling will take several minutes",
 )
 
+p.add_argument(
+    "--save_metadata_to_json",
+    default=False,
+    action=argparse.BooleanOptionalAction,
+    help="flag for whether or not to save a generation information json file with the image.",
+)
+
+p.add_argument(
+    "--write_metadata_to_png",
+    default=False,
+    action=argparse.BooleanOptionalAction,
+    help="flag for whether or not to save generation information in PNG chunk text to generated images.",
+)
+
 ##############################################################################
 ### Web UI flags
 ##############################################################################
diff --git a/apps/stable_diffusion/web/index.py b/apps/stable_diffusion/web/index.py
@@ -148,6 +148,17 @@ def resource_path(relative_path):
                             step=0.1,
                             label="CFG Scale",
                         )
+                    with gr.Row():
+                        save_metadata_to_png = gr.Checkbox(
+                            label="Save prompt information to PNG",
+                            value=False,
+                            interactive=True,
+                        )
+                        save_metadata_to_json = gr.Checkbox(
+                            label="Save prompt information to JSON file",
+                            value=False,
+                            interactive=True,
+                        )
                 with gr.Row():
                     seed = gr.Number(value=-1, precision=0, label="Seed")
                     available_devices = get_available_devices()
@@ -211,6 +222,8 @@ def resource_path(relative_path):
                 precision,
                 device,
                 max_length,
+                save_metadata_to_json,
+                save_metadata_to_png,
             ],
             outputs=[gallery, std_output],
             show_progress=args.progress_bar,
@@ -233,6 +246,8 @@ def resource_path(relative_path):
                 precision,
                 device,
                 max_length,
+                save_metadata_to_json,
+                save_metadata_to_png,
             ],
             outputs=[gallery, std_output],
             show_progress=args.progress_bar,
diff --git a/shark/examples/shark_inference/stable_diffusion/main.py b/shark/examples/shark_inference/stable_diffusion/main.py
@@ -9,7 +9,7 @@
 
 from transformers import CLIPTextModel, CLIPTokenizer
 import torch
-from PIL import Image
+from PIL import Image, PngImagePlugin
 from diffusers import (
     LMSDiscreteScheduler,
     PNDMScheduler,
@@ -329,11 +329,27 @@ def end_profiling(device):
                     progressive=True,
                 )
             else:
-                pil_images[i].save(output_path / f"{img_name}.png", "PNG")
+                pngInfo = PngImagePlugin.PngInfo()
+
+                if args.write_metadata_to_png:
+                    model_name = ""
+                    if args.ckpt_loc:
+                        model_name = Path(args.ckpt_loc).name
+                    else:
+                        model_name = json_store["hf_model_id"]
+                    pngInfo.add_text(
+                        "parameters",
+                        f"{json_store['prompt']}\nNegative prompt: {json_store['negative prompt']}\nSteps:{json_store['steps']}, Sampler: {json_store['scheduler']}, CFG scale: {json_store['guidance_scale']}, Seed: {json_store['seed']}, Size: {args.width}x{args.height}, Model: {model_name}",
+                    )
+
+                pil_images[i].save(
+                    output_path / f"{img_name}.png", "PNG", pnginfo=pngInfo
+                )
                 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"
+                        "Image saved as png instead. Supported formats: png / jpg"
                     )
-            with open(output_path / f"{img_name}.json", "w") as f:
-                f.write(json.dumps(json_store, indent=4))
+            if args.save_metadata_to_json:
+                with open(output_path / f"{img_name}.json", "w") as f:
+                    f.write(json.dumps(json_store, indent=4))
diff --git a/shark/examples/shark_inference/stable_diffusion/stable_args.py b/shark/examples/shark_inference/stable_diffusion/stable_args.py
@@ -283,6 +283,20 @@ def is_valid_file(arg):
     help="flag to clear all mlir and vmfb from common locations. Recompiling will take several minutes",
 )
 
+p.add_argument(
+    "--save_metadata_to_json",
+    default=True,
+    action=argparse.BooleanOptionalAction,
+    help="flag for whether or not to save a generation information json file with the image.",
+)
+
+p.add_argument(
+    "--write_metadata_to_png",
+    default=False,
+    action=argparse.BooleanOptionalAction,
+    help="flag for whether or not to save generation information in PNG chunk text to generated images.",
+)
+
 ##############################################################################
 ### Web UI flags
 ##############################################################################
diff --git a/web/models/stable_diffusion/main.py b/web/models/stable_diffusion/main.py
@@ -1,11 +1,12 @@
 import torch
 import os
-from PIL import Image
+from PIL import Image, PngImagePlugin
 from tqdm.auto import tqdm
 from models.stable_diffusion.cache_objects import model_cache
 from models.stable_diffusion.stable_args import args
 from models.stable_diffusion.utils import disk_space_check
 from random import randint
+import json
 import numpy as np
 import time
 import sys
@@ -92,11 +93,22 @@ def save_output_img(output_img):
         )
     else:
         out_img_path = Path(generated_imgs_path, f"{out_img_name}.png")
-        output_img.save(out_img_path, "PNG")
+        pngInfo = PngImagePlugin.PngInfo()
+
+        if args.write_metadata_to_png:
+            pngInfo.add_text(
+                "parameters",
+                f"{args.prompts}\nNegative prompt: {args.negative_prompts}\nSteps:{args.steps}, Sampler: {args.scheduler}, CFG scale: {args.guidance_scale}, Seed: {args.seed}, Size: {args.width}x{args.height}, Model: {args.variant}",
+            )
+
+        output_img.save(
+            output_path / f"{out_img_name}.png", "PNG", pnginfo=pngInfo
+        )
+
         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"
+                "Image saved as png instead. Supported formats: png / jpg"
             )
 
     new_entry = {
@@ -117,6 +129,11 @@ def save_output_img(output_img):
         dictwriter_obj.writerow(new_entry)
         csv_obj.close()
 
+    if args.save_metadata_to_json:
+        del new_entry["OUTPUT"]
+        with open(f"{output_path}/{out_img_name}.json", "w") as f:
+            json.dump(new_entry, f, indent=4)
+
 
 def stable_diff_inf(
     prompt: str,
@@ -209,6 +226,7 @@ def stable_diff_inf(
 
     avg_ms = 0
     for i, t in tqdm(enumerate(scheduler.timesteps)):
+
         step_start = time.time()
         timestep = torch.tensor([t]).to(dtype).detach().numpy()
         latent_model_input = scheduler.scale_model_input(latents, t)
diff --git a/web/models/stable_diffusion/stable_args.py b/web/models/stable_diffusion/stable_args.py
@@ -226,6 +226,20 @@
     help="flag to clear all mlir and vmfb from common locations. Recompiling will take several minutes",
 )
 
+p.add_argument(
+    "--save_metadata_to_json",
+    default=False,
+    action=argparse.BooleanOptionalAction,
+    help="flag for whether or not to save a generation information json file with the image.",
+)
+
+p.add_argument(
+    "--write_metadata_to_png",
+    default=False,
+    action=argparse.BooleanOptionalAction,
+    help="flag for whether or not to save generation information in PNG chunk text to generated images.",
+)
+
 ##############################################################################
 ### Web UI flags
 ##############################################################################
