Better ZeroGPU utilization

app.py

@@ -41,9 +41,9 @@ _DESCRIPTION = '''
 
 * The demo is still under construction, and more features are expected to be implemented soon.
 
-* The demo takes around 50 seconds on L4.
+* The demo takes around 50 seconds on L4, and about 60 seconds on Huggingface ZeroGPU.
 
-* If the Gradio Demo unfortunately hangs or is very crowded, you can use the Gradio Demo or Online Demo. The Online Demo is free to try, and the registration invitation code is `aiuni24`. However, the Online Demo is slightly different from the Gradio Demo, in that the inference speed is slower, and the generation results is less stable, but the quality of the texture is better.
+* If the Huggingface Demo unfortunately hangs or is very crowded, you can use the Gradio Demo or Online Demo. The Online Demo is free to try, and the registration invitation code is `aiuni24`. However, the Online Demo is slightly different from the Gradio Demo, in that the inference speed is slower, and the generation results is less stable, but the quality of the texture is better.
 
 
 '''
@@ -53,7 +53,7 @@ def launch():
         
     with gr.Blocks(
         title=_TITLE,
-        theme=gr.themes.Monochrome(),
+        # theme=gr.themes.Monochrome(),
     ) as demo:
         with gr.Row():
             with gr.Column(scale=1):

gradio_app/gradio_3dgen.py

@@ -8,19 +8,30 @@ from gradio_app.custom_models.mvimg_prediction import run_mvprediction
 from gradio_app.custom_models.normal_prediction import predict_normals
 from scripts.refine_lr_to_sr import run_sr_fast
 from scripts.utils import save_glb_and_video
-from scripts.multiview_inference import geo_reconstruct
+# from scripts.multiview_inference import geo_reconstruct
+from scripts.multiview_inference import geo_reconstruct_part1, geo_reconstruct_part2, geo_reconstruct_part3
+
+@spaces.GPU
+def run_mv(preview_img, input_processing, seed):
+    if preview_img.size[0] <= 512:
+        preview_img = run_sr_fast([preview_img])[0]
+    rgb_pils, front_pil = run_mvprediction(preview_img, remove_bg=input_processing, seed=int(seed)) # 6s
+    return rgb_pils, front_pil
 
-@spaces.GPU(duration=100)
 def generate3dv2(preview_img, input_processing, seed, render_video=True, do_refine=True, expansion_weight=0.1, init_type="std"):
     if preview_img is None:
-        raise gr.Error("preview_img is none")
+        raise gr.Error("The input image is none!")
     if isinstance(preview_img, str):
         preview_img = Image.open(preview_img)
     
-    if preview_img.size[0] <= 512:
-        preview_img = run_sr_fast([preview_img])[0]
-    rgb_pils, front_pil = run_mvprediction(preview_img, remove_bg=input_processing, seed=int(seed)) # 6s
-    new_meshes = geo_reconstruct(rgb_pils, None, front_pil, do_refine=do_refine, predict_normal=True, expansion_weight=expansion_weight, init_type=init_type)
+    rgb_pils, front_pil = run_mv(preview_img, input_processing, seed)
+
+    vertices, faces, img_list = geo_reconstruct_part1(rgb_pils, None, front_pil, do_refine=do_refine, predict_normal=True, expansion_weight=expansion_weight, init_type=init_type)
+
+    meshes = geo_reconstruct_part2(vertices, faces)
+
+    new_meshes = geo_reconstruct_part3(meshes, img_list)
+
     vertices = new_meshes.verts_packed()
     vertices = vertices / 2 * 1.35
     vertices[..., [0, 2]] = - vertices[..., [0, 2]]
@@ -32,7 +43,7 @@ def generate3dv2(preview_img, input_processing, seed, render_video=True, do_refi
 #######################################
 def create_ui(concurrency_id="wkl"):
     with gr.Row():
-        with gr.Column(scale=2):
+        with gr.Column(scale=1):
             input_image = gr.Image(type='pil', image_mode='RGBA', label='Frontview')
             
             example_folder = os.path.join(os.path.dirname(__file__), "./examples")
@@ -46,7 +57,7 @@ def create_ui(concurrency_id="wkl"):
             )
             
 
-        with gr.Column(scale=3):
+        with gr.Column(scale=1):
             # export mesh display
             output_mesh = gr.Model3D(value=None, label="Mesh Model", show_label=True, height=320)
             output_video = gr.Video(label="Preview", show_label=True, show_share_button=True, height=320, visible=False)

scripts/multiview_inference.py

@@ -95,6 +95,62 @@ def geo_reconstruct(rgb_pils, normal_pils, front_pil, do_refine=False, predict_n
     normal_stg2 = [img.resize((1024, 1024)) for img in rm_normals]  # reduce computation on huggingface demo, use 1024 instead of 2048
 
     vertices, faces = run_mesh_refine(vertices, faces, normal_stg2, steps=100, start_edge_len=0.02, end_edge_len=0.005, decay=0.99, update_normal_interval=20, update_warmup=5, return_mesh=False, process_inputs=False, process_outputs=False)
+
     meshes = simple_clean_mesh(to_pyml_mesh(vertices, faces), apply_smooth=True, stepsmoothnum=1, apply_sub_divide=True, sub_divide_threshold=0.25).to("cuda")
+    
+    new_meshes = multiview_color_projection(meshes, img_list, resolution=1024, device="cuda", complete_unseen=True, confidence_threshold=0.2, cameras_list = get_cameras_list([0, 90, 180, 270], "cuda", focal=1))
+    return new_meshes
+
+########################
+import spaces
+
+@spaces.GPU
+def geo_reconstruct_part1(rgb_pils, normal_pils, front_pil, do_refine=False, predict_normal=True, expansion_weight=0.1, init_type="std"):
+    if front_pil.size[0] <= 512:
+        front_pil = run_sr_fast([front_pil])[0]
+    if do_refine:
+        refined_rgbs = refine_rgb(rgb_pils, front_pil)  # 6s
+    else:
+        refined_rgbs = [rgb.resize((512, 512), resample=Image.LANCZOS) for rgb in rgb_pils]
+    img_list = [front_pil] + run_sr_fast(refined_rgbs[1:])
+    
+    if predict_normal:
+        rm_normals = predict_normals([img.resize((512, 512), resample=Image.LANCZOS) for img in img_list], guidance_scale=1.5)
+    else:
+        rm_normals = simple_remove([img.resize((512, 512), resample=Image.LANCZOS) for img in normal_pils])
+    # transfer the alpha channel of rm_normals to img_list
+    for idx, img in enumerate(rm_normals):
+        if idx == 0 and img_list[0].mode == "RGBA":
+            temp = img_list[0].resize((2048, 2048))
+            rm_normals[0] = Image.fromarray(np.concatenate([np.array(rm_normals[0])[:, :, :3], np.array(temp)[:, :, 3:4]], axis=-1))
+            continue
+        img_list[idx] = Image.fromarray(np.concatenate([np.array(img_list[idx]), np.array(img)[:, :, 3:4]], axis=-1))
+    assert img_list[0].mode == "RGBA"
+    assert np.mean(np.array(img_list[0])[..., 3]) < 250
+    
+    img_list = [img_list[0]] + erode_alpha(img_list[1:])
+    normal_stg1 = [img.resize((512, 512)) for img in rm_normals]
+    if init_type in ["std", "thin"]:
+        meshes = fast_geo(normal_stg1[0], normal_stg1[2], normal_stg1[1], init_type=init_type)
+        _ = multiview_color_projection(meshes, rgb_pils, resolution=512, device="cuda", complete_unseen=False, confidence_threshold=0.1)    # just check for validation, may throw error
+        vertices, faces, _ = from_py3d_mesh(meshes)
+        vertices, faces = reconstruct_stage1(normal_stg1, steps=200, vertices=vertices, faces=faces, start_edge_len=0.1, end_edge_len=0.02, gain=0.05, return_mesh=False, loss_expansion_weight=expansion_weight)
+    elif init_type in ["ball"]:
+        vertices, faces = reconstruct_stage1(normal_stg1, steps=200, end_edge_len=0.01, return_mesh=False, loss_expansion_weight=expansion_weight)
+    
+    normal_stg2 = [img.resize((1024, 1024)) for img in rm_normals]  # reduce computation on huggingface demo, use 1024 instead of 2048
+
+    vertices, faces = run_mesh_refine(vertices, faces, normal_stg2, steps=100, start_edge_len=0.02, end_edge_len=0.005, decay=0.99, update_normal_interval=20, update_warmup=5, return_mesh=False, process_inputs=False, process_outputs=False)
+
+    return vertices, faces, img_list
+
+# no GPU
+def geo_reconstruct_part2(vertices, faces):
+    meshes = simple_clean_mesh(to_pyml_mesh(vertices, faces), apply_smooth=True, stepsmoothnum=1, apply_sub_divide=True, sub_divide_threshold=0.25)
+    return meshes
+
+@spaces.GPU
+def geo_reconstruct_part3(meshes, img_list):
+    meshes = meshes.to("cuda")
     new_meshes = multiview_color_projection(meshes, img_list, resolution=1024, device="cuda", complete_unseen=True, confidence_threshold=0.2, cameras_list = get_cameras_list([0, 90, 180, 270], "cuda", focal=1))
     return new_meshes