Update app.py

app.py

@@ -1,8 +1,16 @@
 # MIT License
 # (see original notice and terms)
 
-import gradio as gr
 import os
+import types
+import zipfile
+import importlib
+from typing import *
+
+import gradio as gr
+import numpy as np
+import torch
+import tempfile
 
 # ---- Force CPU-only environment globally ----
 os.environ["CUDA_VISIBLE_DEVICES"] = "-1"          # hide GPUs from torch
@@ -10,14 +18,6 @@ os.environ.setdefault("ATTN_BACKEND", "sdpa")      # avoid xformers path
 os.environ.setdefault("SPCONV_ALGO", "native")     # safe sparseconv algo
 # ---------------------------------------------
 
-from typing import *
-import torch
-import numpy as np
-import tempfile
-import zipfile
-import types
-import importlib
-
 # ---------------------------------------------------------------------------
 # Ensure bundled hi3dgen sources are available (extracted from hi3dgen.zip)
 # ---------------------------------------------------------------------------
@@ -52,6 +52,7 @@ def _ensure_xformers_stub():
     ops_mod = types.ModuleType('xformers.ops')
 
     def memory_efficient_attention(query, key, value, attn_bias=None):
+        # SDPA fallback
         return F.scaled_dot_product_attention(query, key, value, attn_bias)
 
     ops_mod.memory_efficient_attention = memory_efficient_attention
@@ -62,20 +63,19 @@ def _ensure_xformers_stub():
 _ensure_xformers_stub()
 
 # ---------------------------------------------------------------------------
-# Import pipeline AFTER stubbing xformers, then patch CUDA-hotspots to CPU
+# Patch CUDA hotspots to CPU **BEFORE** importing the pipeline
 # ---------------------------------------------------------------------------
-from hi3dgen.pipelines import Hi3DGenPipeline
-import trimesh
-
-# ---- Force CPU inside hi3dgen (avoid any CUDA paths) ----
 print("[PATCH] Applying CPU monkey-patches to hi3dgen")
 
 # 1) utils_cube.construct_dense_grid(..., device=...) -> force CPU
 uc = importlib.import_module("hi3dgen.representations.mesh.utils_cube")
 if not hasattr(uc, "_CPU_PATCHED"):
-    _orig_construct_dense_grid = uc.construct_dense_grid
+    _uc_orig_construct_dense_grid = uc.construct_dense_grid
+
     def _construct_dense_grid_cpu(res, device=None):
-        return _orig_construct_dense_grid(res, device="cpu")
+        # ignore any requested device, always CPU
+        return _uc_orig_construct_dense_grid(res, device="cpu")
+
     uc.construct_dense_grid = _construct_dense_grid_cpu
     uc._CPU_PATCHED = True
     print("[PATCH] utils_cube.construct_dense_grid -> CPU")
@@ -85,29 +85,44 @@ cm = importlib.import_module("hi3dgen.representations.mesh.cube2mesh")
 M = cm.EnhancedMarchingCubes
 if not hasattr(M, "_CPU_PATCHED"):
     _orig_init = M.__init__
+
     def _init_cpu(self, *args, **kwargs):
-        # Ensure device is CPU regardless of how it's passed
+        # ensure device ends up on CPU regardless of how it's passed
         if "device" in kwargs:
             kwargs["device"] = torch.device("cpu")
         else:
             kwargs.setdefault("device", torch.device("cpu"))
         return _orig_init(self, *args, **kwargs)
+
     M.__init__ = _init_cpu
     M._CPU_PATCHED = True
     print("[PATCH] cube2mesh.EnhancedMarchingCubes.__init__ -> CPU (flex)")
 
-# 3) Belt & suspenders: coerce torch.arange(device='cuda') to CPU if any call slips through
+# 3) IMPORTANT: cube2mesh does "from .utils_cube import construct_dense_grid"
+#    so we must override the BOUND symbol inside cube2mesh as well.
+if getattr(cm, "construct_dense_grid", None) is not _construct_dense_grid_cpu:
+    cm.construct_dense_grid = _construct_dense_grid_cpu
+    print("[PATCH] cube2mesh.construct_dense_grid (bound name) -> CPU")
+
+# 4) Belt & suspenders: coerce torch.arange(device='cuda') to CPU if anything slips through
 if not hasattr(torch, "_ARANGE_CPU_PATCHED"):
     _orig_arange = torch.arange
+
     def _arange_cpu(*args, **kwargs):
         dev = kwargs.get("device", None)
         if dev is not None and str(dev).startswith("cuda"):
             kwargs["device"] = "cpu"
         return _orig_arange(*args, **kwargs)
+
     torch.arange = _arange_cpu
     torch._ARANGE_CPU_PATCHED = True
     print("[PATCH] torch.arange(device='cuda') -> CPU")
-# ----------------------------------------------------------
+
+# ---------------------------------------------------------------------------
+# Now import pipeline (AFTER patches so bound names are already overridden)
+# ---------------------------------------------------------------------------
+from hi3dgen.pipelines import Hi3DGenPipeline
+import trimesh
 
 MAX_SEED = np.iinfo(np.int32).max
 TMP_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'tmp')
@@ -326,7 +341,7 @@ with gr.Blocks(css="footer {visibility: hidden}") as demo:
     gr.Markdown(
         """
         **Acknowledgments**: Hi3DGen is built on the shoulders of giants. We would like to express our gratitude to the open-source research community and the developers of these pioneering projects:
-        - **3D Modeling:** Our 3D Model is finetuned from the SOTA open-source 3D foundation model [Trellis](https://github.com/microsoft/TRELLIS); inspired by Rodin, Tripo, and Dora.
+        - **3D Modeling:** Finetuned from the SOTA open-source 3D foundation model [Trellis].
         - **Normal Estimation:** Builds on StableNormal and GenPercept.
         """
     )