Update app.py

app.py

@@ -74,9 +74,21 @@ config.model.pop("latent_std")
 model_primx = load_from_config(config.model)
 # load rembg
 rembg_session = rembg.new_session()
+current_fg_state = None
+
+# background removal function
+def background_remove_process(input_image):
+    input_image = remove_background(input_image, rembg_session)
+    input_image = resize_foreground(input_image, 0.85)
+    input_cond_preview_pil = input_image
+    raw_image = np.array(input_image)
+    mask = (raw_image[..., -1][..., None] > 0) * 1
+    raw_image = raw_image[..., :3] * mask
+    input_cond = torch.from_numpy(np.array(raw_image)[None, ...]).to(device)
+    return gr.update(interactive=True), input_cond, input_cond_preview_pil
 
 # process function
-def process(input_image, input_num_steps, input_seed=42, input_cfg=6.0):
+def process(input_cond, input_num_steps, input_seed=42, input_cfg=6.0):
     # seed
     torch.manual_seed(input_seed)
 
@@ -91,16 +103,8 @@ def process(input_image, input_num_steps, input_seed=42, input_cfg=6.0):
     fwd_fn = model.forward_with_cfg
 
     # text-conditioned
-    if input_image is None:
+    if input_cond is None:
         raise NotImplementedError
-    # image-conditioned (may also input text, but no text usually works too)
-    else:
-        input_image = remove_background(input_image, rembg_session)
-        input_image = resize_foreground(input_image, 0.85)
-        raw_image = np.array(input_image)
-        mask = (raw_image[..., -1][..., None] > 0) * 1
-        raw_image = raw_image[..., :3] * mask
-        input_cond = torch.from_numpy(np.array(raw_image)[None, ...]).to(device)
     
     with torch.no_grad():
         latent = torch.randn(1, config.model.num_prims, 1, 4, 4, 4)
@@ -178,8 +182,11 @@ with block:
     
     with gr.Row(variant='panel'):
         with gr.Column(scale=1):
-            # input image
-            input_image = gr.Image(label="image", type='pil')
+            with gr.Row():
+                # input image
+                input_image = gr.Image(label="image", type='pil')
+                # background removal
+                removal_previewer = gr.Image(label="Background Removal Preview", type='pil', interactive=False)
             # inference steps
             input_num_steps = gr.Radio(choices=[25, 50, 100, 200], label="DDIM steps", value=25)
             # random seed
@@ -187,7 +194,7 @@ with block:
             # random seed
             input_seed = gr.Slider(label="random seed", minimum=0, maximum=10000, step=1, value=42, info="Try different seed if the result is not satisfying as this is a generative model!")
             # gen button
-            button_gen = gr.Button("Generate")
+            button_gen = gr.Button(value="Generate", interactive=False)
             export_glb_btn = gr.Button(value="Export GLB", interactive=False)
 
         with gr.Column(scale=1):
@@ -231,15 +238,16 @@ with block:
                         outputs=[output_glb],
                     )
 
-        button_gen.click(process, inputs=[input_image, input_num_steps, input_seed, input_cfg], outputs=[output_rgb_video, output_prim_video, output_mat_video, export_glb_btn])
+        input_image.change(background_remove_process, inputs=[input_image], outputs=[button_gen, current_fg_state, removal_previewer])
+
+        button_gen.click(process, inputs=[current_fg_state, input_num_steps, input_seed, input_cfg], outputs=[output_rgb_video, output_prim_video, output_mat_video, export_glb_btn])
         
         export_glb_btn.click(export_mesh, inputs=[], outputs=[output_glb, hdr_row])
     
     gr.Examples(
         examples=[
-            "assets/examples/fruit_elephant.jpg",
-            "assets/examples/mei_ling_panda.png",
-            "assets/examples/shuai_panda_notail.png",
+            os.path.join("assets/examples", f)
+            for f in os.listdir("assets/examples")
         ],
         inputs=[input_image],
         outputs=[output_rgb_video, output_prim_video, output_mat_video, export_glb_btn],