Free memory after inference to reduce potential memory overflow

app.py

@@ -109,14 +109,18 @@ def generate_imgs(x: torch.Tensor,
     #     - torch.Tensor object -> Pil object
     process_img = partial(dinv.utils.plotting.preprocess_img, rescale_mode="clip")
     to_pil = transforms.ToPILImage()
-    x = to_pil(process_img(x)[0].to('cpu'))
-    y = to_pil(process_img(y_plot)[0].to('cpu'))
-    out = to_pil(process_img(out)[0].to('cpu'))
-    out_baseline = to_pil(process_img(out_baseline)[0].to('cpu'))
+    x_pil = to_pil(process_img(x)[0].to('cpu'))
+    y_pil = to_pil(process_img(y_plot)[0].to('cpu'))
+    out_pil = to_pil(process_img(out)[0].to('cpu'))
+    out_baseline_pil = to_pil(process_img(out_baseline)[0].to('cpu'))
 
+
+    # Free memory
+    del x, y, out, out_baseline, y_plot
+    torch.cuda.empty_cache()
     print(torch.cuda.memory_allocated() / 1024**2)
-    return x, y, out, out_baseline, physics.display_saved_params(), metrics_y, metrics_out, metrics_out_baseline
 
+    return x_pil, y_pil, out_pil, out_baseline_pil, physics.display_saved_params(), metrics_y, metrics_out, metrics_out_baseline
 
 get_dataset_on_DEVICE_STR = partial(EvalDataset, device_str=DEVICE_STR)
 get_physics_on_DEVICE_STR = partial(PhysicsWithGenerator, device_str=DEVICE_STR)
@@ -175,10 +179,10 @@ with gr.Blocks(title=title, theme=gr.themes.Glass()) as interface:
 
         # 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)
+            gt_img = gr.Image(label="Ground-truth image", interactive=True, key='gt_img')
+            observed_img = gr.Image(label="Observed image", interactive=False, key='observed_img')
+            model_a_out = gr.Image(label="RAM output", interactive=False, key='ram_out')
+            model_b_out = gr.Image(label="DPIR output", interactive=False, key='dpir_out')
 
         # Manage datasets and display metric values
         with gr.Row():
@@ -187,16 +191,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=4)
+                idx_slider = gr.Slider(minimum=0, maximum=len(dataset)-1, step=1, label="Sample index", key='idx_slider')
                 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="Observed metric", lines=3, key=5)
+                observed_metrics = gr.Textbox(label="Observed metric", lines=3, key='metrics')
             with gr.Column(scale=1, min_width=160):
-                out_a_metric = gr.Textbox(label="RAM output metrics", lines=3, key=6)
+                out_a_metric = gr.Textbox(label="RAM output metrics", lines=3, key='ram_metrics')
             with gr.Column(scale=1, min_width=160):
-                out_b_metric = gr.Textbox(label="DPIR output metrics", lines=3, key=7)
+                out_b_metric = gr.Textbox(label="DPIR output metrics", lines=3, key='dpir_metrics')
 
         # Manage physics
         with gr.Row():
@@ -204,7 +208,7 @@ 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, key=8)
+                use_generator_button = gr.Checkbox(label="Generate physics parameters during inference", value=True, key='use_gen')
             with gr.Column(scale=1):
                 with gr.Row():
                     key_selector = gr.Dropdown(choices=list(physics.saved_params["updatable_params"].keys()),