Update lrm/inferrer.py

lrm/inferrer.py

@@ -45,14 +45,19 @@ class LRMInferrer:
 
     def _load_checkpoint(self, model_name: str, cache_dir = './.cache'):
         # download checkpoint if not exists
-        if not os.path.exists(cache_dir):
-            os.makedirs(cache_dir, exist_ok=True)
-        if not os.path.exists(os.path.join(cache_dir, f'{model_name}.pth')):
-            # TODO: on-the-fly download not supported yet, plz download manually
+        local_dir = os.path.join(cache_dir, model_name)
+        if not os.path.exists(local_dir):
+            os.makedirs(local_dir, exist_ok=True)
+        if not os.path.exists(os.path.join(local_dir, f'model.pth')):
             # os.system(f'wget -O {os.path.join(cache_dir, f"{model_name}.pth")} https://zxhezexin.com/modelzoo/openlrm/{model_name}.pth')
-            raise FileNotFoundError(f"Checkpoint {model_name} not found in {cache_dir}")
-        local_path = os.path.join(cache_dir, f'{model_name}.pth')
-        checkpoint = torch.load(local_path, map_location=self.device)
+            # raise FileNotFoundError(f"Checkpoint {model_name} not found in {cache_dir}")
+            from huggingface_hub import hf_hub_download
+            repo_id = f'zxhezexin/{model_name}'
+            config_path = hf_hub_download(repo_id=repo_id, filename='config.json', local_dir=local_dir)
+            model_path = hf_hub_download(repo_id=repo_id, filename=f'model.pth', local_dir=local_dir)
+        else:
+            model_path = os.path.join(local_dir, f'model.pth')
+        checkpoint = torch.load(model_path, map_location=self.device)
         return checkpoint
 
     def _build_model(self, model_kwargs, model_weights):
@@ -193,11 +198,11 @@ class LRMInferrer:
 
         image = torch.tensor(np.array(Image.open(source_image))).permute(2, 0, 1).unsqueeze(0) / 255.0
         # if RGBA, blend to RGB
-        print(f"[DEBUG] check 1.")
+        # print(f"[DEBUG] check 1.")
         if image.shape[1] == 4:
             image = image[:, :3, ...] * image[:, 3:, ...] + (1 - image[:, 3:, ...])
             print(f"[DEBUG] image.shape={image.shape} and image[0,0,0,0]={image[0,0,0,0]}")
-        print(f"[DEBUG] check 2.")
+        # print(f"[DEBUG] check 2.")
         image = torch.nn.functional.interpolate(image, size=(source_image_size, source_image_size), mode='bicubic', align_corners=True)
         image = torch.clamp(image, 0, 1)
         results = self.infer_single(
@@ -240,11 +245,11 @@ if __name__ == '__main__':
 
     """
     Example usage:
-    python -m lrm.inferrer --model_name lrm-base-obj-v1 --source_image ./assets/sample_input/owl.png --export_video --export_mesh
+    python -m lrm.inferrer --model_name openlrm-base-obj-1.0 --source_image ./assets/sample_input/owl.png --export_video --export_mesh
     """
 
     parser = argparse.ArgumentParser()
-    parser.add_argument('--model_name', type=str, default='lrm-base-obj-v1')
+    parser.add_argument('--model_name', type=str, default='openlrm-base-obj-1.0')
     parser.add_argument('--source_image', type=str, default='./assets/sample_input/owl.png')
     parser.add_argument('--dump_path', type=str, default='./dumps')
     parser.add_argument('--source_size', type=int, default=-1)