Update app.py

app.py

@@ -1,171 +1,136 @@
 # MIT License
-
 # Copyright (c) Microsoft
-
 # Permission is hereby granted, free of charge, to any person obtaining a copy
 # of this software and associated documentation files (the "Software"), to deal
 # in the Software without restriction, including without limitation the rights
 # to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
-# copies of the Software, and to permit persons to whom the Software is
-# furnished to do so, subject to the following conditions:
-
-# The above copyright notice and this permission notice shall be included in all
-# copies or substantial portions of the Software.
-
-# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
-# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
-# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
-# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
-# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
-# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
-# SOFTWARE.
+# copies of the Software, and to permit persons to do so, subject to the following conditions:
+# The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
 
 # Copyright (c) [2025] [Microsoft]
-# Copyright (c) [2025] [Chongjie Ye] 
+# Copyright (c) [2025] [Chongjie Ye]
 # SPDX-License-Identifier: MIT
 # This file has been modified by Chongjie Ye on 2025/04/10
-# Original file was released under MIT, with the full license text # available at https://github.com/atong01/conditional-flow-matching/blob/1.0.7/LICENSE.
+# Original file was released under MIT, with the full license text available at:
+# https://github.com/atong01/conditional-flow-matching/blob/1.0.7/LICENSE
 # This modified file is released under the same license.
 
 import gradio as gr
 import os
 
 # ---- Force CPU-only environment globally ----
-os.environ["CUDA_VISIBLE_DEVICES"] = "-1"          # hide any GPUs from torch
-os.environ.setdefault("ATTN_BACKEND", "sdpa")      # avoid xformers
+os.environ["CUDA_VISIBLE_DEVICES"] = "-1"          # hide GPUs from torch
+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
 
 # ---------------------------------------------------------------------------
-# NOTE
-# The original Hi3DGen implementation expects the `hi3dgen` Python package to
-# reside alongside this app file.  Hugging Face Spaces do not currently
-# support uploading an entire folder via the web interface, so the `hi3dgen`
-# source tree is bundled into a single `hi3dgen.zip` archive.  On startup we
-# extract this archive into the working directory if the `hi3dgen` package is
-# not already present.  This allows the rest of the code to `import hi3dgen` as
-# normal.
+# Ensure bundled hi3dgen sources are available (extracted from hi3dgen.zip)
 # ---------------------------------------------------------------------------
-
 def _ensure_hi3dgen_available():
-    """Unpack hi3dgen.zip into the current directory if the hi3dgen package
-    is missing.  This function is idempotent and safe to call multiple times.
-    """
     pkg_name = 'hi3dgen'
-    pkg_dir = os.path.join(os.path.dirname(__file__), pkg_name)
+    here = os.path.dirname(__file__)
+    pkg_dir = os.path.join(here, pkg_name)
     if os.path.isdir(pkg_dir):
         return
-    archive_path = os.path.join(os.path.dirname(__file__), f"{pkg_name}.zip")
-    if os.path.isfile(archive_path):
-        try:
-            with zipfile.ZipFile(archive_path, 'r') as zf:
-                zf.extractall(os.path.dirname(__file__))
-        except Exception as e:
-            raise RuntimeError(f"Failed to extract {archive_path}: {e}")
-    else:
+    archive_path = os.path.join(here, f"{pkg_name}.zip")
+    if not os.path.isfile(archive_path):
         raise FileNotFoundError(
-            f"Required archive {archive_path} is missing. Make sure to upload the hi3dgen.zip file alongside app.py."
+            f"Required archive {archive_path} is missing. Upload hi3dgen.zip next to app.py."
         )
+    try:
+        with zipfile.ZipFile(archive_path, 'r') as zf:
+            zf.extractall(here)
+    except Exception as e:
+        raise RuntimeError(f"Failed to extract {archive_path}: {e}")
 
-# Make sure the hi3dgen package is available before importing it
 _ensure_hi3dgen_available()
 
 # ---------------------------------------------------------------------------
-# xformers stub
-#
-# Some modules in the Hi3DGen pipeline import `xformers.ops.memory_efficient_attention`
-# to compute multi-head attention.  The official `xformers` library is not
-# installed in this Space (and requires GPU-only build), so we provide a
-# minimal in-memory stub that exposes a compatible API backed by PyTorch's
-# built-in scaled dot-product attention.  This stub is lightweight and
-# CPU-friendly.  It registers both the `xformers` and `xformers.ops` modules
-# in sys.modules so that subsequent imports succeed.
+# xformers stub (CPU-friendly fallback for xformers.ops.memory_efficient_attention)
 # ---------------------------------------------------------------------------
-
 def _ensure_xformers_stub():
-
-    # ---------------------------------------------------------------------------
     import sys
-    # If xformers is already available, do nothing.
     if 'xformers.ops' in sys.modules:
         return
     import torch.nn.functional as F
-    # Create a new module object for xformers and its ops submodule
-    xformers_mod = types.ModuleType('xformers')
+    xf_mod = types.ModuleType('xformers')
     ops_mod = types.ModuleType('xformers.ops')
 
     def memory_efficient_attention(query, key, value, attn_bias=None):
-        """
-        Fallback implementation of memory_efficient_attention for CPU environments.
-        This wraps torch.nn.functional.scaled_dot_product_attention.
-        """
-        # PyTorch expects the attention mask (bias) to be additive with shape
-        # broadcastable to (batch, num_heads, seq_len_query, seq_len_key).  If
-        # attn_bias is provided and is non-zero, pass it through; otherwise
-        # supply None to avoid unnecessary allocations.
         return F.scaled_dot_product_attention(query, key, value, attn_bias)
 
-    # Populate the ops module with our fallback function
     ops_mod.memory_efficient_attention = memory_efficient_attention
-    # Expose ops as an attribute of xformers
-    xformers_mod.ops = ops_mod
-    # Register modules
-    sys.modules['xformers'] = xformers_mod
+    xf_mod.ops = ops_mod
+    sys.modules['xformers'] = xf_mod
     sys.modules['xformers.ops'] = ops_mod
 
-    # Ensure the xformers stub is registered before importing Hi3DGen
 _ensure_xformers_stub()
 
 # ---------------------------------------------------------------------------
-# Monkey-patch Hi3DGen to run on CPU
-#
-# Some utility functions and classes in the Hi3DGen codebase assume the
-# presence of a CUDA device by default.  Specifically, the function
-# `construct_dense_grid` in `hi3dgen.representations.mesh.utils_cube` uses
-# `device='cuda'` as its default argument, and the class
-# `EnhancedMarchingCubes` in `hi3dgen.representations.mesh.cube2mesh` has
-# a constructor that defaults to `device="cuda"`.  On CPU-only Spaces, these
-# defaults cause runtime errors when PyTorch attempts to allocate tensors on
-# a non-existent GPU.  To avoid this, we override the default arguments for
-# these functions to use the CPU instead.  If the patch cannot be applied
-# (for example, if the module structure changes in a future version), we
-# catch any exceptions and log a warning without stopping execution.
-try:
-    from hi3dgen.representations.mesh import utils_cube
-    if hasattr(utils_cube.construct_dense_grid, '__defaults__'):
-        _defaults = list(utils_cube.construct_dense_grid.__defaults__ or ())
-        if _defaults and _defaults[-1] == 'cuda':
-            _defaults[-1] = 'cpu'
-            utils_cube.construct_dense_grid.__defaults__ = tuple(_defaults)
-
-    from hi3dgen.representations.mesh.cube2mesh import EnhancedMarchingCubes
-    if hasattr(EnhancedMarchingCubes.__init__, '__defaults__'):
-        _mc_defaults = list(EnhancedMarchingCubes.__init__.__defaults__ or ())
-        if _mc_defaults and _mc_defaults[-1] == 'cuda':
-            _mc_defaults[-1] = 'cpu'
-            EnhancedMarchingCubes.__init__.__defaults__ = tuple(_mc_defaults)
-except Exception as _e:
-    print(f"Warning: failed to apply CPU device overrides: {_e}")
-
+# Import pipeline AFTER stubbing xformers, then patch CUDA-hotspots to CPU
+# ---------------------------------------------------------------------------
 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
+    def _construct_dense_grid_cpu(res, device=None):
+        return _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")
+
+# 2) cube2mesh.EnhancedMarchingCubes default device -> force CPU
+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, res, device=None):
+        return _orig_init(self, res, device=torch.device("cpu"))
+    M.__init__ = _init_cpu
+    M._CPU_PATCHED = True
+    print("[PATCH] cube2mesh.EnhancedMarchingCubes.__init__ -> CPU")
+
+# 3) Belt & suspenders: coerce torch.arange(device='cuda') to CPU if any call 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")
+# ----------------------------------------------------------
+
 MAX_SEED = np.iinfo(np.int32).max
 TMP_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'tmp')
 WEIGHTS_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'weights')
 os.makedirs(TMP_DIR, exist_ok=True)
 os.makedirs(WEIGHTS_DIR, exist_ok=True)
 
+# ---------------------------------------------------------------------------
+# Weights caching
+# ---------------------------------------------------------------------------
 def cache_weights(weights_dir: str) -> dict:
-    import os
     from huggingface_hub import snapshot_download
-
     os.makedirs(weights_dir, exist_ok=True)
     model_ids = [
         "Stable-X/trellis-normal-v0-1",
@@ -175,49 +140,62 @@ def cache_weights(weights_dir: str) -> dict:
     cached_paths = {}
     for model_id in model_ids:
         print(f"Caching weights for: {model_id}")
-        # Check if the model is already cached
         local_path = os.path.join(weights_dir, model_id.split("/")[-1])
         if os.path.exists(local_path):
             print(f"Already cached at: {local_path}")
             cached_paths[model_id] = local_path
             continue
-        # Download the model and cache it
         print(f"Downloading and caching model: {model_id}")
-        # Use snapshot_download to download the model
-        local_path = snapshot_download(repo_id=model_id, local_dir=os.path.join(weights_dir, model_id.split("/")[-1]), force_download=False)
+        local_path = snapshot_download(
+            repo_id=model_id,
+            local_dir=os.path.join(weights_dir, model_id.split("/")[-1]),
+            force_download=False
+        )
         cached_paths[model_id] = local_path
         print(f"Cached at: {local_path}")
-
     return cached_paths
 
+# ---------------------------------------------------------------------------
+# Pre/Post processing and generation
+# ---------------------------------------------------------------------------
 def preprocess_mesh(mesh_prompt):
     print("Processing mesh")
     trimesh_mesh = trimesh.load_mesh(mesh_prompt)
-    trimesh_mesh.export(mesh_prompt+'.glb')
-    return mesh_prompt+'.glb'
+    out_path = mesh_prompt + '.glb'
+    trimesh_mesh.export(out_path)
+    return out_path
 
 def preprocess_image(image):
     if image is None:
         return None
-    image = hi3dgen_pipeline.preprocess_image(image, resolution=1024)
-    return image
-
-def generate_3d(image, seed=-1,  
-                ss_guidance_strength=3, ss_sampling_steps=50,
-                slat_guidance_strength=3, slat_sampling_steps=6,):
+    return hi3dgen_pipeline.preprocess_image(image, resolution=1024)
+
+def generate_3d(
+    image,
+    seed: int = -1,
+    ss_guidance_strength: float = 3,
+    ss_sampling_steps: int = 50,
+    slat_guidance_strength: float = 3,
+    slat_sampling_steps: int = 6,
+):
     if image is None:
         return None, None, None
 
     if seed == -1:
         seed = np.random.randint(0, MAX_SEED)
-    
+
     image = hi3dgen_pipeline.preprocess_image(image, resolution=1024)
-    normal_image = normal_predictor(image, resolution=768, match_input_resolution=True, data_type='object')
+    normal_image = normal_predictor(
+        image,
+        resolution=768,
+        match_input_resolution=True,
+        data_type='object'
+    )
 
     outputs = hi3dgen_pipeline.run(
         normal_image,
         seed=seed,
-        formats=["mesh",],
+        formats=["mesh"],
         preprocess_image=False,
         sparse_structure_sampler_params={
             "steps": ss_sampling_steps,
@@ -229,80 +207,71 @@ def generate_3d(image, seed=-1,
         },
     )
     generated_mesh = outputs['mesh'][0]
-    
-    # Save outputs
-    import datetime
-
 
+    import datetime
     output_id = datetime.datetime.now().strftime("%Y%m%d%H%M%S")
     os.makedirs(os.path.join(TMP_DIR, output_id), exist_ok=True)
     mesh_path = f"{TMP_DIR}/{output_id}/mesh.glb"
-    
-    # Export mesh
-    trimesh_mesh = generated_mesh.to_trimesh(transform_pose=True)
 
+    trimesh_mesh = generated_mesh.to_trimesh(transform_pose=True)
     trimesh_mesh.export(mesh_path)
 
     return normal_image, mesh_path, mesh_path
 
 def convert_mesh(mesh_path, export_format):
-    """Download the mesh in the selected format."""
     if not mesh_path:
         return None
-    
-    # Create a temporary file to store the mesh data
     temp_file = tempfile.NamedTemporaryFile(suffix=f".{export_format}", delete=False)
     temp_file_path = temp_file.name
-    
-    new_mesh_path = mesh_path.replace(".glb", f".{export_format}")
     mesh = trimesh.load_mesh(mesh_path)
-    mesh.export(temp_file_path)  # Export to the temporary file
-    
-    return temp_file_path # Return the path to the temporary file
+    mesh.export(temp_file_path)
+    return temp_file_path
 
-# Create the Gradio interface with improved layout
+# ---------------------------------------------------------------------------
+# UI
+# ---------------------------------------------------------------------------
 with gr.Blocks(css="footer {visibility: hidden}") as demo:
     gr.Markdown(
         """
         <h1 style='text-align: center;'>Hi3DGen: High-fidelity 3D Geometry Generation from Images via Normal Bridging</h1>
         <p style='text-align: center;'>
             <strong>V0.1, Introduced By 
-            <a href="https://gaplab.cuhk.edu.cn/" target="_blank">GAP Lab</a> from CUHKSZ and 
-            <a href="https://www.nvsgames.cn/" target="_blank">Game-AIGC Team</a> from ByteDance</strong>
+            <a href="https://gaplab.cuhk.edu.cn/" target="_blank">GAP Lab</a> (CUHKSZ) and 
+            <a href="https://www.nvsgames.cn/" target="_blank">Game-AIGC Team</a> (ByteDance)</strong>
         </p>
         """
     )
-    
+
     with gr.Row():
         gr.Markdown("""
-                    <p align="center">
-                    <a title="Website" href="https://stable-x.github.io/Hi3DGen/" target="_blank" rel="noopener noreferrer" style="display: inline-block;">
-                        <img src="https://www.obukhov.ai/img/badges/badge-website.svg">
-                    </a>
-                    <a title="arXiv" href="https://stable-x.github.io/Hi3DGen/hi3dgen_paper.pdf" target="_blank" rel="noopener noreferrer" style="display: inline-block;">
-                        <img src="https://www.obukhov.ai/img/badges/badge-pdf.svg">
-                    </a>
-                    <a title="Github" href="https://github.com/Stable-X/Hi3DGen" target="_blank" rel="noopener noreferrer" style="display: inline-block;">
-                        <img src="https://img.shields.io/github/stars/Stable-X/Hi3DGen?label=GitHub%20%E2%98%85&logo=github&color=C8C" alt="badge-github-stars">
-                    </a>
-                    <a title="Social" href="https://x.com/ychngji6" target="_blank" rel="noopener noreferrer" style="display: inline-block;">
-                        <img src="https://www.obukhov.ai/img/badges/badge-social.svg" alt="social">
-                    </a>
-                    </p>
-                    """)
+            <p align="center">
+            <a title="Website" href="https://stable-x.github.io/Hi3DGen/" target="_blank" rel="noopener noreferrer" style="display: inline-block;">
+                <img src="https://www.obukhov.ai/img/badges/badge-website.svg">
+            </a>
+            <a title="arXiv" href="https://stable-x.github.io/Hi3DGen/hi3dgen_paper.pdf" target="_blank" rel="noopener noreferrer" style="display: inline-block;">
+                <img src="https://www.obukhov.ai/img/badges/badge-pdf.svg">
+            </a>
+            <a title="Github" href="https://github.com/Stable-X/Hi3DGen" target="_blank" rel="noopener noreferrer" style="display: inline-block;">
+                <img src="https://img.shields.io/github/stars/Stable-X/Hi3DGen?label=GitHub%20%E2%98%85&logo=github&color=C8C" alt="badge-github-stars">
+            </a>
+            <a title="Social" href="https://x.com/ychngji6" target="_blank" rel="noopener noreferrer" style="display: inline-block;">
+                <img src="https://www.obukhov.ai/img/badges/badge-social.svg" alt="social">
+            </a>
+            </p>
+        """)
 
     with gr.Row():
         with gr.Column(scale=1):
             with gr.Tabs():
-                
                 with gr.Tab("Single Image"):
                     with gr.Row():
                         image_prompt = gr.Image(label="Image Prompt", image_mode="RGBA", type="pil")
                         normal_output = gr.Image(label="Normal Bridge", image_mode="RGBA", type="pil")
-                        
                 with gr.Tab("Multiple Images"):
-                    gr.Markdown("<div style='text-align: center; padding: 40px; font-size: 24px;'>Multiple Images functionality is coming soon!</div>")
-                        
+                    gr.Markdown(
+                        "<div style='text-align: center; padding: 40px; font-size: 24px;'>Multiple Images functionality is coming soon!</div>"
+                    )
+
             with gr.Accordion("Advanced Settings", open=False):
                 seed = gr.Slider(-1, MAX_SEED, label="Seed", value=0, step=1)
                 gr.Markdown("#### Stage 1: Sparse Structure Generation")
@@ -313,15 +282,14 @@ with gr.Blocks(css="footer {visibility: hidden}") as demo:
                 with gr.Row():
                     slat_guidance_strength = gr.Slider(0.0, 10.0, label="Guidance Strength", value=3.0, step=0.1)
                     slat_sampling_steps = gr.Slider(1, 50, label="Sampling Steps", value=6, step=1)
-                    
+
             with gr.Group():
                 with gr.Row():
                     gen_shape_btn = gr.Button("Generate Shape", size="lg", variant="primary")
-                        
-        # Right column - Output
+
         with gr.Column(scale=1):
             with gr.Column():
-                model_output = gr.Model3D(label="3D Model Preview (Each model is approximately 40MB, may take around 1 minute to load)")
+                model_output = gr.Model3D(label="3D Model Preview (Each model is ~40MB; may take ~1 min to load)")
             with gr.Column():
                 export_format = gr.Dropdown(
                     choices=["obj", "glb", "ply", "stl"],
@@ -335,11 +303,11 @@ with gr.Blocks(css="footer {visibility: hidden}") as demo:
         inputs=[image_prompt],
         outputs=[image_prompt]
     )
-    
+
     gen_shape_btn.click(
         generate_3d,
         inputs=[
-            image_prompt, seed,  
+            image_prompt, seed,
             ss_guidance_strength, ss_sampling_steps,
             slat_guidance_strength, slat_sampling_steps
         ],
@@ -348,15 +316,13 @@ with gr.Blocks(css="footer {visibility: hidden}") as demo:
         lambda: gr.Button(interactive=True),
         outputs=[download_btn],
     )
-    
-    
+
     def update_download_button(mesh_path, export_format):
         if not mesh_path:
             return gr.File.update(value=None, interactive=False)
-        
         download_path = convert_mesh(mesh_path, export_format)
         return download_path
-    
+
     export_format.change(
         update_download_button,
         inputs=[model_output, export_format],
@@ -365,31 +331,19 @@ with gr.Blocks(css="footer {visibility: hidden}") as demo:
         lambda: gr.Button(interactive=True),
         outputs=[download_btn],
     )
-    
+
     examples = None
 
     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) and we draw inspiration from the teams behind [Rodin](https://hyperhuman.deemos.com/rodin), [Tripo](https://www.tripo3d.ai/app/home), and [Dora](https://github.com/Seed3D/Dora).
-        - **Normal Estimation:** Our Normal Estimation Model builds on the leading normal estimation research such as [StableNormal](https://github.com/hugoycj/StableNormal) and [GenPercept](https://github.com/aim-uofa/GenPercept).
-        
-        **Your contributions and collaboration push the boundaries of 3D modeling!**
+        - **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].
+        - **Normal Estimation:** Builds on [StableNormal] and [GenPercept].
         """
     )
 
+# ---------------------------------------------------------------------------
+# Entry
+# ---------------------------------------------------------------------------
 if __name__ == "__main__":
-    # Download and cache the weights
-    cache_weights(WEIGHTS_DIR)
-
-    hi3dgen_pipeline = Hi3DGenPipeline.from_pretrained("weights/trellis-normal-v0-1")
-    hi3dgen_pipeline.cuda()
-
-    # Initialize normal predictor
-    try:
-        normal_predictor = torch.hub.load(os.path.join(torch.hub.get_dir(), 'hugoycj_StableNormal_main'), "StableNormal_turbo", yoso_version='yoso-normal-v1-8-1', source='local', local_cache_dir='./weights', pretrained=True)
-    except:
-        normal_predictor = torch.hub.load("hugoycj/StableNormal", "StableNormal_turbo", trust_repo=True, yoso_version='yoso-normal-v1-8-1', local_cache_dir='./weights')    
-
-    # Launch the app
-    demo.launch(share=False, server_name="0.0.0.0")
+    # Cache model w