Deploy DiffusionGS online generation demos

diffusionGS/pipline_obj.py

@@ -3,7 +3,6 @@ import warnings
 from typing import Callable, List, Optional, Union, Dict, Any
 import PIL.Image
 import trimesh
-import rembg
 import torch
 import numpy as np
 from huggingface_hub import hf_hub_download
@@ -122,6 +121,7 @@ class DiffusionGSPipeline():
                 do_remove = False
             do_remove = do_remove or force
             if do_remove:
+                import rembg
                 image = rembg.remove(image)
 
             # calculate the min bbox of the image

diffusionGS/systems/utils.py

@@ -13,7 +13,6 @@ from torchvision import transforms
 from diffusionGS.utils.structure import Mesh
 import inspect
 from PIL import Image, ImageOps
-import rembg
 import cv2
 def _warmup_beta(beta_start, beta_end, num_diffusion_timesteps, warmup_frac):
     betas = beta_end * np.ones(num_diffusion_timesteps, dtype=np.float64)
@@ -573,6 +572,7 @@ def preprocess(image:Image.Image,session) -> Image.Image:
         rgba = image.copy()
     else:
         # https://github.com/pymatting/pymatting/issues/19
+        import rembg
         rgba = rembg.remove(image, alpha_matting=True, session=session)
 
     alpha = rgba.getchannel('A')

diffusionGS/utils/mesh_utils.py

@@ -1,5 +1,4 @@
 import numpy as np
-import pymeshlab as pml
 
 
 def poisson_mesh_reconstruction(points, normals=None):
@@ -57,6 +56,7 @@ def decimate_mesh(
         solver.simplify_mesh(target_count=target, preserve_border=False, verbose=False)
         verts, faces, normals = solver.getMesh()
     else:
+        import pymeshlab as pml
         m = pml.Mesh(verts, faces)
         ms = pml.MeshSet()
         ms.add_mesh(m, "mesh")  # will copy!
@@ -95,6 +95,7 @@ def clean_mesh(
     remesh=True,
     remesh_size=0.01,
 ):
+    import pymeshlab as pml
     # verts: [N, 3]
     # faces: [N, 3]
 
@@ -144,4 +145,4 @@ def clean_mesh(
         f"[INFO] mesh cleaning: {_ori_vert_shape} --> {verts.shape}, {_ori_face_shape} --> {faces.shape}"
     )
 
-    return verts, faces
+    return verts, faces

diffusionGS/utils/saving.py

@@ -17,7 +17,6 @@ from pytorch_lightning.loggers import WandbLogger
 import torchvision
 from diffusionGS.models.geometry.utils import Mesh
 from diffusionGS.utils.typing import *
-import open3d as o3d
 from einops import rearrange
 from diffusionGS.models.gsrenderer.gs_core import imageseq2video, render_turntable, render_generic
 from diffusionGS.models.gsrenderer.cam_utils import get_interpolated_poses_many
@@ -521,6 +520,7 @@ class SaverMixin:
         return save_path 
     
     def save_vertices_as_ply_open3d(self, vertices, normal, filepath):
+        import open3d as o3d
         # points = [[v[0], v[1], v[2]] for v in vertices]
         points  = vertices
         point_cloud = o3d.geometry.PointCloud()

space_app.py

@@ -7,6 +7,8 @@ from types import SimpleNamespace
 # Gradio worker thread when ZeroGPU starts the GPU process.
 os.environ.setdefault("NUMBA_THREADING_LAYER", "workqueue")
 os.environ.setdefault("NUMBA_NUM_THREADS", "2")
+for variable in ("OMP_NUM_THREADS", "MKL_NUM_THREADS", "OPENBLAS_NUM_THREADS"):
+    os.environ[variable] = "1"
 
 # ZeroGPU must install its CUDA emulation before torch is imported.
 import spaces
@@ -54,7 +56,7 @@ def load_object_pipeline(model_dir):
     return DiffusionGSPipeline("cuda:0", config, system)
 
 
-torch.set_num_threads(2)
+torch.set_num_threads(1)
 model_dir = Path(snapshot_download(
     MODEL_REPO, revision=MODEL_REVISION,
     allow_patterns=["obj_ckpt_512.ckpt", "obj_configs.yaml", "camera_template.pt",
@@ -106,8 +108,9 @@ object_service.generate = generate_object
 object_service.extract_mesh = extract_mesh
 scene_service.generate = generate_scene
 demo = build_object_demo(service=object_service)
+scene_demo = build_scene_demo(service=scene_service)
 with demo.route("Scene Generation", "/scene"):
-    build_scene_demo(service=scene_service).render()
+    scene_demo.render()
 
 if __name__ == "__main__":
     demo.launch(server_name="0.0.0.0", server_port=int(os.environ.get("GRADIO_SERVER_PORT", "7860")),