separate load data and load model, set debug False

scripts/scripts_utils/plotly_interface.py

@@ -244,7 +244,7 @@ def prediction_plot(
     n_samples: int = 1,
     use_biaser: bool = True,
 ) -> go.Figure:
-    range_radius = 70
+    range_radius = 50
     if use_biaser:
         risk_level = float(risk_level)
     else:
@@ -262,8 +262,8 @@ def prediction_plot(
         ),
         title_text="Road Scene",
         hovermode="closest",
-        width=1200,
-        height=600,
+        width=600,
+        height=300,
         updatemenus=[
             dict(
                 type="buttons",
@@ -332,8 +332,7 @@ def update_figure(
 
     return fig
 
-def load_from_huggingface(model_source: str = "TRI-ML/risk_biased_model",  data_source: str = "jmercat/risk_biased_dataset", config_name: str="learning_config.py", checkpoint_name: str = "last.ckpt", device: str = "cpu") -> Tuple[LitTrajectoryPredictor, Dataset]:
-    dataset = load_dataset(data_source, split="test")
+def load_predictor_from_hf(model_source: str = "TRI-ML/risk_biased_model", config_name: str="learning_config.py", checkpoint_name: str = "last.ckpt", device: str = "cpu") -> Tuple[LitTrajectoryPredictor, Dataset]:
     config_file = hf_hub_download(model_source, filename=config_name, use_auth_token=os.getenv('SECRET_AUTH_TOKEN'))
     ckpt = torch.load(hf_hub_download(model_source, filename=checkpoint_name, use_auth_token=os.getenv('SECRET_AUTH_TOKEN')), map_location="cpu")
     cfg = Config.fromfile(config_file)
@@ -342,23 +341,31 @@ def load_from_huggingface(model_source: str = "TRI-ML/risk_biased_model",  data_
     predictor.eval()
     predictor = predictor.to(device)
     
-    return predictor, dataset
+    return predictor
+
+def load_dataset_from_hf(data_source: str = "jmercat/risk_biased_dataset") -> Dataset:
+    dataset = load_dataset(data_source, split="test")
+    return dataset
 
 def main(load_from=None, cfg_path=None):
     # Define the device to use
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print("Getting dataset")
+    dataset = load_dataset_from_hf()
+        
+    if load_from is not None:
+        cfg = Config.fromfile(cfg_path)
+        predictor = get_predictor(cfg, WaymoDataloaders.unnormalize_trajectory)
+        predictor = load_weights(predictor, torch.load(load_from, map_location="cpu"))
+    else:
+        print("Getting model.")
+        predictor = load_predictor_from_hf(device=device)
+
+    ui_update_fn = partial(update_figure, predictor, dataset)
     
     # Do the same thing as above but using the gradio blocks API
     with gr.Blocks() as interface:
-        
-        predictor, dataset = load_from_huggingface(device=device)
-        
-        if load_from is not None:
-            cfg = Config.fromfile(cfg_path)
-            predictor = get_predictor(cfg, WaymoDataloaders.unnormalize_trajectory)
-            predictor = load_weights(predictor, torch.load(load_from, map_location="cpu"))
 
-        ui_update_fn = partial(update_figure, predictor, dataset)
         gr.Markdown(
         """
             # Risk-Aware Prediction
@@ -391,7 +398,7 @@ def main(load_from=None, cfg_path=None):
         # n_samples.change(ui_update_fn, inputs=[index, risk_level, n_samples, image], outputs=image)
         button.click(ui_update_fn, inputs=[index, risk_level, n_samples, image], outputs=image)
 
-        interface.launch(debug=True)
+        interface.launch(debug=False)
 
 
 if __name__ == "__main__":