Working ZeroGPU version (#6)

- Working ZeroGPU version (9ae28d8d321fe12e90eecf0e5506325c245274e3)
- Update freesplatter/webui/runner.py (0d0baebb992391cede34eb2ad8ffed970c5697ef)
- Update freesplatter/webui/tab_img_to_3d.py (0dd1fb37cf6a2c4f916410a260f27ee22b6cf101)
- Create open3d_zerogpu_fix.py (04a21477f8ecbd98eb713d7e1b4b35ca7e9da1bc)
- Update app.py (88c07bd9d42992b9f0d2d820b61494306c3fe465)


Co-authored-by: Charles Bensimon <cbensimon@users.noreply.huggingface.co>

app.py

@@ -1,6 +1,8 @@
+import open3d_zerogpu_fix
 import os
 if 'OMP_NUM_THREADS' not in os.environ:
     os.environ['OMP_NUM_THREADS'] = '16'
+import spaces
 import torch
 import subprocess
 import gradio as gr
@@ -11,30 +13,16 @@ from freesplatter.webui.runner import FreeSplatterRunner
 from freesplatter.webui.tab_img_to_3d import create_interface_img_to_3d
 
 
-def install_cuda_toolkit():
-    CUDA_TOOLKIT_URL = "https://developer.download.nvidia.com/compute/cuda/12.1.0/local_installers/cuda_12.1.0_530.30.02_linux.run"
-    CUDA_TOOLKIT_FILE = "/tmp/%s" % os.path.basename(CUDA_TOOLKIT_URL)
-    subprocess.call(["wget", "-q", CUDA_TOOLKIT_URL, "-O", CUDA_TOOLKIT_FILE])
-    subprocess.call(["chmod", "+x", CUDA_TOOLKIT_FILE])
-    subprocess.call([CUDA_TOOLKIT_FILE, "--silent", "--toolkit"])
-
-    os.environ["CUDA_HOME"] = "/usr/local/cuda"
-    os.environ["PATH"] = "%s/bin:%s" % (os.environ["CUDA_HOME"], os.environ["PATH"])
-    os.environ["LD_LIBRARY_PATH"] = "%s/lib:%s" % (
-        os.environ["CUDA_HOME"],
-        "" if "LD_LIBRARY_PATH" not in os.environ else os.environ["LD_LIBRARY_PATH"],
-    )
-    # Fix: arch_list[-1] += '+PTX'; IndexError: list index out of range
-    os.environ["TORCH_CUDA_ARCH_LIST"] = "8.0;8.6"
-
-install_cuda_toolkit()
-
-
 torch.set_grad_enabled(False)
 device = torch.device('cuda')
 runner = FreeSplatterRunner(device)
 
 
+@spaces.GPU(duration=120)
+def run_img_to_3d(*args):
+    yield from runner.run_img_to_3d(*args, cache_dir=gr.utils.get_upload_folder())
+
+
 _HEADER_ = '''
 # FreeSplatter 🤗 Gradio Demo
 \n\nOfficial demo of the paper [FreeSplatter: Pose-free Gaussian Splatting for Sparse-view 3D Reconstruction](https://arxiv.org/abs/2404.07191). [[Github]](https://github.com/TencentARC/FreeSplatter)  
@@ -82,18 +70,15 @@ with gr.Blocks(analytics_enabled=False, title='FreeSplatter Demo') as demo:
             with gr.Tabs() as sub_tabs_img_to_3d:
                 with gr.TabItem('Hunyuan3D Std', id='tab_hunyuan3d_std'):
                     _, var_img_to_3d_hunyuan3d_std = create_interface_img_to_3d(
-                        runner.run_segmentation,
-                        runner.run_img_to_3d, 
+                        run_img_to_3d, 
                         model='Hunyuan3D Std')
                 with gr.TabItem('Zero123++ v1.1', id='tab_zero123plus_v11'):
                     _, var_img_to_3d_zero123plus_v11 = create_interface_img_to_3d(
-                        runner.run_segmentation,
-                        runner.run_img_to_3d, 
+                        run_img_to_3d, 
                         model='Zero123++ v1.1')
                 with gr.TabItem('Zero123++ v1.2', id='tab_zero123plus_v12'):
                     _, var_img_to_3d_zero123plus_v12 = create_interface_img_to_3d(
-                        runner.run_segmentation,
-                        runner.run_img_to_3d, 
+                        run_img_to_3d, 
                         model='Zero123++ v1.2')
 
     gr.Markdown(_CITE_)

freesplatter/webui/runner.py

@@ -1,4 +1,3 @@
-import spaces
 import os
 import json
 import uuid
@@ -157,17 +156,14 @@ class FreeSplatterRunner:
         image, 
         do_rembg=True,
     ):
-        torch.cuda.empty_cache()
-
         if do_rembg:
             image = remove_background(image, self.rembg)
 
         return image
 
-    @spaces.GPU
     def run_img_to_3d(
         self, 
-        image_rgba, 
+        image,
         model='Zero123++ v1.2', 
         diffusion_steps=30, 
         guidance_scale=4.0,
@@ -177,7 +173,10 @@ class FreeSplatterRunner:
         mesh_reduction=0.5,
         cache_dir=None,
     ):
-        torch.cuda.empty_cache()
+        image_rgba = self.run_segmentation(image)
+
+        res = [image_rgba]
+        yield res + [None] * (6 - len(res))
 
         self.output_dir = os.path.join(cache_dir, f'output_{uuid.uuid4()}')
         os.makedirs(self.output_dir, exist_ok=True)
@@ -226,6 +225,10 @@ class FreeSplatterRunner:
             images = images[[0, 2, 4, 5, 3, 1]]
             alphas = alphas[[0, 2, 4, 5, 3, 1]]
         images_vis = v2.functional.to_pil_image(rearrange(images, 'nm c h w -> c h (nm w)'))
+
+        res += [images_vis]
+        yield res + [None] * (6 - len(res))
+
         images = v2.functional.resize(images, 512, interpolation=3, antialias=True).clamp(0, 1)
         alphas = v2.functional.resize(alphas, 512, interpolation=0, antialias=True).clamp(0, 1)
 
@@ -237,12 +240,12 @@ class FreeSplatterRunner:
         images, alphas = images[view_indices], alphas[view_indices]
         legends = [f'V{i}' if i != 0 else 'Input' for i in view_indices]
 
-        gs_vis_path, video_path, mesh_fine_path, fig = self.run_freesplatter_object(
-            images, alphas, legends=legends, gs_type=gs_type, mesh_reduction=mesh_reduction)
+        for item in self.run_freesplatter_object(
+            images, alphas, legends=legends, gs_type=gs_type, mesh_reduction=mesh_reduction):
+            res += [item]
+            yield res + [None] * (6 - len(res))
 
-        return images_vis, gs_vis_path, video_path, mesh_fine_path, fig
 
-    @spaces.GPU
     def run_views_to_3d(
         self, 
         image_files, 
@@ -251,7 +254,6 @@ class FreeSplatterRunner:
         mesh_reduction=0.5,
         cache_dir=None,
     ):
-        torch.cuda.empty_cache()
 
         self.output_dir = os.path.join(cache_dir, f'output_{uuid.uuid4()}')
         os.makedirs(self.output_dir, exist_ok=True)
@@ -300,7 +302,6 @@ class FreeSplatterRunner:
         gs_type='2DGS', 
         mesh_reduction=0.5,
     ):
-        torch.cuda.empty_cache()
         device = self.device
 
         freesplatter = self.freesplatter_2dgs if gs_type == '2DGS' else self.freesplatter
@@ -316,11 +317,13 @@ class FreeSplatterRunner:
         c2ws_pred, focals_pred = freesplatter.estimate_poses(images, gaussians, masks=alphas, use_first_focal=True, pnp_iter=10)
         fig = self.visualize_cameras_object(images, c2ws_pred, focals_pred, legends=legends)
         t2 = time.time()
+        yield fig
         
         # save gaussians
         gs_vis_path = os.path.join(self.output_dir, 'gs_vis.ply')
         save_gaussian(gaussians, gs_vis_path, freesplatter, opacity_threshold=5e-3, pad_2dgs_scale=True)
         print(f'Save gaussian at {gs_vis_path}')
+        yield gs_vis_path
 
         # render video
         with torch.inference_mode():
@@ -339,6 +342,7 @@ class FreeSplatterRunner:
         save_video(video_frames, video_path, fps=30)
         print(f'Save video at {video_path}')
         t3 = time.time()
+        yield video_path
 
         # extract mesh
         with torch.inference_mode():
@@ -454,7 +458,7 @@ class FreeSplatterRunner:
         print(f'Generate mesh: {t4-t3:.2f} seconds.')
         print(f'Optimize mesh: {t5-t4:.2f} seconds.')
 
-        return gs_vis_path, video_path, mesh_fine_path, fig
+        yield mesh_fine_path
 
     def visualize_cameras_object(
         self, 
@@ -494,14 +498,12 @@ class FreeSplatterRunner:
         return fig
     
     # FreeSplatter-S
-    @spaces.GPU
     def run_views_to_scene(
         self, 
         image1,
         image2,
         cache_dir=None,
     ):
-        torch.cuda.empty_cache()
 
         self.output_dir = os.path.join(cache_dir, f'output_{uuid.uuid4()}')
         os.makedirs(self.output_dir, exist_ok=True)
@@ -531,7 +533,6 @@ class FreeSplatterRunner:
         images, 
         legends=None, 
     ):
-        torch.cuda.empty_cache()
 
         freesplatter = self.freesplatter_scene
 

freesplatter/webui/tab_img_to_3d.py

@@ -5,7 +5,7 @@ from .gradio_custommodel3d import CustomModel3D
 from .gradio_customgs import CustomGS
 
 
-def create_interface_img_to_3d(segmentation_api, freesplatter_api, model='Zero123++ v1.2'):
+def create_interface_img_to_3d(freesplatter_api, model='Zero123++ v1.2'):
     default_views = {
         'Zero123++ v1.1': ['Input', 'V2', 'V3', 'V5'],
         'Zero123++ v1.2': ['V1', 'V2', 'V3', 'V5', 'V6'],
@@ -137,14 +137,8 @@ def create_interface_img_to_3d(segmentation_api, freesplatter_api, model='Zero12
                 )
 
         var_dict['run_btn'].click(
-            fn=segmentation_api, 
-            inputs=var_dict['in_image'],
-            outputs=var_dict['fg_image'], 
-            concurrency_id='default_group',
-            api_name='run_segmentation',
-        ).success(
-            fn=partial(freesplatter_api, cache_dir=interface.GRADIO_CACHE),
-            inputs=[var_dict['fg_image'], 
+            fn=freesplatter_api,
+            inputs=[var_dict['in_image'], 
                     var_dict['model'], 
                     var_dict['diffusion_steps'], 
                     var_dict['guidance_scale'], 
@@ -152,7 +146,7 @@ def create_interface_img_to_3d(segmentation_api, freesplatter_api, model='Zero12
                     var_dict['view_indices'],
                     var_dict['gs_type'], 
                     var_dict['mesh_reduction']],
-            outputs=[var_dict['out_multiview'], var_dict['out_gs_vis'], var_dict['out_video'], var_dict['out_mesh'], var_dict['out_pose']], 
+            outputs=[var_dict['fg_image'], var_dict['out_multiview'], var_dict['out_pose'], var_dict['out_gs_vis'], var_dict['out_video'], var_dict['out_mesh']], 
             concurrency_id='default_group',
             api_name='run_image_to_3d',
         )

open3d_zerogpu_fix.py

@@ -0,0 +1,7 @@
+import fileinput
+import site
+from pathlib import Path
+
+with fileinput.FileInput(f'{site.getsitepackages()[0]}/open3d/__init__.py', inplace=True) as file:
+    for line in file:
+        print(line.replace('_pybind_cuda.open3d_core_cuda_device_count()', '1'), end='')