[SD] Update metadata info and canvas size (huggingface#1109)

* [SD] Save missing metadata in case of img2img and outpaint

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

* [SD] Update the canvas size for inpaint/outpaint

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

* [SD] Update output gallery on each inference

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

---------

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

Author: Shukla-Gaurav

apps/stable_diffusion/scripts/img2img.py

@@ -146,6 +146,7 @@ def img2img_inf(
     generated_imgs = []
     seeds = []
     img_seed = utils.sanitize_seed(seed)
+    extra_info = {"STRENGTH": strength}
     for current_batch in range(batch_count):
         if current_batch > 0:
             img_seed = utils.sanitize_seed(-1)
@@ -165,7 +166,7 @@ def img2img_inf(
             args.use_base_vae,
             cpu_scheduling,
         )
-        save_output_img(out_imgs[0], img_seed)
+        save_output_img(out_imgs[0], img_seed, extra_info)
         generated_imgs.extend(out_imgs)
         seeds.append(img_seed)
         img2img_obj.log += "\n"
@@ -260,5 +261,6 @@ def img2img_inf(
     text_output += img2img_obj.log
     text_output += f"\nTotal image generation time: {total_time:.4f}sec"
 
-    save_output_img(generated_imgs[0], seed)
+    extra_info = {"STRENGTH": args.strength}
+    save_output_img(generated_imgs[0], seed, extra_info)
     print(text_output)

apps/stable_diffusion/scripts/outpaint.py

@@ -35,7 +35,7 @@ class Config:
 def outpaint_inf(
     prompt: str,
     negative_prompt: str,
-    init_image: str,
+    init_image: Image,
     pixels: int,
     mask_blur: int,
     directions: list,
@@ -66,7 +66,7 @@ def outpaint_inf(
     args.guidance_scale = guidance_scale
     args.steps = steps
     args.scheduler = scheduler
-    args.img_path = init_image
+    args.img_path = "not none"
 
     # set ckpt_loc and hf_model_id.
     types = (
@@ -143,7 +143,6 @@ def outpaint_inf(
     generated_imgs = []
     seeds = []
     img_seed = utils.sanitize_seed(seed)
-    image = Image.open(args.img_path)
 
     left = True if "left" in directions else False
     right = True if "right" in directions else False
@@ -156,7 +155,7 @@ def outpaint_inf(
         out_imgs = outpaint_obj.generate_images(
             prompt,
             negative_prompt,
-            image,
+            init_image,
             pixels,
             mask_blur,
             left,
@@ -271,5 +270,22 @@ def outpaint_inf(
         text_output += outpaint_obj.log
         text_output += f"\nTotal image generation time: {total_time:.4f}sec"
 
-        save_output_img(generated_imgs[0], seed)
+        # save this information as metadata of output generated image.
+        directions = []
+        if args.left:
+            directions.append("left")
+        if args.right:
+            directions.append("right")
+        if args.top:
+            directions.append("up")
+        if args.bottom:
+            directions.append("down")
+        extra_info = {
+            "PIXELS": args.pixels,
+            "MASK_BLUR": args.mask_blur,
+            "DIRECTIONS": directions,
+            "NOISE_Q": args.noise_q,
+            "COLOR_VARIATION": args.color_variation,
+        }
+        save_output_img(generated_imgs[0], seed, extra_info)
         print(text_output)

apps/stable_diffusion/scripts/txt2img.py

@@ -158,6 +158,7 @@ def txt2img_inf(
         generated_imgs.extend(out_imgs)
         seeds.append(img_seed)
         txt2img_obj.log += "\n"
+        yield generated_imgs, txt2img_obj.log
 
     total_time = time.time() - start_time
     text_output = f"prompt={args.prompts}"
@@ -168,10 +169,10 @@ def txt2img_inf(
         f"\nsteps={steps}, guidance_scale={guidance_scale}, seed={seeds}"
     )
     text_output += f"\nsize={height}x{width}, batch_count={batch_count}, batch_size={batch_size}, max_length={args.max_length}"
-    text_output += txt2img_obj.log
+    # text_output += txt2img_obj.log
     text_output += f"\nTotal image generation time: {total_time:.4f}sec"
 
-    return generated_imgs, text_output
+    yield generated_imgs, text_output
 
 
 if __name__ == "__main__":

apps/stable_diffusion/src/utils/utils.py

@@ -551,7 +551,7 @@ def clear_all():
 
 
 # save output images and the inputs corresponding to it.
-def save_output_img(output_img, img_seed):
+def save_output_img(output_img, img_seed, extra_info={}):
     output_path = args.output_dir if args.output_dir else Path.cwd()
     generated_imgs_path = Path(
         output_path, "generated_imgs", dt.now().strftime("%Y%m%d")
@@ -604,6 +604,8 @@ def save_output_img(output_img, img_seed):
         "OUTPUT": out_img_path,
     }
 
+    new_entry.update(extra_info)
+
     with open(csv_path, "a") as csv_obj:
         dictwriter_obj = DictWriter(csv_obj, fieldnames=list(new_entry.keys()))
         dictwriter_obj.writerow(new_entry)

apps/stable_diffusion/web/ui/img2img_ui.py

@@ -75,7 +75,9 @@
                         elem_id="negative_prompt_box",
                     )
 
-                init_image = gr.Image(label="Input Image", type="pil")
+                init_image = gr.Image(label="Input Image", type="pil").style(
+                    height=300, full_width=True
+                )
 
                 with gr.Accordion(label="Advanced Options", open=False):
                     with gr.Row():

apps/stable_diffusion/web/ui/inpaint_ui.py

@@ -76,7 +76,7 @@
                     source="upload",
                     tool="sketch",
                     type="pil",
-                )
+                ).style(height=350, full_width=True)
 
                 with gr.Accordion(label="Advanced Options", open=False):
                     with gr.Row():

apps/stable_diffusion/web/ui/outpaint_ui.py

@@ -71,7 +71,9 @@
                         elem_id="negative_prompt_box",
                     )
 
-                init_image = gr.Image(label="Input Image", type="filepath")
+                init_image = gr.Image(label="Input Image", type="pil").style(
+                    height=300, full_width=True
+                )
 
                 with gr.Accordion(label="Advanced Options", open=False):
                     with gr.Row():