add conditional models

open-mmlab · Feb 21, 2023 · a0d47c2 · a0d47c2
1 parent a17e934
commit a0d47c2
Show file tree

Hide file tree

Showing 3 changed files with 5 additions and 2 deletions.
diff --git a/mmedit/apis/inferencers/base_mmedit_inferencer.py b/mmedit/apis/inferencers/base_mmedit_inferencer.py
@@ -12,7 +12,7 @@
 from torchvision import utils
 
 from mmedit.registry import MODELS
-from mmedit.utils import ConfigType, SampleList
+from mmedit.utils import ConfigType, SampleList, register_all_modules
 from .inference_functions import set_random_seed
 
 InputType = Union[str, int, np.ndarray]
@@ -54,6 +54,7 @@ def __init__(self,
             device = torch.device(
                 'cuda' if torch.cuda.is_available() else 'cpu')
         self.device = device
+        register_all_modules()
         super().__init__(config, ckpt, device)
 
         self._init_extra_parameters(extra_parameters)

diff --git a/mmedit/apis/inferencers/conditional_inferencer.py b/mmedit/apis/inferencers/conditional_inferencer.py
@@ -20,7 +20,7 @@ class ConditionalInferencer(BaseMMEditInferencer):
         visualize=['result_out_dir'],
         postprocess=[])
 
-    extra_parameters = dict(num_batches=4, sample_model='ema')
+    extra_parameters = dict(num_batches=4, sample_model='orig')
 
     def preprocess(self, label: InputsType) -> Dict:
         """Process the inputs into a model-feedable format.

diff --git a/mmedit/edit.py b/mmedit/edit.py
@@ -44,6 +44,8 @@ class MMEdit:
 
         # conditional models
         'biggan',
+        'sngan_proj',
+        'sagan',
 
         # unconditional models
         'styleganv1',