update

factories.py

@@ -10,30 +10,6 @@ from datasets import Preprocessed_fastMRI, Preprocessed_LIDCIDRI, LsdirMiniDatas
 from model_factory import get_model
 from physics.blur_generator import GaussianBlurGenerator
 
-DEFAULT_MODEL_PARAMS = {
-    "in_channels": [1, 2, 3],
-    "grayscale": False,
-    "conv_type": "base",
-    "pool_type": "base",
-    "layer_scale_init_value": 1e-6,
-    "init_type": "ortho",
-    "gain_init_conv": 1.0,
-    "gain_init_linear": 1.0,
-    "drop_prob": 0.0,
-    "replk": False,
-    "mult_fact": 4,
-    "antialias": "gaussian",
-    "nc_base": 64,
-    "cond_type": "base",
-    "blind": False,
-    "pretrained_pth": None,
-    "N": 2,
-    "c_mult": 2,
-    "depth_encoding": 2,
-    "relu_in_encoding": False,
-    "skip_in_encoding": True
-}
-
 
 class PhysicsWithGenerator(torch.nn.Module):
     """Interface between Physics, Generator and Gradio."""
@@ -211,14 +187,7 @@ class EvalModel(torch.nn.Module):
         if self.name == "unext_emb_physics_config_C":
             if self.ckpt_pth == "":
                 self.ckpt_pth = "ckpt/ram.pth.tar"
-            self.model = get_model(model_name=self.name,
-                                   device='cpu',
-                                   **DEFAULT_MODEL_PARAMS)
-
-            # load model checkpoint on cpu
-            state_dict = torch.load(self.ckpt_pth, map_location=lambda storage, loc: storage)
-
-            self.model.load_state_dict(state_dict)
+            self.model = get_model()
             self.model.to(device_str)
             self.model.eval()