Merge pull request #60 from catid/main

Add multi-gpu support, model cache, fix README and deprecated containers

README.md

@@ -22,7 +22,7 @@ https://github.com/TencentARC/InstantMesh/assets/20635237/dab3511e-e7c6-4c0b-bab
 - [x] Release inference and training code.
 - [x] Release model weights.
 - [x] Release huggingface gradio demo. Please try it at [demo](https://huggingface.co/spaces/TencentARC/InstantMesh) link.
-- [ ] Add support to low-memory GPU environment.
+- [x] Add support to low-memory GPU environment.
 - [ ] Add support to more multi-view diffusion models.
 
 # ⚙️ Dependencies and Installation
@@ -66,14 +66,14 @@ By default, we use the `instant-mesh-large` reconstruction model variant.
 
 ## Start a local gradio demo
 
-To start a gradio demo in your local machine, simply running:
+To start a gradio demo in your local machine, simply run:
 ```bash
-python app.py
+CUDA_VISIBLE_DEVICES=0 python app.py
 ```
 
 ## Running with command line
 
-To generate 3D meshes from images via command line, simply running:
+To generate 3D meshes from images via command line, simply run:
 ```bash
 python run.py configs/instant-mesh-large.yaml examples/hatsune_miku.png --save_video
 ```

app.py

@@ -23,6 +23,16 @@ 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')
+else:
+    device0 = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    device1 = device0
+
+# Define the cache directory for model files
+model_cache_dir = './models/'
+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):
     """
@@ -76,29 +86,30 @@ pipeline = DiffusionPipeline.from_pretrained(
     "sudo-ai/zero123plus-v1.2", 
     custom_pipeline="zero123plus",
     torch_dtype=torch.float16,
+    cache_dir=model_cache_dir
 )
 pipeline.scheduler = EulerAncestralDiscreteScheduler.from_config(
     pipeline.scheduler.config, timestep_spacing='trailing'
 )
 
 # load custom white-background UNet
-unet_ckpt_path = hf_hub_download(repo_id="TencentARC/InstantMesh", filename="diffusion_pytorch_model.bin", repo_type="model")
+unet_ckpt_path = hf_hub_download(repo_id="TencentARC/InstantMesh", filename="diffusion_pytorch_model.bin", repo_type="model", cache_dir=model_cache_dir)
 state_dict = torch.load(unet_ckpt_path, map_location='cpu')
 pipeline.unet.load_state_dict(state_dict, strict=True)
 
-pipeline = pipeline.to(device)
+pipeline = pipeline.to(device0)
 
 # load reconstruction model
 print('Loading reconstruction model ...')
-model_ckpt_path = hf_hub_download(repo_id="TencentARC/InstantMesh", filename="instant_mesh_large.ckpt", repo_type="model")
+model_ckpt_path = hf_hub_download(repo_id="TencentARC/InstantMesh", filename="instant_mesh_large.ckpt", repo_type="model", cache_dir=model_cache_dir)
 model = instantiate_from_config(model_config)
 state_dict = torch.load(model_ckpt_path, map_location='cpu')['state_dict']
 state_dict = {k[14:]: v for k, v in state_dict.items() if k.startswith('lrm_generator.') and 'source_camera' not in k}
 model.load_state_dict(state_dict, strict=True)
 
-model = model.to(device)
+model = model.to(device1)
 if IS_FLEXICUBES:
-    model.init_flexicubes_geometry(device, fovy=30.0)
+    model.init_flexicubes_geometry(device1, fovy=30.0)
 model = model.eval()
 
 print('Loading Finished!')
@@ -124,7 +135,7 @@ def generate_mvs(input_image, sample_steps, sample_seed):
     seed_everything(sample_seed)
     
     # sampling
-    generator = torch.Generator(device=device)
+    generator = torch.Generator(device=device0)
     z123_image = pipeline(
         input_image, 
         num_inference_steps=sample_steps, 
@@ -172,11 +183,11 @@ def make3d(images):
     images = torch.from_numpy(images).permute(2, 0, 1).contiguous().float()     # (3, 960, 640)
     images = rearrange(images, 'c (n h) (m w) -> (n m) c h w', n=3, m=2)        # (6, 3, 320, 320)
 
-    input_cameras = get_zero123plus_input_cameras(batch_size=1, radius=4.0).to(device)
+    input_cameras = get_zero123plus_input_cameras(batch_size=1, radius=4.0).to(device1)
     render_cameras = get_render_cameras(
-        batch_size=1, radius=4.5, elevation=20.0, is_flexicubes=IS_FLEXICUBES).to(device)
+        batch_size=1, radius=4.5, elevation=20.0, is_flexicubes=IS_FLEXICUBES).to(device1)
 
-    images = images.unsqueeze(0).to(device)
+    images = images.unsqueeze(0).to(device1)
     images = v2.functional.resize(images, (320, 320), interpolation=3, antialias=True).clamp(0, 1)
 
     mesh_fpath = tempfile.NamedTemporaryFile(suffix=f".obj", delete=False).name

docker/Dockerfile

@@ -1,8 +1,10 @@
 # get the development image from nvidia cuda 12.1
-FROM nvidia/cuda:12.1.0-cudnn8-devel-ubuntu20.04
+FROM nvidia/cuda:12.4.1-runtime-ubuntu22.04
 
-LABEL name="instantmesh" \
-    maintainer="instantmesh"
+LABEL name="instantmesh" maintainer="instantmesh"
+
+# Add a volume for downloaded models
+VOLUME /workspace/models
 
 # create workspace folder and set it as working directory
 RUN mkdir -p /workspace/instantmesh
@@ -12,12 +14,12 @@ WORKDIR /workspace
 ENV DEBIAN_FRONTEND=noninteractive
 RUN apt-get update && \
     apt-get install -y tzdata && \
-    ln -fs /usr/share/zoneinfo/Asia/Shanghai /etc/localtime && \
+    ln -fs /usr/share/zoneinfo/America/Chicago /etc/localtime && \ 
     dpkg-reconfigure --frontend noninteractive tzdata
 
 # update package lists and install git, wget, vim, libegl1-mesa-dev, and libglib2.0-0
 RUN apt-get update && \
-    apt-get install -y git wget vim libegl1-mesa-dev libglib2.0-0 unzip
+    apt-get install -y build-essential git wget vim libegl1-mesa-dev libglib2.0-0 unzip
 
 # install conda
 RUN wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh && \
@@ -36,7 +38,7 @@ RUN conda create -n instantmesh python=3.10 && echo "source activate instantmesh
 ENV PATH /workspace/miniconda3/envs/instantmesh/bin:$PATH
 
 RUN conda install Ninja
-RUN conda install cuda -c nvidia/label/cuda-12.1.0 -y
+RUN conda install cuda -c nvidia/label/cuda-12.4.1 -y
 
 RUN pip install torch==2.1.0 torchvision==0.16.0 torchaudio==2.1.0 --index-url https://download.pytorch.org/whl/cu121
 RUN pip install xformers==0.0.22.post7

docker/README.md

@@ -1,13 +1,28 @@
 # Docker setup
 
-This docker setup is tested on WSL(Ubuntu).
+This docker setup is tested on Ubuntu.
 
 make sure you are under directory yourworkspace/instantmesh/
 
-run
+Build docker image:
 
-`docker build -t instantmesh/deploy:cuda12.1 -f docker/Dockerfile .`
+```bash
+docker build -t instantmesh -f docker/Dockerfile .
+```
 
-then run
+Run docker image with a local model cache (so it is fast when container is started next time):
 
-`docker run --gpus all -it instantmesh/deploy:cuda12.1`
+```bash
+mkdir -p $HOME/models/
+export MODEL_DIR=$HOME/models/
+
+docker run -it -p 43839:43839 --platform=linux/amd64 --gpus all -v $MODEL_DIR:/workspace/instantmesh/models instantmesh
+```
+
+To use specific GPUs:
+
+```bash
+docker run -it -p 43839:43839 --platform=linux/amd64 --gpus '"device=0,1"' -v $MODEL_DIR:/workspace/instantmesh/models instantmesh
+```
+
+Navigate to `http://localhost:43839` to use the demo.