Update gradio_app.py

gradio_app.py

@@ -18,8 +18,6 @@ import sys
 sys.path.insert(0, './hy3dshape')
 sys.path.insert(0, './hy3dpaint')
 
-pythonpath = sys.executable
-print(pythonpath)
 
 try:
     from torchvision_fix import apply_fix
@@ -52,7 +50,7 @@ from hy3dpaint.convert_utils import create_glb_with_pbr_materials
 
 
 MAX_SEED = 1e7
-ENV = "Huggingface" # "Huggingface"
+ENV = "Local" # "Huggingface"
 if ENV == 'Huggingface':
     """
     Setup environment for running on Huggingface platform.
@@ -67,65 +65,11 @@ if ENV == 'Huggingface':
         with the specified directory structure.
     """
     import os, spaces, subprocess, sys, shlex
-    from spaces import zero
-
-    def install_cuda_toolkit():
-        # CUDA_TOOLKIT_URL = "https://developer.download.nvidia.com/compute/cuda/11.8.0/local_installers/cuda_11.8.0_520.61.05_linux.run"
-        CUDA_TOOLKIT_URL = "https://developer.download.nvidia.com/compute/cuda/12.2.0/local_installers/cuda_12.2.0_535.54.03_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"
-
-    def prepare_env():
-        # print('install custom')
-        # os.system(f"cd /home/user/app/hy3dpaint/custom_rasterizer && {pythonpath} -m pip install -e .")
-        # os.system(f"cd /home/user/app/hy3dpaint/packages/custom_rasterizer && pip install -e .")
-        subprocess.run(shlex.split("pip install custom_rasterizer-0.1-cp310-cp310-linux_x86_64.whl"), check=True)
-
-        print("cd /home/user/app/hy3dpaint/differentiable_renderer/ && bash compile_mesh_painter.sh")
-        os.system("cd /home/user/app/hy3dpaint/DifferentiableRenderer && bash compile_mesh_painter.sh")
-        # print("wget https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth -P /home/user/app/hy3dpaint/ckpt")
-        # os.system("wget https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth -P /home/user/app/hy3dpaint/ckpt")
-
-    def check():
-        import custom_rasterizer
-        print(type(custom_rasterizer))
-        print(dir(custom_rasterizer))
-        print(getattr(custom_rasterizer, '__file__', None))
-
-        package_dir = None
-        if hasattr(custom_rasterizer, '__file__') and custom_rasterizer.__file__:
-            package_dir = os.path.dirname(custom_rasterizer.__file__)
-        elif hasattr(custom_rasterizer, '__path__'):
-            package_dir = list(custom_rasterizer.__path__)[0]
-        else:
-            raise RuntimeError("Cannot determine package path")
-        print(package_dir)
-
-        for root, dirs, files in os.walk(package_dir):
-            level = root.replace(package_dir, '').count(os.sep)
-            indent = ' ' * 4 * level
-            print(f"{indent}{os.path.basename(root)}/")
-            subindent = ' ' * 4 * (level + 1)
-            for f in files:
-                print(f"{subindent}{f}")
-
-    # print(torch.__version__)
-    # install_cuda_toolkit()
-    print(torch.__version__)
-    prepare_env()
-    check()
-    
+    print("cd /home/user/app/hy3dgen/texgen/differentiable_renderer/ && bash compile_mesh_painter.sh")
+    os.system("cd /home/user/app/hy3dgen/texgen/differentiable_renderer/ && bash compile_mesh_painter.sh")
+    print('install custom')
+    subprocess.run(shlex.split("pip install custom_rasterizer-0.1-cp310-cp310-linux_x86_64.whl"),
+                   check=True)
 else:
     """
     Define a dummy `spaces` module with a GPU decorator class for local environment.
@@ -378,7 +322,7 @@ def _gen_shape(
     main_image = image if not MV_MODE else image['front']
     return mesh, main_image, save_folder, stats, seed
 
-@spaces.GPU(duration=130)
+@spaces.GPU(duration=60)
 def generation_all(
     caption=None,
     image=None,
@@ -395,6 +339,8 @@ def generation_all(
     randomize_seed: bool = False,
 ):
     start_time_0 = time.time()
+    print("DEBUG: Starting generation_all...")
+    sys.stdout.flush()
     mesh, image, save_folder, stats, seed = _gen_shape(
         caption,
         image,
@@ -410,8 +356,12 @@ def generation_all(
         num_chunks=num_chunks,
         randomize_seed=randomize_seed,
     )
+    print("DEBUG: _gen_shape finished.")
+    sys.stdout.flush()
     path = export_mesh(mesh, save_folder, textured=False)
     
+    print("DEBUG: Initial export_mesh finished.")
+    sys.stdout.flush()
 
     print(path)
     print('='*40)
@@ -423,10 +373,18 @@ def generation_all(
     # stats['time']['postprocessing'] = time.time() - tmp_time
 
     tmp_time = time.time()
+    print("DEBUG: Starting face_reduce_worker...")
+    sys.stdout.flush()
     mesh = face_reduce_worker(mesh)
+    print("DEBUG: face_reduce_worker finished.")
+    sys.stdout.flush()
 
     # path = export_mesh(mesh, save_folder, textured=False, type='glb')
+    print("DEBUG: Starting second export_mesh (obj)...")
+    sys.stdout.flush()
     path = export_mesh(mesh, save_folder, textured=False, type='obj') # 这样操作也会 core dump
+    print("DEBUG: Second export_mesh (obj) finished.")
+    sys.stdout.flush()
 
     logger.info("---Face Reduction takes %s seconds ---" % (time.time() - tmp_time))
     stats['time']['face reduction'] = time.time() - tmp_time
@@ -434,7 +392,11 @@ def generation_all(
     tmp_time = time.time()
 
     text_path = os.path.join(save_folder, f'textured_mesh.obj')
+    print("DEBUG: Starting tex_pipeline...")
+    sys.stdout.flush()
     path_textured = tex_pipeline(mesh_path=path, image_path=image, output_mesh_path=text_path, save_glb=False)
+    print("DEBUG: tex_pipeline finished.")
+    sys.stdout.flush()
         
     logger.info("---Texture Generation takes %s seconds ---" % (time.time() - tmp_time))
     stats['time']['texture generation'] = time.time() - tmp_time
@@ -442,7 +404,11 @@ def generation_all(
     tmp_time = time.time()
     # Convert textured OBJ to GLB using obj2gltf with PBR support
     glb_path_textured = os.path.join(save_folder, 'textured_mesh.glb')
+    print("DEBUG: Starting quick_convert_with_obj2gltf...")
+    sys.stdout.flush()
     conversion_success = quick_convert_with_obj2gltf(path_textured, glb_path_textured)
+    print("DEBUG: quick_convert_with_obj2gltf finished.")
+    sys.stdout.flush()
 
     logger.info("---Convert textured OBJ to GLB takes %s seconds ---" % (time.time() - tmp_time))
     stats['time']['convert textured OBJ to GLB'] = time.time() - tmp_time
@@ -797,7 +763,7 @@ if __name__ == '__main__':
     parser.add_argument('--host', type=str, default='0.0.0.0')
     parser.add_argument('--device', type=str, default='cuda')
     parser.add_argument('--mc_algo', type=str, default='mc')
-    parser.add_argument('--cache-path', type=str, default='/root/save_dir')
+    parser.add_argument('--cache-path', type=str, default='./save_dir')
     parser.add_argument('--enable_t23d', action='store_true')
     parser.add_argument('--disable_tex', action='store_true')
     parser.add_argument('--enable_flashvdm', action='store_true')
@@ -869,24 +835,26 @@ if __name__ == '__main__':
             HAS_TEXTUREGEN = True
             
         except Exception as e:
+            import traceback
+            traceback.print_exc()
             print(f"Error loading texture generator: {e}")
             print("Failed to load texture generator.")
             print('Please try to install requirements by following README.md')
             HAS_TEXTUREGEN = False
 
-    # HAS_T2I = True
-    # if args.enable_t23d:
-    #     from hy3dgen.text2image import HunyuanDiTPipeline
+    HAS_T2I = True
+    if args.enable_t23d:
+        from hy3dgen.text2image import HunyuanDiTPipeline
 
-    #     t2i_worker = HunyuanDiTPipeline('Tencent-Hunyuan/HunyuanDiT-v1.1-Diffusers-Distilled')
-    #     HAS_T2I = True
+        t2i_worker = HunyuanDiTPipeline('Tencent-Hunyuan/HunyuanDiT-v1.1-Diffusers-Distilled')
+        HAS_T2I = True
 
     from hy3dshape import FaceReducer, FloaterRemover, DegenerateFaceRemover, MeshSimplifier, \
         Hunyuan3DDiTFlowMatchingPipeline
     from hy3dshape.pipelines import export_to_trimesh
-    from hy3dshape.rembg import BackgroundRemover
+    from hy3dshape.birefnet_rembg import BiRefNetRemover
 
-    rmbg_worker = BackgroundRemover()
+    rmbg_worker = BiRefNetRemover()
     i23d_worker = Hunyuan3DDiTFlowMatchingPipeline.from_pretrained(
         args.model_path,
         subfolder=args.subfolder,
@@ -915,13 +883,6 @@ if __name__ == '__main__':
 
     if args.low_vram_mode:
         torch.cuda.empty_cache()
-        
     demo = build_app()
     app = gr.mount_gradio_app(app, demo, path="/")
-
-    if ENV == 'Huggingface':
-        # for Zerogpu
-        from spaces import zero
-        zero.startup()
-
     uvicorn.run(app, host=args.host, port=args.port)