add new func to show canonical image

app.py

@@ -218,6 +218,7 @@ def edit_with_pnp(input_video, prompt, num_steps, guidance_scale, seed, n_prompt
         ori_size = canonical_image.size
         image = processor_partial(canonical_image.resize((size_, size_)), detect_resolution=size_, image_resolution=size_)
         image = image.resize(ori_size, resample=Image.BILINEAR)
+        image.save("control_map.png")
         
         generator = torch.manual_seed(seed) if seed != -1 else None
         output_images = pipe(
@@ -228,6 +229,8 @@ def edit_with_pnp(input_video, prompt, num_steps, guidance_scale, seed, n_prompt
             negative_prompt=n_prompt,
             generator=generator
         ).images
+
+        output_images[0].save("edited_canonical_image.png")
         # output_images[0] = output_images[0].resize(ori_size, resample=Image.BILINEAR)
     
     else:
@@ -245,6 +248,7 @@ def edit_with_pnp(input_video, prompt, num_steps, guidance_scale, seed, n_prompt
         image = image[:, :, None]
         image = np.concatenate([image, image, image], axis=2)
         image = Image.fromarray(image)
+        image.save("control_map.png")
         
         generator = torch.manual_seed(seed) if seed != -1 else None
         output_images = pipe(
@@ -255,11 +259,18 @@ def edit_with_pnp(input_video, prompt, num_steps, guidance_scale, seed, n_prompt
             negative_prompt=n_prompt,
             generator=generator
         ).images
+
+        output_images[0].save("edited_canonical_image.png")
     
     edit_video_path = NaRCan_make_video(output_images[0], pth_path, frames_path)
+    edit_image_path = [
+        (image_path, "canonical image"), 
+        ("control_map.png", "control map"),
+        ("edited_canonical_image.png", "edited canonical image")
+        ]
 
     # Here we return the first output image as the result
-    return edit_video_path
+    return edit_video_path, edit_image_path
 
 
 ########
@@ -273,7 +284,7 @@ intro = """
    NaRCan - <small>Natural Refined Canonical Image</small>
 </h1>
 <span>[<a target="_blank" href="https://koi953215.github.io/NaRCan_page/">Project page</a>], [<a target="_blank" href="https://huggingface.co/papers/2406.06523">Paper</a>]</span>
-<div style="display:flex; justify-content: center;margin-top: 0.5em">Each edit takes 30 sec ~ 2 min </div>
+<div style="display:flex; justify-content: center;margin-top: 0.5em">Try selecting different control types (Canny or Lineart) in Advanced options!</div>
 </div>
 """
 
@@ -303,6 +314,10 @@ with gr.Blocks(css="style.css") as demo:
                             # placeholder="bear, Van Gogh Style"
                         )
     
+
+    with gr.Row():
+        canonical_result = gr.Gallery(label="Edited Canonical Image", columns=3)
+
                
     with gr.Row():
         run_button = gr.Button("Edit your video!", visible=True)
@@ -351,7 +366,7 @@ with gr.Blocks(css="style.css") as demo:
                                n_prompt,
                                control_type,
                                ],
-                                 outputs = [output_video]
+                                 outputs = [output_video, canonical_result]
                                 )
 
     gr.Examples(