update

app.py

@@ -172,7 +172,7 @@ def get_dataset(dataset_name):
 
 
 # global variables shared by all users
-ram_model = EvalModel("unext_emb_physics_config_C", device_str=DEVICE_STR)
+ram_model = EvalModel(device_str=DEVICE_STR)
 ram_model.eval()
 psnr = Metric.get_list_metrics(["PSNR"], device_str=DEVICE_STR)
 

factories.py

@@ -170,32 +170,17 @@ class PhysicsWithGenerator(torch.nn.Module):
 
 class EvalModel(torch.nn.Module):
     """Eval model.
-
-    Is there a difference with BaselineModel ?
-        -> BaselineModel should be models that are already trained and will have fixed weights.
-        -> Eval model will change depending on differents checkpoints.
     """
-    all_models = ["unext_emb_physics_config_C"]
 
-    def __init__(self, model_name: str, ckpt_pth: str = "", device_str: str = "cpu") -> None:
+    def __init__(self, device_str: str = "cpu") -> None:
         """Load the model we want to evaluate."""
         super().__init__()
-        self.name = model_name
-        self.ckpt_pth = ckpt_pth
-        if self.name not in self.all_models:
-            raise ValueError(f"{self.name} is unavailable.")
-        if self.name == "unext_emb_physics_config_C":
-            if self.ckpt_pth == "":
-                self.ckpt_pth = "ckpt/ram.pth.tar"
-            self.model = get_model()
-            self.model.to(device_str)
-            self.model.eval()
+        self.name = 'RAM'
+        self.model = get_model()
+        self.model.to(device_str)
+        self.model.eval()
 
     def forward(self, y: torch.Tensor, physics: torch.nn.Module) -> torch.Tensor:
-        physics.noise_model.sigma = torch.nn.Parameter(torch.tensor([1e-06]))
-        physics.noise_model.gain = torch.nn.Parameter(torch.tensor([1e-06]))
-        print('sigma = ', physics.noise_model.sigma)
-        print('gain = ', physics.noise_model.gain)
         return self.model(y, physics=physics)
 
 

model_factory.py

@@ -1,6 +1,8 @@
 import torch
 from models.ram import RAM
 
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
+
 def get_model():
     """
     Load the model.
@@ -12,6 +14,5 @@ def get_model():
     :return: model
     """
     model = RAM()
-    state_dict = torch.load('ckpt/ram.pth.tar')
-    model.load_state_dict(state_dict)
+    model.load_state_dict(torch.load(hf_hub_download(repo_id="mterris/ram", filename="ram.pth.tar"), map_location=device))
     return model