From f71a756b3534183b344cf66567008cf92a41fc91 Mon Sep 17 00:00:00 2001
From: bluestyle97 <bluestyle928@gmail.com>
Date: Sun, 5 May 2024 17:14:27 +0800
Subject: [PATCH] change model cache dir and add notification for multi-GPU

---
 README.md         | 5 +++++
 app.py            | 3 ++-
 src/model_mesh.py | 2 +-
 3 files changed, 8 insertions(+), 2 deletions(-)

diff --git a/README.md b/README.md
index 71e2486..6826d36 100644
--- a/README.md
+++ b/README.md
@@ -68,6 +68,11 @@ By default, we use the `instant-mesh-large` reconstruction model variant.
 
 To start a gradio demo in your local machine, simply run:
 ```bash
+python app.py
+```
+
+If you have multiple GPUs in your machine, the demo app will run on two GPUs automatically to save memory. You can also force it to run on a single GPU:
+```bash
 CUDA_VISIBLE_DEVICES=0 python app.py
 ```
 
diff --git a/app.py b/app.py
index c74140d..2f0fa3d 100644
--- a/app.py
+++ b/app.py
@@ -23,6 +23,7 @@ from src.utils.infer_util import remove_background, resize_foreground, images_to
 import tempfile
 from huggingface_hub import hf_hub_download
 
+
 if torch.cuda.is_available() and torch.cuda.device_count() >= 2:
     device0 = torch.device('cuda:0')
     device1 = torch.device('cuda:1')
@@ -31,7 +32,7 @@ else:
     device1 = device0
 
 # Define the cache directory for model files
-model_cache_dir = './models/'
+model_cache_dir = './ckpts/'
 os.makedirs(model_cache_dir, exist_ok=True)
 
 def get_render_cameras(batch_size=1, M=120, radius=2.5, elevation=10.0, is_flexicubes=False):
diff --git a/src/model_mesh.py b/src/model_mesh.py
index 1de154f..22bfd4d 100644
--- a/src/model_mesh.py
+++ b/src/model_mesh.py
@@ -274,7 +274,7 @@ class MVRecon(pl.LightningModule):
 
         loss_reg = sdf_reg_loss_entropy + flexicubes_surface_reg + flexicubes_weights_reg
 
-        loss = loss_mse + loss_lpips + loss_mask + loss_normal + loss_reg
+        loss = loss_mse + loss_lpips + loss_mask + loss_depth + loss_normal + loss_reg
 
         prefix = 'train'
         loss_dict = {}