Update app.py

app.py

@@ -346,4 +346,39 @@ with gr.Blocks(css="footer {visibility: hidden}") as demo:
 # Entry
 # ---------------------------------------------------------------------------
 if __name__ == "__main__":
-    # Cache model w
+    # Cache model weights locally
+    cache_weights(WEIGHTS_DIR)
+
+    # Load pipeline on CPU
+    hi3dgen_pipeline = Hi3DGenPipeline.from_pretrained("weights/trellis-normal-v0-1")
+    try:
+        hi3dgen_pipeline.to("cpu")
+    except Exception:
+        pass  # some pipelines may not implement .to
+
+    # Initialize normal predictor (CPU)
+    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 Exception:
+        normal_predictor = torch.hub.load(
+            "hugoycj/StableNormal",
+            "StableNormal_turbo",
+            trust_repo=True,
+            yoso_version='yoso-normal-v1-8-1',
+            local_cache_dir='./weights'
+        )
+
+    try:
+        normal_predictor.to("cpu")
+    except Exception:
+        pass
+
+    # Launch the Gradio app
+    demo.launch(share=False, server_name="0.0.0.0")