Update app.py

app.py

@@ -22,7 +22,7 @@ from src.utils.camera_util import (
     FOV_to_intrinsics, 
     get_zero123plus_input_cameras,
     get_circular_camera_poses,
-)
+    )
 from src.utils.mesh_util import save_obj, save_glb
 from src.utils.infer_util import remove_background, resize_foreground, images_to_video
 
@@ -57,23 +57,23 @@ def get_render_cameras(batch_size=1, M=120, radius=2.5, elevation=10.0, is_flexi
         intrinsics = FOV_to_intrinsics(50.0).unsqueeze(0).repeat(M, 1, 1).float().flatten(-2)
         cameras = torch.cat([extrinsics, intrinsics], dim=-1)
         cameras = cameras.unsqueeze(0).repeat(batch_size, 1, 1)
-    return cameras
+        return cameras
 
 
-def images_to_video(images, output_path, fps=30):
+        def images_to_video(images, output_path, fps=30):
     # images: (N, C, H, W)
     os.makedirs(os.path.dirname(output_path), exist_ok=True)
     frames = []
     for i in range(images.shape[0]):
         frame = (images[i].permute(1, 2, 0).cpu().numpy() * 255).astype(np.uint8).clip(0, 255)
         assert frame.shape[0] == images.shape[2] and frame.shape[1] == images.shape[3], \
-            f"Frame shape mismatch: {frame.shape} vs {images.shape}"
+        f"Frame shape mismatch: {frame.shape} vs {images.shape}"
         assert frame.min() >= 0 and frame.max() <= 255, \
-            f"Frame value out of range: {frame.min()} ~ {frame.max()}"
+        f"Frame value out of range: {frame.min()} ~ {frame.max()}"
         frames.append(frame)
-    imageio.mimwrite(output_path, np.stack(frames), fps=fps, codec='h264')
+        imageio.mimwrite(output_path, np.stack(frames), fps=fps, codec='h264')
 
-def find_cuda():
+        def find_cuda():
     # Check if CUDA_HOME or CUDA_PATH environment variables are set
     cuda_home = os.environ.get('CUDA_HOME') or os.environ.get('CUDA_PATH')
 
@@ -88,24 +88,24 @@ def find_cuda():
         cuda_path = os.path.dirname(os.path.dirname(nvcc_path))
         return cuda_path
 
-    return None
+        return None
 
-cuda_path = find_cuda()
+        cuda_path = find_cuda()
 
-if cuda_path:
-    print(f"CUDA installation found at: {cuda_path}")
-else:
-    print("CUDA installation not found")
+        if cuda_path:
+            print(f"CUDA installation found at: {cuda_path}")
+        else:
+            print("CUDA installation not found")
 
-config_path = 'configs/instant-mesh-large.yaml'
-config = OmegaConf.load(config_path)
-config_name = os.path.basename(config_path).replace('.yaml', '')
-model_config = config.model_config
-infer_config = config.infer_config
+            config_path = 'configs/instant-mesh-large.yaml'
+            config = OmegaConf.load(config_path)
+            config_name = os.path.basename(config_path).replace('.yaml', '')
+            model_config = config.model_config
+            infer_config = config.infer_config
 
-IS_FLEXICUBES = True if config_name.startswith('instant-mesh') else False
+            IS_FLEXICUBES = True if config_name.startswith('instant-mesh') else False
 
-device = torch.device('cuda')
+            device = torch.device('cuda')
 
 # load diffusion model
 print('Loading diffusion model ...')
@@ -113,10 +113,10 @@ pipeline = DiffusionPipeline.from_pretrained(
     "sudo-ai/zero123plus-v1.2", 
     custom_pipeline="zero123plus",
     torch_dtype=torch.float16,
-)
+    )
 pipeline.scheduler = EulerAncestralDiscreteScheduler.from_config(
     pipeline.scheduler.config, timestep_spacing='trailing'
-)
+    )
 
 # load custom white-background UNet
 unet_ckpt_path = hf_hub_download(repo_id="TencentARC/InstantMesh", filename="diffusion_pytorch_model.bin", repo_type="model")
@@ -143,27 +143,27 @@ def check_input_image(input_image):
         raise gr.Error("No image uploaded!")
 
 
-def preprocess(input_image, do_remove_background):
+        def preprocess(input_image, do_remove_background):
 
-    rembg_session = rembg.new_session() if do_remove_background else None
+            rembg_session = rembg.new_session() if do_remove_background else None
 
-    if do_remove_background:
-        input_image = remove_background(input_image, rembg_session)
-        input_image = resize_foreground(input_image, 0.85)
+            if do_remove_background:
+                input_image = remove_background(input_image, rembg_session)
+                input_image = resize_foreground(input_image, 0.85)
 
-    return input_image
+                return input_image
 
 
-@spaces.GPU
-def generate_mvs(input_image, sample_steps, sample_seed):
+                @spaces.GPU
+                def generate_mvs(input_image, sample_steps, sample_seed):
 
-    seed_everything(sample_seed)
-    
+                    seed_everything(sample_seed)
+                    
     # sampling
     z123_image = pipeline(
         input_image, 
         num_inference_steps=sample_steps
-    ).images[0]
+        ).images[0]
 
     show_image = np.asarray(z123_image, dtype=np.uint8)
     show_image = torch.from_numpy(show_image)     # (960, 640, 3)
@@ -174,15 +174,15 @@ def generate_mvs(input_image, sample_steps, sample_seed):
     return z123_image, show_image
 
 
-@spaces.GPU
-def make3d(images):
+    @spaces.GPU
+    def make3d(images):
 
-    global model
-    if IS_FLEXICUBES:
-        model.init_flexicubes_geometry(device, use_renderer=False)
-    model = model.eval()
+        global model
+        if IS_FLEXICUBES:
+            model.init_flexicubes_geometry(device, use_renderer=False)
+            model = model.eval()
 
-    images = np.asarray(images, dtype=np.float32) / 255.0
+            images = np.asarray(images, dtype=np.float32) / 255.0
     images = torch.from_numpy(images).permute(2, 0, 1).contiguous().float()     # (3, 960, 640)
     images = rearrange(images, 'c (n h) (m w) -> (n m) c h w', n=3, m=2)        # (6, 3, 320, 320)
 
@@ -208,7 +208,7 @@ def make3d(images):
             planes,
             use_texture_map=False,
             **infer_config,
-        )
+            )
 
         vertices, faces, vertex_colors = mesh_out
         vertices = vertices[:, [1, 2, 0]]
@@ -218,7 +218,7 @@ def make3d(images):
         
         print(f"Mesh saved to {mesh_fpath}")
 
-    return mesh_fpath, mesh_glb_fpath
+        return mesh_fpath, mesh_glb_fpath
 
 
 ###############################################################################
@@ -228,14 +228,14 @@ model = load_v2()
 device = torch.device('cuda')
 accelerator = Accelerator(
     mixed_precision="fp16",
-)
+    )
 model = accelerator.prepare(model)
 model.eval()
 print("Model loaded to device")
 
 def wireframe_render(mesh):
     views = [
-        (90, 20), (270, 20)
+    (90, 20), (270, 20)
     ]
     mesh.vertices = mesh.vertices[:, [0, 2, 1]]
 
@@ -260,7 +260,7 @@ def wireframe_render(mesh):
             facecolors=(0.8, 0.5, 0.2, 1.0),  # Brownish yellow
             edgecolors='k',
             linewidths=0.5,
-        ))
+            ))
 
         # Set limits and center the view on the object
         ax.set_xlim(center[0] - scale / 2, center[0] + scale / 2)
@@ -300,13 +300,13 @@ def wireframe_render(mesh):
     plt.close(fig)
     return save_path
 
-@spaces.GPU(duration=360)
-def do_inference(input_3d, sample_seed=0, do_sampling=False, do_marching_cubes=False):
-    set_seed(sample_seed)
-    print("Seed value:", sample_seed)
+    @spaces.GPU(duration=360)
+    def do_inference(input_3d, sample_seed=0, do_sampling=False, do_marching_cubes=False):
+        set_seed(sample_seed)
+        print("Seed value:", sample_seed)
 
-    input_mesh = trimesh.load(input_3d)
-    pc_list, mesh_list = process_mesh_to_pc([input_mesh], marching_cubes = do_marching_cubes)
+        input_mesh = trimesh.load(input_3d)
+        pc_list, mesh_list = process_mesh_to_pc([input_mesh], marching_cubes = do_marching_cubes)
     pc_normal = pc_list[0] # 4096, 6
     mesh = mesh_list[0]
     vertices = mesh.vertices
@@ -330,16 +330,16 @@ def do_inference(input_3d, sample_seed=0, do_sampling=False, do_marching_cubes=F
     try:
         if mesh.visual.vertex_colors is not None:
             orange_color = np.array([255, 165, 0, 255], dtype=np.uint8)
-    
+            
             mesh.visual.vertex_colors = np.tile(orange_color, (mesh.vertices.shape[0], 1))
         else:
             orange_color = np.array([255, 165, 0, 255], dtype=np.uint8)
             mesh.visual.vertex_colors = np.tile(orange_color, (mesh.vertices.shape[0], 1))
-    except Exception as e:
-        print(e)
-    input_save_name = f"processed_input_{int(time.time())}.obj"
-    mesh.export(input_save_name)
-    input_render_res = wireframe_render(mesh)
+        except Exception as e:
+            print(e)
+            input_save_name = f"processed_input_{int(time.time())}.obj"
+            mesh.export(input_save_name)
+            input_render_res = wireframe_render(mesh)
 
     pc_coor = pc_coor / np.abs(pc_coor).max() * 0.99 # input should be from -1 to 1
 
@@ -352,17 +352,17 @@ def do_inference(input_3d, sample_seed=0, do_sampling=False, do_marching_cubes=F
     # with accelerator.autocast():
     with accelerator.autocast():
         outputs = model(input, do_sampling)
-    print("Model inference done")
-    recon_mesh = outputs[0]
+        print("Model inference done")
+        recon_mesh = outputs[0]
 
-    valid_mask = torch.all(~torch.isnan(recon_mesh.reshape((-1, 9))), dim=1)
+        valid_mask = torch.all(~torch.isnan(recon_mesh.reshape((-1, 9))), dim=1)
     recon_mesh = recon_mesh[valid_mask]  # nvalid_face x 3 x 3
     vertices = recon_mesh.reshape(-1, 3).cpu()
     vertices_index = np.arange(len(vertices))  # 0, 1, ..., 3 x face
     triangles = vertices_index.reshape(-1, 3)
 
     artist_mesh = trimesh.Trimesh(vertices=vertices, faces=triangles, force="mesh",
-                                 merge_primitives=True)
+       merge_primitives=True)
 
     artist_mesh.merge_vertices()
     artist_mesh.update_faces(artist_mesh.nondegenerate_faces())
@@ -378,12 +378,12 @@ def do_inference(input_3d, sample_seed=0, do_sampling=False, do_marching_cubes=F
         orange_color = np.array([255, 165, 0, 255], dtype=np.uint8)
         artist_mesh.visual.vertex_colors = np.tile(orange_color, (artist_mesh.vertices.shape[0], 1))
 
-    num_faces = len(artist_mesh.faces)
+        num_faces = len(artist_mesh.faces)
 
-    brown_color = np.array([165, 42, 42, 255], dtype=np.uint8)
-    face_colors = np.tile(brown_color, (num_faces, 1))
+        brown_color = np.array([165, 42, 42, 255], dtype=np.uint8)
+        face_colors = np.tile(brown_color, (num_faces, 1))
 
-    artist_mesh.visual.face_colors = face_colors
+        artist_mesh.visual.face_colors = face_colors
     # add time stamp to avoid cache
     save_name = f"output_{int(time.time())}.obj"
     artist_mesh.export(save_name)
@@ -395,18 +395,18 @@ output_model_obj = gr.Model3D(
     label="Generated Mesh (OBJ Format)",
     display_mode="wireframe",
     clear_color=[1, 1, 1, 1],
-)
+    )
 preprocess_model_obj = gr.Model3D(
     label="Processed Input Mesh (OBJ Format)",
     display_mode="wireframe",
     clear_color=[1, 1, 1, 1],
-)
+    )
 input_image_render = gr.Image(
     label="Wireframe Render of Processed Input Mesh",
-)
+    )
 output_image_render = gr.Image(
     label="Wireframe Render of Generated Mesh",
-)
+    )
 
 ###############################################################################
 # Gradio
@@ -454,137 +454,137 @@ STEP4_HEADER = """
 """
 
 with gr.Blocks() as demo:
-	gr.Markdown(HEADER)
-	gr.Markdown(STEP1_HEADER)
-	with gr.Row(variant = "panel"):
-		with gr.Column():
-			with gr.Row():
-				input_image = gr.Image(
-					label = "Input Image",
-					image_mode = "RGBA",
-					sources = "upload",
-					type="pil",
-					elem_id="content_image"
-				)
-				processed_image = gr.Image(label="Processed Image",
-					image_mode="RGBA",
-					type="pil",
-					interactive=False
-				)
-			with gr.Row():
-				with gr.Group():
-					do_remove_background = gr.Checkbox(
-						label="Remove Background",
-						value=True)
-					sample_seed = gr.Number(
-						value=42, 
-						label="Seed Value", 
-						precision=0
-					)
-					sample_steps = gr.Slider(
-                        label="Sample Steps",
-                        minimum=30,
-                        maximum=75,
-                        value=75,
-                        step=5
+    gr.Markdown(HEADER)
+    gr.Markdown(STEP1_HEADER)
+    with gr.Row(variant = "panel"):
+        with gr.Column():
+            with gr.Row():
+                input_image = gr.Image(
+                    label = "Input Image",
+                    image_mode = "RGBA",
+                    sources = "upload",
+                    type="pil",
+                    elem_id="content_image"
                     )
-			with gr.Row():
-				step1_submit = gr.Button("Generate", elem_id="generate", variant="primary")
-		with gr.Column():
-			with gr.Row():
-				with gr.Column():
-					mv_show_images = gr.Image(
-                        label="Generated Multi-views",
-                        type="pil",
-                        width=379,
-                        interactive=False
+                processed_image = gr.Image(label="Processed Image",
+                    image_mode="RGBA",
+                    type="pil",
+                    interactive=False
                     )
-		with gr.Column():
-			with gr.Tab("OBJ"):
-				output_model_obj = gr.Model3D(
-					label = "Output Model (OBJ Format)",
-					interactive = False,
-				)
-				gr.Markdown("Note: Downloaded object will be flipped in case of .obj export. Export .glb instead or manually flip it before usage.")
-			with gr.Tab("GLB"):
-				output_model_glb = gr.Model3D(
-                    label="Output Model (GLB Format)",
-                    interactive=False,
-				)
-				gr.Markdown("Note: The model shown here has a darker appearance. Download to get correct results.")
-				gr.Markdown('''Try a different <b>seed value</b> if the result is unsatisfying (Default: 42).''')
-	with gr.Row():
-        gr.Markdown(STEP2_HEADER)
-	with gr.Row(variant="panel"):
-		with gr.Column():
-			with gr.Row():
-				input_3d = gr.Model3D(
-                    label="Input Mesh",
-                    display_mode="wireframe",
-                    clear_color=[1,1,1,1],
-                )
-
-		with gr.Row():
-			with gr.Group():
-				do_marching_cubes = gr.Checkbox(label="Preprocess with Marching Cubes", value=False)
-				do_sampling = gr.Checkbox(label="Random Sampling", value=False)
-				sample_seed = gr.Number(value=0, label="Seed Value", precision=0)
-
-			with gr.Row():
-				step2_submit = gr.Button("Generate", elem_id="generate", variant="primary")
-
-			with gr.Row(variant="panel"):
-				mesh_examples = gr.Examples(
-                    examples=[
-                        os.path.join("examples", img_name) for img_name in sorted(os.listdir("examples"))
-                    ],
-                    inputs=input_3d,
-                    outputs=[preprocess_model_obj, input_image_render, output_model_obj, output_image_render],
-                    fn=do_inference,
-                    cache_examples = False,
-                    examples_per_page=10
-                )
-
-		with gr.Column():
-			with gr.Row():
-				input_image_render.render()
-			with gr.Row():
-				with gr.Tab("OBJ"):
-					preprocess_model_obj.render()
-			with gr.Row():
-				output_image_render.render()
-			with gr.Row():
-				with gr.Tab("OBJ"):
-					output_model_obj.render()
-			with gr.Row():
-				gr.Markdown('''Try click random sampling and different <b>Seed Value</b> if the result is unsatisfying''')
-
-	gr.Markdown(STEP3_HEADER)
-	gr.Markdown(STEP4_HEADER)
-
-	mv_images = gr.State()
-
-	step1_submit.click(fn=check_input_image, inputs=[input_image]).success(
-        fn=preprocess,
-        inputs=[input_image, do_remove_background],
-        outputs=[processed_image],
-	).success(
-        fn=generate_mvs,
-        inputs=[processed_image, sample_steps, sample_seed],
-        outputs=[mv_images, mv_show_images],
-	).success(
-        fn=make3d,
-        inputs=[mv_images],
-        outputs=[output_model_obj, output_model_glb]
-	)
-
-	step2_submit.click(
-        fn=do_inference,
-        inputs=[input_3d, sample_seed, do_sampling, do_marching_cubes],
-        outputs=[preprocess_model_obj, input_image_render, output_model_obj, output_image_render],
-	)
-
-
-
-demo.queue(max_size=10)
-demo.launch()
+                with gr.Row():
+                    with gr.Group():
+                        do_remove_background = gr.Checkbox(
+                            label="Remove Background",
+                            value=True)
+                        sample_seed = gr.Number(
+                            value=42, 
+                            label="Seed Value", 
+                            precision=0
+                            )
+                        sample_steps = gr.Slider(
+                            label="Sample Steps",
+                            minimum=30,
+                            maximum=75,
+                            value=75,
+                            step=5
+                            )
+                        with gr.Row():
+                            step1_submit = gr.Button("Generate", elem_id="generate", variant="primary")
+                            with gr.Column():
+                                with gr.Row():
+                                    with gr.Column():
+                                        mv_show_images = gr.Image(
+                                            label="Generated Multi-views",
+                                            type="pil",
+                                            width=379,
+                                            interactive=False
+                                            )
+                                        with gr.Column():
+                                            with gr.Tab("OBJ"):
+                                                output_model_obj = gr.Model3D(
+                                                    label = "Output Model (OBJ Format)",
+                                                    interactive = False,
+                                                    )
+                                                gr.Markdown("Note: Downloaded object will be flipped in case of .obj export. Export .glb instead or manually flip it before usage.")
+                                                with gr.Tab("GLB"):
+                                                    output_model_glb = gr.Model3D(
+                                                        label="Output Model (GLB Format)",
+                                                        interactive=False,
+                                                        )
+                                                    gr.Markdown("Note: The model shown here has a darker appearance. Download to get correct results.")
+                                                    gr.Markdown('''Try a different <b>seed value</b> if the result is unsatisfying (Default: 42).''')
+                                                    with gr.Row():
+                                                        gr.Markdown(STEP2_HEADER)
+                                                        with gr.Row(variant="panel"):
+                                                            with gr.Column():
+                                                                with gr.Row():
+                                                                    input_3d = gr.Model3D(
+                                                                        label="Input Mesh",
+                                                                        display_mode="wireframe",
+                                                                        clear_color=[1,1,1,1],
+                                                                        )
+
+                                                                    with gr.Row():
+                                                                        with gr.Group():
+                                                                            do_marching_cubes = gr.Checkbox(label="Preprocess with Marching Cubes", value=False)
+                                                                            do_sampling = gr.Checkbox(label="Random Sampling", value=False)
+                                                                            sample_seed = gr.Number(value=0, label="Seed Value", precision=0)
+
+                                                                            with gr.Row():
+                                                                                step2_submit = gr.Button("Generate", elem_id="generate", variant="primary")
+
+                                                                                with gr.Row(variant="panel"):
+                                                                                    mesh_examples = gr.Examples(
+                                                                                        examples=[
+                                                                                        os.path.join("examples", img_name) for img_name in sorted(os.listdir("examples"))
+                                                                                        ],
+                                                                                        inputs=input_3d,
+                                                                                        outputs=[preprocess_model_obj, input_image_render, output_model_obj, output_image_render],
+                                                                                        fn=do_inference,
+                                                                                        cache_examples = False,
+                                                                                        examples_per_page=10
+                                                                                        )
+
+                                                                                    with gr.Column():
+                                                                                        with gr.Row():
+                                                                                            input_image_render.render()
+                                                                                            with gr.Row():
+                                                                                                with gr.Tab("OBJ"):
+                                                                                                    preprocess_model_obj.render()
+                                                                                                    with gr.Row():
+                                                                                                        output_image_render.render()
+                                                                                                        with gr.Row():
+                                                                                                            with gr.Tab("OBJ"):
+                                                                                                                output_model_obj.render()
+                                                                                                                with gr.Row():
+                                                                                                                    gr.Markdown('''Try click random sampling and different <b>Seed Value</b> if the result is unsatisfying''')
+
+                                                                                                                    gr.Markdown(STEP3_HEADER)
+                                                                                                                    gr.Markdown(STEP4_HEADER)
+
+                                                                                                                    mv_images = gr.State()
+
+                                                                                                                    step1_submit.click(fn=check_input_image, inputs=[input_image]).success(
+                                                                                                                        fn=preprocess,
+                                                                                                                        inputs=[input_image, do_remove_background],
+                                                                                                                        outputs=[processed_image],
+                                                                                                                        ).success(
+                                                                                                                        fn=generate_mvs,
+                                                                                                                        inputs=[processed_image, sample_steps, sample_seed],
+                                                                                                                        outputs=[mv_images, mv_show_images],
+                                                                                                                        ).success(
+                                                                                                                        fn=make3d,
+                                                                                                                        inputs=[mv_images],
+                                                                                                                        outputs=[output_model_obj, output_model_glb]
+                                                                                                                        )
+
+                                                                                                                        step2_submit.click(
+                                                                                                                            fn=do_inference,
+                                                                                                                            inputs=[input_3d, sample_seed, do_sampling, do_marching_cubes],
+                                                                                                                            outputs=[preprocess_model_obj, input_image_render, output_model_obj, output_image_render],
+                                                                                                                            )
+
+
+
+                                                                                                                        demo.queue(max_size=10)
+                                                                                                                        demo.launch()