Revert "update"

This reverts commit e27ec6896bcd31b9c627966cf9bcf72aae7cbadc.

app.py

@@ -5,31 +5,40 @@ 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()
 
-subprocess.run(
-    shlex.split(
-        "pip install ./custom_diffusers --force-reinstall --no-deps"
+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 --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
@@ -52,7 +61,9 @@ def install_cuda_toolkit():
     os.environ["TORCH_CUDA_ARCH_LIST"] = "8.0;8.6"
     print("==> finished installation")
 
-install_cuda_toolkit()
+# Only execute CUDA installation in non-UI debug mode
+if not UI_ONLY_MODE:
+    install_cuda_toolkit()
 
 @spaces.GPU
 def check_gpu():
@@ -63,9 +74,17 @@ 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())      
+    print("Device count:", torch.cuda.device_count()) 
+
+    # test nvdiffrast
+    import nvdiffrast.torch as dr
+    dr.RasterizeCudaContext(device="cuda:0")
+    print("nvdiffrast initialized successfully")       
     
-check_gpu()
+
+# Only check GPU in non-UI debug mode
+if not UI_ONLY_MODE:
+    check_gpu()
 
 
 import base64
@@ -89,19 +108,23 @@ import random
 import time
 import numpy as np
 
-from video_render import render_video_from_obj
+# 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
 
-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"
 
-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
+# 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
 
 TEMP_MESH_ADDRESS=''
 
@@ -143,6 +166,11 @@ 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)
@@ -212,7 +240,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=30)
+@spaces.GPU(duration=120)
 def bundle_image_to_mesh(
         gen_3d_bundle_image, 
         camera_radius=3.5,
@@ -300,7 +328,8 @@ def image_to_base64(image_path):
 
 # def main():
 
-torch.set_grad_enabled(False)
+if not UI_ONLY_MODE:
+    torch.set_grad_enabled(False)
 
 # Convert the logo image to base64
 logo_base64 = image_to_base64(LOGO_PATH)

models/DiMeR/models/DiMeR.py

@@ -86,8 +86,11 @@ 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,9 +24,13 @@ 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,10 +72,15 @@ 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,26 +40,35 @@ 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)