Upload app.py

app.py

@@ -94,113 +94,118 @@ def segment_and_overlay(image, points, sam_predictor):
     return overlaid, visible_mask
 
 
-# def image_to_3d(images,masks,seed,ss_guidance_strength,ss_sampling_steps,slat_guidance_strength,slat_sampling_steps):
-#     """
-#     将图像转换为 3D 模型。
-#     """
-#     print(images.shape, masks.shape)
-#     print(type(images), type(masks))
-#     print(seed)
-#     print(ss_guidance_strength, ss_sampling_steps, slat_guidance_strength, slat_sampling_steps)
-#     # user_dir = os.path.join(TMP_DIR, str(req.session_hash))
-#     test_tensor = torch.zeros((1, 3, 64, 64))
-#     print(test_tensor.shape, test_tensor.device)
-#     # outputs = pipeline.run_multi_image(
-#     #     [images],
-#     #     [masks],
-#     #     seed=seed,
-#     #     formats=["gaussian", "mesh"],
-#     #     preprocess_image=False,
-#     #     sparse_structure_sampler_params={
-#     #         "steps": ss_sampling_steps,
-#     #         "cfg_strength": ss_guidance_strength,
-#     #     },
-#     #     slat_sampler_params={
-#     #         "steps": slat_sampling_steps,
-#     #         "cfg_strength": slat_guidance_strength,
-#     #     },
-#     #     mode="stochastic",
-#     # )
-#     # video = render_utils.render_video(outputs['gaussian'][0], num_frames=120, bg_color=(1, 1, 1))['color']
-#     # video_geo = render_utils.render_video(outputs['mesh'][0], num_frames=120, bg_color=(1, 1, 1))['normal']
-#     # video = [np.concatenate([video[i], video_geo[i]], axis=1) for i in range(len(video))]
-#     # video_path = os.path.join(user_dir, 'sample.mp4')
-#     # imageio.mimsave(video_path, video, fps=15)
-#     # state = pack_state(outputs['gaussian'][0], outputs['mesh'][0])
-#     # torch.cuda.empty_cache()
-#     state = None
-#     video_path = None
-#     return state, video_path
-
-# @spaces.GPU(duration=90)
-# def extract_glb(
-#     state: dict,
-#     mesh_simplify: float,
-#     texture_size: int,
-#     req: gr.Request,
-# ) -> tuple:
-#     """
-#     从生成的 3D 模型中提取 GLB 文件。
-#     """
-#     user_dir = os.path.join(TMP_DIR, str(req.session_hash))
-#     gs, mesh = unpack_state(state)
-#     glb = postprocessing_utils.to_glb(gs, mesh, simplify=mesh_simplify, texture_size=texture_size, verbose=False)
-#     glb_path = os.path.join(user_dir, 'sample.glb')
-#     glb.export(glb_path)
-#     torch.cuda.empty_cache()
-#     return glb_path, glb_path
-
-# @spaces.GPU
-# def extract_gaussian(state: dict, req: gr.Request) -> tuple:
-#     """
-#     从生成的 3D 模型中提取 Gaussian 文件。
-#     """
-#     user_dir = os.path.join(TMP_DIR, str(req.session_hash))
-#     gs, _ = unpack_state(state)
-#     gaussian_path = os.path.join(user_dir, 'sample.ply')
-#     gs.save_ply(gaussian_path)
-#     torch.cuda.empty_cache()
-#     return gaussian_path, gaussian_path
-
-
-# def pack_state(gs: Gaussian, mesh: MeshExtractResult) -> dict:
-#     return {
-#         'gaussian': {
-#             **gs.init_params,
-#             '_xyz': gs._xyz.cpu().numpy(),
-#             '_features_dc': gs._features_dc.cpu().numpy(),
-#             '_scaling': gs._scaling.cpu().numpy(),
-#             '_rotation': gs._rotation.cpu().numpy(),
-#             '_opacity': gs._opacity.cpu().numpy(),
-#         },
-#         'mesh': {
-#             'vertices': mesh.vertices.cpu().numpy(),
-#             'faces': mesh.faces.cpu().numpy(),
-#         },
-#     }
+@spaces.GPU
+def image_to_3d(
+    image: List[tuple],
+    masks: List[np.ndarray],
+    seed: int,
+    ss_guidance_strength: float,
+    ss_sampling_steps: int,
+    slat_guidance_strength: float,
+    slat_sampling_steps: int,
+    multiimage_algo: str,
+    req: gr.Request,
+) -> tuple:
+    """
+    将图像转换为 3D 模型。
+    """
+    user_dir = os.path.join(TMP_DIR, str(req.session_hash))
+    outputs = pipeline.run_multi_image(
+        [img[0] for img in image],
+        [mask[0] for mask in masks],
+        seed=seed,
+        formats=["gaussian", "mesh"],
+        preprocess_image=False,
+        sparse_structure_sampler_params={
+            "steps": ss_sampling_steps,
+            "cfg_strength": ss_guidance_strength,
+        },
+        slat_sampler_params={
+            "steps": slat_sampling_steps,
+            "cfg_strength": slat_guidance_strength,
+        },
+        mode=multiimage_algo,
+    )
+    video = render_utils.render_video(outputs['gaussian'][0], num_frames=120)['color']
+    video_geo = render_utils.render_video(outputs['mesh'][0], num_frames=120)['normal']
+    video = [np.concatenate([video[i], video_geo[i]], axis=1) for i in range(len(video))]
+    video_path = os.path.join(user_dir, 'sample.mp4')
+    imageio.mimsave(video_path, video, fps=15)
+    state = pack_state(outputs['gaussian'][0], outputs['mesh'][0])
+    torch.cuda.empty_cache()
+    return state, video_path
+
+
+@spaces.GPU(duration=90)
+def extract_glb(
+    state: dict,
+    mesh_simplify: float,
+    texture_size: int,
+    req: gr.Request,
+) -> tuple:
+    """
+    从生成的 3D 模型中提取 GLB 文件。
+    """
+    user_dir = os.path.join(TMP_DIR, str(req.session_hash))
+    gs, mesh = unpack_state(state)
+    glb = postprocessing_utils.to_glb(gs, mesh, simplify=mesh_simplify, texture_size=texture_size, verbose=False)
+    glb_path = os.path.join(user_dir, 'sample.glb')
+    glb.export(glb_path)
+    torch.cuda.empty_cache()
+    return glb_path, glb_path
+
+
+@spaces.GPU
+def extract_gaussian(state: dict, req: gr.Request) -> tuple:
+    """
+    从生成的 3D 模型中提取 Gaussian 文件。
+    """
+    user_dir = os.path.join(TMP_DIR, str(req.session_hash))
+    gs, _ = unpack_state(state)
+    gaussian_path = os.path.join(user_dir, 'sample.ply')
+    gs.save_ply(gaussian_path)
+    torch.cuda.empty_cache()
+    return gaussian_path, gaussian_path
+
+
+def pack_state(gs: Gaussian, mesh: MeshExtractResult) -> dict:
+    return {
+        'gaussian': {
+            **gs.init_params,
+            '_xyz': gs._xyz.cpu().numpy(),
+            '_features_dc': gs._features_dc.cpu().numpy(),
+            '_scaling': gs._scaling.cpu().numpy(),
+            '_rotation': gs._rotation.cpu().numpy(),
+            '_opacity': gs._opacity.cpu().numpy(),
+        },
+        'mesh': {
+            'vertices': mesh.vertices.cpu().numpy(),
+            'faces': mesh.faces.cpu().numpy(),
+        },
+    }
     
     
-# def unpack_state(state: dict) -> tuple:
-#     gs = Gaussian(
-#         aabb=state['gaussian']['aabb'],
-#         sh_degree=state['gaussian']['sh_degree'],
-#         mininum_kernel_size=state['gaussian']['mininum_kernel_size'],
-#         scaling_bias=state['gaussian']['scaling_bias'],
-#         opacity_bias=state['gaussian']['opacity_bias'],
-#         scaling_activation=state['gaussian']['scaling_activation'],
-#     )
-#     gs._xyz = torch.tensor(state['gaussian']['_xyz'], device='cuda')
-#     gs._features_dc = torch.tensor(state['gaussian']['_features_dc'], device='cuda')
-#     gs._scaling = torch.tensor(state['gaussian']['_scaling'], device='cuda')
-#     gs._rotation = torch.tensor(state['gaussian']['_rotation'], device='cuda')
-#     gs._opacity = torch.tensor(state['gaussian']['_opacity'], device='cuda')
+def unpack_state(state: dict) -> tuple:
+    gs = Gaussian(
+        aabb=state['gaussian']['aabb'],
+        sh_degree=state['gaussian']['sh_degree'],
+        mininum_kernel_size=state['gaussian']['mininum_kernel_size'],
+        scaling_bias=state['gaussian']['scaling_bias'],
+        opacity_bias=state['gaussian']['opacity_bias'],
+        scaling_activation=state['gaussian']['scaling_activation'],
+    )
+    gs._xyz = torch.tensor(state['gaussian']['_xyz'], device='cuda')
+    gs._features_dc = torch.tensor(state['gaussian']['_features_dc'], device='cuda')
+    gs._scaling = torch.tensor(state['gaussian']['_scaling'], device='cuda')
+    gs._rotation = torch.tensor(state['gaussian']['_rotation'], device='cuda')
+    gs._opacity = torch.tensor(state['gaussian']['_opacity'], device='cuda')
     
-#     mesh = edict(
-#         vertices=torch.tensor(state['mesh']['vertices'], device='cuda'),
-#         faces=torch.tensor(state['mesh']['faces'], device='cuda'),
-#     )
+    mesh = edict(
+        vertices=torch.tensor(state['mesh']['vertices'], device='cuda'),
+        faces=torch.tensor(state['mesh']['faces'], device='cuda'),
+    )
     
-#     return gs, mesh
+    return gs, mesh
 
 def get_sam_predictor():
     sam_checkpoint = hf_hub_download("ybelkada/segment-anything", "checkpoints/sam_vit_h_4b8939.pth")
@@ -425,22 +430,22 @@ with gr.Blocks(delete_cache=(600, 600)) as demo:
         * Different random seeds can be tried in "Generation Settings", if you think the results are not ideal.
         * If the reconstruction 3D asset is satisfactory, you can extract the GLB file and download it.
         """)
-    with gr.Row():
-        with gr.Column():
-            with gr.Accordion(label="Generation Settings", open=True):
-                seed = gr.Slider(0, MAX_SEED, label="Seed", value=1, step=1)
-                randomize_seed = gr.Checkbox(label="Randomize Seed", value=True)
-                gr.Markdown("Stage 1: Sparse Structure Generation")
-                with gr.Row():
-                    ss_guidance_strength = gr.Slider(0.0, 10.0, label="Guidance Strength", value=7.5, step=0.1)
-                    ss_sampling_steps = gr.Slider(1, 50, label="Sampling Steps", value=12, step=1)
-                gr.Markdown("Stage 2: Structured Latent Generation")
-                with gr.Row():
-                    slat_guidance_strength = gr.Slider(0.0, 10.0, label="Guidance Strength", value=3.0, step=0.1)
-                    slat_sampling_steps = gr.Slider(1, 50, label="Sampling Steps", value=12, step=1)
-            generate_btn = gr.Button("Generate")
-        with gr.Column():
-            video_output = gr.Video(label="Generated 3D Asset", autoplay=True, loop=True, height=300)
+    # with gr.Row():
+    #     with gr.Column():
+    #         with gr.Accordion(label="Generation Settings", open=True):
+    #             seed = gr.Slider(0, MAX_SEED, label="Seed", value=1, step=1)
+    #             randomize_seed = gr.Checkbox(label="Randomize Seed", value=True)
+    #             gr.Markdown("Stage 1: Sparse Structure Generation")
+    #             with gr.Row():
+    #                 ss_guidance_strength = gr.Slider(0.0, 10.0, label="Guidance Strength", value=7.5, step=0.1)
+    #                 ss_sampling_steps = gr.Slider(1, 50, label="Sampling Steps", value=12, step=1)
+    #             gr.Markdown("Stage 2: Structured Latent Generation")
+    #             with gr.Row():
+    #                 slat_guidance_strength = gr.Slider(0.0, 10.0, label="Guidance Strength", value=3.0, step=0.1)
+    #                 slat_sampling_steps = gr.Slider(1, 50, label="Sampling Steps", value=12, step=1)
+    #         generate_btn = gr.Button("Generate")
+    #     with gr.Column():
+    #         video_output = gr.Video(label="Generated 3D Asset", autoplay=True, loop=True, height=300)
     
     # # Handlers
     # demo.load(start_session)
@@ -540,6 +545,7 @@ with gr.Blocks(delete_cache=(600, 600)) as demo:
 # 启动 Gradio App
 if __name__ == "__main__":
     pipeline = Amodal3RImageTo3DPipeline.from_pretrained("Sm0kyWu/Amodal3R")
+    pipeline.cuda()
     try:
         pipeline.preprocess_image(Image.fromarray(np.zeros((512, 512, 3), dtype=np.uint8)))
     except: