change model cache dir and add notification for multi-GPU

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
 ```
 

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):

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 = {}