auto download model weights

run.py

@@ -9,6 +9,7 @@
 from omegaconf import OmegaConf
 from einops import rearrange, repeat
 from tqdm import tqdm
+from huggingface_hub import hf_hub_download
 from diffusers import DiffusionPipeline, EulerAncestralDiscreteScheduler
 
 from src.utils.train_util import instantiate_from_config
@@ -106,15 +107,23 @@ def render_frames(model, planes, render_cameras, render_size=512, chunk_size=1,
 
 # load custom white-background UNet
 print('Loading custom white-background unet ...')
-state_dict = torch.load(infer_config.unet_path, map_location='cpu')
+if os.path.exists(infer_config.unet_path):
+    unet_ckpt_path = infer_config.unet_path
+else:
+    unet_ckpt_path = hf_hub_download(repo_id="TencentARC/InstantMesh", filename="diffusion_pytorch_model.bin", repo_type="model")
+state_dict = torch.load(unet_ckpt_path, map_location='cpu')
 pipeline.unet.load_state_dict(state_dict, strict=True)
 
 pipeline = pipeline.to(device)
 
 # load reconstruction model
 print('Loading reconstruction model ...')
 model = instantiate_from_config(model_config)
-state_dict = torch.load(infer_config.model_path, map_location='cpu')['state_dict']
+if os.path.exists(infer_config.model_path):
+    model_ckpt_path = infer_config.model_path
+else:
+    model_ckpt_path = hf_hub_download(repo_id="TencentARC/InstantMesh", filename=f"{config_name.replace('-', '_')}.ckpt", repo_type="model")
+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.')}
 model.load_state_dict(state_dict, strict=True)