Predict normal for multiviews by controlnet

gradio_app.py

@@ -1,10 +1,11 @@
 import os
+import copy
 import torch
 import fire
 import gradio as gr
 from PIL import Image
 from functools import partial
-from diffusers import DiffusionPipeline, EulerAncestralDiscreteScheduler
+from diffusers import DiffusionPipeline, EulerAncestralDiscreteScheduler, ControlNetModel
 from share_btn import community_icon_html, loading_icon_html, share_js
 
 import cv2
@@ -122,7 +123,7 @@ def save_image(image, original_image):
     os.remove(out_path)
     os.remove(in_path)
 
-def gen_multiview(pipeline, predictor, input_image, scale_slider, steps_slider, seed, output_processing=False, original_image=None):
+def gen_multiview(pipeline, pipeline_normal, predictor, input_image, scale_slider, steps_slider, seed, output_processing=False, original_image=None, out_normal=True):
     seed = int(seed)
     torch.manual_seed(seed)
     image = pipeline(input_image, 
@@ -131,6 +132,15 @@ def gen_multiview(pipeline, predictor, input_image, scale_slider, steps_slider,
                     generator=torch.Generator(pipeline.device).manual_seed(seed)).images[0]
     side_len = image.width//2
     subimages = [image.crop((x, y, x + side_len, y+side_len)) for y in range(0, image.height, side_len) for x in range(0, image.width, side_len)]
+
+    # normal images
+    out_images_normal = [gr.Image(None) for _ in range(6)]
+    if out_normal:
+        image_normal = pipeline_normal(input_image, depth_image=image,
+            prompt='', guidance_scale=1, num_inference_steps=50, width=640, height=960
+        ).images[0]
+        subimages_normal = [image_normal.crop((x, y, x + side_len, y+side_len)) for y in range(0, image_normal.height, side_len) for x in range(0, image_normal.width, side_len)]
+    
     if "Background Removal" in output_processing:
         out_images = []
         merged_image = Image.new('RGB', (640, 960))
@@ -142,9 +152,17 @@ def gen_multiview(pipeline, predictor, input_image, scale_slider, steps_slider,
             y = (i % 3) * 320
             merged_image.paste(sub_image, (x, y))
         save_image(merged_image, original_image)
-        return out_images + [merged_image]
+
+        if out_normal:
+            out_images_normal = []
+            # merged_image_normal = Image.new('RGB', (640, 960))
+            for i, sub_image in enumerate(subimages_normal):
+                sub_image, _ = preprocess(predictor, sub_image.convert('RGB'), segment=True, rescale=False)
+                out_images_normal.append(sub_image)
+
+        return out_images + [merged_image] + out_images_normal
     save_image(image, original_image)
-    return subimages + [image]
+    return subimages + [image] + out_images_normal
 
 
 def run_demo():
@@ -159,6 +177,14 @@ def run_demo():
     )
     pipeline.to(f'cuda:{_GPU_ID}')
 
+    normal_pipeline = copy.copy(pipeline)
+    controlnet = ControlNetModel.from_pretrained(
+        "sudo-ai/controlnet-zp12-normal-gen-v1", 
+        torch_dtype=torch.float16, use_auth_token=os.environ["HF_TOKEN"]
+    )
+    normal_pipeline.add_controlnet(controlnet, conditioning_scale=1.0)
+    normal_pipeline.to(f'cuda:{_GPU_ID}')
+
     predictor = sam_init()
 
     custom_theme = gr.themes.Soft(primary_hue="blue").set(
@@ -187,6 +213,8 @@ def run_demo():
                     label='Examples (click one of the images below to start)',
                     examples_per_page=10
                 )
+                with gr.Row():
+                    out_normal = gr.Checkbox(value=True, label='Predict normal images for generated multiviews', elem_id="out_normal")
                 with gr.Accordion('Advanced options', open=False):
                     with gr.Row():
                         with gr.Column():
@@ -213,6 +241,14 @@ def run_demo():
                     view_4 = gr.Image(interactive=False, height=240, show_label=False)
                     view_5 = gr.Image(interactive=False, height=240, show_label=False)
                     view_6 = gr.Image(interactive=False, height=240, show_label=False)
+                with gr.Row():
+                    norm_1 = gr.Image(interactive=False, height=240, show_label=False)
+                    norm_2 = gr.Image(interactive=False, height=240, show_label=False)
+                    norm_3 = gr.Image(interactive=False, height=240, show_label=False)
+                with gr.Row():
+                    norm_4 = gr.Image(interactive=False, height=240, show_label=False)
+                    norm_5 = gr.Image(interactive=False, height=240, show_label=False)
+                    norm_6 = gr.Image(interactive=False, height=240, show_label=False)
                 full_view = gr.Image(visible=False, interactive=False, elem_id="six_view")
                 with gr.Group(elem_id="share-btn-container", visible=False) as share_group:
                     community_icon = gr.HTML(community_icon_html)
@@ -227,9 +263,10 @@ def run_demo():
             ).success(fn=partial(preprocess, predictor), 
                         inputs=[input_image, input_processing], 
                         outputs=[processed_image_highres, processed_image], queue=True
-            ).success(fn=partial(gen_multiview, pipeline, predictor), 
-                        inputs=[processed_image_highres, scale_slider, steps_slider, seed, output_processing, input_image],
-                        outputs=[view_1, view_2, view_3, view_4, view_5, view_6, full_view], queue=True
+            ).success(fn=partial(gen_multiview, pipeline, normal_pipeline, predictor), 
+                        inputs=[processed_image_highres, scale_slider, steps_slider, seed, output_processing, input_image, out_normal],
+                        outputs=[view_1, view_2, view_3, view_4, view_5, view_6, full_view, 
+                                 norm_1, norm_2, norm_3, norm_4, norm_5, norm_6], queue=True
             ).success(show_share_btn, outputs=share_group, queue=False)
         
         share_button.click(None, [], [], _js=share_js)