Define global variables across all users -> reduce loading time

app.py

@@ -16,50 +16,52 @@ from factories import PhysicsWithGenerator, EvalModel, BaselineModel, EvalDatase
 
 
 ### Config
-DEVICE_STR = 'cuda' if torch.cuda.is_available() else 'cpu'            # run model inference on NVIDIA gpu
-torch.set_grad_enabled(False)  # stops tracking values for gradients
+# run model inference on NVIDIA gpu if available
+DEVICE_STR = 'cuda' if torch.cuda.is_available() else 'cpu'
+# stops tracking values for gradients
+torch.set_grad_enabled(False)
 
 
 ### Gradio Utils
 def generate_imgs_from_user(image,
-                            model: EvalModel, baseline: BaselineModel,
                             physics: PhysicsWithGenerator, use_gen: bool,
+                            baseline: BaselineModel, model: EvalModel,
                             metrics: List[Metric]):
+    # Happens when user image is missing
     if image is None:
         return None, None, None, None, None, None, None, None
 
-    # PIL image -> torch.Tensor
+    # PIL image -> torch.Tensor / (1, C, H, W) / move to DEVICE_STR
     x = transforms.ToTensor()(image).unsqueeze(0).to(DEVICE_STR)
 
-    return generate_imgs(x, model, baseline, physics, use_gen, metrics)
+    return generate_imgs(x, physics, use_gen, baseline, model, metrics)
 
 def generate_imgs_from_dataset(dataset: EvalDataset, idx: int,
-                               model: EvalModel, baseline: BaselineModel,
                                physics: PhysicsWithGenerator, use_gen: bool,
+                               baseline: BaselineModel, model: EvalModel,
                                metrics: List[Metric]):
     ### Load 1 image
     x = dataset[idx]    # shape : (C, H, W)
     x = x.unsqueeze(0)  # shape : (1, C, H, W)
 
-    return generate_imgs(x, model, baseline, physics, use_gen, metrics)
+    return generate_imgs(x, physics, use_gen, baseline, model, metrics)
 
 def generate_random_imgs_from_dataset(dataset: EvalDataset,
-                                      model: EvalModel,
-                                      baseline: BaselineModel,
                                       physics: PhysicsWithGenerator,
                                       use_gen: bool,
+                                      baseline: BaselineModel,
+                                      model: EvalModel,
                                       metrics: List[Metric]):
     idx = random.randint(0, len(dataset)-1)
     x, y, out, out_baseline, saved_params_str, metrics_y, metrics_out, metrics_out_baseline = generate_imgs_from_dataset(
-        dataset, idx, model, baseline, physics, use_gen, metrics
+        dataset, idx, physics, use_gen, baseline, model, metrics
         )
     return idx, x, y, out, out_baseline, saved_params_str, metrics_y, metrics_out, metrics_out_baseline
     
 def generate_imgs(x: torch.Tensor,
-                  model: EvalModel, baseline: BaselineModel,
                   physics: PhysicsWithGenerator, use_gen: bool,
+                  baseline: BaselineModel, model: EvalModel,
                   metrics: List[Metric]):
-
     ### Compute y
     y = physics(x, use_gen)  # possible reduction in img shape due to Blurring
 
@@ -114,11 +116,9 @@ def generate_imgs(x: torch.Tensor,
     return x, y, out, out_baseline, physics.display_saved_params(), metrics_y, metrics_out, metrics_out_baseline
 
 
-get_list_metrics_on_DEVICE_STR = partial(Metric.get_list_metrics, device_str=DEVICE_STR)
-get_eval_model_on_DEVICE_STR = partial(EvalModel, device_str=DEVICE_STR)
-get_baseline_model_on_DEVICE_STR = partial(BaselineModel, device_str=DEVICE_STR)
 get_dataset_on_DEVICE_STR = partial(EvalDataset, device_str=DEVICE_STR)
 get_physics_on_DEVICE_STR = partial(PhysicsWithGenerator, device_str=DEVICE_STR)
+get_baseline_model_on_DEVICE_STR = partial(BaselineModel, device_str=DEVICE_STR)
 
 def get_dataset(dataset_name):
     if dataset_name == 'MRI':
@@ -142,37 +142,42 @@ def get_dataset(dataset_name):
     return dataset, idx, physics, baseline, available_physics
 
 
+# global variables shared by all users
+ram_model = EvalModel("unext_emb_physics_config_C", device_str=DEVICE_STR)
+psnr = Metric.get_list_metrics(["PSNR"], device_str=DEVICE_STR)
+
+generate_imgs_from_user_partial = partial(generate_imgs_from_user, model=ram_model, metrics=psnr)
+generate_imgs_from_dataset_partial = partial(generate_imgs_from_dataset, model=ram_model, metrics=psnr)
+generate_random_imgs_from_dataset_partial = partial(generate_random_imgs_from_dataset, model=ram_model, metrics=psnr)
+
+
 ### Gradio Blocks interface
 
 title = "Inverse problem playground"  # displayed on gradio tab and in the gradio page
 with gr.Blocks(title=title, theme=gr.themes.Glass()) as interface:
     gr.Markdown("## " + title)
 
-    ### DEFAULT VALUES
-    # Issue: giving directly a `torch.nn.module` to `gr.State(...)` since it has __call__ method
-    # Solution: using lambda expression
-    model_a_placeholder = gr.State(lambda: get_eval_model_on_DEVICE_STR("unext_emb_physics_config_C", ""))
-    model_b_placeholder = gr.State(lambda: get_baseline_model_on_DEVICE_STR("DPIR"))
-    metrics_placeholder = gr.State(get_list_metrics_on_DEVICE_STR(["PSNR"]))
-
+    ### USER-SPECIFIC VARIABLES
     dataset_placeholder = gr.State(get_dataset_on_DEVICE_STR("Natural"))
-    physics_placeholder = gr.State(lambda: get_physics_on_DEVICE_STR("MotionBlur_easy"))
     available_physics_placeholder = gr.State(['MotionBlur_easy', 'MotionBlur_medium', 'MotionBlur_hard',
                                               'GaussianBlur_easy', 'GaussianBlur_medium', 'GaussianBlur_hard'])
+    # Issue giving directly a `torch.nn.module` to `gr.State(...)` since it has __call__ method
+    # Solution: using lambda expression
+    physics_placeholder = gr.State(lambda: get_physics_on_DEVICE_STR("MotionBlur_easy"))
+    model_b_placeholder = gr.State(lambda: get_baseline_model_on_DEVICE_STR("DPIR"))
 
 
-    ### LAYOUT
-    # Display images
-    with gr.Row():
-        gt_img = gr.Image(label="Ground-truth image", interactive=True)
-        observed_img = gr.Image(label="Observed image", interactive=False)
-        model_a_out = gr.Image(label="RAM output", interactive=False)
-        model_b_out = gr.Image(label="DPIR output", interactive=False)
-
     @gr.render(inputs=[dataset_placeholder, physics_placeholder, available_physics_placeholder])
     def dynamic_layout(dataset, physics, available_physics):
         ### LAYOUT
 
+        # Display images
+        with gr.Row():
+            gt_img = gr.Image(label="Ground-truth image", interactive=True, key=0)
+            observed_img = gr.Image(label="Observed image", interactive=False, key=1)
+            model_a_out = gr.Image(label="RAM output", interactive=False, key=2)
+            model_b_out = gr.Image(label="DPIR output", interactive=False, key=3)
+
         # Manage datasets and display metric values
         with gr.Row():
             with gr.Column(scale=1, min_width=160):
@@ -180,19 +185,16 @@ with gr.Blocks(title=title, theme=gr.themes.Glass()) as interface:
                 choose_dataset = gr.Radio(choices=EvalDataset.all_datasets,
                                           label="Datasets",
                                           value=dataset.name)
-                idx_slider = gr.Slider(minimum=0, maximum=len(dataset)-1, step=1, label="Sample index", key=0)
+                idx_slider = gr.Slider(minimum=0, maximum=len(dataset)-1, step=1, label="Sample index", key=4)
                 with gr.Row():
                     load_button = gr.Button("Run on index image from dataset", size='md')
                     load_random_button = gr.Button("Run on random image from dataset", size='md')
             with gr.Column(scale=1, min_width=160):
-                observed_metrics = gr.Textbox(label="PSNR Observed",
-                                              lines=1)
+                observed_metrics = gr.Textbox(label="Observed metric", lines=1, key=5)
             with gr.Column(scale=1, min_width=160):
-                out_a_metric = gr.Textbox(label="PSNR RAM output",
-                                          lines=1)
+                out_a_metric = gr.Textbox(label="RAM output metrics", lines=1, key=6)
             with gr.Column(scale=1, min_width=160):
-                out_b_metric = gr.Textbox(label="PSNR DPIR",
-                                          lines=1)
+                out_b_metric = gr.Textbox(label="DPIR output metrics", lines=1, key=7)
 
         # Manage physics
         with gr.Row():
@@ -200,19 +202,18 @@ with gr.Blocks(title=title, theme=gr.themes.Glass()) as interface:
                 choose_physics = gr.Radio(choices=available_physics,
                                           label="Physics",
                                           value=physics.name)
-                use_generator_button = gr.Checkbox(label="Generate physics parameters during inference", value=True)
+                use_generator_button = gr.Checkbox(label="Generate physics parameters during inference", value=True, key=8)
             with gr.Column(scale=1):
                 with gr.Row():
                     key_selector = gr.Dropdown(choices=list(physics.saved_params["updatable_params"].keys()),
-                                               label="Updatable Parameter Key")
+                                               label="Updatable Key")
                     value_text = gr.Textbox(label="Update Value")
-                update_button = gr.Button("Manually update parameter value")
+                update_button = gr.Button("Manually update parameter value", size='md')
             with gr.Column(scale=2):
                 physics_params = gr.Textbox(label="Physics parameters",
                                             lines=5,
                                             value=physics.display_saved_params())  
 
-
         ### Event listeners
 
         choose_dataset.change(fn=get_dataset,
@@ -223,34 +224,28 @@ with gr.Blocks(title=title, theme=gr.themes.Glass()) as interface:
                               outputs=[physics_placeholder])
         update_button.click(fn=physics.update_and_display_params,
                             inputs=[key_selector, value_text], outputs=physics_params)
-        run_button.click(fn=generate_imgs_from_user,
+        run_button.click(fn=generate_imgs_from_user_partial,
                          inputs=[gt_img,
-                                 model_a_placeholder,
-                                 model_b_placeholder,
                                  physics_placeholder,
                                  use_generator_button,
-                                 metrics_placeholder],
+                                 model_b_placeholder],
                          outputs=[gt_img, observed_img, model_a_out, model_b_out,
                                   physics_params, observed_metrics, out_a_metric, out_b_metric])
-        load_button.click(fn=generate_imgs_from_dataset,
+        load_button.click(fn=generate_imgs_from_dataset_partial,
                           inputs=[dataset_placeholder,
                                   idx_slider,
-                                  model_a_placeholder,
-                                  model_b_placeholder,
                                   physics_placeholder,
                                   use_generator_button,
-                                  metrics_placeholder],
+                                  model_b_placeholder],
                           outputs=[gt_img, observed_img, model_a_out, model_b_out,
                                    physics_params, observed_metrics, out_a_metric, out_b_metric])
-        load_random_button.click(fn=generate_random_imgs_from_dataset,
+        load_random_button.click(fn=generate_random_imgs_from_dataset_partial,
                                  inputs=[dataset_placeholder,
-                                         model_a_placeholder,
-                                         model_b_placeholder,
                                          physics_placeholder,
                                          use_generator_button,
-                                         metrics_placeholder],
+                                         model_b_placeholder],
                                  outputs=[idx_slider, gt_img, observed_img, model_a_out, model_b_out,
                                           physics_params, observed_metrics, out_a_metric, out_b_metric])
 
 
-interface.launch()
+interface.launch()