Update app.py

app.py

@@ -50,8 +50,8 @@ amp = True
 precision_dtype = torch.float16
 
 rm = RayMarcher(
-    config.image_height,
-    config.image_width,
+    256,
+    256,
     **config.rm,
 ).to(device)
 
@@ -143,12 +143,16 @@ def process(input_cond, input_num_steps, input_seed=42, input_cfg=6.0):
 
     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):
+def export_mesh(remesh="No", mc_resolution=256, decimate=100000):
     # exporting GLB mesh
     output_glb_path = os.path.join(config.output_dir, GRADIO_GLB_PATH)
-    config.inference.remesh = remesh
+    if remesh == "No":
+        config.inference.remesh = False
+    elif remesh == "Yes":
+        config.inference.remesh = True
     config.inference.decimate = decimate
     config.inference.mc_resolution = mc_resolution
+    config.inference.batch_size = 8192
     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)
@@ -195,6 +199,9 @@ with block:
             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(value="Generate", interactive=False)
+            with gr.Row():
+                input_mc_resolution = gr.Radio(choices=[32, 64, 128, 256], label="MC Resolution", value=128, info="Cube resolution for mesh extraction")
+                input_remesh = gr.Radio(choices=["No", "Yes"], label="Remesh", value="No", info="Remesh or not?")
             export_glb_btn = gr.Button(value="Export GLB", interactive=False)
 
         with gr.Column(scale=1):
@@ -240,7 +247,7 @@ with block:
 
     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])
+    export_glb_btn.click(export_mesh, inputs=[input_remesh, input_mc_resolution], outputs=[output_glb, hdr_row])
     
     gr.Examples(
         examples=[