[Fix] support inferencer in pip install

open-mmlab · May 26, 2023 · 099544c · 099544c
1 parent 3cdc213
commit 099544c
Showing 1 changed file with 14 additions and 6 deletions.
diff --git a/mmagic/apis/mmagic_inferencer.py b/mmagic/apis/mmagic_inferencer.py
@@ -1,5 +1,4 @@
 # Copyright (c) OpenMMLab. All rights reserved.
-import os
 import os.path as osp
 import warnings
 from typing import Dict, List, Optional, Union
@@ -155,8 +154,13 @@ def _get_inferencer_kwargs(self, model_name: Optional[str],
                 setting_to_use = model_setting
             config_dir = cfgs['settings'][setting_to_use]['Config']
             config_dir = config_dir[config_dir.find('configs'):]
-            kwargs['config'] = os.path.join(
-                osp.dirname(__file__), '..', '..', config_dir)
+            if osp.exists(
+                    osp.join(osp.dirname(__file__), '..', '..', config_dir)):
+                kwargs['config'] = osp.join(
+                    osp.dirname(__file__), '..', '..', config_dir)
+            else:
+                kwargs['config'] = osp.join(
+                    osp.dirname(__file__), '..', '.mim', config_dir)
             if 'Weights' in cfgs['settings'][setting_to_use].keys():
                 kwargs['ckpt'] = cfgs['settings'][setting_to_use]['Weights']
 
@@ -233,9 +237,13 @@ def get_model_config(self, model_name: str) -> Dict:
     @staticmethod
     def init_inference_supported_models_cfg() -> None:
         if not MMagicInferencer.inference_supported_models_cfg_inited:
-            all_cfgs_dir = osp.join(
-                osp.dirname(__file__), '..', '..', 'configs')
-
+            if osp.exists(
+                    osp.join(osp.dirname(__file__), '..', '..', 'configs')):
+                all_cfgs_dir = osp.join(
+                    osp.dirname(__file__), '..', '..', 'configs')
+            else:
+                all_cfgs_dir = osp.join(
+                    osp.dirname(__file__), '..', '.mim', 'configs')
             for model_name in MMagicInferencer.inference_supported_models:
                 meta_file_dir = osp.join(all_cfgs_dir, model_name,
                                          'metafile.yml')