update

app.py

@@ -5,40 +5,31 @@ import spaces
 import ctypes
 import shlex
 import torch
-import argparse
 print(f'gradio version: {gr.__version__}')
 
-# Add command line argument parsing
-parser = argparse.ArgumentParser(description='DiMeR Demo')
-parser.add_argument('--ui_only', action='store_true', help='Only load the UI interface, do not initialize models (for UI debugging)')
-args = parser.parse_args()
 
-UI_ONLY_MODE = args.ui_only
-print(f"UI_ONLY_MODE: {UI_ONLY_MODE}")
-
-if not UI_ONLY_MODE:
-    subprocess.run(
-        shlex.split(
-            "pip install ./custom_diffusers --force-reinstall --no-deps"
-        )
+subprocess.run(
+    shlex.split(
+        "pip install ./custom_diffusers --force-reinstall --no-deps"
     )
-    subprocess.run(
-        shlex.split(
-            "pip install --no-index --no-cache-dir pytorch3d -f https://dl.fbaipublicfiles.com/pytorch3d/packaging/wheels/py310_cu121_pyt240/download.html"
-        )
+)
+subprocess.run(
+    shlex.split(
+        "pip install --no-index --no-cache-dir pytorch3d -f https://dl.fbaipublicfiles.com/pytorch3d/packaging/wheels/py310_cu121_pyt240/download.html"
     )
+)
 
-    subprocess.run(
-        shlex.split(
-            "pip install ./extension/nvdiffrast-0.3.1+torch-py3-none-any.whl --force-reinstall --no-deps"
-        )
+subprocess.run(
+    shlex.split(
+        "pip install ./extension/nvdiffrast-0.3.1+torch-py3-none-any.whl --force-reinstall --no-deps"
     )
+)
 
-    subprocess.run(
-        shlex.split(
-            "pip install ./extension/renderutils_plugin-0.1.0-cp310-cp310-linux_x86_64.whl --force-reinstall --no-deps"
-        )
+subprocess.run(
+    shlex.split(
+        "pip install ./extension/renderutils_plugin-0.1.0-cp310-cp310-linux_x86_64.whl --force-reinstall --no-deps"
     )
+)
 
 # Status variables for tracking if detailed prompt and image have been generated
 generated_detailed_prompt = False
@@ -61,9 +52,7 @@ def install_cuda_toolkit():
     os.environ["TORCH_CUDA_ARCH_LIST"] = "8.0;8.6"
     print("==> finished installation")
 
-# Only execute CUDA installation in non-UI debug mode
-if not UI_ONLY_MODE:
-    install_cuda_toolkit()
+install_cuda_toolkit()
 
 @spaces.GPU
 def check_gpu():
@@ -74,17 +63,9 @@ def check_gpu():
     os.environ['LD_LIBRARY_PATH'] = "/usr/local/cuda-12.1/lib64:" + os.environ.get('LD_LIBRARY_PATH', '')
     subprocess.run(['nvidia-smi'])  # Test if CUDA is available
     print(f"torch.cuda.is_available:{torch.cuda.is_available()}")
-    print("Device count:", torch.cuda.device_count()) 
-
-    # test nvdiffrast
-    import nvdiffrast.torch as dr
-    dr.RasterizeCudaContext(device="cuda:0")
-    print("nvdiffrast initialized successfully")       
+    print("Device count:", torch.cuda.device_count())      
     
-
-# Only check GPU in non-UI debug mode
-if not UI_ONLY_MODE:
-    check_gpu()
+check_gpu()
 
 
 import base64
@@ -108,23 +89,19 @@ import random
 import time
 import numpy as np
 
-# Only import video rendering module and initialize models in non-UI debug mode
-if not UI_ONLY_MODE:
-    from video_render import render_video_from_obj
+from video_render import render_video_from_obj
 
-    access_token = os.getenv("HUGGINGFACE_TOKEN")
-    from pipeline.kiss3d_wrapper import init_wrapper_from_config, run_text_to_3d, run_image_to_3d, image2mesh_preprocess, image2mesh_main
+access_token = os.getenv("HUGGINGFACE_TOKEN")
+from pipeline.kiss3d_wrapper import init_wrapper_from_config, run_text_to_3d, run_image_to_3d, image2mesh_preprocess, image2mesh_main
 
 # Add logo file path and hyperlinks
 LOGO_PATH = "app_assets/logo_temp_.png"  # Update this to the actual path of your logo
 ARXIV_LINK = "https://arxiv.org/pdf/2504.17670"
 GITHUB_LINK = "https://github.com/lutao2021/DiMeR"
 
-# Only initialize models in non-UI debug mode
-if not UI_ONLY_MODE:
-    k3d_wrapper = init_wrapper_from_config('./pipeline/pipeline_config/default.yaml')
-    from models.ISOMER.scripts.utils import fix_vert_color_glb
-    torch.backends.cuda.matmul.allow_tf32 = True
+k3d_wrapper = init_wrapper_from_config('./pipeline/pipeline_config/default.yaml')
+from models.ISOMER.scripts.utils import fix_vert_color_glb
+torch.backends.cuda.matmul.allow_tf32 = True
 
 TEMP_MESH_ADDRESS=''
 
@@ -166,11 +143,6 @@ def save_py3dmesh_with_trimesh_fast(meshes, save_glb_path=TEMP_MESH_ADDRESS, app
 
 @spaces.GPU
 def text_to_detailed(prompt, seed=None):
-    # test nvdiffrast
-    import nvdiffrast.torch as dr
-    dr.RasterizeCudaContext(device="cuda:0")
-    print("nvdiffrast initialized successfully")
-
     print(f"torch.cuda.is_available():{torch.cuda.is_available()}")
     # print(f"Before text_to_detailed: {torch.cuda.memory_allocated() / 1024**3} GB")
     return k3d_wrapper.get_detailed_prompt(prompt, seed)
@@ -240,7 +212,7 @@ def image2mesh_main_(reference_3d_bundle_image, caption, seed, strength1=0.5, st
         return gen_save_path, recon_mesh_path, mesh_cache
     # return gen_save_path, recon_mesh_path
 
-@spaces.GPU(duration=120)
+@spaces.GPU(duration=30)
 def bundle_image_to_mesh(
         gen_3d_bundle_image, 
         camera_radius=3.5,
@@ -328,8 +300,7 @@ def image_to_base64(image_path):
 
 # def main():
 
-if not UI_ONLY_MODE:
-    torch.set_grad_enabled(False)
+torch.set_grad_enabled(False)
 
 # Convert the logo image to base64
 logo_base64 = image_to_base64(LOGO_PATH)

models/DiMeR/models/DiMeR.py

@@ -86,11 +86,8 @@ class DiMeR(nn.Module):
         
     @spaces.GPU
     def init_flexicubes_geometry(self, device, fovy=50.0):
-        print(1)
         camera = PerspectiveCamera(fovy=fovy, device=device)
-        print(2)
         renderer = NeuralRender(device, camera_model=camera)
-        print(3)
         self.geometry = FlexiCubesGeometry(
             grid_res=self.grid_res, 
             scale=self.grid_scale, 

models/DiMeR/models/geometry/camera/perspective_camera.py

@@ -24,13 +24,9 @@ def projection(x=0.1, n=1.0, f=50.0, near_plane=None):
 class PerspectiveCamera(Camera):
     def __init__(self, fovy=49.0, device='cuda'):
         super(PerspectiveCamera, self).__init__()
-        print(1.1)
         self.device = device
-        print(1.2)
         focal = np.tan(fovy / 180.0 * np.pi * 0.5)
-        print(1.3)
         self.proj_mtx = torch.from_numpy(projection(x=focal, f=1000.0, n=1.0, near_plane=0.1)).to(self.device).unsqueeze(dim=0)
-        print(1.4)
 
     def project(self, points_bxnx4):
         out = torch.matmul(

models/DiMeR/models/geometry/render/neural_render.py

@@ -72,15 +72,10 @@ class NeuralRender(Renderer):
     @spaces.GPU
     def __init__(self, device='cuda', camera_model=None):
         super(NeuralRender, self).__init__()
-        print(2.1)
         self.device = device
-        print(2.2)
         self.ctx = dr.RasterizeCudaContext(device=device)
-        print(2.3)
         self.projection_mtx = None
-        print(2.4)
         self.camera = camera_model
-        print(2.5)
         
     # ==============================================================================================
     #  pixel shader

models/DiMeR/models/geometry/rep_3d/flexicubes_geometry.py

@@ -40,35 +40,26 @@ class FlexiCubesGeometry(Geometry):
             self, grid_res=64, scale=2.0, device='cuda', renderer=None,
             render_type='neural_render', args=None):
         super(FlexiCubesGeometry, self).__init__()
-        print(3.1)
         self.grid_res = grid_res
         self.device = device
         self.args = args
-        print(3.2)
         self.fc = FlexiCubes(device, weight_scale=0.5)
-        print(3.3)
         self.verts, self.indices = self.fc.construct_voxel_grid(grid_res)
-        print(3.4)
         if isinstance(scale, list):
             self.verts[:, 0] = self.verts[:, 0] * scale[0]
             self.verts[:, 1] = self.verts[:, 1] * scale[1]
             self.verts[:, 2] = self.verts[:, 2] * scale[1]
         else:
             self.verts = self.verts * scale
-        print(3.5)
             
         all_edges = self.indices[:, self.fc.cube_edges].reshape(-1, 2)
         self.all_edges = torch.unique(all_edges, dim=0)
 
         # Parameters used for fix boundary sdf
-        print(3.6)
         self.center_indices, self.boundary_indices = get_center_boundary_index(self.grid_res, device)
-        print(3.7)
         self.renderer = renderer
         self.render_type = render_type
-        print(3.8)
         self.ctx = dr.RasterizeCudaContext(device=device)
-        print(3.9)
 
         # self.verts.requires_grad_(True)