Update app.py

app.py

@@ -84,7 +84,6 @@ def process(input_image, input_num_steps, input_seed=42, input_cfg=6.0):
     output_rgb_video_path = os.path.join(config.output_dir, GRADIO_RGB_VIDEO_PATH)
     output_prim_video_path = os.path.join(config.output_dir, GRADIO_PRIM_VIDEO_PATH)
     output_mat_video_path = os.path.join(config.output_dir, GRADIO_MAT_VIDEO_PATH)
-    output_glb_path = os.path.join(config.output_dir, GRADIO_GLB_PATH)
 
     respacing = "ddim{}".format(input_num_steps)
     diffusion = create_diffusion(timestep_respacing=respacing, **config.diffusion)
@@ -139,7 +138,14 @@ def process(input_image, input_num_steps, input_seed=42, input_cfg=6.0):
         prim_params = {'srt_param': recon_srt_param[0].detach().cpu(), 'feat_param': recon_feat_param[0].detach().cpu()}
         torch.save({'model_state_dict': prim_params}, "{}/denoised.pt".format(config.output_dir))
 
+    return output_rgb_video_path, output_prim_video_path, output_mat_video_path, gr.update(interactive=True)
+
+def export_mesh(remesh=False, decimate=100000, mc_resolution=256):
     # exporting GLB mesh
+    output_glb_path = os.path.join(config.output_dir, GRADIO_GLB_PATH)
+    config.inference.remesh = remesh
+    config.inference.decimate = decimate
+    config.inference.mc_resolution = mc_resolution
     denoise_param_path = os.path.join(config.output_dir, 'denoised.pt')
     primx_ckpt_weight = torch.load(denoise_param_path, map_location='cpu')['model_state_dict']
     model_primx.load_state_dict(primx_ckpt_weight)
@@ -148,8 +154,7 @@ def process(input_image, input_num_steps, input_seed=42, input_cfg=6.0):
     with torch.no_grad():
         model_primx.srt_param[:, 1:4] *= 0.85
         extract_texmesh(config.inference, model_primx, config.output_dir, device)
-
-    return output_rgb_video_path, output_prim_video_path, output_mat_video_path, output_glb_path
+    return output_glb_path, gr.update(visible=True)
 
 # gradio UI
 _TITLE = '''3DTopia-XL'''
@@ -183,15 +188,15 @@ with block:
             input_seed = gr.Slider(label="random seed", minimum=0, maximum=10000, step=1, value=42)
             # gen button
             button_gen = gr.Button("Generate")
+            export_glb_btn = gr.Button(value="Export GLB", interactive=False)
 
         with gr.Column(scale=1):
-            with gr.Tab("Video"):
-                with gr.Row():
-                    # final video results
-                    output_rgb_video = gr.Video(label="video")
-                    output_prim_video = gr.Video(label="video")
-                    output_mat_video = gr.Video(label="video")
-            with gr.Tab("GLB"):
+            with gr.Row():
+                # final video results
+                output_rgb_video = gr.Video(label="RGB")
+                output_prim_video = gr.Video(label="Primitives")
+                output_mat_video = gr.Video(label="Material")
+            with gr.Row():
                 # glb file
                 output_glb = LitModel3D(
                     label="3D GLB Model",
@@ -201,8 +206,33 @@ with block:
                     contrast=1.0,
                     scale=1.0,
                 )
-
-        button_gen.click(process, inputs=[input_image, input_num_steps, input_seed, input_cfg], outputs=[output_rgb_video, output_prim_video, output_mat_video, output_glb])
+            with gr.Column(visible=False, scale=1.0) as hdr_row:
+                gr.Markdown("""## HDR Environment Map
+                
+                Select / Upload an HDR environment map to light the 3D model.
+                """)
+                with gr.Row():
+                    hdr_illumination_file = gr.File(
+                        label="HDR Envmap", file_types=[".hdr"], file_count="single"
+                    )
+                    example_hdris = [
+                        os.path.join("assets/hdri", f)
+                        for f in os.listdir("assets/hdri")
+                    ]
+                    hdr_illumination_example = gr.Examples(
+                        examples=example_hdris,
+                        inputs=hdr_illumination_file,
+                    )
+
+                    hdr_illumination_file.change(
+                        lambda x: gr.update(env_map=x.name if x is not None else None),
+                        inputs=hdr_illumination_file,
+                        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])
+        
+        export_glb_btn.click(export_mesh, inputs=[], outputs=[output_glb, hdr_row])
     
     gr.Examples(
         examples=[
@@ -211,7 +241,7 @@ with block:
             "assets/examples/shuai_panda_notail.png",
         ],
         inputs=[input_image],
-        outputs=[output_rgb_video, output_prim_video, output_mat_video, output_glb],
+        outputs=[output_rgb_video, output_prim_video, output_mat_video, export_glb_btn],
         fn=lambda x: process(input_image=x),
         cache_examples=False,
         label='Single Image to 3D PBR Asset'