From c030efb5e92a95780907f2a7bcc532ea0a066b85 Mon Sep 17 00:00:00 2001 From: catid Date: Fri, 3 May 2024 04:47:01 +0000 Subject: [PATCH 1/7] add 2gpu support --- app.py | 21 ++++++++++++++------- 1 file changed, 14 insertions(+), 7 deletions(-) diff --git a/app.py b/app.py index a805d39..9baab6e 100644 --- a/app.py +++ b/app.py @@ -23,6 +23,13 @@ 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 + def get_render_cameras(batch_size=1, M=120, radius=2.5, elevation=10.0, is_flexicubes=False): """ @@ -86,7 +93,7 @@ unet_ckpt_path = hf_hub_download(repo_id="TencentARC/InstantMesh", filename="dif 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 ...') @@ -96,9 +103,9 @@ 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 +131,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 +179,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 From 5b18954620b2c49e11eb0378a25b99488388e461 Mon Sep 17 00:00:00 2001 From: catid Date: Fri, 3 May 2024 17:42:56 +0000 Subject: [PATCH 2/7] add models volume cache --- app.py | 9 ++++++--- docker/Dockerfile | 5 ++++- 2 files changed, 10 insertions(+), 4 deletions(-) diff --git a/app.py b/app.py index 9baab6e..66b883c 100644 --- a/app.py +++ b/app.py @@ -30,6 +30,9 @@ 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): """ @@ -89,7 +92,7 @@ pipeline.scheduler = EulerAncestralDiscreteScheduler.from_config( ) # 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) @@ -97,7 +100,7 @@ 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} @@ -375,4 +378,4 @@ with gr.Blocks() as demo: ) demo.queue(max_size=10) -demo.launch(server_name="0.0.0.0", server_port=43839) +demo.launch(server_name="0.0.0.0", server_port=43839, share=True) diff --git a/docker/Dockerfile b/docker/Dockerfile index f65d5d2..a7c2754 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -51,5 +51,8 @@ RUN pip install -r requirements.txt COPY . /workspace/instantmesh +# Add a volume for downloaded models +VOLUME /workspace/models + # Run the command when the container starts -CMD ["python", "app.py"] \ No newline at end of file +CMD ["python", "app.py"] From 8d0ee3eac5c915296fef9a0299f245bf77f5e17f Mon Sep 17 00:00:00 2001 From: catid Date: Fri, 3 May 2024 18:13:23 +0000 Subject: [PATCH 3/7] Update from deprecated version of Nvidia container --- docker/Dockerfile | 17 ++++++++--------- 1 file changed, 8 insertions(+), 9 deletions(-) diff --git a/docker/Dockerfile b/docker/Dockerfile index a7c2754..5f12169 100644 --- a/docker/Dockerfile +++ b/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 @@ -51,8 +53,5 @@ RUN pip install -r requirements.txt COPY . /workspace/instantmesh -# Add a volume for downloaded models -VOLUME /workspace/models - # Run the command when the container starts CMD ["python", "app.py"] From 697c96e7504378213f4aa8fd02298d767822aba0 Mon Sep 17 00:00:00 2001 From: catid Date: Fri, 3 May 2024 18:25:14 +0000 Subject: [PATCH 4/7] no share] --- app.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/app.py b/app.py index 66b883c..032f112 100644 --- a/app.py +++ b/app.py @@ -378,4 +378,4 @@ with gr.Blocks() as demo: ) demo.queue(max_size=10) -demo.launch(server_name="0.0.0.0", server_port=43839, share=True) +demo.launch(server_name="0.0.0.0", server_port=43839) From 368393a745841245bf75dbcfe16bd9335e25691e Mon Sep 17 00:00:00 2001 From: catid Date: Fri, 3 May 2024 19:00:52 +0000 Subject: [PATCH 5/7] model cache --- README.md | 2 +- app.py | 1 + docker/README.md | 19 ++++++++++++++----- 3 files changed, 16 insertions(+), 6 deletions(-) diff --git a/README.md b/README.md index 59887cc..57f0ede 100644 --- a/README.md +++ b/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 diff --git a/app.py b/app.py index 032f112..c74140d 100644 --- a/app.py +++ b/app.py @@ -86,6 +86,7 @@ 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' diff --git a/docker/README.md b/docker/README.md index 3062135..68646d0 100644 --- a/docker/README.md +++ b/docker/README.md @@ -1,13 +1,22 @@ # 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` \ No newline at end of file +```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 +``` + +Navigate to `http://localhost:43839` to use the demo. From 3fb546398c7c247b07ca52a9ae59ed0e759ca169 Mon Sep 17 00:00:00 2001 From: catid Date: Fri, 3 May 2024 19:11:18 +0000 Subject: [PATCH 6/7] document gpu selection --- README.md | 2 +- docker/README.md | 6 ++++++ 2 files changed, 7 insertions(+), 1 deletion(-) diff --git a/README.md b/README.md index 57f0ede..6e54f41 100644 --- a/README.md +++ b/README.md @@ -68,7 +68,7 @@ By default, we use the `instant-mesh-large` reconstruction model variant. To start a gradio demo in your local machine, simply running: ```bash -python app.py +CUDA_VISIBLE_DEVICES=0 python app.py ``` ## Running with command line diff --git a/docker/README.md b/docker/README.md index 68646d0..9b2ee12 100644 --- a/docker/README.md +++ b/docker/README.md @@ -19,4 +19,10 @@ 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. From b671a6802b3401f005f067d0cc151fee8183140d Mon Sep 17 00:00:00 2001 From: catid Date: Fri, 3 May 2024 19:18:41 +0000 Subject: [PATCH 7/7] correct typo --- README.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/README.md b/README.md index 6e54f41..71e2486 100644 --- a/README.md +++ b/README.md @@ -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 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 ```